From 84bd87dcc989b7aae44d7ae5b0a898f6546438ec Mon Sep 17 00:00:00 2001 From: zhangwenjian Date: Mon, 7 Sep 2026 17:10:13 +0800 Subject: [PATCH] =?UTF-8?q?fix=F0=9F=90=9B:=20drain=20the=20queue=20on=20t?= =?UTF-8?q?he=20way=20out=20instead=20of=20abandoning=20it?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Nothing stopped the queue when the process exited. core v2.7.0 made the drain work - Memory.Shutdown closes the queue and waits for every consumer to finish what it holds, and the legacy adapter cancels its context and closes the underlying queue - but no call site ever reached it. The only Shutdown() in this repository applies to the previous adapter during a reload, so the installed one was simply left. The login log, the operation log and the API sync all publish through it, so a rolling restart dropped whatever had not been consumed, on the path where the process exits 0 and reports "Server exiting". Setup now registers a BeforeExit callback that shuts down the adapter this package installed. Three things it has to get right, each with a test. It reads `installed` when it runs, not when it registers. A reload replaces the adapter, and the one from start-up is a queue nobody has published to since. It never goes through sdk.Runtime.GetQueueAdapter. That accessor never returns nil - with no queue section configured it wraps the runtime's own fallback - so it would look like it worked while closing a queue this package neither built nor started. That is the same trap setupQueue already had to drop an `if q != nil` for. It registers once. Setup is re-run on every configuration change, and a callback per reload would leave the shutdown phase holding a row of identical entries, each eligible to be named as the one that overran the budget. That last one needed a seam. shutdownQueue takes the adapter on its first run, so the second and third callbacks find nothing and return - three registrations produce exactly the same observable result as one, and a test going through the effect passes either way. It did: the counter-proof for "register on every reload" came back green until the registration was counted at the seam instead. The wait is bounded here rather than left to the phase. Shutdown takes no context, so a consumer that never finishes would hold the process until SIGKILL; the callback gives up and says what is being lost, which the phase's generic overrun message cannot. Ordering falls out of the phase rather than being arranged: callbacks run in reverse registration order, this one registers during setup and the job scheduler's registers on AfterListen, so the schedulers stop before the queue drains. Verified against core v2.7.0 rather than read off the source. Closes #911. --- common/storage/initialize.go | 87 ++++++++++ common/storage/queue_drain_test.go | 267 +++++++++++++++++++++++++++++ 2 files changed, 354 insertions(+) create mode 100644 common/storage/queue_drain_test.go diff --git a/common/storage/initialize.go b/common/storage/initialize.go index 305fc780..24d04bbf 100644 --- a/common/storage/initialize.go +++ b/common/storage/initialize.go @@ -8,10 +8,12 @@ package storage import ( + "context" "log" "sync" "github.com/go-admin-team/go-admin-core/v2/captcha" + corelog "github.com/go-admin-team/go-admin-core/v2/logger" "github.com/go-admin-team/go-admin-core/v2/sdk" "github.com/go-admin-team/go-admin-core/v2/sdk/config" ) @@ -21,6 +23,7 @@ func Setup() { setupCache() setupCaptcha() setupQueue() + registerQueueDrain() } func setupCache() { @@ -41,8 +44,92 @@ var ( // shut it down, and counted so a consumer can tell one from the next. installed interface{ Shutdown() } installedGen uint64 + // drainRegistered records that the BeforeExit callback is on the runtime, + // so that a reload does not add another one. + drainRegistered bool ) +// setShutdown is sdk.Runtime.SetShutdown, indirected so that registering can +// be observed. +// +// It has to be: the runtime does not report how many callbacks a phase holds, +// and shutdownQueue takes the adapter on its first run, so every registration +// after the first returns immediately and changes nothing anybody can see. A +// reload adding one callback per round would therefore be invisible from the +// outside - which is exactly how it would survive. +var setShutdown = func(f func(context.Context)) { sdk.Runtime.SetShutdown(f) } + +// registerQueueDrain puts shutdownQueue on the BeforeExit phase, once. +// +// Setup is one of the callbacks bootstrap.SetupConfig re-runs on every +// configuration change, so registering from it without a guard would leave one +// callback per reload - each shutting down the same adapter, each reported +// separately when the budget runs out. +// +// A flag under the existing mutex rather than a sync.Once: the tests in this +// package already save and restore installed and installedGen to keep one test +// from deciding what the next one sees, and a sync.Once cannot be put back. +func registerQueueDrain() { + queueMu.Lock() + first := !drainRegistered + drainRegistered = true + queueMu.Unlock() + + if first { + setShutdown(shutdownQueue) + } +} + +// shutdownQueue drains the queue this package installed, on the way out. +// +// Nothing used to. core's Memory.Shutdown closes the queue and waits for every +// consumer to finish what it is holding, and the legacy adapter cancels its +// context and closes the underlying queue - but neither ran at exit, so the +// process left with the login log, the operation log and the API sync still +// buffered, and left reporting success. +// +// The adapter is read here rather than captured at registration because a +// reload replaces it. Registration happens once per process; this runs against +// whatever is current when the signal arrives. +// +// Only an adapter this package installed. sdk.Runtime.GetQueueAdapter never +// returns nil - with no queue section configured the runtime wraps its own +// fallback queue - so going through that accessor would shut down a queue this +// package neither built nor started. +// +// The adapter is taken, not read: after this the package owns nothing, so a +// reload arriving mid-shutdown builds a new one instead of being handed a +// closed one to shut down again. Both implementations tolerate a second +// Shutdown, so this is about who owns it rather than about a crash. +func shutdownQueue(ctx context.Context) { + queueMu.Lock() + q := installed + installed = nil + queueMu.Unlock() + + if q == nil { + return + } + + done := make(chan struct{}) + go func() { + defer close(done) + q.Shutdown() + }() + + select { + case <-done: + corelog.Info("queue: drained") + case <-ctx.Done(): + // The wait is what the budget bounds, not the work: Shutdown takes no + // context and is still running in that goroutine. Saying so here names + // what is being lost, which the generic overrun message cannot. + corelog.Warnf("queue: the shutdown budget ran out while the queue was still draining - " + + "whatever it had not delivered goes with the process. Raise extend.shutdown.cleanup " + + "if this recurs.") + } +} + // QueueGeneration reports how many times this package has installed a queue // adapter. It changes every time setupQueue builds a new one, which is on // every configuration reload, and stays 0 for as long as the configuration has diff --git a/common/storage/queue_drain_test.go b/common/storage/queue_drain_test.go new file mode 100644 index 00000000..12b8ddf9 --- /dev/null +++ b/common/storage/queue_drain_test.go @@ -0,0 +1,267 @@ +package storage + +import ( + "context" + "sync/atomic" + "testing" + "time" + + "github.com/go-admin-team/go-admin-core/v2/sdk" + "github.com/go-admin-team/go-admin-core/v2/sdk/config" + "github.com/go-admin-team/go-admin-core/v2/sdk/runtime" + corestorage "github.com/go-admin-team/go-admin-core/v2/storage" +) + +// countingQueue stands in for an installed adapter. Only Shutdown is exercised +// - the drain callback never publishes or consumes - so the rest of +// AdapterQueue is deliberately absent: `installed` is typed on Shutdown alone, +// and widening the fake would only invite it to be used for something else. +type countingQueue struct { + calls atomic.Int32 + block chan struct{} +} + +func (q *countingQueue) Shutdown() { + q.calls.Add(1) + if q.block != nil { + <-q.block + } +} + +// isolate gives the test its own runtime and its own view of what this package +// has installed, and puts the process-wide state back afterwards. +// +// Same isolation as TestSetupBumpsTheQueueGenerationOnEveryReload, plus +// drainRegistered: it is what stops a reload registering a second callback, so +// leaving it set would make every later test in this binary see a package that +// has already registered. +func isolate(t *testing.T) { + t.Helper() + + prevQ, prevC := config.QueueConfig, config.CacheConfig + prevRuntime := sdk.Runtime + queueMu.Lock() + prevInstalled, prevGen, prevRegistered := installed, installedGen, drainRegistered + queueMu.Unlock() + + t.Cleanup(func() { + config.QueueConfig, config.CacheConfig = prevQ, prevC + sdk.Runtime = prevRuntime + queueMu.Lock() + installed, installedGen, drainRegistered = prevInstalled, prevGen, prevRegistered + queueMu.Unlock() + }) + + sdk.Runtime = runtime.NewConfig() + queueMu.Lock() + installed, drainRegistered = nil, false + queueMu.Unlock() +} + +// setInstalled puts a fake where setupQueue would have left the real adapter. +// +// Legitimate because the callback reads `installed` when it runs rather than +// capturing it at registration - that is the property that lets a reload +// replace the adapter and still have the right one drained. +func setInstalled(q interface{ Shutdown() }) { + queueMu.Lock() + installed = q + queueMu.Unlock() +} + +func currentInstalled() interface{ Shutdown() } { + queueMu.Lock() + defer queueMu.Unlock() + return installed +} + +// Issue #911: nothing shut the queue down at exit, so whatever was buffered +// went with the process. +// +// This is the half that matters most - a callback is on BeforeExit and it +// reaches the adapter this package installed. It says nothing about how many +// times the callback was registered; see the test below for that. +func TestSetupPutsTheQueueDrainOnBeforeExit(t *testing.T) { + isolate(t) + config.CacheConfig = &config.Cache{Memory: struct{}{}} + config.QueueConfig = &config.Queue{Memory: &config.QueueMemory{PoolSize: 10}} + + Setup() + Setup() + Setup() + + q := &countingQueue{} + setInstalled(q) + + if err := sdk.Runtime.RunShutdown(context.Background()); err != nil { + t.Fatalf("RunShutdown: %v", err) + } + + if got := q.calls.Load(); got != 1 { + t.Errorf("Shutdown called %d times, want 1 - 0 means nothing registered the drain", got) + } +} + +// Setup is re-run on every configuration change, so registering from it has to +// be guarded: a callback per reload would leave the shutdown phase holding a +// row of identical entries, each timed and each eligible to be named as the one +// that overran the budget. +// +// Counted at the seam rather than through the effect. The test above cannot see +// this - shutdownQueue takes the adapter on its first run, so the second and +// third callbacks find nothing and return, and three registrations produce +// exactly the same observable result as one. That is a good property of the +// callback and a blind spot for any test that goes through it. +func TestSetupRegistersTheDrainOncePerProcessHoweverManyReloads(t *testing.T) { + isolate(t) + + previous := setShutdown + t.Cleanup(func() { setShutdown = previous }) + registrations := 0 + setShutdown = func(func(context.Context)) { registrations++ } + + config.CacheConfig = &config.Cache{Memory: struct{}{}} + config.QueueConfig = &config.Queue{Memory: &config.QueueMemory{PoolSize: 10}} + + Setup() + Setup() + Setup() + + if registrations != 1 { + t.Errorf("three reloads registered the drain %d times, want 1", registrations) + } +} + +// The drain reaches the adapter that is current when the signal arrives, not +// one captured while wiring up. A reload replaces the adapter, and draining the +// one that was installed at start-up would drain something nobody has published +// to since. +func TestTheDrainRunsAgainstTheAdapterInstalledLast(t *testing.T) { + isolate(t) + config.CacheConfig = &config.Cache{Memory: struct{}{}} + config.QueueConfig = &config.Queue{Memory: &config.QueueMemory{PoolSize: 10}} + Setup() + + first, second := &countingQueue{}, &countingQueue{} + setInstalled(first) + setInstalled(second) + + if err := sdk.Runtime.RunShutdown(context.Background()); err != nil { + t.Fatalf("RunShutdown: %v", err) + } + + if first.calls.Load() != 0 { + t.Error("the adapter that was replaced was shut down; the callback captured it instead of " + + "reading it when it ran") + } + if second.calls.Load() != 1 { + t.Errorf("the current adapter was shut down %d times, want 1", second.calls.Load()) + } +} + +// Nothing installed is the shipped default: settings.yml has no queue section, +// so setupQueue returns early and the runtime's own fallback queue is what +// callers get. Shutting that down would close a queue this package neither +// built nor started. +func TestTheDrainDoesNothingWhenThisPackageInstalledNothing(t *testing.T) { + isolate(t) + config.CacheConfig = &config.Cache{Memory: struct{}{}} + config.QueueConfig = &config.Queue{} + Setup() + + if currentInstalled() != nil { + t.Fatal("an empty queue configuration installed an adapter, so this test asserts nothing") + } + if err := sdk.Runtime.RunShutdown(context.Background()); err != nil { + t.Fatalf("RunShutdown: %v", err) + } +} + +// The budget bounds the wait, not the work. A consumer that never finishes must +// not hold the process past its grace period - SIGKILL would arrive mid-write +// instead of at a point of the process's choosing. +func TestTheDrainStopsWaitingWhenTheBudgetIsGone(t *testing.T) { + isolate(t) + + blocked := &countingQueue{block: make(chan struct{})} + t.Cleanup(func() { close(blocked.block) }) + setInstalled(blocked) + + ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) + defer cancel() + + returned := make(chan struct{}) + go func() { + defer close(returned) + shutdownQueue(ctx) + }() + + select { + case <-returned: + case <-time.After(5 * time.Second): + t.Fatal("shutdownQueue did not return after its context expired - it waits on a Shutdown " + + "that takes no context, so the wait has to be bounded here") + } + + if blocked.calls.Load() != 1 { + t.Errorf("Shutdown called %d times, want 1 - the drain has to be attempted even when it "+ + "cannot be waited out", blocked.calls.Load()) + } +} + +// After the drain this package owns nothing. A reload arriving mid-shutdown +// then builds a new adapter rather than being handed a closed one as its +// `previous` to shut down again. +func TestTheDrainGivesUpOwnershipOfTheAdapter(t *testing.T) { + isolate(t) + setInstalled(&countingQueue{}) + + shutdownQueue(context.Background()) + + if got := currentInstalled(); got != nil { + t.Errorf("installed is %T after the drain, want nil", got) + } +} + +// runtimeQueue is a full AdapterQueue, so it can be handed to the runtime +// rather than only to this package's own record of what it installed. +type runtimeQueue struct { + countingQueue +} + +func (q *runtimeQueue) String() string { return "runtime-fake" } +func (q *runtimeQueue) Append(corestorage.Messager) error { return nil } +func (q *runtimeQueue) Register(string, corestorage.ConsumerFunc) {} +func (q *runtimeQueue) Run() {} + +// The drain must not reach a queue this package did not install. +// +// sdk.Runtime.GetQueueAdapter never returns nil: with no queue section +// configured it wraps the runtime's own fallback, and the wrapper's Shutdown +// forwards. Reaching for the accessor would therefore look like it worked and +// would close a queue this package neither built nor started - the same shape +// as the `if q := GetQueueAdapter(); q != nil` that setupQueue already had to +// drop. +// +// The previous test's empty-configuration case cannot see this: it only checks +// that RunShutdown returns, which it would either way. +func TestTheDrainNeverReachesTheRuntimesOwnQueue(t *testing.T) { + isolate(t) + + onTheRuntime := &runtimeQueue{} + sdk.Runtime.SetQueueAdapter(onTheRuntime) + + if currentInstalled() != nil { + t.Fatal("this package installed something, so the distinction under test is not set up") + } + if sdk.Runtime.GetQueueAdapter() == nil { + t.Fatal("the accessor returned nil, so it is no longer the trap this guards") + } + + shutdownQueue(context.Background()) + + if got := onTheRuntime.calls.Load(); got != 0 { + t.Errorf("the runtime's queue was shut down %d times - the drain went through "+ + "GetQueueAdapter instead of the adapter this package installed", got) + } +}