Harden Scale to Zero Lib

server/cmd/api/api/chromium.go

@@ -28,8 +28,11 @@ func (s *ApiService) UploadExtensionsAndRestart(ctx context.Context, request oap
 	start := time.Now()
 	log.Info("upload extensions: begin")
 
-	s.stz.Disable(ctx)
-	defer s.stz.Enable(ctx)
+	if err := s.stz.Disable(ctx); err != nil {
+		log.Error("failed to disable scale-to-zero", "error", err)
+		return oapi.UploadExtensionsAndRestart500JSONResponse{InternalErrorJSONResponse: oapi.InternalErrorJSONResponse{Message: "failed to disable scale-to-zero"}}, nil
+	}
+	defer s.stz.Enable(context.WithoutCancel(ctx))
 
 	if request.Body == nil {
 		return oapi.UploadExtensionsAndRestart400JSONResponse{BadRequestErrorJSONResponse: oapi.BadRequestErrorJSONResponse{Message: "request body required"}}, nil
@@ -275,8 +278,11 @@ func (s *ApiService) PatchChromiumFlags(ctx context.Context, request oapi.PatchC
 	start := time.Now()
 	log.Info("patch chromium flags: begin")
 
-	s.stz.Disable(ctx)
-	defer s.stz.Enable(ctx)
+	if err := s.stz.Disable(ctx); err != nil {
+		log.Error("failed to disable scale-to-zero", "error", err)
+		return oapi.PatchChromiumFlags500JSONResponse{InternalErrorJSONResponse: oapi.InternalErrorJSONResponse{Message: "failed to disable scale-to-zero"}}, nil
+	}
+	defer s.stz.Enable(context.WithoutCancel(ctx))
 
 	if request.Body == nil {
 		return oapi.PatchChromiumFlags400JSONResponse{BadRequestErrorJSONResponse: oapi.BadRequestErrorJSONResponse{Message: "request body required"}}, nil

server/cmd/api/api/computer.go

@@ -18,8 +18,11 @@ import (
 func (s *ApiService) MoveMouse(ctx context.Context, request oapi.MoveMouseRequestObject) (oapi.MoveMouseResponseObject, error) {
 	log := logger.FromContext(ctx)
 
-	s.stz.Disable(ctx)
-	defer s.stz.Enable(ctx)
+	if err := s.stz.Disable(ctx); err != nil {
+		log.Error("failed to disable scale-to-zero", "error", err)
+		return oapi.MoveMouse500JSONResponse{InternalErrorJSONResponse: oapi.InternalErrorJSONResponse{Message: "failed to disable scale-to-zero"}}, nil
+	}
+	defer s.stz.Enable(context.WithoutCancel(ctx))
 
 	// serialize input operations to avoid overlapping xdotool commands
 	s.inputMu.Lock()
@@ -90,8 +93,11 @@ func (s *ApiService) MoveMouse(ctx context.Context, request oapi.MoveMouseReques
 func (s *ApiService) ClickMouse(ctx context.Context, request oapi.ClickMouseRequestObject) (oapi.ClickMouseResponseObject, error) {
 	log := logger.FromContext(ctx)
 
-	s.stz.Disable(ctx)
-	defer s.stz.Enable(ctx)
+	if err := s.stz.Disable(ctx); err != nil {
+		log.Error("failed to disable scale-to-zero", "error", err)
+		return oapi.ClickMouse500JSONResponse{InternalErrorJSONResponse: oapi.InternalErrorJSONResponse{Message: "failed to disable scale-to-zero"}}, nil
+	}
+	defer s.stz.Enable(context.WithoutCancel(ctx))
 
 	// serialize input operations to avoid overlapping xdotool commands
 	s.inputMu.Lock()
@@ -211,8 +217,11 @@ func (s *ApiService) ClickMouse(ctx context.Context, request oapi.ClickMouseRequ
 func (s *ApiService) TakeScreenshot(ctx context.Context, request oapi.TakeScreenshotRequestObject) (oapi.TakeScreenshotResponseObject, error) {
 	log := logger.FromContext(ctx)
 
-	s.stz.Disable(ctx)
-	defer s.stz.Enable(ctx)
+	if err := s.stz.Disable(ctx); err != nil {
+		log.Error("failed to disable scale-to-zero", "error", err)
+		return oapi.TakeScreenshot500JSONResponse{InternalErrorJSONResponse: oapi.InternalErrorJSONResponse{Message: "failed to disable scale-to-zero"}}, nil
+	}
+	defer s.stz.Enable(context.WithoutCancel(ctx))
 
 	// serialize input operations to avoid race with other input/screen actions
 	s.inputMu.Lock()
@@ -331,8 +340,11 @@ func (s *ApiService) TakeScreenshot(ctx context.Context, request oapi.TakeScreen
 func (s *ApiService) TypeText(ctx context.Context, request oapi.TypeTextRequestObject) (oapi.TypeTextResponseObject, error) {
 	log := logger.FromContext(ctx)
 
-	s.stz.Disable(ctx)
-	defer s.stz.Enable(ctx)
+	if err := s.stz.Disable(ctx); err != nil {
+		log.Error("failed to disable scale-to-zero", "error", err)
+		return oapi.TypeText500JSONResponse{InternalErrorJSONResponse: oapi.InternalErrorJSONResponse{Message: "failed to disable scale-to-zero"}}, nil
+	}
+	defer s.stz.Enable(context.WithoutCancel(ctx))
 
 	// serialize input operations to avoid overlapping xdotool commands
 	s.inputMu.Lock()
@@ -377,8 +389,11 @@ func (s *ApiService) TypeText(ctx context.Context, request oapi.TypeTextRequestO
 func (s *ApiService) PressKey(ctx context.Context, request oapi.PressKeyRequestObject) (oapi.PressKeyResponseObject, error) {
 	log := logger.FromContext(ctx)
 
-	s.stz.Disable(ctx)
-	defer s.stz.Enable(ctx)
+	if err := s.stz.Disable(ctx); err != nil {
+		log.Error("failed to disable scale-to-zero", "error", err)
+		return oapi.PressKey500JSONResponse{InternalErrorJSONResponse: oapi.InternalErrorJSONResponse{Message: "failed to disable scale-to-zero"}}, nil
+	}
+	defer s.stz.Enable(context.WithoutCancel(ctx))
 
 	// serialize input operations to avoid overlapping xdotool commands
 	s.inputMu.Lock()
@@ -487,8 +502,11 @@ func (s *ApiService) PressKey(ctx context.Context, request oapi.PressKeyRequestO
 func (s *ApiService) Scroll(ctx context.Context, request oapi.ScrollRequestObject) (oapi.ScrollResponseObject, error) {
 	log := logger.FromContext(ctx)
 
-	s.stz.Disable(ctx)
-	defer s.stz.Enable(ctx)
+	if err := s.stz.Disable(ctx); err != nil {
+		log.Error("failed to disable scale-to-zero", "error", err)
+		return oapi.Scroll500JSONResponse{InternalErrorJSONResponse: oapi.InternalErrorJSONResponse{Message: "failed to disable scale-to-zero"}}, nil
+	}
+	defer s.stz.Enable(context.WithoutCancel(ctx))
 
 	// serialize input operations to avoid overlapping xdotool commands
 	s.inputMu.Lock()
@@ -576,8 +594,11 @@ func (s *ApiService) Scroll(ctx context.Context, request oapi.ScrollRequestObjec
 func (s *ApiService) DragMouse(ctx context.Context, request oapi.DragMouseRequestObject) (oapi.DragMouseResponseObject, error) {
 	log := logger.FromContext(ctx)
 
-	s.stz.Disable(ctx)
-	defer s.stz.Enable(ctx)
+	if err := s.stz.Disable(ctx); err != nil {
+		log.Error("failed to disable scale-to-zero", "error", err)
+		return oapi.DragMouse500JSONResponse{InternalErrorJSONResponse: oapi.InternalErrorJSONResponse{Message: "failed to disable scale-to-zero"}}, nil
+	}
+	defer s.stz.Enable(context.WithoutCancel(ctx))
 
 	// serialize input operations to avoid overlapping xdotool commands
 	s.inputMu.Lock()

server/cmd/api/api/display.go

@@ -21,8 +21,11 @@ import (
 func (s *ApiService) PatchDisplay(ctx context.Context, req oapi.PatchDisplayRequestObject) (oapi.PatchDisplayResponseObject, error) {
 	log := logger.FromContext(ctx)
 
-	s.stz.Disable(ctx)
-	defer s.stz.Enable(ctx)
+	if err := s.stz.Disable(ctx); err != nil {
+		logger.FromContext(ctx).Error("failed to disable scale-to-zero", "error", err)
+		return oapi.PatchDisplay500JSONResponse{InternalErrorJSONResponse: oapi.InternalErrorJSONResponse{Message: "failed to disable scale-to-zero"}}, nil
+	}
+	defer s.stz.Enable(context.WithoutCancel(ctx))
 
 	if req.Body == nil {
 		return oapi.PatchDisplay400JSONResponse{BadRequestErrorJSONResponse: oapi.BadRequestErrorJSONResponse{Message: "missing request body"}}, nil

server/cmd/api/main.go

@@ -69,7 +69,7 @@ func main() {
 		slogger.Error("invalid default recording parameters", "err", err)
 		os.Exit(1)
 	}
-	stz := scaletozero.NewUnikraftCloudController()
+	stz := scaletozero.NewDebouncedController(scaletozero.NewUnikraftCloudController())
 
 	// DevTools WebSocket upstream manager: tail Chromium supervisord log
 	const chromiumLogPath = "/var/log/supervisord/chromium"

server/lib/devtoolsproxy/proxy.go

@@ -198,7 +198,11 @@ func (u *UpstreamManager) runTailOnce(ctx context.Context) {
 // If logCDPMessages is true, all CDP messages will be logged with their direction.
 func WebSocketProxyHandler(mgr *UpstreamManager, logger *slog.Logger, logCDPMessages bool, ctrl scaletozero.Controller) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		ctrl.Disable(context.WithoutCancel(r.Context()))
+		if err := ctrl.Disable(context.WithoutCancel(r.Context())); err != nil {
+			logger.Error("failed to disable scale-to-zero", "error", err)
+			http.Error(w, "failed to disable scale-to-zero", http.StatusInternalServerError)
+			return
+		}
 		defer ctrl.Enable(context.WithoutCancel(r.Context()))
 
 		upstreamCurrent := mgr.Current()

server/lib/recorder/ffmpeg.go

@@ -152,7 +152,7 @@ func (fr *FFmpegRecorder) Start(ctx context.Context) error {
 
 	args, err := ffmpegArgs(fr.params, fr.outputPath)
 	if err != nil {
-		_ = fr.stz.Enable(ctx)
+		_ = fr.stz.Enable(context.WithoutCancel(ctx))
 		fr.cmd = nil
 		close(fr.exited)
 		fr.mu.Unlock()
@@ -170,7 +170,7 @@ func (fr *FFmpegRecorder) Start(ctx context.Context) error {
 	fr.mu.Unlock()
 
 	if err := cmd.Start(); err != nil {
-		_ = fr.stz.Enable(ctx)
+		_ = fr.stz.Enable(context.WithoutCancel(ctx))
 		fr.mu.Lock()
 		fr.ffmpegErr = err
 		fr.cmd = nil // reset cmd on failure to start so IsRecording() remains correct
@@ -194,7 +194,7 @@ func (fr *FFmpegRecorder) Start(ctx context.Context) error {
 
 // Stop gracefully stops the recording using a multi-phase shutdown process.
 func (fr *FFmpegRecorder) Stop(ctx context.Context) error {
-	defer fr.stz.Enable(ctx)
+	defer fr.stz.Enable(context.WithoutCancel(ctx))
 	// This isn't scientific - give ffmpeg a long time to complete since encoding pipelines can
 	// be complex and we care more about the recording than performance. In cases where ffmpeg
 	// "falls behind" (e.g. it's resource constrained) it's better for our use case to wait for
@@ -211,7 +211,7 @@ func (fr *FFmpegRecorder) Stop(ctx context.Context) error {
 
 // ForceStop immediately terminates the recording process.
 func (fr *FFmpegRecorder) ForceStop(ctx context.Context) error {
-	defer fr.stz.Enable(ctx)
+	defer fr.stz.Enable(context.WithoutCancel(ctx))
 	err := fr.shutdownInPhases(ctx, []shutdownPhase{
 		{"kill", []syscall.Signal{syscall.SIGKILL}, 100 * time.Millisecond, "immediate kill"},
 	})
@@ -348,7 +348,7 @@ func ffmpegArgs(params FFmpegRecordingParams, outputPath string) ([]string, erro
 // waitForCommand should be run in the background to wait for the ffmpeg process to complete and
 // update the internal state accordingly.
 func (fr *FFmpegRecorder) waitForCommand(ctx context.Context) {
-	defer fr.stz.Enable(ctx)
+	defer fr.stz.Enable(context.WithoutCancel(ctx))
 
 	log := logger.FromContext(ctx)
 

server/lib/scaletozero/scaletozero.go

@@ -84,3 +84,53 @@ func (o *Oncer) Enable(ctx context.Context) error {
 	o.enableOnce.Do(func() { o.enableErr = o.ctrl.Enable(ctx) })
 	return o.enableErr
 }
+
+type DebouncedController struct {
+	ctrl        Controller
+	mu          sync.Mutex
+	disabled    bool
+	activeCount int
+}
+
+func NewDebouncedController(ctrl Controller) Controller {
+	return &DebouncedController{ctrl: ctrl}
+}
+
+func (c *DebouncedController) Disable(ctx context.Context) error {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	c.activeCount++
+	if c.disabled {
+		return nil
+	}
+
+	if err := c.ctrl.Disable(ctx); err != nil {
+		c.activeCount--
+		return err
+	}
+
+	c.disabled = true
+	return nil
+}
+
+func (c *DebouncedController) Enable(ctx context.Context) error {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	if c.activeCount > 0 {
+		c.activeCount--
+	}
+
+	// nothing to do
+	if c.activeCount > 0 || !c.disabled {
+		return nil
+	}
+
+	if err := c.ctrl.Enable(ctx); err != nil {
+		return err
+	}
+
+	c.disabled = false
+	return nil
+}