Tweak shutdown behavior

Author: nevillelyh

cmd/cog/main.go

@@ -97,12 +97,14 @@ func serverCommand() *ff.Command {
 			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())
+				for {
+					sig := <-ch
+					if sig == syscall.SIGTERM && 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) {

internal/tests/shutdown_test.go

@@ -10,6 +10,23 @@ import (
 	"github.com/replicate/cog-runtime/internal/server"
 )
 
+func TestShutdownByServerSigInt(t *testing.T) {
+	if *legacyCog {
+		// Compat: legacy Cog doesn't handle SIGINT properly without a TTY
+		t.SkipNow()
+	}
+	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.SIGINT))
+	assert.NoError(t, ct.Cleanup())
+	assert.Equal(t, 0, ct.cmd.ProcessState.ExitCode())
+}
+
 func TestShutdownByServerSigTerm(t *testing.T) {
 	ct := NewCogTest(t, "sleep")
 	assert.NoError(t, ct.Start())
@@ -32,11 +49,18 @@ func TestShutdownIgnoreSignal(t *testing.T) {
 	assert.Equal(t, server.StatusReady.String(), hc.Status)
 	assert.Equal(t, server.SetupSucceeded, hc.Setup.Status)
 
+	// Ignore SIGTERM
 	must.Do(syscall.Kill(ct.ServerPid(), syscall.SIGTERM))
 	assert.Nil(t, ct.cmd.ProcessState)
 	assert.Equal(t, server.StatusReady.String(), ct.HealthCheck().Status)
 
-	ct.Shutdown()
+	if *legacyCog {
+		// Compat: legacy Cog doesn't handle SIGINT properly without a TTY
+		ct.Shutdown()
+	} else {
+		// Handle SIGINT
+		must.Do(syscall.Kill(ct.ServerPid(), syscall.SIGINT))
+	}
 	assert.NoError(t, ct.Cleanup())
 	assert.Equal(t, 0, ct.cmd.ProcessState.ExitCode())
 }

python/coglet/file_runner.py

@@ -108,6 +108,9 @@ def _cancel_handler(signum, _) -> None:
         else:
             # Blocking predict, use SIGUSR1 to cancel
             signal.signal(signal.SIGUSR1, _cancel_handler)
+        # When running in a shell, Ctrl-C sends SIGINT to all processes in the process group
+        # Ignore it here and require shutdown from parent Go server
+        signal.signal(signal.SIGINT, signal.SIG_IGN)
 
         ready = True
         self._send_ipc(FileRunner.IPC_READY)