kra-new/internal/data/runtime_clients.go

package data

import (
	"context"
	"sync"
	"sync/atomic"
	"time"

	"github.com/redis/go-redis/v9"
	"go.mongodb.org/mongo-driver/mongo"
	"gorm.io/gorm"
)

// 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
}

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)
		}
	})
}

// 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 {
	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.retired.retire(old, closeGormDB)
	}
}

func (r *reloadableDB) close() {
	all := append([]*gorm.DB{r.current.Load()}, r.retired.drain()...)
	seen := map[*gorm.DB]struct{}{}
	for _, db := range all {
		if db == nil {
			continue
		}
		if _, ok := seen[db]; ok {
			continue
		}
		seen[db] = struct{}{}
		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 retiredSet[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
	r.mu.Unlock()
	if old != nil && old != client {
		r.retired.retire(old, closeRedisClient)
	}
}

func (r *reloadableRedis) close() {
	r.mu.Lock()
	current := r.current
	r.current = nil
	r.mu.Unlock()
	for _, client := range append([]redis.UniversalClient{current}, r.retired.drain()...) {
		if client != nil {
			closeRedisClient(client)
		}
	}
}

func closeRedisClient(client redis.UniversalClient) {
	_ = client.Close()
}