Skip to content

fix(api): don't kill sandboxes when snapshot requests are cancelled mid-flight #3302

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
mishushakov wants to merge 1 commit into
base: main
Choose a base branch
from
+41 −16
Draft

fix(api): don't kill sandboxes when snapshot requests are cancelled mid-flight #3302

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 16 additions & 9 deletions packages/api/internal/orchestrator/checkpoint_instance.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,14 +44,20 @@ 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)
Comment on lines +47 to +51

@gemini-code-assist gemini-code-assist Bot Jul 17, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Detaching the context from the caller's cancellation using context.WithoutCancel also strips any associated deadlines or timeouts. Without a safety timeout, the subsequent gRPC and database calls could hang indefinitely if the node or network becomes unresponsive. It is highly recommended to wrap the detached context with a reasonable safety timeout.

Suggested change
// 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)
// 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.
// Use a safety timeout to prevent the orchestrator from hanging indefinitely if the node is unresponsive.
ctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Minute)
defer cancel()


// 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
// RemoveSandbox can transition directly to Killing.
var once sync.Once
finish := func(err error) {
once.Do(func() {
finishSnapshotting(context.WithoutCancel(ctx), err)
finishSnapshotting(ctx, err)
})
}
defer finish(nil)
Expand All @@ -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
Expand All @@ -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)
Expand All @@ -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)

@gemini-code-assist gemini-code-assist Bot Jul 17, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

critical

The helper method failSnapshotBuild (defined in snapshot_template.go at line 143) contains a compilation error at line 146: FinishedAt: new(time.Now()). In Go, new() takes a type argument, not a value expression. This will cause the build to fail.

Please fix failSnapshotBuild in snapshot_template.go by passing a pointer to a local time.Time variable instead:

func (o *Orchestrator) failSnapshotBuild(ctx context.Context, buildID uuid.UUID, cause error) {
	now := time.Now()
	err := o.sqlcDB.UpdateEnvBuildStatus(ctx, queries.UpdateEnvBuildStatusParams{
		Status:     types.BuildStatusFailed,
		FinishedAt: &now,
		Reason:     types.BuildReason{Message: cause.Error()},
		BuildID:    buildID,
	})
	if err != nil {
		telemetry.ReportError(ctx, "error failing build", 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")

Expand Down
17 changes: 12 additions & 5 deletions packages/api/internal/orchestrator/delete_instance.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)

@gemini-code-assist gemini-code-assist Bot Jul 17, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Since ctx has been detached from caller cancellation and has no timeout, calling removeSandboxFromNode could block indefinitely if the node is unresponsive. To prevent this, wrap the context passed to removeSandboxFromNode with a safety timeout. This keeps the other cleanup operations (like store removal and analytics) safe from early cancellation while ensuring the node call doesn't hang forever.

nodeCtx, cancel := context.WithTimeout(ctx, 3*time.Minute)
defer cancel()
err = o.removeSandboxFromNode(nodeCtx, sbx, opts.Action, opts.Reason, opts.FilesystemOnly)

if err != nil {
fields := []zap.Field{
Expand Down
15 changes: 13 additions & 2 deletions packages/api/internal/orchestrator/snapshot_template.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,14 +52,20 @@ 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)
Comment on lines +55 to +59

@gemini-code-assist gemini-code-assist Bot Jul 17, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Detaching the context from the caller's cancellation using context.WithoutCancel also strips any associated deadlines or timeouts. Without a safety timeout, the subsequent gRPC and database calls could hang indefinitely if the node or network becomes unresponsive. It is highly recommended to wrap the detached context with a reasonable safety timeout.

Suggested change
// 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)
// 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.
// Use a safety timeout to prevent the orchestrator from hanging indefinitely if the node is unresponsive.
ctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Minute)
defer cancel()


// 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
// RemoveSandbox can transition directly to Killing.
var once sync.Once
finish := func(err error) {
once.Do(func() {
finishSnapshotting(context.WithoutCancel(ctx), err)
finishSnapshotting(ctx, err)
})
}
defer finish(nil)
Expand Down Expand Up @@ -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")

Expand Down
Loading