fix: speed up provisioning shutdown

Author: dobrac

packages/orchestrator/internal/sandbox/block/cache.go

@@ -112,12 +112,16 @@ func (m *Cache) ExportToDiff(out io.Writer) (*header.DiffMetadata, error) {
 		}
 
 		dirty.Set(uint(blockIdx))
-		_, err = out.Write(block)
+		n, err := out.Write(block)
 		if err != nil {
 			zap.L().Error("error writing to out", zap.Error(err))
 
 			return nil, err
 		}
+
+		if int64(n) != m.blockSize {
+			return nil, fmt.Errorf("short write: %d != %d", int64(n), m.blockSize)
+		}
 	}
 
 	return &header.DiffMetadata{

packages/orchestrator/internal/sandbox/sandbox.go

@@ -655,6 +655,61 @@ func (s *Sandbox) FirecrackerVersions() fc.FirecrackerVersions {
 	return s.process.Versions
 }
 
+func (s *Sandbox) Shutdown(ctx context.Context) error {
+	ctx, span := tracer.Start(ctx, "shutdown sandbox")
+	defer span.End()
+
+	// Stop the health check before pausing the VM
+	s.Checks.Stop()
+
+	if err := s.process.Pause(ctx); err != nil {
+		return fmt.Errorf("failed to pause VM: %w", err)
+	}
+
+	if err := s.memory.Disable(); err != nil {
+		return fmt.Errorf("failed to disable uffd: %w", err)
+	}
+
+	// This is required because the FC API doesn't support passing /dev/null
+	tf, err := storage.TemplateFiles{
+		BuildID:            uuid.New().String(),
+		KernelVersion:      s.Template.Files().KernelVersion,
+		FirecrackerVersion: s.Template.Files().FirecrackerVersion,
+	}.CacheFiles(s.config)
+	if err != nil {
+		return fmt.Errorf("failed to create template files: %w", err)
+	}
+	defer tf.Close(s.config)
+
+	// The snapfile is required only because the FC API doesn't support passing /dev/null
+	snapfile := template.NewLocalFileLink(tf.CacheSnapfilePath(s.config))
+	defer snapfile.Close()
+
+	// The memfile is required only because the FC API doesn't support passing /dev/null
+	memfile, err := storage.AcquireTmpMemfile(ctx, s.config, tf.BuildID)
+	if err != nil {
+		return fmt.Errorf("failed to acquire memfile snapshot: %w", err)
+	}
+	defer memfile.Close()
+
+	err = s.process.CreateSnapshot(
+		ctx,
+		snapfile.Path(),
+		memfile.Path(),
+	)
+	if err != nil {
+		return fmt.Errorf("error creating snapshot: %w", err)
+	}
+
+	// This should properly flush rootfs to the underlying device.
+	err = s.Stop(ctx)
+	if err != nil {
+		return fmt.Errorf("error stopping sandbox: %w", err)
+	}
+
+	return nil
+}
+
 func (s *Sandbox) Pause(
 	ctx context.Context,
 	m metadata.Template,

packages/orchestrator/internal/template/build/phases/base/provision.go

@@ -14,7 +14,6 @@ import (
 	tt "text/template"
 	"time"
 
-	"github.com/google/uuid"
 	"go.uber.org/zap"
 	"go.uber.org/zap/zapio"
 
@@ -27,8 +26,6 @@ import (
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/template/build/core/rootfs"
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/template/build/writer"
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/template/constants"
-	"github.com/e2b-dev/infra/packages/orchestrator/internal/template/metadata"
-	"github.com/e2b-dev/infra/packages/shared/pkg/storage"
 	"github.com/e2b-dev/infra/packages/shared/pkg/telemetry"
 	"github.com/e2b-dev/infra/packages/shared/pkg/utils"
 )
@@ -150,17 +147,10 @@ func (bb *BaseBuilder) provisionSandbox(
 
 	userLogger.Info("Provisioning was successful, cleaning up")
 
-	snapshot, err := sbx.Pause(ctx, metadata.Template{
-		Template: storage.TemplateFiles{
-			BuildID:            uuid.NewString(),
-			KernelVersion:      fcVersions.KernelVersion,
-			FirecrackerVersion: fcVersions.FirecrackerVersion,
-		},
-	})
+	err = sbx.Shutdown(ctx)
 	if err != nil {
-		return fmt.Errorf("error pausing provisioned sandbox: %w", err)
+		return fmt.Errorf("error shutting down provisioned sandbox: %w", err)
 	}
-	defer snapshot.Close(context.WithoutCancel(ctx))
 
 	err = filesystem.RemoveFile(ctx, rootfsPath, provisionScriptResultPath)
 	if err != nil {

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

@@ -112,7 +112,7 @@ func (sb *StepBuilder) Layer(
 	if !forceBuild {
 		m, err := sb.index.LayerMetaFromHash(ctx, hash)
 		if err != nil {
-			sb.logger.Info("layer not found in cache, building new base layer", zap.Error(err), zap.String("hash", hash))
+			sb.logger.Info("layer not found in cache, building new step layer", zap.Error(err), zap.String("hash", hash))
 		} else {
 			// Check if the layer is cached
 			meta, err := sb.index.Cached(ctx, m.Template.BuildID)

packages/shared/pkg/storage/template_cache.go

@@ -47,3 +47,7 @@ func (c TemplateCacheFiles) CacheMetadataPath(config BuilderConfig) string {
 func (c TemplateCacheFiles) cacheDir(config BuilderConfig) string {
 	return filepath.Join(config.GetTemplateCacheDir(), c.BuildID, "cache", c.CacheIdentifier)
 }
+
+func (c TemplateCacheFiles) Close(config BuilderConfig) error {
+	return os.RemoveAll(c.cacheDir(config))
+}