Optimize snapshots storage

packages/api/.gitignore

@@ -1,3 +1,4 @@
 bin
 mkfcenv.tar.gz
-.env
+.env
+.shared

packages/nomad/loki.hcl

@@ -57,8 +57,8 @@ job "loki" {
 
       resources {
         memory_max = 2048
-        memory = 1024
-        cpu    = 512
+        memory = 256
+        cpu    = 256
       }
 
       template {

packages/nomad/otel-collector.hcl

@@ -194,6 +194,15 @@ processors:
           - "nomad_client_unallocated_memory"
           - "orchestrator.*"
           - "api.*"
+  metricstransform:
+    transforms:
+      - include: "nomad_client_host_cpu_idle"
+        match_type: strict
+        action: update
+        operations:
+          - action: aggregate_labels
+            aggregation_type: sum
+            label_set: [instance, node_id, node_status, node_pool]
   attributes/session-proxy:
     actions:
       - key: service.name
@@ -249,7 +258,7 @@ service:
       receivers:
         - prometheus
         - otlp
-      processors: [filter, batch]
+      processors: [filter, batch, metricstransform]
       exporters:
         - prometheusremotewrite/grafana_cloud_metrics
     # metrics/session-proxy:

packages/nomad/proxies/client.conf

@@ -73,7 +73,7 @@ server {
   location / {
     if ($node_ip = "") {
       # If you set any text, the header will be set to `application/octet-stream` and then browser won't be able to render the content
-      return 404;
+      return 404; # Invalid sandbox url
     }
 
 
@@ -85,6 +85,14 @@ server {
   }
 }
 
+# Mock for sandbox server when the sandbox is not running, 127.0.0.1 is returned by the DNS resolver
+server {
+  listen 3003;
+
+  default_type text/plain;
+  return 404 'Sandbox is not running or not found.';
+}
+
 server {
   listen 3001;
   location /health {

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

@@ -69,14 +69,7 @@ func main() {
 	fmt.Printf("=======\n")
 
 	for _, mapping := range h.Mapping {
-		rangeMessage := fmt.Sprintf("%d-%d", mapping.Offset/h.Metadata.BlockSize, (mapping.Offset+mapping.Length-1)/h.Metadata.BlockSize)
-
-		fmt.Printf(
-			"%-14s [%11d,%11d) = [%11d,%11d) in %s, %d B\n",
-			rangeMessage,
-			mapping.Offset, mapping.Offset+mapping.Length,
-			mapping.BuildStorageOffset, mapping.BuildStorageOffset+mapping.Length, mapping.BuildId.String(), mapping.Length,
-		)
+		fmt.Println(mapping.Format(h.Metadata.BlockSize))
 	}
 
 	fmt.Printf("\nMAPPING SUMMARY\n")

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

@@ -0,0 +1,158 @@
+package main
+
+import (
+	"context"
+	"flag"
+	"fmt"
+	"log"
+	"os"
+
+	"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() {
+	baseBuildId := flag.String("base", "", "base build id")
+	diffBuildId := flag.String("diff", "", "diff build id")
+	kind := flag.String("kind", "", "'memfile' or 'rootfs'")
+	visualize := flag.Bool("visualize", false, "visualize the headers")
+
+	flag.Parse()
+
+	baseTemplate := storage.NewTemplateFiles(
+		"",
+		*baseBuildId,
+		"",
+		"",
+		false,
+	)
+
+	diffTemplate := storage.NewTemplateFiles(
+		"",
+		*diffBuildId,
+		"",
+		"",
+		false,
+	)
+
+	var baseStoragePath string
+	var diffStoragePath string
+
+	if *kind == "memfile" {
+		baseStoragePath = baseTemplate.StorageMemfileHeaderPath()
+		diffStoragePath = diffTemplate.StorageMemfileHeaderPath()
+	} else if *kind == "rootfs" {
+		baseStoragePath = baseTemplate.StorageRootfsHeaderPath()
+		diffStoragePath = diffTemplate.StorageRootfsHeaderPath()
+	} else {
+		log.Fatalf("invalid kind: %s", *kind)
+	}
+
+	ctx := context.Background()
+
+	baseObj := gcs.NewObject(ctx, gcs.TemplateBucket, baseStoragePath)
+	diffObj := gcs.NewObject(ctx, gcs.TemplateBucket, diffStoragePath)
+
+	baseHeader, err := header.Deserialize(baseObj)
+	if err != nil {
+		log.Fatalf("failed to deserialize base header: %s", err)
+	}
+
+	diffHeader, err := header.Deserialize(diffObj)
+	if err != nil {
+		log.Fatalf("failed to deserialize diff header: %s", err)
+	}
+
+	fmt.Printf("\nBASE METADATA\n")
+	fmt.Printf("Storage path       %s/%s\n", gcs.TemplateBucket.BucketName(), baseStoragePath)
+	fmt.Printf("========\n")
+
+	for _, mapping := range baseHeader.Mapping {
+		fmt.Println(mapping.Format(baseHeader.Metadata.BlockSize))
+	}
+
+	if *visualize {
+		bottomLayers := header.Layers(baseHeader.Mapping)
+		delete(*bottomLayers, baseHeader.Metadata.BaseBuildId)
+
+		fmt.Println("")
+		fmt.Println(
+			header.Visualize(
+				baseHeader.Mapping,
+				baseHeader.Metadata.Size,
+				baseHeader.Metadata.BlockSize,
+				128,
+				bottomLayers,
+				&map[uuid.UUID]struct{}{
+					baseHeader.Metadata.BuildId: {},
+				},
+			),
+		)
+	}
+
+	if err := header.ValidateMappings(baseHeader.Mapping, baseHeader.Metadata.Size, baseHeader.Metadata.BlockSize); err != nil {
+		log.Fatalf("failed to validate base header: %s", err)
+	}
+
+	fmt.Printf("\nDIFF METADATA\n")
+	fmt.Printf("Storage path       %s/%s\n", gcs.TemplateBucket.BucketName(), diffStoragePath)
+	fmt.Printf("========\n")
+
+	onlyDiffMappings := make([]*header.BuildMap, 0)
+
+	for _, mapping := range diffHeader.Mapping {
+		if mapping.BuildId == diffHeader.Metadata.BuildId {
+			onlyDiffMappings = append(onlyDiffMappings, mapping)
+		}
+	}
+
+	for _, mapping := range onlyDiffMappings {
+		fmt.Println(mapping.Format(baseHeader.Metadata.BlockSize))
+	}
+
+	if *visualize {
+		fmt.Println("")
+		fmt.Println(
+			header.Visualize(
+				onlyDiffMappings,
+				baseHeader.Metadata.Size,
+				baseHeader.Metadata.BlockSize,
+				128,
+				nil,
+				header.Layers(onlyDiffMappings),
+			),
+		)
+	}
+
+	mergedHeader := header.MergeMappings(baseHeader.Mapping, onlyDiffMappings)
+
+	fmt.Printf("\n\nMERGED METADATA\n")
+	fmt.Printf("========\n")
+
+	for _, mapping := range mergedHeader {
+		fmt.Println(mapping.Format(baseHeader.Metadata.BlockSize))
+	}
+
+	if *visualize {
+		bottomLayers := header.Layers(baseHeader.Mapping)
+		delete(*bottomLayers, baseHeader.Metadata.BaseBuildId)
+
+		fmt.Println("")
+		fmt.Println(
+			header.Visualize(
+				mergedHeader,
+				baseHeader.Metadata.Size,
+				baseHeader.Metadata.BlockSize,
+				128,
+				bottomLayers,
+				header.Layers(onlyDiffMappings),
+			),
+		)
+	}
+
+	if err := header.ValidateMappings(mergedHeader, baseHeader.Metadata.Size, baseHeader.Metadata.BlockSize); err != nil {
+		fmt.Fprintf(os.Stderr, "\n\n[VALIDATION ERROR]: failed to validate merged header: %s", err)
+	}
+}

packages/orchestrator/go.mod

@@ -15,7 +15,6 @@ require (
 	github.com/hashicorp/consul/api v1.30.0
 	github.com/jellydator/ttlcache/v3 v3.3.0
 	github.com/loopholelabs/userfaultfd-go v0.1.2
-	github.com/miekg/dns v1.1.62
 	github.com/pojntfx/go-nbd v0.3.2
 	github.com/shirou/gopsutil/v4 v4.24.10
 	github.com/vishvananda/netlink v1.3.0
@@ -93,6 +92,7 @@ require (
 	github.com/mdlayher/genetlink v1.3.2 // indirect
 	github.com/mdlayher/netlink v1.7.2 // indirect
 	github.com/mdlayher/socket v0.5.1 // indirect
+	github.com/miekg/dns v1.1.62 // indirect
 	github.com/mitchellh/go-homedir v1.1.0 // indirect
 	github.com/mitchellh/mapstructure v1.5.1-0.20220423185008-bf980b35cac4 // indirect
 	github.com/oklog/ulid v1.3.1 // indirect

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

@@ -102,6 +102,7 @@ func (b *File) Slice(off, length int64) ([]byte, error) {
 		return nil, fmt.Errorf("failed to get mapping: %w", err)
 	}
 
+	// Pass empty huge page when the build id is nil.
 	if *buildID == uuid.Nil {
 		return header.EmptyHugePage, nil
 	}

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

@@ -11,7 +11,7 @@ import (
 	"github.com/e2b-dev/infra/packages/shared/pkg/storage/gcs"
 )
 
-const buildExpiration = time.Hour * 25
+const buildExpiration = time.Hour * 48
 
 const cachePath = "/orchestrator/build"
 

packages/orchestrator/internal/sandbox/sandbox.go

@@ -84,7 +84,6 @@ func NewSandbox(
 		config.KernelVersion,
 		config.FirecrackerVersion,
 		config.HugePages,
-		isSnapshot,
 	)
 	if err != nil {
 		return nil, cleanup, fmt.Errorf("failed to get template snapshot data: %w", err)
@@ -393,7 +392,13 @@ func (s *Sandbox) Snapshot(
 		return nil, fmt.Errorf("failed to create memfile diff file: %w", err)
 	}
 
-	err = header.CreateDiff(sourceFile, s.files.MemfilePageSize(), memfileDirtyPages, memfileDiffFile)
+	memfileDirtyPages, emptyDirtyPages, err := header.CreateDiff(
+		sourceFile,
+		s.files.MemfilePageSize(),
+		memfileDirtyPages,
+		originalMemfile,
+		memfileDiffFile,
+	)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create memfile diff: %w", err)
 	}
@@ -402,15 +407,32 @@ func (s *Sandbox) Snapshot(
 
 	releaseLock()
 
-	memfileMapping := header.CreateMapping(
-		memfileMetadata,
+	var memfileMappings []*header.BuildMap
+
+	memfileEmptyMapping := header.CreateMapping(
+		&uuid.Nil,
+		emptyDirtyPages,
+		memfileMetadata.BlockSize,
+	)
+
+	if memfileEmptyMapping != nil {
+		memfileMappings = header.MergeMappings(
+			originalMemfile.Header().Mapping,
+			memfileEmptyMapping,
+		)
+
+		memfileMappings = header.NormalizeMappings(memfileMappings)
+	}
+
+	memfileDirtyMappings := header.CreateMapping(
 		&buildId,
 		memfileDirtyPages,
+		memfileMetadata.BlockSize,
 	)
 
-	memfileMappings := header.MergeMappings(
+	memfileMappings = header.MergeMappings(
 		originalMemfile.Header().Mapping,
-		memfileMapping,
+		memfileDirtyMappings,
 	)
 
 	snapfile, err := template.NewLocalFile(snapshotTemplateFiles.CacheSnapfilePath())
@@ -469,9 +491,9 @@ func (s *Sandbox) Snapshot(
 	}
 
 	rootfsMapping := header.CreateMapping(
-		rootfsMetadata,
 		&buildId,
 		rootfsDirtyBlocks,
+		rootfsMetadata.BlockSize,
 	)
 
 	rootfsMappings := header.MergeMappings(

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

@@ -14,7 +14,7 @@ import (
 
 // How long to keep the template in the cache since the last access.
 // Should be longer than the maximum possible sandbox lifetime.
-const templateExpiration = time.Hour * 25
+const templateExpiration = time.Hour * 72
 
 type Cache struct {
 	cache      *ttlcache.Cache[string, Template]
@@ -62,15 +62,13 @@ func (c *Cache) GetTemplate(
 	kernelVersion,
 	firecrackerVersion string,
 	hugePages bool,
-	isSnapshot bool,
 ) (Template, error) {
 	storageTemplate, err := newTemplateFromStorage(
 		templateId,
 		buildId,
 		kernelVersion,
 		firecrackerVersion,
 		hugePages,
-		isSnapshot,
 		nil,
 		nil,
 		c.bucket,
@@ -125,7 +123,6 @@ func (c *Cache) AddSnapshot(
 		kernelVersion,
 		firecrackerVersion,
 		hugePages,
-		true,
 		memfileHeader,
 		rootfsHeader,
 		c.bucket,