feat: add PawSQL docs examples and image CI
All checks were successful
Build and Push Image / docker-build-and-push (push) Successful in 2m28s
All checks were successful
Build and Push Image / docker-build-and-push (push) Successful in 2m28s
This commit is contained in:
253
internal/postgres/resolver.go
Normal file
253
internal/postgres/resolver.go
Normal file
@@ -0,0 +1,253 @@
|
||||
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)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user