Skip to content

Commit 643d726

Browse files
mishushakovclaude
andauthored
feat(api): add sandbox fork endpoint (#3202)
## Summary - Adds `POST /sandboxes/{sandboxID}/fork`: checkpoints the running sandbox in place and creates one or more new sandboxes from that snapshot under fresh IDs. - The checkpoint reuses the mechanism behind snapshot templates (`Sandbox.Checkpoint` gRPC): the sandbox is briefly paused on its node, snapshotted to its own `snapshots` row, and resumed under the same execution ID — the original never leaves the sandbox store, so its ID, expiration, node placement, and concurrency reservation are untouched. - `SandboxForkRequest` supports an optional `timeout` (TTL for the forked sandboxes, defaults to the standard 15s sandbox default) and `count` (number of forks, default 1). The snapshot is captured once regardless of count; all forks boot from it in parallel. - Per-fork results: the 201 response is a list of `SandboxForkResult`, each holding either the created `sandbox` or the `error` that prevented it from starting (Promise.allSettled-style). A non-201 status means the request failed before any fork was attempted. Checkpoint always captures full memory state; forking requires a snapshot-capable envd (400 otherwise). - Splits the resume data fetcher into `buildResumeSandboxData` (resume from own snapshot) and `buildResumeSandboxDataFromSnapshot` (boot a different ID from a snapshot), so each forked sandbox's envd access token is generated for the ID it actually runs under. - Threads `SnapshotSandboxID` through `SandboxMetadata` so the resume origin-node remap on placement timeout targets the snapshot's row (previously a silent no-op when the started ID differed from the snapshot's). - Regenerated `packages/api/internal/api/api.gen.go` and the integration test client from the updated OpenAPI spec. ## Test plan - [x] `go build`/`go vet`/`go test -race` pass for `packages/api/...` and `tests/integration` - [x] `make fmt` and scoped `golangci-lint run` on touched packages are clean - [x] Integration tests: happy-path fork (original and fork both running afterwards), multi-fork with `count: 2` (unique IDs, all running, no per-fork errors), invalid `count: 0` (400), forking a paused sandbox (409), forking a killed sandbox (404), cross-team access (404) 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>
1 parent 07df385 commit 643d726

8 files changed

Lines changed: 1245 additions & 196 deletions

File tree

packages/api/internal/api/api.gen.go

Lines changed: 442 additions & 190 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.
Lines changed: 204 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,204 @@
1+
package handlers
2+
3+
import (
4+
"context"
5+
"errors"
6+
"fmt"
7+
"net/http"
8+
"time"
9+
10+
"github.com/gin-gonic/gin"
11+
"github.com/google/uuid"
12+
"go.opentelemetry.io/otel/trace"
13+
"go.uber.org/zap"
14+
"golang.org/x/sync/errgroup"
15+
16+
"github.com/e2b-dev/infra/packages/api/internal/api"
17+
"github.com/e2b-dev/infra/packages/api/internal/orchestrator"
18+
"github.com/e2b-dev/infra/packages/api/internal/sandbox"
19+
"github.com/e2b-dev/infra/packages/api/internal/utils"
20+
"github.com/e2b-dev/infra/packages/auth/pkg/auth"
21+
"github.com/e2b-dev/infra/packages/shared/pkg/ginutils"
22+
"github.com/e2b-dev/infra/packages/shared/pkg/id"
23+
sbxlogger "github.com/e2b-dev/infra/packages/shared/pkg/logger/sandbox"
24+
"github.com/e2b-dev/infra/packages/shared/pkg/telemetry"
25+
sharedUtils "github.com/e2b-dev/infra/packages/shared/pkg/utils"
26+
)
27+
28+
// maxForkCount caps how many sandboxes a single fork request can create,
29+
// bounding the parallel boots (and the result allocation) per request.
30+
const maxForkCount = 100
31+
32+
// PostSandboxesSandboxIDFork forks a running sandbox: it checkpoints the
33+
// sandbox in place (snapshot it and resume it on its node, so the original
34+
// keeps running with its ID and expiration untouched) and creates count new
35+
// sandboxes from that snapshot under fresh IDs. Each fork succeeds or fails
36+
// independently: the response carries one result per requested fork, holding
37+
// either the created sandbox or the error that prevented it from starting.
38+
func (a *APIStore) PostSandboxesSandboxIDFork(c *gin.Context, sandboxID api.SandboxID) {
39+
ctx := c.Request.Context()
40+
41+
teamInfo := auth.MustGetTeamInfo(c)
42+
teamID := teamInfo.Team.ID
43+
44+
sandboxID, err := utils.ShortID(sandboxID)
45+
if err != nil {
46+
a.sendAPIStoreError(c, http.StatusBadRequest, "Invalid sandbox ID")
47+
48+
return
49+
}
50+
51+
span := trace.SpanFromContext(ctx)
52+
span.SetAttributes(telemetry.WithSandboxID(sandboxID))
53+
54+
traceID := span.SpanContext().TraceID().String()
55+
c.Set("traceID", traceID)
56+
57+
body, err := ginutils.ParseOptionalBody[api.PostSandboxesSandboxIDForkJSONRequestBody](ctx, c)
58+
if err != nil {
59+
a.sendAPIStoreError(c, http.StatusBadRequest, fmt.Sprintf("Error when parsing request: %s", err))
60+
61+
return
62+
}
63+
64+
forkTimeout := sandbox.SandboxTimeoutDefault
65+
if body.Timeout != nil {
66+
forkTimeout = time.Duration(*body.Timeout) * time.Second
67+
68+
if forkTimeout > time.Duration(teamInfo.Limits.MaxLengthHours)*time.Hour {
69+
a.sendAPIStoreError(c, http.StatusBadRequest, fmt.Sprintf("Timeout cannot be greater than %d hours", teamInfo.Limits.MaxLengthHours))
70+
71+
return
72+
}
73+
}
74+
75+
forkCount := 1
76+
if body.Count != nil {
77+
forkCount = int(*body.Count)
78+
}
79+
80+
if forkCount < 1 {
81+
a.sendAPIStoreError(c, http.StatusBadRequest, "Count must be at least 1")
82+
83+
return
84+
}
85+
86+
if forkCount > maxForkCount {
87+
a.sendAPIStoreError(c, http.StatusBadRequest, fmt.Sprintf("Count cannot be greater than %d", maxForkCount))
88+
89+
return
90+
}
91+
92+
// The original sandbox keeps running and holds one slot, so more forks
93+
// than the concurrency limit can never succeed.
94+
if int64(forkCount) >= teamInfo.Limits.SandboxConcurrency {
95+
a.sendAPIStoreError(c, http.StatusBadRequest, fmt.Sprintf("Count must be lower than the maximum number of concurrent sandboxes (%d)", teamInfo.Limits.SandboxConcurrency))
96+
97+
return
98+
}
99+
100+
original, err := a.orchestrator.GetSandbox(ctx, teamID, sandboxID)
101+
if err != nil {
102+
if errors.Is(err, sandbox.ErrNotFound) {
103+
apiErr := forkHandleNotRunningSandbox(ctx, a, sandboxID, teamID)
104+
a.sendAPIStoreError(c, apiErr.Code, apiErr.ClientMsg)
105+
106+
return
107+
}
108+
109+
telemetry.ReportError(ctx, "error getting sandbox for fork", err, telemetry.WithSandboxID(sandboxID))
110+
a.sendAPIStoreError(c, http.StatusInternalServerError, "Error forking sandbox")
111+
112+
return
113+
}
114+
115+
if err := sharedUtils.CheckEnvdVersionForSnapshot(original.EnvdVersion); err != nil {
116+
a.sendAPIStoreError(c, http.StatusBadRequest, err.Error())
117+
118+
return
119+
}
120+
121+
// Checkpoint the sandbox in place: it is briefly paused on its node,
122+
// snapshotted, and resumed under the same execution ID, so the original
123+
// keeps its ID, expiration, and concurrency slot.
124+
err = a.orchestrator.CheckpointSandbox(ctx, teamID, sandboxID)
125+
var transErr *sandbox.InvalidStateTransitionError
126+
127+
switch {
128+
case err == nil:
129+
case errors.Is(err, sandbox.ErrNotFound):
130+
apiErr := forkHandleNotRunningSandbox(ctx, a, sandboxID, teamID)
131+
a.sendAPIStoreError(c, apiErr.Code, apiErr.ClientMsg)
132+
133+
return
134+
case errors.As(err, &transErr):
135+
a.sendAPIStoreError(c, http.StatusConflict, fmt.Sprintf("Sandbox '%s' cannot be forked while in '%s' state", sandboxID, transErr.CurrentState))
136+
137+
return
138+
case errors.Is(err, orchestrator.PauseQueueExhaustedError{}):
139+
a.sendAPIStoreError(c, http.StatusServiceUnavailable, fmt.Sprintf("Sandbox '%s' cannot be forked right now because its node is busy, please retry", sandboxID))
140+
141+
return
142+
default:
143+
telemetry.ReportError(ctx, "error checkpointing sandbox for fork", err, telemetry.WithSandboxID(sandboxID))
144+
a.sendAPIStoreError(c, http.StatusInternalServerError, "Error forking sandbox")
145+
146+
return
147+
}
148+
149+
sbxlogger.E(&sbxlogger.SandboxMetadata{
150+
SandboxID: sandboxID,
151+
TemplateID: original.TemplateID,
152+
TeamID: teamID.String(),
153+
}).Debug(ctx, "Creating forked sandboxes from snapshot", zap.Int("count", forkCount))
154+
155+
// All forks boot in parallel from the same immutable snapshot, each
156+
// succeeding or failing independently.
157+
results := make([]api.SandboxForkResult, forkCount)
158+
159+
wg := errgroup.Group{}
160+
for i := range forkCount {
161+
wg.Go(func() error {
162+
forkedSandboxID := InstanceIDPrefix + id.Generate()
163+
164+
forkedSbx, createErr := a.startSandbox(
165+
ctx,
166+
forkedSandboxID,
167+
forkTimeout,
168+
teamInfo,
169+
a.buildResumeSandboxDataFromSnapshot(sandboxID, forkedSandboxID, nil),
170+
&c.Request.Header,
171+
true,
172+
nil, // mcp
173+
)
174+
if createErr != nil {
175+
telemetry.ReportError(ctx, "error creating forked sandbox", createErr.Err, telemetry.WithSandboxID(forkedSandboxID))
176+
results[i] = api.SandboxForkResult{Error: &api.Error{Code: int32(createErr.Code), Message: createErr.ClientMsg}}
177+
178+
//nolint:nilerr // per-fork errors are reported in the result entry, not propagated
179+
return nil
180+
}
181+
182+
results[i] = api.SandboxForkResult{Sandbox: forkedSbx}
183+
184+
return nil
185+
})
186+
}
187+
_ = wg.Wait()
188+
189+
c.JSON(http.StatusCreated, results)
190+
}
191+
192+
// forkHandleNotRunningSandbox classifies a fork request for a sandbox that is
193+
// not running: 409 if it is paused (a snapshot exists), 404 otherwise.
194+
func forkHandleNotRunningSandbox(ctx context.Context, a *APIStore, sandboxID string, teamID uuid.UUID) api.APIError {
195+
apiErr := pauseHandleNotRunningSandbox(ctx, a.snapshotCache, sandboxID, teamID)
196+
switch apiErr.Code {
197+
case http.StatusConflict:
198+
apiErr.ClientMsg = fmt.Sprintf("Sandbox '%s' is paused and cannot be forked; resume it first", sandboxID)
199+
case http.StatusInternalServerError:
200+
apiErr.ClientMsg = "Error forking sandbox"
201+
}
202+
203+
return apiErr
204+
}

packages/api/internal/handlers/sandbox_resume.go

Lines changed: 14 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -186,17 +186,25 @@ func convertDatabaseMountsToOrchestratorMounts(volumes []*types.SandboxVolumeMou
186186
return results
187187
}
188188

189-
// buildResumeSandboxData returns a SandboxDataFetcher that fetches snapshot data
190-
// from the cache and builds SandboxMetadata for resume operations.
191-
// The returned callback is called inside the sandbox lock to prevent race conditions.
189+
// buildResumeSandboxData returns a SandboxDataFetcher for resuming a sandbox
190+
// from its own snapshot.
192191
func (a *APIStore) buildResumeSandboxData(sandboxID string, autoPauseOverride *bool) orchestrator.SandboxDataFetcher {
192+
return a.buildResumeSandboxDataFromSnapshot(sandboxID, sandboxID, autoPauseOverride)
193+
}
194+
195+
// buildResumeSandboxDataFromSnapshot returns a SandboxDataFetcher that fetches
196+
// snapshot data for snapshotSandboxID from the cache and builds SandboxMetadata
197+
// for resume operations. sandboxID is the ID the sandbox will run under — it
198+
// differs from snapshotSandboxID when forking — and scopes the envd access token.
199+
// The returned callback is called inside the sandbox lock to prevent race conditions.
200+
func (a *APIStore) buildResumeSandboxDataFromSnapshot(snapshotSandboxID, sandboxID string, autoPauseOverride *bool) orchestrator.SandboxDataFetcher {
193201
return func(ctx context.Context) (orchestrator.SandboxMetadata, *api.APIError) {
194-
lastSnapshot, err := a.snapshotCache.Get(ctx, sandboxID)
202+
lastSnapshot, err := a.snapshotCache.Get(ctx, snapshotSandboxID)
195203
if err != nil {
196204
return orchestrator.SandboxMetadata{}, &api.APIError{
197205
Code: http.StatusInternalServerError,
198206
ClientMsg: "Error when getting snapshot",
199-
Err: fmt.Errorf("error getting last snapshot for sandbox '%s': %w", sandboxID, err),
207+
Err: fmt.Errorf("error getting last snapshot for sandbox '%s': %w", snapshotSandboxID, err),
200208
}
201209
}
202210

@@ -253,6 +261,7 @@ func (a *APIStore) buildResumeSandboxData(sandboxID string, autoPauseOverride *b
253261
VolumeMounts: convertDatabaseMountsToOrchestratorMounts(volumes),
254262
EnvdAccessToken: envdAccessToken,
255263
NodeID: &nodeID,
264+
SnapshotSandboxID: snapshotSandboxID,
256265
}, nil
257266
}
258267
}
Lines changed: 130 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,130 @@
1+
package orchestrator
2+
3+
import (
4+
"context"
5+
"fmt"
6+
"sync"
7+
"time"
8+
9+
"github.com/gogo/status"
10+
"github.com/google/uuid"
11+
"google.golang.org/grpc/codes"
12+
13+
"github.com/e2b-dev/infra/packages/api/internal/sandbox"
14+
"github.com/e2b-dev/infra/packages/db/pkg/types"
15+
"github.com/e2b-dev/infra/packages/db/queries"
16+
"github.com/e2b-dev/infra/packages/shared/pkg/grpc/orchestrator"
17+
"github.com/e2b-dev/infra/packages/shared/pkg/storage/storageopts"
18+
"github.com/e2b-dev/infra/packages/shared/pkg/telemetry"
19+
)
20+
21+
// CheckpointSandbox snapshots a running sandbox in place: the sandbox is
22+
// briefly paused on its node, snapshotted, and resumed under the same
23+
// execution ID, so it keeps running with its ID, expiration, and reservation
24+
// untouched. The snapshot is written to the sandbox's own snapshots row, so
25+
// it can immediately be resumed or forked from.
26+
func (o *Orchestrator) CheckpointSandbox(ctx context.Context, teamID uuid.UUID, sandboxID string) error {
27+
ctx, span := tracer.Start(ctx, "checkpoint-sandbox")
28+
defer span.End()
29+
30+
sbx, alreadyDone, finishSnapshotting, err := o.sandboxStore.StartRemoving(ctx, teamID, sandboxID, sandbox.RemoveOpts{Action: sandbox.StateActionSnapshot})
31+
if err != nil {
32+
return fmt.Errorf("failed to start snapshotting: %w", err)
33+
}
34+
35+
// alreadyDone conflates joining a concurrent checkpoint that just
36+
// succeeded with finding the sandbox stuck in Snapshotting after a failed
37+
// one, where no fresh snapshot exists. Treating it as success could fork
38+
// stale state, so report a conflict and let the caller retry (same
39+
// behavior as CreateSnapshotTemplate).
40+
if alreadyDone {
41+
return &sandbox.InvalidStateTransitionError{
42+
CurrentState: sandbox.StateSnapshotting,
43+
TargetState: sandbox.StateSnapshotting,
44+
}
45+
}
46+
47+
// finish completes the snapshotting transition exactly once.
48+
// On success (nil) it restores the sandbox to Running.
49+
// On error it leaves the state as Snapshotting so that
50+
// RemoveSandbox can transition directly to Killing.
51+
var once sync.Once
52+
finish := func(err error) {
53+
once.Do(func() {
54+
finishSnapshotting(context.WithoutCancel(ctx), err)
55+
})
56+
}
57+
defer finish(nil)
58+
59+
node := o.getOrConnectNode(ctx, sbx.ClusterID, sbx.NodeID)
60+
if node == nil {
61+
return fmt.Errorf("node '%s' not found", sbx.NodeID)
62+
}
63+
64+
upsertResult, err := o.throttledUpsertSnapshot(ctx, buildUpsertSnapshotParams(sbx, node, false))
65+
if err != nil {
66+
return fmt.Errorf("error upserting snapshot: %w", err)
67+
}
68+
69+
// Checkpoint pauses the sandbox, snapshots it, and resumes it on the
70+
// orchestrator with the same ExecutionID. Once the pause has started, the
71+
// orchestrator stops the old sandbox itself on error; RemoveSandbox is
72+
// still needed to clean up API-side state (store, routing, analytics).
73+
client, childCtx := node.GetClient(ctx)
74+
_, err = client.Sandbox.Checkpoint(childCtx, &orchestrator.SandboxCheckpointRequest{
75+
SandboxId: sbx.SandboxID,
76+
BuildId: upsertResult.BuildID.String(),
77+
Metadata: map[string]string{storageopts.ObjectMetadataTemplateID: upsertResult.TemplateID},
78+
})
79+
if err != nil {
80+
// Cleanup must run even when the checkpoint failed because this
81+
// request's context was cancelled (e.g. client disconnect mid-fork).
82+
cleanupCtx := context.WithoutCancel(ctx)
83+
84+
o.failSnapshotBuild(cleanupCtx, upsertResult.BuildID, err)
85+
86+
// The orchestrator rejects these before pausing the VM (envd too old,
87+
// starting-sandboxes queue full), so the sandbox is still running
88+
// healthy on its node: restore it to Running instead of killing it.
89+
if st, ok := status.FromError(err); ok {
90+
switch st.Code() {
91+
case codes.FailedPrecondition:
92+
finish(nil)
93+
94+
return fmt.Errorf("checkpoint rejected: %w", err)
95+
case codes.ResourceExhausted:
96+
finish(nil)
97+
98+
return PauseQueueExhaustedError{}
99+
}
100+
}
101+
102+
// Complete the snapshotting transition with error — leaves state as
103+
// Snapshotting (no restore to Running) and clears the transition key
104+
// so RemoveSandbox can proceed without deadlock.
105+
finish(err)
106+
107+
if killErr := o.RemoveSandbox(cleanupCtx, teamID, sandboxID, sandbox.RemoveOpts{Action: sandbox.StateActionKill}); killErr != nil {
108+
telemetry.ReportError(cleanupCtx, "error killing sandbox after failed checkpoint", killErr)
109+
}
110+
111+
return fmt.Errorf("checkpoint failed: %w", err)
112+
}
113+
114+
now := time.Now()
115+
err = o.sqlcDB.UpdateEnvBuildStatus(ctx, queries.UpdateEnvBuildStatusParams{
116+
Status: types.BuildStatusSuccess,
117+
FinishedAt: &now,
118+
Reason: types.BuildReason{},
119+
BuildID: upsertResult.BuildID,
120+
})
121+
if err != nil {
122+
return fmt.Errorf("error updating build status: %w", err)
123+
}
124+
125+
o.snapshotCache.Invalidate(context.WithoutCancel(ctx), sandboxID)
126+
127+
telemetry.ReportEvent(ctx, "Checkpointed sandbox")
128+
129+
return nil
130+
}

packages/api/internal/orchestrator/create_instance.go

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,9 @@ type SandboxMetadata struct {
5454
VolumeMounts []*orchestrator.SandboxVolumeMount
5555
EnvdAccessToken *string
5656
NodeID *string
57+
// SnapshotSandboxID is the sandbox ID the resume snapshot is stored under.
58+
// It differs from the ID of the sandbox being started when forking.
59+
SnapshotSandboxID string
5760
}
5861

5962
// buildEgressConfig constructs the orchestrator egress configuration from
@@ -322,7 +325,13 @@ func (o *Orchestrator) CreateSandbox(
322325
placed, err := placement.PlaceSandbox(ctx, o.placementAlgorithm, clusterNodes, node, sbxRequest, builds.ToMachineInfo(sbxData.Build), labelFilteringEnabled, allLabels)
323326
if err != nil {
324327
if isResume && placed.TimedOut {
325-
o.maybeRemapResumeOriginNode(ctx, sandboxID, team, sbxData.NodeID, placed.WarmedNode)
328+
// Remap by the snapshot's own sandbox ID: when forking, the started
329+
// sandbox ID has no snapshot row and the remap would silently no-op.
330+
snapshotSandboxID := sandboxID
331+
if sbxData.SnapshotSandboxID != "" {
332+
snapshotSandboxID = sbxData.SnapshotSandboxID
333+
}
334+
o.maybeRemapResumeOriginNode(ctx, snapshotSandboxID, team, sbxData.NodeID, placed.WarmedNode)
326335
}
327336

328337
return sandbox.Sandbox{}, &api.APIError{

0 commit comments

Comments
 (0)