diff --git a/packages/api/internal/orchestrator/orchestrator.go b/packages/api/internal/orchestrator/orchestrator.go index 74e6479049..934f228d5f 100644 --- a/packages/api/internal/orchestrator/orchestrator.go +++ b/packages/api/internal/orchestrator/orchestrator.go @@ -167,7 +167,7 @@ func New( go redisbackend.NewCleaner(redisStorage).Start(ctx) case cfg.SandboxStorageBackendRedis: - reservationStorage = redisreservations.NewReservationStorage(redisClient) + reservationStorage = redisreservations.NewReservationStorage(redisClient, redisStorage.Notifier()) sandboxStorage = redisStorage logger.L().Info(ctx, "Using redis sandbox storage backend") default: diff --git a/packages/api/internal/sandbox/reservations/redis/README.md b/packages/api/internal/sandbox/reservations/redis/README.md new file mode 100644 index 0000000000..67c48b0f72 --- /dev/null +++ b/packages/api/internal/sandbox/reservations/redis/README.md @@ -0,0 +1,20 @@ +# Redis Reservation Storage + +This package coordinates sandbox creation reservations across API instances. + +## Keys + +- Storage index: `sandbox:storage:{teamID}:index` +- Pending zset: `sandbox:storage:{teamID}:reservations:pending` +- Result key: `sandbox:storage:{teamID}:reservations:{sandboxID}:result` +- PubSub routing key: `sandbox:storage:{teamID}:reservations:{sandboxID}:notify` + +## Flow + +`Reserve` runs a Lua script that atomically removes stale pending entries, checks whether the sandbox already exists or has already pending start, enforces the team limit using `SCARD(storage index) + ZCARD(pending zset)`, deletes any stale result key, and adds the sandbox ID to the pending zset. + +When creation completes, it removes the sandbox from the pending zset, writes a TTL result key containing either the sandbox or the creation error, and publishes the routing key. + +A waiter subscribes to the routing key, probes the result key immediately, then waits for PubSub notifications or the 1 second fallback ticker. PubSub is best-effort; the fallback ticker is required for correctness. + +`Release` is called when the sandbox is removed from storage (`Store.Remove`). It removes the sandbox from the pending zset, deletes the result key, and publishes the routing key. diff --git a/packages/api/internal/sandbox/reservations/redis/reservation.go b/packages/api/internal/sandbox/reservations/redis/reservation.go index 29068ecd7c..922057cb85 100644 --- a/packages/api/internal/sandbox/reservations/redis/reservation.go +++ b/packages/api/internal/sandbox/reservations/redis/reservation.go @@ -15,8 +15,11 @@ import ( ) const ( - resultTTL = 30 * time.Second - retryInterval = 20 * time.Millisecond + resultTTL = 30 * time.Second + + // fallbackPollInterval is how often the waiter re-checks Redis when no + // PubSub wakeup arrives + fallbackPollInterval = 1 * time.Second // staleTTL is the maximum age of a pending entry before it is considered stale // and cleaned up. This handles the case where an API instance crashes mid-creation. @@ -26,13 +29,21 @@ const ( var _ sandbox.ReservationStorage = (*ReservationStorage)(nil) +// Publish is fire-and-forget — drops on queue saturation are recovered by the waiter's fallback ticker. +type Notifier interface { + Subscribe(routingKey string) (<-chan struct{}, func()) + Publish(ctx context.Context, routingKey string) +} + type ReservationStorage struct { redisClient redis.UniversalClient + notifier Notifier } -func NewReservationStorage(redisClient redis.UniversalClient) *ReservationStorage { +func NewReservationStorage(redisClient redis.UniversalClient, notifier Notifier) *ReservationStorage { return &ReservationStorage{ redisClient: redisClient, + notifier: notifier, } } @@ -76,21 +87,29 @@ func (s *ReservationStorage) Release(ctx context.Context, teamID uuid.UUID, sand pendingSetKey := getPendingSetKey(teamIDStr) resultKeyStr := getResultKey(teamIDStr, sandboxID) - err := releaseScript.Run(ctx, s.redisClient, []string{pendingSetKey, resultKeyStr}, sandboxID).Err() + err := releaseScript.Run(ctx, s.redisClient, + []string{pendingSetKey, resultKeyStr}, + sandboxID, + ).Err() if err != nil { return fmt.Errorf("failed to run release script: %w", err) } + // Wake any in-process waiter so it checks the pending set immediately. + s.notifier.Publish(ctx, getReservationRoutingKey(teamIDStr, sandboxID)) + return nil } // createFinishStart returns a callback that completes the reservation. -// It removes the sandbox from the pending zset and stores the result for cross-instance waiters. func (s *ReservationStorage) createFinishStart(ctx context.Context, teamID uuid.UUID, sandboxID string) func(sandbox.Sandbox, error) { return func(sbx sandbox.Sandbox, startErr error) { teamIDStr := teamID.String() pendingSetKey := getPendingSetKey(teamIDStr) resultKeyStr := getResultKey(teamIDStr, sandboxID) + routingKey := getReservationRoutingKey(teamIDStr, sandboxID) + + bgCtx := context.WithoutCancel(ctx) resultData, encodeErr := encodeResult(sbx, startErr) if encodeErr != nil { @@ -99,14 +118,17 @@ func (s *ReservationStorage) createFinishStart(ctx context.Context, teamID uuid. logger.WithSandboxID(sandboxID), ) - // Still try to remove from pending even if encoding fails - _ = s.redisClient.ZRem(context.WithoutCancel(ctx), pendingSetKey, sandboxID).Err() + // Still try to remove from pending even if encoding fails. + _ = s.redisClient.ZRem(bgCtx, pendingSetKey, sandboxID).Err() + + // Wake waiters so they can observe that the reservation is gone. + s.notifier.Publish(bgCtx, routingKey) return } ttlSeconds := int(resultTTL.Seconds()) - err := finishStartScript.Run(context.WithoutCancel(ctx), s.redisClient, + err := finishStartScript.Run(bgCtx, s.redisClient, []string{pendingSetKey, resultKeyStr}, sandboxID, resultData, ttlSeconds, ).Err() @@ -115,7 +137,13 @@ func (s *ReservationStorage) createFinishStart(ctx context.Context, teamID uuid. zap.Error(err), logger.WithSandboxID(sandboxID), ) + + return } + + // Wake any in-process waiter immediately. Drop-tolerant: the + // fallback ticker covers a saturated publish queue. + s.notifier.Publish(bgCtx, routingKey) } } @@ -126,39 +154,79 @@ func (s *ReservationStorage) createWaitForStart(teamID uuid.UUID, sandboxID stri teamIDStr := teamID.String() resultKeyStr := getResultKey(teamIDStr, sandboxID) pendingSetKey := getPendingSetKey(teamIDStr) + routingKey := getReservationRoutingKey(teamIDStr, sandboxID) - for { - // Check for result - data, err := s.redisClient.Get(ctx, resultKeyStr).Bytes() - if err == nil { - return decodeResult(data) - } - if !errors.Is(err, redis.Nil) { - return sandbox.Sandbox{}, fmt.Errorf("failed to check result key: %w", err) - } + ch, cleanup := s.notifier.Subscribe(routingKey) + defer cleanup() - // No result yet — check if still pending (ZSCORE returns nil if not a member) - err = s.redisClient.ZScore(ctx, pendingSetKey, sandboxID).Err() - if errors.Is(err, redis.Nil) { - // Not pending anymore, final check - data, err = s.redisClient.Get(ctx, resultKeyStr).Bytes() - if err == nil { - return decodeResult(data) - } + // Initial probe: the producer may have finished before we subscribed, + // or we may be a late waiter joining after the result was already set. + if done, sbx, err := s.tryReadResult(ctx, resultKeyStr, pendingSetKey, sandboxID); done { + return sbx, err + } - return sandbox.Sandbox{}, fmt.Errorf("sandbox %s is no longer pending and has no result", sandboxID) - } - if err != nil { - return sandbox.Sandbox{}, fmt.Errorf("failed to check pending set: %w", err) - } + ticker := time.NewTicker(fallbackPollInterval) + defer ticker.Stop() - // Wait before next poll + for { select { case <-ctx.Done(): return sandbox.Sandbox{}, ctx.Err() - case <-time.After(retryInterval): - // continue polling + case <-ch: + case <-ticker.C: + } + + if done, sbx, err := s.tryReadResult(ctx, resultKeyStr, pendingSetKey, sandboxID); done { + return sbx, err } } } } + +// tryReadResult performs a single probe of the reservation state. +// +// Returns done=true when the wait is over: +// - the result key holds an encoded terminal result, or +// - the sandbox vanished from the pending set without a result, or +// - the Redis call itself failed. +// +// Returns done=false when the reservation is still pending and the caller +// should wait for the next wakeup. +func (s *ReservationStorage) tryReadResult( + ctx context.Context, + resultKey, pendingSetKey, sandboxID string, +) (done bool, sbx sandbox.Sandbox, err error) { + data, getErr := s.redisClient.Get(ctx, resultKey).Bytes() + if getErr == nil { + sbx, err = decodeResult(data) + + return true, sbx, err + } + if !errors.Is(getErr, redis.Nil) { + return true, sandbox.Sandbox{}, fmt.Errorf("failed to check result key: %w", getErr) + } + + // No result yet, so check whether another instance is still creating the sandbox. + scoreErr := s.redisClient.ZScore(ctx, pendingSetKey, sandboxID).Err() + if errors.Is(scoreErr, redis.Nil) { + // Re-read the result in case finishStart or a new Release wrote it + // between the initial GET and the legacy pending-set check. + data, getErr = s.redisClient.Get(ctx, resultKey).Bytes() + if getErr == nil { + sbx, err = decodeResult(data) + + return true, sbx, err + } + if !errors.Is(getErr, redis.Nil) { + return true, sandbox.Sandbox{}, fmt.Errorf("failed to check result key: %w", getErr) + } + + return true, sandbox.Sandbox{}, fmt.Errorf("sandbox %s is no longer pending and has no result", sandboxID) + } + if scoreErr != nil { + return true, sandbox.Sandbox{}, fmt.Errorf("failed to check pending set: %w", scoreErr) + } + + // Still pending, no result yet. + return false, sandbox.Sandbox{}, nil +} diff --git a/packages/api/internal/sandbox/reservations/redis/reservation_pubsub_test.go b/packages/api/internal/sandbox/reservations/redis/reservation_pubsub_test.go new file mode 100644 index 0000000000..d93322b1b3 --- /dev/null +++ b/packages/api/internal/sandbox/reservations/redis/reservation_pubsub_test.go @@ -0,0 +1,435 @@ +package redis + +import ( + "context" + "errors" + "sync" + "testing" + "time" + + "github.com/google/uuid" + goredis "github.com/redis/go-redis/v9" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/e2b-dev/infra/packages/api/internal/sandbox" + storage_redis "github.com/e2b-dev/infra/packages/api/internal/sandbox/storage/redis" + "github.com/e2b-dev/infra/packages/shared/pkg/consts" + redis_utils "github.com/e2b-dev/infra/packages/shared/pkg/redis" +) + +// testSandbox is the canonical successful-finish payload used across pubsub tests. +func testSandbox(teamID uuid.UUID, sandboxID string) sandbox.Sandbox { + return sandbox.Sandbox{ + ClientID: consts.ClientID, + SandboxID: sandboxID, + TemplateID: "test", + TeamID: teamID, + StartTime: time.Now(), + EndTime: time.Now().Add(time.Hour), + MaxInstanceLength: time.Hour, + } +} + +// setupReservationStorageWithoutSubManager wires the reservation store so that +// PubSub messages are never delivered in-process. Used to exercise the safety-net path. +func setupReservationStorageWithoutSubManager(t *testing.T) (*ReservationStorage, goredis.UniversalClient) { + t.Helper() + + client := redis_utils.SetupInstance(t) + + storageInstance := newTestSandboxStorage(t, client) + t.Cleanup(func() { storageInstance.Close(context.WithoutCancel(t.Context())) }) + + storage := NewReservationStorage(client, storageInstance.Notifier()) + + return storage, client +} + +func TestWaitForStart_WokenByFinishStartPublish(t *testing.T) { + t.Parallel() + + storage, _ := setupTestReservationStorage(t) + teamID := uuid.New() + sbxID := "pubsub-finish" + + finishStart, _, err := storage.Reserve(t.Context(), teamID, sbxID, 10) + require.NoError(t, err) + require.NotNil(t, finishStart) + + _, waitForStart, err := storage.Reserve(t.Context(), teamID, sbxID, 10) + require.NoError(t, err) + require.NotNil(t, waitForStart) + + waiterDone := make(chan struct{}) + var got sandbox.Sandbox + var waitErr error + go func() { + got, waitErr = waitForStart(t.Context()) + close(waiterDone) + }() + + // Let the waiter subscribe before we finish. + time.Sleep(50 * time.Millisecond) + + start := time.Now() + finishStart(testSandbox(teamID, sbxID), nil) + + select { + case <-waiterDone: + elapsed := time.Since(start) + require.NoError(t, waitErr) + assert.Equal(t, sbxID, got.SandboxID) + assert.Less(t, elapsed, 500*time.Millisecond, + "waiter should wake via PubSub, not the fallback ticker") + case <-time.After(3 * time.Second): + require.FailNow(t, "waiter did not wake in time") + } +} + +func TestWaitForStart_WokenByReleasePublish(t *testing.T) { + t.Parallel() + + storage, _ := setupTestReservationStorage(t) + teamID := uuid.New() + sbxID := "pubsub-release" + + finishStart, _, err := storage.Reserve(t.Context(), teamID, sbxID, 10) + require.NoError(t, err) + require.NotNil(t, finishStart) + + _, waitForStart, err := storage.Reserve(t.Context(), teamID, sbxID, 10) + require.NoError(t, err) + require.NotNil(t, waitForStart) + + waiterErr := make(chan error, 1) + go func() { + _, err := waitForStart(t.Context()) + waiterErr <- err + }() + time.Sleep(50 * time.Millisecond) + + start := time.Now() + require.NoError(t, storage.Release(t.Context(), teamID, sbxID)) + + select { + case err := <-waiterErr: + require.Error(t, err) + assert.Contains(t, err.Error(), "no longer pending") + assert.Less(t, time.Since(start), 500*time.Millisecond, + "release should wake the waiter via PubSub") + case <-time.After(3 * time.Second): + require.FailNow(t, "waiter did not wake in time") + } +} + +func TestWaitForStart_FallbackTickerWhenPubSubMissed(t *testing.T) { + t.Parallel() + + storage, _ := setupReservationStorageWithoutSubManager(t) + teamID := uuid.New() + sbxID := "pubsub-fallback" + + finishStart, _, err := storage.Reserve(t.Context(), teamID, sbxID, 10) + require.NoError(t, err) + require.NotNil(t, finishStart) + + _, waitForStart, err := storage.Reserve(t.Context(), teamID, sbxID, 10) + require.NoError(t, err) + + // Set the result directly via the finishStartScript path; with the + // subManager not running, no fan-out occurs. The waiter must rely on + // the fallback ticker. + resultData, err := encodeResult(testSandbox(teamID, sbxID), nil) + require.NoError(t, err) + teamIDStr := teamID.String() + err = finishStartScript.Run(t.Context(), storage.redisClient, + []string{getPendingSetKey(teamIDStr), getResultKey(teamIDStr, sbxID)}, + sbxID, resultData, int(resultTTL.Seconds()), + ).Err() + require.NoError(t, err) + + done := make(chan error, 1) + go func() { + _, err := waitForStart(t.Context()) + done <- err + }() + + select { + case err := <-done: + require.NoError(t, err) + case <-time.After(fallbackPollInterval + 2*time.Second): + require.FailNow(t, "fallback ticker did not resolve the wait") + } +} + +// TestWaitForStart_ResultLandedBeforeSubscribe covers the race where the +// producer finishes BEFORE the waiter calls waitForStart. The initial +// post-subscribe GET must catch the already-present result rather than +// blocking until a wakeup that will never come. +func TestWaitForStart_ResultLandedBeforeSubscribe(t *testing.T) { + t.Parallel() + + storage, _ := setupTestReservationStorage(t) + teamID := uuid.New() + sbxID := "pubsub-race" + + finishStart, _, err := storage.Reserve(t.Context(), teamID, sbxID, 10) + require.NoError(t, err) + + _, waitForStart, err := storage.Reserve(t.Context(), teamID, sbxID, 10) + require.NoError(t, err) + + // Finish BEFORE the waiter ever subscribes. + finishStart(testSandbox(teamID, sbxID), nil) + // Give the in-process publisher worker a moment to drain. + time.Sleep(50 * time.Millisecond) + + done := make(chan struct{}) + var got sandbox.Sandbox + var waitErr error + start := time.Now() + go func() { + got, waitErr = waitForStart(t.Context()) + close(done) + }() + + select { + case <-done: + require.NoError(t, waitErr) + assert.Equal(t, sbxID, got.SandboxID) + assert.Less(t, time.Since(start), 200*time.Millisecond, + "initial post-subscribe check must catch a pre-existing result") + case <-time.After(3 * time.Second): + require.FailNow(t, "waiter did not return") + } +} + +// TestWaitForStart_ContextCancellation asserts ctx.Done is the only path +// out under normal load; cancellation must unblock immediately. +func TestWaitForStart_ContextCancellation(t *testing.T) { + t.Parallel() + + storage, _ := setupTestReservationStorage(t) + teamID := uuid.New() + sbxID := "pubsub-cancel" + + _, _, err := storage.Reserve(t.Context(), teamID, sbxID, 10) + require.NoError(t, err) + + _, waitForStart, err := storage.Reserve(t.Context(), teamID, sbxID, 10) + require.NoError(t, err) + + ctx, cancel := context.WithTimeout(t.Context(), 100*time.Millisecond) + defer cancel() + + start := time.Now() + _, err = waitForStart(ctx) + elapsed := time.Since(start) + + require.ErrorIs(t, err, context.DeadlineExceeded) + assert.Less(t, elapsed, 500*time.Millisecond, "cancellation should be immediate") +} + +// TestWaitForStart_MultipleWaitersOnePublish confirms the fan-out wakes +// every concurrent waiter from a single producer publish. +func TestWaitForStart_MultipleWaitersOnePublish(t *testing.T) { + t.Parallel() + + storage, _ := setupTestReservationStorage(t) + teamID := uuid.New() + sbxID := "pubsub-multi" + const numWaiters = 10 + + finishStart, _, err := storage.Reserve(t.Context(), teamID, sbxID, 50) + require.NoError(t, err) + + waiters := make([]func(ctx context.Context) (sandbox.Sandbox, error), numWaiters) + for i := range numWaiters { + _, w, err := storage.Reserve(t.Context(), teamID, sbxID, 50) + require.NoError(t, err) + require.NotNil(t, w) + waiters[i] = w + } + + var wg sync.WaitGroup + errs := make([]error, numWaiters) + completions := make([]time.Duration, numWaiters) + for i, w := range waiters { + wg.Add(1) + go func(i int, w func(ctx context.Context) (sandbox.Sandbox, error)) { + defer wg.Done() + start := time.Now() + _, errs[i] = w(t.Context()) + completions[i] = time.Since(start) + }(i, w) + } + + // Let everyone subscribe. + time.Sleep(100 * time.Millisecond) + finishStart(testSandbox(teamID, sbxID), nil) + + done := make(chan struct{}) + go func() { wg.Wait(); close(done) }() + + select { + case <-done: + for i := range numWaiters { + require.NoError(t, errs[i]) + assert.Less(t, completions[i], 600*time.Millisecond, + "waiter %d should be woken by PubSub", i) + } + case <-time.After(3 * time.Second): + require.FailNow(t, "not all waiters completed") + } +} + +// TestWaitForStart_FailedStartPropagatesPromptly proves a producer-side +// failure round-trips back through the result key and wakes the waiter +// via PubSub. +func TestWaitForStart_FailedStartPropagatesPromptly(t *testing.T) { + t.Parallel() + + storage, _ := setupTestReservationStorage(t) + teamID := uuid.New() + sbxID := "pubsub-failed" + + finishStart, _, err := storage.Reserve(t.Context(), teamID, sbxID, 10) + require.NoError(t, err) + + _, waitForStart, err := storage.Reserve(t.Context(), teamID, sbxID, 10) + require.NoError(t, err) + + waiterErr := make(chan error, 1) + go func() { + _, err := waitForStart(t.Context()) + waiterErr <- err + }() + time.Sleep(50 * time.Millisecond) + + start := time.Now() + finishStart(sandbox.Sandbox{}, errors.New("boom")) + + select { + case err := <-waiterErr: + require.Error(t, err) + assert.Contains(t, err.Error(), "boom") + assert.Less(t, time.Since(start), 500*time.Millisecond) + case <-time.After(3 * time.Second): + require.FailNow(t, "waiter did not return") + } +} + +func TestWaitForStart_RedisOpsBounded(t *testing.T) { + t.Parallel() + + client := redis_utils.SetupInstance(t) + + counter := &cmdCounter{} + client.AddHook(counter) + + storageInstance := newTestSandboxStorage(t, client) + go storageInstance.Start(t.Context()) + t.Cleanup(func() { storageInstance.Close(context.WithoutCancel(t.Context())) }) + + storage := NewReservationStorage(client, storageInstance.Notifier()) + + teamID := uuid.New() + sbxID := "pubsub-bounded" + + finishStart, _, err := storage.Reserve(t.Context(), teamID, sbxID, 10) + require.NoError(t, err) + + _, waitForStart, err := storage.Reserve(t.Context(), teamID, sbxID, 10) + require.NoError(t, err) + + waiterDone := make(chan error, 1) + go func() { + _, err := waitForStart(t.Context()) + waiterDone <- err + }() + + time.Sleep(50 * time.Millisecond) + counter.Reset() // start counting from after the waiter has subscribed + + time.Sleep(1500 * time.Millisecond) + finishStart(testSandbox(teamID, sbxID), nil) + + select { + case err := <-waiterDone: + require.NoError(t, err) + case <-time.After(3 * time.Second): + require.FailNow(t, "waiter did not return") + } + + // Allow at most a handful of read ops: initial GET (+ ZSCORE legacy + // safety net), maybe one fallback tick if it fired just before the + // publish landed, and the post-wakeup GET. 10 is a generous ceiling + // for the new design; the old design would blow past 50. + got := counter.Reads() + assert.LessOrEqual(t, got, 10, + "expected bounded reads; got %d (regression: polling is back)", got) +} + +// cmdCounter is a redis.Hook that counts read-side operations on the +// client. We only care about the waiter's reads, not the producer's +// writes, so we only count GET and ZSCORE. +type cmdCounter struct { + mu sync.Mutex + get int + zs int +} + +func (c *cmdCounter) Reset() { + c.mu.Lock() + defer c.mu.Unlock() + c.get = 0 + c.zs = 0 +} + +func (c *cmdCounter) Reads() int { + c.mu.Lock() + defer c.mu.Unlock() + + return c.get + c.zs +} + +func (c *cmdCounter) DialHook(next goredis.DialHook) goredis.DialHook { + return next +} + +func (c *cmdCounter) ProcessHook(next goredis.ProcessHook) goredis.ProcessHook { + return func(ctx context.Context, cmd goredis.Cmder) error { + c.mu.Lock() + switch cmd.Name() { + case "get": + c.get++ + case "zscore": + c.zs++ + } + c.mu.Unlock() + + return next(ctx, cmd) + } +} + +func (c *cmdCounter) ProcessPipelineHook(next goredis.ProcessPipelineHook) goredis.ProcessPipelineHook { + return func(ctx context.Context, cmds []goredis.Cmder) error { + c.mu.Lock() + for _, cmd := range cmds { + switch cmd.Name() { + case "get": + c.get++ + case "zscore": + c.zs++ + } + } + c.mu.Unlock() + + return next(ctx, cmds) + } +} + +// Compile-time assertion that the storage Notifier satisfies the local +// Notifier seam. If this breaks, the orchestrator wiring needs to change too. +var _ Notifier = (*storage_redis.Notifier)(nil) diff --git a/packages/api/internal/sandbox/reservations/redis/reservation_test.go b/packages/api/internal/sandbox/reservations/redis/reservation_test.go index 27c3c007cb..bb252f054b 100644 --- a/packages/api/internal/sandbox/reservations/redis/reservation_test.go +++ b/packages/api/internal/sandbox/reservations/redis/reservation_test.go @@ -13,9 +13,11 @@ import ( goredis "github.com/redis/go-redis/v9" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "go.opentelemetry.io/otel/metric/noop" "golang.org/x/sync/errgroup" "github.com/e2b-dev/infra/packages/api/internal/sandbox" + storage_redis "github.com/e2b-dev/infra/packages/api/internal/sandbox/storage/redis" "github.com/e2b-dev/infra/packages/shared/pkg/consts" redis_utils "github.com/e2b-dev/infra/packages/shared/pkg/redis" ) @@ -29,11 +31,25 @@ var testTeamID = uuid.New() func setupTestReservationStorage(t *testing.T) (*ReservationStorage, goredis.UniversalClient) { t.Helper() client := redis_utils.SetupInstance(t) - storage := NewReservationStorage(client) + + storageInstance := newTestSandboxStorage(t, client) + go storageInstance.Start(t.Context()) + t.Cleanup(func() { storageInstance.Close(context.WithoutCancel(t.Context())) }) + + storage := NewReservationStorage(client, storageInstance.Notifier()) return storage, client } +func newTestSandboxStorage(t *testing.T, client goredis.UniversalClient) *storage_redis.Storage { + t.Helper() + + storageInstance, err := storage_redis.NewStorage(client, noop.NewMeterProvider()) + require.NoError(t, err) + + return storageInstance +} + func TestReservation(t *testing.T) { t.Parallel() storage, _ := setupTestReservationStorage(t) @@ -496,7 +512,10 @@ func TestReservation_StalePendingCleanup(t *testing.T) { assert.Equal(t, int64(1), count) // Create a new storage instance (simulating a fresh/restarted API) - storage := NewReservationStorage(client) + storageInstance := newTestSandboxStorage(t, client) + go storageInstance.Start(t.Context()) + t.Cleanup(func() { storageInstance.Close(context.WithoutCancel(t.Context())) }) + storage := NewReservationStorage(client, storageInstance.Notifier()) // Reserve with limit=1 — this should succeed because the stale entry // gets cleaned up by the reserveScript before counting diff --git a/packages/api/internal/sandbox/reservations/redis/utils.go b/packages/api/internal/sandbox/reservations/redis/utils.go index ab19d25d03..043be1f928 100644 --- a/packages/api/internal/sandbox/reservations/redis/utils.go +++ b/packages/api/internal/sandbox/reservations/redis/utils.go @@ -9,6 +9,7 @@ const ( reservationsKey = "reservations" pendingKey = "pending" resultKey = "result" + notifySuffix = "notify" ) // getStorageIndexKey returns the existing storage team index key (read-only). @@ -33,3 +34,9 @@ func getPendingSetKey(teamID string) string { func getResultKey(teamID, sandboxID string) string { return redis_utils.CreateKey(getReservationPrefix(teamID), sandboxID, resultKey) } + +// getReservationRoutingKey is PubSub routing key for reservation completion notifications. +// e.g. sandbox:storage:{teamID}:reservations:sandboxID:notify +func getReservationRoutingKey(teamID, sandboxID string) string { + return redis_utils.CreateKey(getReservationPrefix(teamID), sandboxID, notifySuffix) +} diff --git a/packages/api/internal/sandbox/storage/redis/notifier.go b/packages/api/internal/sandbox/storage/redis/notifier.go new file mode 100644 index 0000000000..c731d3c2d8 --- /dev/null +++ b/packages/api/internal/sandbox/storage/redis/notifier.go @@ -0,0 +1,30 @@ +package redis + +import ( + "context" +) + +// Notifier is the public seam onto the shared storage pub/sub infrastructure. +type Notifier struct { + sub *subscriptionManager + pub *publisher +} + +// Subscribe registers interest in routingKey. The returned channel is +// signaled (non-blocking, drop-on-full) whenever a matching message +// arrives on the shared notify channel. The caller MUST invoke cleanup +// when done to avoid a memory leak. +func (n *Notifier) Subscribe(routingKey string) (<-chan struct{}, func()) { + return n.sub.subscribe(routingKey) +} + +// Publish enqueues routingKey for asynchronous PUBLISH on the shared +// notify channel. Every consumer should use a fallback ticker. +func (n *Notifier) Publish(ctx context.Context, routingKey string) { + n.pub.Publish(ctx, routingKey) +} + +// Notifier returns the cross-package pub/sub seam +func (s *Storage) Notifier() *Notifier { + return &Notifier{sub: s.subManager, pub: s.publisher} +} diff --git a/packages/api/internal/sandbox/store.go b/packages/api/internal/sandbox/store.go index cb988090ec..d5a5aa633c 100644 --- a/packages/api/internal/sandbox/store.go +++ b/packages/api/internal/sandbox/store.go @@ -220,7 +220,3 @@ func (s *Store) Reserve(ctx context.Context, teamID uuid.UUID, sandboxID string, return finishStart, waitForStart, nil } - -func (s *Store) Release(ctx context.Context, teamID uuid.UUID, sandboxID string) error { - return s.reservations.Release(ctx, teamID, sandboxID) -}