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