From 9dcd7f575afc391a8eae6a362ee7c138ed32ba33 Mon Sep 17 00:00:00 2001 From: AdaAibaby <8872412+adababys@users.noreply.github.com> Date: Thu, 22 Oct 2026 15:08:18 +0800 Subject: [PATCH] fix(api): drain in-flight background tails on shutdown via existing TrackWork The API shutdown path drains HTTP/gRPC handlers and then waits on the orchestrators tracked work (apiStore.Drain -> o.work). But several goroutines spawned during request handling with a detached context are launched with a bare go and are tracked by nothing, so shutdown abandons them: - store.go: AsyncNewlyCreatedSandbox (create analytics + posthog event) - delete_instance.go x2: analyticsRemove (billing-critical delete event) - create_instance.go: removeSandboxFromNode (kills the orphaned sandbox left when adding it to the store fails) - volume_delete.go / volume_create.go: deleteVolume (removes volume data on the node) Rather than add a second parallel WaitGroup (as the original #3375 did, before main grew TrackWork/Drain), route these through the mechanism that already exists. New GoTrackedTail registers the goroutine with TrackWork when the drain has not started, so Drain waits for it; once the drain has started TrackWork refuses new counts (adding to the WaitGroup after Wait would panic), so the goroutine still runs untracked and fire-and-forget - dropping it would lose a billing event or leak a sandbox. The sandbox store reaches it through a new optional GoTracked callback, keeping the package free of an orchestrator import. No new shutdown step and no new public WaitGroup: the tails now ride the apiStore.Drain already in main.go. Adds tests for both paths (Drain waits for a pre-drain tail; a tail started mid-drain still runs and never panics). Rebased follow-up to #3375 at the maintainers invitation. Closes #3357. --- .../api/internal/handlers/volume_create.go | 9 ++- .../api/internal/handlers/volume_delete.go | 9 ++- .../internal/orchestrator/create_instance.go | 4 +- .../internal/orchestrator/delete_instance.go | 6 +- .../orchestrator/go_tracked_tail_test.go | 77 +++++++++++++++++++ .../api/internal/orchestrator/orchestrator.go | 23 ++++++ packages/api/internal/sandbox/store.go | 12 ++- 7 files changed, 127 insertions(+), 13 deletions(-) create mode 100644 packages/api/internal/orchestrator/go_tracked_tail_test.go diff --git a/packages/api/internal/handlers/volume_create.go b/packages/api/internal/handlers/volume_create.go index 9f8ba66a0c..0df744e7a3 100644 --- a/packages/api/internal/handlers/volume_create.go +++ b/packages/api/internal/handlers/volume_create.go @@ -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) { diff --git a/packages/api/internal/handlers/volume_delete.go b/packages/api/internal/handlers/volume_delete.go index 647eb3de4b..905b16306c 100644 --- a/packages/api/internal/handlers/volume_delete.go +++ b/packages/api/internal/handlers/volume_delete.go @@ -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) } diff --git a/packages/api/internal/orchestrator/create_instance.go b/packages/api/internal/orchestrator/create_instance.go index 67b294afcf..c941ec5428 100644 --- a/packages/api/internal/orchestrator/create_instance.go +++ b/packages/api/internal/orchestrator/create_instance.go @@ -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, diff --git a/packages/api/internal/orchestrator/delete_instance.go b/packages/api/internal/orchestrator/delete_instance.go index ad9ff60df6..1ece751284 100644 --- a/packages/api/internal/orchestrator/delete_instance.go +++ b/packages/api/internal/orchestrator/delete_instance.go @@ -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 { diff --git a/packages/api/internal/orchestrator/go_tracked_tail_test.go b/packages/api/internal/orchestrator/go_tracked_tail_test.go new file mode 100644 index 0000000000..d81d6a4778 --- /dev/null +++ b/packages/api/internal/orchestrator/go_tracked_tail_test.go @@ -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") + } +} diff --git a/packages/api/internal/orchestrator/orchestrator.go b/packages/api/internal/orchestrator/orchestrator.go index 1adac0f6a6..b4bd68e205 100644 --- a/packages/api/internal/orchestrator/orchestrator.go +++ b/packages/api/internal/orchestrator/orchestrator.go @@ -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() } diff --git a/packages/api/internal/sandbox/store.go b/packages/api/internal/sandbox/store.go index 2cddfe4994..bd79787095 100644 --- a/packages/api/internal/sandbox/store.go +++ b/packages/api/internal/sandbox/store.go @@ -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