From 1030bc1546786a6aac0d338a02a7718de147b1f6 Mon Sep 17 00:00:00 2001 From: Damilola Edwards Date: Wed, 22 Jul 2026 22:43:39 +0100 Subject: [PATCH] Close the session cap race under concurrent creation The session limit was enforced by checking the current count and then creating the session afterward, with nothing tying the two together. Concurrent requests could all see room under the cap before any of their sessions existed yet, so a burst of creations could sail past max_sessions. The docker backend's implicit new-session path used by execute_python had no cap check on it at all. Add a reservation step that holds a slot in memory from the moment a creation is allowed until the attempt finishes, so concurrent callers can no longer all pass the same stale count. Wire it into both the explicit session-create path and the docker backend's implicit new-session path. --- pkg/sandbox/direct.go | 4 +- pkg/sandbox/docker.go | 21 +++++++- pkg/sandbox/session.go | 70 +++++++++++++++++++++++- pkg/sandbox/session_test.go | 103 ++++++++++++++++++++++++++++++++++++ 4 files changed, 194 insertions(+), 4 deletions(-) diff --git a/pkg/sandbox/direct.go b/pkg/sandbox/direct.go index f8bd3174..00d59d7b 100644 --- a/pkg/sandbox/direct.go +++ b/pkg/sandbox/direct.go @@ -436,7 +436,9 @@ func (b *DirectBackend) CreateSession(ctx context.Context, ownerID string, _ map return "", fmt.Errorf("sessions are disabled") } - canCreate, count, maxAllowed := b.sessionManager.CanCreateSession(ctx, ownerID) + canCreate, release, count, maxAllowed := b.sessionManager.ReserveSession(ctx, ownerID) + defer release() + if !canCreate { return "", fmt.Errorf( "maximum sessions limit reached (%d/%d). Use manage_session with operation 'list' to see sessions, then 'destroy' to free up a slot", diff --git a/pkg/sandbox/docker.go b/pkg/sandbox/docker.go index fbe8aa6d..f3069706 100644 --- a/pkg/sandbox/docker.go +++ b/pkg/sandbox/docker.go @@ -338,6 +338,19 @@ func (b *DockerBackend) executeWithNewSession(ctx context.Context, req ExecuteRe timeout = time.Duration(b.cfg.Timeout) * time.Second } + // Reserve a slot for the new session before creating it, closing the gap + // where a burst of concurrent new-session executions could all pass a + // point-in-time cap check before any of their containers exist. + canCreate, release, count, maxAllowed := b.sessionManager.ReserveSession(ctx, req.OwnerID) + defer release() + + if !canCreate { + return nil, fmt.Errorf( + "maximum sessions limit reached (%d/%d). Use manage_session with operation 'list' to see sessions, then 'destroy' to free up a slot", + count, maxAllowed, + ) + } + // Generate session ID upfront so it can be stored in container labels. sessionID := b.sessionManager.GenerateSessionID() @@ -1177,8 +1190,12 @@ func (b *DockerBackend) CreateSession(ctx context.Context, ownerID string, env m return "", fmt.Errorf("sessions are disabled") } - // Check if we can create a new session. - canCreate, count, maxAllowed := b.sessionManager.CanCreateSession(ctx, ownerID) + // Reserve a slot for the new session. Held until this function returns, so + // a concurrent caller cannot also pass the cap check before this container + // exists. + canCreate, release, count, maxAllowed := b.sessionManager.ReserveSession(ctx, ownerID) + defer release() + if !canCreate { return "", fmt.Errorf( "maximum sessions limit reached (%d/%d). Use manage_session with operation 'list' to see sessions, then 'destroy' to free up a slot", diff --git a/pkg/sandbox/session.go b/pkg/sandbox/session.go index 77ff3a15..bb0a3b37 100644 --- a/pkg/sandbox/session.go +++ b/pkg/sandbox/session.go @@ -46,7 +46,12 @@ type SessionManager struct { lastUsed map[string]time.Time // activeExecs tracks sessions with running executions to prevent TTL purging mid-execution. activeExecs map[string]int - mu sync.RWMutex + // reserved tracks in-flight session creations that have passed the cap + // check but are not yet visible via store.List, keyed by owner ("" for the + // unauthenticated/global bucket). This closes the gap between deciding a + // session can be created and the store reflecting it. + reserved map[string]int + mu sync.RWMutex done chan struct{} stopOnce sync.Once @@ -63,6 +68,7 @@ func NewSessionManager(cfg config.SessionConfig, log logrus.FieldLogger, store S log: log.WithField("component", "session-manager"), lastUsed: make(map[string]time.Time, cfg.MaxSessions), activeExecs: make(map[string]int, cfg.MaxSessions), + reserved: make(map[string]int), done: make(chan struct{}), store: store, } @@ -346,6 +352,68 @@ func (m *SessionManager) CanCreateSession(ctx context.Context, ownerID string) ( return count < maxSessions, count, maxSessions } +// ReserveSession atomically checks the session cap and reserves a slot for a +// new session, for backends that actually create one. Unlike CanCreateSession, +// which only inspects a point-in-time count, ReserveSession accounts for +// creations that are already in flight but not yet visible via store.List, so +// concurrent callers cannot all observe room under the cap and overshoot it. +// +// On success the caller owns the returned release function and must call it +// exactly once, whether or not the session was actually created (a defer right +// after a successful reservation covers both the success and failure paths). +func (m *SessionManager) ReserveSession(ctx context.Context, ownerID string) (canCreate bool, release func(), count int, maxAllowed int) { + noop := func() {} + + if !m.cfg.IsEnabled() { + return false, noop, 0, 0 + } + + maxSessions := m.cfg.MaxSessions + if maxSessions <= 0 { + // No limit configured. + return true, noop, 0, 0 + } + + sessions, err := m.store.List(ctx) + if err != nil { + m.log.WithError(err).Warn("Failed to list sessions for limit check") + // Be conservative and allow creation on error. + return true, noop, 0, maxSessions + } + + known := 0 + for _, s := range sessions { + if ownerID == "" || s.OwnerID == ownerID { + known++ + } + } + + m.mu.Lock() + defer m.mu.Unlock() + + total := known + m.reserved[ownerID] + if total >= maxSessions { + return false, noop, total, maxSessions + } + + m.reserved[ownerID]++ + + var once sync.Once + release = func() { + once.Do(func() { + m.mu.Lock() + defer m.mu.Unlock() + + m.reserved[ownerID]-- + if m.reserved[ownerID] <= 0 { + delete(m.reserved, ownerID) + } + }) + } + + return true, release, total, maxSessions +} + // MaxSessions returns the configured maximum number of sessions. func (m *SessionManager) MaxSessions() int { return m.cfg.MaxSessions diff --git a/pkg/sandbox/session_test.go b/pkg/sandbox/session_test.go index 711935bb..46a8d302 100644 --- a/pkg/sandbox/session_test.go +++ b/pkg/sandbox/session_test.go @@ -2,7 +2,9 @@ package sandbox import ( "context" + "sync" "testing" + "time" "github.com/sirupsen/logrus" "github.com/stretchr/testify/require" @@ -53,3 +55,104 @@ func TestSessionManagerRemoveSessionThenUnmarkDoesNotRecreateState(t *testing.T) require.NotContains(t, m.lastUsed, sessionID) require.NotContains(t, m.activeExecs, sessionID) } + +// laggingStore models a backend whose List does not see a session the moment +// it is created, mirroring how a freshly started container or workspace takes +// a little while to show up wherever the backend enumerates live sessions. +type laggingStore struct { + mu sync.Mutex + sessions []*Session +} + +func (s *laggingStore) commit(session *Session) { + s.mu.Lock() + defer s.mu.Unlock() + + s.sessions = append(s.sessions, session) +} + +func (s *laggingStore) List(context.Context) ([]*Session, error) { + s.mu.Lock() + defer s.mu.Unlock() + + out := make([]*Session, len(s.sessions)) + copy(out, s.sessions) + + return out, nil +} + +func (s *laggingStore) Get(context.Context, string) (*Session, error) { return nil, nil } +func (s *laggingStore) Remove(context.Context, *Session) error { return nil } + +func TestReserveSessionHoldsCapUnderConcurrentBurst(t *testing.T) { + const maxSessions = 3 + const burst = 12 + + store := &laggingStore{} + enabled := true + m := NewSessionManager(config.SessionConfig{ + Enabled: &enabled, + MaxSessions: maxSessions, + }, logrus.New(), store) + + ctx := context.Background() + + var ( + wg sync.WaitGroup + mu sync.Mutex + created int + ) + + for i := 0; i < burst; i++ { + wg.Add(1) + + go func(n int) { + defer wg.Done() + + canCreate, release, _, _ := m.ReserveSession(ctx, "") + defer release() + + if !canCreate { + return + } + + // The backing resource takes a moment to exist and become visible + // to List, same as a real container starting up. + time.Sleep(20 * time.Millisecond) + store.commit(&Session{ID: string(rune('a' + n)), CreatedAt: time.Now()}) + + mu.Lock() + created++ + mu.Unlock() + }(i) + } + + wg.Wait() + + require.LessOrEqualf(t, created, maxSessions, + "expected at most %d sessions created under a concurrent burst, got %d", maxSessions, created) +} + +func TestReserveSessionReleaseFreesSlotForReuse(t *testing.T) { + enabled := true + m := NewSessionManager(config.SessionConfig{ + Enabled: &enabled, + MaxSessions: 1, + }, logrus.New(), emptyStore{}) + + ctx := context.Background() + + canCreate, release, _, _ := m.ReserveSession(ctx, "") + require.True(t, canCreate) + + canCreate, _, count, max := m.ReserveSession(ctx, "") + require.False(t, canCreate, "the single slot is already reserved") + require.Equal(t, 1, count) + require.Equal(t, 1, max) + + release() + + canCreate, release, _, _ = m.ReserveSession(ctx, "") + require.True(t, canCreate, "releasing the first reservation should free the slot") + release() +}