Write to cache when building a template

.mockery.yaml

@@ -9,8 +9,13 @@ packages:
 
   github.com/e2b-dev/infra/packages/shared/pkg/storage:
     interfaces:
+      StorageProvider:
+        config:
+          dir: packages/shared/pkg/storage
+          filename: mockstorageprovider_test.go
+          pkgname: storage
       StorageObjectProvider:
         config:
-          dir: packages/shared/pkg/storage/mocks
-          filename: mocks.go
-          pkgname: storagemocks
+          dir: packages/shared/pkg/storage
+          filename: mockstorageobjectprovider_test.go
+          pkgname: storage

packages/orchestrator/benchmark_test.go

@@ -231,6 +231,7 @@ func BenchmarkBaseImageLaunch(b *testing.B) {
 		sandboxes,
 		templateCache,
 		buildMetrics,
+		featureFlags,
 	)
 
 	buildPath := filepath.Join(os.Getenv("LOCAL_TEMPLATE_STORAGE_BASE_PATH"), buildID, "rootfs.ext4")

packages/orchestrator/cmd/build-template/main.go

@@ -186,6 +186,7 @@ func buildTemplate(
 		sandboxes,
 		templateCache,
 		buildMetrics,
+		featureFlags,
 	)
 
 	logger = logger.

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

