105 lines
2.1 KiB
Go
105 lines
2.1 KiB
Go
package data
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
"sync/atomic"
|
|
|
|
"github.com/redis/go-redis/v9"
|
|
"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
|
|
}
|
|
|
|
func newReloadableDB(db *gorm.DB) *reloadableDB {
|
|
r := &reloadableDB{}
|
|
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) {
|
|
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 {
|
|
current atomic.Pointer[redis.Client]
|
|
mu sync.Mutex
|
|
retired []*redis.Client
|
|
}
|
|
|
|
func newReloadableRedis(client *redis.Client) *reloadableRedis {
|
|
r := &reloadableRedis{}
|
|
if client != nil {
|
|
r.current.Store(client)
|
|
}
|
|
return r
|
|
}
|
|
|
|
func (r *reloadableRedis) load() *redis.Client { return r.current.Load() }
|
|
|
|
func (r *reloadableRedis) replace(client *redis.Client) {
|
|
old := r.current.Swap(client)
|
|
if old != nil && old != client {
|
|
r.mu.Lock()
|
|
r.retired = append(r.retired, old)
|
|
r.mu.Unlock()
|
|
}
|
|
}
|
|
|
|
func (r *reloadableRedis) close() {
|
|
current := r.current.Load()
|
|
r.mu.Lock()
|
|
all := append([]*redis.Client{current}, r.retired...)
|
|
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()
|
|
}
|
|
}
|