Implement graceful shutdown with proper prediction completion

This implements a comprehensive graceful shutdown mechanism that waits for
in-flight predictions to complete before stopping runners and the service.

Key changes:

**Runner-level graceful shutdown:**
- Add shutdownWhenIdle atomic flag and readyForShutdown channel to Runner
- GracefulShutdown() signals runners to shutdown when idle
- updateStatus() automatically closes readyForShutdown when becoming READY with no pending predictions
- Add nil check with warning for test compatibility

**Handler-level prediction rejection:**
- Add gracefulShutdown atomic flag to reject new predictions during shutdown
- Handler.Stop() sets flag and waits for manager shutdown
- Predict() returns 503 Service Unavailable during shutdown

**Manager-level coordinated shutdown:**
- Manager.Stop() signals all runners for graceful shutdown
- Use WaitGroup.Go() for independent parallel runner shutdowns
- Respect RunnerShutdownGracePeriod timeout before force stopping
- Wait on runner.readyForShutdown channel or timeout

**Service-level errgroup coordination:**
- Fix errgroup goroutines to exit on shutdown signal
- Add shutdown case to force shutdown monitor goroutine
- Signal handler already had proper shutdown case
- Add contextcheck nolint for long-lived errgroup context

**Test coverage:**
- Add E2E test for 503 rejection of new predictions during shutdown
- Verify graceful shutdown waits for in-flight predictions
- Test service properly stops after shutdown completes

This restores the graceful shutdown behavior from commit 575d218 that was
lost during the server refactor, ensuring predictions complete naturally
during the grace period rather than being immediately force-killed.

Author: tempusfrangit

internal/runner/manager.go

