pawsql/internal/postgres/resolver.go

package postgres

import (
	"context"
	"log/slog"
	"sync"
	"time"

	"github.com/barkstack/pawsql/internal/config"
	"github.com/barkstack/pawsql/internal/router"
)

// DatabaseController starts and stops managed PostgreSQL databases.
type DatabaseController interface {
	EnsureDatabase(context.Context, string, config.PostgresConfig) (string, error)
	StopDatabase(context.Context, string) error
}

// Resolver lazily starts managed PostgreSQL containers when a route is selected.
// Concurrent connections to one database share a single start or stop operation.
type Resolver struct {
	routes     router.BackendResolver
	controller DatabaseController
	managed    map[string]config.PostgresConfig

	locksMu sync.Mutex
	locks   map[string]*sync.Mutex

	statesMu sync.Mutex
	states   map[string]*leaseState
}

type leaseState struct {
	active            int
	idleGeneration    uint64
	idleTimer         *time.Timer
	trafficGeneration uint64
	trafficTimer      *time.Timer
	lastActivity      time.Time
	clientBytes       uint64
	backendBytes      uint64
}

// TrafficStats is a point-in-time aggregate for one managed database.
type TrafficStats struct {
	ActiveConnections int
	ClientBytes       uint64
	BackendBytes      uint64
	LastActivity      time.Time
}

// NewResolver wraps static route matching with lazy PostgreSQL provisioning.
func NewResolver(routes router.BackendResolver, databases []config.DatabaseConfig, controller DatabaseController) *Resolver {
	managed := make(map[string]config.PostgresConfig)
	for _, database := range databases {
		if database.Postgres != nil {
			managed[database.Name] = *database.Postgres
		}
	}
	return &Resolver{
		routes: routes, controller: controller, managed: managed,
		locks: make(map[string]*sync.Mutex), states: make(map[string]*leaseState),
	}
}

// Resolve selects a hostname route and starts its PostgreSQL container if needed.
func (r *Resolver) Resolve(ctx context.Context, hostname string) (router.Backend, error) {
	backend, err := r.routes.Resolve(ctx, hostname)
	if err != nil {
		return router.Backend{}, err
	}
	return r.ensure(ctx, backend)
}

// ResolveDatabase selects a database route and starts its PostgreSQL container if needed.
func (r *Resolver) ResolveDatabase(ctx context.Context, database string) (router.Backend, error) {
	backend, err := r.routes.ResolveDatabase(ctx, database)
	if err != nil {
		return router.Backend{}, err
	}
	return r.ensure(ctx, backend)
}

func (r *Resolver) ensure(ctx context.Context, backend router.Backend) (router.Backend, error) {
	postgres, managed := r.managed[backend.DatabaseName]
	if !managed {
		return backend, nil
	}
	lock := r.lockFor(backend.DatabaseName)
	lock.Lock()
	defer lock.Unlock()
	address, err := r.controller.EnsureDatabase(ctx, backend.DatabaseName, postgres)
	if err != nil {
		return router.Backend{}, err
	}
	r.acquire(backend.DatabaseName, postgres)
	backend.Address = address
	return backend, nil
}

func (r *Resolver) lockFor(database string) *sync.Mutex {
	r.locksMu.Lock()
	defer r.locksMu.Unlock()
	lock := r.locks[database]
	if lock == nil {
		lock = &sync.Mutex{}
		r.locks[database] = lock
	}
	return lock
}

// ReleaseConnection records a proxied managed-database session ending. Once the
// final session closes, an idle_timeout countdown begins.
func (r *Resolver) ReleaseConnection(database string) {
	postgres, managed := r.managed[database]
	if !managed {
		return
	}
	r.statesMu.Lock()
	state := r.states[database]
	if state == nil || state.active == 0 {
		r.statesMu.Unlock()
		return
	}
	state.active--
	if state.active == 0 {
		state.trafficGeneration++
		if state.trafficTimer != nil {
			state.trafficTimer.Stop()
			state.trafficTimer = nil
		}
		if postgres.IdleTimeout > 0 {
			state.idleGeneration++
			generation := state.idleGeneration
			state.idleTimer = time.AfterFunc(postgres.IdleTimeout, func() {
				r.stopAfterConnectionIdle(database, generation)
			})
		}
	}
	r.statesMu.Unlock()
}

