Keep Mongo operations online across Vault credential rotation

2026-09-12 22:19:36 +02:00
1 parent 80d6917aa7
commit 23f7b78c3b
@@ -0,0 +1,253 @@
+package source
+
+import (
+	"context"
+	"sync"
+	"sync/atomic"
+	"time"
+
+	merrors "git.sendico.io/tech/sxchnge/shared/errors"
+	mlogger "git.sendico.io/tech/sxchnge/shared/logging"
+	"go.mongodb.org/mongo-driver/v2/mongo"
+	"go.mongodb.org/mongo-driver/v2/mongo/options"
+	"go.mongodb.org/mongo-driver/v2/mongo/readpref"
+	"go.uber.org/zap"
+)
+
+// Lease owns a credential and its revocation capability. Neither may be logged.
+type Lease struct {
+	URI       string
+	ExpiresAt time.Time
+	Revoke    func(context.Context) error
+}
+type Issuer interface {
+	Issue(context.Context) (Lease, error)
+}
+
+type generation struct {
+	client  *mongo.Client
+	lease   Lease
+	refs    int
+	retired bool
+	drained chan struct{}
+}
+type pin struct {
+	manager    *Manager
+	generation *generation
+	active     atomic.Bool
+}
+type pinKey struct{}
+
+// Manager publishes a verified replacement before retiring the previous pool.
+// Every cursor, operation and transaction retains its originating generation.
+type Manager struct {
+	mu         sync.Mutex
+	current    *generation
+	closed     bool
+	issuer     Issuer
+	database   string
+	logger     mlogger.Logger
+	cancel     context.CancelFunc
+	workerDone chan struct{}
+	draining   sync.WaitGroup
+}
+
+func New(ctx context.Context, logger mlogger.Logger, database string, issuer Issuer) (*Manager, error) {
+	if database == "" || issuer == nil {
+		return nil, merrors.InvalidArgument("database and credential issuer are required")
+	}
+	life, cancel := context.WithCancel(context.Background())
+	m := &Manager{issuer: issuer, database: database, logger: logger, cancel: cancel, workerDone: make(chan struct{})}
+	g, err := m.prepare(ctx)
+	if err != nil {
+		cancel()
+		return nil, err
+	}
+	m.current = g
+	go m.run(life)
+	return m, nil
+}
+
+func (m *Manager) Acquire(ctx context.Context) (*mongo.Database, context.Context, func(), error) {
+	if err := ctx.Err(); err != nil {
+		return nil, ctx, nil, err
+	}
+	if p, ok := ctx.Value(pinKey{}).(*pin); ok && p.manager == m {
+		if !p.active.Load() {
+			return nil, ctx, nil, merrors.FailedPrecondition("database operation scope has ended")
+		}
+		return p.generation.client.Database(m.database), ctx, func() {}, nil
+	}
+	m.mu.Lock()
+	g := m.current
+	if m.closed || g == nil || !time.Now().Before(g.lease.ExpiresAt) {
+		m.mu.Unlock()
+		return nil, ctx, nil, merrors.ServiceUnavailable("no valid database credential is available")
+	}
+	if session := mongo.SessionFromContext(ctx); session != nil && session.Client() != g.client {
+		m.mu.Unlock()
+		return nil, ctx, nil, merrors.FailedPrecondition("transaction has no matching database operation scope")
+	}
+	g.refs++
+	m.mu.Unlock()
+	p := &pin{manager: m, generation: g}
+	p.active.Store(true)
+	bounded, cancel := context.WithDeadline(ctx, g.lease.ExpiresAt)
+	scoped := context.WithValue(bounded, pinKey{}, p)
+	release := func() {
+		if !p.active.CompareAndSwap(true, false) {
+			return
+		}
+		cancel()
+		m.mu.Lock()
+		g.refs--
+		if g.retired && g.refs == 0 {
+			close(g.drained)
+		}
+		m.mu.Unlock()
+	}
+	return g.client.Database(m.database), scoped, release, nil
+}
+
+func (m *Manager) prepare(ctx context.Context) (*generation, error) {
+	ctx, cancel := context.WithTimeout(ctx, 10*time.Second)
+	defer cancel()
+	lease, err := m.issuer.Issue(ctx)
+	if err != nil {
+		return nil, err
+	}
+	var client *mongo.Client
+	success := false
+	defer func() {
+		if !success {
+			if client != nil {
+				cleanup, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+				_ = client.Disconnect(cleanup)
+				cancel()
+			}
+			m.draining.Add(1)
+			go func() { defer m.draining.Done(); m.revoke(lease) }()
+		}
+	}()
+	if lease.Revoke == nil || time.Until(lease.ExpiresAt) < 2*time.Second {
+		return nil, merrors.FailedPrecondition("database credential lifetime is too short")
+	}
+	client, err = mongo.Connect(options.Client().ApplyURI(lease.URI).SetServerSelectionTimeout(5 * time.Second))
+	if err != nil {
+		return nil, merrors.OpaqueWrap(merrors.ErrServiceUnavailable, err, "create replacement database pool")
+	}
+	if err = client.Ping(ctx, readpref.Primary()); err != nil {
+		return nil, merrors.OpaqueWrap(merrors.ErrServiceUnavailable, err, "validate replacement database pool")
+	}
+	if time.Until(lease.ExpiresAt) < time.Second {
+		return nil, merrors.FailedPrecondition("replacement database credential expired during validation")
+	}
+	success = true
+	return &generation{client: client, lease: lease, drained: make(chan struct{})}, nil
+}
+
+func (m *Manager) retireLocked(g *generation) {
+	g.retired = true
+	if g.refs == 0 {
+		close(g.drained)
+	}
+	m.draining.Add(1)
+	go func() {
+		defer m.draining.Done()
+		<-g.drained
+		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+		defer cancel()
+		if err := g.client.Disconnect(ctx); err != nil {
+			m.logger.Warn("Retired database pool close failed")
+		}
+		m.revoke(g.lease)
+		m.logger.Info("Retired database pool drained")
+	}()
+}
+
+func (m *Manager) revoke(lease Lease) {
+	if lease.Revoke == nil {
+		return
+	}
+	// Expiration is the backstop when Vault is unavailable during bounded cleanup.
+	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
+	defer cancel()
+	if err := lease.Revoke(ctx); err != nil {
+		m.logger.Warn("Database credential cleanup failed; Vault expiry will revoke it")
+	}
+
+}
+
+func (m *Manager) run(ctx context.Context) {
+	defer close(m.workerDone)
+	m.mu.Lock()
+	wait := time.Until(m.current.lease.ExpiresAt) / 2
+	m.mu.Unlock()
+	backoff := 250 * time.Millisecond
+	failureReported, expiryReported := false, false
+	for {
+		timer := time.NewTimer(wait)
+		select {
+		case <-ctx.Done():
+			timer.Stop()
+			return
+		case <-timer.C:
+		}
+		next, err := m.prepare(ctx)
+		if err != nil {
+			if ctx.Err() != nil {
+				return
+			}
+			m.mu.Lock()
+			remaining := time.Until(m.current.lease.ExpiresAt)
+			m.mu.Unlock()
+			if !failureReported {
+				m.logger.Warn("Database credential refresh failed; retaining current pool", zap.Duration("valid_for", max(remaining, 0)))
+				failureReported = true
+			}
+			if remaining <= 0 && !expiryReported {
+				m.logger.Error("Database credentials expired; rejecting new operations until refresh succeeds")
+				expiryReported = true
+			}
+			wait = backoff
+			if remaining > 0 {
+				wait = max(100*time.Millisecond, min(wait, remaining/4))
+			}
+			backoff = min(5*time.Second, backoff*2)
+
+			continue
+		}
+		m.mu.Lock()
+		if m.closed {
+			m.retireLocked(next)
+			m.mu.Unlock()
+			return
+		}
+		old := m.current
+		m.current = next
+		m.retireLocked(old)
+		m.mu.Unlock()
+		failureReported, expiryReported = false, false
+		backoff = 250 * time.Millisecond
+		m.logger.Info("Database pool replaced without restarting service")
+		wait = time.Until(next.lease.ExpiresAt) / 2
+	}
+}
+
+func (m *Manager) Close(ctx context.Context) error {
+	m.mu.Lock()
+	if !m.closed {
+		m.closed = true
+		m.cancel()
+		m.retireLocked(m.current)
+	}
+	m.mu.Unlock()
+	done := make(chan struct{})
+	go func() { <-m.workerDone; m.draining.Wait(); close(done) }()
+	select {
+	case <-done:
+		return nil
+	case <-ctx.Done():
+		return ctx.Err()
+	}
+}
@@ -0,0 +1,30 @@
+// Package source owns Mongo connection generations and operation lifetimes.
+package source
+
+import (
+	"context"
+
+	merrors "git.sendico.io/tech/sxchnge/shared/errors"
+	"go.mongodb.org/mongo-driver/v2/mongo"
+)
+
+// Source pins one database/client until release. The returned context must be
+// passed to nested operations so transactions never cross client generations.
+type Source interface {
+	Acquire(context.Context) (*mongo.Database, context.Context, func(), error)
+}
+
+type fixed struct{ database *mongo.Database }
+
+// Fixed adapts an externally owned database, e.g. a test fixture. It never closes it.
+func Fixed(database *mongo.Database) Source { return &fixed{database: database} }
+
+func (s *fixed) Acquire(ctx context.Context) (*mongo.Database, context.Context, func(), error) {
+	if s.database == nil {
+		return nil, ctx, nil, merrors.FailedPrecondition("database is not configured")
+	}
+	if session := mongo.SessionFromContext(ctx); session != nil && session.Client() != s.database.Client() {
+		return nil, ctx, nil, merrors.FailedPrecondition("transaction belongs to another database client")
+	}
+	return s.database, ctx, func() {}, nil
+}