Merge branch 'main' of github.com:e2b-dev/infra into lev-compression-final

Author: levb

.github/workflows/build-and-upload-job.yml

@@ -42,8 +42,15 @@ jobs:
         with:
           ref: ${{ inputs.commit_sha }}
 
+      - name: Checkout actions
+        uses: actions/checkout@v5
+        with:
+          sparse-checkout: |
+            .github/actions
+          path: actions-from-ref
+
       - name: Setup environment
-        uses: ./.github/actions/deploy-setup
+        uses: ./actions-from-ref/.github/actions/deploy-setup
         with:
           environment: ${{ inputs.environment }}
           install_gcloud: "true"

.github/workflows/deploy-infra.yml

@@ -40,8 +40,15 @@ jobs:
         with:
           ref: ${{ inputs.commit_sha }}
 
+      - name: Checkout actions
+        uses: actions/checkout@v5
+        with:
+          sparse-checkout: |
+            .github/actions
+          path: actions-from-ref
+
       - name: Setup environment
-        uses: ./.github/actions/deploy-setup
+        uses: ./actions-from-ref/.github/actions/deploy-setup
         with:
           environment: ${{ inputs.environment }}
           install_gcloud: "true"

.github/workflows/deploy-job.yml

@@ -43,8 +43,15 @@ jobs:
         with:
           ref: ${{ inputs.commit_sha }}
 
+      - name: Checkout actions
+        uses: actions/checkout@v5
+        with:
+          sparse-checkout: |
+            .github/actions
+          path: actions-from-ref
+
       - name: Setup environment
-        uses: ./.github/actions/deploy-setup
+        uses: ./actions-from-ref/.github/actions/deploy-setup
         with:
           environment: ${{ inputs.environment }}
           infisical_machine_identity_id: ${{ vars.INFISICAL_MACHINE_IDENTITY_ID }}

.github/workflows/pr-tests.yml

@@ -39,8 +39,6 @@ jobs:
     steps:
       - name: Checkout repository
         uses: actions/checkout@v5
-        with:
-          ref: ${{ inputs.commit_sha }}
 
       - name: Setup Go
         uses: ./.github/actions/go-setup-cache

packages/orchestrator/benchmarks/benchmark_test.go

@@ -80,7 +80,7 @@ func BenchmarkBaseImageLaunch(b *testing.B) {
 		return utils.Must(filepath.Abs(s))
 	}
 
