Update envd on first layer build (#941)

Author: dobrac

packages/orchestrator/internal/template/build/build_layer.go

@@ -15,6 +15,7 @@ import (
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/fc"
 	sbxtemplate "github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/template"
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/template/build/config"
+	"github.com/e2b-dev/infra/packages/orchestrator/internal/template/build/envd"
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/template/build/sandboxtools"
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/template/build/writer"
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/template/constants"
@@ -36,6 +37,7 @@ func (b *Builder) buildLayer(
 	sourceMeta LayerMetadata,
 	exportTemplate storage.TemplateFiles,
 	resumeSandbox bool,
+	updateEnvd bool,
 	fn func(ctx context.Context, sbx *sandbox.Sandbox) (sandboxtools.CommandMetadata, error),
 ) (LayerMetadata, error) {
 	ctx, childSpan := b.tracer.Start(ctx, "run-in-sandbox")
@@ -170,6 +172,14 @@ func (b *Builder) buildLayer(
 		b.proxy.RemoveFromPool(sbx.Metadata.Config.ExecutionId)
 	}()
 
+	// Update envd binary to the latest version
+	if updateEnvd {
+		err = b.updateEnvdInSandbox(ctx, b.tracer, postProcessor, sbx)
+		if err != nil {
+			return LayerMetadata{}, fmt.Errorf("failed to update envd in sandbox: %w", err)
+		}
+	}
+
 	meta, err := fn(ctx, sbx)
 	if err != nil {
 		return LayerMetadata{}, fmt.Errorf("error running action in sandbox: %w", err)
@@ -199,6 +209,83 @@ func (b *Builder) buildLayer(
 	return exportMeta, nil
 }
 
+// updateEnvdInSandbox updates the envd binary in the sandbox to the latest version.
+func (b *Builder) updateEnvdInSandbox(
+	ctx context.Context,
+	tracer trace.Tracer,
+	postProcessor *writer.PostProcessor,
+	sbx *sandbox.Sandbox,
+) error {
+	ctx, childSpan := tracer.Start(ctx, "update-envd")
+	defer childSpan.End()
+
+	envdVersion, err := envd.GetEnvdVersion(ctx)
+	if err != nil {
+		return fmt.Errorf("error getting envd version: %w", err)
+	}
+	postProcessor.Debug(fmt.Sprintf("Updating envd to version v%s", envdVersion))
+
+	// Step 1: Copy the updated envd binary from host to /tmp in sandbox
+	tmpEnvdPath := "/tmp/envd_updated"
+	err = sandboxtools.CopyFile(
+		ctx,
+		b.tracer,
+		b.proxy,
+		sbx.Metadata.Config.SandboxId,
+		"root",
+		storage.HostEnvdPath,
+		tmpEnvdPath,
+	)
+	if err != nil {
+		return fmt.Errorf("failed to copy envd binary to sandbox: %w", err)
+	}
+
+	// Step 2: Replace the binary
+	replaceEnvdCmd := fmt.Sprintf(`
+		# Replace the binary and set permissions
+		chmod +x %s
+		mv -f %s %s
+	`, tmpEnvdPath, tmpEnvdPath, storage.GuestEnvdPath)
+
+	err = sandboxtools.RunCommandWithLogger(
+		ctx,
+		b.tracer,
+		b.proxy,
+		postProcessor,
+		zap.DebugLevel,
+		"update-envd-replace",
+		sbx.Metadata.Config.SandboxId,
+		replaceEnvdCmd,
+		sandboxtools.CommandMetadata{User: "root"},
+	)
+	if err != nil {
+		return fmt.Errorf("failed to replace envd binary: %w", err)
+	}
+
+	// Step 3: Restart the systemd envd service
+	// Error is ignored because it's expected the envd connection will be lost
+	_ = sandboxtools.RunCommand(
+		ctx,
+		b.tracer,
+		b.proxy,
+		sbx.Metadata.Config.SandboxId,
+		"systemctl restart envd",
+		sandboxtools.CommandMetadata{User: "root"},
+	)
+
+	// Step 4: Wait for envd to initialize
+	err = sbx.WaitForEnvd(
+		ctx,
+		b.tracer,
+		waitEnvdTimeout,
+	)
+	if err != nil {
+		return fmt.Errorf("failed to wait for envd initialization after update: %w", err)
+	}
+
+	return nil
+}
+
 func pauseAndUpload(
 	ctx context.Context,
 	tracer trace.Tracer,

packages/orchestrator/internal/template/build/builder.go

@@ -163,7 +163,7 @@ func (b *Builder) Build(ctx context.Context, finalMetadata storage.TemplateFiles
 		return nil, fmt.Errorf("error getting base hash: %w", err)
 	}
 
-	lastCached, baseMetadata, err := b.setupBase(ctx, cacheScope, finalMetadata, template, lastHash, isV1Build)
+	isLastLayerCached, baseMetadata, err := b.setupBase(ctx, cacheScope, finalMetadata, template, lastHash, isV1Build)
 	if err != nil {
 		return nil, fmt.Errorf("error setting up build: %w", err)
 	}
@@ -183,10 +183,10 @@ func (b *Builder) Build(ctx context.Context, finalMetadata storage.TemplateFiles
 		}
 		baseSource = "FROM " + fromImage
 	}
-	postProcessor.Info(layerInfo(lastCached, "base", baseSource, lastHash))
+	postProcessor.Info(layerInfo(isLastLayerCached, "base", baseSource, lastHash))
 
 	// Build the base layer if not cached
-	if !lastCached {
+	if !isLastLayerCached {
 		templateBuildDir := filepath.Join(templatesDirectory, finalMetadata.BuildID)
 		err = os.MkdirAll(templateBuildDir, 0o777)
 		if err != nil {
@@ -349,9 +349,6 @@ func (b *Builder) Build(ctx context.Context, finalMetadata storage.TemplateFiles
 
 	sourceMetadata := baseMetadata
 
-	// First start is create (to change CPU, etc), subsequent starts are resume.
-	shouldResumeSandbox := false
-
 	// Build Steps
 	for i, step := range template.Steps {
 		layerIndex := i + 1
@@ -385,7 +382,6 @@ func (b *Builder) Build(ctx context.Context, finalMetadata storage.TemplateFiles
 			return nil, fmt.Errorf("error checking if layer is cached: %w", err)
 		}
 		isCached := !force && found
-		lastCached = isCached
 
 		prefix := fmt.Sprintf("builder %d/%d", layerIndex, len(template.Steps))
 		cmd := fmt.Sprintf("%s %s", strings.ToUpper(step.Type), strings.Join(step.Args, " "))
@@ -411,7 +407,9 @@ func (b *Builder) Build(ctx context.Context, finalMetadata storage.TemplateFiles
 				lastHash,
 				sourceMetadata,
 				stepMetadata.Template,
-				shouldResumeSandbox,
+				// First not cached layer is create (to change CPU, etc), subsequent are resumes.
+				!isLastLayerCached,
+				isLastLayerCached,
 				func(ctx context.Context, sbx *sandbox.Sandbox) (sandboxtools.CommandMetadata, error) {
 					postProcessor.Debug(fmt.Sprintf("Running action in: %s/%s", sourceMetadata.Template.TemplateID, sourceMetadata.Template.BuildID))
 
@@ -438,13 +436,15 @@ func (b *Builder) Build(ctx context.Context, finalMetadata storage.TemplateFiles
 			}
 			stepMetadata = meta
 
-			if !shouldResumeSandbox {
+			// If the last layer is cached, update the base metadata to the step metadata
+			// This is needed to properly resume the sandbox for the next step
+			if isLastLayerCached {
 				baseMetadata = stepMetadata
-				shouldResumeSandbox = true
 			}
 		}
 
 		sourceMetadata = stepMetadata
+		isLastLayerCached = isCached
 	}
 	// Build Steps
 
@@ -487,7 +487,9 @@ func (b *Builder) Build(ctx context.Context, finalMetadata storage.TemplateFiles
 		lastHash,
 		sourceMetadata,
 		finalMetadata,
+		// Always restart the sandbox for the final layer to properly wire the rootfs path for the final template.
 		false,
+		isLastLayerCached,
 		b.postProcessingFn(postProcessor, finalMetadata, sourceMetadata.Metadata, startMetadata),
 	)
 	if err != nil {

packages/orchestrator/internal/template/build/hash_index.go

@@ -10,7 +10,6 @@ import (
 	"strings"
 
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/template/build/config"
-	"github.com/e2b-dev/infra/packages/orchestrator/internal/template/build/envd"
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/template/build/layerstorage"
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/template/build/sandboxtools"
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/template/build/utils"
@@ -84,11 +83,6 @@ func hashKeys(baseKey string, keys ...string) string {
 }
 
 func hashBase(template config.TemplateConfig) (string, error) {
-	envdHash, err := envd.GetEnvdHash()
-	if err != nil {
-		return "", fmt.Errorf("error getting envd binary hash: %w", err)
-	}
-
 	var baseSource string
 	if template.FromTemplate != nil {
 		// When building from template, use the base template metadata
@@ -103,7 +97,7 @@ func hashBase(template config.TemplateConfig) (string, error) {
 		baseSource = template.FromImage
 	}
 
-	return hashKeys(hashingVersion, envdHash, provisionScriptFile, strconv.FormatInt(template.DiskSizeMB, 10), baseSource), nil
+	return hashKeys(hashingVersion, provisionScriptFile, strconv.FormatInt(template.DiskSizeMB, 10), baseSource), nil
 }
 
 func hashStep(previousHash string, step *templatemanager.TemplateStep) string {