fix(orchestrator): drain in-flight snapshot uploads on shutdown

Instead of cancelling the detached upload goroutine when the server shuts
down, wait for in-flight uploads to finish so a graceful restart doesn't
drop a snapshot that is still uploading.

- Track async uploads with a WaitGroup (uploadsWG).
- Server.Close now takes a context and waits on the WaitGroup, bounded by
  that context so a forced stop (cancelled close context) still exits
  promptly. The gRPC server is closed before this, so no new uploads start
  during the wait.

Author: jakubno

packages/orchestrator/pkg/factories/run.go

@@ -575,8 +575,8 @@ func run(config cfg.Config, opts Options) (success bool) {
 	if err != nil {
 		logger.L().Fatal(ctx, "failed to create orchestrator server", zap.Error(err))
 	}
-	closers = append(closers, closer{"orchestrator server", func(context.Context) error {
-		return orchestratorService.Close()
+	closers = append(closers, closer{"orchestrator server", func(ctx context.Context) error {
+		return orchestratorService.Close(ctx)
 	}})
 
 	// template manager sandbox logger

packages/orchestrator/pkg/server/main.go

@@ -60,6 +60,10 @@ type Server struct {
 	sandboxKilledCounter  metric.Int64Counter
 	uploadFailedCounter   metric.Int64Counter
 
+	// uploadsWG tracks in-flight async snapshot uploads so a graceful shutdown
+	// can wait for them to finish instead of dropping them.
+	uploadsWG sync.WaitGroup
+
 	done      chan struct{}
 	closeOnce sync.Once
 }
@@ -163,11 +167,26 @@ func New(ctx context.Context, cfg ServiceConfig) (*Server, error) {
 	return server, nil
 }
 
-func (s *Server) Close() error {
+func (s *Server) Close(ctx context.Context) error {
 	s.closeOnce.Do(func() {
 		close(s.done)
 	})
 
+	// Wait for in-flight snapshot uploads to finish so a graceful shutdown
+	// doesn't drop a snapshot that is still uploading. ctx is cancelled on a
+	// forced stop, in which case we stop waiting and let the process exit.
+	uploadsDone := make(chan struct{})
+	go func() {
+		s.uploadsWG.Wait()
+		close(uploadsDone)
+	}()
+
+	select {
+	case <-uploadsDone:
+	case <-ctx.Done():
+		logger.L().Warn(ctx, "shutting down with snapshot uploads still in flight", zap.Error(context.Cause(ctx)))
+	}
+
 	s.uploadedBuilds.Stop()
 
 	return nil

packages/orchestrator/pkg/server/sandboxes.go

@@ -904,23 +904,14 @@ func (s *Server) snapshotAndCacheSandbox(
 // background and cleans up the Redis peer key once done. Used by the Pause
 // handler where no prefetch data is available.
 func (s *Server) uploadSnapshotAsync(ctx context.Context, sbx *sandbox.Sandbox, res *snapshotResult) {
-	// Detach from the request (the upload retries for up to uploadTotalBudget),
-	// but cancel on server shutdown so the background loop doesn't outlive the
-	// process.
-	uploadCtx, cancel := context.WithCancel(context.WithoutCancel(ctx))
+	// Detach from the request: the upload retries for up to uploadTotalBudget.
+	// A graceful shutdown waits for it to finish (see Server.Close via uploadsWG)
+	// rather than cancelling, so an in-flight snapshot isn't dropped on restart.
+	uploadCtx := context.WithoutCancel(ctx)
 
-	// Watcher: cancel on shutdown. Exits when the upload finishes (uploadCtx
-	// cancelled by the worker's defer), so it never leaks.
+	s.uploadsWG.Add(1)
 	go func() {
-		select {
-		case <-s.done:
-			cancel()
-		case <-uploadCtx.Done():
-		}
-	}()
-
-	go func() {
-		defer cancel()
+		defer s.uploadsWG.Done()
 
 		spanCtx, span := tracer.Start(uploadCtx, "upload snapshot")
 		defer span.End()