feat(picod): support configurable shutdown timeout and cap execution request timeout

Signed-off-by: Abhinav Singh <abhinavsingh717073@gmail.com>

Author: Abhinav-kodes

cmd/picod/main.go

@@ -21,6 +21,7 @@ import (
 	"flag"
 	"os/signal"
 	"syscall"
+	"time"
 
 	"k8s.io/klog/v2"
 
@@ -30,14 +31,16 @@ import (
 func main() {
 	port := flag.Int("port", 8080, "Port for the PicoD server to listen on")
 	workspace := flag.String("workspace", "", "Root directory for file operations (default: current working directory)")
+	shutdownTimeout := flag.Duration("shutdown-timeout", 90*time.Second, "Grace period for graceful shutdown")
 
 	// Initialize klog flags
 	klog.InitFlags(nil)
 	flag.Parse()
 
 	config := picod.Config{
-		Port:      *port,
-		Workspace: *workspace,
+		Port:            *port,
+		Workspace:       *workspace,
+		ShutdownTimeout: *shutdownTimeout,
 	}
 
 	// Create a context that is canceled on SIGINT or SIGTERM.

pkg/picod/execute.go

@@ -27,6 +27,7 @@ import (
 	"time"
 
 	"github.com/gin-gonic/gin"
+	"k8s.io/klog/v2"
 )
 
 const (
@@ -84,6 +85,16 @@ func (s *Server) ExecuteHandler(c *gin.Context) {
 		}
 	}
 
+	// Cap timeout to shutdown grace period
+	shutdownTimeout := s.config.ShutdownTimeout
+	if shutdownTimeout <= 0 {
+		shutdownTimeout = 90 * time.Second
+	}
+	if timeoutDuration > shutdownTimeout {
+		klog.Infof("Requested timeout %v exceeds shutdown grace period; capping to %v", timeoutDuration, shutdownTimeout)
+		timeoutDuration = shutdownTimeout
+	}
+
 	// Create context with timeout
 	ctx, cancel := context.WithTimeout(context.Background(), timeoutDuration)
 	defer cancel()

pkg/picod/execute_test.go

@@ -395,6 +395,40 @@ func TestExecuteHandler_TimeoutHandling(t *testing.T) {
 	assert.Contains(t, resp.Stderr, "Command timed out")
 }
 
+func TestExecuteHandler_TimeoutCapping(t *testing.T) {
+	server, tmpDir := setupExecuteTestServer(t)
+	defer os.RemoveAll(tmpDir)
+	defer os.Unsetenv(PublicKeyEnvVar)
+
+	// Set server shutdown timeout very low
+	server.config.ShutdownTimeout = 200 * time.Millisecond
+
+	// Command would take 2 seconds, but timeout is requested as 10s.
+	// Since 10s exceeds ShutdownTimeout (200ms), it should be capped to 200ms
+	// and therefore time out.
+	req := ExecuteRequest{
+		Command: []string{"sleep", "2"},
+		Timeout: "10s",
+	}
+	body, _ := json.Marshal(req)
+
+	w := httptest.NewRecorder()
+	c, _ := gin.CreateTestContext(w)
+	c.Request, _ = http.NewRequest("POST", "/api/execute", bytes.NewBuffer(body))
+	c.Request.Header.Set("Content-Type", "application/json")
+
+	server.ExecuteHandler(c)
+
+	assert.Equal(t, http.StatusOK, w.Code)
+
+	var resp ExecuteResponse
+	err := json.Unmarshal(w.Body.Bytes(), &resp)
+	require.NoError(t, err)
+	// It should have timed out due to capping
+	assert.Equal(t, TimeoutExitCode, resp.ExitCode)
+	assert.Contains(t, resp.Stderr, "Command timed out")
+}
+
 func TestExecuteHandler_EnvironmentVariables(t *testing.T) {
 	server, tmpDir := setupExecuteTestServer(t)
 	defer os.RemoveAll(tmpDir)

pkg/picod/server.go

@@ -29,8 +29,9 @@ import (
 
 // Config defines server configuration
 type Config struct {
-	Port      int    `json:"port"`
-	Workspace string `json:"workspace"`
+	Port            int           `json:"port"`
+	Workspace       string        `json:"workspace"`
+	ShutdownTimeout time.Duration `json:"shutdown_timeout"`
 }
 
 // Server defines the PicoD HTTP server
@@ -44,6 +45,9 @@ type Server struct {
 
 // NewServer creates a new PicoD server instance
 func NewServer(config Config) *Server {
+	if config.ShutdownTimeout <= 0 {
+		config.ShutdownTimeout = 90 * time.Second
+	}
 	s := &Server{
 		config:      config,
 		startTime:   time.Now(),
@@ -110,8 +114,12 @@ func (s *Server) Run(ctx context.Context) error {
 	idleConnsClosed := make(chan struct{})
 	go func() {
 		<-ctx.Done()
-		// Allow enough time for in-flight commands to finish (default timeout is 60s).
-		shutdownCtx, cancel := context.WithTimeout(context.Background(), 90*time.Second)
+		// Allow enough time for in-flight commands to finish.
+		shutdownTimeout := s.config.ShutdownTimeout
+		if shutdownTimeout <= 0 {
+			shutdownTimeout = 90 * time.Second
+		}
+		shutdownCtx, cancel := context.WithTimeout(context.Background(), shutdownTimeout)
 		defer cancel()
 		if err := server.Shutdown(shutdownCtx); err != nil {
 			klog.Errorf("HTTP server shutdown error: %v", err)