AdaAibaby · Oct 22, 2026
@@ -288,11 +288,12 @@ func isValidVolumeName(name string) bool {
 // created before the database row could be persisted, so we don't leak the
 // underlying directory. It runs in the background and only logs on failure.
 func (a *APIStore) cleanupOrchestratorVolume(ctx context.Context, clusterID uuid.UUID, volume queries.Volume) {
-	go func(ctx context.Context) {
-		if err := a.deleteVolume(ctx, clusterID, volume); err != nil {
-			telemetry.ReportCriticalError(ctx, "failed to clean up volume after failing to persist it", err)
+	bgCtx := context.WithoutCancel(ctx)
+	a.orchestrator.GoTrackedTail(func() {
+		if err := a.deleteVolume(bgCtx, clusterID, volume); err != nil {
+			telemetry.ReportCriticalError(bgCtx, "failed to clean up volume after failing to persist it", err)
 		}
-	}(context.WithoutCancel(ctx))
+	})
 }
 
 func (a *APIStore) createVolume(ctx context.Context, clusterID uuid.UUID, volume queries.Volume) (response *orchestrator.CreateVolumeResponse, err error) {

@@ -42,12 +42,13 @@ func (a *APIStore) DeleteVolumesVolumeID(c *gin.Context, volumeID api.VolumeID)
 		Set("volume_name", volume.Name),
 	)
 
-	go func(ctx context.Context) {
+	bgCtx := context.WithoutCancel(ctx)
+	a.orchestrator.GoTrackedTail(func() {
 		// if this fails, we can clean it up later
-		if err := a.deleteVolume(ctx, clusterID, volume); err != nil {
-			telemetry.ReportCriticalError(ctx, fmt.Sprintf("failed to delete data in volume %q", volume.ID.String()), err)
+		if err := a.deleteVolume(bgCtx, clusterID, volume); err != nil {
+			telemetry.ReportCriticalError(bgCtx, fmt.Sprintf("failed to delete data in volume %q", volume.ID.String()), err)
 		}
-	}(context.WithoutCancel(ctx))
+	})
 
 	c.Status(http.StatusNoContent)
 }

@@ -630,7 +630,7 @@ func (o *Orchestrator) CreateSandbox(
 		// Clean up the sandbox from the node
 		// Copy to a new variable to avoid race conditions
 		sbxToRemove := sbx
-		go func() {
+		o.GoTrackedTail(func() {
 			killErr := o.removeSandboxFromNode(
 				context.WithoutCancel(ctx),
 				sbxToRemove,
@@ -645,7 +645,7 @@ func (o *Orchestrator) CreateSandbox(
 					zap.String("kill_reason", sandbox.KillReasonUnknown.String()),
 				)
 			}
-		}()
+		})
 
 		return sandbox.Sandbox{}, &api.APIError{
 			Code:      http.StatusInternalServerError,

@@ -122,7 +122,8 @@ func (o *Orchestrator) RemoveSandbox(ctx context.Context, teamID uuid.UUID, sand
 
 		if time.Since(sbx.EndTime) > sandbox.StaleCutoff && opts.Action.Effect == sandbox.TransitionExpires {
 			o.sandboxStore.Remove(context.WithoutCancel(ctx), teamID, sandboxID)
-			go o.analyticsRemove(context.WithoutCancel(ctx), sbx, opts.Action)
+			bgCtx := context.WithoutCancel(ctx)
+			o.GoTrackedTail(func() { o.analyticsRemove(bgCtx, sbx, opts.Action) })
 		}
 
 		return nil
@@ -147,7 +148,8 @@ func (o *Orchestrator) RemoveSandbox(ctx context.Context, teamID uuid.UUID, sand
 			return
 		}
 		o.sandboxStore.Remove(context.WithoutCancel(ctx), teamID, sandboxID)
-		go o.analyticsRemove(context.WithoutCancel(ctx), sbx, opts.Action)
+		bgCtx := context.WithoutCancel(ctx)
+		o.GoTrackedTail(func() { o.analyticsRemove(bgCtx, sbx, opts.Action) })
 	}()
 	err = o.removeSandboxFromNode(ctx, sbx, opts.Action, opts.Reason, opts.FilesystemOnly)
 	if err != nil {

@@ -0,0 +1,77 @@
+package orchestrator
+
+import (
+	"context"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// A zero-value Orchestrator is enough: GoTrackedTail/TrackWork/Drain only touch
+// drainMu, draining, work, and outstandingWork, all usable at their zero value.
+
+// TestGoTrackedTail_DrainWaitsForTail proves a tail started before the drain is
+// registered with TrackWork, so Drain blocks until it finishes.
+func TestGoTrackedTail_DrainWaitsForTail(t *testing.T) {
+	t.Parallel()
+
+	var o Orchestrator
+	var done atomic.Bool
+	release := make(chan struct{})
+
+	o.GoTrackedTail(func() {
+		<-release
+		done.Store(true)
+	})
+
+	// Outstanding work is visible synchronously (TrackWork.Add ran before return).
+	require.Equal(t, int64(1), o.OutstandingWork(), "tail must be tracked before drain")
+
+	drainReturned := make(chan error, 1)
+	go func() {
+		drainReturned <- o.Drain(context.Background())
+	}()
+
+	// Drain must still be blocked while the tail is in flight.
+	select {
+	case <-drainReturned:
+		t.Fatal("Drain returned before the tracked tail finished")
+	case <-time.After(50 * time.Millisecond):
+	}
+
+	close(release)
+
+	select {
+	case err := <-drainReturned:
+		require.NoError(t, err)
+		assert.True(t, done.Load(), "tail should have completed")
+	case <-time.After(2 * time.Second):
+		t.Fatal("Drain did not return after the tail finished")
+	}
+}
+
+// TestGoTrackedTail_RunsUntrackedDuringDrain proves that once the drain has
+// started (TrackWork refuses new counts) a tail still runs — fire-and-forget,
+// never dropped — and never panics by adding to the WaitGroup after Wait.
+func TestGoTrackedTail_RunsUntrackedDuringDrain(t *testing.T) {
+	t.Parallel()
+
+	var o Orchestrator
+
+	// Drain with nothing in flight returns immediately and flips draining=true.
+	require.NoError(t, o.Drain(context.Background()))
+
+	ran := make(chan struct{})
+	assert.NotPanics(t, func() {
+		o.GoTrackedTail(func() { close(ran) })
+	})
+
+	select {
+	case <-ran:
+	case <-time.After(2 * time.Second):
+		t.Fatal("tail started during drain never ran")
+	}
+}
@@ -218,6 +218,7 @@ func New(
 		sandbox.Callbacks{
 			AsyncNewlyCreatedSandbox: o.handleNewlyCreatedSandbox,
 			KillOrphanSandbox:        o.killOrphanSandbox,
+			GoTracked:                o.GoTrackedTail,
 		},
 	)
 
@@ -330,6 +331,28 @@ func (o *Orchestrator) TrackWork() (func(), bool) {
 	}, true
 }
 
+// GoTrackedTail runs fn in a background goroutine that is the tail of an
+// already-admitted operation (a billing/analytics event, an orphan-sandbox
+// kill, a volume cleanup) rather than new work. When the drain has not started
+// the goroutine is registered with TrackWork, so Drain waits for it. Once the
+// drain has started TrackWork refuses new counts (adding to the WaitGroup after
+// Drain called Wait would panic), so fn still runs — untracked, fire-and-forget
+// — because dropping it would lose a billing event or leak a sandbox. Callers
+// must bind any request-scoped context with context.WithoutCancel before
+// calling, since the caller may return first.
+func (o *Orchestrator) GoTrackedTail(fn func()) {
+	if release, ok := o.TrackWork(); ok {
+		go func() {
+			defer release()
+			fn()
+		}()
+
+		return
+	}
+
+	go fn()
+}
+
 func (o *Orchestrator) OutstandingWork() int64 {
 	return o.outstandingWork.Load()
 }

@@ -45,6 +45,10 @@ type Callbacks struct {
 	// KillOrphanSandbox kills an orphaned sandbox on the orchestrator node via gRPC.
 	// Used during sync when the Redis backend detects sandboxes running on a node but not present in the store.
 	KillOrphanSandbox OrphanCallback
+	// GoTracked, if non-nil, runs the asynchronous AsyncNewlyCreatedSandbox event
+	// as a drain-tracked tail so the shutdown path waits for in-flight creation
+	// analytics instead of abandoning them. Nil falls back to a bare goroutine.
+	GoTracked func(func())
 }
 
 type Store struct {
@@ -88,7 +92,13 @@ func (s *Store) Add(ctx context.Context, sandbox Sandbox, creation *CreationMeta
 
 	if creation != nil {
 		meta := *creation
-		go s.callbacks.AsyncNewlyCreatedSandbox(context.WithoutCancel(ctx), sandbox, meta)
+		bgCtx := context.WithoutCancel(ctx)
+		event := func() { s.callbacks.AsyncNewlyCreatedSandbox(bgCtx, sandbox, meta) }
+		if s.callbacks.GoTracked != nil {
+			s.callbacks.GoTracked(event)
+		} else {
+			go event()
+		}
 	}
 
 	return nil