// RecordTraffic meters proxied bytes and resets the traffic-idle countdown.
func (r *Resolver) RecordTraffic(database string, clientToBackend bool, bytes int64) {
	postgres, managed := r.managed[database]
	if !managed || bytes <= 0 {
		return
	}
	r.statesMu.Lock()
	state := r.stateFor(database)
	if clientToBackend {
		state.clientBytes += uint64(bytes)
	} else {
		state.backendBytes += uint64(bytes)
	}
	state.lastActivity = time.Now()
	if state.active > 0 && postgres.TrafficIdleTimeout > 0 {
		r.scheduleTrafficIdleLocked(database, state, postgres.TrafficIdleTimeout)
	}
	r.statesMu.Unlock()
}

// TrafficStats reports aggregate traffic collected since PawSQL started.
func (r *Resolver) TrafficStats(database string) (TrafficStats, bool) {
	if _, managed := r.managed[database]; !managed {
		return TrafficStats{}, false
	}
	r.statesMu.Lock()
	defer r.statesMu.Unlock()
	state := r.states[database]
	if state == nil {
		return TrafficStats{}, true
	}
	return TrafficStats{
		ActiveConnections: state.active,
		ClientBytes:       state.clientBytes,
		BackendBytes:      state.backendBytes,
		LastActivity:      state.lastActivity,
	}, true
}

func (r *Resolver) acquire(database string, postgres config.PostgresConfig) {
	r.statesMu.Lock()
	defer r.statesMu.Unlock()
	state := r.stateFor(database)
	state.active++
	state.lastActivity = time.Now()
	state.idleGeneration++
	if state.idleTimer != nil {
		state.idleTimer.Stop()
		state.idleTimer = nil
	}
	if postgres.TrafficIdleTimeout > 0 {
		r.scheduleTrafficIdleLocked(database, state, postgres.TrafficIdleTimeout)
	}
}

func (r *Resolver) stateFor(database string) *leaseState {
	state := r.states[database]
	if state == nil {
		state = &leaseState{}
		r.states[database] = state
	}
	return state
}

func (r *Resolver) scheduleTrafficIdleLocked(database string, state *leaseState, timeout time.Duration) {
	state.trafficGeneration++
	generation := state.trafficGeneration
	if state.trafficTimer != nil {
		state.trafficTimer.Stop()
	}
	state.trafficTimer = time.AfterFunc(timeout, func() {
		r.stopAfterTrafficIdle(database, generation)
	})
}

func (r *Resolver) stopAfterConnectionIdle(database string, generation uint64) {
	r.stopIfIdle(database, generation, false)
}

func (r *Resolver) stopAfterTrafficIdle(database string, generation uint64) {
	r.stopIfIdle(database, generation, true)
}

func (r *Resolver) stopIfIdle(database string, generation uint64, traffic bool) {
	lock := r.lockFor(database)
	lock.Lock()
	defer lock.Unlock()
	r.statesMu.Lock()
	state := r.states[database]
	if state == nil {
		r.statesMu.Unlock()
		return
	}
	if traffic {
		if state.active == 0 || state.trafficGeneration != generation {
			r.statesMu.Unlock()
			return
		}
		state.trafficTimer = nil
	} else {
		if state.active != 0 || state.idleGeneration != generation {
			r.statesMu.Unlock()
			return
		}
		state.idleTimer = nil
	}
	r.statesMu.Unlock()
	if err := r.controller.StopDatabase(context.Background(), database); err != nil {
		slog.Error("stop idle PostgreSQL container", "database", database, "traffic_idle", traffic, "error", err)
	}
}