package data import ( "context" "sync" "time" "github.com/redis/go-redis/v9" "go.mongodb.org/mongo-driver/mongo" "gorm.io/gorm" ) // retireGrace bounds how long a client replaced by a hot reload stays open. // In-flight operations still hold the old handle, so it cannot be closed at // swap time; keeping it until shutdown would leak one full pool per reload. const retireGrace = 5 * time.Minute // retiredSet holds clients replaced by a hot reload and closes each of them // once the grace period expires. Whatever is still pending at shutdown is // drained by the owner's close. type retiredSet[T comparable] struct { mu sync.Mutex items []T } func (s *retiredSet[T]) retire(item T, closeItem func(T)) { s.mu.Lock() s.items = append(s.items, item) s.mu.Unlock() time.AfterFunc(retireGrace, func() { if s.take(item) { closeItem(item) } }) } // take removes item and reports whether this caller now owns closing it. func (s *retiredSet[T]) take(item T) bool { s.mu.Lock() defer s.mu.Unlock() for index, existing := range s.items { if existing == item { s.items = append(s.items[:index], s.items[index+1:]...) return true } } return false } // forget removes every pending retirement for item without closing it. A // reload can legitimately publish a previously retired handle again (for // example A -> B -> A); its old timers must no longer own that live handle. func (s *retiredSet[T]) forget(item T) { s.mu.Lock() defer s.mu.Unlock() for index := 0; index < len(s.items); { if s.items[index] != item { index++ continue } s.items = append(s.items[:index], s.items[index+1:]...) } } func (s *retiredSet[T]) drain() []T { s.mu.Lock() defer s.mu.Unlock() items := s.items s.items = nil return items } // reloadable owns one hot-swappable storage client. The database, Redis and // Mongo handles all need the same three things — a racy-free read on every // request, an atomic swap on reload, and a close that also drains whatever the // grace period has not reaped yet — so they share this one implementation. type reloadable[T comparable] struct { mu sync.RWMutex current T retired retiredSet[T] closeOne func(T) } func newReloadable[T comparable](current T, closeOne func(T)) *reloadable[T] { return &reloadable[T]{current: current, closeOne: closeOne} } // load returns the active client, or the zero value when the client was never // configured, so optional backends need no separate presence flag. func (r *reloadable[T]) load() T { var zero T if r == nil { return zero } r.mu.RLock() defer r.mu.RUnlock() return r.current } // replace publishes next and retires the handle it displaced. func (r *reloadable[T]) replace(next T) { var zero T if next != zero { r.retired.forget(next) } r.mu.Lock() old := r.current r.current = next r.mu.Unlock() if old != zero && old != next { r.retired.retire(old, func(item T) { // Serialize the final liveness check with replace: a handle that was // re-published while its timer fired must never be closed underneath // the new active configuration. r.mu.Lock() defer r.mu.Unlock() if r.current != item { r.closeOne(item) } }) } } func (r *reloadable[T]) close() { var zero T r.mu.Lock() current := r.current r.current = zero r.mu.Unlock() seen := make(map[T]struct{}) for _, item := range append([]T{current}, r.retired.drain()...) { if item == zero { continue } if _, ok := seen[item]; ok { continue } seen[item] = struct{}{} r.closeOne(item) } } // rollback collects the cleanup for every client a multi-step reload opened // before it knows whether the reload will succeed. Any early return closes // exactly what was opened, in reverse order; commit hands every handle over to // the process and cancels all of it. One flag replaces one accepted-boolean plus // one deferred closure per client. type rollback struct { cleanups []func() committed bool } func (r *rollback) add(cleanup func()) { r.cleanups = append(r.cleanups, cleanup) } func (r *rollback) commit() { r.committed = true } func (r *rollback) run() { if r.committed { return } for index := len(r.cleanups) - 1; index >= 0; index-- { r.cleanups[index]() } } // newReloadableDB owns only the hot-swappable database handle. Domain-specific // callbacks are attached by Data when a pool becomes active. func newReloadableDB(db *gorm.DB) *reloadable[*gorm.DB] { return newReloadable(db, closeGormDB) } func closeGormDB(db *gorm.DB) { if sqlDB, err := db.DB(); err == nil { _ = sqlDB.Close() } } func closeMongoClient(client *mongo.Client) { _ = client.Disconnect(context.Background()) } func closeRedisClient(client redis.UniversalClient) { _ = client.Close() }