From 85762ba4b990c12653ac1753e1de5b1ffa67a41b Mon Sep 17 00:00:00 2001
From: Dane Everitt <dane@daneeveritt.com>
Date: Sun, 24 Nov 2019 21:31:31 -0800
Subject: [PATCH] Server state handling; auto start servers if they were
 running when the daemon stopped

---
 server/server.go |  3 +-
 wings.go         | 74 +++++++++++++++++++++++++++++++++++++-----------
 2 files changed, 59 insertions(+), 18 deletions(-)

diff --git a/server/server.go b/server/server.go
index 3941026..ffa1ddb 100644
--- a/server/server.go
+++ b/server/server.go
@@ -29,7 +29,7 @@ type Server struct {
 	Suspended bool `json:"suspended"`
 
 	// The power state of the server.
-	State string `json:"state" yaml:"-"`
+	State string `default:"offline" json:"state"`
 
 	// The command that should be used when booting up the server instance.
 	Invocation string `json:"invocation"`
@@ -224,7 +224,6 @@ func FromConfiguration(data []byte, cfg *config.SystemConfiguration) (*Server, e
 		return nil, err
 	}
 
-	s.State = ProcessOfflineState
 	s.Environment = env
 	s.Cache = cache.New(time.Minute*10, time.Minute*15)
 	s.Filesystem = &Filesystem{
diff --git a/wings.go b/wings.go
index b0d22a7..0c97a40 100644
--- a/wings.go
+++ b/wings.go
@@ -5,8 +5,10 @@ import (
 	"flag"
 	"fmt"
 	"github.com/gorilla/websocket"
+	"github.com/pkg/errors"
 	"github.com/pterodactyl/wings/config"
 	"github.com/pterodactyl/wings/server"
+	"github.com/remeh/sizedwaitgroup"
 	"go.uber.org/zap"
 	"net/http"
 )
@@ -68,26 +70,66 @@ func main() {
 		return
 	}
 
+	// Just for some nice log output.
 	for _, s := range servers {
 		zap.S().Infow("loaded configuration for server", zap.String("server", s.Uuid))
-		zap.S().Infow("ensuring envrionment exists", zap.String("server", s.Uuid))
-
-		if err := s.Environment.Create(); err != nil {
-			zap.S().Errorw("failed to create an environment for server", zap.String("server", s.Uuid), zap.Error(err))
-		}
-
-		if r, err := s.Environment.IsRunning(); err != nil {
-			zap.S().Errorw("error checking server environment status", zap.String("server", s.Uuid), zap.Error(err))
-		} else if r {
-			zap.S().Infow("detected server is running, re-attaching to process", zap.String("server", s.Uuid))
-			s.SetState(server.ProcessRunningState)
-			if err := s.Environment.Attach(); err != nil {
-				zap.S().Errorw("error attaching to server environment", zap.String("server", s.Uuid), zap.Error(err))
-				s.SetState(server.ProcessOfflineState)
-			}
-		}
 	}
 
+	// Create a new WaitGroup that limits us to 4 servers being bootstrapped at a time
+	// on Wings. This allows us to ensure the environment exists, write configurations,
+	// and reboot processes without causing a slow-down due to sequential booting.
+	wg := sizedwaitgroup.New(4)
+
+	for _, serv := range servers {
+		go func(s *server.Server) {
+			defer wg.Done()
+
+			// Create a server environment if none exists currently. This allows us to recover from Docker
+			// being reinstalled on the host system for example.
+			zap.S().Infow("ensuring envrionment exists", zap.String("server", s.Uuid))
+			if err := s.Environment.Create(); err != nil {
+				zap.S().Errorw("failed to create an environment for server", zap.String("server", s.Uuid), zap.Error(err))
+			}
+
+			if r, err := s.Environment.IsRunning(); err != nil {
+				zap.S().Errorw("error checking server environment status", zap.String("server", s.Uuid), zap.Error(err))
+			} else if r {
+				// If the server is currently running on Docker, mark the process as being in that state.
+				// We never want to stop an instance that is currently running external from Wings since
+				// that is a good way of keeping things running even if Wings gets in a very corrupted state.
+				zap.S().Infow("detected server is running, re-attaching to process", zap.String("server", s.Uuid))
+				s.SetState(server.ProcessRunningState)
+
+				// If we cannot attach to the environment go ahead and mark the processs as being offline.
+				if err := s.Environment.Attach(); err != nil {
+					zap.S().Warnw("error attaching to server environment", zap.String("server", s.Uuid), zap.Error(err))
+					s.SetState(server.ProcessOfflineState)
+				}
+			} else if !r {
+				// If the server is not in a running state right now but according to the configuration it
+				// should be, we want to go ahead and restart the instance.
+				if s.State == server.ProcessRunningState || s.State == server.ProcessStartingState {
+					if err := s.Environment.Start(); err != nil {
+						zap.S().Warnw(
+							"failed to put server instance back in running state",
+							zap.String("server", s.Uuid),
+							zap.Error(errors.WithStack(err)),
+						)
+					}
+				} else {
+					if s.State == "" {
+						// Addresses potentially invalid data in the stored file that can cause Wings to lose
+						// track of what the actual server state is.
+						s.SetState(server.ProcessOfflineState)
+					}
+				}
+			}
+		}(serv)
+	}
+
+	// Wait until all of the servers are ready to go before we fire up the HTTP server.
+	wg.Wait()
+
 	r := &Router{
 		Servers: servers,
 		token:   c.AuthenticationToken,