fix: 🐛 Stop live and draining pipelines on shutdown

cmd/cortex-envoy/main.go

@@ -355,8 +355,7 @@ func main() {
 		sessionAPISrv.Shutdown(shutdownCtx)
 	}
 
-	outboundPipeline.Stop(shutdownCtx)
-	inboundPipeline.Stop(shutdownCtx)
+	rld.Close(shutdownCtx)
 
 	if sessions != nil {
 		sessions.Close()

cmd/cortex/main.go

@@ -1285,11 +1285,11 @@ func main() {
 			_ = sessionAPISrv.Server().Close()
 		}
 	}
-	outboundPipeline.Stop(shutdownCtx)
-	inboundPipeline.Stop(shutdownCtx)
+	rld.Close(shutdownCtx)
 
-	// Flushed AFTER both pipelines have stopped and before the store closes, so the
-	// final minute includes every event a draining request still produced. Any earlier
+	// Flushed AFTER rld.Close (every pipeline stopped, or the shutdown budget spent)
+	// and before the store closes, so the final minute includes every event a
+	// draining request still produced. Any earlier
 	// and a request finishing during the pipeline drain would land in a minute already
 	// written; any later and the store is gone.
 	//

core/reloader/reloader.go

@@ -13,8 +13,8 @@
 // builds fresh pipelines via the caller-supplied PipelineBuilder,
 // runs their Start hooks under a bounded context, and Stores them
 // into the Holders. Each old pipeline keeps running until no request
-// holds it, or the drain window expires, then Stop is invoked in a
-// background goroutine.
+// holds it, the drain window expires, or Close is called, then Stop is
+// invoked in a background goroutine. Close also Stops the live pipelines.
 //
 // Validation failure at any stage leaves the current pipelines
 // untouched and records the error in Status so operators can see
@@ -135,6 +135,17 @@ type Reloader struct {
 	lastHash string
 	// startData is the config file as New read it.
 	startData []byte
+
+	// cancelLoop and loopDone stop and await the watch goroutine; set by Start.
+	cancelLoop context.CancelFunc
+	loopDone   chan struct{}
+	// Close sets closeCtx, then closes closing, to stop every old pipeline still
+	// draining on closeCtx; drains counts their goroutines. commitMu orders a
+	// reload's commit against Close.
+	closing  chan struct{}
+	closeCtx context.Context
+	drains   sync.WaitGroup
+	commitMu sync.Mutex
 }
 
 // New constructs a Reloader. initialCfg must be the config the caller
