优化结构

internal/data/runtime_clients.go

@@ -4,60 +4,63 @@ import (
 	"context"
 	"sync"
 	"sync/atomic"
+	"time"
 
 	"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
+// 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
 }
 
-type reloadableMongo struct {
-	mu      sync.RWMutex
-	current *mongo.Client
-	retired []*mongo.Client
+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)
+		}
+	})
 }
 
-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) load() *mongo.Client {
-	if r == nil {
-		return nil
-	}
-	r.mu.RLock()
-	client := r.current
-	r.mu.RUnlock()
-	return client
-}
-
-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())
+// 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
+}
+
+func (s *retiredSet[T]) drain() []T {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	items := s.items
+	s.items = nil
+	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]
 }
 
 func newReloadableDB(db *gorm.DB, enqueue dataScopeAuditEnqueue) *reloadableDB {
@@ -77,18 +80,12 @@ 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()
+		r.retired.retire(old, closeGormDB)
 	}
 }
 
 func (r *reloadableDB) close() {
-	current := r.current.Load()
-	r.mu.Lock()
-	all := append([]*gorm.DB{current}, r.retired...)
-	r.retired = nil
-	r.mu.Unlock()
+	all := append([]*gorm.DB{r.current.Load()}, r.retired.drain()...)
 	seen := map[*gorm.DB]struct{}{}
 	for _, db := range all {
 		if db == nil {
@@ -98,16 +95,65 @@ func (r *reloadableDB) close() {
 			continue
 		}
 		seen[db] = struct{}{}
-		if sqlDB, err := db.DB(); err == nil {
-			_ = sqlDB.Close()
+		closeGormDB(db)
+	}
+}
+
+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 []redis.UniversalClient
+	retired retiredSet[redis.UniversalClient]
 }
 
 func newReloadableRedis(client redis.UniversalClient) *reloadableRedis {
@@ -124,22 +170,24 @@ 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()
+	if old != nil && old != client {
+		r.retired.retire(old, closeRedisClient)
+	}
 }
 
 func (r *reloadableRedis) close() {
 	r.mu.Lock()
-	all := append([]redis.UniversalClient{r.current}, r.retired...)
+	current := r.current
 	r.current = nil
-	r.retired = nil
 	r.mu.Unlock()
-	for _, client := range all {
-		if client == nil {
-			continue
+	for _, client := range append([]redis.UniversalClient{current}, r.retired.drain()...) {
+		if client != nil {
+			closeRedisClient(client)
 		}
-		_ = client.Close()
 	}
 }
+
+func closeRedisClient(client redis.UniversalClient) {
+	_ = client.Close()
+}