diff --git a/server/cmd/api/api/process.go b/server/cmd/api/api/process.go index c8f5d4af..16fc6280 100644 --- a/server/cmd/api/api/process.go +++ b/server/cmd/api/api/process.go @@ -240,6 +240,15 @@ func (s *ApiService) ProcessSpawn(ctx context.Context, request oapi.ProcessSpawn return oapi.ProcessSpawn500JSONResponse{InternalErrorJSONResponse: oapi.InternalErrorJSONResponse{Message: "failed to start process"}}, nil } + // Disable scale-to-zero while the process is running. + // Track success so we only re-enable if disable succeeded. + stzDisabled := false + if err := s.stz.Disable(ctx); err != nil { + log.Error("failed to disable scale-to-zero", "err", err) + } else { + stzDisabled = true + } + id := openapi_types.UUID(uuid.New()) h := &processHandle{ id: id, @@ -294,8 +303,10 @@ func (s *ApiService) ProcessSpawn(ctx context.Context, request oapi.ProcessSpawn } }() - // Waiter goroutine - go func() { + // Waiter goroutine - use context without cancel since HTTP request may complete + // before the process exits + stzCtx := context.WithoutCancel(ctx) + go func(stzWasDisabled bool) { err := cmd.Wait() code := 0 if err != nil { @@ -310,6 +321,15 @@ func (s *ApiService) ProcessSpawn(ctx context.Context, request oapi.ProcessSpawn } } h.setExited(code) + + // Re-enable scale-to-zero now that the process has exited, + // but only if we successfully disabled it earlier + if stzWasDisabled { + if err := s.stz.Enable(stzCtx); err != nil { + log.Error("failed to enable scale-to-zero", "err", err) + } + } + // Send exit event evt := oapi.ProcessStreamEventEvent("exit") h.outCh <- oapi.ProcessStreamEvent{Event: &evt, ExitCode: &code} @@ -325,7 +345,7 @@ func (s *ApiService) ProcessSpawn(ctx context.Context, request oapi.ProcessSpawn delete(s.procs, procID) s.procMu.Unlock() }(id.String()) - }() + }(stzDisabled) startedAt := h.started pid := h.pid diff --git a/server/cmd/api/api/process_test.go b/server/cmd/api/api/process_test.go index 2204eba9..844a8aff 100644 --- a/server/cmd/api/api/process_test.go +++ b/server/cmd/api/api/process_test.go @@ -15,6 +15,7 @@ import ( "github.com/google/uuid" openapi_types "github.com/oapi-codegen/runtime/types" oapi "github.com/onkernel/kernel-images/server/lib/oapi" + "github.com/onkernel/kernel-images/server/lib/scaletozero" "github.com/stretchr/testify/require" ) @@ -43,7 +44,7 @@ func TestProcessExec(t *testing.T) { func TestProcessSpawnStatusAndStream(t *testing.T) { t.Parallel() ctx := context.Background() - svc := &ApiService{procs: make(map[string]*processHandle)} + svc := &ApiService{procs: make(map[string]*processHandle), stz: scaletozero.NewNoopController()} // Spawn a short-lived process that emits stdout and stderr then exits cmd := "sh" @@ -111,7 +112,7 @@ func TestProcessSpawnStatusAndStream(t *testing.T) { func TestProcessStdinAndExit(t *testing.T) { t.Parallel() ctx := context.Background() - svc := &ApiService{procs: make(map[string]*processHandle)} + svc := &ApiService{procs: make(map[string]*processHandle), stz: scaletozero.NewNoopController()} // Spawn a process that reads exactly 3 bytes then exits cmd := "sh" @@ -150,7 +151,7 @@ func TestProcessStdinAndExit(t *testing.T) { func TestProcessKill(t *testing.T) { t.Parallel() ctx := context.Background() - svc := &ApiService{procs: make(map[string]*processHandle)} + svc := &ApiService{procs: make(map[string]*processHandle), stz: scaletozero.NewNoopController()} cmd := "sh" args := []string{"-c", "sleep 5"}