优化结构
This commit is contained in:
Yvan
@@ -3,7 +3,6 @@ package data
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/redis/go-redis/v9"
|
||||
@@ -48,6 +47,21 @@ func (s *retiredSet[T]) take(item T) bool {
|
||||
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()
|
||||
@@ -56,138 +70,116 @@ func (s *retiredSet[T]) drain() []T {
|
||||
return items
|
||||
}
|
||||
|
||||
// reloadableDB makes the pointer swap atomic. Replaced pools stay open for
|
||||
// retireGrace so in-flight GORM operations remain valid.
|
||||
type reloadableDB struct {
|
||||
current atomic.Pointer[gorm.DB]
|
||||
retired retiredSet[*gorm.DB]
|
||||
// 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 newReloadableDB(db *gorm.DB, enqueue dataScopeAuditEnqueue) *reloadableDB {
|
||||
r := &reloadableDB{}
|
||||
registerDataScopeCallbacks(db, enqueue)
|
||||
r.current.Store(db)
|
||||
return r
|
||||
func newReloadable[T comparable](current T, closeOne func(T)) *reloadable[T] {
|
||||
return &reloadable[T]{current: current, closeOne: closeOne}
|
||||
}
|
||||
|
||||
func (r *reloadableDB) WithContext(ctx context.Context) *gorm.DB {
|
||||
return r.current.Load().WithContext(ctx)
|
||||
// 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
|
||||
}
|
||||
|
||||
func (r *reloadableDB) DB() *gorm.DB { return r.current.Load() }
|
||||
|
||||
func (r *reloadableDB) replace(db *gorm.DB, enqueue dataScopeAuditEnqueue) {
|
||||
registerDataScopeCallbacks(db, enqueue)
|
||||
old := r.current.Swap(db)
|
||||
if old != nil && old != db {
|
||||
r.retired.retire(old, closeGormDB)
|
||||
// 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 *reloadableDB) close() {
|
||||
all := append([]*gorm.DB{r.current.Load()}, r.retired.drain()...)
|
||||
seen := map[*gorm.DB]struct{}{}
|
||||
for _, db := range all {
|
||||
if db == nil {
|
||||
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[db]; ok {
|
||||
if _, ok := seen[item]; ok {
|
||||
continue
|
||||
}
|
||||
seen[db] = struct{}{}
|
||||
closeGormDB(db)
|
||||
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 attaches the row-level data scope callbacks the pool needs
|
||||
// before it becomes readable. Data.replacePrimaryDB is the reload counterpart.
|
||||
func newReloadableDB(db *gorm.DB, enqueue dataScopeAuditEnqueue) *reloadable[*gorm.DB] {
|
||||
registerDataScopeCallbacks(db, enqueue)
|
||||
return newReloadable(db, closeGormDB)
|
||||
}
|
||||
|
||||
func closeGormDB(db *gorm.DB) {
|
||||
if sqlDB, err := db.DB(); err == nil {
|
||||
_ = sqlDB.Close()
|
||||
}
|
||||
}
|
||||
|
||||
type reloadableMongo struct {
|
||||
mu sync.RWMutex
|
||||
current *mongo.Client
|
||||
retired retiredSet[*mongo.Client]
|
||||
}
|
||||
|
||||
func newReloadableMongo(client *mongo.Client) *reloadableMongo {
|
||||
return &reloadableMongo{current: client}
|
||||
}
|
||||
|
||||
func (r *reloadableMongo) replace(client *mongo.Client) {
|
||||
r.mu.Lock()
|
||||
old := r.current
|
||||
r.current = client
|
||||
r.mu.Unlock()
|
||||
if old != nil && old != client {
|
||||
r.retired.retire(old, closeMongoClient)
|
||||
}
|
||||
}
|
||||
|
||||
func (r *reloadableMongo) load() *mongo.Client {
|
||||
if r == nil {
|
||||
return nil
|
||||
}
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
return r.current
|
||||
}
|
||||
|
||||
func (r *reloadableMongo) close() {
|
||||
r.mu.Lock()
|
||||
current := r.current
|
||||
r.current = nil
|
||||
r.mu.Unlock()
|
||||
for _, client := range append([]*mongo.Client{current}, r.retired.drain()...) {
|
||||
if client != nil {
|
||||
closeMongoClient(client)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func closeMongoClient(client *mongo.Client) {
|
||||
_ = client.Disconnect(context.Background())
|
||||
}
|
||||
|
||||
type reloadableRedis struct {
|
||||
mu sync.RWMutex
|
||||
current redis.UniversalClient
|
||||
retired retiredSet[redis.UniversalClient]
|
||||
}
|
||||
|
||||
func newReloadableRedis(client redis.UniversalClient) *reloadableRedis {
|
||||
return &reloadableRedis{current: client}
|
||||
}
|
||||
|
||||
func (r *reloadableRedis) load() redis.UniversalClient {
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
return r.current
|
||||
}
|
||||
|
||||
func (r *reloadableRedis) replace(client redis.UniversalClient) {
|
||||
r.mu.Lock()
|
||||
old := r.current
|
||||
r.current = client
|
||||
r.mu.Unlock()
|
||||
if old != nil && old != client {
|
||||
r.retired.retire(old, closeRedisClient)
|
||||
}
|
||||
}
|
||||
|
||||
func (r *reloadableRedis) close() {
|
||||
r.mu.Lock()
|
||||
current := r.current
|
||||
r.current = nil
|
||||
r.mu.Unlock()
|
||||
for _, client := range append([]redis.UniversalClient{current}, r.retired.drain()...) {
|
||||
if client != nil {
|
||||
closeRedisClient(client)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func closeRedisClient(client redis.UniversalClient) {
|
||||
_ = client.Close()
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user