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

Author: jakubno

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,22 @@
+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 {
+	return make(chan error)
+}

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/server/create_template.go

@@ -3,14 +3,14 @@ package server
 import (
 	"context"
 	"fmt"
-	"go.uber.org/zap"
-	"go.uber.org/zap/zapcore"
 	"os/exec"
 	"strconv"
 	"strings"
 	"time"
 
 	"go.opentelemetry.io/otel/attribute"
+	"go.uber.org/zap"
+	"go.uber.org/zap/zapcore"
 	"google.golang.org/grpc/metadata"
 
 	template_manager "github.com/e2b-dev/infra/packages/shared/pkg/grpc/template-manager"
@@ -63,7 +63,7 @@ func (s *serverStore) TemplateCreate(templateRequest *template_manager.TemplateC
 		BuildLogsWriter: logsWriter,
 	}
 
-	buildStorage := s.templateStorage.NewBuild(template.TemplateFiles)
+	buildStorage := storage.NewTemplateBuild(nil, nil, template.TemplateFiles)
 
 	var err error
 

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

@@ -13,7 +13,7 @@ func (s *serverStore) TemplateBuildDelete(ctx context.Context, in *template_mana
 	childCtx, childSpan := s.tracer.Start(ctx, "template-delete-request")
 	defer childSpan.End()
 
-	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
 	}

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

@@ -7,6 +7,7 @@ import (
 	artifactregistry "cloud.google.com/go/artifactregistry/apiv1"
 	"github.com/docker/docker/client"
 	docker "github.com/fsouza/go-dockerclient"
+
 	"github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors/logging"
 	"github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors/recovery"
 	"github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors/selector"
@@ -22,8 +23,8 @@ import (
 	e2bgrpc "github.com/e2b-dev/infra/packages/shared/pkg/grpc"
 	templatemanager "github.com/e2b-dev/infra/packages/shared/pkg/grpc/template-manager"
 	l "github.com/e2b-dev/infra/packages/shared/pkg/logger"
+	"github.com/e2b-dev/infra/packages/shared/pkg/storage"
 	"github.com/e2b-dev/infra/packages/template-manager/internal/constants"
-	"github.com/e2b-dev/infra/packages/template-manager/internal/template"
 )
 
 type serverStore struct {
@@ -35,7 +36,7 @@ type serverStore struct {
 	dockerClient       *client.Client
 	legacyDockerClient *docker.Client
 	artifactRegistry   *artifactregistry.Client
-	templateStorage    *template.Storage
+	templateBuild      storage.TemplateBuild
 }
 
 func New(logger *zap.Logger, buildLogger *zap.Logger) *grpc.Server {
@@ -81,16 +82,13 @@ func New(logger *zap.Logger, buildLogger *zap.Logger) *grpc.Server {
 		panic(err)
 	}
 
-	templateStorage := template.NewStorage(ctx)
-
 	templatemanager.RegisterTemplateServiceServer(s, &serverStore{
 		tracer:             otel.Tracer(constants.ServiceName),
 		logger:             logger,
 		buildLogger:        buildLogger,
 		dockerClient:       dockerClient,
 		legacyDockerClient: legacyClient,
 		artifactRegistry:   artifactRegistry,
-		templateStorage:    templateStorage,
 	})
 
 	grpc_health_v1.RegisterHealthServer(s, health.NewServer())

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

@@ -10,9 +10,11 @@ import (
 	"go.opentelemetry.io/otel/trace"
 	"google.golang.org/grpc/codes"
 
+	"github.com/gogo/status"
+
 	"github.com/e2b-dev/infra/packages/shared/pkg/consts"
+	"github.com/e2b-dev/infra/packages/shared/pkg/storage"
 	"github.com/e2b-dev/infra/packages/shared/pkg/telemetry"
-	"github.com/gogo/status"
 )
 
 func GetDockerImageURL(templateID string) string {
@@ -24,13 +26,13 @@ func Delete(
 	ctx context.Context,
 	tracer trace.Tracer,
 	artifactRegistry *artifactregistry.Client,
-	templateStorage *Storage,
+	templateBuild storage.TemplateBuild,
 	buildId string,
 ) error {
 	childCtx, childSpan := tracer.Start(ctx, "delete-template")
 	defer childSpan.End()
 
-	err := templateStorage.Remove(ctx, buildId)
+	err := templateBuild.Remove(ctx)
 	if err != nil {
 		return fmt.Errorf("error when deleting template objects: %w", err)
 	}