Fix orphaned process on shutdown

Author: nevillelyh

cmd/cog/main.go

@@ -78,32 +78,31 @@ func serverCommand() *ff.Command {
 				"upload-url", cfg.UploadUrl,
 			)
 
-			ctx, cancel := context.WithCancel(ctx)
-			go func() {
-				ch := make(chan os.Signal, 1)
-				signal.Notify(ch, os.Interrupt, syscall.SIGTERM)
-				s := <-ch
-				if cfg.AwaitExplicitShutdown {
-					log.Warnw("ignoring signal to stop", "signal", s)
-				} else {
-					log.Infow("stopping Cog HTTP server", "signal", s)
-					cancel()
-				}
-			}()
-
 			addr := fmt.Sprintf("%s:%d", cfg.Host, cfg.Port)
 			log.Infow("starting Cog HTTP server", "addr", addr, "version", util.Version(), "pid", os.Getpid())
 			serverCfg := server.Config{
 				UseProcedureMode:      cfg.UseProcedureMode,
 				AwaitExplicitShutdown: cfg.AwaitExplicitShutdown,
 				UploadUrl:             cfg.UploadUrl,
 			}
+			ctx, cancel := context.WithCancel(ctx)
 			h := server.NewHandler(serverCfg, cancel)
 			s := server.NewServer(addr, h, cfg.UseProcedureMode)
 			go func() {
 				<-ctx.Done()
 				must.Do(s.Shutdown(ctx))
 			}()
+			go func() {
+				ch := make(chan os.Signal, 1)
+				signal.Notify(ch, os.Interrupt, syscall.SIGTERM)
+				sig := <-ch
+				if cfg.AwaitExplicitShutdown {
+					log.Warnw("ignoring signal to stop", "signal", sig)
+				} else {
+					log.Infow("stopping Cog HTTP server", "signal", sig)
+					must.Do(h.Stop())
+				}
+			}()
 			if err := s.ListenAndServe(); errors.Is(err, http.ErrServerClosed) {
 				exitCode := h.ExitCode()
 				if exitCode == 0 {

internal/server/server.go

@@ -107,25 +107,32 @@ func (h *Handler) OpenApi(w http.ResponseWriter, r *http.Request) {
 }
 
 func (h *Handler) Shutdown(w http.ResponseWriter, r *http.Request) {
+	err := h.Stop()
+	if err != nil {
+		http.Error(w, err.Error(), http.StatusInternalServerError)
+	} else {
+		w.WriteHeader(http.StatusOK)
+	}
+}
+
+func (h *Handler) Stop() error {
 	// Procedure mode and no runner yet
 	if h.runner == nil {
-		w.WriteHeader(http.StatusOK)
 		// Shut down immediately
 		h.shutdown()
-		return
+		return nil
 	}
 
 	// Request runner stop
 	if err := h.runner.Stop(); err != nil {
-		http.Error(w, err.Error(), http.StatusInternalServerError)
-	} else {
-		w.WriteHeader(http.StatusOK)
+		return err
 	}
 	go func() {
 		// Shut down once runner exists
 		h.runner.WaitForStop()
 		h.shutdown()
 	}()
+	return nil
 }
 
 func (h *Handler) updateRunner(srcDir string) error {

internal/tests/shutdown_test.go

@@ -23,19 +23,6 @@ func TestShutdownByServerSigTerm(t *testing.T) {
 	assert.Equal(t, 0, ct.cmd.ProcessState.ExitCode())
 }
 
-func TestShutdownByRunnerSigTerm(t *testing.T) {
-	ct := NewCogTest(t, "sleep")
-	assert.NoError(t, ct.Start())
-
-	hc := ct.WaitForSetup()
-	assert.Equal(t, server.StatusReady.String(), hc.Status)
-	assert.Equal(t, server.SetupSucceeded, hc.Setup.Status)
-
-	must.Do(syscall.Kill(ct.ServerPid(), syscall.SIGTERM))
-	assert.NoError(t, ct.Cleanup())
-	assert.Equal(t, 0, ct.cmd.ProcessState.ExitCode())
-}
-
 func TestShutdownIgnoreSignal(t *testing.T) {
 	ct := NewCogTest(t, "sleep")
 	ct.AppendArgs("--await-explicit-shutdown=true")