perf(api): wake reservation waiters via pub/sub instead of 20ms polling [ENG-4070]

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:

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.

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
+}