Merge branch 'main' of github.com:e2b-dev/infra into lev-block-cache-bitmap

Author: levb

packages/api/internal/handlers/proxy_grpc.go

@@ -116,8 +116,7 @@ func (s *SandboxService) ResumeSandbox(ctx context.Context, req *proxygrpc.Sandb
 		return nil, status.Error(codes.InvalidArgument, "invalid sandbox ID")
 	}
 
-	var autoResume *dbtypes.SandboxAutoResumeConfig
-	snap, _, err := s.getAutoResumeSnapshot(ctx, sandboxID)
+	snap, autoResume, err := s.getAutoResumeSnapshot(ctx, sandboxID)
 	if err != nil {
 		return nil, err
 	}
@@ -155,15 +154,6 @@ func (s *SandboxService) ResumeSandbox(ctx context.Context, req *proxygrpc.Sandb
 		}
 	}
 
-	// Reload snapshot metadata after orchestrator checks so we do not resume from stale
-	// pre-pause snapshot data.
-	snap, autoResume, err = s.getAutoResumeSnapshot(ctx, sandboxID)
-	if err != nil {
-		return nil, err
-	}
-
-	teamID = snap.Snapshot.TeamID
-
 	team, err := dbapi.GetTeamByID(ctx, s.api.authDB, teamID)
 	if err != nil {
 		return nil, status.Errorf(codes.Internal, "failed to get team: %v", err)
@@ -172,15 +162,7 @@ func (s *SandboxService) ResumeSandbox(ctx context.Context, req *proxygrpc.Sandb
 
 	timeout := calculateAutoResumeTimeout(autoResume, minAutoResumeTimeout, team)
 
-	autoPause := snap.Snapshot.AutoPause
-	nodeID := &snap.Snapshot.OriginNodeID
-
-	alias := ""
-	if len(snap.Aliases) > 0 {
-		alias = snap.Aliases[0]
-	}
-
-	var envdAccessToken *string = nil
+	var envdAccessToken *string
 	if snap.Snapshot.EnvSecure {
 		accessToken, tokenErr := s.api.getEnvdAccessToken(snap.EnvBuild.EnvdVersion, sandboxID)
 		if tokenErr != nil {
@@ -228,25 +210,13 @@ func (s *SandboxService) ResumeSandbox(ctx context.Context, req *proxygrpc.Sandb
 	headers := http.Header{}
 	sbx, apiErr := s.api.startSandboxInternal(
 		ctx,
-		snap.Snapshot.SandboxID,
+		sandboxID,
 		timeout,
-		nil,
-		snap.Snapshot.Metadata,
-		alias,
 		team,
-		snap.EnvBuild,
+		s.api.buildResumeSandboxData(sandboxID, nil),
 		&headers,
 		true,
-		nodeID,
-		snap.Snapshot.EnvID,
-		snap.Snapshot.BaseEnvID,
-		autoPause,
-		autoResume,
-		envdAccessToken,
-		snap.Snapshot.AllowInternetAccess,
-		network,
 		nil, // mcp
-		nil, // volumeMounts
 	)
 	if apiErr != nil {
 		return nil, status.Error(sharedutils.GRPCCodeFromHTTPStatus(apiErr.Code), apiErr.ClientMsg)

packages/api/internal/handlers/sandbox.go

@@ -14,11 +14,9 @@ import (
 
 	"github.com/e2b-dev/infra/packages/api/internal/api"
 	"github.com/e2b-dev/infra/packages/api/internal/middleware/otel/tracing"
+	"github.com/e2b-dev/infra/packages/api/internal/orchestrator"
 	"github.com/e2b-dev/infra/packages/api/internal/sandbox"
 	typesteam "github.com/e2b-dev/infra/packages/auth/pkg/types"
-	"github.com/e2b-dev/infra/packages/db/pkg/types"
-	"github.com/e2b-dev/infra/packages/db/queries"
-	"github.com/e2b-dev/infra/packages/shared/pkg/grpc/orchestrator"
 	sbxlogger "github.com/e2b-dev/infra/packages/shared/pkg/logger/sandbox"
 	"github.com/e2b-dev/infra/packages/shared/pkg/telemetry"
 )
@@ -29,45 +27,21 @@ func (a *APIStore) startSandbox(
 	ctx context.Context,
 	sandboxID string,
 	timeout time.Duration,
-	envVars map[string]string,
-	metadata map[string]string,
-	alias string,
 	team *typesteam.Team,
-	build queries.EnvBuild,
+	getSandboxData orchestrator.SandboxDataFetcher,
 	requestHeader *http.Header,
 	isResume bool,
-	nodeID *string,
-	templateID string,
-	baseTemplateID string,
-	autoPause bool,
-	autoResume *types.SandboxAutoResumeConfig,
-	envdAccessToken *string,
-	allowInternetAccess *bool,
-	network *types.SandboxNetworkConfig,
 	mcp api.Mcp,
-	volumeMounts []*orchestrator.SandboxVolumeMount,
 ) (*api.Sandbox, *api.APIError) {
 	sbx, apiErr := a.startSandboxInternal(
 		ctx,
 		sandboxID,
 		timeout,
-		envVars,
-		metadata,
-		alias,
 		team,
-		build,
+		getSandboxData,
 		requestHeader,
 		isResume,
-		nodeID,
-		templateID,
-		baseTemplateID,
-		autoPause,
-		autoResume,
-		envdAccessToken,
-		allowInternetAccess,
-		network,
 		mcp,
-		volumeMounts,
 	)
 	if apiErr != nil {
 		return nil, apiErr
@@ -81,56 +55,27 @@ func (a *APIStore) startSandboxInternal(
 	ctx context.Context,
 	sandboxID string,
 	timeout time.Duration,
-	envVars map[string]string,
-	metadata map[string]string,
-	alias string,
 	team *typesteam.Team,
-	build queries.EnvBuild,
+	getSandboxData orchestrator.SandboxDataFetcher,
 	requestHeader *http.Header,
 	isResume bool,
-	nodeID *string,
-	templateID string,
-	baseTemplateID string,
-	autoPause bool,
-	autoResume *types.SandboxAutoResumeConfig,
-	envdAccessToken *string,
-	allowInternetAccess *bool,
-	network *types.SandboxNetworkConfig,
 	mcp api.Mcp,
-	volumeMounts []*orchestrator.SandboxVolumeMount,
 ) (sandbox.Sandbox, *api.APIError) {
 	startTime := time.Now()
 	endTime := startTime.Add(timeout)
 
-	autoResumePolicy := "unset"
-	if autoResume != nil {
-		autoResumePolicy = string(autoResume.Policy)
-	}
-
 	// Unique ID for the execution (from start/resume to stop/pause)
 	executionID := uuid.New().String()
 	sbx, instanceErr := a.orchestrator.CreateSandbox(
 		ctx,
 		sandboxID,
 		executionID,
-		alias,
 		team,
-		build,
-		metadata,
-		envVars,
+		getSandboxData,
 		startTime,
 		endTime,
 		timeout,
 		isResume,
-		nodeID,
-		templateID,
-		baseTemplateID,
-		autoPause,
-		autoResume,
-		envdAccessToken,
-		allowInternetAccess,
-		network,
-		volumeMounts,
 	)
 	if instanceErr != nil {
 		telemetry.ReportError(ctx, "error when creating instance", instanceErr.Err)
@@ -146,15 +91,15 @@ func (a *APIStore) startSandboxInternal(
 	props := properties.
 		Set("environment", sbx.TemplateID).
 		Set("instance_id", sbx.SandboxID).
-		Set("alias", alias).
+		Set("alias", sbx.Alias).
 		Set("resume", isResume).
 		Set("build_id", sbx.BuildID).
-		Set("envd_version", build.EnvdVersion).
+		Set("envd_version", sbx.EnvdVersion).
 		Set("node_id", sbx.NodeID).
 		Set("vcpu", sbx.VCpu).
 		Set("ram_mb", sbx.RamMB).
 		Set("total_disk_size_mb", sbx.TotalDiskSizeMB).
-		Set("auto_pause", autoPause)
+		Set("auto_pause", sbx.AutoPause)
 
 	// Calculate the time it took for the sandbox to start from request receipt
 	if requestStartTime, ok := tracing.GetRequestStartTime(ctx); ok {
@@ -171,7 +116,7 @@ func (a *APIStore) startSandboxInternal(
 	telemetry.ReportEvent(ctx, "Created analytics event")
 
 	go func() {
-		a.templateSpawnCounter.IncreaseTemplateSpawnCount(baseTemplateID, time.Now())
+		a.templateSpawnCounter.IncreaseTemplateSpawnCount(sbx.BaseTemplateID, time.Now())
 	}()
 
 	telemetry.SetAttributes(ctx,
@@ -184,13 +129,19 @@ func (a *APIStore) startSandboxInternal(
 		TeamID:     team.ID.String(),
 	}
 	sbxlogger.E(logMetadata).Info(ctx, "Sandbox created", zap.String("end_time", endTime.Format("2006-01-02 15:04:05 -07:00")))
+
+	autoResumePolicy := "unset"
+	if sbx.AutoResume != nil {
+		autoResumePolicy = string(sbx.AutoResume.Policy)
+	}
+
 	sbxlogger.I(logMetadata).Info(
 		ctx,
 		"Sandbox created details",
 		zap.String("end_time", endTime.Format("2006-01-02 15:04:05 -07:00")),
 		zap.String("auto_resume_policy", autoResumePolicy),
-		zap.Bool("auto_pause", autoPause),
-		zap.String("parent_template_id", baseTemplateID),
+		zap.Bool("auto_pause", sbx.AutoPause),
+		zap.String("parent_template_id", sbx.BaseTemplateID),
 	)
 
 	return sbx, nil

packages/api/internal/handlers/sandbox_connect.go

@@ -15,7 +15,6 @@ import (
 	"github.com/e2b-dev/infra/packages/api/internal/sandbox"
 	"github.com/e2b-dev/infra/packages/api/internal/utils"
 	"github.com/e2b-dev/infra/packages/auth/pkg/auth"
-	"github.com/e2b-dev/infra/packages/db/pkg/types"
 	"github.com/e2b-dev/infra/packages/shared/pkg/ginutils"
 	"github.com/e2b-dev/infra/packages/shared/pkg/logger"
 	sbxlogger "github.com/e2b-dev/infra/packages/shared/pkg/logger/sandbox"
@@ -139,76 +138,21 @@ func (a *APIStore) PostSandboxesSandboxIDConnect(c *gin.Context, sandboxID api.S
 		return
 	}
 
-	autoPause := lastSnapshot.Snapshot.AutoPause
-	snap := lastSnapshot.Snapshot
-	build := lastSnapshot.EnvBuild
-
-	nodeID := &snap.OriginNodeID
-
-	alias := ""
-	if len(lastSnapshot.Aliases) > 0 {
-		alias = lastSnapshot.Aliases[0]
-	}
-
 	sbxlogger.E(&sbxlogger.SandboxMetadata{
 		SandboxID:  sandboxID,
-		TemplateID: snap.EnvID,
+		TemplateID: lastSnapshot.Snapshot.EnvID,
 		TeamID:     teamID.String(),
 	}).Debug(ctx, "Started resuming sandbox")
 
-	var envdAccessToken *string = nil
-	if snap.EnvSecure {
-		accessToken, tokenErr := a.getEnvdAccessToken(build.EnvdVersion, sandboxID)
-		if tokenErr != nil {
-			telemetry.ReportErrorByCode(ctx, tokenErr.Code, "Secure envd access token error", tokenErr.Err,
-				telemetry.WithTemplateID(snap.EnvID),
-				telemetry.WithBuildID(build.ID.String()),
-				telemetry.WithSandboxID(sandboxID),
-				telemetry.WithTeamID(teamID.String()),
-			)
-			a.sendAPIStoreError(c, tokenErr.Code, tokenErr.ClientMsg)
-
-			return
-		}
-
-		envdAccessToken = &accessToken
-	}
-
-	var network *types.SandboxNetworkConfig
-	if snap.Config != nil {
-		network = snap.Config.Network
-	}
-	var autoResume *types.SandboxAutoResumeConfig
-	if snap.Config != nil {
-		autoResume = snap.Config.AutoResume
-	}
-
-	var volumes []*types.SandboxVolumeMountConfig
-	if snap.Config != nil {
-		volumes = snap.Config.VolumeMounts
-	}
-
 	sbx, createErr := a.startSandbox(
 		ctx,
-		snap.SandboxID,
+		sandboxID,
 		timeout,
-		nil,
-		snap.Metadata,
-		alias,
 		teamInfo,
-		build,
+		a.buildResumeSandboxData(sandboxID, nil),
 		&c.Request.Header,
 		true,
-		nodeID,
-		snap.EnvID,
-		snap.BaseEnvID,
-		autoPause,
-		autoResume,
-		envdAccessToken,
-		snap.AllowInternetAccess,
-		network,
 		nil, // mcp
-		convertDatabaseMountsToOrchestratorMounts(volumes), // volumes
 	)
 	if createErr != nil {
 		a.sendAPIStoreError(c, createErr.Code, createErr.ClientMsg)

packages/api/internal/handlers/sandbox_create.go

@@ -22,6 +22,7 @@ import (
 	"github.com/e2b-dev/infra/packages/api/internal/api"
 	templatecache "github.com/e2b-dev/infra/packages/api/internal/cache/templates"
 	"github.com/e2b-dev/infra/packages/api/internal/middleware/otel/metrics"
+	apiorch "github.com/e2b-dev/infra/packages/api/internal/orchestrator"
 	"github.com/e2b-dev/infra/packages/api/internal/sandbox"
 	"github.com/e2b-dev/infra/packages/auth/pkg/auth"
 	sqlcdb "github.com/e2b-dev/infra/packages/db/client"
@@ -219,27 +220,34 @@ func (a *APIStore) PostSandboxes(c *gin.Context) {
 		return
 	}
 
+	getSandboxData := func(_ context.Context) (apiorch.SandboxMetadata, *api.APIError) {
+		// The data can't be influenced by action on the same sandbox as other operations,
+		// so it's safe to reuse the data
+		return apiorch.SandboxMetadata{
+			Metadata:            metadata,
+			EnvVars:             envVars,
+			Build:               *build,
+			AllowInternetAccess: allowInternetAccess,
+			Network:             network,
+			Alias:               alias,
+			TemplateID:          env.TemplateID,
+			BaseTemplateID:      env.TemplateID,
+			AutoPause:           autoPause,
+			AutoResume:          autoResume,
+			VolumeMounts:        sbxVolumeMounts,
+			EnvdAccessToken:     envdAccessToken,
+		}, nil
+	}
+
 	sbx, createErr := a.startSandbox(
 		ctx,
 		sandboxID,
 		timeout,
-		envVars,
-		metadata,
-		alias,
 		teamInfo,
-		*build,
+		getSandboxData,
 		&c.Request.Header,
 		false,
-		nil,
-		env.TemplateID,
-		env.TemplateID,
-		autoPause,
-		autoResume,
-		envdAccessToken,
-		allowInternetAccess,
-		network,
 		mcp,
-		sbxVolumeMounts,
 	)
 	if createErr != nil {
 		a.sendAPIStoreError(c, createErr.Code, createErr.ClientMsg)