Pass a context through to the start/stop/terminate actions

environment/docker/power.go

@@ -138,9 +138,7 @@ func (e *Environment) Start(ctx context.Context) error {
 // You most likely want to be using WaitForStop() rather than this function,
 // since this will return as soon as the command is sent, rather than waiting
 // for the process to be completed stopped.
-//
+func (e *Environment) Stop(ctx context.Context) error {
-// TODO: pass context through from the server instance.
-func (e *Environment) Stop() error {
 	e.mu.RLock()
 	s := e.meta.Stop
 	e.mu.RUnlock()
@@ -164,7 +162,7 @@ func (e *Environment) Stop() error {
 		case "SIGTERM":
 			signal = syscall.SIGTERM
 		}
-		return e.Terminate(signal)
+		return e.Terminate(ctx, signal)
 	}
 
 	// If the process is already offline don't switch it back to stopping. Just leave it how
@@ -180,7 +178,7 @@ func (e *Environment) Stop() error {
 	}
 
 	t := time.Second * 30
-	if err := e.client.ContainerStop(context.Background(), e.Id, &t); err != nil {
+	if err := e.client.ContainerStop(ctx, e.Id, &t); err != nil {
 		// If the container does not exist just mark the process as stopped and return without
 		// an error.
 		if client.IsErrNotFound(err) {
@@ -199,27 +197,30 @@ func (e *Environment) Stop() error {
 // be returned, or the instance will be terminated forcefully depending on the
 // value of the second argument.
 func (e *Environment) WaitForStop(seconds uint, terminate bool) error {
-	if err := e.Stop(); err != nil {
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second * time.Duration(seconds))
+	defer cancel()
+
+	return e.WaitForStopWithContext(ctx, terminate)
+}
+
+func (e *Environment) WaitForStopWithContext(ctx context.Context, terminate bool) error {
+	if err := e.Stop(ctx); err != nil {
 		return err
 	}
 
-	ctx, cancel := context.WithTimeout(context.Background(), time.Duration(seconds)*time.Second)
-	defer cancel()
-
 	// Block the return of this function until the container as been marked as no
 	// longer running. If this wait does not end by the time seconds have passed,
 	// attempt to terminate the container, or return an error.
 	ok, errChan := e.client.ContainerWait(ctx, e.Id, container.WaitConditionNotRunning)
 	select {
 	case <-ctx.Done():
-		if ctxErr := ctx.Err(); ctxErr != nil {
+		if err := ctx.Err(); err != nil {
 			if terminate {
 				log.WithField("container_id", e.Id).Info("server did not stop in time, executing process termination")
 
-				return e.Terminate(os.Kill)
+				return e.Terminate(ctx, os.Kill)
 			}
-
+			return err
-			return ctxErr
 		}
 	case err := <-errChan:
 		// If the error stems from the container not existing there is no point in wasting
@@ -233,7 +234,7 @@ func (e *Environment) WaitForStop(seconds uint, terminate bool) error {
 					l.WithField("error", err).Warn("error while waiting for container stop; terminating process")
 				}
 
-				return e.Terminate(os.Kill)
+				return e.Terminate(ctx, os.Kill)
 			}
 			return errors.WrapIf(err, "environment/docker: error waiting on container to enter \"not-running\" state")
 		}
@@ -244,8 +245,8 @@ func (e *Environment) WaitForStop(seconds uint, terminate bool) error {
 }
 
 // Terminate forcefully terminates the container using the signal provided.
-func (e *Environment) Terminate(signal os.Signal) error {
+func (e *Environment) Terminate(ctx context.Context, signal os.Signal) error {
-	c, err := e.ContainerInspect(context.Background())
+	c, err := e.ContainerInspect(ctx)
 	if err != nil {
 		// Treat missing containers as an okay error state, means it is obviously
 		// already terminated at this point.
@@ -270,7 +271,7 @@ func (e *Environment) Terminate(signal os.Signal) error {
 	// We set it to stopping than offline to prevent crash detection from being triggered.
 	e.SetState(environment.ProcessStoppingState)
 	sig := strings.TrimSuffix(strings.TrimPrefix(signal.String(), "signal "), "ed")
-	if err := e.client.ContainerKill(context.Background(), e.Id, sig); err != nil && !client.IsErrNotFound(err) {
+	if err := e.client.ContainerKill(ctx, e.Id, sig); err != nil && !client.IsErrNotFound(err) {
 		return errors.WithStack(err)
 	}
 	e.SetState(environment.ProcessOfflineState)

environment/environment.go

@@ -58,18 +58,24 @@ type ProcessEnvironment interface {
 	// can be started an error should be returned.
 	Start(ctx context.Context) error
 
-	// Stops a server instance. If the server is already stopped an error should
+	// Stop stops a server instance. If the server is already stopped an error will
-	// not be returned.
+	// not be returned, this function will act as a no-op.
-	Stop() error
+	Stop(ctx context.Context) error
 
-	// Waits for a server instance to stop gracefully. If the server is still detected
+	// WaitForStop waits for a server instance to stop gracefully. If the server is
-	// as running after seconds, an error will be returned, or the server will be terminated
+	// still detected as running after seconds, an error will be returned, or the server
-	// depending on the value of the second argument.
+	// will be terminated depending on the value of the second argument.
 	WaitForStop(seconds uint, terminate bool) error
 
-	// Terminates a running server instance using the provided signal. If the server
+	// WaitForStopWithContext works in the same fashion as WaitForStop, but accepts a
-	// is not running no error should be returned.
+	// context value and will stop waiting when the context is canceled. If the terminate
-	Terminate(signal os.Signal) error
+	// option is true, the server will be terminate when the context is canceled if the
+	// server has not stopped at that point.
+	WaitForStopWithContext(ctx context.Context, terminate bool) error
+
+	// Terminate stops a running server instance using the provided signal. This function
+	// is a no-op if the server is already stopped.
+	Terminate(ctx context.Context, signal os.Signal) error
 
 	// Destroys the environment removing any containers that were created (in Docker
 	// environments at least).

server/power.go

@@ -133,11 +133,14 @@ func (s *Server) HandlePowerAction(action PowerAction, waitSeconds ...int) error
 
 		return s.Environment.Start(s.Context())
 	case PowerActionStop:
-		// We're specifically waiting for the process to be stopped here, otherwise the lock is released
+		fallthrough
-		// too soon, and you can rack up all sorts of issues.
-		return s.Environment.WaitForStop(10*60, true)
 	case PowerActionRestart:
-		if err := s.Environment.WaitForStop(10*60, true); err != nil {
+		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 {
 			// 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.
@@ -149,6 +152,15 @@ func (s *Server) HandlePowerAction(action PowerAction, waitSeconds ...int) error
 			return err
 		}
 
+		if action == PowerActionStop {
+			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
@@ -156,7 +168,7 @@ func (s *Server) HandlePowerAction(action PowerAction, waitSeconds ...int) error
 
 		return s.Environment.Start(s.Context())
 	case PowerActionTerminate:
-		return s.Environment.Terminate(os.Kill)
+		return s.Environment.Terminate(s.Context(), os.Kill)
 	}
 
 	return errors.New("attempting to handle unknown power action")
@@ -207,5 +219,6 @@ func (s *Server) onBeforeStart() error {
 		}
 	}
 
+	s.Log().Info("completed server preflight, starting boot process...")
 	return nil
 }