Improve naming, add to sandbox table on the first insert

Author: jakubno

packages/api/internal/orchestrator/lifecycle.go

@@ -11,7 +11,7 @@ import (
 	e2bcatalog "github.com/e2b-dev/infra/packages/shared/pkg/sandbox-catalog"
 )
 
-func (o *Orchestrator) addToSandboxRoutingTable(ctx context.Context, sandbox sandbox.Sandbox) {
+func (o *Orchestrator) addSandboxToRoutingTable(ctx context.Context, sandbox sandbox.Sandbox) {
 	node := o.GetNode(sandbox.ClusterID, sandbox.NodeID)
 	if node == nil {
 		logger.L().Error(ctx, "failed to get node", logger.WithNodeID(sandbox.NodeID))

packages/api/internal/orchestrator/orchestrator.go

@@ -145,7 +145,7 @@ func New(
 		sandboxStorage,
 		reservationStorage,
 		sandbox.Callbacks{
-			AddToSandboxRoutingTable: o.addToSandboxRoutingTable,
+			AddSandboxToRoutingTable: o.addSandboxToRoutingTable,
 			AsyncSandboxCounter:      o.sandboxCounterInsert,
 			AsyncNewlyCreatedSandbox: o.handleNewlyCreatedSandbox,
 		},

packages/api/internal/sandbox/store.go

@@ -51,8 +51,8 @@ func WithOnlyExpired(isExpired bool) ItemsOption {
 }
 
 type Callbacks struct {
-	// AddToSandboxRoutingTable should be called sync to prevent race conditions where we would know where to route the sandbox
-	AddToSandboxRoutingTable InsertCallback
+	// AddSandboxToRoutingTable should be called sync to prevent race conditions where we would know where to route the sandbox
+	AddSandboxToRoutingTable InsertCallback
 	// AsyncSandboxCounter should be called async to prevent blocking the main goroutine
 	AsyncSandboxCounter InsertCallback
 	// AsyncNewlyCreatedSandbox should be called async to prevent blocking the main goroutine
@@ -94,6 +94,7 @@ func (s *Store) Add(ctx context.Context, sandbox Sandbox, newlyCreated bool) err
 	err := s.storage.Add(ctx, sandbox)
 	if err == nil {
 		// Count only newly added sandboxes to the store
+		s.callbacks.AddSandboxToRoutingTable(ctx, sandbox)
 		go s.callbacks.AsyncSandboxCounter(context.WithoutCancel(ctx), sandbox)
 	} else {
 		// There's a race condition when the sandbox is added from node sync
@@ -115,9 +116,6 @@ func (s *Store) Add(ctx context.Context, sandbox Sandbox, newlyCreated bool) err
 		finishStart(sandbox, nil)
 	}
 
-	// Add sandbox info to node for routing
-	s.callbacks.AddToSandboxRoutingTable(ctx, sandbox)
-
 	if newlyCreated {
 		go s.callbacks.AsyncNewlyCreatedSandbox(context.WithoutCancel(ctx), sandbox)
 	}

packages/api/internal/sandbox/store_test.go

@@ -189,7 +189,7 @@ func TestAdd_NewSandbox(t *testing.T) {
 
 		tracker := NewCallbackTracker(3) // Expect 3 callbacks
 		callbacks := sandbox.Callbacks{
-			AddToSandboxRoutingTable: tracker.Track("AddToSandboxRoutingTable"),
+			AddSandboxToRoutingTable: tracker.Track("AddSandboxToRoutingTable"),
 			AsyncSandboxCounter:      tracker.Track("AsyncSandboxCounter"),
 			AsyncNewlyCreatedSandbox: tracker.Track("AsyncNewlyCreatedSandbox"),
 		}
@@ -207,7 +207,7 @@ func TestAdd_NewSandbox(t *testing.T) {
 		require.NoError(t, err)
 
 		// Verify all callbacks called exactly once
-		tracker.AssertCallCount(t, "AddToSandboxRoutingTable", 1)
+		tracker.AssertCallCount(t, "AddSandboxToRoutingTable", 1)
 		tracker.AssertCallCount(t, "AsyncSandboxCounter", 1)
 		tracker.AssertCallCount(t, "AsyncNewlyCreatedSandbox", 1)
 
@@ -228,7 +228,7 @@ func TestAdd_AlreadyInCache(t *testing.T) {
 		// First add with all 3 callbacks
 		tracker1 := NewCallbackTracker(3)
 		callbacks1 := sandbox.Callbacks{
-			AddToSandboxRoutingTable: tracker1.Track("AddToSandboxRoutingTable"),
+			AddSandboxToRoutingTable: tracker1.Track("AddSandboxToRoutingTable"),
 			AsyncSandboxCounter:      tracker1.Track("AsyncSandboxCounter"),
 			AsyncNewlyCreatedSandbox: tracker1.Track("AsyncNewlyCreatedSandbox"),
 		}
@@ -241,9 +241,9 @@ func TestAdd_AlreadyInCache(t *testing.T) {
 
 		// Second add with newlyCreated=true, only 2 callbacks
 		// (AsyncSandboxCounter is NOT called because already in cache)
-		tracker2 := NewCallbackTracker(2)
+		tracker2 := NewCallbackTracker(1)
 		callbacks2 := sandbox.Callbacks{
-			AddToSandboxRoutingTable: tracker2.Track("AddToSandboxRoutingTable"),
+			AddSandboxToRoutingTable: tracker2.Track("AddSandboxToRoutingTable"),
 			AsyncSandboxCounter:      tracker2.Track("AsyncSandboxCounter"),
 			AsyncNewlyCreatedSandbox: tracker2.Track("AsyncNewlyCreatedSandbox"),
 		}
@@ -253,12 +253,12 @@ func TestAdd_AlreadyInCache(t *testing.T) {
 		tracker2.WaitForCalls(t, 2*time.Second)
 
 		require.NoError(t, err)
-		tracker2.AssertCallCount(t, "AddToSandboxRoutingTable", 1)
+		tracker2.AssertNotCalled(t, "AddSandboxToRoutingTable")
 		tracker2.AssertNotCalled(t, "AsyncSandboxCounter") // NOT called when already in cache!
 		tracker2.AssertCallCount(t, "AsyncNewlyCreatedSandbox", 1)
 	})
 
-	t.Run("newlyCreated=false - only AddToSandboxRoutingTable called when already in cache", func(t *testing.T) {
+	t.Run("newlyCreated=false - only AddSandboxToRoutingTable called when already in cache", func(t *testing.T) {
 		ctx := t.Context()
 
 		storage := memory.NewStorage()
@@ -267,7 +267,7 @@ func TestAdd_AlreadyInCache(t *testing.T) {
 		// First add with newlyCreated=true
 		tracker1 := NewCallbackTracker(3)
 		callbacks1 := sandbox.Callbacks{
-			AddToSandboxRoutingTable: tracker1.Track("AddToSandboxRoutingTable"),
+			AddSandboxToRoutingTable: tracker1.Track("AddSandboxToRoutingTable"),
 			AsyncSandboxCounter:      tracker1.Track("AsyncSandboxCounter"),
 			AsyncNewlyCreatedSandbox: tracker1.Track("AsyncNewlyCreatedSandbox"),
 		}
@@ -280,9 +280,9 @@ func TestAdd_AlreadyInCache(t *testing.T) {
 
 		// Second add with newlyCreated=false, only 1 callback
 		// (AsyncSandboxCounter NOT called because already in cache)
-		tracker2 := NewCallbackTracker(1)
+		tracker2 := NewCallbackTracker(0)
 		callbacks2 := sandbox.Callbacks{
-			AddToSandboxRoutingTable: tracker2.Track("AddToSandboxRoutingTable"),
+			AddSandboxToRoutingTable: tracker2.Track("AddSandboxToRoutingTable"),
 			AsyncSandboxCounter:      tracker2.Track("AsyncSandboxCounter"),
 			AsyncNewlyCreatedSandbox: tracker2.Track("AsyncNewlyCreatedSandbox"),
 		}
@@ -292,14 +292,14 @@ func TestAdd_AlreadyInCache(t *testing.T) {
 		tracker2.WaitForCalls(t, 2*time.Second)
 
 		require.NoError(t, err)
-		tracker2.AssertCallCount(t, "AddToSandboxRoutingTable", 1)
+		tracker2.AssertNotCalled(t, "AddSandboxToRoutingTable")
 		tracker2.AssertNotCalled(t, "AsyncSandboxCounter") // NOT called when already in cache
 		tracker2.AssertNotCalled(t, "AsyncNewlyCreatedSandbox")
 	})
 }
 
 func TestAdd_NotNewlyCreated(t *testing.T) {
-	t.Run("not in cache - AddToSandboxRoutingTable and AsyncSandboxCounter called", func(t *testing.T) {
+	t.Run("not in cache - AddSandboxToRoutingTable and AsyncSandboxCounter called", func(t *testing.T) {
 		ctx := t.Context()
 
 		storage := memory.NewStorage()
@@ -308,7 +308,7 @@ func TestAdd_NotNewlyCreated(t *testing.T) {
 		// Add with newlyCreated=false, expect 2 callbacks
 		tracker := NewCallbackTracker(2)
 		callbacks := sandbox.Callbacks{
-			AddToSandboxRoutingTable: tracker.Track("AddToSandboxRoutingTable"),
+			AddSandboxToRoutingTable: tracker.Track("AddSandboxToRoutingTable"),
 			AsyncSandboxCounter:      tracker.Track("AsyncSandboxCounter"),
 			AsyncNewlyCreatedSandbox: tracker.Track("AsyncNewlyCreatedSandbox"),
 		}
@@ -319,12 +319,12 @@ func TestAdd_NotNewlyCreated(t *testing.T) {
 		tracker.WaitForCalls(t, 2*time.Second)
 
 		require.NoError(t, err)
-		tracker.AssertCallCount(t, "AddToSandboxRoutingTable", 1)
+		tracker.AssertCallCount(t, "AddSandboxToRoutingTable", 1)
 		tracker.AssertCallCount(t, "AsyncSandboxCounter", 1)
 		tracker.AssertNotCalled(t, "AsyncNewlyCreatedSandbox")
 	})
 
-	t.Run("already in cache - only AddToSandboxRoutingTable called", func(t *testing.T) {
+	t.Run("already in cache - only AddSandboxToRoutingTable called", func(t *testing.T) {
 		ctx := t.Context()
 
 		storage := memory.NewStorage()
@@ -333,7 +333,7 @@ func TestAdd_NotNewlyCreated(t *testing.T) {
 		// First add
 		tracker1 := NewCallbackTracker(2)
 		callbacks1 := sandbox.Callbacks{
-			AddToSandboxRoutingTable: tracker1.Track("AddToSandboxRoutingTable"),
+			AddSandboxToRoutingTable: tracker1.Track("AddSandboxToRoutingTable"),
 			AsyncSandboxCounter:      tracker1.Track("AsyncSandboxCounter"),
 			AsyncNewlyCreatedSandbox: tracker1.Track("AsyncNewlyCreatedSandbox"),
 		}
@@ -345,9 +345,9 @@ func TestAdd_NotNewlyCreated(t *testing.T) {
 		require.NoError(t, err)
 
 		// Second add with same sandbox, newlyCreated=false, only 1 callback
-		tracker2 := NewCallbackTracker(1)
+		tracker2 := NewCallbackTracker(0)
 		callbacks2 := sandbox.Callbacks{
-			AddToSandboxRoutingTable: tracker2.Track("AddToSandboxRoutingTable"),
+			AddSandboxToRoutingTable: tracker2.Track("AddSandboxToRoutingTable"),
 			AsyncSandboxCounter:      tracker2.Track("AsyncSandboxCounter"),
 			AsyncNewlyCreatedSandbox: tracker2.Track("AsyncNewlyCreatedSandbox"),
 		}
@@ -357,7 +357,7 @@ func TestAdd_NotNewlyCreated(t *testing.T) {
 		tracker2.WaitForCalls(t, 2*time.Second)
 
 		require.NoError(t, err)
-		tracker2.AssertCallCount(t, "AddToSandboxRoutingTable", 1)
+		tracker2.AssertNotCalled(t, "AddSandboxToRoutingTable")
 		tracker2.AssertNotCalled(t, "AsyncSandboxCounter") // NOT called when already in cache
 		tracker2.AssertNotCalled(t, "AsyncNewlyCreatedSandbox")
 	})
@@ -377,7 +377,7 @@ func TestAdd_StorageErrors(t *testing.T) {
 		// Expect 0 callbacks since error should be returned immediately
 		tracker := NewCallbackTracker(1)
 		callbacks := sandbox.Callbacks{
-			AddToSandboxRoutingTable: tracker.Track("AddToSandboxRoutingTable"),
+			AddSandboxToRoutingTable: tracker.Track("AddSandboxToRoutingTable"),
 			AsyncSandboxCounter:      tracker.Track("AsyncSandboxCounter"),
 			AsyncNewlyCreatedSandbox: tracker.Track("AsyncNewlyCreatedSandbox"),
 		}
@@ -394,7 +394,7 @@ func TestAdd_StorageErrors(t *testing.T) {
 		time.Sleep(100 * time.Millisecond)
 
 		// No callbacks should have been called
-		tracker.AssertNotCalled(t, "AddToSandboxRoutingTable")
+		tracker.AssertNotCalled(t, "AddSandboxToRoutingTable")
 		tracker.AssertNotCalled(t, "AsyncSandboxCounter")
 		tracker.AssertNotCalled(t, "AsyncNewlyCreatedSandbox")
 	})
@@ -411,7 +411,7 @@ func TestAdd_ConcurrentCalls(t *testing.T) {
 		tracker := NewCallbackTracker(numGoroutines * 3) // Each add calls 3 callbacks
 
 		callbacks := sandbox.Callbacks{
-			AddToSandboxRoutingTable: tracker.Track("AddToSandboxRoutingTable"),
+			AddSandboxToRoutingTable: tracker.Track("AddSandboxToRoutingTable"),
 			AsyncSandboxCounter:      tracker.Track("AsyncSandboxCounter"),
 			AsyncNewlyCreatedSandbox: tracker.Track("AsyncNewlyCreatedSandbox"),
 		}
@@ -448,7 +448,7 @@ func TestAdd_ConcurrentCalls(t *testing.T) {
 		tracker.WaitForCalls(t, 5*time.Second)
 
 		// Verify all callbacks were called expected number of times
-		tracker.AssertCallCount(t, "AddToSandboxRoutingTable", numGoroutines)
+		tracker.AssertCallCount(t, "AddSandboxToRoutingTable", numGoroutines)
 		tracker.AssertCallCount(t, "AsyncSandboxCounter", numGoroutines)
 		tracker.AssertCallCount(t, "AsyncNewlyCreatedSandbox", numGoroutines)
 
@@ -475,7 +475,7 @@ func TestAdd_ConcurrentCalls(t *testing.T) {
 		tracker := NewCallbackTracker(3 + (numGoroutines-1)*2)
 
 		callbacks := sandbox.Callbacks{
-			AddToSandboxRoutingTable: tracker.Track("AddToSandboxRoutingTable"),
+			AddSandboxToRoutingTable: tracker.Track("AddSandboxToRoutingTable"),
 			AsyncSandboxCounter:      tracker.Track("AsyncSandboxCounter"),
 			AsyncNewlyCreatedSandbox: tracker.Track("AsyncNewlyCreatedSandbox"),
 		}
@@ -503,7 +503,7 @@ func TestAdd_ConcurrentCalls(t *testing.T) {
 		tracker.WaitForCalls(t, 5*time.Second)
 
 		// Verify callbacks
-		tracker.AssertCallCount(t, "AddToSandboxRoutingTable", numGoroutines)
+		tracker.AssertCallCount(t, "AddSandboxToRoutingTable", numGoroutines)
 		tracker.AssertCallCount(t, "AsyncSandboxCounter", 1)                  // Only called once (first successful add)
 		tracker.AssertCallCount(t, "AsyncNewlyCreatedSandbox", numGoroutines) // All calls have newlyCreated=true