diff --git a/packages/api/internal/orchestrator/checkpoint_instance.go b/packages/api/internal/orchestrator/checkpoint_instance.go index b4d5d537b2..5ce99ccea3 100644 --- a/packages/api/internal/orchestrator/checkpoint_instance.go +++ b/packages/api/internal/orchestrator/checkpoint_instance.go @@ -44,6 +44,12 @@ func (o *Orchestrator) CheckpointSandbox(ctx context.Context, teamID uuid.UUID, } } + // Past this point the checkpoint is committed: the sandbox will be paused + // on its node, and cancelling the checkpoint mid-snapshot kills it. Detach + // from the caller's cancellation (e.g. client disconnect mid-fork) so the + // checkpoint always runs to completion once started. + ctx = context.WithoutCancel(ctx) + // finish completes the snapshotting transition exactly once. // On success (nil) it restores the sandbox to Running. // On error it leaves the state as Snapshotting so that @@ -51,7 +57,7 @@ func (o *Orchestrator) CheckpointSandbox(ctx context.Context, teamID uuid.UUID, var once sync.Once finish := func(err error) { once.Do(func() { - finishSnapshotting(context.WithoutCancel(ctx), err) + finishSnapshotting(ctx, err) }) } defer finish(nil) @@ -77,11 +83,7 @@ func (o *Orchestrator) CheckpointSandbox(ctx context.Context, teamID uuid.UUID, Metadata: map[string]string{storageopts.ObjectMetadataTemplateID: upsertResult.TemplateID}, }) if err != nil { - // Cleanup must run even when the checkpoint failed because this - // request's context was cancelled (e.g. client disconnect mid-fork). - cleanupCtx := context.WithoutCancel(ctx) - - o.failSnapshotBuild(cleanupCtx, upsertResult.BuildID, err) + o.failSnapshotBuild(ctx, upsertResult.BuildID, err) // The orchestrator rejects these before pausing the VM (envd too old, // starting-sandboxes queue full), so the sandbox is still running @@ -104,8 +106,8 @@ func (o *Orchestrator) CheckpointSandbox(ctx context.Context, teamID uuid.UUID, // so RemoveSandbox can proceed without deadlock. finish(err) - if killErr := o.RemoveSandbox(cleanupCtx, teamID, sandboxID, sandbox.RemoveOpts{Action: sandbox.StateActionKill}); killErr != nil { - telemetry.ReportError(cleanupCtx, "error killing sandbox after failed checkpoint", killErr) + if killErr := o.RemoveSandbox(ctx, teamID, sandboxID, sandbox.RemoveOpts{Action: sandbox.StateActionKill}); killErr != nil { + telemetry.ReportError(ctx, "error killing sandbox after failed checkpoint", killErr) } return fmt.Errorf("checkpoint failed: %w", err) @@ -119,10 +121,15 @@ func (o *Orchestrator) CheckpointSandbox(ctx context.Context, teamID uuid.UUID, BuildID: upsertResult.BuildID, }) if err != nil { + // The sandbox itself resumed fine (finish(nil) restores it to Running), + // but without a success status the snapshot is unusable: mark the build + // failed so it doesn't dangle in Snapshotting forever. + o.failSnapshotBuild(ctx, upsertResult.BuildID, err) + return fmt.Errorf("error updating build status: %w", err) } - o.snapshotCache.Invalidate(context.WithoutCancel(ctx), sandboxID) + o.snapshotCache.Invalidate(ctx, sandboxID) telemetry.ReportEvent(ctx, "Checkpointed sandbox") diff --git a/packages/api/internal/orchestrator/delete_instance.go b/packages/api/internal/orchestrator/delete_instance.go index fb9b2c7056..48e7939df6 100644 --- a/packages/api/internal/orchestrator/delete_instance.go +++ b/packages/api/internal/orchestrator/delete_instance.go @@ -85,24 +85,31 @@ func (o *Orchestrator) RemoveSandbox(ctx context.Context, teamID uuid.UUID, sand return ErrSandboxOperationFailed } } + // Past this point the removal is committed. Detach from the caller's + // cancellation (e.g. client disconnect) so a pause can't be aborted + // mid-snapshot — the node would kill the sandbox without a persisted + // snapshot — and a kill can't leave the VM running on the node after + // its store entry is removed below. + ctx = context.WithoutCancel(ctx) + defer func() { - finish(context.WithoutCancel(ctx), err) + finish(ctx, err) }() if alreadyDone { logger.L().Info(ctx, "Sandbox was already in the process of being removed", logger.WithSandboxID(sandboxID), zap.String("state", string(sbx.State))) 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) + o.sandboxStore.Remove(ctx, teamID, sandboxID) + go o.analyticsRemove(ctx, sbx, opts.Action) } return nil } - defer func() { go o.analyticsRemove(context.WithoutCancel(ctx), sbx, opts.Action) }() + defer func() { go o.analyticsRemove(ctx, sbx, opts.Action) }() // Once we start the removal process, we want to make sure it gets removed from the store - defer o.sandboxStore.Remove(context.WithoutCancel(ctx), teamID, sandboxID) + defer o.sandboxStore.Remove(ctx, teamID, sandboxID) err = o.removeSandboxFromNode(ctx, sbx, opts.Action, opts.Reason, opts.FilesystemOnly) if err != nil { fields := []zap.Field{ diff --git a/packages/api/internal/orchestrator/snapshot_template.go b/packages/api/internal/orchestrator/snapshot_template.go index 7c648e2eb4..b6f71be4af 100644 --- a/packages/api/internal/orchestrator/snapshot_template.go +++ b/packages/api/internal/orchestrator/snapshot_template.go @@ -52,6 +52,12 @@ func (o *Orchestrator) CreateSnapshotTemplate(ctx context.Context, teamID uuid.U } } + // Past this point the snapshot is committed: the sandbox will be paused + // on its node, and cancelling the checkpoint mid-snapshot kills it. Detach + // from the caller's cancellation (e.g. client disconnect) so the snapshot + // always runs to completion once started. + ctx = context.WithoutCancel(ctx) + // finish completes the snapshotting transition exactly once. // On success (nil) it restores the sandbox to Running. // On error it leaves the state as Snapshotting so that @@ -59,7 +65,7 @@ func (o *Orchestrator) CreateSnapshotTemplate(ctx context.Context, teamID uuid.U var once sync.Once finish := func(err error) { once.Do(func() { - finishSnapshotting(context.WithoutCancel(ctx), err) + finishSnapshotting(ctx, err) }) } defer finish(nil) @@ -116,10 +122,15 @@ func (o *Orchestrator) CreateSnapshotTemplate(ctx context.Context, teamID uuid.U BuildID: upsertResult.BuildID, }) if err != nil { + // The sandbox itself resumed fine (finish(nil) restores it to Running), + // but without an uploaded status the snapshot is unusable: mark the + // build failed so it doesn't dangle in Snapshotting forever. + o.failSnapshotBuild(ctx, upsertResult.BuildID, err) + return SnapshotTemplateResult{}, fmt.Errorf("error updating build status: %w", err) } - o.snapshotCache.Invalidate(context.WithoutCancel(ctx), sandboxID) + o.snapshotCache.Invalidate(ctx, sandboxID) telemetry.ReportEvent(ctx, "Snapshot template completed")