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) + } +}