kra-new/internal/data/runtime_clients.go

package data

import (
	"context"
	"sync"
	"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
}

// forget removes every pending retirement for item without closing it. A
// reload can legitimately publish a previously retired handle again (for
// example A -> B -> A); its old timers must no longer own that live handle.
func (s *retiredSet[T]) forget(item T) {
	s.mu.Lock()
	defer s.mu.Unlock()
	for index := 0; index < len(s.items); {
		if s.items[index] != item {
			index++
			continue
		}
		s.items = append(s.items[:index], s.items[index+1:]...)
	}
}

func (s *retiredSet[T]) drain() []T {
	s.mu.Lock()
	defer s.mu.Unlock()
	items := s.items
	s.items = nil
	return items
}

// reloadable owns one hot-swappable storage client. The database, Redis and
// Mongo handles all need the same three things — a racy-free read on every
// request, an atomic swap on reload, and a close that also drains whatever the
// grace period has not reaped yet — so they share this one implementation.
type reloadable[T comparable] struct {
	mu       sync.RWMutex
	current  T
	retired  retiredSet[T]
	closeOne func(T)
}

func newReloadable[T comparable](current T, closeOne func(T)) *reloadable[T] {
	return &reloadable[T]{current: current, closeOne: closeOne}
}

// load returns the active client, or the zero value when the client was never
// configured, so optional backends need no separate presence flag.
func (r *reloadable[T]) load() T {
	var zero T
	if r == nil {
		return zero
	}
	r.mu.RLock()
	defer r.mu.RUnlock()
	return r.current
}

// replace publishes next and retires the handle it displaced.
func (r *reloadable[T]) replace(next T) {
	var zero T
	if next != zero {
		r.retired.forget(next)
	}
	r.mu.Lock()
	old := r.current
	r.current = next
	r.mu.Unlock()
	if old != zero && old != next {
		r.retired.retire(old, func(item T) {
			// Serialize the final liveness check with replace: a handle that was
			// re-published while its timer fired must never be closed underneath
			// the new active configuration.
			r.mu.Lock()
			defer r.mu.Unlock()
			if r.current != item {
				r.closeOne(item)
			}
		})
	}
}

func (r *reloadable[T]) close() {
	var zero T
	r.mu.Lock()
	current := r.current
	r.current = zero
	r.mu.Unlock()
	seen := make(map[T]struct{})
	for _, item := range append([]T{current}, r.retired.drain()...) {
		if item == zero {
			continue
		}
		if _, ok := seen[item]; ok {
			continue
		}
		seen[item] = struct{}{}
		r.closeOne(item)
	}
}

// rollback collects the cleanup for every client a multi-step reload opened
// before it knows whether the reload will succeed. Any early return closes
// exactly what was opened, in reverse order; commit hands every handle over to
// the process and cancels all of it. One flag replaces one accepted-boolean plus
// one deferred closure per client.
type rollback struct {
	cleanups  []func()
	committed bool
}

func (r *rollback) add(cleanup func()) { r.cleanups = append(r.cleanups, cleanup) }

func (r *rollback) commit() { r.committed = true }

func (r *rollback) run() {
	if r.committed {
		return
	}
	for index := len(r.cleanups) - 1; index >= 0; index-- {
		r.cleanups[index]()
	}
}

// newReloadableDB attaches the row-level data scope callbacks the pool needs
// before it becomes readable. Data.replacePrimaryDB is the reload counterpart.
func newReloadableDB(db *gorm.DB, enqueue dataScopeAuditEnqueue) *reloadable[*gorm.DB] {
	registerDataScopeCallbacks(db, enqueue)
	return newReloadable(db, closeGormDB)
}

func closeGormDB(db *gorm.DB) {
	if sqlDB, err := db.DB(); err == nil {
		_ = sqlDB.Close()
	}
}

func closeMongoClient(client *mongo.Client) {
	_ = client.Disconnect(context.Background())
}

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