@@ -30,6 +30,8 @@ import (
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/template/constants"
 	artifactsregistry "github.com/e2b-dev/infra/packages/shared/pkg/artifacts-registry"
 	"github.com/e2b-dev/infra/packages/shared/pkg/dockerhub"
+	"github.com/e2b-dev/infra/packages/shared/pkg/env"
+	featureflags "github.com/e2b-dev/infra/packages/shared/pkg/feature-flags"
 	"github.com/e2b-dev/infra/packages/shared/pkg/smap"
 	"github.com/e2b-dev/infra/packages/shared/pkg/storage"
 	"github.com/e2b-dev/infra/packages/shared/pkg/storage/header"
@@ -52,6 +54,9 @@ type Builder struct {
 	sandboxes           *smap.Map[*sandbox.Sandbox]
 	templateCache       *sbxtemplate.Cache
 	metrics             *metrics.BuildMetrics
+	flags               *featureflags.Client
+
+	rootCachePath string
 }
 
 func NewBuilder(
@@ -65,6 +70,7 @@ func NewBuilder(
 	sandboxes *smap.Map[*sandbox.Sandbox],
 	templateCache *sbxtemplate.Cache,
 	buildMetrics *metrics.BuildMetrics,
+	flags *featureflags.Client,
 ) *Builder {
 	return &Builder{
 		logger:              logger,
@@ -77,6 +83,8 @@ func NewBuilder(
 		sandboxes:           sandboxes,
 		templateCache:       templateCache,
 		metrics:             buildMetrics,
+		flags:               flags,
+		rootCachePath:       env.GetEnv("SHARED_CHUNK_CACHE_PATH", ""),
 	}
 }
 
@@ -189,21 +197,40 @@ func (b *Builder) Build(ctx context.Context, template storage.TemplateFiles, con
 	return runBuild(ctx, logger, buildContext, b)
 }
 
+func (b *Builder) useNFSCache(ctx context.Context) bool {
+	if b.rootCachePath == "" {
+		// can't enable cache if we don't have a cache path
+		return false
+	}
+
+	flag, err := b.flags.BoolFlag(ctx, featureflags.BuildingFeatureFlagName)
+	if err != nil {
+		zap.L().Error("failed to get nfs cache feature flag", zap.Error(err))
+	}
+
+	return flag
+}
+
 func runBuild(
 	ctx context.Context,
 	userLogger *zap.Logger,
 	bc buildcontext.BuildContext,
 	builder *Builder,
 ) (*Result, error) {
-	index := cache.NewHashIndex(bc.CacheScope, builder.buildStorage, builder.templateStorage)
+	templateStorage := builder.templateStorage
+	if builder.useNFSCache(ctx) {
+		templateStorage = storage.NewCachedProvider(builder.rootCachePath, templateStorage)
+	}
+
+	index := cache.NewHashIndex(bc.CacheScope, builder.buildStorage, templateStorage)
 
 	layerExecutor := layer.NewLayerExecutor(
 		bc,
 		builder.logger,
 		builder.templateCache,
 		builder.proxy,
 		builder.sandboxes,
-		builder.templateStorage,
+		templateStorage,
 		builder.buildStorage,
 		index,
 	)
@@ -212,7 +239,7 @@ func runBuild(
 		bc,
 		builder.logger,
 		builder.proxy,
-		builder.templateStorage,
+		templateStorage,
 		builder.artifactRegistry,
 		builder.dockerhubRepository,
 		layerExecutor,
@@ -241,7 +268,7 @@ func runBuild(
 	postProcessingBuilder := finalize.New(
 		bc,
 		builder.sandboxFactory,
-		builder.templateStorage,
+		templateStorage,
 		builder.proxy,
 		layerExecutor,
 	)
@@ -267,7 +294,7 @@ func runBuild(
 	// Get the base rootfs size from the template files
 	// This is the size of the rootfs after provisioning and before building the layers
 	// (as they don't change the rootfs size)
-	rootfsSize, err := getRootfsSize(ctx, builder.templateStorage, lastLayerResult.Metadata.Template)
+	rootfsSize, err := getRootfsSize(ctx, templateStorage, lastLayerResult.Metadata.Template)
 	if err != nil {
 		return nil, fmt.Errorf("error getting rootfs size: %w", err)
 	}

packages/orchestrator/internal/template/server/main.go

@@ -21,6 +21,7 @@ import (
 	artifactsregistry "github.com/e2b-dev/infra/packages/shared/pkg/artifacts-registry"
 	"github.com/e2b-dev/infra/packages/shared/pkg/dockerhub"
 	"github.com/e2b-dev/infra/packages/shared/pkg/env"
+	featureflags "github.com/e2b-dev/infra/packages/shared/pkg/feature-flags"
 	templatemanager "github.com/e2b-dev/infra/packages/shared/pkg/grpc/template-manager"
 	"github.com/e2b-dev/infra/packages/shared/pkg/limit"
 	"github.com/e2b-dev/infra/packages/shared/pkg/smap"
@@ -61,6 +62,7 @@ func New(
 	templatePersistence storage.StorageProvider,
 	limiter *limit.Limiter,
 	info *service.ServiceInfo,
+	flags *featureflags.Client,
 ) (s *ServerStore, e error) {
 	logger.Info("Initializing template manager")
 
@@ -110,6 +112,7 @@ func New(
 		sandboxes,
 		templateCache,
 		buildMetrics,
+		flags,
 	)
 
 	store := &ServerStore{

packages/orchestrator/main.go

@@ -399,6 +399,7 @@ func run(config cfg.Config) (success bool) {
 			persistence,
 			limiter,
 			serviceInfo,
+			featureFlags,
 		)
 		if err != nil {
 			zap.L().Fatal("failed to create template manager", zap.Error(err))

packages/shared/pkg/feature-flags/flags.go

@@ -47,6 +47,7 @@ var (
 	SandboxLifeCycleEventsWriteFlagName = newBoolFlag("sandbox-lifecycle-events-write", env.IsDevelopment())
 	SnapshotFeatureFlagName             = newBoolFlag("use-nfs-for-snapshots", env.IsDevelopment())
 	TemplateFeatureFlagName             = newBoolFlag("use-nfs-for-templates", env.IsDevelopment())
+	BuildingFeatureFlagName             = newBoolFlag("use-nfs-for-building-templates", env.IsDevelopment())
 	SandboxEventsPublishFlagName        = newBoolFlag("sandbox-events-publish", env.IsDevelopment())
 	BestOfKPlacementAlgorithm           = newBoolFlag("best-of-k-placement-algorithm", env.IsDevelopment())
 	BestOfKCanFit                       = newBoolFlag("best-of-k-can-fit", true)