135 lines
2.7 KiB
Go
135 lines
2.7 KiB
Go
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) *reloadableDB {
|
|
r := &reloadableDB{}
|
|
registerDataScopeCallbacks(db)
|
|
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) {
|
|
registerDataScopeCallbacks(db)
|
|
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()
|
|
}
|
|
}
|