@@ -149,6 +160,7 @@ func New(configPath string, inbound, outbound *pipeline.Holder, build PipelineBu
 		drainWindow:  10 * time.Minute,
 		debounce:     250 * time.Millisecond,
 		startTimeout: 60 * time.Second,
+		closing:      make(chan struct{}),
 	}
 	for _, opt := range opts {
 		opt(r)
@@ -170,7 +182,7 @@ func New(configPath string, inbound, outbound *pipeline.Holder, build PipelineBu
 
 // Start arms the watcher. Blocks until the watcher is ready or fails
 // to arm (invalid directory, fsnotify init error). Returns nil on
-// success; the watch goroutine then runs until ctx is cancelled.
+// success; the watch goroutine then runs until ctx is cancelled or Close.
 //
 // Safe to call at most once. A Reloader whose Start returned an error
 // is not usable — construct a new one.
@@ -240,7 +252,12 @@ func (r *Reloader) Start(ctx context.Context) error {
 		r.seedHash()
 	}
 
-	go r.watchLoop(ctx, watcher)
+	ctx, r.cancelLoop = context.WithCancel(ctx)
+	r.loopDone = make(chan struct{})
+	go func() {
+		defer close(r.loopDone)
+		r.watchLoop(ctx, watcher)
+	}()
 	slog.Info("reloader: watching for config changes",
 		"path", r.configPath, "dir", dir, "drainWindow", r.drainWindow)
 	return nil
@@ -405,7 +422,7 @@ func (r *Reloader) reloadOnce(parent context.Context, watcher *fsnotify.Watcher)
 	}
 	if err := newOut.Start(startCtx); err != nil {
 		// Inbound started cleanly; unwind it before bailing.
-		unwindCtx, unwindCancel := context.WithTimeout(context.Background(), 15*time.Second)
+		unwindCtx, unwindCancel := r.stopCtx()
 		defer unwindCancel()
 		newIn.Stop(unwindCtx)
 		r.recordFailure(attempt, fmt.Errorf("outbound Start: %w", err))
@@ -415,12 +432,23 @@ func (r *Reloader) reloadOnce(parent context.Context, watcher *fsnotify.Watcher)
 	// Commit: swap Holders and update active-config snapshot. From
 	// here on, requests drain onto the new pipelines; old in-flight
 	// requests keep their existing *Pipeline reference.
+	r.commitMu.Lock()
+	select {
+	case <-r.closing:
+		r.commitMu.Unlock()
+		newOut.Stop(r.closeCtx)
+		newIn.Stop(r.closeCtx)
+		return
+	default:
+	}
 	oldIn := r.inbound.Load()
 	oldOut := r.outbound.Load()
 	r.inbound.Store(newIn)
 	r.outbound.Store(newOut)
 	r.activeCfg.Store(newCfg)
 	r.lastHash = hash
+	r.drains.Add(1)
+	r.commitMu.Unlock()
 
 	// After the swap, so a hook cannot observe a config that was refused, and before the log line, so
 	// "pipelines swapped" means everything this reload applies has been applied.
@@ -432,16 +460,16 @@ func (r *Reloader) reloadOnce(parent context.Context, watcher *fsnotify.Watcher)
 		"sha256", hash[:12],
 		"drainWindow", r.drainWindow)
 
-	// Drain old pipelines in the background. Using context.Background
-	// deliberately: parent may cancel during shutdown, and we still
-	// want a bounded Stop on the old plugins so they flush cleanly.
+	// Drain old pipelines in the background, on their own Stop budget rather than
+	// parent's; Close cuts the drain short at shutdown.
 	go func() {
+		defer r.drains.Done()
 		deadline := time.Now().Add(r.drainWindow)
 		var remaining atomic.Int64
 		var wg sync.WaitGroup
 		for _, p := range []*pipeline.Pipeline{oldIn, oldOut} {
 			if p != nil {
-				wg.Go(func() { remaining.Add(stopWhenUnheld(p, deadline)) })
+				wg.Go(func() { remaining.Add(r.stopWhenUnheld(p, deadline)) })
 			}
 		}
 		wg.Wait()
@@ -461,23 +489,67 @@ func (r *Reloader) reloadOnce(parent context.Context, watcher *fsnotify.Watcher)
 	r.recordSuccess(attempt, hash)
 }
 
-// stopWhenUnheld Stops p once no request holds it, or at deadline, and returns how
-// many requests still held it then.
-func stopWhenUnheld(p *pipeline.Pipeline, deadline time.Time) int64 {
+// stopWhenUnheld Stops p once no request holds it, at deadline, or on Close, and
+// returns how many requests still held it then.
+func (r *Reloader) stopWhenUnheld(p *pipeline.Pipeline, deadline time.Time) int64 {
 	timer := time.NewTimer(time.Until(deadline))
 	defer timer.Stop()
 	var n int64
 	select {
 	case <-p.Drained():
 	case <-timer.C:
 		n = p.Held()
+	case <-r.closing:
+		n = p.Held()
 	}
-	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
+	ctx, cancel := r.stopCtx()
 	defer cancel()
 	p.Stop(ctx)
 	return n
 }
 
+// stopCtx is the budget for a Stop: Close's ctx once Close has begun, else 15s.
+func (r *Reloader) stopCtx() (context.Context, context.CancelFunc) {
+	select {
+	case <-r.closing:
+		return r.closeCtx, func() {}
+	default:
+		return context.WithTimeout(context.Background(), 15*time.Second)
+	}
+}
+
+// Close stops reloading, then Stops the live pipelines and every old one still
+// draining; a reload still building Stops its pipelines instead. It returns when
+// they are stopped or ctx is done, whichever is first; a plugin still shutting
+// down then finishes in the background. Call it once, after the listeners have
+// stopped.
+func (r *Reloader) Close(ctx context.Context) {
+	r.closeCtx = ctx
+	r.commitMu.Lock()
+	close(r.closing)
+	liveIn, liveOut := r.inbound.Load(), r.outbound.Load()
+	r.commitMu.Unlock()
+	if r.cancelLoop != nil {
+		r.cancelLoop()
+	}
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		var live sync.WaitGroup
+		live.Go(func() { liveOut.Stop(ctx) })
+		live.Go(func() { liveIn.Stop(ctx) })
+		live.Wait()
+		r.drains.Wait()
+		if r.loopDone != nil {
+			<-r.loopDone
+		}
+	}()
+	select {
+	case <-done:
+	case <-ctx.Done():
+	}
+}
+
 // hash is what decides whether a reload has anything to do: the config's bytes,
 // and, with WithExtraWatch, the fingerprint of everything else the pipelines are
 // built from. While the fingerprint is empty it is the config's sha256, as it

core/reloader/reloader_test.go

@@ -31,8 +31,12 @@ type fakeBuilder struct {
 	// `next` at call time.
 	next atomic.Value // holds builderResult
 	cfg  atomic.Value // holds *config.Config
+	// gate, when set, makes build signal entered and wait for release.
+	gate atomic.Pointer[buildGate]
 }
 
+type buildGate struct{ entered, release chan struct{} }
+
 type builderResult struct {
 	inbound  *pipeline.Pipeline
 	outbound *pipeline.Pipeline
@@ -48,6 +52,10 @@ func (b *fakeBuilder) set(r builderResult) {
 }
 
 func (b *fakeBuilder) build() (*pipeline.Pipeline, *pipeline.Pipeline, *config.Config, error) {
+	if g := b.gate.Load(); g != nil {
+		close(g.entered)
+		<-g.release
+	}
 	r := b.next.Load().(builderResult)
 	return r.inbound, r.outbound, r.cfg, r.err
 }
@@ -971,11 +979,15 @@ func TestValidateReloadable_RefusesMemoryEdit(t *testing.T) {
 	}
 }
 
-// shutdownPlugin records how many requests held its pipeline when it was stopped.
+// shutdownPlugin records how many requests held its pipeline when it was stopped
+// and whether its ctx was already done, and with block set, stays in Shutdown
+// until block is closed.
 type shutdownPlugin struct {
 	p       *pipeline.Pipeline
 	held    int64
+	ctxErr  error
 	stopped chan struct{}
+	block   chan struct{}
 }
 
 func (*shutdownPlugin) Name() string { return "shutdown" }
@@ -988,16 +1000,19 @@ func (*shutdownPlugin) OnRequest(context.Context, *pipeline.Context) pipeline.Ac
 func (*shutdownPlugin) OnResponse(context.Context, *pipeline.Context) pipeline.Action {
 	return pipeline.Action{Type: pipeline.Continue}
 }
-func (s *shutdownPlugin) Shutdown(context.Context) error {
-	s.held = s.p.Held()
+func (s *shutdownPlugin) Shutdown(ctx context.Context) error {
+	s.held, s.ctxErr = s.p.Held(), ctx.Err()
 	close(s.stopped)
+	if s.block != nil {
+		<-s.block
+	}
 	return nil
 }
 
 // storeShutdownPipeline Stores a pipeline of one shutdownPlugin in h.
-func storeShutdownPipeline(t *testing.T, h *pipeline.Holder) *shutdownPlugin {
+func storeShutdownPipeline(t *testing.T, h *pipeline.Holder, block chan struct{}) *shutdownPlugin {
 	t.Helper()
-	sp := &shutdownPlugin{stopped: make(chan struct{})}
+	sp := &shutdownPlugin{stopped: make(chan struct{}), block: block}
 	p, err := pipeline.New([]pipeline.Plugin{sp})
 	if err != nil {
 		t.Fatalf("pipeline.New: %v", err)
@@ -1055,7 +1070,7 @@ func TestReloader_StopsAnOldPipelineWhenNoRequestHoldsIt(t *testing.T) {
 			committed := make(chan struct{}, 1)
 			_, b, cfgPath, inH, outH := setupWith(t, WithDrainWindow(tc.window),
 				WithOnCommit(func(*config.Config) { committed <- struct{}{} }))
-			idle, busy := storeShutdownPipeline(t, inH), storeShutdownPipeline(t, outH)
+			idle, busy := storeShutdownPipeline(t, inH, nil), storeShutdownPipeline(t, outH, nil)
 
 			pctx := &pipeline.Context{Direction: pipeline.Outbound}
 			outH.Run(context.Background(), pctx)
@@ -1095,6 +1110,93 @@ func TestReloader_StopsAnOldPipelineWhenNoRequestHoldsIt(t *testing.T) {
 	}
 }
 
+// Close stops the live pipelines and an old one still draining under an open
+// request, each once, before it returns.
+func TestReloader_CloseStopsLiveAndDrainingPipelines(t *testing.T) {
+	committed := make(chan struct{}, 1)
+	r, b, cfgPath, _, outH := setupWith(t, WithDrainWindow(time.Hour),
+		WithOnCommit(func(*config.Config) { committed <- struct{}{} }))
+	draining := storeShutdownPipeline(t, outH, nil)
+	outH.Run(context.Background(), &pipeline.Context{Direction: pipeline.Outbound})
+
+	liveIn, liveOut := pipeline.NewHolder(emptyPipeline(t)), pipeline.NewHolder(emptyPipeline(t))
+	in, out := storeShutdownPipeline(t, liveIn, nil), storeShutdownPipeline(t, liveOut, nil)
+	b.set(builderResult{inbound: liveIn.Load(), outbound: liveOut.Load(), cfg: &config.Config{Mode: "envoy-sidecar"}})
+	writeConfig(t, cfgPath, "mode: envoy-sidecar\n# reloaded\n")
+	<-committed
+
+	// A second Stop of any of them would panic closing its stopped channel.
+	r.Close(context.Background())
+	for name, sp := range map[string]*shutdownPlugin{"inbound": in, "outbound": out, "draining": draining} {
+		select {
+		case <-sp.stopped:
+		default:
+			t.Errorf("the %s pipeline was not stopped by Close", name)
+		}
+	}
+	if draining.held != 1 {
+		t.Errorf("the draining pipeline stopped with %d requests holding it, want 1", draining.held)
+	}
+}
+
+// Close returns by its ctx even while a live pipeline's plugin is still shutting
+// down, and that plugin does not keep the other live pipeline from stopping.
+func TestReloader_CloseDoesNotOutwaitItsContext(t *testing.T) {
+	r, _, _, inH, outH := setup(t)
+	block := make(chan struct{})
+	t.Cleanup(func() { close(block) })
+	storeShutdownPipeline(t, outH, block)
+	in := storeShutdownPipeline(t, inH, nil)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	r.Close(ctx)
+	<-in.stopped
+}
+
+// Close stops the live pipelines without waiting for a reload still building, and
+// that reload stops its own pipelines instead of committing them.
+func TestReloader_CloseDuringABuildStopsLiveAndBuiltPipelines(t *testing.T) {
+	r, b, cfgPath, inH, outH := setup(t)
+	live := storeShutdownPipeline(t, outH, nil)
+	g := &buildGate{entered: make(chan struct{}), release: make(chan struct{})}
+	b.gate.Store(g)
+	builtIn := storeShutdownPipeline(t, pipeline.NewHolder(emptyPipeline(t)), nil)
+	builtOut := storeShutdownPipeline(t, pipeline.NewHolder(emptyPipeline(t)), nil)
+	b.set(builderResult{inbound: builtIn.p, outbound: builtOut.p, cfg: &config.Config{Mode: "envoy-sidecar"}})
+	writeConfig(t, cfgPath, "mode: envoy-sidecar\n# building\n")
+	<-g.entered
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	r.Close(ctx)
+	<-live.stopped // while the build is still blocked
+
+	close(g.release)
+	<-builtIn.stopped
+	<-builtOut.stopped
+	if inH.Load() == builtIn.p || outH.Load() == builtOut.p {
+		t.Error("a reload committed its pipelines after Close began")
+	}
+}
+
+// Once Close has begun, an old pipeline is stopped on Close's ctx, even when its
+// drain is what woke it.
+func TestReloader_StopsOnTheCloseContextOnceClosing(t *testing.T) {
+	r, _, _, _, _ := setup(t)
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	r.Close(ctx)
+	// Repeated because select picks among the ready drain and close cases at random.
+	for range 20 {
+		sp := storeShutdownPipeline(t, pipeline.NewHolder(emptyPipeline(t)), nil)
+		r.stopWhenUnheld(sp.p, time.Now().Add(time.Hour))
+		if sp.ctxErr == nil {
+			t.Fatal("stopped on a fresh budget after Close began")
+		}
+	}
+}
+
 // Without WithDrainWindow, an old pipeline waits up to 10 minutes for its requests.
 func TestNew_DrainWindowDefaultsToTenMinutes(t *testing.T) {
 	r := New("config.yaml", pipeline.NewHolder(emptyPipeline(t)), pipeline.NewHolder(emptyPipeline(t)),

docs/framework-architecture.md

@@ -460,8 +460,7 @@ if err := outboundPipeline.Start(initCtx); err != nil {
 // ... serve traffic ...
 
 // After listeners have drained (a pod) or closed (a local install) on SIGTERM:
-outboundPipeline.Stop(shutdownCtx) // reverse order within each pipeline
-inboundPipeline.Stop(shutdownCtx)
+rld.Close(shutdownCtx) // the live pipelines, and any a reload is still draining
 ```
 
 Semantics:
@@ -853,7 +852,7 @@ curl http://localhost:9093/config               # now-active config
 5. Compare the new config to the active one on unreloadable fields (`Mode`, `Listener.*`, `Session.*`, `CostLedger.*`, `MTLS.*`); refuse if they differ (see below).
 6. `Start` the new pipelines with a 60 s budget. On Start failure, `Stop` any partially-started pipelines so their goroutines don't leak.
 7. `inboundH.Store(newIn)` + `outboundH.Store(newOut)` — new requests now route to the new pipelines.
-8. Background goroutine: each old pipeline waits until no request holds it (`Pipeline.Drained`; a request holds from `Run` to `RunFinish`) or the drain window (default 10 min) expires → `oldPipeline.Stop(ctx)` with a 15 s budget. In-flight requests pinned to an old pipeline finish against it; a WARN names how many were still open when the window stopped it. A CONNECT tunnel or WebSocket holds its pipeline until the connection closes, so with a client that keeps one open, the full window and the WARN are the usual outcome, and rapid reloads keep several old pipelines' plugins running until their windows end.
+8. Background goroutine: each old pipeline waits until no request holds it (`Pipeline.Drained`; a request holds from `Run` to `RunFinish`) or the drain window (default 10 min) expires → `oldPipeline.Stop(ctx)` with a 15 s budget. In-flight requests pinned to an old pipeline finish against it; a WARN names how many were still open when the window stopped it. A CONNECT tunnel or WebSocket holds its pipeline until the connection closes, so with a client that keeps one open, the full window and the WARN are the usual outcome, and rapid reloads keep several old pipelines' plugins running until their windows end. `Reloader.Close` at shutdown stops them at once.
 9. `Status.LastSuccess`, `ReloadsOK`, `ActiveConfigSHA256` update atomically.
 
 **What's reloadable, what isn't**