Abstract the storage to allow for local storage in tests and developmnet

.github/actions/build-sandbox-template/action.yml

@@ -14,6 +14,7 @@ runs:
         GCP_DOCKER_REPOSITORY_NAME: "custom-environments-tests"
         # Not used, but required to be defined
         GOOGLE_SERVICE_ACCOUNT_BASE64: "required-but-not-needed"
+        TEMPLATE_STORAGE: "local"
       run: |
         echo "TESTS_SANDBOX_TEMPLATE_ID=${TEMPLATE_ID}" >> .env.test
         echo "TESTS_SANDBOX_BUILD_ID=${BUILD_ID}" >> .env.test

.github/actions/start-services/action.yml

@@ -46,9 +46,11 @@ runs:
       env:
         ENVD_TIMEOUT: "60s"
         SUPABASE_JWT_SECRETS: "supabasejwtsecretsupabasejwtsecret"
+        TEMPLATE_STORAGE: "local"
       run: |
         echo "ENVD_TIMEOUT=${ENVD_TIMEOUT}" >> .env.test
         echo "SUPABASE_JWT_SECRETS=${SUPABASE_JWT_SECRETS}" >> .env.test
+        echo "TEMPLATE_STORAGE=${TEMPLATE_STORAGE}" >> .env.test
 
         mkdir -p ~/logs
 

packages/orchestrator/cmd/simulate-headers-merge/main.go

