From 956e87eb9310497cee3d566f17111ef4f704c73a Mon Sep 17 00:00:00 2001
From: Dane Everitt <dane@daneeveritt.com>
Date: Tue, 18 Aug 2020 20:43:28 -0700
Subject: [PATCH] Handle restarting correctly, ensure power state propagates
 correctly

---
 environment/docker/power.go | 14 +++++++++++++-
 server/power.go             | 28 ++++++++++++----------------
 2 files changed, 25 insertions(+), 17 deletions(-)

diff --git a/environment/docker/power.go b/environment/docker/power.go
index 6187d83..72efddc 100644
--- a/environment/docker/power.go
+++ b/environment/docker/power.go
@@ -134,7 +134,11 @@ func (e *Environment) Stop() error {
 		return e.Terminate(os.Kill)
 	}
 
-	e.setState(system.ProcessStoppingState)
+	// If the process is already offline don't switch it back to stopping. Just leave it how
+	// it is and continue through to the stop handling for the process.
+	if e.State() != system.ProcessOfflineState {
+		e.setState(system.ProcessStoppingState)
+	}
 
 	// Only attempt to send the stop command to the instance if we are actually attached to
 	// the instance. If we are not for some reason, just send the container stop event.
@@ -202,6 +206,14 @@ func (e *Environment) Terminate(signal os.Signal) error {
 	}
 
 	if !c.State.Running {
+		// If the container is not running but we're not already in a stopped state go ahead
+		// and update things to indicate we should be completely stopped now. Set to stopping
+		// first so crash detection is not triggered.
+		if e.State() != system.ProcessOfflineState {
+			e.setState(system.ProcessStoppingState)
+			e.setState(system.ProcessOfflineState)
+		}
+
 		return nil
 	}
 
diff --git a/server/power.go b/server/power.go
index ab868fb..f6ea50d 100644
--- a/server/power.go
+++ b/server/power.go
@@ -95,22 +95,18 @@ func (s *Server) HandlePowerAction(action PowerAction, waitSeconds ...int) error
 	case PowerActionStop:
 		// We're specificially waiting for the process to be stopped here, otherwise the lock is released
 		// too soon, and you can rack up all sorts of issues.
-		return s.Environment.WaitForStop(10 * 60, true)
+		return s.Environment.WaitForStop(10*60, true)
 	case PowerActionRestart:
-		// Only try to wait for stop if the process is currently running, otherwise just skip right to the
-		// start event.
-		if r, _ := s.Environment.IsRunning(); !r {
-			if err := s.Environment.WaitForStop(10 * 60, 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.
-				//
-				// However, if terminate is not passed you'll get a context deadline error. We could
-				// probably handle that nicely here, but I'd rather just pass it back up the stack for now.
-				// Either way, any type of error indicates we should not attempt to start the server back
-				// up.
-				return err
-			}
+		if err := s.Environment.WaitForStop(10*60, 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.
+			//
+			// However, if terminate is not passed you'll get a context deadline error. We could
+			// probably handle that nicely here, but I'd rather just pass it back up the stack for now.
+			// Either way, any type of error indicates we should not attempt to start the server back
+			// up.
+			return err
 		}
 
 		// Now actually try to start the process by executing the normal pre-boot logic.
@@ -158,4 +154,4 @@ func (s *Server) onBeforeStart() error {
 	}
 
 	return nil
-}
\ No newline at end of file
+}