From dd54ec83fa5dba35e9e7db6e9291bdca81133aed Mon Sep 17 00:00:00 2001 From: PierrunoYT Date: Mon, 20 Jul 2026 17:03:37 +0200 Subject: [PATCH 1/2] fix(swarm): close lifecycle admission on shutdown --- internal/swarm/lifecycle.go | 93 +++++++++++++++------ internal/swarm/lifecycle_test.go | 136 +++++++++++++++++++++++++++++++ internal/swarm/scheduler.go | 25 +++--- internal/swarm/team.go | 73 ++++++++++++++--- 4 files changed, 278 insertions(+), 49 deletions(-) diff --git a/internal/swarm/lifecycle.go b/internal/swarm/lifecycle.go index aef9c34a9..7e220c133 100644 --- a/internal/swarm/lifecycle.go +++ b/internal/swarm/lifecycle.go @@ -12,6 +12,11 @@ import ( // team is at its slot cap the member is queued and launches when a slot frees // (it stays pending in the coordinator until then). func (s *Swarm) Spawn(pol Policy, teamName, agentType, task, cwd string) (string, error) { + if err := s.beginLifecycleAdmission(); err != nil { + return "", err + } + defer s.lifecycleMu.RUnlock() + def, err := s.registry.Lookup(agentType) if err != nil { return "", err @@ -23,15 +28,25 @@ func (s *Swarm) Spawn(pol Policy, teamName, agentType, task, cwd string) (string } s.rememberCwd(id, cwd) spec := s.buildSpec(pol, id, id, team, def, task, cwd) - s.dispatch(spec) + s.dispatchAdmitted(spec) return id, nil } +func (s *Swarm) beginLifecycleAdmission() error { + s.lifecycleMu.RLock() + if s.closed { + s.lifecycleMu.RUnlock() + return ErrSwarmClosed + } + return nil +} + // dispatch admits a spec to its team (launching now or queuing for a slot). -func (s *Swarm) dispatch(spec MemberSpec) { +// The caller must hold lifecycleMu for reading. +func (s *Swarm) dispatchAdmitted(spec MemberSpec) { t := s.team(spec.Team) if t.admit(spec) { - s.launch(t, spec) + s.launchAdmitted(t, spec) } // Otherwise the spec is queued; the coordinator task stays pending until a // slot frees and afterExit launches it. @@ -39,17 +54,22 @@ func (s *Swarm) dispatch(spec MemberSpec) { // launch starts a member for spec and supervises it. A synchronous launch // failure fails the task and frees the slot. -func (s *Swarm) launch(t *Team, spec MemberSpec) { +// The caller must hold lifecycleMu for reading. +func (s *Swarm) launchAdmitted(t *Team, spec MemberSpec) { handle, err := s.launcher.Launch(s.baseCtx, spec) if err != nil { _ = s.coord.Fail(spec.TaskID, "launch: "+err.Error()) - s.afterExit(t) + s.afterExitAdmitted(t) return } m := &Member{ID: spec.ID, AgentType: spec.AgentType, TaskID: spec.TaskID, handle: handle} t.addMember(m) _ = s.coord.SetStatus(spec.TaskID, StatusRunning) - go s.watch(t, m, spec) + s.watchers.Add(1) + go func() { + defer s.watchers.Done() + s.watch(t, m, spec) + }() } // watch awaits a member, applies bounded relaunch on temporary failures, records @@ -60,22 +80,28 @@ func (s *Swarm) launch(t *Team, spec MemberSpec) { // life (Member.ID == MemberSpec.ID); a retry never reuses the struct for a // different spec. func (s *Swarm) watch(t *Team, m *Member, spec MemberSpec) { - res, err := m.handle.Wait() - if err != nil { - if isRetryable(err) && m.restarts < maxMemberRestarts { - if nh, relErr := s.launcher.Launch(s.baseCtx, spec); relErr == nil { - m.restarts++ - m.handle = nh - go s.watch(t, m, spec) - return + for { + res, err := m.handle.Wait() + if err != nil { + if isRetryable(err) && m.restarts < maxMemberRestarts { + if s.beginLifecycleAdmission() == nil { + nh, relErr := s.launcher.Launch(s.baseCtx, spec) + s.lifecycleMu.RUnlock() + if relErr == nil { + m.restarts++ + m.handle = nh + continue + } + } + // Fall through if shutdown started or relaunch failed. } - // fall through to record the original error if relaunch fails + // res.SessionID is preserved by the handle even on error, so a member that + // ran then failed stays drillable; a pure launch error carries an empty id. + _ = s.coord.FailWithSession(m.TaskID, memberError(err), res.SessionID) + } else { + _ = s.coord.CompleteWithSession(m.TaskID, res.Result, res.SessionID) } - // res.SessionID is preserved by the handle even on error, so a member that - // ran then failed stays drillable; a pure launch error carries an empty id. - _ = s.coord.FailWithSession(m.TaskID, memberError(err), res.SessionID) - } else { - _ = s.coord.CompleteWithSession(m.TaskID, res.Result, res.SessionID) + break } t.removeMember(m.ID) s.afterExit(t) @@ -85,17 +111,33 @@ func (s *Swarm) watch(t *Team, m *Member, spec MemberSpec) { // any. Each exit drains at most one queued member; that member's own exit drains // the next, so the queue empties one-per-slot without unbounded recursion. func (s *Swarm) afterExit(t *Team) { + if s.beginLifecycleAdmission() != nil { + t.releaseSlot() + return + } + defer s.lifecycleMu.RUnlock() + s.afterExitAdmitted(t) +} + +// afterExitAdmitted drains at most one queued spec while lifecycle admission is +// held open. The caller must hold lifecycleMu for reading. +func (s *Swarm) afterExitAdmitted(t *Team) { next, ok := t.onExit() if !ok { return } - s.launch(t, next) + s.launchAdmitted(t, next) } // Handoff transfers a task to a fresh member of toAgentType, delivering a note to // the new member's inbox and marking the original task handed-off. It returns the // new task id. A handoff of an already-terminal task is rejected (fail closed). func (s *Swarm) Handoff(pol Policy, teamName, taskID, toAgentType, note string) (string, error) { + if err := s.beginLifecycleAdmission(); err != nil { + return "", err + } + defer s.lifecycleMu.RUnlock() + task, ok := s.coord.Get(taskID) if !ok { return "", fmt.Errorf("%w: %s", ErrUnknownTask, taskID) @@ -131,7 +173,7 @@ func (s *Swarm) Handoff(pol Policy, teamName, taskID, toAgentType, note string) // Retire the original task (it has been re-delegated). _ = s.coord.SetStatus(taskID, StatusHandedOff) spec := s.buildSpec(pol, newID, newID, team, def, handoffTask, cwd) - s.dispatch(spec) + s.dispatchAdmitted(spec) return newID, nil } @@ -139,6 +181,11 @@ func (s *Swarm) Handoff(pol Policy, teamName, taskID, toAgentType, note string) // (e.g. a crashed worker) onto fresh members of toAgentType, returning the // adopted task ids. Terminal tasks and tasks with a live owner are left alone. func (s *Swarm) AdoptOrphans(pol Policy, teamName, toAgentType string) ([]string, error) { + if err := s.beginLifecycleAdmission(); err != nil { + return nil, err + } + defer s.lifecycleMu.RUnlock() + def, err := s.registry.Lookup(toAgentType) if err != nil { return nil, err @@ -161,7 +208,7 @@ func (s *Swarm) AdoptOrphans(pol Policy, teamName, toAgentType string) ([]string cwd := s.cwdFor(task.ID) s.rememberCwd(task.ID, cwd) spec := s.buildSpec(pol, newAgent, task.ID, team, def, task.Description, cwd) - s.dispatch(spec) + s.dispatchAdmitted(spec) adopted = append(adopted, task.ID) } return adopted, nil diff --git a/internal/swarm/lifecycle_test.go b/internal/swarm/lifecycle_test.go index 8f8dd072d..9bca1b784 100644 --- a/internal/swarm/lifecycle_test.go +++ b/internal/swarm/lifecycle_test.go @@ -2,6 +2,7 @@ package swarm import ( "context" + "errors" "strings" "sync" "testing" @@ -217,6 +218,141 @@ func TestPermanentErrorNoRetry(t *testing.T) { } } +func TestLifecycleAdmissionRejectedAfterClose(t *testing.T) { + l := newLauncher(okFor) + sw := newSwarmFor(t, l) + sw.Close() + + if _, err := sw.Spawn(Policy{}, "team", "teammate", "task", ""); !errors.Is(err, ErrSwarmClosed) { + t.Fatalf("Spawn after Close error = %v, want ErrSwarmClosed", err) + } + if _, err := sw.Handoff(Policy{}, "team", "task", "teammate", "note"); !errors.Is(err, ErrSwarmClosed) { + t.Fatalf("Handoff after Close error = %v, want ErrSwarmClosed", err) + } + if _, err := sw.AdoptOrphans(Policy{}, "team", "teammate"); !errors.Is(err, ErrSwarmClosed) { + t.Fatalf("AdoptOrphans after Close error = %v, want ErrSwarmClosed", err) + } + if got := len(l.recorded()); got != 0 { + t.Fatalf("launches after Close = %d, want 0", got) + } +} + +func TestCloseDoesNotLaunchQueuedMembers(t *testing.T) { + gate := make(chan struct{}) + l := newLauncher(okFor) + l.gate = gate + sw := newSwarmFor(t, l) + + var ids []string + for i := 0; i < 3; i++ { + id, err := sw.Spawn(Policy{}, "team", "teammate", "task", "") + if err != nil { + t.Fatalf("Spawn %d: %v", i, err) + } + ids = append(ids, id) + } + if got := len(l.recorded()); got != 2 { + t.Fatalf("initial launches = %d, want 2 with one queued", got) + } + + sw.Close() + if got := len(l.recorded()); got != 2 { + t.Fatalf("launches after Close = %d, want queued member not launched", got) + } + for _, id := range ids { + task, ok := sw.Coordinator().Get(id) + if !ok || !task.Status.terminal() { + t.Fatalf("task %s after Close = %+v, want terminal", id, task) + } + } +} + +func TestClosePreventsMemberRetry(t *testing.T) { + started := make(chan struct{}, maxMemberRestarts+1) + release := make(chan struct{}) + l := FuncLauncher{Run: func(context.Context, MemberSpec) (MemberResult, error) { + started <- struct{}{} + <-release + return MemberResult{}, ErrMemberTemporary + }} + sw := newSwarmFor(t, l) + _, err := sw.Spawn(Policy{}, "team", "teammate", "task", "") + if err != nil { + t.Fatalf("Spawn: %v", err) + } + select { + case <-started: + case <-time.After(3 * time.Second): + t.Fatal("initial member did not start") + } + + closed := make(chan struct{}) + go func() { + sw.Close() + close(closed) + }() + waitFor(t, "swarm closed state", func() bool { + sw.lifecycleMu.RLock() + defer sw.lifecycleMu.RUnlock() + return sw.closed + }) + close(release) + select { + case <-closed: + case <-time.After(3 * time.Second): + t.Fatal("Close did not return after retryable member exit") + } + select { + case <-started: + t.Fatal("member retried after Close") + default: + } +} + +func TestCloseWaitsForMemberWatchers(t *testing.T) { + release := make(chan struct{}) + l := FuncLauncher{Run: func(context.Context, MemberSpec) (MemberResult, error) { + <-release + return MemberResult{Result: "done"}, nil + }} + sw := newSwarmFor(t, l) + id, err := sw.Spawn(Policy{}, "team", "teammate", "task", "") + if err != nil { + t.Fatalf("Spawn: %v", err) + } + waitFor(t, "task running", func() bool { + task, ok := sw.Coordinator().Get(id) + return ok && task.Status == StatusRunning + }) + + const callers = 3 + closing := make(chan struct{}, callers) + closed := make(chan struct{}, callers) + for i := 0; i < callers; i++ { + go func() { + closing <- struct{}{} + sw.Close() + closed <- struct{}{} + }() + } + for i := 0; i < callers; i++ { + <-closing + } + select { + case <-closed: + t.Fatal("a Close caller returned before the member watcher exited") + case <-time.After(50 * time.Millisecond): + } + close(release) + for i := 0; i < callers; i++ { + select { + case <-closed: + case <-time.After(3 * time.Second): + t.Fatal("all Close callers did not return after the member watcher exited") + } + } +} + func TestHandoffDeliversNoteAndRetiresOriginal(t *testing.T) { gate := make(chan struct{}) l := newLauncher(okFor) diff --git a/internal/swarm/scheduler.go b/internal/swarm/scheduler.go index 923fb58ab..d078dee82 100644 --- a/internal/swarm/scheduler.go +++ b/internal/swarm/scheduler.go @@ -121,11 +121,12 @@ type Scheduler struct { ctx context.Context cancel context.CancelFunc - mu sync.Mutex - jobs map[string]*scheduledJob - closed bool - seq atomic.Uint64 - wg sync.WaitGroup + mu sync.Mutex + jobs map[string]*scheduledJob + closed bool + seq atomic.Uint64 + wg sync.WaitGroup + closeOnce sync.Once } // newScheduler builds a Scheduler bound to sw. Its context derives from the @@ -229,15 +230,13 @@ func (s *Scheduler) List() []JobStatus { // Close cancels every job and waits for their loops to exit. Safe to call more // than once. func (s *Scheduler) Close() { - s.mu.Lock() - if s.closed { + s.closeOnce.Do(func() { + s.mu.Lock() + s.closed = true + s.cancel() s.mu.Unlock() - return - } - s.closed = true - s.cancel() - s.mu.Unlock() - s.wg.Wait() + s.wg.Wait() + }) } // run is one job's timing loop. It requests a fresh one-shot ticker per interval diff --git a/internal/swarm/team.go b/internal/swarm/team.go index c913f987f..5e78357d3 100644 --- a/internal/swarm/team.go +++ b/internal/swarm/team.go @@ -64,6 +64,11 @@ type Swarm struct { baseCtx context.Context cancel context.CancelFunc + lifecycleMu sync.RWMutex + closed bool + watchers sync.WaitGroup + closeOnce sync.Once + mu sync.Mutex teams map[string]*Team taskCwd map[string]string // taskID -> cwd, for handoff/adoption relaunch @@ -137,31 +142,57 @@ func New(opts Options) (*Swarm, error) { } // Close cancels every running member's context and releases resources. It is -// safe to call more than once. The scheduler is closed first so no new spawn -// fires after shutdown begins. +// safe to call more than once. It closes lifecycle admission before canceling +// members, fails queued tasks, and waits for every member watcher to exit. func (s *Swarm) Close() { - s.mu.Lock() - sched := s.scheduler - s.mu.Unlock() - if sched != nil { - sched.Close() - } - if s.cancel != nil { - s.cancel() - } + s.closeOnce.Do(func() { + s.lifecycleMu.Lock() + s.closed = true + if s.cancel != nil { + s.cancel() + } + s.mu.Lock() + sched := s.scheduler + teams := make([]*Team, 0, len(s.teams)) + for _, team := range s.teams { + teams = append(teams, team) + } + s.mu.Unlock() + s.lifecycleMu.Unlock() + + if sched != nil { + sched.Close() + } + for _, team := range teams { + for _, spec := range team.clearQueue() { + _ = s.coord.Fail(spec.TaskID, ErrSwarmClosed.Error()) + } + } + s.watchers.Wait() + }) } // Scheduler returns the swarm's recurring-spawn scheduler, creating it on first // use. Scheduling is opt-in: until a job is added the scheduler does nothing. func (s *Swarm) Scheduler() *Scheduler { + s.lifecycleMu.RLock() + closed := s.closed s.mu.Lock() - defer s.mu.Unlock() if s.scheduler == nil { s.scheduler = newScheduler(s) } - return s.scheduler + sched := s.scheduler + s.mu.Unlock() + s.lifecycleMu.RUnlock() + if closed { + sched.Close() + } + return sched } +// ErrSwarmClosed reports lifecycle work submitted after shutdown begins. +var ErrSwarmClosed = errors.New("swarm: closed") + // rememberCwd records a task's working dir so a handoff/adoption relaunch keeps it. func (s *Swarm) rememberCwd(taskID, cwd string) { s.mu.Lock() @@ -288,6 +319,22 @@ func (t *Team) onExit() (MemberSpec, bool) { return MemberSpec{}, false } +func (t *Team) releaseSlot() { + t.mu.Lock() + if t.running > 0 { + t.running-- + } + t.mu.Unlock() +} + +func (t *Team) clearQueue() []MemberSpec { + t.mu.Lock() + defer t.mu.Unlock() + queued := t.queue + t.queue = nil + return queued +} + func (t *Team) addMember(m *Member) { t.mu.Lock() t.members[m.ID] = m From 3e18b0b1e9d3f6faa241c4ea69842e717028ad6e Mon Sep 17 00:00:00 2001 From: PierrunoYT Date: Sat, 25 Jul 2026 16:40:24 +0200 Subject: [PATCH 2/2] fix(swarm): admit lifecycle work without holding the lock across Launch MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses the P1 on #776 plus the two documentation P3s. [P1] beginLifecycleAdmission returned with lifecycleMu.RLock() held for the caller's entire duration — including the call into the injected MemberLauncher.Launch. Close needs the write lock before it can call s.cancel(), so a launcher that blocks in Launch waiting on its context (the ordinary shape of a launcher that waits for a slot, a daemon connection, or a sandbox handshake) deadlocked every Close caller: the launch held the read lock while Close waited for the write lock to reach the very cancellation the launch needed to unblock. beginLifecycleAdmission is now a counted ticket: it holds lifecycleMu only across the closed check and a lifecycleWork.Add(1), then returns a release func (lifecycleWork.Done) the caller defers. Close flips closed and releases the lock BEFORE calling cancel, then waits out lifecycleWork (admitted-but-unfinished work) followed by watchers (goroutines admitted work may still be starting) — so it remains a true completion barrier without holding any lock across external launcher code. Every call site (Spawn, Handoff, AdoptOrphans, the retry path in watch, and afterExit) moved to the ticket pattern. [P3] beginLifecycleAdmission's contract and the Scheduler/Close handoff are now both documented at the point that matters: the func-returning signature makes the release contract explicit in the type itself, and Scheduler's comment spells out why the read/write lock ordering guarantees exactly one of Scheduler or Close closes whatever scheduler ends up installed. Tests: TestCloseCancelsLauncherBlockedInLaunch reproduces the exact deadlock (spawn parked in Launch, Close in a second goroutine, both must complete once cancellation reaches the launcher) — verified it hangs under the old held-lock pattern by temporarily restoring it. TestCloseWaitsForAdmittedSpawn is the converse: Close must not return while an admitted spawn is still running. Validated with `go test -race ./internal/swarm/... -count=1` under WSL/Linux (this session's Windows host has no C toolchain for -race directly), plus `go vet ./...` under WSL and go build/vet/test on Windows. --- internal/swarm/lifecycle.go | 57 +++++++++++----- internal/swarm/lifecycle_test.go | 111 +++++++++++++++++++++++++++++++ internal/swarm/team.go | 45 +++++++++++-- 3 files changed, 191 insertions(+), 22 deletions(-) diff --git a/internal/swarm/lifecycle.go b/internal/swarm/lifecycle.go index 7e220c133..be1a89005 100644 --- a/internal/swarm/lifecycle.go +++ b/internal/swarm/lifecycle.go @@ -12,10 +12,11 @@ import ( // team is at its slot cap the member is queued and launches when a slot frees // (it stays pending in the coordinator until then). func (s *Swarm) Spawn(pol Policy, teamName, agentType, task, cwd string) (string, error) { - if err := s.beginLifecycleAdmission(); err != nil { + release, err := s.beginLifecycleAdmission() + if err != nil { return "", err } - defer s.lifecycleMu.RUnlock() + defer release() def, err := s.registry.Lookup(agentType) if err != nil { @@ -32,17 +33,36 @@ func (s *Swarm) Spawn(pol Policy, teamName, agentType, task, cwd string) (string return id, nil } -func (s *Swarm) beginLifecycleAdmission() error { +// beginLifecycleAdmission admits one unit of lifecycle work if shutdown has not +// begun. On success it returns a release func the caller MUST call exactly once +// (a deferred call is the norm) — Close blocks until every admitted unit has +// released. On shutdown it returns a nil func and ErrSwarmClosed, and there is +// nothing to release. +// +// It deliberately does NOT keep lifecycleMu held for the caller's duration, which +// an earlier revision did. Close needs the write lock before it can cancel, so a +// launcher that blocks in Launch until its context is cancelled would deadlock +// against every Close caller: the launch holds the read lock while Close waits +// for the write lock to issue the cancellation neither side can reach. Holding the +// lock only across the flag check and the counter increment keeps the ordering +// guarantee (no admission once closed is set) without letting external launcher +// code sit inside the critical section. +// +// The read lock is enough to make the WaitGroup safe: Close flips closed under +// the write lock, so an Add here either completes before Close acquires it (and is +// therefore visible to the later Wait) or observes closed and is refused. +func (s *Swarm) beginLifecycleAdmission() (func(), error) { s.lifecycleMu.RLock() + defer s.lifecycleMu.RUnlock() if s.closed { - s.lifecycleMu.RUnlock() - return ErrSwarmClosed + return nil, ErrSwarmClosed } - return nil + s.lifecycleWork.Add(1) + return s.lifecycleWork.Done, nil } // dispatch admits a spec to its team (launching now or queuing for a slot). -// The caller must hold lifecycleMu for reading. +// The caller must hold an admission ticket from beginLifecycleAdmission. func (s *Swarm) dispatchAdmitted(spec MemberSpec) { t := s.team(spec.Team) if t.admit(spec) { @@ -54,7 +74,7 @@ func (s *Swarm) dispatchAdmitted(spec MemberSpec) { // launch starts a member for spec and supervises it. A synchronous launch // failure fails the task and frees the slot. -// The caller must hold lifecycleMu for reading. +// The caller must hold an admission ticket from beginLifecycleAdmission. func (s *Swarm) launchAdmitted(t *Team, spec MemberSpec) { handle, err := s.launcher.Launch(s.baseCtx, spec) if err != nil { @@ -84,9 +104,9 @@ func (s *Swarm) watch(t *Team, m *Member, spec MemberSpec) { res, err := m.handle.Wait() if err != nil { if isRetryable(err) && m.restarts < maxMemberRestarts { - if s.beginLifecycleAdmission() == nil { + if release, admitErr := s.beginLifecycleAdmission(); admitErr == nil { nh, relErr := s.launcher.Launch(s.baseCtx, spec) - s.lifecycleMu.RUnlock() + release() if relErr == nil { m.restarts++ m.handle = nh @@ -111,16 +131,17 @@ func (s *Swarm) watch(t *Team, m *Member, spec MemberSpec) { // any. Each exit drains at most one queued member; that member's own exit drains // the next, so the queue empties one-per-slot without unbounded recursion. func (s *Swarm) afterExit(t *Team) { - if s.beginLifecycleAdmission() != nil { + release, err := s.beginLifecycleAdmission() + if err != nil { t.releaseSlot() return } - defer s.lifecycleMu.RUnlock() + defer release() s.afterExitAdmitted(t) } // afterExitAdmitted drains at most one queued spec while lifecycle admission is -// held open. The caller must hold lifecycleMu for reading. +// held open. The caller must hold an admission ticket from beginLifecycleAdmission. func (s *Swarm) afterExitAdmitted(t *Team) { next, ok := t.onExit() if !ok { @@ -133,10 +154,11 @@ func (s *Swarm) afterExitAdmitted(t *Team) { // the new member's inbox and marking the original task handed-off. It returns the // new task id. A handoff of an already-terminal task is rejected (fail closed). func (s *Swarm) Handoff(pol Policy, teamName, taskID, toAgentType, note string) (string, error) { - if err := s.beginLifecycleAdmission(); err != nil { + release, err := s.beginLifecycleAdmission() + if err != nil { return "", err } - defer s.lifecycleMu.RUnlock() + defer release() task, ok := s.coord.Get(taskID) if !ok { @@ -181,10 +203,11 @@ func (s *Swarm) Handoff(pol Policy, teamName, taskID, toAgentType, note string) // (e.g. a crashed worker) onto fresh members of toAgentType, returning the // adopted task ids. Terminal tasks and tasks with a live owner are left alone. func (s *Swarm) AdoptOrphans(pol Policy, teamName, toAgentType string) ([]string, error) { - if err := s.beginLifecycleAdmission(); err != nil { + release, err := s.beginLifecycleAdmission() + if err != nil { return nil, err } - defer s.lifecycleMu.RUnlock() + defer release() def, err := s.registry.Lookup(toAgentType) if err != nil { diff --git a/internal/swarm/lifecycle_test.go b/internal/swarm/lifecycle_test.go index 9bca1b784..881594e82 100644 --- a/internal/swarm/lifecycle_test.go +++ b/internal/swarm/lifecycle_test.go @@ -500,3 +500,114 @@ func TestSpawnUnknownAgentType(t *testing.T) { t.Fatal("Spawn with unknown agent type must error") } } + +// blockingLauncher blocks inside Launch until its context is cancelled — the +// shape of a real launcher that waits on a slot, a daemon connection, or a +// sandbox handshake before it can return a handle. +type blockingLauncher struct { + entered chan struct{} + once sync.Once +} + +func (l *blockingLauncher) Launch(ctx context.Context, spec MemberSpec) (MemberHandle, error) { + l.once.Do(func() { close(l.entered) }) + <-ctx.Done() + return nil, ctx.Err() +} + +// TestCloseCancelsLauncherBlockedInLaunch is the regression test for holding +// lifecycle admission across MemberLauncher.Launch. When admission was a read lock +// held for the caller's whole duration, this deadlocked: the spawn sat inside +// Launch waiting for its context, while Close waited for the write lock it needed +// before it could cancel that very context. Neither side could progress, so both +// the spawn and every Close caller hung. +// +// Admission is now a counted ticket, so Close can flip the flag, cancel, and then +// wait the ticket out. +func TestCloseCancelsLauncherBlockedInLaunch(t *testing.T) { + launcher := &blockingLauncher{entered: make(chan struct{})} + sw := newSwarmFor(t, launcher) + + spawned := make(chan error, 1) + go func() { + _, err := sw.Spawn(Policy{}, "team", "teammate", "task", "") + spawned <- err + }() + select { + case <-launcher.entered: + case <-time.After(3 * time.Second): + t.Fatal("launcher never entered Launch") + } + + closed := make(chan struct{}) + go func() { + sw.Close() + close(closed) + }() + select { + case <-closed: + case <-time.After(5 * time.Second): + t.Fatal("Close deadlocked against a launcher blocked in Launch") + } + select { + case <-spawned: + case <-time.After(3 * time.Second): + t.Fatal("Spawn never returned after Close cancelled the launch context") + } +} + +// TestCloseWaitsForAdmittedSpawn pins the other half of the contract: Close is a +// barrier, so it must not return while an admitted spawn is still running. Without +// the lifecycleWork wait, Close could finish while this launch was mid-flight. +func TestCloseWaitsForAdmittedSpawn(t *testing.T) { + entered := make(chan struct{}) + finish := make(chan struct{}) + launcher := &gatedLaunchLauncher{entered: entered, finish: finish} + sw := newSwarmFor(t, launcher) + + go func() { + _, _ = sw.Spawn(Policy{}, "team", "teammate", "task", "") + }() + select { + case <-entered: + case <-time.After(3 * time.Second): + t.Fatal("launcher never entered Launch") + } + + closed := make(chan struct{}) + go func() { + sw.Close() + close(closed) + }() + select { + case <-closed: + t.Fatal("Close returned while an admitted spawn was still in flight") + case <-time.After(200 * time.Millisecond): + } + close(finish) + select { + case <-closed: + case <-time.After(5 * time.Second): + t.Fatal("Close did not return after the admitted spawn finished") + } +} + +// gatedLaunchLauncher blocks in Launch until the test releases it, ignoring +// context cancellation so the test controls exactly when the admitted work ends. +type gatedLaunchLauncher struct { + entered chan struct{} + finish chan struct{} + once sync.Once +} + +func (l *gatedLaunchLauncher) Launch(_ context.Context, spec MemberSpec) (MemberHandle, error) { + l.once.Do(func() { close(l.entered) }) + <-l.finish + return &funcHandle{id: spec.ID, done: closedChan()}, nil +} + +func closedChan() chan struct{} { + ch := make(chan struct{}) + close(ch) + return ch +} diff --git a/internal/swarm/team.go b/internal/swarm/team.go index 5e78357d3..35bbcc53d 100644 --- a/internal/swarm/team.go +++ b/internal/swarm/team.go @@ -64,10 +64,17 @@ type Swarm struct { baseCtx context.Context cancel context.CancelFunc + // lifecycleMu guards closed and orders admission against shutdown. It is held + // only for the flag check plus the lifecycleWork increment — never across a + // launcher call, which would deadlock Close (see beginLifecycleAdmission). lifecycleMu sync.RWMutex closed bool - watchers sync.WaitGroup - closeOnce sync.Once + // lifecycleWork counts admitted-but-unfinished lifecycle operations (spawn, + // handoff, adoption, retry, queue drain). Close waits it out so no admitted + // operation is still touching the swarm when Close returns. + lifecycleWork sync.WaitGroup + watchers sync.WaitGroup + closeOnce sync.Once mu sync.Mutex teams map[string]*Team @@ -146,11 +153,14 @@ func New(opts Options) (*Swarm, error) { // members, fails queued tasks, and waits for every member watcher to exit. func (s *Swarm) Close() { s.closeOnce.Do(func() { + // Closing admission and cancelling are separate steps on purpose. The flag + // flips under lifecycleMu so no further work is admitted, and the lock is + // released BEFORE cancelling: a member already being launched holds no lock + // here, but it may be blocked inside MemberLauncher.Launch waiting for the + // context this cancel provides, and it can only release its admission ticket + // once that cancellation lands. s.lifecycleMu.Lock() s.closed = true - if s.cancel != nil { - s.cancel() - } s.mu.Lock() sched := s.scheduler teams := make([]*Team, 0, len(s.teams)) @@ -160,6 +170,9 @@ func (s *Swarm) Close() { s.mu.Unlock() s.lifecycleMu.Unlock() + if s.cancel != nil { + s.cancel() + } if sched != nil { sched.Close() } @@ -168,12 +181,34 @@ func (s *Swarm) Close() { _ = s.coord.Fail(spec.TaskID, ErrSwarmClosed.Error()) } } + // Admitted work first, then the watchers it may still be starting: a spawn + // in flight can add a watcher, so waiting on watchers alone could miss one. + s.lifecycleWork.Wait() s.watchers.Wait() }) } // Scheduler returns the swarm's recurring-spawn scheduler, creating it on first // use. Scheduling is opt-in: until a job is added the scheduler does nothing. +// +// The post-close handoff between this method and Close relies on the ORDER in +// which each reads/writes s.scheduler and s.closed, not on holding one lock across +// both: Close takes lifecycleMu.Lock() (excluding this RLock), sets closed, reads +// whatever s.scheduler currently is, and only THEN releases lifecycleMu and calls +// sched.Close() on what it captured. So a call here either: +// - runs entirely before Close's write lock — sees closed == false, and the +// scheduler it creates/returns is exactly the one Close's later snapshot will +// capture and close; or +// - runs entirely after Close's write lock — sees closed == true (memory model: +// the RLock/Lock pair orders the read after Close's write) and closes the +// scheduler itself here, because Close already captured its own snapshot +// (possibly nil, if this call is what first creates one) and will not look +// again. +// +// Either way exactly one of the two call sites closes the scheduler that ends up +// installed, and neither can observe a "created after close but never closed" +// scheduler — that would require reading closed == false here while also reading +// a scheduler Close had already passed over, which the lock ordering rules out. func (s *Swarm) Scheduler() *Scheduler { s.lifecycleMu.RLock() closed := s.closed