Split seekable files from non-seekable files

.mockery.yaml

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

packages/orchestrator/cmd/inspect-data/main.go

@@ -24,14 +24,17 @@ func main() {
 
 	var storagePath string
 	var blockSize int64
+	var objectType storage.SeekableObjectType
 
 	switch *kind {
 	case "memfile":
 		storagePath = template.StorageMemfilePath()
 		blockSize = 2097152
+		objectType = storage.MemfileObjectType
 	case "rootfs":
 		storagePath = template.StorageRootfsPath()
 		blockSize = 4096
+		objectType = storage.RootFSObjectType
 	default:
 		log.Fatalf("invalid kind: %s", *kind)
 	}
@@ -47,7 +50,7 @@ func main() {
 		log.Fatalf("failed to get storage provider: %s", err)
 	}
 
-	obj, err := storage.OpenObject(ctx, storagePath)
+	obj, err := storage.OpenSeekableObject(ctx, storagePath, objectType)
 	if err != nil {
 		log.Fatalf("failed to open object: %s", err)
 	}

packages/orchestrator/cmd/inspect-header/main.go

@@ -22,6 +22,7 @@ func main() {
 	}
 
 	var storagePath string
+	var objectType storage.ObjectType
 
 	switch *kind {
 	case "memfile":
@@ -30,6 +31,8 @@ func main() {
 		storagePath = template.StorageRootfsHeaderPath()
 	default:
 		log.Fatalf("invalid kind: %s", *kind)
+
+		return
 	}
 
 	ctx := context.Background()
@@ -38,7 +41,7 @@ func main() {
 		log.Fatalf("failed to get storage provider: %s", err)
 	}
 
-	obj, err := storage.OpenObject(ctx, storagePath)
+	obj, err := storage.OpenObject(ctx, storagePath, objectType)
 	if err != nil {
 		log.Fatalf("failed to open object: %s", err)
 	}

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

@@ -31,14 +31,17 @@ func main() {
 
 	var baseStoragePath string
 	var diffStoragePath string
+	var objectType storage.ObjectType
 
 	switch *kind {
 	case "memfile":
 		baseStoragePath = baseTemplate.StorageMemfileHeaderPath()
 		diffStoragePath = diffTemplate.StorageMemfileHeaderPath()
+		objectType = storage.MemfileHeaderObjectType
 	case "rootfs":
 		baseStoragePath = baseTemplate.StorageRootfsHeaderPath()
 		diffStoragePath = diffTemplate.StorageRootfsHeaderPath()
+		objectType = storage.RootFSHeaderObjectType
 	default:
 		log.Fatalf("invalid kind: %s", *kind)
 	}
@@ -50,12 +53,12 @@ func main() {
 		log.Fatalf("failed to get storage provider: %s", err)
 	}
 
-	baseObj, err := storageProvider.OpenObject(ctx, baseStoragePath)
+	baseObj, err := storageProvider.OpenObject(ctx, baseStoragePath, objectType)
 	if err != nil {
 		log.Fatalf("failed to open object: %s", err)
 	}
 
-	diffObj, err := storageProvider.OpenObject(ctx, diffStoragePath)
+	diffObj, err := storageProvider.OpenObject(ctx, diffStoragePath, objectType)
 	if err != nil {
 		log.Fatalf("failed to open object: %s", err)
 	}

packages/orchestrator/internal/sandbox/build/storage_diff.go

@@ -18,10 +18,12 @@ func storagePath(buildId string, diffType DiffType) string {
 }
 
 type StorageDiff struct {
-	chunker     *utils.SetOnce[*block.Chunker]
-	cachePath   string
-	cacheKey    DiffStoreKey
-	storagePath string
+	chunker           *utils.SetOnce[*block.Chunker]
+	cachePath         string
+	cacheKey          DiffStoreKey
+	storagePath       string
+	storageObjectType storage.SeekableObjectType
+
 	blockSize   int64
 	metrics     blockmetrics.Metrics
 	persistence storage.StorageProvider
@@ -40,17 +42,30 @@ func newStorageDiff(
 	cachePathSuffix := id.Generate()
 
 	storagePath := storagePath(buildId, diffType)
+	storageObjectType := storageObjectType(diffType)
 	cacheFile := fmt.Sprintf("%s-%s-%s", buildId, diffType, cachePathSuffix)
 	cachePath := filepath.Join(basePath, cacheFile)
 
 	return &StorageDiff{
-		storagePath: storagePath,
-		cachePath:   cachePath,
-		chunker:     utils.NewSetOnce[*block.Chunker](),
-		blockSize:   blockSize,
-		metrics:     metrics,
-		persistence: persistence,
-		cacheKey:    GetDiffStoreKey(buildId, diffType),
+		storagePath:       storagePath,
+		storageObjectType: storageObjectType,
+		cachePath:         cachePath,
+		chunker:           utils.NewSetOnce[*block.Chunker](),
+		blockSize:         blockSize,
+		metrics:           metrics,
+		persistence:       persistence,
+		cacheKey:          GetDiffStoreKey(buildId, diffType),
+	}
+}
+
+func storageObjectType(diffType DiffType) storage.SeekableObjectType {
+	switch diffType {
+	case Memfile:
+		return storage.MemfileObjectType
+	case Rootfs:
+		return storage.RootFSObjectType
+	default:
+		panic(fmt.Sprintf("unknown diff type: %s", diffType))
 	}
 }
 
@@ -59,7 +74,7 @@ func (b *StorageDiff) CacheKey() DiffStoreKey {
 }
 
 func (b *StorageDiff) Init(ctx context.Context) error {
-	obj, err := b.persistence.OpenObject(ctx, b.storagePath)
+	obj, err := b.persistence.OpenSeekableObject(ctx, b.storagePath, b.storageObjectType)
 	if err != nil {
 		return err
 	}

packages/orchestrator/internal/sandbox/template/storage.go

@@ -23,6 +23,28 @@ type Storage struct {
 	source *build.File
 }
 
+func storageHeaderObjectType(diffType build.DiffType) storage.ObjectType {
+	switch diffType {
+	case build.Memfile:
+		return storage.MemfileHeaderObjectType
+	case build.Rootfs:
+		return storage.RootFSHeaderObjectType
+	default:
+		panic(fmt.Sprintf("unknown diff type: %s", diffType))
+	}
+}
+
+func objectType(diffType build.DiffType) storage.SeekableObjectType {
+	switch diffType {
+	case build.Memfile:
+		return storage.MemfileObjectType
+	case build.Rootfs:
+		return storage.RootFSObjectType
+	default:
+		panic(fmt.Sprintf("unknown diff type: %s", diffType))
+	}
+}
+
 func NewStorage(
 	ctx context.Context,
 	store *build.DiffStore,
@@ -34,7 +56,8 @@ func NewStorage(
 ) (*Storage, error) {
 	if h == nil {
 		headerObjectPath := buildId + "/" + string(fileType) + storage.HeaderSuffix
-		headerObject, err := persistence.OpenObject(ctx, headerObjectPath)
+		headerObjectType := storageHeaderObjectType(fileType)
+		headerObject, err := persistence.OpenObject(ctx, headerObjectPath, headerObjectType)
 		if err != nil {
 			return nil, err
 		}
@@ -54,7 +77,8 @@ func NewStorage(
 	// If we can't find the diff header in storage, we try to find the "old" style template without a header as a fallback.
 	if h == nil {
 		objectPath := buildId + "/" + string(fileType)
-		object, err := persistence.OpenObject(ctx, objectPath)
+		objectType := objectType(fileType)
+		object, err := persistence.OpenSeekableObject(ctx, objectPath, objectType)
 		if err != nil {
 			return nil, err
 		}

packages/orchestrator/internal/sandbox/template/storage_file.go

@@ -18,6 +18,7 @@ func newStorageFile(
 	persistence storage.StorageProvider,
 	objectPath string,
 	path string,
+	objectType storage.ObjectType,
 ) (*storageFile, error) {
 	f, err := os.Create(path)
 	if err != nil {
@@ -26,7 +27,7 @@ func newStorageFile(
 
 	defer f.Close()
 
-	object, err := persistence.OpenObject(ctx, objectPath)
+	object, err := persistence.OpenObject(ctx, objectPath, objectType)
 	if err != nil {
 		return nil, err
 	}

packages/orchestrator/internal/sandbox/template/storage_template.go

@@ -87,6 +87,7 @@ func (t *storageTemplate) Fetch(ctx context.Context, buildStore *build.DiffStore
 			t.persistence,
 			t.files.StorageSnapfilePath(),
 			t.files.CacheSnapfilePath(),
+			storage.SnapfileObjectType,
 		)
 		if snapfileErr != nil {
 			errMsg := fmt.Errorf("failed to fetch snapfile: %w", snapfileErr)
@@ -119,6 +120,7 @@ func (t *storageTemplate) Fetch(ctx context.Context, buildStore *build.DiffStore
 			t.persistence,
 			t.files.StorageMetadataPath(),
 			t.files.CacheMetadataPath(),
+			storage.MetadataObjectType,
 		)
 		if err != nil && !errors.Is(err, storage.ErrObjectNotExist) {
 			sourceErr := fmt.Errorf("failed to fetch metafile: %w", err)

packages/orchestrator/internal/sandbox/template_build.go

@@ -38,7 +38,7 @@ func (t *TemplateBuild) Remove(ctx context.Context) error {
 }
 
 func (t *TemplateBuild) uploadMemfileHeader(ctx context.Context, h *headers.Header) error {
-	object, err := t.persistence.OpenObject(ctx, t.files.StorageMemfileHeaderPath())
+	object, err := t.persistence.OpenObject(ctx, t.files.StorageMemfileHeaderPath(), storage.MemfileHeaderObjectType)
 	if err != nil {
 		return err
 	}
@@ -57,7 +57,7 @@ func (t *TemplateBuild) uploadMemfileHeader(ctx context.Context, h *headers.Head
 }
 
 func (t *TemplateBuild) uploadMemfile(ctx context.Context, memfilePath string) error {
-	object, err := t.persistence.OpenObject(ctx, t.files.StorageMemfilePath())
+	object, err := t.persistence.OpenSeekableObject(ctx, t.files.StorageMemfilePath(), storage.MemfileObjectType)
 	if err != nil {
 		return err
 	}
@@ -71,7 +71,7 @@ func (t *TemplateBuild) uploadMemfile(ctx context.Context, memfilePath string) e
 }
 
 func (t *TemplateBuild) uploadRootfsHeader(ctx context.Context, h *headers.Header) error {
-	object, err := t.persistence.OpenObject(ctx, t.files.StorageRootfsHeaderPath())
+	object, err := t.persistence.OpenObject(ctx, t.files.StorageRootfsHeaderPath(), storage.RootFSHeaderObjectType)
 	if err != nil {
 		return err
 	}
@@ -90,7 +90,7 @@ func (t *TemplateBuild) uploadRootfsHeader(ctx context.Context, h *headers.Heade
 }
 
 func (t *TemplateBuild) uploadRootfs(ctx context.Context, rootfsPath string) error {
-	object, err := t.persistence.OpenObject(ctx, t.files.StorageRootfsPath())
+	object, err := t.persistence.OpenSeekableObject(ctx, t.files.StorageRootfsPath(), storage.RootFSObjectType)
 	if err != nil {
 		return err
 	}
@@ -105,7 +105,7 @@ func (t *TemplateBuild) uploadRootfs(ctx context.Context, rootfsPath string) err
 
 // Snap-file is small enough so we don't use composite upload.
 func (t *TemplateBuild) uploadSnapfile(ctx context.Context, path string) error {
-	object, err := t.persistence.OpenObject(ctx, t.files.StorageSnapfilePath())
+	object, err := t.persistence.OpenSeekableObject(ctx, t.files.StorageSnapfilePath(), storage.SnapfileName)
 	if err != nil {
 		return err
 	}
@@ -119,7 +119,7 @@ func (t *TemplateBuild) uploadSnapfile(ctx context.Context, path string) error {
 
 // Metadata is small enough so we don't use composite upload.
 func (t *TemplateBuild) uploadMetadata(ctx context.Context, path string) error {
-	object, err := t.persistence.OpenObject(ctx, t.files.StorageMetadataPath())
+	object, err := t.persistence.OpenObject(ctx, t.files.StorageMetadataPath(), storage.MetadataObjectType)
 	if err != nil {
 		return err
 	}

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

@@ -332,7 +332,7 @@ func getRootfsSize(
 	s storage.StorageProvider,
 	metadata storage.TemplateFiles,
 ) (uint64, error) {
-	obj, err := s.OpenObject(ctx, metadata.StorageRootfsHeaderPath())
+	obj, err := s.OpenObject(ctx, metadata.StorageRootfsHeaderPath(), storage.RootFSHeaderObjectType)
 	if err != nil {
 		return 0, fmt.Errorf("error opening rootfs header object: %w", err)
 	}

packages/orchestrator/internal/template/build/commands/copy.go

@@ -125,7 +125,7 @@ func (c *Copy) Execute(
 	}
 
 	// 1) Download the layer tar file from the storage to the local filesystem
-	obj, err := c.FilesStorage.OpenObject(ctx, paths.GetLayerFilesCachePath(c.CacheScope, step.GetFilesHash()))
+	obj, err := c.FilesStorage.OpenObject(ctx, paths.GetLayerFilesCachePath(c.CacheScope, step.GetFilesHash()), storage.LayerObjectType)
 	if err != nil {
 		return metadata.Context{}, fmt.Errorf("failed to open files object from storage: %w", err)
 	}

packages/orchestrator/internal/template/build/storage/cache/cache.go

@@ -63,7 +63,7 @@ func (h *HashIndex) LayerMetaFromHash(ctx context.Context, hash string) (LayerMe
 	ctx, span := tracer.Start(ctx, "get layer_metadata")
 	defer span.End()
 
-	obj, err := h.indexStorage.OpenObject(ctx, paths.HashToPath(h.cacheScope, hash))
+	obj, err := h.indexStorage.OpenObject(ctx, paths.HashToPath(h.cacheScope, hash), storage.LayerMetadataObjectType)
 	if err != nil {
 		return LayerMetadata{}, fmt.Errorf("error opening object for layer metadata: %w", err)
 	}
@@ -91,7 +91,7 @@ func (h *HashIndex) SaveLayerMeta(ctx context.Context, hash string, template Lay
 	ctx, span := tracer.Start(ctx, "save layer_metadata")
 	defer span.End()
 
-	obj, err := h.indexStorage.OpenObject(ctx, paths.HashToPath(h.cacheScope, hash))
+	obj, err := h.indexStorage.OpenObject(ctx, paths.HashToPath(h.cacheScope, hash), storage.LayerMetadataObjectType)
 	if err != nil {
 		return fmt.Errorf("error creating object for saving UUID: %w", err)
 	}

packages/orchestrator/internal/template/metadata/template_metadata.go

@@ -138,7 +138,7 @@ func fromTemplate(ctx context.Context, s storage.StorageProvider, files storage.
 	ctx, span := tracer.Start(ctx, "from template")
 	defer span.End()
 
-	obj, err := s.OpenObject(ctx, files.StorageMetadataPath())
+	obj, err := s.OpenObject(ctx, files.StorageMetadataPath(), storage.MetadataObjectType)
 	if err != nil {
 		return Template{}, fmt.Errorf("error opening object for template metadata: %w", err)
 	}

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

@@ -7,6 +7,7 @@ import (
 
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/template/build/storage/paths"
 	templatemanager "github.com/e2b-dev/infra/packages/shared/pkg/grpc/template-manager"
+	"github.com/e2b-dev/infra/packages/shared/pkg/storage"
 )
 
 const signedUrlExpiration = time.Minute * 30
@@ -22,7 +23,7 @@ func (s *ServerStore) InitLayerFileUpload(ctx context.Context, in *templatemanag
 	}
 
 	path := paths.GetLayerFilesCachePath(cacheScope, in.GetHash())
-	obj, err := s.buildStorage.OpenObject(ctx, path)
+	obj, err := s.buildStorage.OpenObject(ctx, path, storage.LayerObjectType)
 	if err != nil {
 		return nil, fmt.Errorf("failed to open layer files cache: %w", err)
 	}
@@ -32,8 +33,7 @@ func (s *ServerStore) InitLayerFileUpload(ctx context.Context, in *templatemanag
 		return nil, fmt.Errorf("failed to get signed url: %w", err)
 	}
 
-	_, err = obj.Size(ctx)
-	if err != nil {
+	if exists, _ := obj.Exists(ctx); !exists {
 		return &templatemanager.InitLayerFileUploadResponse{
 			Present: false,
 			Url:     &signedUrl,