/* * @Author: zhangwenjian * @Date: 2025/04/13 22:03 * @Last Modified by: zhangwenjian * @Last Modified time: 2025/04/13 22:03 */ 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" ) // Setup 配置storage组件 func Setup() { setupCache() setupCaptcha() setupQueue() registerQueueDrain() } func setupCache() { cacheAdapter, err := config.CacheConfig.Setup() if err != nil { log.Fatalf("cache setup error, %s\n", err.Error()) } sdk.Runtime.SetCacheAdapter(cacheAdapter) } func setupCaptcha() { captcha.SetStore(captcha.NewCacheStore(sdk.Runtime.GetCacheAdapter(), 600)) } var ( queueMu sync.Mutex // installed is the adapter setupQueue built, kept so the next reload can // 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) } } // drainedInTime shuts q down and reports whether it finished before ctx expired. func drainedInTime(ctx context.Context, q interface{ Shutdown() }) bool { done := make(chan struct{}) go func() { defer close(done) q.Shutdown() }() return finishedBeforeDeadline(ctx, done) } // finishedBeforeDeadline waits for done or for ctx, and resolves a tie in // favour of done. // // The tie is the reason this is a function of its own rather than one select // inline. Both channels can be ready when the select runs, select picks at // random among ready cases, and so a single look reports an overrun for a // drain that completed - about half the times it lands there, which is exactly // often enough to be dismissed as noise. core's own RunShutdown re-checks for // this reason. // // Taking channels rather than a queue is what makes it testable: a closed done // and an expired ctx can be handed in together, which is the state a race // would otherwise have to be caught in. func finishedBeforeDeadline(ctx context.Context, done <-chan struct{}) bool { select { case <-done: return true case <-ctx.Done(): } select { case <-done: return true default: return false } } // 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 } if drainedInTime(ctx, q) { corelog.Info("queue: drained") return } // 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 // no queue section at all - in which case nothing is installed and callers are // working with the runtime's own fallback queue. // // It exists because there is no way to ask for the adapter's identity from the // outside. sdk.Runtime.GetQueueAdapter and GetQueuePrefix build a fresh // runtime.Queue wrapper on every call, so comparing what two calls return // compares two wrappers and never matches, however many times the underlying // adapter has been replaced. This package creates the adapter, so this is the // only place that knows. A counter rather than the adapter itself keeps the // comparison on a uint64: an adapter type that is not comparable would panic // an `==` between two interface values. func QueueGeneration() uint64 { queueMu.Lock() defer queueMu.Unlock() return installedGen } func setupQueue() { if config.QueueConfig.Empty() { return } queueMu.Lock() defer queueMu.Unlock() queueAdapter, err := config.QueueConfig.Setup() if err != nil { log.Fatalf("queue setup error, %s\n", err.Error()) } previous := installed sdk.Runtime.SetQueueAdapter(queueAdapter) installed = queueAdapter installedGen++ // The previous adapter goes down after the new one is installed, not // before. Shutdown waits for its consumers to deliver what it still holds, // and for that whole wait the runtime would otherwise be handing producers // a queue that has stopped accepting: every Append in the window comes back // ErrQueueClosed, and both call sites in common/middleware log it. Swapping // first leaves no such window - a producer gets the new queue or the old // one, and both work. // // Only an adapter this package installed. GetQueueAdapter never returns // nil - with nothing configured the runtime falls back to its own memory // queue and wraps that - so the `if q := GetQueueAdapter(); q != nil` this // replaces was always true, and shut down the fallback queue on the very // first start, before anything had used it. if previous != nil { previous.Shutdown() } // Deliberately not started here. Run has to come after the consumers have // registered: the contract implementations refuse a registration once the // queue is running (storage.ErrQueueAlreadyStarted), and the legacy // adapter this repository still goes through swallows that error rather // than reporting it - its own comment says the interface gives it no way // to tell the caller. Starting here and registering afterwards is // therefore a race that loses consumers in silence. Whoever registers is // the one that starts it. }