diff --git a/internal/daemon/pool.go b/internal/daemon/pool.go index 6de38315b..235d57a91 100644 --- a/internal/daemon/pool.go +++ b/internal/daemon/pool.go @@ -92,13 +92,15 @@ type Pool struct { opts PoolOptions slots chan struct{} - mu sync.Mutex - draining bool - active map[int]WorkerHandle // worker id -> handle, for drain/kill + status - nextID int - - drainOnce sync.Once - drained chan struct{} + mu sync.Mutex + draining bool + active map[int]WorkerHandle // worker id -> handle, for drain/kill + status + launching int // launchers in progress; Drain must not mistake these for idle + nextID int + + drainStartOnce sync.Once + drainOnce sync.Once + drained chan struct{} } // workerStat tracks one in-flight request's restart count (local to Run). @@ -197,12 +199,23 @@ func (p *Pool) Run(ctx context.Context, spec WorkerSpec, sink Sink) (int, error) return 0, ErrPoolDraining } code, err := p.runOnce(ctx, stat.id, spec, sink) + // A run can observe a normal worker result just as Drain starts. Check + // again before classifying it so shutdown remains terminal rather than + // entering a retry path or reporting ErrPermanent. + if p.isDraining() { + return 0, ErrPoolDraining + } switch { case err != nil: lastErr = err if ctx.Err() != nil { return 0, ctx.Err() } + // Drain is terminal: do not backoff/retry, and do not wrap the + // shutdown error as ErrPermanent when attempts are exhausted. + if errors.Is(err, ErrPoolDraining) { + return 0, ErrPoolDraining + } p.logf("worker %d launch/run error: %v", stat.id, err) case code == 0: return 0, nil // clean success @@ -213,6 +226,9 @@ func (p *Pool) Run(ctx context.Context, spec WorkerSpec, sink Sink) (int, error) lastErr = fmt.Errorf("worker %d tempfail (code=%d)", stat.id, code) p.logf("worker %d tempfail — retry after %s", stat.id, p.opts.TempfailDelay) if !p.sleep(ctx, p.opts.TempfailDelay) { + if p.isDraining() { + return 0, ErrPoolDraining + } return 0, ctx.Err() } continue // tempfail retries do not count against the crash backoff @@ -227,6 +243,9 @@ func (p *Pool) Run(ctx context.Context, spec WorkerSpec, sink Sink) (int, error) delay := p.opts.Backoff(stat.restarts) p.logf("worker %d restart %d after backoff %s", stat.id, stat.restarts, delay) if !p.sleep(ctx, delay) { + if p.isDraining() { + return 0, ErrPoolDraining + } return 0, ctx.Err() } } @@ -239,11 +258,41 @@ func (p *Pool) Run(ctx context.Context, spec WorkerSpec, sink Sink) (int, error) // runOnce launches a single worker, pumps its output to sink, and returns its // exit code. The worker handle is tracked so Drain can kill it. func (p *Pool) runOnce(ctx context.Context, id int, spec WorkerSpec, sink Sink) (int, error) { + p.mu.Lock() + if p.draining { + p.mu.Unlock() + return 0, ErrPoolDraining + } + p.launching++ + p.mu.Unlock() + inLaunch := true + defer func() { + if inLaunch { + p.mu.Lock() + p.launching-- + p.mu.Unlock() + } + }() handle, err := p.opts.Launcher(ctx, spec) + p.mu.Lock() + draining := p.draining + if err == nil && !draining { + p.active[id] = handle + p.launching-- + inLaunch = false + } + p.mu.Unlock() if err != nil { + if draining { + return 0, ErrPoolDraining + } return 0, err } - p.track(id, handle) + if draining { + _ = handle.Kill() + _, _ = handle.Wait() + return 0, ErrPoolDraining + } defer p.untrack(id) // Pump stdout lines until the stream ends. @@ -281,15 +330,6 @@ func (p *Pool) newStat() *workerStat { return &workerStat{id: p.nextID} } -// track/untrack key the active set by the pool's monotonic worker id, not the OS -// pid: the OS can reuse a pid the instant a worker exits, so a pid key could collide -// a finished worker with a freshly-launched one and drop the wrong handle (D10). -func (p *Pool) track(id int, h WorkerHandle) { - p.mu.Lock() - p.active[id] = h - p.mu.Unlock() -} - func (p *Pool) untrack(id int) { p.mu.Lock() delete(p.active, id) @@ -315,25 +355,37 @@ func (p *Pool) sleep(ctx context.Context, d time.Duration) bool { return true case <-ctx.Done(): return false + case <-p.drained: + return false } } -// Drain stops accepting new work, gives in-flight workers a grace window -// (KillTimeout) to finish on their own, then force-kills any straggler. It -// returns as soon as the pool is idle (graceful) or the window elapses. Safe to -// call once; subsequent calls are no-ops. -func (p *Pool) Drain() { - p.drainOnce.Do(func() { +// beginDrain publishes the terminal pool state before callers cancel work that +// may be waiting in Run. Publishing this separately from Drain's bounded +// cleanup makes the shutdown result deterministic for every Run wakeup. +func (p *Pool) beginDrain() { + p.drainStartOnce.Do(func() { p.mu.Lock() p.draining = true p.mu.Unlock() close(p.drained) + }) +} +// Drain stops accepting new work, gives in-flight workers a grace window +// (KillTimeout) to finish on their own, then force-kills any straggler. It +// returns as soon as the pool is idle, the grace window elapses and existing +// workers are force-killed, or one separately bounded late-launch cleanup wait +// elapses. A launcher that ignores cancellation may finish its cleanup after +// Drain returns. Safe to call once; subsequent calls are no-ops. +func (p *Pool) Drain() { + p.beginDrain() + p.drainOnce.Do(func() { // Grace window: poll until idle or the deadline. deadline := time.Now().Add(p.opts.KillTimeout) for time.Now().Before(deadline) { p.mu.Lock() - n := len(p.active) + n := len(p.active) + p.launching p.mu.Unlock() if n == 0 { return // all workers drained gracefully @@ -352,6 +404,21 @@ func (p *Pool) Drain() { p.logf("drain: killing straggler worker pid=%d", h.Pid()) _ = h.Kill() } + + // Force-kill only covers handles already in active. A launcher still + // inside Launcher has no handle yet; wait one more KillTimeout for that + // late-launch path to finish kill+wait. Do not wait forever: a Launcher + // that ignores ctx would otherwise wedge shutdown. + deadline = time.Now().Add(p.opts.KillTimeout) + for time.Now().Before(deadline) { + p.mu.Lock() + n := p.launching + p.mu.Unlock() + if n == 0 { + return + } + time.Sleep(5 * time.Millisecond) + } }) } diff --git a/internal/daemon/pool_test.go b/internal/daemon/pool_test.go index 7ced2bef4..aa22818ba 100644 --- a/internal/daemon/pool_test.go +++ b/internal/daemon/pool_test.go @@ -3,6 +3,7 @@ package daemon import ( "context" "errors" + "strings" "sync" "sync/atomic" "testing" @@ -38,6 +39,10 @@ type fakeWorker struct { exitCode int killed int32 waitCh chan struct{} // when non-nil, Wait blocks until closed (drain tests) + killCh chan struct{} // when non-nil, Kill signals before any blocked Wait returns + // waitAfterKill keeps Wait blocked after Kill so drain tests can prove the + // pool waits for reaping rather than merely dispatching a kill signal. + waitAfterKill bool } func (w *fakeWorker) Stdout() Lines { return &fakeLines{lines: w.out, err: w.outErr} } @@ -49,7 +54,14 @@ func (w *fakeWorker) Wait() (int, error) { } func (w *fakeWorker) Kill() error { atomic.StoreInt32(&w.killed, 1) - if w.waitCh != nil { + if w.killCh != nil { + select { + case <-w.killCh: + default: + close(w.killCh) + } + } + if w.waitCh != nil && !w.waitAfterKill { select { case <-w.waitCh: default: @@ -224,7 +236,8 @@ func TestPoolDrainKillsStraggler(t *testing.T) { _, _ = pool.Run(context.Background(), WorkerSpec{Session: "a"}, &collectSink{}) close(runDone) }() - waitFor(t, func() bool { return pool.QueueDepth() == 1 }) + // Wait until the worker is tracked; Drain reads the active set, not slot occupancy. + waitFor(t, func() bool { return len(pool.WorkerStats()) == 1 }) pool.Drain() // KillTimeout elapses, straggler is force-killed if atomic.LoadInt32(&straggler.killed) != 1 { @@ -242,6 +255,182 @@ func TestPoolDrainKillsStraggler(t *testing.T) { } } +func TestPoolDrainKillsWorkerLaunchedAfterDrainStarts(t *testing.T) { + launchStarted := make(chan struct{}) + releaseLaunch := make(chan struct{}) + releaseWait := make(chan struct{}) + straggler := &fakeWorker{pid: 1, waitCh: releaseWait, killCh: make(chan struct{}), waitAfterKill: true} + pool, _ := NewPool(PoolOptions{Size: 1, MaxAttempts: 1, KillTimeout: 2 * time.Second, Backoff: func(int) time.Duration { return 0 }, Launcher: func(context.Context, WorkerSpec) (WorkerHandle, error) { + close(launchStarted) + <-releaseLaunch + return straggler, nil + }}) + runResult := make(chan error, 1) + go func() { + _, err := pool.Run(context.Background(), WorkerSpec{Session: "a"}, &collectSink{}) + runResult <- err + }() + select { + case <-launchStarted: + case <-time.After(2 * time.Second): + t.Fatal("Run did not start the launcher") + } + drained := make(chan struct{}) + go func() { pool.Drain(); close(drained) }() + waitFor(t, pool.isDraining) + select { + case <-drained: + t.Fatal("Drain returned while a launcher was still in progress") + default: + } + close(releaseLaunch) + select { + case <-straggler.killCh: + case <-time.After(2 * time.Second): + t.Fatal("Drain did not kill the worker launched after draining began") + } + select { + case <-drained: + t.Fatal("Drain returned before the late worker was reaped") + default: + } + close(releaseWait) + select { + case <-drained: + case <-time.After(2 * time.Second): + t.Fatal("Drain did not finish") + } + select { + case err := <-runResult: + if !errors.Is(err, ErrPoolDraining) { + t.Fatalf("Run error = %v, want ErrPoolDraining", err) + } + case <-time.After(2 * time.Second): + t.Fatal("Run did not finish after Drain") + } + if atomic.LoadInt32(&straggler.killed) != 1 { + t.Fatal("Drain must kill a worker whose launch completed after draining began") + } +} + +func TestPoolDrainBoundsBlockedLauncher(t *testing.T) { + const killTimeout = 100 * time.Millisecond + launchStarted := make(chan struct{}) + releaseLaunch := make(chan struct{}) + lateWorker := &fakeWorker{pid: 1} + pool, err := NewPool(PoolOptions{ + Size: 1, + MaxAttempts: 1, + KillTimeout: killTimeout, + Backoff: func(int) time.Duration { return 0 }, + Launcher: func(context.Context, WorkerSpec) (WorkerHandle, error) { + close(launchStarted) + <-releaseLaunch + return lateWorker, nil + }, + }) + if err != nil { + t.Fatalf("NewPool: %v", err) + } + runResult := make(chan error, 1) + go func() { + _, err := pool.Run(context.Background(), WorkerSpec{Session: "a"}, &collectSink{}) + runResult <- err + }() + select { + case <-launchStarted: + case <-time.After(2 * time.Second): + t.Fatal("Run did not start the launcher") + } + drainElapsed := make(chan time.Duration, 1) + go func() { + start := time.Now() + pool.Drain() + drainElapsed <- time.Since(start) + }() + // The first timeout accounts for the launch in progress. The second, separate + // timeout is what keeps Drain bounded once no handle exists to kill yet. + select { + case elapsed := <-drainElapsed: + if elapsed < 2*killTimeout { + t.Fatalf("Drain returned after %s, want at least both bounded waits (%s)", elapsed, 2*killTimeout) + } + case <-time.After(2 * time.Second): + t.Fatal("Drain did not return after the bounded blocked-launch wait") + } + close(releaseLaunch) + select { + case err := <-runResult: + if !errors.Is(err, ErrPoolDraining) { + t.Fatalf("Run error = %v, want ErrPoolDraining", err) + } + case <-time.After(2 * time.Second): + t.Fatal("Run did not finish after releasing the blocked launcher") + } + if atomic.LoadInt32(&lateWorker.killed) != 1 { + t.Fatal("late worker was not killed after the bounded Drain return") + } +} + +func TestPoolDrainInterruptsRetryDelays(t *testing.T) { + cases := []struct { + name string + exitCode int + maxAttempts int + }{ + {name: "backoff", exitCode: 1, maxAttempts: 2}, + {name: "tempfail", exitCode: ExitTempfail, maxAttempts: 1}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + delaying := make(chan struct{}) + pool, err := NewPool(PoolOptions{ + Size: 1, + MaxAttempts: tc.maxAttempts, + KillTimeout: time.Second, + TempfailDelay: time.Hour, + Backoff: func(int) time.Duration { + return time.Hour + }, + Log: func(message string) { + if strings.Contains(message, "retry after") || strings.Contains(message, "restart") { + select { + case <-delaying: + default: + close(delaying) + } + } + }, + Launcher: func(context.Context, WorkerSpec) (WorkerHandle, error) { + return &fakeWorker{pid: 1, exitCode: tc.exitCode}, nil + }, + }) + if err != nil { + t.Fatalf("NewPool: %v", err) + } + result := make(chan error, 1) + go func() { + _, err := pool.Run(context.Background(), WorkerSpec{Session: "a"}, &collectSink{}) + result <- err + }() + select { + case <-delaying: + case <-time.After(2 * time.Second): + t.Fatal("Run did not enter the retry delay") + } + pool.Drain() + select { + case err := <-result: + if !errors.Is(err, ErrPoolDraining) { + t.Fatalf("Run error = %v, want ErrPoolDraining", err) + } + case <-time.After(2 * time.Second): + t.Fatal("Run remained in retry delay after Drain") + } + }) + } +} + func waitFor(t *testing.T, cond func() bool) { t.Helper() deadline := time.Now().Add(2 * time.Second) diff --git a/internal/daemon/server.go b/internal/daemon/server.go index 21c44e0d6..24160354a 100644 --- a/internal/daemon/server.go +++ b/internal/daemon/server.go @@ -170,6 +170,10 @@ func (s *Server) untrackConn(c net.Conn) { func (s *Server) Shutdown() { s.shutdownOnce.Do(func() { close(s.done) + // Publish pool shutdown before cancelling session contexts. Otherwise a + // Run woken from a retry delay can observe context cancellation before + // Drain marks the pool terminal and incorrectly report context.Canceled. + s.opts.Pool.beginDrain() s.cancel() // stop in-flight pool runs s.mu.Lock() if s.listener != nil { diff --git a/internal/daemon/server_test.go b/internal/daemon/server_test.go index 0fa6c99f7..96c274969 100644 --- a/internal/daemon/server_test.go +++ b/internal/daemon/server_test.go @@ -1,9 +1,11 @@ package daemon import ( + "context" "errors" "os" "path/filepath" + "strings" "testing" "time" ) @@ -178,3 +180,72 @@ func TestServerRejectsUnknownCommand(t *testing.T) { t.Fatal("run with empty session id must return an error") } } + +func TestServerShutdownMakesRetryDelaysDrainTerminal(t *testing.T) { + cases := []struct { + name string + exitCode int + maxAttempts int + }{ + {name: "backoff", exitCode: 1, maxAttempts: 2}, + {name: "tempfail", exitCode: ExitTempfail, maxAttempts: 1}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + delaying := make(chan struct{}) + pool, err := NewPool(PoolOptions{ + Size: 1, + MaxAttempts: tc.maxAttempts, + KillTimeout: 20 * time.Millisecond, + TempfailDelay: time.Hour, + Backoff: func(int) time.Duration { return time.Hour }, + Log: func(message string) { + if strings.Contains(message, "retry after") || strings.Contains(message, "restart") { + select { + case <-delaying: + default: + close(delaying) + } + } + }, + Launcher: func(context.Context, WorkerSpec) (WorkerHandle, error) { + return &fakeWorker{pid: 1, exitCode: tc.exitCode}, nil + }, + }) + if err != nil { + t.Fatalf("NewPool: %v", err) + } + mgr, err := NewSessionManager(SessionManagerOptions{Pool: pool}) + if err != nil { + t.Fatalf("NewSessionManager: %v", err) + } + dir := t.TempDir() + srv, err := NewServer(ServerOptions{ + Paths: Paths{Socket: filepath.Join(dir, "d.sock"), Lock: filepath.Join(dir, "d.lock"), Status: filepath.Join(dir, "d.status")}, + Manager: mgr, + Pool: pool, + }) + if err != nil { + t.Fatalf("NewServer: %v", err) + } + sess, err := mgr.Start(srv.ctx, WorkerSpec{Session: "a"}) + if err != nil { + t.Fatalf("Start: %v", err) + } + select { + case <-delaying: + case <-time.After(2 * time.Second): + t.Fatal("session did not enter retry delay") + } + srv.Shutdown() + select { + case <-sess.Done(): + if !errors.Is(sess.Err(), ErrPoolDraining) { + t.Fatalf("session error = %v, want ErrPoolDraining", sess.Err()) + } + case <-time.After(2 * time.Second): + t.Fatal("session did not finish after shutdown") + } + }) + } +}