diff --git a/internal/swarm/lifecycle.go b/internal/swarm/lifecycle.go index aef9c34a9..be1a89005 100644 --- a/internal/swarm/lifecycle.go +++ b/internal/swarm/lifecycle.go @@ -12,6 +12,12 @@ 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) { + release, err := s.beginLifecycleAdmission() + if err != nil { + return "", err + } + defer release() + def, err := s.registry.Lookup(agentType) if err != nil { return "", err @@ -23,15 +29,44 @@ 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 } +// 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 { + return nil, ErrSwarmClosed + } + s.lifecycleWork.Add(1) + return s.lifecycleWork.Done, 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 an admission ticket from beginLifecycleAdmission. +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 +74,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 an admission ticket from beginLifecycleAdmission. +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 +100,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 release, admitErr := s.beginLifecycleAdmission(); admitErr == nil { + nh, relErr := s.launcher.Launch(s.baseCtx, spec) + release() + 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 +131,35 @@ 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) { + release, err := s.beginLifecycleAdmission() + if err != nil { + t.releaseSlot() + return + } + defer release() + s.afterExitAdmitted(t) +} + +// afterExitAdmitted drains at most one queued spec while lifecycle admission is +// held open. The caller must hold an admission ticket from beginLifecycleAdmission. +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) { + release, err := s.beginLifecycleAdmission() + if err != nil { + return "", err + } + defer release() + task, ok := s.coord.Get(taskID) if !ok { return "", fmt.Errorf("%w: %s", ErrUnknownTask, taskID) @@ -131,7 +195,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 +203,12 @@ 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) { + release, err := s.beginLifecycleAdmission() + if err != nil { + return nil, err + } + defer release() + def, err := s.registry.Lookup(toAgentType) if err != nil { return nil, err @@ -161,7 +231,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..881594e82 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) @@ -364,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/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..35bbcc53d 100644 --- a/internal/swarm/team.go +++ b/internal/swarm/team.go @@ -64,6 +64,18 @@ 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 + // 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 taskCwd map[string]string // taskID -> cwd, for handoff/adoption relaunch @@ -137,31 +149,85 @@ 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() { + // 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 + 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 s.cancel != nil { + s.cancel() + } + if sched != nil { + sched.Close() + } + for _, team := range teams { + for _, spec := range team.clearQueue() { + _ = 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 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 +354,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