fix(orchestrator): stop nbd read flood for permanently-missing build objects

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

@@ -4,6 +4,7 @@ package build
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"math"
 	"os"
@@ -19,6 +20,7 @@ import (
 
 	"github.com/e2b-dev/infra/packages/orchestrator/pkg/cfg"
 	"github.com/e2b-dev/infra/packages/shared/pkg/featureflags"
+	"github.com/e2b-dev/infra/packages/shared/pkg/storage"
 	"github.com/e2b-dev/infra/packages/shared/pkg/logger"
 	"github.com/e2b-dev/infra/packages/shared/pkg/units"
 	"github.com/e2b-dev/infra/packages/shared/pkg/utils"
@@ -39,6 +41,14 @@ type deleteDiff struct {
 	closeOnce sync.Once
 }
 
+const (
+	// negativeCacheTTL is how long a permanently-missing build object is held in
+	// the negative cache before being rechecked. Keeps the kernel-retry flood from
+	// hammering GCS/S3 with repeated 404s while still recovering when a build is
+	// eventually uploaded or a header arrives.
+	negativeCacheTTL = 60 * time.Second
+)
+
 type DiffStore struct {
 	cachePath string
 	cache     *ttlcache.Cache[DiffStoreKey, Diff]
@@ -54,6 +64,10 @@ type DiffStore struct {
 	pdDelay time.Duration
 
 	insertionTimes sync.Map // map[DiffStoreKey]time.Time — tracks when each diff was cached
+	// negativeCache holds keys for build objects that are permanently missing from
+	// the store. ttlcache handles automatic expiration so entries never accumulate
+	// unboundedly.
+	negativeCache *ttlcache.Cache[DiffStoreKey, struct{}]
 }
 
 func NewDiffStore(
@@ -71,14 +85,19 @@ func NewDiffStore(
 		ttlcache.WithTTL[DiffStoreKey, Diff](ttl),
 	)
 
+	negCache := ttlcache.New(
+		ttlcache.WithTTL[DiffStoreKey, struct{}](negativeCacheTTL),
+	)
+
 	ds := &DiffStore{
-		cachePath: cachePath,
-		cache:     cache,
-		cancel:    func() {},
-		config:    config,
-		flags:     flags,
-		pdSizes:   make(map[DiffStoreKey]*deleteDiff),
-		pdDelay:   delay,
+		cachePath:     cachePath,
+		cache:         cache,
+		negativeCache: negCache,
+		cancel:        func() {},
+		config:        config,
+		flags:         flags,
+		pdSizes:       make(map[DiffStoreKey]*deleteDiff),
+		pdDelay:       delay,
 	}
 
 	cache.OnEviction(func(ctx context.Context, _ ttlcache.EvictionReason, item *ttlcache.Item[DiffStoreKey, Diff]) {
@@ -111,12 +130,14 @@ func (s *DiffStore) Start(ctx context.Context) {
 	s.cancel = cancel
 
 	go s.cache.Start()
+	go s.negativeCache.Start()
 	go s.startDiskSpaceEviction(ctx, s.config, s.flags)
 }
 
 func (s *DiffStore) Close() {
 	s.cancel()
 	s.cache.Stop()
+	s.negativeCache.Stop()
 }
 
 // Get returns the cached Diff for key, refreshing TTL and cancelling any
@@ -135,15 +156,26 @@ func (s *DiffStore) Get(key DiffStoreKey) (Diff, bool) {
 // singleflight to construct + cache a new one. The create closure is invoked
 // at most once per key across concurrent callers; on success the returned Diff
 // is cached and its insertion time recorded.
+//
+// When create fails with storage.ErrObjectNotExist the key is placed in the
+// negative cache for negativeCacheTTL, short-circuiting all subsequent calls
+// for that key and preventing repeated storage 404s under kernel NBD retries.
 func (s *DiffStore) GetOrCreate(ctx context.Context, key DiffStoreKey, create func(context.Context) (Diff, error)) (Diff, error) {
 	s.resetDelete(key)
 
+	if s.negativeCache.Has(key) {
+		return nil, storage.ErrObjectNotExist
+	}
+
 	if item := s.cache.Get(key); item != nil {
 		return item.Value(), nil
 	}
 
 	v, err, _ := s.initGroup.Do(string(key), func() (any, error) {
-		// Double-check: another goroutine may have cached it while we waited.
+		// Double-check: another goroutine may have populated either cache while we waited.
+		if s.negativeCache.Has(key) {
+			return nil, storage.ErrObjectNotExist
+		}
 		if item := s.cache.Get(key); item != nil {
 			return item.Value(), nil
 		}
@@ -152,6 +184,16 @@ func (s *DiffStore) GetOrCreate(ctx context.Context, key DiffStoreKey, create fu
 
 		diff, err := create(ctx)
 		if err != nil {
+			if errors.Is(err, storage.ErrObjectNotExist) {
+				// Recheck the positive cache before installing the negative entry.
+				// Add may have raced with this storage probe: it deletes the negative
+				// entry and populates the positive cache while create was in-flight.
+				// If that happened, installing a negative entry here would shadow the
+				// available diff for up to negativeCacheTTL.
+				if s.cache.Get(key) == nil {
+					s.negativeCache.Set(key, struct{}{}, ttlcache.DefaultTTL)
+				}
+			}
 			return nil, err
 		}
 
@@ -168,6 +210,8 @@ func (s *DiffStore) GetOrCreate(ctx context.Context, key DiffStoreKey, create fu
 }
 
 func (s *DiffStore) Add(d Diff) {
+	// Clear any negative cache entry so a re-uploaded build is immediately visible.
+	s.negativeCache.Delete(d.CacheKey())
 	s.resetDelete(d.CacheKey())
 	s.cache.Set(d.CacheKey(), d, ttlcache.DefaultTTL)
 	s.insertionTimes.LoadOrStore(d.CacheKey(), time.Now())

packages/orchestrator/pkg/sandbox/build/cache_test.go

@@ -15,6 +15,8 @@ package build
 // causing a race when closing the cancel channel.
 
 import (
+	"context"
+	"errors"
 	"fmt"
 	"sync"
 	"testing"
@@ -30,6 +32,7 @@ import (
 	"github.com/e2b-dev/infra/packages/orchestrator/pkg/cfg"
 	blockmetrics "github.com/e2b-dev/infra/packages/orchestrator/pkg/sandbox/block/metrics"
 	"github.com/e2b-dev/infra/packages/shared/pkg/featureflags"
+	"github.com/e2b-dev/infra/packages/shared/pkg/storage"
 	"github.com/e2b-dev/infra/packages/shared/pkg/storage/header"
 )
 
@@ -644,3 +647,156 @@ func mustParseCfg(t *testing.T) cfg.Config {
 
 	return c
 }
+
+func TestGetOrCreate_NegativeCacheShortCircuits(t *testing.T) {
+	store, err := NewDiffStore(
+		mustParseCfg(t),
+		flagsWithMaxBuildCachePercentage(t, 90),
+		t.TempDir(),
+		time.Hour,
+		time.Minute,
+	)
+	require.NoError(t, err)
+	store.Start(t.Context())
+	t.Cleanup(store.Close)
+
+	key := DiffStoreKey(uuid.New().String())
+	calls := 0
+	create := func(ctx context.Context) (Diff, error) {
+		calls++
+		return nil, storage.ErrObjectNotExist
+	}
+
+	// First call: create is invoked, ErrObjectNotExist is cached negatively.
+	_, err = store.GetOrCreate(t.Context(), key, create)
+	require.ErrorIs(t, err, storage.ErrObjectNotExist)
+	require.Equal(t, 1, calls, "create should be called once on first miss")
+
+	// Subsequent calls within TTL must short-circuit without invoking create.
+	for i := 0; i < 3; i++ {
+		_, err = store.GetOrCreate(t.Context(), key, create)
+		require.ErrorIs(t, err, storage.ErrObjectNotExist)
+	}
+	require.Equal(t, 1, calls, "create must not be called again while negative cache is hot")
+}
+
+func TestGetOrCreate_NegativeCacheNotSetForOtherErrors(t *testing.T) {
+	store, err := NewDiffStore(
+		mustParseCfg(t),
+		flagsWithMaxBuildCachePercentage(t, 90),
+		t.TempDir(),
+		time.Hour,
+		time.Minute,
+	)
+	require.NoError(t, err)
+	store.Start(t.Context())
+	t.Cleanup(store.Close)
+
+	key := DiffStoreKey(uuid.New().String())
+	sentinel := errors.New("transient network error")
+	calls := 0
+	create := func(ctx context.Context) (Diff, error) {
+		calls++
+		return nil, sentinel
+	}
+
+	// Two calls for a non-ErrObjectNotExist error — no negative caching expected.
+	_, err = store.GetOrCreate(t.Context(), key, create)
+	require.ErrorIs(t, err, sentinel)
+	_, err = store.GetOrCreate(t.Context(), key, create)
+	require.ErrorIs(t, err, sentinel)
+	require.Equal(t, 2, calls, "create should be retried after non-ErrObjectNotExist errors")
+}
+
+func TestGetOrCreate_AddClearsNegativeCache(t *testing.T) {
+	cachePath := t.TempDir()
+	store, err := NewDiffStore(
+		mustParseCfg(t),
+		flagsWithMaxBuildCachePercentage(t, 90),
+		cachePath,
+		time.Hour,
+		time.Minute,
+	)
+	require.NoError(t, err)
+	store.Start(t.Context())
+	t.Cleanup(store.Close)
+
+	diff := newRootFSDiff(t, cachePath, uuid.New().String())
+	key := diff.CacheKey()
+
+	calls := 0
+	notFound := func(ctx context.Context) (Diff, error) {
+		calls++
+		return nil, storage.ErrObjectNotExist
+	}
+
+	// Populate negative cache.
+	_, err = store.GetOrCreate(t.Context(), key, notFound)
+	require.ErrorIs(t, err, storage.ErrObjectNotExist)
+	require.Equal(t, 1, calls)
+
+	// Explicit Add (simulating a build upload completing) clears the negative cache.
+	store.Add(diff)
+
+	returnDiff := func(ctx context.Context) (Diff, error) {
+		calls++
+		return diff, nil
+	}
+	got, err := store.GetOrCreate(t.Context(), key, returnDiff)
+	// After Add, the positive cache holds the diff — create should not be called.
+	require.NoError(t, err)
+	require.Equal(t, diff, got)
+	require.Equal(t, 1, calls, "create must not be called after Add populates the positive cache")
+}
+
+func TestGetOrCreate_AddRaceWithStorageProbe(t *testing.T) {
+	// Regression: if Add() populates the positive cache while create's storage
+	// probe is still in-flight and later returns ErrObjectNotExist, the negative
+	// cache must NOT shadow the diff that Add already installed.
+	cachePath := t.TempDir()
+	store, err := NewDiffStore(
+		mustParseCfg(t),
+		flagsWithMaxBuildCachePercentage(t, 90),
+		cachePath,
+		time.Hour,
+		time.Minute,
+	)
+	require.NoError(t, err)
+	store.Start(t.Context())
+	t.Cleanup(store.Close)
+
+	diff := newRootFSDiff(t, cachePath, uuid.New().String())
+	key := diff.CacheKey()
+
+	// ready gates the storage probe so we can inject Add() before it returns.
+	ready := make(chan struct{})
+
+	// Start a GetOrCreate whose create blocks until we signal it.
+	errc := make(chan error, 1)
+	go func() {
+		_, err := store.GetOrCreate(t.Context(), key, func(ctx context.Context) (Diff, error) {
+			<-ready // wait for Add to fire first
+			return nil, storage.ErrObjectNotExist
+		})
+		errc <- err
+	}()
+
+	// Give the goroutine a moment to enter initGroup.Do and block on <-ready.
+	time.Sleep(20 * time.Millisecond)
+
+	// Add publishes the diff into the positive cache (simulates a concurrent upload).
+	store.Add(diff)
+
+	// Unblock create — it returns ErrObjectNotExist, but Add already ran.
+	close(ready)
+
+	require.ErrorIs(t, <-errc, storage.ErrObjectNotExist)
+
+	// The diff added by Add must be reachable immediately via GetOrCreate.
+	got, err := store.GetOrCreate(t.Context(), key, func(ctx context.Context) (Diff, error) {
+		t.Fatal("create must not be called: diff is in the positive cache")
+		return nil, nil
+	})
+	require.NoError(t, err)
+	require.Equal(t, diff, got)
+}

packages/orchestrator/pkg/sandbox/nbd/dispatch.go

@@ -16,6 +16,7 @@ import (
 	"go.uber.org/zap"
 
 	"github.com/e2b-dev/infra/packages/shared/pkg/logger"
+	"github.com/e2b-dev/infra/packages/shared/pkg/storage"
 	"github.com/e2b-dev/infra/packages/shared/pkg/utils"
 )
 
@@ -333,14 +334,29 @@ func (d *Dispatch) cmdRead(ctx context.Context, cmdHandle uint64, cmdFrom uint64
 			// Per-request backend failure: signal it to the NBD client via the
 			// response error byte and keep the dispatch loop alive. Only
 			// writeResponse errors (dead NBD socket) escalate through d.fatal.
-			logger.L().Error(ctx, "nbd backend read failed",
-				zap.Error(readErr),
-				zap.String("nbd_op", "read"),
-				zap.String("nbd_provider", d.provName),
-				zap.Uint64("nbd_handle", handle),
-				zap.Uint64("nbd_offset", from),
-				zap.Uint32("nbd_length", length),
-			)
+			//
+			// ErrObjectNotExist is a permanent failure (the build is gone); the
+			// DiffStore negative cache absorbs subsequent retries. Log at WARN to
+			// avoid flooding dashboards with expected 404s while still surfacing them.
+			if errors.Is(readErr, storage.ErrObjectNotExist) {
+				logger.L().Warn(ctx, "nbd backend read failed: build object not found in store",
+					zap.Error(readErr),
+					zap.String("nbd_op", "read"),
+					zap.String("nbd_provider", d.provName),
+					zap.Uint64("nbd_handle", handle),
+					zap.Uint64("nbd_offset", from),
+					zap.Uint32("nbd_length", length),
+				)
+			} else {
+				logger.L().Error(ctx, "nbd backend read failed",
+					zap.Error(readErr),
+					zap.String("nbd_op", "read"),
+					zap.String("nbd_provider", d.provName),
+					zap.Uint64("nbd_handle", handle),
+					zap.Uint64("nbd_offset", from),
+					zap.Uint32("nbd_length", length),
+				)
+			}
 
 			return d.writeResponse(1, handle, []byte{})
 		}