@@ -7,10 +7,11 @@ import (
 	"log"
 	"os"
 
+	"github.com/google/uuid"
+
 	"github.com/e2b-dev/infra/packages/shared/pkg/storage"
 	"github.com/e2b-dev/infra/packages/shared/pkg/storage/gcs"
 	"github.com/e2b-dev/infra/packages/shared/pkg/storage/header"
-	"github.com/google/uuid"
 )
 
 func main() {

packages/shared/pkg/storage/template.go

@@ -2,13 +2,15 @@ package storage
 
 import (
 	"fmt"
+	"os"
 	"path/filepath"
 
 	"github.com/e2b-dev/infra/packages/shared/pkg/storage/header"
 )
 
 const (
-	EnvsDisk = "/mnt/disks/fc-envs/v1"
+	EnvsDisk        = "/mnt/disks/fc-envs/v1"
+	LocalStorageDir = "/tmp"
 
 	KernelsDir     = "/fc-kernels"
 	KernelMountDir = "/fc-vm"
@@ -37,13 +39,21 @@ const (
 // Path to the directory where the kernel can be accessed inside when the dirs are mounted.
 var KernelMountedPath = filepath.Join(KernelMountDir, KernelName)
 
+type Type string
+
+const (
+	LocalStorage  Type = "local"
+	BucketStorage Type = "bucket"
+)
+
 type TemplateFiles struct {
 	TemplateId         string
 	BuildId            string
 	KernelVersion      string
 	FirecrackerVersion string
 
-	hugePages bool
+	hugePages   bool
+	StorageType Type
 }
 
 func NewTemplateFiles(
@@ -53,12 +63,22 @@ func NewTemplateFiles(
 	firecrackerVersion string,
 	hugePages bool,
 ) *TemplateFiles {
+	// Choose where are the template build data stored. Default to bucket storage.
+	var storageType Type
+	switch os.Getenv("TEMPLATE_STORAGE") {
+	case "local":
+		storageType = LocalStorage
+	default:
+		storageType = BucketStorage
+	}
+
 	return &TemplateFiles{
 		TemplateId:         templateId,
 		BuildId:            buildId,
 		KernelVersion:      kernelVersion,
 		FirecrackerVersion: firecrackerVersion,
 		hugePages:          hugePages,
+		StorageType:        storageType,
 	}
 }
 
@@ -112,7 +132,15 @@ func (t *TemplateFiles) StorageSnapfilePath() string {
 }
 
 func (t *TemplateFiles) BuildDir() string {
-	return filepath.Join(EnvsDisk, t.TemplateId, buildDirName, t.BuildId)
+	var baseDir string
+	switch t.StorageType {
+	case LocalStorage:
+		baseDir = LocalStorageDir
+	default:
+		baseDir = EnvsDisk
+	}
+
+	return filepath.Join(baseDir, t.TemplateId, buildDirName, t.BuildId)
 }
 
 func (t *TemplateFiles) BuildMemfilePath() string {

packages/shared/pkg/storage/template_bucket.go

@@ -0,0 +1,31 @@
+package storage
+
+import (
+	"context"
+
+	"github.com/e2b-dev/infra/packages/shared/pkg/storage/gcs"
+	"github.com/e2b-dev/infra/packages/shared/pkg/storage/header"
+)
+
+type TemplateBuild interface {
+	Remove(ctx context.Context) error
+	Upload(ctx context.Context, snapfilePath string, memfilePath *string, rootfsPath *string) chan error
+}
+
+func NewTemplateBuild(
+	memfileHeader *header.Header,
+	rootfsHeader *header.Header,
+	files *TemplateFiles,
+) TemplateBuild {
+	switch files.StorageType {
+	case LocalStorage:
+		return &TemplateLocalBuild{files: files}
+	default:
+		return &TemplateBucketBuild{
+			bucket:        gcs.GetTemplateBucket(),
+			memfileHeader: memfileHeader,
+			rootfsHeader:  rootfsHeader,
+			files:         files,
+		}
+	}
+}

packages/shared/pkg/storage/template_build.go → packages/shared/pkg/storage/template_build_bucket.go

@@ -12,7 +12,7 @@ import (
 	"github.com/e2b-dev/infra/packages/shared/pkg/storage/header"
 )
 
-type TemplateBuild struct {
+type TemplateBucketBuild struct {
 	files *TemplateFiles
 
 	memfileHeader *header.Header
@@ -21,20 +21,7 @@ type TemplateBuild struct {
 	bucket *gcs.BucketHandle
 }
 
-func NewTemplateBuild(
-	memfileHeader *header.Header,
-	rootfsHeader *header.Header,
-	files *TemplateFiles,
-) *TemplateBuild {
-	return &TemplateBuild{
-		bucket:        gcs.GetTemplateBucket(),
-		memfileHeader: memfileHeader,
-		rootfsHeader:  rootfsHeader,
-		files:         files,
-	}
-}
-
-func (t *TemplateBuild) Remove(ctx context.Context) error {
+func (t *TemplateBucketBuild) Remove(ctx context.Context) error {
 	err := gcs.RemoveDir(ctx, t.bucket, t.files.StorageDir())
 	if err != nil {
 		return fmt.Errorf("error when removing template build '%s': %w", t.files.StorageDir(), err)
@@ -43,7 +30,7 @@ func (t *TemplateBuild) Remove(ctx context.Context) error {
 	return nil
 }
 
-func (t *TemplateBuild) uploadMemfileHeader(ctx context.Context, h *header.Header) error {
+func (t *TemplateBucketBuild) uploadMemfileHeader(ctx context.Context, h *header.Header) error {
 	object := gcs.NewObject(ctx, t.bucket, t.files.StorageMemfileHeaderPath())
 
 	serialized, err := header.Serialize(h.Metadata, h.Mapping)
@@ -59,7 +46,7 @@ func (t *TemplateBuild) uploadMemfileHeader(ctx context.Context, h *header.Heade
 	return nil
 }
 
-func (t *TemplateBuild) uploadMemfile(ctx context.Context, memfilePath string) error {
+func (t *TemplateBucketBuild) uploadMemfile(ctx context.Context, memfilePath string) error {
 	object := gcs.NewObject(ctx, t.bucket, t.files.StorageMemfilePath())
 
 	err := object.UploadWithCli(ctx, memfilePath)
@@ -70,7 +57,7 @@ func (t *TemplateBuild) uploadMemfile(ctx context.Context, memfilePath string) e
 	return nil
 }
 
-func (t *TemplateBuild) uploadRootfsHeader(ctx context.Context, h *header.Header) error {
+func (t *TemplateBucketBuild) uploadRootfsHeader(ctx context.Context, h *header.Header) error {
 	object := gcs.NewObject(ctx, t.bucket, t.files.StorageRootfsHeaderPath())
 
 	serialized, err := header.Serialize(h.Metadata, h.Mapping)
@@ -86,7 +73,7 @@ func (t *TemplateBuild) uploadRootfsHeader(ctx context.Context, h *header.Header
 	return nil
 }
 
-func (t *TemplateBuild) uploadRootfs(ctx context.Context, rootfsPath string) error {
+func (t *TemplateBucketBuild) uploadRootfs(ctx context.Context, rootfsPath string) error {
 	object := gcs.NewObject(ctx, t.bucket, t.files.StorageRootfsPath())
 
 	err := object.UploadWithCli(ctx, rootfsPath)
@@ -98,7 +85,7 @@ func (t *TemplateBuild) uploadRootfs(ctx context.Context, rootfsPath string) err
 }
 
 // Snapfile is small enough so we dont use composite upload.
-func (t *TemplateBuild) uploadSnapfile(ctx context.Context, snapfile io.Reader) error {
+func (t *TemplateBucketBuild) uploadSnapfile(ctx context.Context, snapfile io.Reader) error {
 	object := gcs.NewObject(ctx, t.bucket, t.files.StorageSnapfilePath())
 
 	n, err := object.ReadFrom(snapfile)
@@ -109,7 +96,7 @@ func (t *TemplateBuild) uploadSnapfile(ctx context.Context, snapfile io.Reader)
 	return nil
 }
 
-func (t *TemplateBuild) Upload(
+func (t *TemplateBucketBuild) Upload(
 	ctx context.Context,
 	snapfilePath string,
 	memfilePath *string,

packages/shared/pkg/storage/template_build_local.go

@@ -0,0 +1,24 @@
+package storage
+
+import (
+	"context"
+)
+
+type TemplateLocalBuild struct {
+	files *TemplateFiles
+}
+
+func (t *TemplateLocalBuild) Remove(ctx context.Context) error {
+	return nil
+}
+
+func (t *TemplateLocalBuild) Upload(
+	ctx context.Context,
+	snapfilePath string,
+	memfilePath *string,
+	rootfsPath *string,
+) chan error {
+	chanErr := make(chan error, 1)
+	close(chanErr)
+	return chanErr
+}

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

@@ -15,7 +15,6 @@ import (
 
 	"github.com/e2b-dev/infra/packages/shared/pkg/storage"
 	"github.com/e2b-dev/infra/packages/template-manager/internal/build"
-	"github.com/e2b-dev/infra/packages/template-manager/internal/template"
 )
 
 func main() {
@@ -72,9 +71,7 @@ func Build(ctx context.Context, kernelVersion, fcVersion, templateID, buildID st
 		return fmt.Errorf("error building template: %w", err)
 	}
 
-	tempStorage := template.NewStorage(ctx)
-
-	buildStorage := tempStorage.NewBuild(t.TemplateFiles)
+	buildStorage := storage.NewTemplateBuild(nil, nil, t.TemplateFiles)
 
 	memfilePath := t.BuildMemfilePath()
 	rootfsPath := t.BuildRootfsPath()

packages/template-manager/internal/build/build.go

@@ -3,23 +3,24 @@ package build
 import (
 	"context"
 	"fmt"
+	"os"
+	"os/exec"
+	"strings"
+	"time"
+
 	"github.com/bits-and-blooms/bitset"
 	"github.com/docker/docker/client"
+	docker "github.com/fsouza/go-dockerclient"
+	"github.com/google/uuid"
+	"go.opentelemetry.io/otel/trace"
+
 	template_manager "github.com/e2b-dev/infra/packages/shared/pkg/grpc/template-manager"
 	"github.com/e2b-dev/infra/packages/shared/pkg/storage"
 	"github.com/e2b-dev/infra/packages/shared/pkg/storage/header"
 	"github.com/e2b-dev/infra/packages/shared/pkg/telemetry"
 	"github.com/e2b-dev/infra/packages/template-manager/internal/cache"
-	"github.com/e2b-dev/infra/packages/template-manager/internal/template"
-	docker "github.com/fsouza/go-dockerclient"
-	"github.com/google/uuid"
-	"go.opentelemetry.io/otel/trace"
-	"os"
 
 	"go.uber.org/zap"
-	"os/exec"
-	"strings"
-	"time"
 )
 
 type TemplateBuilder struct {
@@ -30,20 +31,18 @@ type TemplateBuilder struct {
 	buildLogger        *zap.Logger
 	dockerClient       *client.Client
 	legacyDockerClient *docker.Client
-	templateStorage    *template.Storage
 }
 
 const cleanupTimeout = time.Second * 10
 
-func NewBuilder(logger *zap.Logger, buildLogger *zap.Logger, tracer trace.Tracer, dockerClient *client.Client, legacyDockerClient *docker.Client, templateStorage *template.Storage, buildCache *cache.BuildCache) *TemplateBuilder {
+func NewBuilder(logger *zap.Logger, buildLogger *zap.Logger, tracer trace.Tracer, dockerClient *client.Client, legacyDockerClient *docker.Client, buildCache *cache.BuildCache) *TemplateBuilder {
 	return &TemplateBuilder{
 		logger:             logger,
 		tracer:             tracer,
 		buildCache:         buildCache,
 		buildLogger:        buildLogger,
 		dockerClient:       dockerClient,
 		legacyDockerClient: legacyDockerClient,
-		templateStorage:    templateStorage,
 	}
 }
 
@@ -86,20 +85,6 @@ func (b *TemplateBuilder) Build(ctx context.Context, template *Env, envID string
 		return err
 	}
 
-	// Remove build files if build fails or times out
-	defer func() {
-		if err != nil {
-			removeCtx, cancel := context.WithTimeout(context.Background(), cleanupTimeout)
-			defer cancel()
-
-			removeErr := b.templateStorage.Remove(removeCtx, buildID)
-			if removeErr != nil {
-				b.logger.Error("Error while removing build files", zap.Error(removeErr))
-				telemetry.ReportError(ctx, removeErr)
-			}
-		}
-	}()
-
 	// MEMFILE
 	memfilePath := template.BuildMemfilePath()
 	memfileDiffPath := template.BuildMemfileDiffPath()
@@ -221,6 +206,20 @@ func (b *TemplateBuilder) Build(ctx context.Context, template *Env, envID string
 		template.TemplateFiles,
 	)
 
+	// Remove build files if build fails or times out
+	defer func() {
+		if err != nil {
+			removeCtx, cancel := context.WithTimeout(context.Background(), cleanupTimeout)
+			defer cancel()
+
+			removeErr := templateBuild.Remove(removeCtx)
+			if removeErr != nil {
+				b.logger.Error("Error while removing build files", zap.Error(removeErr))
+				telemetry.ReportError(ctx, removeErr)
+			}
+		}
+	}()
+
 	upload := templateBuild.Upload(
 		ctx,
 		template.BuildSnapfilePath(),

packages/template-manager/internal/server/create_template.go

@@ -3,6 +3,7 @@ package server
 import (
 	"context"
 	"fmt"
+
 	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/trace"
 	"go.uber.org/zap"

packages/template-manager/internal/server/delete_template.go

@@ -16,7 +16,7 @@ func (s *ServerStore) TemplateBuildDelete(ctx context.Context, in *template_mana
 	s.wg.Add(1)
 	defer s.wg.Done()
 
-	err := template.Delete(childCtx, s.tracer, s.artifactRegistry, s.templateStorage, in.BuildID)
+	err := template.Delete(childCtx, s.tracer, s.artifactRegistry, s.templateBuild, in.BuildID)
 	if err != nil {
 		return nil, err
 	}