Merge branch 'shutdown-fixes' into cleanup-tmp

Author: tempusfrangit

internal/runner/runner.go

@@ -707,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
 	}
 }
 

internal/runner/runner_test.go

@@ -1017,7 +1017,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()")
@@ -1428,30 +1428,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/tests/shutdown_test.go

@@ -25,7 +25,7 @@ func TestShutdownEndpointE2E(t *testing.T) {
 		module:                    "async_sleep",
 		predictorClass:            "Predictor",
 		concurrencyMax:            1,
-		runnerShutdownGracePeriod: 10 * time.Second,
+		runnerShutdownGracePeriod: 100 * time.Millisecond,
 	})
 	defer httpTestServer.Close()
 
@@ -51,7 +51,7 @@ func TestShutdownEndpointE2E(t *testing.T) {
 	// Service should shutdown gracefully
 	require.Eventually(t, func() bool {
 		return svc.IsStopped()
-	}, 10*time.Second, 100*time.Millisecond, "service should have stopped after shutdown")
+	}, 1*time.Second, 10*time.Millisecond, "service should have stopped after shutdown")
 
 	// Service should no longer be running
 	assert.False(t, svc.IsRunning())
@@ -73,7 +73,7 @@ func TestShutdownEndpointWaitsForInflightPredictions(t *testing.T) {
 		module:                    "async_sleep", // Use async predictor for cancellation support
 		predictorClass:            "Predictor",
 		concurrencyMax:            1,
-		runnerShutdownGracePeriod: 30 * time.Second, // Allow time for graceful shutdown
+		runnerShutdownGracePeriod: 200 * time.Millisecond, // Allow time for graceful shutdown
 	})
 	defer httpTestServer.Close()
 
@@ -108,7 +108,7 @@ func TestShutdownEndpointWaitsForInflightPredictions(t *testing.T) {
 	case webhookEvent := <-receiverServer.webhookReceiverChan:
 		assert.Equal(t, runner.PredictionProcessing, webhookEvent.Response.Status)
 		assert.Equal(t, predictionID, webhookEvent.Response.ID)
-	case <-time.After(10 * time.Second):
+	case <-time.After(5 * time.Second):
 		t.Fatal("did not receive prediction started webhook")
 	}
 
@@ -143,15 +143,15 @@ func TestShutdownEndpointWaitsForInflightPredictions(t *testing.T) {
 	case webhookEvent := <-receiverServer.webhookReceiverChan:
 		assert.Equal(t, predictionID, webhookEvent.Response.ID)
 		assert.Equal(t, runner.PredictionCanceled, webhookEvent.Response.Status)
-	case <-time.After(15 * time.Second):
+	case <-time.After(5 * time.Second):
 		t.Fatal("did not receive prediction canceled webhook")
 	}
 
 	// The key test: verify the SERVICE itself has shut down gracefully
 	// Wait for service to stop (it should stop automatically after shutdown)
 	require.Eventually(t, func() bool {
 		return svc.IsStopped()
-	}, 10*time.Second, 10*time.Millisecond, "service should have stopped after shutdown")
+	}, 1*time.Second, 10*time.Millisecond, "service should have stopped after shutdown")
 
 	// Service should no longer be running
 	assert.False(t, svc.IsRunning())
@@ -187,7 +187,7 @@ func TestShutdownEndpointFailsInNonAwaitMode(t *testing.T) {
 	assert.Equal(t, http.StatusOK, shutdownResp.StatusCode)
 
 	// Give it a moment for shutdown to process
-	time.Sleep(1 * time.Second)
+	time.Sleep(50 * time.Millisecond)
 
 	// Verify server has shut down by attempting health check
 	// Should fail after shutdown completes since the server should be down
@@ -197,7 +197,7 @@ func TestShutdownEndpointFailsInNonAwaitMode(t *testing.T) {
 			resp.Body.Close()
 		}
 		return err != nil // Server should be down
-	}, 10*time.Second, 500*time.Millisecond, "server should have shut down")
+	}, 2*time.Second, 50*time.Millisecond, "server should have shut down")
 }
 
 func TestShutdownEndpointMultipleCallsIdempotent(t *testing.T) {
@@ -238,7 +238,7 @@ func TestShutdownEndpointMultipleCallsIdempotent(t *testing.T) {
 	// Service should shutdown gracefully (only once)
 	require.Eventually(t, func() bool {
 		return svc.IsStopped()
-	}, 10*time.Second, 100*time.Millisecond, "service should have stopped after shutdown")
+	}, 1*time.Second, 10*time.Millisecond, "service should have stopped after shutdown")
 
 	assert.False(t, svc.IsRunning())
 	assert.True(t, svc.IsStopped())