From 3e24ef6145208cefc17e6819d0e35bd22388b146 Mon Sep 17 00:00:00 2001 From: Jared Lunde Date: Fri, 26 Jun 2026 09:38:21 -0700 Subject: [PATCH] feat(template): pre-bake initialized PGDATA to skip first-boot initdb MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A fresh Postgres primary ran `initdb` plus the full `CREATE EXTENSION` suite in the guest before accepting connections — ~1-2s+ squarely on the cold-boot critical path. Bake an already-initialized cluster (initdb + extensions) into the rootfs at image-build time and copy it onto the fresh data volume on first boot instead. - `beyond-pg build-template ` (new subcommand): runs the canonical `pg::initdb` flag set into /main, starts a build-time postgres with the same filtered `shared_preload_libraries`, and creates the same REQUIRED+OPTIONAL extension set the supervisor would (reusing its lists + `extension_installed` filter) — single source, so the baked cluster can't drift from what the runtime expects. Points at the shipped pg_hba (local peer) so the build-time psql authenticates without a password. - `maybe_initdb` materializes the template when present (atomic staging + rename, pg_wal symlink repointed to the runtime waldir, targeted fsync), else falls back to runtime initdb. PG_VERSION only appears once a complete, correct PGDATA is published, so the skip-on-PG_VERSION check stays safe and an interrupted materialize is simply redone next boot. - Per-instance state (superuser/replicator passwords, roles) stays OUT of the template; `post_start` applies it every boot, so no shared build-time secret is exposed and `CREATE EXTENSION IF NOT EXISTS` becomes a no-op on first boot. Verified: chroot build-template produces a 40M template (initdb + beyond_queue, pg_cron, pg_stat_statements, pg_trgm, hypopg, pg_repack) and a clean delta; unit tests cover materialize copy/symlink/idempotency/leftover-staging recovery. Co-Authored-By: Claude Opus 4.8 (1M context) --- packer/scripts/10-pgdata-template.sh | 30 +++ src/boot.rs | 262 ++++++++++++++++++++++++++- src/config.rs | 20 ++ src/main.rs | 30 +++ src/pg.rs | 20 +- src/supervisor.rs | 8 +- src/template.rs | 238 ++++++++++++++++++++++++ 7 files changed, 597 insertions(+), 11 deletions(-) create mode 100755 packer/scripts/10-pgdata-template.sh create mode 100644 src/template.rs diff --git a/packer/scripts/10-pgdata-template.sh b/packer/scripts/10-pgdata-template.sh new file mode 100755 index 0000000..a1e0ebb --- /dev/null +++ b/packer/scripts/10-pgdata-template.sh @@ -0,0 +1,30 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Pre-bake an initialized PGDATA template into the image. +# +# `beyond-pg build-template` runs `initdb` (the canonical runtime flag set) and +# then the full `CREATE EXTENSION` suite against a throwaway build-time postgres, +# leaving a cluster at /usr/local/share/beyond-pg/pgdata-template. At first boot, +# `beyond-pg-init` copies this onto the fresh data volume instead of running +# initdb + CREATE EXTENSION in the guest — taking both off the cold-boot path. +# +# Runs LAST: it needs the postgres server + every extension `.so` installed +# (02/03/04) and the beyond-pg binary in place (06). Building on the cleaned +# image (post-09) mirrors the runtime environment exactly (same trimmed locales, +# so the en_US.UTF-8 initdb locale resolves identically). +# +# Per-instance state is NOT baked in (superuser/replicator passwords, roles) — +# the supervisor's post_start applies those on every boot. + +TEMPLATE_DIR="/usr/local/share/beyond-pg/pgdata-template" + +echo "==> Building pre-initialized PGDATA template at ${TEMPLATE_DIR}..." +/usr/local/bin/beyond-pg build-template "${TEMPLATE_DIR}" + +# Sanity: the template must carry a complete cluster (PG_VERSION present) or the +# first-boot materialize would silently fall back to runtime initdb. +test -f "${TEMPLATE_DIR}/main/PG_VERSION" \ + || { echo "FATAL: template missing PG_VERSION" >&2; exit 1; } + +echo "==> 10-pgdata-template done ($(du -sh "${TEMPLATE_DIR}" | cut -f1))" diff --git a/src/boot.rs b/src/boot.rs index 9d3cd0d..69655a7 100644 --- a/src/boot.rs +++ b/src/boot.rs @@ -18,7 +18,7 @@ use crate::mmds::{MmdsConfig, MmdsError, PgTier}; use crate::pg::{self, PGDATA}; const PG_WAL_LINK: &str = "/var/lib/postgresql/18/main/pg_wal"; -const PG_WAL_TARGET: &str = "/var/lib/postgresql/18/wal"; +const PG_WAL_TARGET: &str = crate::pg::PG_WALDIR; const HOOKS_PRE_START: &str = "/etc/postgresql/18/hooks/pre-start.d"; #[derive(Debug, thiserror::Error)] @@ -216,10 +216,146 @@ async fn maybe_initdb(cfg: &MmdsConfig) -> Result<(), BootError> { chown_data_tree(); + // Fast path: a pre-baked PGDATA template shipped in the rootfs (image build + // ran `beyond-pg build-template`). Copy it onto the fresh volume instead of + // running initdb (+ first-boot CREATE EXTENSION) in the guest — both come off + // the cold-boot critical path. The supervisor's `post_start` still resets the + // per-instance superuser password and runs idempotent `CREATE EXTENSION IF NOT + // EXISTS` (no-ops against the baked extensions), so the result is identical to + // a runtime initdb. + if template_available() { + info!("materializing PGDATA from template, skipping initdb"); + materialize_template()?; + return Ok(()); + } + run_initdb(&cfg.postgres_password).await?; Ok(()) } +/// True iff a complete pre-baked PGDATA template is present in the rootfs. +fn template_available() -> bool { + template_available_in(crate::template::TEMPLATE_DIR) +} + +fn template_available_in(template_dir: &str) -> bool { + Path::new(&format!("{template_dir}/main/PG_VERSION")).exists() +} + +/// Copy the baked PGDATA template onto the fresh data volume, then chown it. +fn materialize_template() -> Result<(), BootError> { + materialize_template_into(crate::template::TEMPLATE_DIR, PGDATA, PG_WAL_TARGET)?; + chown_data_tree(); + info!("materialized PGDATA from template"); + Ok(()) +} + +/// Copy the template at `template_dir` (`main/` + `wal/`) onto `pgdata` + `wal_target`. +/// +/// Atomic + idempotent (per CLAUDE.md): the template is copied into staging dirs +/// alongside the destinations, the `pg_wal` symlink is repointed to `wal_target`, +/// the staged data is flushed, and only then are the staging dirs renamed into +/// place. `pgdata/PG_VERSION` therefore becomes visible only once a complete, +/// correct PGDATA is published — so `maybe_initdb`'s skip-on-`PG_VERSION` is safe +/// and an interrupted materialize is simply redone on the next boot. +fn materialize_template_into( + template_dir: &str, + pgdata: &str, + wal_target: &str, +) -> Result<(), BootError> { + let tmpl_main = format!("{template_dir}/main"); + let tmpl_wal = format!("{template_dir}/wal"); + let main_staging = format!("{pgdata}.staging"); + let wal_staging = format!("{wal_target}.staging"); + + // Idempotent clean slate. `pgdata` may exist empty (partial-PGDATA cleanup + // recreated it above); a prior interrupted materialize may have left a wal + // dir or *.staging dirs. + for p in [&main_staging, &wal_staging, &pgdata.to_string(), &wal_target.to_string()] { + rm_rf(p)?; + } + + // `cp -a` preserves the pg_wal symlink, file modes, and postgres ownership. + cp_a(&tmpl_wal, &wal_staging)?; + cp_a(&tmpl_main, &main_staging)?; + + // Repoint the baked `pg_wal` symlink (template-relative) to the runtime WAL + // dir BEFORE publishing, so a crash after the rename never leaves a wrong + // target that `verify_wal_symlink` would reject. + let staged_link = format!("{main_staging}/pg_wal"); + match std::fs::symlink_metadata(&staged_link) { + Ok(_) => std::fs::remove_file(&staged_link)?, + Err(e) if e.kind() == std::io::ErrorKind::NotFound => {} + Err(e) => return Err(BootError::Io(e)), + } + std::os::unix::fs::symlink(wal_target, &staged_link)?; + + // Flush the staged tree (file contents + dir entries) so the data is durable + // BEFORE it is published. Then rename atomically — WAL first, then PGDATA, so + // `PG_VERSION` only appears once `main` is renamed, by which point both the + // symlink and the data are correct and durable. Finally fsync the parent dir + // to make the renames themselves durable. Targeted fsync (not a global + // sync(2)) keeps this fast and non-blocking under concurrent I/O. + fsync_tree(Path::new(&wal_staging))?; + fsync_tree(Path::new(&main_staging))?; + std::fs::rename(&wal_staging, wal_target)?; + std::fs::rename(&main_staging, pgdata)?; + if let Some(parent) = Path::new(pgdata).parent() { + fsync_dir(parent)?; + } + + Ok(()) +} + +/// Recursively fsync every file and directory under `path` (depth-first, so each +/// directory is fsynced after its entries). Symlinks are not followed — the +/// containing directory's fsync makes the symlink entry durable. +fn fsync_tree(path: &Path) -> Result<(), BootError> { + let meta = std::fs::symlink_metadata(path)?; + if meta.is_dir() { + for entry in std::fs::read_dir(path)? { + fsync_tree(&entry?.path())?; + } + fsync_dir(path)?; + } else if meta.is_file() { + std::fs::File::open(path)?.sync_all()?; + } + Ok(()) +} + +/// fsync a directory (durably commit its entries). On Linux a directory can be +/// opened read-only and `fsync`'d via `sync_all`. +fn fsync_dir(path: &Path) -> Result<(), BootError> { + std::fs::File::open(path)?.sync_all()?; + Ok(()) +} + +/// `rm -rf` a path (file, symlink, or directory). Idempotent — a missing path is +/// not an error. +fn rm_rf(path: &str) -> Result<(), BootError> { + match std::fs::symlink_metadata(path) { + Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(()), + Err(e) => Err(BootError::Io(e)), + Ok(meta) if meta.is_dir() => Ok(std::fs::remove_dir_all(path)?), + Ok(_) => Ok(std::fs::remove_file(path)?), + } +} + +/// `cp -a src dst` — recursive copy preserving symlinks, modes, and ownership. +fn cp_a(src: &str, dst: &str) -> Result<(), BootError> { + let status = std::process::Command::new("cp") + .args(["-a", src, dst]) + .status()?; + if status.success() { + Ok(()) + } else { + Err(BootError::Io(std::io::Error::other(format!( + "cp -a {src} {dst} failed: {status}" + )))) + } +} + + /// chown `/var/lib/postgresql` → postgres recursively. A fresh durable volume /// mounts root-owned over the image dir, and a root-run `pg_basebackup` (replica /// seeding) writes a root-owned PGDATA — either way postgres dies at startup with @@ -258,7 +394,7 @@ async fn run_initdb(password: &str) -> Result<(), BootError> { .to_str() .ok_or_else(|| BootError::Io(std::io::Error::other("tempfile path is not UTF-8")))?; info!("running initdb"); - pg::initdb(PGDATA, path_str) + pg::initdb(PGDATA, PG_WAL_TARGET, path_str) .await .map_err(|e| BootError::InitDb(e.to_string())) // pwfile is dropped here — tempfile removes it from disk @@ -465,6 +601,128 @@ mod tests { ); } + // ----------------------------------------------------------------------- + // PGDATA template materialize unit tests + // ----------------------------------------------------------------------- + + /// Build a minimal fake template at `dir`: `main/` with PG_VERSION + a + /// template-relative `pg_wal` symlink, and `wal/` with a segment file. + fn make_fake_template(dir: &Path) { + let tmain = dir.join("main"); + let twal = dir.join("wal"); + std::fs::create_dir_all(&tmain).unwrap(); + std::fs::create_dir_all(&twal).unwrap(); + std::fs::write(tmain.join("PG_VERSION"), "18\n").unwrap(); + std::fs::write(tmain.join("postgresql.conf"), "# initdb default\n").unwrap(); + std::fs::write(twal.join("000000010000000000000001"), b"wal-seg").unwrap(); + // Baked symlink points template-relative — WRONG for runtime; the + // materialize must repoint it to the real wal target. + std::os::unix::fs::symlink(&twal, tmain.join("pg_wal")).unwrap(); + } + + #[test] + fn template_available_detects_pg_version() { + let tmp = tempfile::tempdir().unwrap(); + let dir = tmp.path().join("template"); + assert!(!template_available_in(dir.to_str().unwrap())); + std::fs::create_dir_all(dir.join("main")).unwrap(); + assert!( + !template_available_in(dir.to_str().unwrap()), + "empty main/ is not a usable template" + ); + std::fs::write(dir.join("main/PG_VERSION"), "18").unwrap(); + assert!(template_available_in(dir.to_str().unwrap())); + } + + #[test] + fn materialize_copies_template_and_repoints_wal_symlink() { + let tmp = tempfile::tempdir().unwrap(); + let template = tmp.path().join("template"); + make_fake_template(&template); + + let data = tmp.path().join("data"); + std::fs::create_dir_all(&data).unwrap(); + let pgdata = data.join("main"); + let wal_target = data.join("wal"); + + materialize_template_into( + template.to_str().unwrap(), + pgdata.to_str().unwrap(), + wal_target.to_str().unwrap(), + ) + .unwrap(); + + // PGDATA + WAL contents present. + assert!(pgdata.join("PG_VERSION").exists()); + assert!(pgdata.join("postgresql.conf").exists()); + assert!(wal_target.join("000000010000000000000001").exists()); + // pg_wal repointed to the absolute runtime WAL target (what + // verify_wal_symlink expects), not the template-relative path. + assert_eq!( + std::fs::read_link(pgdata.join("pg_wal")).unwrap(), + wal_target + ); + // No staging dirs left behind. + assert!(!data.join("main.staging").exists()); + assert!(!data.join("wal.staging").exists()); + } + + #[test] + fn materialize_is_idempotent() { + let tmp = tempfile::tempdir().unwrap(); + let template = tmp.path().join("template"); + make_fake_template(&template); + let data = tmp.path().join("data"); + std::fs::create_dir_all(&data).unwrap(); + let pgdata = data.join("main"); + let wal_target = data.join("wal"); + + let run = || { + materialize_template_into( + template.to_str().unwrap(), + pgdata.to_str().unwrap(), + wal_target.to_str().unwrap(), + ) + }; + run().unwrap(); + // A second materialize over a published PGDATA succeeds (clean-slate + // removes the prior copy) and yields the same correct layout. + run().unwrap(); + assert!(pgdata.join("PG_VERSION").exists()); + assert_eq!( + std::fs::read_link(pgdata.join("pg_wal")).unwrap(), + wal_target + ); + } + + #[test] + fn materialize_recovers_from_leftover_staging() { + let tmp = tempfile::tempdir().unwrap(); + let template = tmp.path().join("template"); + make_fake_template(&template); + let data = tmp.path().join("data"); + std::fs::create_dir_all(&data).unwrap(); + let pgdata = data.join("main"); + let wal_target = data.join("wal"); + + // Simulate an interrupted prior run: stale staging dirs + an empty PGDATA + // (as partial-PGDATA cleanup would have recreated). + std::fs::create_dir_all(data.join("main.staging")).unwrap(); + std::fs::write(data.join("main.staging/garbage"), b"x").unwrap(); + std::fs::create_dir_all(data.join("wal.staging")).unwrap(); + std::fs::create_dir_all(&pgdata).unwrap(); + + materialize_template_into( + template.to_str().unwrap(), + pgdata.to_str().unwrap(), + wal_target.to_str().unwrap(), + ) + .unwrap(); + + assert!(pgdata.join("PG_VERSION").exists()); + assert!(!pgdata.join("garbage").exists(), "stale staging must not leak in"); + } + // ----------------------------------------------------------------------- // run_hook_scripts unit tests // ----------------------------------------------------------------------- diff --git a/src/config.rs b/src/config.rs index 6ef08e9..49bbb8a 100644 --- a/src/config.rs +++ b/src/config.rs @@ -45,6 +45,26 @@ pub fn beyond_conf() -> String { filter_shared_preload_libraries(BEYOND_CONF, PKGLIBDIR) } +/// The `shared_preload_libraries` value (bare comma list, no quotes) from +/// `00-beyond.conf`, filtered to the libraries actually installed. The template +/// builder passes this to a build-time postgres (`-c shared_preload_libraries=…`) +/// so `CREATE EXTENSION pg_cron` / `beyond_queue` — which refuse to load unless +/// preloaded — succeed against the same set the runtime preloads. Empty string +/// if `00-beyond.conf` lists no installed preload libraries. +pub fn preload_libraries() -> String { + for line in BEYOND_CONF.lines() { + if let Some(rewritten) = filter_preload_line(line, "shared_preload_libraries", PKGLIBDIR) { + // `rewritten` is `shared_preload_libraries = '...'` — pull out the + // single-quoted value. + return rewritten + .split_once('=') + .map(|(_, v)| v.trim().trim_matches('\'').to_string()) + .unwrap_or_default(); + } + } + String::new() +} + /// Returns true iff `{pkglibdir}/{lib}.so` exists. Core-postgres libraries /// (`pg_stat_statements`, `auto_explain`) are present in any standard install. fn library_installed(pkglibdir: &str, lib: &str) -> bool { diff --git a/src/main.rs b/src/main.rs index b1435ec..0b665b8 100644 --- a/src/main.rs +++ b/src/main.rs @@ -12,6 +12,7 @@ mod pg; mod rpc; mod sql; mod supervisor; +mod template; mod tls; mod vsock; mod wal_forwarder; @@ -23,6 +24,14 @@ enum Cmd { Supervisor, /// Run idempotent boot setup and exit Boot, + /// Build a pre-initialized PGDATA template (initdb + extensions) at . + /// Run at image-build time; baked into the rootfs and copied onto the data + /// volume on first boot instead of running initdb in the guest. + BuildTemplate { + /// Output directory for the template (default: the rootfs template dir). + #[arg(default_value = template::TEMPLATE_DIR)] + dir: String, + }, } fn main() { @@ -56,5 +65,26 @@ fn main() { .expect("tokio runtime") .block_on(boot::run()); } + Cmd::BuildTemplate { dir } => { + // The Debian postgresql server binaries (initdb, postgres) live in + // /usr/lib/postgresql/18/bin and are NOT on the default PATH — at + // runtime PID 1 puts them there, but the image-build chroot does not. + // Set it before the runtime starts (still single-threaded). + // SAFETY: main thread, before any tokio worker exists. + unsafe { + std::env::set_var( + "PATH", + format!( + "/usr/lib/postgresql/18/bin:{}", + std::env::var("PATH").unwrap_or_default() + ), + ); + } + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("tokio runtime") + .block_on(template::run(&dir)); + } } } diff --git a/src/pg.rs b/src/pg.rs index 05aac46..a0dfa74 100644 --- a/src/pg.rs +++ b/src/pg.rs @@ -10,7 +10,7 @@ use tracing::{debug, warn}; /// pg_hba.conf uses peer auth for Unix socket connections; all psql calls must /// run as the postgres OS user so peer auth matches the "postgres" database user. /// No-op on non-Linux (macOS dev uses host postgres which handles auth itself). -fn drop_to_postgres_user( +pub(crate) fn drop_to_postgres_user( #[cfg_attr(not(target_os = "linux"), allow(unused_variables))] cmd: &mut Command, ) { #[cfg(target_os = "linux")] @@ -27,6 +27,10 @@ fn drop_to_postgres_user( } pub const PGDATA: &str = "/var/lib/postgresql/18/main"; +/// Runtime WAL directory. `initdb --waldir` points here and the cluster's +/// `main/pg_wal` symlink targets it. Kept off `PGDATA` so the WAL has its own +/// path on the data volume (matched by `boot::PG_WAL_TARGET`). +pub const PG_WALDIR: &str = "/var/lib/postgresql/18/wal"; pub const PG_SOCKET_DIR: &str = "/var/run/postgresql"; pub const PG_PORT: u16 = 5433; // Postgres direct; PgBouncer is 5432 (separate process) pub const POSTGRES_USER: &str = "postgres"; @@ -194,15 +198,21 @@ pub async fn reload() -> Result<(), PgError> { /// Run `initdb` to initialize a new database cluster. /// /// `pwfile_path` must point to a `0o600` file containing only the superuser -/// password (created by the caller as a `tempfile::NamedTempFile`). -pub async fn initdb(pgdata: &str, pwfile_path: &str) -> Result<(), PgError> { - debug!("running initdb in {pgdata}"); +/// password (created by the caller as a `tempfile::NamedTempFile`). `waldir` is +/// where the WAL lives (runtime: [`PG_WALDIR`]; the template builder targets a +/// template-relative dir so the baked `pg_wal` symlink is repointed on copy). +/// +/// This is the single source of the canonical `initdb` flag set — both the +/// runtime first-boot path and the build-time template builder call it, so the +/// pre-baked cluster can never drift from what the runtime would have produced. +pub async fn initdb(pgdata: &str, waldir: &str, pwfile_path: &str) -> Result<(), PgError> { + debug!("running initdb in {pgdata} (waldir {waldir})"); let mut cmd = Command::new("initdb"); cmd.args([ "-D", pgdata, "--waldir", - "/var/lib/postgresql/18/wal", + waldir, "--auth=scram-sha-256", "--encoding=UTF8", "--locale=en_US.UTF-8", diff --git a/src/supervisor.rs b/src/supervisor.rs index d0dc600..4a69796 100644 --- a/src/supervisor.rs +++ b/src/supervisor.rs @@ -1863,23 +1863,23 @@ async fn setup_pgbouncer_auth() -> Result<(), crate::pg::PgError> { // authz_check* C functions itself via a raw `CREATE FUNCTION ... AS 'beyond_auth'` // migration (loading the shipped beyond_auth.so on demand). Auto-creating it as // an extension would collide with that migration ("function already exists"). -const REQUIRED_EXTENSIONS: &[&str] = &["beyond_queue", "pg_cron"]; +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. -const EXTENSION_PKGLIBDIR: &str = "/usr/lib/postgresql/18/lib"; +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. -fn extension_installed(ext: &str) -> bool { +pub(crate) fn extension_installed(ext: &str) -> bool { std::path::Path::new(EXTENSION_PKGLIBDIR) .join(format!("{ext}.so")) .exists() } -const OPTIONAL_EXTENSIONS: &[&str] = &[ +pub(crate) const OPTIONAL_EXTENSIONS: &[&str] = &[ "pg_stat_statements", "auto_explain", "pg_trgm", diff --git a/src/template.rs b/src/template.rs new file mode 100644 index 0000000..0052caa --- /dev/null +++ b/src/template.rs @@ -0,0 +1,238 @@ +//! Build-time PGDATA template builder (`beyond-pg build-template `). +//! +//! Produces an already-initialized PostgreSQL cluster — `initdb` output PLUS the +//! full `CREATE EXTENSION` suite — baked into the rootfs at image-build time. At +//! first boot, [`crate::boot::maybe_initdb`] copies this template onto the fresh +//! data volume instead of running `initdb` + `CREATE EXTENSION` in the guest, +//! removing both from the cold-boot critical path. +//! +//! The template layout mirrors what the runtime expects on the data volume: +//! +//! ```text +//! /main → PGDATA contents (becomes /var/lib/postgresql/18/main) +//! /wal → WAL contents (becomes /var/lib/postgresql/18/wal) +//! ``` +//! +//! `/main/pg_wal` is an absolute symlink to `/wal` here; the materialize +//! step repoints it to the runtime [`crate::pg::PG_WALDIR`] after copying. +//! +//! Per-instance state is intentionally NOT baked in: the superuser/replicator +//! passwords and the pgbouncer/replicator roles are applied on every boot by the +//! supervisor's `post_start`, so a shared build-time password is never exposed. + +use std::os::unix::fs::PermissionsExt; +use std::path::Path; +use std::time::Duration; + +use tokio::process::Command; +use tracing::{info, warn}; + +use crate::pg; + +/// Where the baked template lives in the rootfs. Outside `/var/lib/postgresql` +/// (which the data volume shadows at runtime) so it survives into the image. +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 +/// never reaches a running instance. +const BUILD_PASSWORD: &str = "beyond-build-template"; + +#[derive(Debug, thiserror::Error)] +pub enum TemplateError { + #[error("I/O error: {0}")] + Io(#[from] std::io::Error), + #[error("initdb failed: {0}")] + InitDb(String), + #[error("postgres did not become ready within {0:?}")] + NotReady(Duration), + #[error("required extension {ext} failed: {source}")] + RequiredExtension { ext: String, source: pg::PgError }, + #[error("postgres shutdown failed: {0}")] + Shutdown(String), +} + +/// Entry point for `beyond-pg build-template `. +pub async fn run(dir: &str) { + tracing_subscriber::fmt() + .with_env_filter( + tracing_subscriber::EnvFilter::from_default_env() + .add_directive(tracing::Level::INFO.into()), + ) + .init(); + + if let Err(e) = build(dir).await { + eprintln!("[build-template] FATAL: {e}"); + std::process::exit(1); + } + info!("build-template complete: {dir}"); +} + +/// Build the PGDATA template at `dir`. Idempotent: removes any prior contents +/// and rebuilds from scratch. +pub async fn build(dir: &str) -> Result<(), TemplateError> { + let main = format!("{dir}/main"); + let wal = format!("{dir}/wal"); + + // Clean rebuild — initdb refuses a non-empty -D / --waldir. + if Path::new(dir).exists() { + std::fs::remove_dir_all(dir)?; + } + std::fs::create_dir_all(dir)?; + chown_postgres(dir); + + // initdb (and the build-time postgres) run as the postgres OS user; the + // socket dir must exist and be postgres-writable, same as runtime boot. + ensure_socket_dir(); + + info!("build-template: initdb → {main} (waldir {wal})"); + run_initdb(&main, &wal).await?; + // initdb creates the cluster as postgres; keep the whole tree postgres-owned. + chown_postgres(dir); + + info!("build-template: starting postgres to create extensions"); + let mut child = spawn_build_postgres(&main)?; + + let extensions_result = create_extensions().await; + + // Always attempt a clean (fast) shutdown so the template carries a + // shutdown checkpoint, even if extension creation failed. + let shutdown_result = shutdown_postgres(&mut child).await; + + extensions_result?; + shutdown_result?; + + chown_postgres(dir); + info!("build-template: cluster + extensions baked at {dir}"); + Ok(()) +} + +async fn run_initdb(main: &str, wal: &str) -> Result<(), TemplateError> { + let pwfile = tempfile::Builder::new() + .prefix("pg-build-pwfile-") + .tempfile_in("/run/") + .map_err(TemplateError::Io)?; + std::fs::set_permissions(pwfile.path(), std::fs::Permissions::from_mode(0o600))?; + std::fs::write(pwfile.path(), BUILD_PASSWORD)?; + let _ = std::process::Command::new("chown") + .args(["postgres:postgres", pwfile.path().to_str().unwrap_or("")]) + .status(); + + let path_str = pwfile + .path() + .to_str() + .ok_or_else(|| TemplateError::Io(std::io::Error::other("pwfile path is not UTF-8")))?; + + // Same canonical flag set as the runtime first-boot path (single source in + // pg::initdb), but WAL is template-relative — the symlink is repointed on copy. + pg::initdb(main, wal, path_str) + .await + .map_err(|e| TemplateError::InitDb(e.to_string())) +} + +/// Spawn a build-time postgres on the unix socket only, with the same filtered +/// `shared_preload_libraries` the runtime uses (so `pg_cron` / `beyond_queue` +/// `CREATE EXTENSION` succeed). TLS is off (no certs at build time) and there is +/// no TCP listener — this postgres is reachable only via the local socket that +/// `pg::psql` targets. +fn spawn_build_postgres(main: &str) -> Result { + let preload = crate::config::preload_libraries(); + info!("build-template: shared_preload_libraries='{preload}'"); + let port = pg::PG_PORT.to_string(); + + let mut cmd = Command::new("postgres"); + cmd.arg("-D").arg(main); + for (k, v) in [ + ("shared_preload_libraries", preload.as_str()), + ("listen_addresses", ""), + ("ssl", "off"), + ("cron.database_name", "postgres"), + ("logging_collector", "off"), + ("port", port.as_str()), + ("unix_socket_directories", pg::PG_SOCKET_DIR), + // `initdb --auth=scram-sha-256` made the template's own pg_hba require a + // password for local socket connections; point at the shipped hba (same + // file the runtime uses) whose `local all all peer` lets the build-time + // `psql` authenticate as the postgres OS user without a password. + ("hba_file", crate::config::PG_HBA_PATH), + ] { + cmd.arg("-c").arg(format!("{k}={v}")); + } + cmd.stdout(std::process::Stdio::inherit()) + .stderr(std::process::Stdio::inherit()) + .stdin(std::process::Stdio::null()); + // postgres refuses to run as root; drop to the postgres OS user. + pg::drop_to_postgres_user(&mut cmd); + cmd.spawn().map_err(TemplateError::Io) +} + +async fn create_extensions() -> Result<(), TemplateError> { + let timeout = Duration::from_secs(60); + if !pg::wait_until_ready(timeout).await { + return Err(TemplateError::NotReady(timeout)); + } + + // Mirror supervisor::post_start exactly: required extensions are fatal when + // their `.so` is installed; optional extensions are best-effort. Reusing the + // same lists + `extension_installed` filter keeps the baked set identical to + // what the runtime would auto-create (so post_start's IF-NOT-EXISTS calls + // become no-ops on first boot). + for ext in crate::supervisor::REQUIRED_EXTENSIONS { + if !crate::supervisor::extension_installed(ext) { + warn!("build-template: required extension {ext} not installed; skipping"); + continue; + } + pg::psql(&format!("CREATE EXTENSION IF NOT EXISTS {ext}")) + .await + .map_err(|e| TemplateError::RequiredExtension { + ext: (*ext).to_string(), + source: e, + })?; + info!("build-template: created extension {ext}"); + } + + for ext in crate::supervisor::OPTIONAL_EXTENSIONS { + if !crate::supervisor::extension_installed(ext) { + continue; + } + match pg::psql(&format!("CREATE EXTENSION IF NOT EXISTS {ext}")).await { + Ok(()) => info!("build-template: created extension {ext}"), + Err(e) => warn!("build-template: optional extension {ext} failed: {e}"), + } + } + Ok(()) +} + +/// Fast-shutdown the build postgres (SIGINT → checkpoint + clean stop) and wait +/// for it to exit, so the template carries a clean shutdown state. +async fn shutdown_postgres(child: &mut tokio::process::Child) -> Result<(), TemplateError> { + if let Some(pid) = child.id() { + // SIGINT = "fast" shutdown: roll back open txns, checkpoint, exit cleanly. + // SAFETY: kill(2) with a known child pid and a fixed signal number. + unsafe { + libc::kill(pid as libc::pid_t, libc::SIGINT); + } + } + let status = child.wait().await.map_err(TemplateError::Io)?; + if !status.success() { + return Err(TemplateError::Shutdown(format!("exit status {status}"))); + } + Ok(()) +} + +/// `mkdir -p` + `chown postgres` the unix-socket dir (same as `boot::ensure_socket_dir`). +fn ensure_socket_dir() { + let dir = pg::PG_SOCKET_DIR; + if let Err(e) = std::fs::create_dir_all(dir) { + warn!("build-template: create socket dir {dir}: {e}"); + return; + } + chown_postgres(dir); +} + +fn chown_postgres(path: &str) { + let _ = std::process::Command::new("chown") + .args(["-R", "postgres:postgres", path]) + .status(); +}