From e05f9f940ea0d5e205c44770a04ac90930687f7e Mon Sep 17 00:00:00 2001 From: Jared Lunde Date: Fri, 26 Jun 2026 16:16:33 -0700 Subject: [PATCH] perf(supervisor): fix extension filtering + batch post_start into one psql MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two compounding problems made post_start ~1.6s of mostly-waste: 1. `extension_installed` checked for `{ext}.so` in pkglibdir, but `CREATE EXTENSION ` actually needs `{name}.control`. So it mis-judged installability: `postgis` (lib is `postgis-3.so`) read as "not installed". Because the build-time bake (template.rs) filters by `extension_installed` but the runtime post_start did NOT, postgis was SKIPPED at bake then created FRESH every boot (~900ms+) — the real long pole, not the no-ops. 2. The OPTIONAL list carried entries that can never succeed (`auto_explain` is preload-only/no control file; `pgvectorscale`/`pg_jsonschema`/`pg_search` aren't built for noble+pg18) and a misnamed one (`pgvector` → the extension is `vector`). Each was a wasted psql round-trip failing "not available". Fixes: - `extension_installed` checks the control file (`/usr/share/postgresql/18/ extension/{ext}.control`) — the correct, bake-vs-runtime-consistent test. - OPTIONAL_EXTENSIONS = only entries with a control file: pg_stat_statements, pg_trgm, vector, pg_partman, hypopg, pg_repack, postgis. Now postgis is baked and runtime is a no-op; the dead ones are gone. - post_start assembles ONE SQL script (password + pgbouncer auth + filtered CREATE EXTENSIONs + conditional replicator/CDC) and runs it in a single `psql -f -` with ON_ERROR_STOP=1, replacing ~15 process spawns + reconnects. New `pg::psql_script` + `pg::alter_role_password_sql`; pgbouncer SQL moved to `sql::PGBOUNCER_AUTH_SQL`; script assembly extracted to a pure, unit-tested `build_post_start_script`. Tests: post_start_script_* (assembly + replicator/CDC conditionals) + full unit suite green. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/pg.rs | 66 +++++++++---- src/sql.rs | 24 +++++ src/supervisor.rs | 244 +++++++++++++++++++++++++++------------------- src/template.rs | 2 +- 4 files changed, 217 insertions(+), 119 deletions(-) diff --git a/src/pg.rs b/src/pg.rs index a0dfa74..88e5a47 100644 --- a/src/pg.rs +++ b/src/pg.rs @@ -149,25 +149,57 @@ pub async fn psql_env(sql: &str, env: &[(&str, &str)]) -> Result<(), PgError> { Err(PgError::NonZeroExit { code, stderr }) } -/// Set the superuser password using dollar-quoting to avoid SQL injection. -/// -/// Dollar-quoting avoids single-quote injection. The password is MMDS-controlled -/// but correct quoting is still the right habit. -pub async fn set_superuser_password(pw: &str) -> Result<(), PgError> { - set_role_password("postgres", pw).await +/// Build the dollar-quoted `ALTER ROLE … WITH PASSWORD` statement (no trailing +/// semicolon, so it composes into a larger script). Dollar-quoting guards against +/// single-quote injection; the `$_beyond_$` tag is unlikely to appear in a +/// password (if it does, this breaks — a documented unsupported MMDS edge case). +pub fn alter_role_password_sql(role: &str, pw: &str) -> String { + format!("ALTER ROLE {role} WITH PASSWORD $_beyond_${pw}$_beyond_$") } -/// Set a role's password using dollar-quoting (same injection guard as -/// `set_superuser_password`). `role` is a fixed internal identifier (never user -/// input); `pw` is MMDS-controlled and dollar-quoted. -pub async fn set_role_password(role: &str, pw: &str) -> Result<(), PgError> { - // Dollar-quote tag chosen to be unlikely to appear in a password. If the - // password itself contains `$_beyond_$`, this breaks — documented as an - // unsupported edge case in the MMDS field. - psql(&format!( - "ALTER ROLE {role} WITH PASSWORD $_beyond_${pw}$_beyond_$" - )) - .await +/// Run a multi-statement SQL script in a SINGLE `psql` invocation, aborting on +/// the first error (`ON_ERROR_STOP=1`). Collapses the per-statement process +/// spawn + reconnect overhead of calling [`psql`] once per statement — used by +/// post-start setup, which is a fixed batch of idempotent statements. +pub async fn psql_script(sql: &str) -> Result<(), PgError> { + use tokio::io::AsyncWriteExt as _; + let mut cmd = Command::new("psql"); + cmd.args([ + "-U", + POSTGRES_USER, + "-h", + PG_SOCKET_DIR, + "-p", + &PG_PORT.to_string(), + "-d", + "postgres", + "-v", + "ON_ERROR_STOP=1", + "-f", + "-", + ]) + .stdin(std::process::Stdio::piped()) + .stdout(std::process::Stdio::null()) + .stderr(std::process::Stdio::piped()); + // Same peer-auth requirement as `psql`: drop to the postgres OS user. + drop_to_postgres_user(&mut cmd); + + let mut child = cmd.spawn()?; + { + let mut stdin = child + .stdin + .take() + .ok_or_else(|| PgError::Io(std::io::Error::other("psql stdin unavailable")))?; + stdin.write_all(sql.as_bytes()).await?; + stdin.shutdown().await?; + } + let out = child.wait_with_output().await?; + if out.status.success() { + return Ok(()); + } + let code = out.status.code().ok_or(PgError::Signal)?; + let stderr = String::from_utf8_lossy(&out.stderr).into_owned(); + Err(PgError::NonZeroExit { code, stderr }) } // --------------------------------------------------------------------------- diff --git a/src/sql.rs b/src/sql.rs index 5103013..1defa55 100644 --- a/src/sql.rs +++ b/src/sql.rs @@ -31,3 +31,27 @@ BEGIN END IF; END $$"; + +/// pgbouncer auth: the `pgbouncer` login role + the SECURITY DEFINER lookup +/// function PgBouncer's `auth_query` calls, plus the schema USAGE the role needs +/// to call it. Idempotent (CREATE IF NOT EXISTS / OR REPLACE), and every +/// statement is `;`-terminated so it composes into the post-start batch script. +/// +/// USAGE on the schema is REQUIRED — without it `auth_query` fails with +/// "permission denied for schema pgbouncer" for every client and no one can +/// connect through the pooler (EXECUTE on the function alone is not enough). +pub const PGBOUNCER_AUTH_SQL: &str = "DO $$ +BEGIN + IF NOT EXISTS (SELECT FROM pg_roles WHERE rolname = 'pgbouncer') THEN + CREATE ROLE pgbouncer LOGIN PASSWORD NULL; + END IF; +END +$$; +CREATE SCHEMA IF NOT EXISTS pgbouncer; +CREATE OR REPLACE FUNCTION pgbouncer.get_auth(p_user text) + RETURNS TABLE(username text, password text) + SECURITY DEFINER LANGUAGE sql AS $$ + SELECT usename::text, passwd::text FROM pg_shadow WHERE usename = p_user + $$; +GRANT USAGE ON SCHEMA pgbouncer TO pgbouncer; +GRANT EXECUTE ON FUNCTION pgbouncer.get_auth(text) TO pgbouncer;"; diff --git a/src/supervisor.rs b/src/supervisor.rs index 4a69796..47b7f6c 100644 --- a/src/supervisor.rs +++ b/src/supervisor.rs @@ -1752,70 +1752,37 @@ async fn wait_for_sync_standby(name: &str, timeout: std::time::Duration) { // --------------------------------------------------------------------------- async fn post_start(cfg: &MmdsConfig) -> Result<(), Box> { - pg::set_superuser_password(&cfg.postgres_password) - .await - .map_err(|e| format!("failed to set postgres password: {e}"))?; - - setup_pgbouncer_auth() - .await - .map_err(|e| format!("failed to set up pgbouncer auth: {e}"))?; - - // Required extensions — fail the supervisor if CREATE EXTENSION fails for - // an extension whose shared object IS installed; the process manager will - // restart and retry rather than running in a degraded state. - // - // But the standalone postgres primitive ships without the auth/queue - // milestone (beyond_auth/beyond_queue) and pgdg's pg_cron (version pin - // drift). When an extension's .so isn't installed, treat it as a warning, - // not a fatal — the same self-adapting posture as the - // shared_preload_libraries filter (see config::beyond_conf). With the - // extensions present, the behavior is unchanged (still hard-required). + // Resolve the extensions to (idempotently) create: required + optional, + // filtered to those whose control file is present. A REQUIRED extension + // missing its control file is downgraded to a warning (the standalone image + // may ship without the auth/queue milestone) and dropped from the batch. + // Same filter the build-time bake (template.rs) uses ⇒ every CREATE EXTENSION + // is a no-op against the pre-baked cluster. + let mut extensions: Vec<&str> = Vec::new(); for ext in REQUIRED_EXTENSIONS { - if !extension_installed(ext) { + if extension_installed(ext) { + extensions.push(ext); + } else { warn!( - "required extension {ext} not installed (no {ext}.so in {EXTENSION_PKGLIBDIR}); \ - skipping (auth/queue milestone not in this image)" + "required extension {ext} not installed (no {ext}.control in \ + {EXTENSION_CONTROL_DIR}); skipping (auth/queue milestone not in this image)" ); - continue; - } - pg::psql(&format!("CREATE EXTENSION IF NOT EXISTS {ext}")) - .await - .map_err(|e| format!("required extension {ext} failed: {e}"))?; - } - - for ext in OPTIONAL_EXTENSIONS { - if let Err(e) = pg::psql(&format!("CREATE EXTENSION IF NOT EXISTS {ext}")).await { - warn!("optional extension {ext} failed: {e}"); - } - } - - // Create the `replicator` role when anything needs to stream WAL from us: a - // remote replica (replication_password set), a WAL sink, or CDC. Remote - // consumers authenticate via scram (pg_hba `host replication replicator`), - // so when a replication password is configured we set it — without it the - // role is PASSWORD NULL and scram can't succeed. - if cfg.replication_password.is_some() || cfg.wal_sink.is_some() || cfg.cdc_enabled { - pg::psql(crate::sql::REPLICATOR_ROLE_SQL) - .await - .map_err(|e| format!("failed to create replicator role: {e}"))?; - if let Some(pw) = &cfg.replication_password { - pg::set_role_password("replicator", pw) - .await - .map_err(|e| format!("failed to set replicator password: {e}"))?; } } + extensions.extend(OPTIONAL_EXTENSIONS.iter().copied().filter(|e| extension_installed(e))); if cfg.cdc_enabled { warn!("CDC enabled: replication slot 'cdc' will accumulate WAL until a consumer connects"); - pg::psql(crate::sql::CDC_SLOT_SQL) - .await - .map_err(|e| format!("failed to create CDC replication slot: {e}"))?; - - pg::psql(crate::sql::CDC_PUBLICATION_SQL) - .await - .map_err(|e| format!("failed to create CDC publication: {e}"))?; } + // ONE psql invocation for the whole batch (ON_ERROR_STOP) instead of ~15 + // process spawns + reconnects. Every statement is idempotent and, on a + // well-formed image, a no-op or fast role op — so a failure is a real problem + // and rightly aborts post_start (→ restart). + pg::psql_script(&build_post_start_script(cfg, &extensions)) + .await + .map_err(|e| format!("post-start setup failed: {e}"))?; + boot::run_hook_scripts("/etc/postgresql/18/hooks/post-start.d") .await .map_err(|e| format!("post-start hook failed: {e}"))?; @@ -1824,39 +1791,49 @@ async fn post_start(cfg: &MmdsConfig) -> Result<(), Box> Ok(()) } -async fn setup_pgbouncer_auth() -> Result<(), crate::pg::PgError> { - // Create the pgbouncer role and the SECURITY DEFINER auth lookup function. - // Idempotent: CREATE IF NOT EXISTS / CREATE OR REPLACE. - pg::psql( - "DO $$ - BEGIN - IF NOT EXISTS (SELECT FROM pg_roles WHERE rolname = 'pgbouncer') THEN - CREATE ROLE pgbouncer LOGIN PASSWORD NULL; - END IF; - END - $$", - ) - .await?; - - pg::psql("CREATE SCHEMA IF NOT EXISTS pgbouncer").await?; - - pg::psql( - "CREATE OR REPLACE FUNCTION pgbouncer.get_auth(p_user text) - RETURNS TABLE(username text, password text) - SECURITY DEFINER LANGUAGE sql AS $$ - SELECT usename::text, passwd::text FROM pg_shadow WHERE usename = p_user - $$", - ) - .await?; - - // USAGE on the schema is REQUIRED for the pgbouncer role to call get_auth at - // all — without it auth_query fails with "permission denied for schema - // pgbouncer" for every client, and no one can connect through the pooler. - // (EXECUTE on the function is not enough; the caller also needs schema USAGE.) - pg::psql("GRANT USAGE ON SCHEMA pgbouncer TO pgbouncer").await?; - pg::psql("GRANT EXECUTE ON FUNCTION pgbouncer.get_auth(text) TO pgbouncer").await?; +/// Build the single post-start SQL script (pure — no I/O, so it's unit-testable). +/// Order: superuser password → pgbouncer auth → `CREATE EXTENSION IF NOT EXISTS` +/// per `extensions` → (when WAL consumers exist) replicator role + password → +/// (when CDC enabled) slot + publication. Every statement is `;`-terminated so +/// the whole thing runs as one `psql -f -` with `ON_ERROR_STOP=1`. +fn build_post_start_script(cfg: &MmdsConfig, extensions: &[&str]) -> String { + let mut s = String::new(); - Ok(()) + // 1. Reset the superuser password from the per-instance MMDS secret. + s.push_str(&pg::alter_role_password_sql("postgres", &cfg.postgres_password)); + s.push_str(";\n"); + + // 2. pgbouncer auth (role + SECURITY DEFINER lookup function + grants). + s.push_str(crate::sql::PGBOUNCER_AUTH_SQL); + s.push('\n'); + + // 3. Extensions (already filtered to installable by the caller). + for ext in extensions { + s.push_str(&format!("CREATE EXTENSION IF NOT EXISTS {ext};\n")); + } + + // 4. replicator role (+ password) when a replica / WAL sink / CDC will stream + // from us. Remote consumers scram-auth (pg_hba `host replication + // replicator`); without a password the role is PASSWORD NULL and scram + // can't succeed. + if cfg.replication_password.is_some() || cfg.wal_sink.is_some() || cfg.cdc_enabled { + s.push_str(crate::sql::REPLICATOR_ROLE_SQL); + s.push_str(";\n"); + if let Some(pw) = &cfg.replication_password { + s.push_str(&pg::alter_role_password_sql("replicator", pw)); + s.push_str(";\n"); + } + } + + // 5. CDC slot + publication. + if cfg.cdc_enabled { + s.push_str(crate::sql::CDC_SLOT_SQL); + s.push_str(";\n"); + s.push_str(crate::sql::CDC_PUBLICATION_SQL); + s.push_str(";\n"); + } + + s } // beyond_auth is NOT auto-created here: the beyond-auth service installs its @@ -1865,32 +1842,43 @@ async fn setup_pgbouncer_auth() -> Result<(), crate::pg::PgError> { // an extension would collide with that migration ("function already exists"). pub(crate) const REQUIRED_EXTENSIONS: &[&str] = &["beyond_queue", "pg_cron"]; -/// Directory holding PostgreSQL extension shared objects (PG18 Debian layout, -/// `pg_config --pkglibdir`). Mirrors `config`'s PKGLIBDIR. -pub(crate) const EXTENSION_PKGLIBDIR: &str = "/usr/lib/postgresql/18/lib"; - -/// True iff the extension's shared object is present in the image. An extension -/// listed in [`REQUIRED_EXTENSIONS`] but with no installed `.so` (e.g. the -/// future auth/queue milestone, or a dropped pgdg `pg_cron`) is downgraded from -/// fatal to a warning so the standalone primitive still boots. +/// Directory holding PostgreSQL extension control files (PG18 Debian layout, +/// `pg_config --sharedir` + `/extension`). `CREATE EXTENSION ` resolves +/// `{name}.control` here. +pub(crate) const EXTENSION_CONTROL_DIR: &str = "/usr/share/postgresql/18/extension"; + +/// True iff `CREATE EXTENSION {ext}` can find it — i.e. its control file is +/// present. This is the CORRECT installability test, not the bare `.so`: +/// preload-only modules (`auto_explain`) ship a `.so` but NO control file, and +/// `postgis` ships `postgis-3.so` but `postgis.control`. Filtering on the control +/// file (vs the lib) is what lets the build-time bake (`template.rs`) and runtime +/// `post_start` agree on the SAME installable set, so every runtime CREATE +/// EXTENSION is a no-op against the pre-baked cluster. A REQUIRED extension whose +/// control file is absent is downgraded from fatal to a warning so the standalone +/// primitive still boots. pub(crate) fn extension_installed(ext: &str) -> bool { - std::path::Path::new(EXTENSION_PKGLIBDIR) - .join(format!("{ext}.so")) + std::path::Path::new(EXTENSION_CONTROL_DIR) + .join(format!("{ext}.control")) .exists() } +/// Extensions auto-created on boot when present (no-ops against the pre-baked +/// cluster). Every entry must be the extension's REAL name (the `.control` +/// basename) — e.g. pgvector's extension is `vector`. Entries that aren't built +/// for this image (no control file) are filtered out by [`extension_installed`]; +/// they're omitted here too so the list reflects what actually ships. Preload-only +/// modules (`auto_explain`, also `pg_stat_statements`) are loaded via +/// `shared_preload_libraries`, not `CREATE EXTENSION` — `pg_stat_statements` is +/// kept because it ALSO ships a control file (it exposes a view), but it's the +/// only such dual case. pub(crate) const OPTIONAL_EXTENSIONS: &[&str] = &[ "pg_stat_statements", - "auto_explain", "pg_trgm", - "pgvector", - "pgvectorscale", + "vector", // pgvector — the extension is named `vector` "pg_partman", - "pg_jsonschema", "hypopg", "pg_repack", "postgis", - "pg_search", ]; // --------------------------------------------------------------------------- @@ -1951,6 +1939,60 @@ fn parse_memtotal_kib(content: &str) -> Option { mod tests { use super::*; + // --- post_start batch script (pure) --- + + fn pg_cfg(replication_password: Option<&str>, wal_sink: Option<&str>, cdc: bool) -> MmdsConfig { + MmdsConfig { + pg_tier: crate::mmds::PgTier::Single, + ephemeral: false, + postgres_password: "s3cr3t".into(), + postgres_database: "postgres".into(), + wal_sink: wal_sink.map(str::to_owned), + cdc_enabled: cdc, + recovery_target_time: None, + primary_conninfo: None, + replication_password: replication_password.map(str::to_owned), + ram_bytes: 4 * 1024 * 1024 * 1024, + vcpus: 2, + } + } + + #[test] + fn post_start_script_is_one_terminated_batch() { + let s = build_post_start_script(&pg_cfg(None, None, false), &["vector", "postgis"]); + // superuser password (dollar-quoted) + pgbouncer auth, each terminated. + assert!(s.contains("ALTER ROLE postgres WITH PASSWORD $_beyond_$s3cr3t$_beyond_$;")); + assert!(s.contains("CREATE SCHEMA IF NOT EXISTS pgbouncer;")); + assert!(s.contains("GRANT EXECUTE ON FUNCTION pgbouncer.get_auth(text) TO pgbouncer;")); + // each filtered extension, semicolon-terminated. + assert!(s.contains("CREATE EXTENSION IF NOT EXISTS vector;")); + assert!(s.contains("CREATE EXTENSION IF NOT EXISTS postgis;")); + // nothing replication/CDC when not configured. + assert!(!s.contains("replicator")); + assert!(!s.contains("PUBLICATION cdc")); + // every statement is terminated: no trailing un-`;`d fragment. + assert!(s.trim_end().ends_with(';')); + } + + #[test] + fn post_start_script_adds_replicator_and_cdc_when_enabled() { + // WAL sink ⇒ replicator role (no password set); CDC ⇒ slot + publication. + let s = build_post_start_script(&pg_cfg(Some("rpw"), Some("http://sink:9000"), true), &[]); + assert!(s.contains("CREATE ROLE replicator LOGIN REPLICATION PASSWORD NULL")); + assert!(s.contains("ALTER ROLE replicator WITH PASSWORD $_beyond_$rpw$_beyond_$;")); + assert!(s.contains("pg_create_logical_replication_slot('cdc'")); + assert!(s.contains("CREATE PUBLICATION cdc")); + assert!(s.trim_end().ends_with(';')); + } + + #[test] + fn post_start_script_replicator_without_password_when_only_sink() { + // wal_sink alone creates the role but sets NO password (PASSWORD NULL). + let s = build_post_start_script(&pg_cfg(None, Some("http://sink:9000"), false), &[]); + assert!(s.contains("CREATE ROLE replicator")); + assert!(!s.contains("ALTER ROLE replicator WITH PASSWORD")); + } + // --- PgBouncer scaler decision (pure) --- const COOL: Duration = Duration::from_secs(3600); // past every cooldown diff --git a/src/template.rs b/src/template.rs index 0052caa..6554ad5 100644 --- a/src/template.rs +++ b/src/template.rs @@ -35,7 +35,7 @@ pub const TEMPLATE_DIR: &str = "/usr/local/share/beyond-pg/pgdata-template"; /// Throwaway superuser password used only while building the template. The /// runtime resets the superuser password from the MMDS secret on every boot -/// (`supervisor::post_start` → `pg::set_superuser_password`), so this value +/// (`supervisor::post_start` → the post_start password reset), so this value /// never reaches a running instance. const BUILD_PASSWORD: &str = "beyond-build-template";