优化结构

internal/data/runtime_clients.go

@@ -0,0 +1,134 @@
+package data
+
+import (
+	"context"
+	"sync"
+	"sync/atomic"
+
+	"github.com/redis/go-redis/v9"
+	"go.mongodb.org/mongo-driver/mongo"
+	"gorm.io/gorm"
+)
+
+// reloadableDB makes the pointer swap atomic. Replaced pools are retained
+// until application shutdown so in-flight GORM operations remain valid.
+type reloadableDB struct {
+	current atomic.Pointer[gorm.DB]
+	mu      sync.Mutex
+	retired []*gorm.DB
+}
+
+type reloadableMongo struct {
+	mu      sync.RWMutex
+	current *mongo.Client
+	retired []*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
+	if old != nil && old != client {
+		r.retired = append(r.retired, old)
+	}
+	r.mu.Unlock()
+}
+func (r *reloadableMongo) close() {
+	r.mu.Lock()
+	all := append([]*mongo.Client{r.current}, r.retired...)
+	r.current = nil
+	r.retired = nil
+	r.mu.Unlock()
+	for _, client := range all {
+		if client != nil {
+			_ = client.Disconnect(context.Background())
+		}
+	}
+}
+
+func newReloadableDB(db *gorm.DB, enqueue dataScopeAuditEnqueue) *reloadableDB {
+	r := &reloadableDB{}
+	registerDataScopeCallbacks(db, enqueue)
+	r.current.Store(db)
+	return r
+}
+
+func (r *reloadableDB) WithContext(ctx context.Context) *gorm.DB {
+	return r.current.Load().WithContext(ctx)
+}
+
+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.mu.Lock()
+		r.retired = append(r.retired, old)
+		r.mu.Unlock()
+	}
+}
+
+func (r *reloadableDB) close() {
+	current := r.current.Load()
+	r.mu.Lock()
+	all := append([]*gorm.DB{current}, r.retired...)
+	r.retired = nil
+	r.mu.Unlock()
+	seen := map[*gorm.DB]struct{}{}
+	for _, db := range all {
+		if db == nil {
+			continue
+		}
+		if _, ok := seen[db]; ok {
+			continue
+		}
+		seen[db] = struct{}{}
+		if sqlDB, err := db.DB(); err == nil {
+			_ = sqlDB.Close()
+		}
+	}
+}
+
+type reloadableRedis struct {
+	mu      sync.RWMutex
+	current redis.UniversalClient
+	retired []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
+	if old != nil && old != client {
+		r.retired = append(r.retired, old)
+	}
+	r.mu.Unlock()
+}
+
+func (r *reloadableRedis) close() {
+	r.mu.Lock()
+	all := append([]redis.UniversalClient{r.current}, r.retired...)
+	r.current = nil
+	r.retired = nil
+	r.mu.Unlock()
+	for _, client := range all {
+		if client == nil {
+			continue
+		}
+		_ = client.Close()
+	}
+}