Janady13/missionmatch / internal/platform/tenantpool/manager.go

1// Package tenantpool is the bounded, lazy registry of per tenant connection pools for the database per 2// tenant model. After auth validates a request to an organization, the router asks this manager for that 3// tenant's pool; the manager resolves the tenant's database from the control plane directory and its 4// credential from the secrets port, opens a pool (lazily, once, through database.NewBoundedPool so the 5// pgvector codecs and the NOBYPASSRLS role assertion run on every tenant pool), verifies it with a Ping, 6// caches it, and hands back a request scoped tenant.DB plus a release function. 7// 8// The design follows the pgx pool per tenant guidance: a BOUNDED number of live pools (LRU evicted), each 9// small (so MaxPools times MaxConns stays under the server max_connections), a Ping before caching so an 10// unreachable tenant fails fast with a typed error instead of hanging, a pin or refcount so an in flight 11// transaction is never closed under it, singleflight so concurrent first callers open exactly one pool, a 12// nonzero connection lifetime jitter, whole pool idle eviction (distinct from pgx's per connection idle 13// close), and typed errors that NEVER fall back to another tenant. RLS stays forced inside each tenant 14// database as defense in depth; the per request org GUC is still set by RunInTenantTx on the tenant pool. 15package tenantpool 16 17import ( 18 "context" 19 "errors" 20 "fmt" 21 "sync" 22 "time" 23 24 "github.com/google/uuid" 25 "github.com/jackc/pgx/v5/pgconn" 26 "golang.org/x/sync/singleflight" 27 28 "github.com/missionworkforce/missionmatch/internal/platform/database" 29 "github.com/missionworkforce/missionmatch/internal/platform/directory" 30 "github.com/missionworkforce/missionmatch/internal/platform/secrets" 31 "github.com/missionworkforce/missionmatch/internal/platform/tenant" 32) 33 34// ErrTenantUnavailable means the tenant exists in the directory but its status does not permit serving 35// (it is provisioning, migrating, suspended, or offline). It is a fail closed denial, never a fallback. 36var ErrTenantUnavailable = errors.New("tenantpool: tenant database is not available for serving") 37 38// ErrTenantUnreachable means the tenant's database could not be opened or verified (a bad credential, a 39// network failure, or the wrong role). The caller maps it to a typed 503, never a hang and never another 40// tenant's data. 41var ErrTenantUnreachable = errors.New("tenantpool: tenant database is unreachable") 42 43// ErrClosed means the manager has been shut down. 44var ErrClosed = errors.New("tenantpool: manager is closed") 45 46// Config bounds the registry. MaxPools caps the live pools (LRU evicted beyond it); MaxConnsPerPool keeps 47// each pool small; IdleEvict closes a whole pool after this long with no use. Role is the database role 48// each pool must authenticate as (empty defaults to the NOBYPASSRLS app role); the compliance registry 49// asserts the compliance role instead. SecretRefSuffix is appended to the directory secret reference 50// before resolving, so the compliance registry resolves a distinct compliance credential for the same 51// tenant database (the app and compliance roles never share a connection). 52type Config struct { 53 MaxPools int 54 MaxConnsPerPool int 55 IdleEvict time.Duration 56 Role string 57 SecretRefSuffix string 58} 59 60// opened is a live tenant pool plus the function that closes it. close is separated from the pool so the 61// registry logic is unit testable without a real database. 62type opened struct { 63 db *tenant.DB 64 close func() 65 ref string 66} 67 68type entry struct { 69 o *opened 70 pins int 71 last time.Time 72 org uuid.UUID 73} 74 75// Manager is the per tenant pool registry. lookup and build are fields so the registry mechanics can be 76// unit tested with fakes; New wires the real directory lookup and the real open and verify path. 77type Manager struct { 78 lookup func(ctx context.Context, org uuid.UUID) (directory.TenantConn, error) 79 build func(ctx context.Context, tc directory.TenantConn) (*opened, error) 80 maxPools int 81 idle time.Duration 82 now func() time.Time 83 84 mu sync.Mutex 85 pools map[uuid.UUID]*entry 86 sf singleflight.Group 87 closed bool 88 stop chan struct{} 89} 90 91// New builds a manager over the control plane directory (read with the app pool) and the secrets 92// resolver. It starts the idle sweep when IdleEvict is positive. 93func New(dir *directory.Directory, sec secrets.Resolver, cfg Config) *Manager { 94 mp := cfg.MaxPools 95 if mp < 1 { 96 mp = 1 97 } 98 role := cfg.Role 99 if role == "" { 100 role = "missionmatch_app" 101 } 102 m := &Manager{ 103 lookup: dir.Lookup, 104 maxPools: mp, 105 idle: cfg.IdleEvict, 106 now: time.Now, 107 pools: map[uuid.UUID]*entry{}, 108 stop: make(chan struct{}), 109 } 110 m.build = defaultBuild(sec, int32(cfg.MaxConnsPerPool), role, cfg.SecretRefSuffix) 111 if m.idle > 0 { 112 go m.sweepLoop() 113 } 114 return m 115} 116 117// Acquire returns the tenant's pool (opening it lazily and once) plus a release function the caller MUST 118// call when the request is done, so the pool is unpinned and may be evicted when idle. A request with an 119// unknown organization gets directory.ErrUnknownTenant; an inactive or unreachable tenant gets a typed 120// error; never a different tenant's pool. 121func (m *Manager) Acquire(ctx context.Context, org uuid.UUID) (*tenant.DB, func(), error) { 122 m.mu.Lock() 123 if m.closed { 124 m.mu.Unlock() 125 return nil, nil, ErrClosed 126 } 127 if e, ok := m.pools[org]; ok { 128 e.pins++ 129 e.last = m.now() 130 db := e.o.db 131 m.mu.Unlock() 132 return db, m.releaser(e), nil 133 } 134 m.mu.Unlock() 135 136 // Cold path: singleflight so concurrent first callers open exactly one pool for this tenant. 137 v, err, _ := m.sf.Do(org.String(), func() (any, error) { 138 m.mu.Lock() 139 if e, ok := m.pools[org]; ok { // opened by a racing caller while we waited 140 m.mu.Unlock() 141 return e, nil 142 } 143 closed := m.closed 144 m.mu.Unlock() 145 if closed { 146 return nil, ErrClosed 147 } 148 tc, err := m.lookup(ctx, org) 149 if err != nil { 150 return nil, err // directory.ErrUnknownTenant propagates unchanged 151 } 152 if tc.Status != directory.StatusActive { 153 return nil, fmt.Errorf("%w: status %q", ErrTenantUnavailable, tc.Status) 154 } 155 o, err := m.build(ctx, tc) 156 if err != nil { 157 return nil, err 158 } 159 m.mu.Lock() 160 if m.closed { // shut down while we were opening 161 m.mu.Unlock() 162 o.close() 163 return nil, ErrClosed 164 } 165 e := &entry{o: o, last: m.now(), org: org} 166 m.pools[org] = e 167 victims := m.evictLocked(org) 168 m.mu.Unlock() 169 closeAll(victims) 170 return e, nil 171 }) 172 if err != nil { 173 return nil, nil, err 174 } 175 e := v.(*entry) 176 177 m.mu.Lock() 178 if m.closed { 179 m.mu.Unlock() 180 return nil, nil, ErrClosed 181 } 182 cur, ok := m.pools[org] 183 if !ok || cur != e { 184 // Evicted in the small window between open and pin (rare, the new entry is most recently used so 185 // eviction targets others first). Retry the whole acquire. 186 m.mu.Unlock() 187 return m.Acquire(ctx, org) 188 } 189 e.pins++ 190 e.last = m.now() 191 db := e.o.db 192 m.mu.Unlock() 193 return db, m.releaser(e), nil 194} 195 196// releaser returns an idempotent unpin for one acquire. 197func (m *Manager) releaser(e *entry) func() { 198 var once sync.Once 199 return func() { 200 once.Do(func() { 201 m.mu.Lock() 202 if e.pins > 0 { 203 e.pins-- 204 } 205 e.last = m.now() 206 m.mu.Unlock() 207 }) 208 } 209} 210 211// evictLocked closes the least recently used UNPINNED pools while over the cap. A pinned pool (a live 212// request) is never closed, and the just inserted pool (protect) is never closed even though it is not 213// pinned yet, otherwise the caller would find it gone and retry forever. If nothing else is evictable the 214// registry stays temporarily over cap. Caller holds m.mu and MUST close the returned victims outside the 215// lock. 216func (m *Manager) evictLocked(protect uuid.UUID) []*entry { 217 var victims []*entry 218 for len(m.pools) > m.maxPools { 219 var lru *entry 220 for org, e := range m.pools { 221 if e.pins > 0 || org == protect { 222 continue 223 } 224 if lru == nil || e.last.Before(lru.last) { 225 lru = e 226 } 227 } 228 if lru == nil { 229 break // everything else is pinned or protected, accept temporary over cap 230 } 231 delete(m.pools, lru.org) 232 victims = append(victims, lru) 233 } 234 return victims 235} 236 237// sweepOnce closes whole pools idle past IdleEvict (separate from pgx's per connection idle close, which 238// only shrinks a pool toward MinConns, never frees the pool object). 239func (m *Manager) sweepOnce() { 240 now := m.now() 241 m.mu.Lock() 242 var victims []*entry 243 for org, e := range m.pools { 244 if e.pins == 0 && now.Sub(e.last) >= m.idle { 245 delete(m.pools, org) 246 victims = append(victims, e) 247 } 248 } 249 m.mu.Unlock() 250 closeAll(victims) 251} 252 253func (m *Manager) sweepLoop() { 254 interval := m.idle / 2 255 if interval < time.Second { 256 interval = time.Second 257 } 258 t := time.NewTicker(interval) 259 defer t.Stop() 260 for { 261 select { 262 case <-m.stop: 263 return 264 case <-t.C: 265 m.sweepOnce() 266 } 267 } 268} 269 270// Len reports the number of live tenant pools (for tests and metrics). 271func (m *Manager) Len() int { 272 m.mu.Lock() 273 defer m.mu.Unlock() 274 return len(m.pools) 275} 276 277// Close shuts the manager down and closes every live pool. It is idempotent. 278func (m *Manager) Close() { 279 m.mu.Lock() 280 if m.closed { 281 m.mu.Unlock() 282 return 283 } 284 m.closed = true 285 close(m.stop) 286 victims := make([]*entry, 0, len(m.pools)) 287 for org, e := range m.pools { 288 victims = append(victims, e) 289 delete(m.pools, org) 290 } 291 m.mu.Unlock() 292 closeAll(victims) 293} 294 295// closeAll closes evicted pools outside the registry lock (Close blocks until the pool drains, so it must 296// never run under m.mu). An evicted pool is unpinned, so its drain is fast. 297func closeAll(es []*entry) { 298 for _, e := range es { 299 e.o.close() 300 } 301} 302 303// defaultBuild resolves the tenant credential and opens and verifies the pool. On an authentication 304// failure (the credential rotation window) it forgets the cached secret and retries once with a fresh 305// resolve, per the AWS rotation guidance, before returning a typed unreachable error. 306func defaultBuild(sec secrets.Resolver, maxConns int32, role, refSuffix string) func(context.Context, directory.TenantConn) (*opened, error) { 307 return func(ctx context.Context, tc directory.TenantConn) (*opened, error) { 308 ref := tc.SecretRef + refSuffix 309 dsn, err := sec.ResolveDSN(ctx, ref) 310 if err != nil { 311 return nil, fmt.Errorf("%w: resolve credential: %v", ErrTenantUnreachable, err) 312 } 313 o, err := openVerify(ctx, dsn, maxConns, ref, role) 314 if err != nil && isAuthError(err) { 315 sec.Forget(ref) 316 if dsn2, e2 := sec.ResolveDSN(ctx, ref); e2 == nil { 317 o, err = openVerify(ctx, dsn2, maxConns, ref, role) 318 } 319 } 320 if err != nil { 321 return nil, fmt.Errorf("%w: %v", ErrTenantUnreachable, err) 322 } 323 return o, nil 324 } 325} 326 327// openVerify opens a bounded pool, Pings it under a deadline (pgxpool.New returns BEFORE connecting, so an 328// unreachable database looks open until first use), and asserts the expected NOBYPASSRLS role (the app 329// role, or the compliance role for the eeo registry). Any failure closes the pool and returns the error 330// so a dead pool is never cached. 331func openVerify(ctx context.Context, dsn string, maxConns int32, ref, role string) (*opened, error) { 332 pool, err := database.NewBoundedPool(ctx, dsn, maxConns) 333 if err != nil { 334 return nil, err 335 } 336 verifyCtx, cancel := context.WithTimeout(ctx, 5*time.Second) 337 defer cancel() 338 if err := pool.Ping(verifyCtx); err != nil { 339 pool.Close() 340 return nil, err 341 } 342 if err := database.AssertRole(verifyCtx, pool, role); err != nil { 343 pool.Close() 344 return nil, err 345 } 346 return &opened{db: tenant.New(pool), close: pool.Close, ref: ref}, nil 347} 348 349// isAuthError reports a Postgres authentication failure (invalid password or invalid authorization), the 350// rotation window signal that warrants a secret refresh and retry. 351func isAuthError(err error) bool { 352 var pg *pgconn.PgError 353 if errors.As(err, &pg) { 354 return pg.Code == "28P01" || pg.Code == "28000" 355 } 356 return false 357} 358