-	endpoint := os.Getenv("OTEL_EXPORTER_OTLP_ENDPOINT")
+	endpoint := os.Getenv("OTEL_COLLECTOR_GRPC_ENDPOINT")
 	if endpoint != "" {
 		spanExporter, err := telemetry.NewSpanExporter(b.Context(),
 			otlptracegrpc.WithEndpoint(endpoint),

packages/orchestrator/benchmarks/concurrent_benchmark_test.go

@@ -33,8 +33,10 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc"
 	"go.opentelemetry.io/otel/metric/noop"
+	"go.opentelemetry.io/otel/trace"
 
 	"github.com/e2b-dev/infra/packages/clickhouse/pkg/hoststats"
 	"github.com/e2b-dev/infra/packages/orchestrator/pkg/cfg"
@@ -137,11 +139,20 @@ func BenchmarkConcurrentResume(b *testing.B) {
 		kernelVersion   = "vmlinux-6.1.158"
 		fcVersion       = featureflags.DefaultFirecrackerVersion
 		templateID      = "fcb33d09-3141-42c4-8d3b-c2df411681db"
-		buildID         = "ba6aae36-74f7-487a-b6f7-74fd7c94e479"
-		useHugePages    = false
 		templateVersion = "v2.0.0"
+
+		buildIDNormal    = "ba6aae36-74f7-487a-b6f7-74fd7c94e479"
+		buildIDHugePages = "ba6aae36-74f7-487a-b6f7-74fd7c94e480"
 	)
 
+	disableHP, _ := strconv.ParseBool(os.Getenv("DISABLE_HUGE_PAGES"))
+	useHugePages := !disableHP
+	buildID := buildIDHugePages
+	if !useHugePages {
+		buildID = buildIDNormal
+		b.Log("huge pages disabled")
+	}
+
 	// cache & ephemeral paths
 	persistenceDir := getPersistenceDir()
 	kernelsDir := filepath.Join(persistenceDir, "kernels")
@@ -155,7 +166,7 @@ func BenchmarkConcurrentResume(b *testing.B) {
 	}
 
 	// optional OTEL tracing
-	endpoint := os.Getenv("OTEL_EXPORTER_OTLP_ENDPOINT")
+	endpoint := os.Getenv("OTEL_COLLECTOR_GRPC_ENDPOINT")
 	if endpoint != "" {
 		spanExporter, err := telemetry.NewSpanExporter(b.Context(),
 			otlptracegrpc.WithEndpoint(endpoint),
@@ -423,9 +434,17 @@ func runConcurrentResume(
 			// Wait for the barrier.
 			<-gate
 
+			ctx, span := tracer.Start(b.Context(), "bench-resume",
+				trace.WithAttributes(
+					attribute.Int("concurrency", n),
+					attribute.Int("sandbox.index", i),
+					attribute.String("sandbox.id", runtime.SandboxID),
+				),
+			)
+
 			start := time.Now()
 			sbx, err := factory.ResumeSandbox(
-				b.Context(),
+				ctx,
 				tmpl,
 				config,
 				runtime,
@@ -434,6 +453,7 @@ func runConcurrentResume(
 				nil,
 			)
 			elapsed := time.Since(start)
+			span.End()
 
 			results[i] = concurrencyResult{
 				sandboxID: runtime.SandboxID,

packages/orchestrator/cmd/show-build-diff/main.go

@@ -116,7 +116,7 @@ func main() {
 	fmt.Printf("Storage path       %s\n", diffSource)
 	fmt.Printf("========\n")
 
-	onlyDiffMappings := make([]*header.BuildMap, 0)
+	onlyDiffMappings := make([]header.BuildMap, 0)
 
 	for _, mapping := range diffHeader.Mapping {
 		if mapping.BuildId == diffHeader.Metadata.BuildId {

packages/orchestrator/pkg/sandbox/fc/process.go

@@ -489,7 +489,10 @@ func (p *Process) Resume(
 	})
 
 	eg.Go(func() error {
-		err := socket.Wait(egCtx, uffdSocketPath)
+		ctx, uffdSpan := tracer.Start(egCtx, "wait-uffd-socket")
+		err := socket.Wait(ctx, uffdSocketPath)
+		uffdSpan.End()
+
 		if err != nil {
 			return fmt.Errorf("error waiting for uffd socket: %w", err)
 		}
@@ -500,7 +503,10 @@ func (p *Process) Resume(
 	})
 
 	eg.Go(func() error {
+		_, rootfsSpan := tracer.Start(egCtx, "wait-rootfs-path")
 		rootfsPath, err := p.rootfsProvider.Path()
+		rootfsSpan.End()
+
 		if err != nil {
 			return fmt.Errorf("error getting rootfs path: %w", err)
 		}

packages/orchestrator/pkg/sandbox/nbd/testutils/zero_device.go

@@ -24,7 +24,7 @@ func NewZeroDevice(size int64, blockSize int64) (*ZeroDevice, error) {
 		uint64(blockSize),
 		uint64(size),
 	),
-		[]*header.BuildMap{
+		[]header.BuildMap{
 			{
 				Offset:             0,
 				Length:             uint64(size),

packages/shared/pkg/storage/header/header.go

@@ -6,8 +6,8 @@ import (
 	"fmt"
 	"maps"
 	"slices"
+	"sort"
 
-	"github.com/bits-and-blooms/bitset"
 	"github.com/google/uuid"
 	"go.uber.org/zap"
 
@@ -31,27 +31,19 @@ type Header struct {
 	// This means every V4 header has a complete map of all builds referenced
 	// in its Mapping. V3 headers have no BuildFiles; the read path falls back
 	// to a Size() RPC for those.
-	BuildFiles  map[uuid.UUID]BuildFileInfo
-	blockStarts *bitset.BitSet
-	startMap    map[int64]*BuildMap
+	BuildFiles map[uuid.UUID]BuildFileInfo
 
-	Mapping []*BuildMap
+	Mapping []BuildMap
 }
 
 // CloneForUpload returns a clone with copied Mapping and BuildFiles, safe to
 // mutate for serialization without racing with concurrent readers of the
-// original. Only serialization-relevant fields are populated (Metadata,
-// Mapping, BuildFiles); lookup indices (blockStarts, startMap) are left nil.
+// original.
 func (t *Header) CloneForUpload() *Header {
-	mappings := make([]*BuildMap, len(t.Mapping))
-	for i, m := range t.Mapping {
-		mappings[i] = m.Copy()
-	}
-
 	metaCopy := *t.Metadata
 	clone := &Header{
 		Metadata: &metaCopy,
-		Mapping:  mappings,
+		Mapping:  slices.Clone(t.Mapping),
 	}
 
 	if t.BuildFiles != nil {
@@ -62,37 +54,23 @@ func (t *Header) CloneForUpload() *Header {
 	return clone
 }
 
-func NewHeader(metadata *Metadata, mapping []*BuildMap) (*Header, error) {
+func NewHeader(metadata *Metadata, mapping []BuildMap) (*Header, error) {
 	if metadata.BlockSize == 0 {
 		return nil, fmt.Errorf("block size cannot be zero")
 	}
 
 	if len(mapping) == 0 {
-		mapping = []*BuildMap{{
+		mapping = []BuildMap{{
 			Offset:             0,
 			Length:             metadata.Size,
 			BuildId:            metadata.BuildId,
 			BuildStorageOffset: 0,
 		}}
 	}
 
-	blocks := TotalBlocks(int64(metadata.Size), int64(metadata.BlockSize))
-
-	intervals := bitset.New(uint(blocks))
-	startMap := make(map[int64]*BuildMap, len(mapping))
-
-	for _, m := range mapping {
-		block := BlockIdx(int64(m.Offset), int64(metadata.BlockSize))
-
-		intervals.Set(uint(block))
-		startMap[block] = m
-	}
-
 	return &Header{
-		blockStarts: intervals,
-		Metadata:    metadata,
-		Mapping:     mapping,
-		startMap:    startMap,
+		Metadata: metadata,
+		Mapping:  mapping,
 	}, nil
 }
 
@@ -147,7 +125,6 @@ func (t *Header) GetShiftedMapping(ctx context.Context, offset int64) (BuildMap,
 	return b, nil
 }
 
-// TODO: Maybe we can optimize mapping by automatically assuming the mapping is uuid.Nil if we don't find it + stopping storing the nil mapping.
 func (t *Header) getMapping(ctx context.Context, offset int64) (*BuildMap, int64, error) {
 	if offset < 0 || offset >= int64(t.Metadata.Size) {
 		if t.IsNormalizeFixApplied() {
@@ -172,30 +149,26 @@ func (t *Header) getMapping(ctx context.Context, offset int64) (*BuildMap, int64
 		)
 	}
 
-	block := BlockIdx(offset, int64(t.Metadata.BlockSize))
+	i := sort.Search(len(t.Mapping), func(i int) bool {
+		return int64(t.Mapping[i].Offset) > offset
+	})
 
-	start, ok := t.blockStarts.PreviousSet(uint(block))
-	if !ok {
+	if i == 0 {
 		return nil, 0, fmt.Errorf("no source found for offset %d", offset)
 	}
 
-	mapping, ok := t.startMap[int64(start)]
-	if !ok {
-		return nil, 0, fmt.Errorf("no mapping found for offset %d", offset)
-	}
-
-	shift := (block - int64(start)) * int64(t.Metadata.BlockSize)
+	mapping := &t.Mapping[i-1]
+	shift := offset - int64(mapping.Offset)
 
 	// Verify that the offset falls within this mapping's range
 	if shift >= int64(mapping.Length) {
 		if t.IsNormalizeFixApplied() {
-			return nil, 0, fmt.Errorf("offset %d (block %d) is beyond the end of mapping at offset %d (ends at %d)",
-				offset, block, mapping.Offset, mapping.Offset+mapping.Length)
+			return nil, 0, fmt.Errorf("offset %d is beyond the end of mapping at offset %d (ends at %d)",
+				offset, mapping.Offset, mapping.Offset+mapping.Length)
 		}
 
 		logger.L().Warn(ctx, "offset is beyond the end of mapping, but normalize fix is not applied",
 			zap.Int64("offset", offset),
-			zap.Int64("block", block),
 			zap.Uint64("mappingOffset", mapping.Offset),
 			zap.Uint64("mappingEnd", mapping.Offset+mapping.Length),
 			logger.WithBuildID(t.Metadata.BuildId.String()),
@@ -228,9 +201,8 @@ func ValidateHeader(h *Header) error {
 	}
 
 	// Sort mappings by offset to check for gaps/overlaps
-	sortedMappings := make([]*BuildMap, len(h.Mapping))
-	copy(sortedMappings, h.Mapping)
-	slices.SortFunc(sortedMappings, func(a, b *BuildMap) int {
+	sortedMappings := slices.Clone(h.Mapping)
+	slices.SortFunc(sortedMappings, func(a, b BuildMap) int {
 		return cmp.Compare(a.Offset, b.Offset)
 	})