diff --git a/internal/daemon/session.go b/internal/daemon/session.go index 092f7c7df..142731138 100644 --- a/internal/daemon/session.go +++ b/internal/daemon/session.go @@ -32,6 +32,9 @@ const ( // ErrSessionExists is returned when starting a session ID already in use. var ErrSessionExists = errors.New("daemon: session already exists") +// ErrSessionOverloaded means the manager's unfinished-session limit is reached. +var ErrSessionOverloaded = errors.New("daemon: pending session limit reached") + // ErrSessionNotFound is returned by Get/Attach for an unknown session ID. var ErrSessionNotFound = errors.New("daemon: session not found") @@ -219,6 +222,9 @@ func (s *Session) isFinished() bool { // daemon does not accumulate finished sessions without limit. const defaultMaxSessions = 256 +// defaultMaxPending bounds queued and running sessions independently of history. +const defaultMaxPending = 256 + // SessionManager owns the live session registry and routes each session to a // pool worker. Mirrors session-manager.js + roster.js. Finished sessions are // retained (so a late `attach` sees history) up to MaxSessions, past which the @@ -227,8 +233,10 @@ type SessionManager struct { pool *Pool maxBuffer int maxSessions int + maxPending int mu sync.Mutex + pending int // queued + running; protected by mu sessions map[string]*Session order []string // creation order, for FIFO eviction of finished sessions } @@ -240,6 +248,9 @@ type SessionManagerOptions struct { // MaxSessions caps the retained session registry; 0 => default. The oldest // FINISHED sessions are evicted once the cap is exceeded. MaxSessions int + // MaxPending caps unfinished sessions (queued plus running), independently + // of retained history. Non-positive values use the default of 256. + MaxPending int } // NewSessionManager builds a manager over pool. @@ -251,10 +262,15 @@ func NewSessionManager(opts SessionManagerOptions) (*SessionManager, error) { if maxSessions <= 0 { maxSessions = defaultMaxSessions } + maxPending := opts.MaxPending + if maxPending <= 0 { + maxPending = defaultMaxPending + } return &SessionManager{ pool: opts.Pool, maxBuffer: opts.MaxBuffer, maxSessions: maxSessions, + maxPending: maxPending, sessions: map[string]*Session{}, }, nil } @@ -262,13 +278,23 @@ func NewSessionManager(opts SessionManagerOptions) (*SessionManager, error) { // Start creates a session for spec and dispatches it to the pool. It returns the // Session immediately (non-blocking); the run proceeds in the background and the // session transitions queued -> running -> done/failed. A duplicate ID is -// rejected with ErrSessionExists. +// rejected with ErrSessionExists; a full admission limit returns +// ErrSessionOverloaded without creating a session or background goroutine. func (m *SessionManager) Start(ctx context.Context, spec WorkerSpec) (*Session, error) { m.mu.Lock() if _, exists := m.sessions[spec.Session]; exists { m.mu.Unlock() return nil, ErrSessionExists } + if err := ctx.Err(); err != nil { + m.mu.Unlock() + return nil, err + } + if m.pending >= m.maxPending { + m.mu.Unlock() + return nil, ErrSessionOverloaded + } + m.pending++ sess := newSession(spec.Session, spec.Cwd, m.maxBuffer) m.sessions[spec.Session] = sess m.order = append(m.order, spec.Session) @@ -277,7 +303,11 @@ func (m *SessionManager) Start(ctx context.Context, spec WorkerSpec) (*Session, go func() { code, err := m.pool.Run(ctx, spec, sess) + m.mu.Lock() sess.finish(code, err) + m.pending-- + m.pruneLocked() + m.mu.Unlock() }() return sess, nil } diff --git a/internal/daemon/session_test.go b/internal/daemon/session_test.go index a2caf1c1a..510fbaab0 100644 --- a/internal/daemon/session_test.go +++ b/internal/daemon/session_test.go @@ -3,6 +3,8 @@ package daemon import ( "context" "errors" + "fmt" + "sync" "sync/atomic" "testing" "time" @@ -208,6 +210,130 @@ func TestSessionManagerKeepsRunningOverCap(t *testing.T) { } } +func TestSessionManagerAdmissionBound(t *testing.T) { + block := make(chan struct{}) + launcher, _ := seqLauncher(&fakeWorker{pid: 1, waitCh: block}) + pool, _ := NewPool(PoolOptions{Size: 1, Launcher: launcher}) + mgr, _ := NewSessionManager(SessionManagerOptions{Pool: pool, MaxSessions: 2}) + ctx, cancel := context.WithCancel(context.Background()) + defer func() { + cancel() + close(block) + testutil.WaitFor(t, "all admitted sessions finished", func() bool { + mgr.mu.Lock() + defer mgr.mu.Unlock() + return mgr.pending == 0 + }) + }() + var wg sync.WaitGroup + var accepted atomic.Int32 + for i := 0; i < 300; i++ { + wg.Add(1) + go func() { + defer wg.Done() + if _, err := mgr.Start(ctx, WorkerSpec{Session: fmt.Sprintf("s%d", i)}); err == nil { + accepted.Add(1) + } + }() + } + wg.Wait() + if got := accepted.Load(); got != 256 { + t.Fatalf("accepted %d unfinished sessions, want default bound 256", got) + } + if got := len(mgr.Statuses()); got != 256 { + t.Fatalf("registered %d sessions, want 256", got) + } +} + +func TestSessionManagerCancellationReleasesAdmission(t *testing.T) { + block := make(chan struct{}) + var release sync.Once + defer release.Do(func() { close(block) }) + launcher, _ := seqLauncher(&fakeWorker{pid: 1, waitCh: block}, &fakeWorker{pid: 2, exitCode: ExitPermanent}, &fakeWorker{pid: 3}) + pool, _ := NewPool(PoolOptions{Size: 1, Launcher: launcher}) + mgr, _ := NewSessionManager(SessionManagerOptions{Pool: pool, MaxPending: 2, MaxSessions: 10}) + first, err := mgr.Start(context.Background(), WorkerSpec{Session: "running"}) + if err != nil { + t.Fatal(err) + } + testutil.WaitFor(t, "first running", func() bool { return first.State() == SessionRunning }) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + queued, err := mgr.Start(ctx, WorkerSpec{Session: "queued"}) + if err != nil { + t.Fatal(err) + } + if s, err := mgr.Start(context.Background(), WorkerSpec{Session: "retry"}); s != nil || !errors.Is(err, ErrSessionOverloaded) { + t.Fatalf("overloaded Start = (%v, %v)", s, err) + } + if _, ok := mgr.Get("retry"); ok { + t.Fatal("rejected session was registered") + } + if _, err := mgr.Start(context.Background(), WorkerSpec{Session: "running"}); !errors.Is(err, ErrSessionExists) { + t.Fatalf("duplicate at capacity: %v", err) + } + cancel() + testutil.WaitFor(t, "queued cancellation", func() bool { return queued.State() == SessionFailed }) + if !errors.Is(queued.Err(), context.Canceled) { + t.Fatalf("queued error = %v", queued.Err()) + } + if s, err := mgr.Start(ctx, WorkerSpec{Session: "canceled"}); s != nil || !errors.Is(err, context.Canceled) { + t.Fatalf("canceled Start = (%v, %v)", s, err) + } + retry, err := mgr.Start(context.Background(), WorkerSpec{Session: "retry"}) + if err != nil { + t.Fatalf("capacity not released after cancellation: %v", err) + } + release.Do(func() { close(block) }) + testutil.WaitFor(t, "first completed", func() bool { return first.State() == SessionDone }) + testutil.WaitFor(t, "retry failed", func() bool { return retry.State() == SessionFailed }) + if !errors.Is(retry.Err(), ErrPermanent) { + t.Fatalf("retry error = %v", retry.Err()) + } + last, err := mgr.Start(context.Background(), WorkerSpec{Session: "last"}) + if err != nil { + t.Fatal(err) + } + testutil.WaitFor(t, "last completed", func() bool { return last.State() == SessionDone }) + mgr.mu.Lock() + defer mgr.mu.Unlock() + if mgr.pending != 0 || len(mgr.sessions) != 4 { + t.Fatalf("pending=%d retained=%d, want 0 and 4", mgr.pending, len(mgr.sessions)) + } +} + +func TestSessionManagerPrunesOnCompletion(t *testing.T) { + block := make(chan struct{}) + var release sync.Once + defer release.Do(func() { close(block) }) + launcher, _ := seqLauncher( + &fakeWorker{pid: 1, waitCh: block}, + &fakeWorker{pid: 2}, + &fakeWorker{pid: 3, exitCode: ExitPermanent}, + ) + pool, _ := NewPool(PoolOptions{Size: 1, Launcher: launcher}) + mgr, _ := NewSessionManager(SessionManagerOptions{Pool: pool, MaxSessions: 2}) + sessions := make([]*Session, 3) + for i := range sessions { + var err error + sessions[i], err = mgr.Start(context.Background(), WorkerSpec{Session: fmt.Sprintf("s%d", i)}) + if err != nil { + t.Fatal(err) + } + } + release.Do(func() { close(block) }) + for _, s := range sessions { + select { + case <-s.Done(): + case <-time.After(2 * time.Second): + t.Fatal("session did not finish") + } + } + if got := len(mgr.Statuses()); got != 2 { + t.Fatalf("retained %d completed sessions without another Start, want 2", got) + } +} + func TestSessionStatuses(t *testing.T) { launcher, _ := seqLauncher(&fakeWorker{pid: 1, out: []string{"l"}, exitCode: 0}) pool, _ := NewPool(PoolOptions{Size: 1, Launcher: launcher})