Fix logic for context based environment stopping

Uses dual contexts to handle stopping using a timed context, and also terminating the entire process loop if the parent context gets canceled.

server/backup.go

@@ -142,7 +142,7 @@ func (s *Server) RestoreBackup(b backup.BackupInterface, reader io.ReadCloser) (
 	// instance, otherwise you'll likely hit all types of write errors due to the
 	// server being suspended.
 	if s.Environment.State() != environment.ProcessOfflineState {
-		if err = s.Environment.WaitForStop(120, false); err != nil {
+		if err = s.Environment.WaitForStop(s.Context(), time.Minute*2, false); err != nil {
 			if !client.IsErrNotFound(err) {
 				return errors.WrapIf(err, "server/backup: restore: failed to wait for container stop")
 			}

server/install.go

@@ -10,6 +10,7 @@ import (
 	"path/filepath"
 	"strconv"
 	"strings"
+	"time"
 
 	"emperror.dev/errors"
 	"github.com/apex/log"
@@ -79,7 +80,7 @@ func (s *Server) Install(sync bool) error {
 func (s *Server) Reinstall() error {
 	if s.Environment.State() != environment.ProcessOfflineState {
 		s.Log().Debug("waiting for server instance to enter a stopped state")
-		if err := s.Environment.WaitForStop(10, true); err != nil {
+		if err := s.Environment.WaitForStop(s.Context(), time.Second*10, true); err != nil {
 			return err
 		}
 	}

server/listeners.go

@@ -5,6 +5,7 @@ import (
 	"regexp"
 	"strconv"
 	"sync"
+	"time"
 
 	"github.com/apex/log"
 
@@ -44,7 +45,7 @@ func (dsl *diskSpaceLimiter) Reset() {
 func (dsl *diskSpaceLimiter) Trigger() {
 	dsl.o.Do(func() {
 		dsl.server.PublishConsoleOutputFromDaemon("Server is exceeding the assigned disk space limit, stopping process now.")
-		if err := dsl.server.Environment.WaitForStop(60, true); err != nil {
+		if err := dsl.server.Environment.WaitForStop(dsl.server.Context(), time.Minute, true); err != nil {
 			dsl.server.Log().WithField("error", err).Error("failed to stop server after exceeding space limit!")
 		}
 	})

server/power.go

@@ -135,12 +135,9 @@ func (s *Server) HandlePowerAction(action PowerAction, waitSeconds ...int) error
 	case PowerActionStop:
 		fallthrough
 	case PowerActionRestart:
-		ctx, cancel := context.WithTimeout(s.Context(), time.Second)
-		defer cancel()
-
 		// We're specifically waiting for the process to be stopped here, otherwise the lock is
 		// released too soon, and you can rack up all sorts of issues.
-		if err := s.Environment.WaitForStopWithContext(ctx, true); err != nil {
+		if err := s.Environment.WaitForStop(s.Context(), time.Minute*10, true); err != nil {
 			// Even timeout errors should be bubbled back up the stack. If the process didn't stop
 			// nicely, but the terminate argument was passed then the server is stopped without an
 			// error being returned.
@@ -156,11 +153,6 @@ func (s *Server) HandlePowerAction(action PowerAction, waitSeconds ...int) error
 			return nil
 		}
 
-		// Release the resources we acquired for the initial timer context since we don't
-		// need them anymore at this point, and the start process can take quite awhile to
-		// complete.
-		cancel()
-
 		// Now actually try to start the process by executing the normal pre-boot logic.
 		if err := s.onBeforeStart(); err != nil {
 			return err

server/update.go

@@ -1,6 +1,8 @@
 package server
 
 import (
+	"time"
+
 	"github.com/pterodactyl/wings/environment/docker"
 
 	"github.com/pterodactyl/wings/environment"
@@ -58,7 +60,7 @@ func (s *Server) SyncWithEnvironment() {
 			s.Log().Info("server suspended with running process state, terminating now")
 
 			go func(s *Server) {
-				if err := s.Environment.WaitForStop(60, true); err != nil {
+				if err := s.Environment.WaitForStop(s.Context(), time.Minute, true); err != nil {
 					s.Log().WithField("error", err).Warn("failed to terminate server environment after suspension")
 				}
 			}(s)