diff --git a/internal/daemon/pool.go b/internal/daemon/pool.go index 6de38315b..1ea586bc2 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,15 +199,28 @@ 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) + // Fully completed work remains successful even during the drain grace + // window. Failed or killed workers must still take the terminal path. + if err == nil && code == 0 { + return 0, nil + } + // A failed run can finish just as Drain starts. Recheck before retrying + // or classifying the failure as permanent. + 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 case code == ExitPermanent: p.logf("worker %d exited permanently (code=%d) — not retrying", stat.id, code) return code, ErrPermanent @@ -213,6 +228,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,9 +245,18 @@ 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() } } + // On the final tempfail attempt, sleep can select an expired timer while + // the drain channel is also ready. There is no next iteration to recheck + // shutdown, so do it before reporting attempt exhaustion. + if p.isDraining() { + return 0, ErrPoolDraining + } if lastErr == nil { lastErr = ErrPermanent } @@ -239,11 +266,43 @@ 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 { + // Key by the monotonic worker id, not the reusable OS pid, so an older + // worker's untrack cannot remove a newly launched handle (D10). + 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 +340,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 +365,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 +414,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 2270bbf49..b3d41e90d 100644 --- a/internal/daemon/pool_test.go +++ b/internal/daemon/pool_test.go @@ -3,6 +3,8 @@ package daemon import ( "context" "errors" + "fmt" + "strings" "sync" "sync/atomic" "testing" @@ -40,6 +42,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} } @@ -51,7 +57,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: @@ -226,7 +239,8 @@ func TestPoolDrainKillsStraggler(t *testing.T) { _, _ = pool.Run(context.Background(), WorkerSpec{Session: "a"}, &collectSink{}) close(runDone) }() - testutil.WaitFor(t, "queue depth 1", func() bool { return pool.QueueDepth() == 1 }) + // Wait until the worker is tracked; Drain reads the active set, not slot occupancy. + testutil.WaitFor(t, "worker tracked", func() bool { return len(pool.WorkerStats()) == 1 }) pool.Drain() // KillTimeout elapses, straggler is force-killed if atomic.LoadInt32(&straggler.killed) != 1 { @@ -244,6 +258,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) }() + testutil.WaitFor(t, "pool draining", 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 TestPoolRunSurfacesStdoutReadError(t *testing.T) { // A worker stdout read error must surface as a failure (not be swallowed and // reported as clean success), and the worker is killed so it can't block on an @@ -261,3 +451,69 @@ func TestPoolRunSurfacesStdoutReadError(t *testing.T) { t.Error("the worker should be killed on a read error to avoid a pipe-block hang") } } + +// Drain must preserve work that finished cleanly, while failed or killed workers +// still report shutdown rather than a retryable or permanent worker failure. +func TestPoolDrainClassifiesCompletedWorker(t *testing.T) { + for _, code := range []int{0, -1, ExitPermanent} { + t.Run(fmt.Sprint(code), func(t *testing.T) { + releaseWait := make(chan struct{}) + release := sync.OnceFunc(func() { close(releaseWait) }) + worker := &fakeWorker{pid: 1, exitCode: code, waitCh: releaseWait} + launcher, _ := seqLauncher(worker) + pool, err := NewPool(PoolOptions{Size: 1, Launcher: launcher}) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { release(); pool.Drain() }) + result := make(chan error, 1) + go func() { + _, err := pool.Run(context.Background(), WorkerSpec{Session: "a"}, nil) + result <- err + }() + testutil.WaitFor(t, "worker tracked", func() bool { return len(pool.WorkerStats()) == 1 }) + pool.beginDrain() + release() + select { + case err := <-result: + if code == 0 && err != nil { + t.Fatalf("completed worker error = %v, want nil", err) + } + if code != 0 && !errors.Is(err, ErrPoolDraining) { + t.Fatalf("failed worker error = %v, want ErrPoolDraining", err) + } + case <-time.After(2 * time.Second): + t.Fatal("Run did not finish after worker exit") + } + }) + } +} + +func TestPoolDrainDuringFinalTempfailDelay(t *testing.T) { + // An expired timer and the drain channel may both be ready. Exercise both + // select outcomes: exhausting the last attempt must not turn drain permanent. + for iteration := 0; iteration < 100; iteration++ { + var pool *Pool + launcher, calls := seqLauncher(&fakeWorker{pid: 1, exitCode: ExitTempfail}) + var err error + pool, err = NewPool(PoolOptions{ + Size: 1, Launcher: launcher, MaxAttempts: 1, TempfailDelay: time.Nanosecond, + Log: func(message string) { + if strings.Contains(message, "retry after") { + pool.beginDrain() + } + }, + }) + if err != nil { + t.Fatal(err) + } + _, err = pool.Run(context.Background(), WorkerSpec{Session: "a"}, nil) + pool.Drain() + if !errors.Is(err, ErrPoolDraining) { + t.Fatalf("iteration %d: Run error = %v, want ErrPoolDraining", iteration, err) + } + if atomic.LoadInt32(calls) != 1 { + t.Fatal("Run retried after drain began") + } + } +} diff --git a/internal/daemon/server.go b/internal/daemon/server.go index 34fb2e9f6..b91e4bca4 100644 --- a/internal/daemon/server.go +++ b/internal/daemon/server.go @@ -257,6 +257,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 77a23987d..d0bd9e3c0 100644 --- a/internal/daemon/server_test.go +++ b/internal/daemon/server_test.go @@ -1,12 +1,14 @@ package daemon import ( + "context" "errors" "net" "os" "path/filepath" "runtime" "strings" + "sync" "testing" "time" @@ -496,3 +498,104 @@ 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) + } + closeEntered := make(chan struct{}) + releaseClose := make(chan struct{}) + release := sync.OnceFunc(func() { close(releaseClose) }) + local, peer := net.Pipe() + srv.trackConn(&shutdownCloseGate{Conn: local, entered: closeEntered, release: releaseClose}) + t.Cleanup(func() { release(); srv.Shutdown(); _ = local.Close(); _ = peer.Close() }) + 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") + } + shutdownDone := make(chan struct{}) + go func() { srv.Shutdown(); close(shutdownDone) }() + select { + case <-closeEntered: + case <-time.After(2 * time.Second): + t.Fatal("Shutdown did not reach connection cleanup") + } + 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") + } + release() + select { + case <-shutdownDone: + case <-time.After(2 * time.Second): + t.Fatal("Shutdown did not finish after connection cleanup") + } + }) + } +} + +// shutdownCloseGate parks Shutdown after cancellation but before Pool.Drain. +// Embedding a pipe preserves net.Conn behavior outside that ordering boundary. +type shutdownCloseGate struct { + net.Conn + entered chan struct{} + release chan struct{} +} + +func (c *shutdownCloseGate) Close() error { + close(c.entered) + <-c.release + return c.Conn.Close() +}