@@ -17,7 +17,6 @@ import (
 	"time"
 
 	"go.uber.org/zap"
-	"golang.org/x/sync/errgroup"
 
 	"github.com/replicate/cog-runtime/internal/config"
 	"github.com/replicate/cog-runtime/internal/webhook"
@@ -747,43 +746,48 @@ func (m *Manager) Stop() error {
 		log.Info("stopping runner manager")
 
 		m.mu.Lock()
-		defer m.mu.Unlock()
-
-		// Stop all runners
-		for i, runner := range m.runners {
+		runnerList := make([]*Runner, 0, len(m.runners))
+		for _, runner := range m.runners {
 			if runner != nil {
-				log.Infow("stopping runner", "name", runner.runnerCtx.id, "slot", i)
-				if err := runner.Stop(); err != nil {
-					log.Errorw("error stopping runner", "name", runner.runnerCtx.id, "error", err)
-					if stopErr == nil {
-						stopErr = err
-					}
-				}
+				runnerList = append(runnerList, runner)
 			}
 		}
+		m.mu.Unlock()
 
-		// Wait for runners to stop concurrently
-		eg := errgroup.Group{}
-		for i, runner := range m.runners {
-			if runner != nil {
-				name := runner.runnerCtx.id
-				eg.Go(func() error {
-					log.Infow("waiting for runner to stop", "name", name, "slot", i)
-					runner.WaitForStop()
-					return nil
-				})
-			}
+		// Signal all runners for graceful shutdown
+		for _, runner := range runnerList {
+			runner.GracefulShutdown()
 		}
 
-		if err := eg.Wait(); err != nil {
-			log.Errorw("error waiting for runners to stop", "error", err)
-			if stopErr == nil {
-				stopErr = err
-			}
-		} else {
-			log.Info("all runners stopped successfully")
+		// Wait for runners to become idle or timeout using WaitGroup
+		gracePeriod := m.cfg.RunnerShutdownGracePeriod
+		log.Infow("grace period configuration", "grace_period", gracePeriod)
+		graceCtx, cancel := context.WithTimeout(m.ctx, gracePeriod)
+		defer cancel()
+
+		var wg sync.WaitGroup
+		for _, runner := range runnerList {
+			wg.Go(func() {
+				log.Debugw("waiting for runner to become idle", "name", runner.runnerCtx.id, "grace_period", gracePeriod)
+				// Wait for this runner to become idle OR timeout
+				select {
+				case <-runner.readyForShutdown:
+					log.Infow("runner became idle naturally", "name", runner.runnerCtx.id)
+				case <-graceCtx.Done():
+					log.Warnw("grace period expired for runner", "name", runner.runnerCtx.id, "context_err", graceCtx.Err())
+				}
+
+				// Always try to stop, handle errors independently
+				if err := runner.Stop(); err != nil {
+					log.Errorw("failed to stop runner gracefully", "name", runner.runnerCtx.id, "error", err)
+				}
+			})
 		}
 
+		// Wait for all runners to complete shutdown (success or failure)
+		wg.Wait()
+
+		log.Info("all runners stopped successfully")
 		close(m.stopped)
 	})
 

internal/runner/runner.go

@@ -13,6 +13,7 @@ import (
 	"regexp"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"syscall"
 	"time"
 
@@ -303,6 +304,8 @@ type Runner struct {
 	procedureHash      string
 	mu                 sync.RWMutex
 	stopped            chan bool
+	shutdownWhenIdle   atomic.Bool
+	readyForShutdown   chan struct{} // closed when idle and ready to be stopped
 	setupComplete      chan struct{} // closed on first READY after setup
 	webhookSender      webhook.Sender
 	logCaptureComplete chan struct{} // closed when both stdout/stderr capture complete
@@ -345,6 +348,34 @@ func (r *Runner) WaitForStop() {
 	<-r.stopped
 }
 
+func (r *Runner) GracefulShutdown() {
+	log := r.logger.Sugar()
+	if !r.shutdownWhenIdle.CompareAndSwap(false, true) {
+		log.Debugw("graceful shutdown already initiated", "runner_id", r.runnerCtx.id)
+		return
+	}
+
+	r.mu.RLock()
+	shouldSignal := (r.status == StatusReady && len(r.pending) == 0)
+	r.mu.RUnlock()
+
+	log.Debugw("graceful shutdown initiated", "runner_id", r.runnerCtx.id, "status", r.status, "pending_count", len(r.pending), "should_signal", shouldSignal)
+
+	if shouldSignal {
+		if r.readyForShutdown == nil {
+			log.Warnw("readyForShutdown channel is nil, cannot signal shutdown readiness", "runner_id", r.runnerCtx.id)
+		} else {
+			select {
+			case <-r.readyForShutdown:
+				log.Debugw("readyForShutdown already closed", "runner_id", r.runnerCtx.id)
+			default:
+				log.Debugw("closing readyForShutdown channel", "runner_id", r.runnerCtx.id)
+				close(r.readyForShutdown)
+			}
+		}
+	}
+}
+
 func (r *Runner) Start(ctx context.Context) error {
 	log := r.logger.Sugar()
 	r.mu.Lock()
@@ -676,44 +707,32 @@ func (r *Runner) ForceKill() {
 
 func (r *Runner) verifyProcessCleanup(pid int) {
 	log := r.logger.Sugar()
-	const checkInterval = 10 * time.Millisecond
-
 	log.Infow("starting process cleanup verification", "pid", pid)
 
 	timeout := r.cleanupTimeout
 	if timeout == 0 {
-		timeout = 10 * time.Second // Default fallback
+		timeout = 10 * time.Second
 	}
 
 	timer := time.NewTimer(timeout)
 	defer timer.Stop()
 
-	ticker := time.NewTicker(checkInterval)
-	defer ticker.Stop()
-
-	for {
+	select {
+	case <-r.stopped:
+		log.Infow("process cleanup verified successfully", "pid", pid)
 		select {
-		case <-timer.C:
-			log.Errorw("process cleanup timeout exceeded, forcing server exit",
-				"pid", pid, "timeout", timeout)
-			// Signal forced shutdown - this is idempotent and never blocks
-			if r.forceShutdown.TriggerForceShutdown() {
-				log.Errorw("triggered force shutdown signal")
-			}
-			return
+		case r.cleanupSlot <- struct{}{}:
+		default:
+		}
+		return
 
-		case <-ticker.C:
-			// Verify if process group has been terminated
-			if err := r.verifyFn(pid); err == nil {
-				log.Infow("process cleanup verified successfully", "pid", pid)
-				// Return cleanup token to allow future cleanup
-				select {
-				case r.cleanupSlot <- struct{}{}:
-				default:
-				}
-				return
-			}
+	case <-timer.C:
+		log.Errorw("process cleanup timeout exceeded, forcing server exit",
+			"pid", pid, "timeout", timeout)
+		if r.forceShutdown.TriggerForceShutdown() {
+			log.Errorw("triggered force shutdown signal")
 		}
+		return
 	}
 }
 
@@ -806,6 +825,17 @@ func (r *Runner) updateStatus(statusStr string) error {
 		return err
 	}
 	r.status = status
+
+	// Close readyForShutdown channel when idle and shutdown requested
+	if status == StatusReady && r.shutdownWhenIdle.Load() && len(r.pending) == 0 {
+		select {
+		case <-r.readyForShutdown:
+			// Already closed
+		default:
+			close(r.readyForShutdown)
+		}
+	}
+
 	return nil
 }
 
@@ -909,14 +939,14 @@ func (r *Runner) updateSetupResult() {
 	switch r.setupResult.Status {
 	case SetupSucceeded:
 		r.status = StatusReady
-		log.Debug("setup succeeded", "status", r.status.String())
+		log.Debugw("setup succeeded", "status", r.status.String())
 	case SetupFailed:
 		r.status = StatusSetupFailed
-		log.Debug("setup failed", "status", r.status.String())
+		log.Debugw("setup failed", "status", r.status.String())
 	default:
 		r.setupResult.Status = SetupFailed
 		r.status = StatusSetupFailed
-		log.Debug("unknown setup status, defaulting to failed", "status", r.status.String())
+		log.Debugw("unknown setup status, defaulting to failed", "status", r.status.String())
 	}
 }
 
@@ -972,6 +1002,7 @@ func NewRunner(ctx context.Context, ctxCancel context.CancelFunc, runnerCtx Runn
 		verifyFn:           verifyProcessGroupTerminated,
 		cleanupSlot:        make(chan struct{}, 1),
 		stopped:            make(chan bool),
+		readyForShutdown:   make(chan struct{}),
 		setupComplete:      make(chan struct{}),
 		logCaptureComplete: make(chan struct{}),
 		cleanupTimeout:     cleanupTimeout,

internal/runner/runner_test.go

@@ -1013,7 +1013,7 @@ func TestRunnerTempDirectoryCleanup(t *testing.T) {
 		err = r.Stop()
 		require.NoError(t, err, "Runner.Stop() should not error")
 
-		time.Sleep(100 * time.Millisecond)
+		time.Sleep(1 * time.Millisecond)
 
 		_, err = os.Stat(tmpDir)
 		assert.True(t, os.IsNotExist(err), "Temp directory should be cleaned up after Stop()")
@@ -1424,30 +1424,25 @@ func TestForceKillCleanupFailures(t *testing.T) {
 
 	t.Run("procedure mode cleanup success returns token", func(t *testing.T) {
 		forceShutdown := config.NewForceShutdownSignal()
-		verifyCallCount := 0
 
 		r := &Runner{
-			cleanupTimeout: 100 * time.Millisecond,
+			cleanupTimeout: 1 * time.Millisecond,
 			forceShutdown:  forceShutdown,
 			cleanupSlot:    make(chan struct{}, 1),
-			verifyFn: func(pid int) error {
-				verifyCallCount++
-				if verifyCallCount >= 3 {
-					return nil // Success after a few attempts
-				}
-				return fmt.Errorf("process still exists")
-			},
-			logger: zaptest.NewLogger(t),
+			stopped:        make(chan bool),
+			logger:         zaptest.NewLogger(t),
 		}
 
 		// Start verification process
 		go r.verifyProcessCleanup(12345)
 
-		// Wait for verification to complete
-		time.Sleep(150 * time.Millisecond)
+		// Signal process stopped to trigger cleanup completion
+		close(r.stopped)
+
+		// Give a moment for cleanup to complete
+		time.Sleep(1 * time.Millisecond)
 
 		// Verify cleanup token was returned
 		assert.Len(t, r.cleanupSlot, 1)
-		assert.GreaterOrEqual(t, verifyCallCount, 3)
 	})
 }

internal/server/mux.go

@@ -18,7 +18,6 @@ func NewServeMux(handler *Handler, useProcedureMode bool) *http.ServeMux {
 	serveMux.HandleFunc("GET /{$}", handler.Root)
 	serveMux.HandleFunc("GET /health-check", handler.HealthCheck)
 	serveMux.HandleFunc("GET /openapi.json", handler.OpenAPI)
-	serveMux.HandleFunc("POST /shutdown", handler.Shutdown)
 
 	if useProcedureMode {
 		serveMux.HandleFunc("POST /procedures", handler.Predict)

internal/server/server.go

@@ -10,6 +10,7 @@ import (
 	"net/http"
 	"os"
 	"path"
+	"sync/atomic"
 	"time"
 
 	"go.uber.org/zap"
@@ -37,22 +38,21 @@ type IPC struct {
 }
 
 type Handler struct {
-	cfg           config.Config
-	shutdown      context.CancelFunc
-	startedAt     time.Time
-	runnerManager *runner.Manager
+	cfg              config.Config
+	startedAt        time.Time
+	runnerManager    *runner.Manager
+	gracefulShutdown atomic.Bool
 
 	cwd string
 
 	logger *zap.Logger
 }
 
-func NewHandler(ctx context.Context, cfg config.Config, shutdown context.CancelFunc, baseLogger *zap.Logger) (*Handler, error) {
+func NewHandler(ctx context.Context, cfg config.Config, baseLogger *zap.Logger) (*Handler, error) {
 	runnerManager := runner.NewManager(ctx, cfg, baseLogger)
 
 	h := &Handler{
 		cfg:           cfg,
-		shutdown:      shutdown,
 		startedAt:     time.Now(),
 		runnerManager: runnerManager,
 		cwd:           cfg.WorkingDirectory,
@@ -133,30 +133,21 @@ func (h *Handler) OpenAPI(w http.ResponseWriter, r *http.Request) {
 	h.writeBytes(w, []byte(schema))
 }
 
-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)
-	}
-}
-
 // ForceKillAll immediately force-kills all runners (for test cleanup)
 func (h *Handler) ForceKillAll() {
 	h.runnerManager.ForceKillAll()
 }
 
 func (h *Handler) Stop() error {
-	// Stop the runner manager and handle shutdown in background
-	go func() {
-		log := h.logger.Sugar()
-		if err := h.runnerManager.Stop(); err != nil {
-			log.Errorw("failed to stop runner manager", "error", err)
-			os.Exit(1)
-		}
-		h.shutdown()
-	}()
+	// Set graceful shutdown flag to reject new predictions
+	h.gracefulShutdown.Store(true)
+
+	// Stop the runner manager synchronously
+	log := h.logger.Sugar()
+	if err := h.runnerManager.Stop(); err != nil {
+		log.Errorw("failed to stop runner manager", "error", err)
+		return err
+	}
 	return nil
 }
 
@@ -207,6 +198,13 @@ func (h *Handler) HandleIPC(w http.ResponseWriter, r *http.Request) {
 
 func (h *Handler) Predict(w http.ResponseWriter, r *http.Request) {
 	log := h.logger.Sugar()
+
+	// Reject new predictions during graceful shutdown
+	if h.gracefulShutdown.Load() {
+		http.Error(w, "server shutting down", http.StatusServiceUnavailable)
+		return
+	}
+
 	if r.Header.Get("Content-Type") != "application/json" {
 		http.Error(w, "invalid content type", http.StatusUnsupportedMediaType)
 		return