fix

internal/data/runtime_clients.go

@@ -6,6 +6,7 @@ import (
 	"sync/atomic"
 
 	"github.com/redis/go-redis/v9"
+	"go.mongodb.org/mongo-driver/mongo"
 	"gorm.io/gorm"
 )
 
@@ -17,8 +18,40 @@ type reloadableDB struct {
 	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) *reloadableDB {
 	r := &reloadableDB{}
+	registerDataScopeCallbacks(db)
 	r.current.Store(db)
 	return r
 }
@@ -30,6 +63,7 @@ func (r *reloadableDB) WithContext(ctx context.Context) *gorm.DB {
 func (r *reloadableDB) DB() *gorm.DB { return r.current.Load() }
 
 func (r *reloadableDB) replace(db *gorm.DB) {
+	registerDataScopeCallbacks(db)
 	old := r.current.Swap(db)
 	if old != nil && old != db {
 		r.mu.Lock()
@@ -60,45 +94,41 @@ func (r *reloadableDB) close() {
 }
 
 type reloadableRedis struct {
-	current atomic.Pointer[redis.Client]
-	mu      sync.Mutex
-	retired []*redis.Client
+	mu      sync.RWMutex
+	current redis.UniversalClient
+	retired []redis.UniversalClient
 }
 
-func newReloadableRedis(client *redis.Client) *reloadableRedis {
-	r := &reloadableRedis{}
-	if client != nil {
-		r.current.Store(client)
-	}
-	return r
+func newReloadableRedis(client redis.UniversalClient) *reloadableRedis {
+	return &reloadableRedis{current: client}
 }
 
-func (r *reloadableRedis) load() *redis.Client { return r.current.Load() }
+func (r *reloadableRedis) load() redis.UniversalClient {
+	r.mu.RLock()
+	defer r.mu.RUnlock()
+	return r.current
+}
 
-func (r *reloadableRedis) replace(client *redis.Client) {
-	old := r.current.Swap(client)
+func (r *reloadableRedis) replace(client redis.UniversalClient) {
+	r.mu.Lock()
+	old := r.current
+	r.current = client
 	if old != nil && old != client {
-		r.mu.Lock()
 		r.retired = append(r.retired, old)
-		r.mu.Unlock()
 	}
+	r.mu.Unlock()
 }
 
 func (r *reloadableRedis) close() {
-	current := r.current.Load()
 	r.mu.Lock()
-	all := append([]*redis.Client{current}, r.retired...)
+	all := append([]redis.UniversalClient{r.current}, r.retired...)
+	r.current = nil
 	r.retired = nil
 	r.mu.Unlock()
-	seen := map[*redis.Client]struct{}{}
 	for _, client := range all {
 		if client == nil {
 			continue
 		}
-		if _, ok := seen[client]; ok {
-			continue
-		}
-		seen[client] = struct{}{}
 		_ = client.Close()
 	}
 }