From 0a657563dc5f2e62da4d8e98cbd17809f3bec65d Mon Sep 17 00:00:00 2001 From: Jared Lunde Date: Sat, 11 Jul 2026 10:38:34 -0700 Subject: [PATCH 1/2] feat(pgbouncer): peer so_reuseport workers to forward cross-worker query cancels MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When beyond-pg runs more than one so_reuseport pooler worker (PgbScaler scales 1..N on many-vCPU boxes), a Postgres CancelRequest arrives on a fresh TCP connection the kernel may route to a worker that isn't holding the session. PgBouncer had so_reuseport on but no peering configured, so those misrouted cancels were silently dropped — Ctrl-C / statement-timeout / PQcancel would no-op on a busy box with probability (N-1)/N. Configure PgBouncer peering: each worker gets a stable peer_id, a distinct unix_socket_dir, and a shared [peers] map pre-declaring all workers (1..=cap) so the receiving worker forwards a cancel to the peer that owns the session. The map is pre-declared for the full cap so the scaler never rewrites config as it grows/shrinks; forwarding to an idle (not-running) peer just fails, exactly as a dropped cancel would have before, so pre-declaring is safe. - config.rs: pgbouncer_ini takes a peer_id and emits unix_socket_dir/peer_id/ [peers] when max_workers > 1; single-worker boxes keep the empty unix_socket_dir (no worker can misroute a cancel). pgb_peer_id / pgb_peer_socket_dir / pgbouncer_ini_path helpers; peer_id 1 keeps the canonical pgbouncer.ini. - boot.rs: write one config per possible worker and create the per-worker peer socket dirs (owned by postgres) up front. - supervisor.rs: spawn_pgbouncer selects the per-worker config by stable name. - tests/pgbouncer_peering.rs: e2e proof against real pgbouncer 1.25 — a cancel sent to the wrong worker cancels the query with peering ON and is dropped with peering OFF (control). mise test:peering wires it into CI. Verified empirically against pgbouncer 1.25.2: peering ON -> misrouted cancel cancels the query; peering OFF -> query runs to completion (cancel dropped). Co-Authored-By: Claude Opus 4.8 (1M context) --- README.md | 2 +- mise.toml | 4 + packer/files/pgbouncer/pgbouncer.ini | 11 +- src/boot.rs | 46 +++- src/config.rs | 163 +++++++++++- src/supervisor.rs | 12 +- tests/pgbouncer_peering.rs | 368 +++++++++++++++++++++++++++ 7 files changed, 577 insertions(+), 29 deletions(-) create mode 100644 tests/pgbouncer_peering.rs diff --git a/README.md b/README.md index 6e06339..fd11991 100644 --- a/README.md +++ b/README.md @@ -17,7 +17,7 @@ Forks boot in seconds against a CoW snapshot of `/var/lib/postgresql`. Postgres ## What it does - **Forks with the substrate.** Every byte under `/var/lib/postgresql` snapshots atomically; the new box boots on a CoW copy. No `pg_dump`, no replication topology, no wait. -- **PgBouncer on the front, sized to load.** Transaction pooling and TLS termination on `:5432`; Postgres listens on the loopback only. The pooler runs one worker when idle and adds `so_reuseport` workers across cores as connection-handshake load rises, then reaps them when it falls. The worker count survives restarts. +- **PgBouncer on the front, sized to load.** Transaction pooling and TLS termination on `:5432`; Postgres listens on the loopback only. The pooler runs one worker when idle and adds `so_reuseport` workers across cores as connection-handshake load rises, then reaps them when it falls. The worker count survives restarts. Workers peer, so a query-cancel the kernel routes to the wrong worker is forwarded to the one holding the session. - **Logical decoding from day one.** `wal_level = logical`, `max_wal_senders = 10`, `max_replication_slots = 10`. No primary restart to enable CDC later. - **Standard extensions, pinned.** pgvector, pgvectorscale, PostGIS, pg_cron, pg_partman, pg_jsonschema, hypopg, pg_repack, pg_search, pg_stat_statements, pg_trgm, auto_explain. - **Beyond extensions on the same volume.** `beyond-auth` and `beyond-queue` ship in the image and live under their own schemas in your database. Forking your DB forks their state automatically. diff --git a/mise.toml b/mise.toml index f1b8016..9e1d017 100644 --- a/mise.toml +++ b/mise.toml @@ -27,6 +27,10 @@ run = "cargo test --test boot -- --include-ignored --nocapture" description = "Run supervisor lifecycle tests (requires Docker + musl target + beyond-pg-test image)" run = "cargo test --test supervisor -- --include-ignored --nocapture" +[tasks."test:peering"] +description = "Run PgBouncer so_reuseport peering (cross-worker cancel) e2e test (requires Docker + beyond-pg-test image)" +run = "cargo test --test pgbouncer_peering -- --include-ignored --nocapture --test-threads=1" + [tasks."test:handoff"] description = "Run beyond-pg-init → supervisor handoff e2e tests (requires Docker + musl target + beyond-pg-test image)" run = "cargo test --release --test handoff_e2e -- --include-ignored --nocapture --test-threads=1" diff --git a/packer/files/pgbouncer/pgbouncer.ini b/packer/files/pgbouncer/pgbouncer.ini index c06fd8b..145f30a 100644 --- a/packer/files/pgbouncer/pgbouncer.ini +++ b/packer/files/pgbouncer/pgbouncer.ini @@ -33,8 +33,9 @@ ignore_startup_parameters = extra_float_digits,search_path logfile = pidfile = -# No per-process unix socket or pidfile: beyond-pg runs 1..N so_reuseport workers that -# share the TCP :5432 port, and pgbouncer defaults unix_socket_dir to /tmp — leaving it -# set would make every worker collide on /tmp/.s.PGSQL.5432. Clients use TCP :5432; the -# direct Postgres path is :5433. -unix_socket_dir = +# so_reuseport, unix_socket_dir, peer_id and the peer-routing map are appended per +# worker at boot by beyond-pg (config.rs::pgbouncer_ini). beyond-pg runs 1..N +# so_reuseport workers sharing TCP :5432; each gets a DISTINCT unix_socket_dir so +# their peer sockets don't collide, which is what lets a query-cancel that the +# kernel routes to the "wrong" worker be forwarded to the one holding the session. +# Clients use TCP :5432; the direct Postgres path is :5433. diff --git a/src/boot.rs b/src/boot.rs index 9d3cd0d..f16c9f2 100644 --- a/src/boot.rs +++ b/src/boot.rs @@ -98,6 +98,26 @@ fn ensure_socket_dir() { } } +/// Create a per-worker PgBouncer peer socket directory owned by `postgres`. +/// Idempotent: create_dir_all is a no-op when it exists, and re-chowning is safe. +/// The socket inside (`.s.PGSQL.5432`) is created by pgbouncer after it drops to +/// the postgres user, so the directory must be writable by that user. +fn ensure_peer_socket_dir(peer_id: usize) { + let dir = config::pgb_peer_socket_dir(peer_id); + if let Err(e) = std::fs::create_dir_all(&dir) { + warn!("ensure_peer_socket_dir: create {dir}: {e}"); + return; + } + match std::process::Command::new("chown") + .args(["postgres:postgres", &dir]) + .status() + { + Ok(s) if s.success() => info!("pgbouncer peer socket dir {dir} ready (owner postgres)"), + Ok(s) => warn!("chown {dir} exited {s}; continuing"), + Err(e) => warn!("chown {dir} failed to spawn: {e}; continuing"), + } +} + async fn do_boot_primary(cfg: &MmdsConfig) -> Result<(), BootError> { info!("boot: step 1/8 maybe_initdb"); maybe_initdb(cfg).await?; @@ -416,8 +436,7 @@ mod tests { let cfg = test_cfg(Some("http://10.0.0.5:9000"), Some("2026-05-14 03:00:00")); write_pitr_config_into(&cfg, pgdata.path()).unwrap(); - let content = - std::fs::read_to_string(pgdata.path().join("conf.d/05-pitr.conf")).unwrap(); + let content = std::fs::read_to_string(pgdata.path().join("conf.d/05-pitr.conf")).unwrap(); assert!( content.contains("restore_command = 'curl -f -s http://10.0.0.5:9000/%f -o %p'"), "restore_command wrong: {content}" @@ -798,15 +817,26 @@ fn write_config_files(cfg: &MmdsConfig, tls: &crate::tls::TlsConfig) -> Result<( }; write_atomic(Path::new(PG_HBA_PATH), &hba)?; - // pgbouncer.ini — overwritten every boot, with client_tls_* keys pointing - // at the same cert Postgres uses. + // pgbouncer.ini — one config per possible so_reuseport worker, overwritten every + // boot, with client_tls_* keys pointing at the same cert Postgres uses. Workers + // are spawned lazily by the scaler, but all their configs (and peer socket dirs) + // are laid down up front so the shared [peers] map is stable and a scaled-up + // worker never waits on config. peer_id 1 keeps the canonical pgbouncer.ini. if let Some(parent) = Path::new(PGBOUNCER_INI_PATH).parent() { std::fs::create_dir_all(parent)?; } - write_atomic( - Path::new(PGBOUNCER_INI_PATH), - &config::pgbouncer_ini(tls, cfg.ram_bytes, cfg.vcpus), - )?; + let max_workers = config::pgbouncer_max_workers(cfg.vcpus); + for peer_id in 1..=max_workers { + write_atomic( + Path::new(&config::pgbouncer_ini_path(peer_id)), + &config::pgbouncer_ini(tls, cfg.ram_bytes, cfg.vcpus, peer_id), + )?; + // Peer socket dir must exist and be writable by the postgres user pgbouncer + // drops to (only created when peering is active, i.e. max_workers > 1). + if max_workers > 1 { + ensure_peer_socket_dir(peer_id); + } + } Ok(()) } diff --git a/src/config.rs b/src/config.rs index 6ef08e9..e553064 100644 --- a/src/config.rs +++ b/src/config.rs @@ -147,8 +147,60 @@ pub fn pgbouncer_max_workers(vcpus: u32) -> usize { ((vcpus as usize) / 4).clamp(1, 8) } -/// pgbouncer.ini: static base + boot-computed pool sizing + TLS termination keys. -pub fn pgbouncer_ini(tls: &crate::tls::TlsConfig, ram_bytes: u64, vcpus: u32) -> String { +/// Parent directory for per-worker PgBouncer peer sockets. Each so_reuseport worker +/// gets its own `pgb-` subdirectory here so their `.s.PGSQL.5432` sockets +/// don't collide; the sockets exist only so peers can forward cancels to each other. +/// On tmpfs (`/run`), so it survives a supervisor handoff but is recreated on cold boot. +pub const PGB_PEER_DIR: &str = "/run/beyond-pg"; + +/// PgBouncer `peer_id` (1-based; `0` disables peering) for a worker's stable name. +/// Worker 0 keeps the legacy name "pgbouncer" (→ 1); extras are "pgbouncer-N" (→ N+1). +/// Peer ids are stable across restarts so the shared `[peers]` map is fixed and the +/// scaler never has to rewrite config as it grows/shrinks the live worker set. +pub fn pgb_peer_id(name: &str) -> usize { + match name.strip_prefix("pgbouncer-") { + Some(n) => n.parse::().map(|n| n + 1).unwrap_or(1), + None => 1, // "pgbouncer" + } +} + +/// Per-worker unix socket directory for the given `peer_id` (`PGB_PEER_DIR/pgb-`). +pub fn pgb_peer_socket_dir(peer_id: usize) -> String { + format!("{PGB_PEER_DIR}/pgb-{peer_id}") +} + +/// Config file path for the worker with the given `peer_id`. Worker 1 keeps the +/// canonical `/etc/pgbouncer/pgbouncer.ini`; extras get `pgbouncer-.ini`. +pub fn pgbouncer_ini_path(peer_id: usize) -> String { + if peer_id <= 1 { + PGBOUNCER_INI_PATH.to_string() + } else { + format!("/etc/pgbouncer/pgbouncer-{peer_id}.ini") + } +} + +/// pgbouncer.ini for the worker with the given `peer_id` (1-based): static base + +/// boot-computed pool sizing + TLS termination + so_reuseport peering. +/// +/// When the box can run more than one worker (`pgbouncer_max_workers > 1`), each +/// worker gets a DISTINCT `unix_socket_dir` + unique `peer_id` and an identical +/// `[peers]` map covering every possible worker (1..=max). That map is what lets a +/// query-cancel request the kernel routes (via so_reuseport) to a worker that isn't +/// holding the session be forwarded over the peer socket to the one that is — +/// without it, misdirected cancels are silently dropped. The map is pre-declared +/// for the full worker cap so the scaler never rewrites config; forwarding to a +/// peer that isn't currently running simply fails, exactly as a dropped cancel +/// would have before, so pre-declaring idle peers is safe. +/// +/// Single-worker boxes (`max == 1`) can never misroute a cancel, so peering is +/// omitted and `unix_socket_dir` stays empty (no client-facing unix socket; clients +/// use TCP :5432, the direct Postgres path is :5433). +pub fn pgbouncer_ini( + tls: &crate::tls::TlsConfig, + ram_bytes: u64, + vcpus: u32, + peer_id: usize, +) -> String { let mut out = String::from(PGBOUNCER_INI_BASE); if !out.ends_with('\n') { out.push('\n'); @@ -190,6 +242,28 @@ pub fn pgbouncer_ini(tls: &crate::tls::TlsConfig, ram_bytes: u64, vcpus: u32) -> // client-cert behavior (sslmode=allow neither requests nor verifies them). out.push_str(&format!("client_tls_ca_file = {}\n", tls.cert.display())); } + + // so_reuseport peering — cross-worker query cancellation. + let max_workers = pgbouncer_max_workers(vcpus); + if max_workers > 1 { + out.push_str( + "\n# Peering — forward query-cancel requests the kernel routes to the wrong\n\ + # so_reuseport worker to the one holding the session (generated by beyond-pg).\n", + ); + out.push_str(&format!( + "unix_socket_dir = {}\n", + pgb_peer_socket_dir(peer_id) + )); + out.push_str(&format!("peer_id = {peer_id}\n")); + out.push_str("\n[peers]\n"); + for id in 1..=max_workers { + out.push_str(&format!("{id} = host={}\n", pgb_peer_socket_dir(id))); + } + } else { + // Single-worker box: no peer can ever hold a foreign session, so keep the + // client-facing unix socket disabled (clients use TCP :5432). + out.push_str("\nunix_socket_dir =\n"); + } out } @@ -541,8 +615,14 @@ mod tests { "missing libs should be dropped: {out}" ); assert!(!out.contains("pg_cron"), "pg_cron not installed: {out}"); - assert!(!out.contains("beyond_auth"), "beyond_auth not installed: {out}"); - assert!(!out.contains("beyond_queue"), "beyond_queue not installed: {out}"); + assert!( + !out.contains("beyond_auth"), + "beyond_auth not installed: {out}" + ); + assert!( + !out.contains("beyond_queue"), + "beyond_queue not installed: {out}" + ); // Other lines untouched. assert!(out.contains("foo = 1")); } @@ -556,7 +636,9 @@ mod tests { let pkglibdir = dir.path().to_str().unwrap(); let conf = "shared_preload_libraries = 'pg_stat_statements,auto_explain,pg_cron'\n"; let out = filter_shared_preload_libraries(conf, pkglibdir); - assert!(out.contains("shared_preload_libraries = 'pg_stat_statements,auto_explain,pg_cron'")); + assert!( + out.contains("shared_preload_libraries = 'pg_stat_statements,auto_explain,pg_cron'") + ); } #[test] @@ -631,7 +713,7 @@ mod tests { key: PathBuf::from("/run/beyond/tls/key.pem"), ca: Some(PathBuf::from("/run/beyond/tls/ca.pem")), }; - let out = pgbouncer_ini(&tls, 8 * 1024 * 1024 * 1024, 4); + let out = pgbouncer_ini(&tls, 8 * 1024 * 1024 * 1024, 4, 1); assert!( out.starts_with(PGBOUNCER_INI_BASE), "must preserve base config" @@ -654,7 +736,7 @@ mod tests { key: PathBuf::from("/var/lib/postgresql/18/main/beyond/server.key"), ca: None, }; - let out = pgbouncer_ini(&tls, 8 * 1024 * 1024 * 1024, 4); + let out = pgbouncer_ini(&tls, 8 * 1024 * 1024 * 1024, 4, 1); assert!(out.contains("client_tls_sslmode = allow")); assert!(out.contains("client_tls_cert_file")); assert!( @@ -679,23 +761,32 @@ mod tests { ca: None, }; // so_reuseport is ALWAYS on now (so the supervisor can add/reap workers live). - let small = pgbouncer_ini(&tls, 4 * 1024 * 1024 * 1024, 4); + // 4 vCPU -> max_workers 1: single worker can't misroute a cancel, so peering + // is off and the client-facing unix socket stays disabled (TCP :5432 only). + let small = pgbouncer_ini(&tls, 4 * 1024 * 1024 * 1024, 4, 1); assert!(small.contains("so_reuseport = 1")); assert!(small.contains("default_pool_size = 16")); // pool 16 (4 vCPU * 3, floored) - // Empty unix_socket_dir so multiple so_reuseport workers share only TCP :5432. assert!( - small.contains("unix_socket_dir ="), - "needs empty unix_socket_dir: {small}" + small.contains("unix_socket_dir =\n"), + "single-worker box keeps empty unix_socket_dir: {small}" + ); + assert!( + !small.contains("[peers]"), + "no peering on 1-worker box: {small}" + ); + assert!( + !small.contains("peer_id ="), + "no peer_id on 1-worker box: {small}" ); // Mid box: full pool 48 (16 vCPU * 3), not divided across workers. - let mid = pgbouncer_ini(&tls, 32u64 * 1024 * 1024 * 1024, 16); + let mid = pgbouncer_ini(&tls, 32u64 * 1024 * 1024 * 1024, 16, 2); assert!( mid.contains("default_pool_size = 48"), "16 vCPU * 3 = full 48: {mid}" ); // Big box: FULL pool per worker (a single live worker gets the whole 192, not // 192/8=24 — the scaler adds workers for handshake CPU, not query concurrency). - let big = pgbouncer_ini(&tls, 256u64 * 1024 * 1024 * 1024, 64); + let big = pgbouncer_ini(&tls, 256u64 * 1024 * 1024 * 1024, 64, 1); assert!(big.contains("so_reuseport = 1")); assert!( big.contains("default_pool_size = 192"), @@ -706,6 +797,52 @@ mod tests { assert!(tuning.contains("max_connections = 330")); // 192 + 64*2 + 10 } + #[test] + fn pgb_peer_id_maps_stable_names() { + assert_eq!(pgb_peer_id("pgbouncer"), 1); + assert_eq!(pgb_peer_id("pgbouncer-1"), 2); + assert_eq!(pgb_peer_id("pgbouncer-7"), 8); + } + + #[test] + fn pgbouncer_ini_path_keeps_canonical_for_worker_one() { + assert_eq!(pgbouncer_ini_path(1), PGBOUNCER_INI_PATH); + assert_eq!(pgbouncer_ini_path(2), "/etc/pgbouncer/pgbouncer-2.ini"); + assert_eq!(pgbouncer_ini_path(8), "/etc/pgbouncer/pgbouncer-8.ini"); + } + + #[test] + fn pgbouncer_ini_emits_peering_on_multiworker_box() { + let tls = TlsConfig { + source: TlsSource::SelfSigned, + cert: PathBuf::from("/c"), + key: PathBuf::from("/k"), + ca: None, + }; + // 16 vCPU -> max_workers 4. Worker 2's config must carry its own identity + // plus the FULL [peers] map (all 4 workers) so cancels route to any peer. + let w2 = pgbouncer_ini(&tls, 32u64 * 1024 * 1024 * 1024, 16, 2); + assert!(w2.contains("peer_id = 2"), "worker 2 identity: {w2}"); + assert!( + w2.contains("unix_socket_dir = /run/beyond-pg/pgb-2"), + "worker 2 socket dir: {w2}" + ); + assert!(w2.contains("[peers]"), "peers section: {w2}"); + for id in 1..=4 { + assert!( + w2.contains(&format!("{id} = host=/run/beyond-pg/pgb-{id}")), + "peer {id} declared: {w2}" + ); + } + // The map covers exactly the worker cap — no peer 5 on a 4-worker box. + assert!(!w2.contains("5 = host="), "no over-declared peer: {w2}"); + // Every worker on the box shares an identical [peers] map (only identity differs). + let w1 = pgbouncer_ini(&tls, 32u64 * 1024 * 1024 * 1024, 16, 1); + let peers_of = |s: &str| s.split("\n[peers]\n").nth(1).unwrap().to_string(); + assert_eq!(peers_of(&w1), peers_of(&w2), "shared [peers] map"); + assert!(w1.contains("peer_id = 1")); + } + #[test] fn replica_conf_no_wal_sink() { let conf = replica_conf("host=10.0.0.1 port=5433 user=replicator", None); diff --git a/src/supervisor.rs b/src/supervisor.rs index fa655e3..ad6e998 100644 --- a/src/supervisor.rs +++ b/src/supervisor.rs @@ -1379,8 +1379,12 @@ fn spawn_pgbouncer( state: &mut ChildState, log_tx: &mpsc::Sender, ) -> Result<(), Box> { + // Each so_reuseport worker runs its own config (distinct peer_id + unix_socket_dir); + // the stable name maps to a fixed peer_id, so respawns/handoffs keep the same slot. + let peer_id = crate::config::pgb_peer_id(state.name); + let ini = crate::config::pgbouncer_ini_path(peer_id); let mut cmd = Command::new("pgbouncer"); - cmd.arg("/etc/pgbouncer/pgbouncer.ini") + cmd.arg(&ini) .stdout(std::process::Stdio::piped()) .stderr(std::process::Stdio::piped()) .stdin(std::process::Stdio::null()); @@ -1401,7 +1405,11 @@ fn spawn_pgbouncer( ); spawn_async_reader_task(ExecStream::Stderr, stderr, log_tx.clone(), execution_id); - info!("pgbouncer spawned (pid={})", child.id().unwrap_or(0)); + info!( + "pgbouncer '{}' spawned (pid={}, peer_id={peer_id}, ini={ini})", + state.name, + child.id().unwrap_or(0) + ); state.record_start(child); Ok(()) } diff --git a/tests/pgbouncer_peering.rs b/tests/pgbouncer_peering.rs new file mode 100644 index 0000000..23179f9 --- /dev/null +++ b/tests/pgbouncer_peering.rs @@ -0,0 +1,368 @@ +//! PgBouncer so_reuseport peering end-to-end test: cross-worker query cancellation. +//! +//! When beyond-pg runs more than one so_reuseport pooler worker (see +//! `src/config.rs::pgbouncer_ini` / `PgbScaler`), a Postgres CancelRequest arrives +//! on a *fresh* TCP connection that the kernel may route to a different worker than +//! the one holding the session. Without peering that cancel is silently dropped; +//! with peering the receiving worker forwards it to the peer that owns the session. +//! +//! This test proves the mechanism against a real PgBouncer using the exact config +//! keys `pgbouncer_ini` emits (`peer_id`, per-worker `unix_socket_dir`, and the +//! shared `[peers]` map — asserted in the `config` unit tests). To make the routing +//! deterministic in CI, the two workers listen on *distinct* ports and the cancel is +//! sent to the "wrong" one on purpose; so_reuseport port-sharing is orthogonal to the +//! forwarding logic under test and is already exercised in production. +//! +//! - [`peering_forwards_misrouted_cancel`] — peering ON → misrouted cancel cancels the query +//! - [`no_peering_drops_misrouted_cancel`] — peering OFF → misrouted cancel is dropped (control) +//! +//! # Prerequisites +//! Test image: `mise run build:test-image` +//! (`beyond-pg-test:latest` = postgres:18 + pgbouncer 1.25 + …) +//! +//! Run: `cargo test --test pgbouncer_peering -- --ignored --nocapture` + +use std::io::{Read, Write}; +use std::net::TcpStream; +use std::sync::Mutex; +use std::time::{Duration, Instant}; + +// Serialize Docker launches (matches the other e2e suites on macOS Docker Desktop). +static DOCKER: Mutex<()> = Mutex::new(()); + +const IMAGE: &str = "beyond-pg-test:latest"; +const CANCEL_CODE: u32 = 80877102; +const PROTO_3: u32 = 196608; + +// --------------------------------------------------------------------------- +// Minimal Docker RAII helper +// --------------------------------------------------------------------------- + +struct Container(String); + +impl Container { + /// Start postgres (trust auth) with the two pooler ports published to the host. + /// Returns None if the test image is absent so the test skips instead of failing. + fn start() -> Option { + let present = std::process::Command::new("docker") + .args(["image", "inspect", IMAGE]) + .output() + .ok()? + .status + .success(); + if !present { + eprintln!("skipping: {IMAGE} not found — run `mise run build:test-image`"); + return None; + } + let out = std::process::Command::new("docker") + .args([ + "run", + "-d", + "--rm", + "-e", + "POSTGRES_HOST_AUTH_METHOD=trust", + "-e", + "POSTGRES_PASSWORD=x", + "-p", + "127.0.0.1::6432", + "-p", + "127.0.0.1::6433", + IMAGE, + ]) + .output() + .expect("docker run"); + assert!( + out.status.success(), + "docker run failed: {}", + String::from_utf8_lossy(&out.stderr) + ); + Some(Self( + String::from_utf8_lossy(&out.stdout).trim().to_string(), + )) + } + + fn exec(&self, args: &[&str]) -> std::process::Output { + let mut a = vec!["exec", &self.0]; + a.extend_from_slice(args); + std::process::Command::new("docker") + .args(&a) + .output() + .expect("docker exec") + } + + fn exec_u(&self, user: &str, args: &[&str]) -> std::process::Output { + let mut a = vec!["exec", "-u", user, &self.0]; + a.extend_from_slice(args); + std::process::Command::new("docker") + .args(&a) + .output() + .expect("docker exec -u") + } + + /// Host port mapped to the container's `container_port`. + fn host_port(&self, container_port: u16) -> u16 { + let out = std::process::Command::new("docker") + .args(["port", &self.0, &format!("{container_port}/tcp")]) + .output() + .expect("docker port"); + let s = String::from_utf8_lossy(&out.stdout); + let line = s.lines().next().unwrap_or_default(); + line.rsplit(':') + .next() + .and_then(|p| p.trim().parse().ok()) + .unwrap_or_else(|| panic!("no host port for {container_port}: {s:?}")) + } + + fn wait_postgres_ready(&self, timeout: Duration) -> bool { + let deadline = Instant::now() + timeout; + while Instant::now() < deadline { + if self.exec(&["pg_isready", "-q"]).status.success() { + return true; + } + std::thread::sleep(Duration::from_secs(1)); + } + false + } + + fn logs(&self) -> String { + let out = std::process::Command::new("docker") + .args(["logs", &self.0]) + .output() + .expect("docker logs"); + format!( + "stdout:\n{}\nstderr:\n{}", + String::from_utf8_lossy(&out.stdout), + String::from_utf8_lossy(&out.stderr) + ) + } +} + +impl Drop for Container { + fn drop(&mut self) { + let _ = std::process::Command::new("docker") + .args(["rm", "-f", &self.0]) + .status(); + } +} + +/// Write the two worker configs and (re)start both pgbouncers. `peering` toggles +/// the `peer_id` + `[peers]` keys; both cases keep distinct `unix_socket_dir`. +fn configure_and_start(c: &Container, peering: bool) { + // trust auth still requires the user to exist in auth_file (pgbouncer 1.25). + let mk = c.exec_u( + "postgres", + &["bash", "-c", "echo '\"postgres\" \"\"' > /tmp/userlist.txt"], + ); + assert!( + mk.status.success(), + "write userlist: {}", + String::from_utf8_lossy(&mk.stderr) + ); + + for id in 1..=2u16 { + let port = 6431 + id; + let peers = if peering { + "\n[peers]\n1 = host=/tmp/pgb1\n2 = host=/tmp/pgb2" + } else { + "" + }; + let ident = if peering { + format!("\nunix_socket_dir = /tmp/pgb{id}\npeer_id = {id}") + } else { + format!("\nunix_socket_dir = /tmp/pgb{id}") + }; + let ini = format!( + "[databases]\n\ + * = host=/var/run/postgresql port=5432\n\ + [pgbouncer]\n\ + listen_addr = 0.0.0.0\n\ + listen_port = {port}\n\ + auth_type = trust\n\ + auth_file = /tmp/userlist.txt\n\ + pool_mode = transaction\n\ + logfile = /tmp/pgb{id}.log\n\ + pidfile = /tmp/pgb{id}.pid{ident}{peers}\n" + ); + let script = format!("mkdir -p /tmp/pgb{id} && cat > /tmp/pgb{id}.ini <<'INI'\n{ini}INI"); + let w = c.exec_u("postgres", &["bash", "-c", &script]); + assert!( + w.status.success(), + "write ini {id}: {}", + String::from_utf8_lossy(&w.stderr) + ); + } + + // Stop any prior workers, then start fresh. + let _ = c.exec(&[ + "bash", + "-c", + "kill $(cat /tmp/pgb1.pid /tmp/pgb2.pid 2>/dev/null) 2>/dev/null; sleep 1; true", + ]); + for id in 1..=2u16 { + let s = c.exec_u( + "postgres", + &["bash", "-c", &format!("pgbouncer -d /tmp/pgb{id}.ini")], + ); + assert!( + s.status.success(), + "start pgbouncer {id}: {}\n{}", + String::from_utf8_lossy(&s.stderr), + c.logs() + ); + } + std::thread::sleep(Duration::from_secs(2)); +} + +// --------------------------------------------------------------------------- +// Raw pgwire probe +// --------------------------------------------------------------------------- + +fn read_exact(s: &mut TcpStream, n: usize) -> std::io::Result> { + let mut buf = vec![0u8; n]; + s.read_exact(&mut buf)?; + Ok(buf) +} + +/// Read one backend message: (tag, body-without-length-prefix). +fn read_msg(s: &mut TcpStream) -> std::io::Result<(u8, Vec)> { + let tag = read_exact(s, 1)?[0]; + let len = u32::from_be_bytes(read_exact(s, 4)?.try_into().unwrap()) as usize; + let body = read_exact(s, len - 4)?; + Ok((tag, body)) +} + +/// Send StartupMessage (user/database=postgres) and return this session's +/// (backend pid, secret key) from BackendKeyData. Trust auth → no password. +fn startup(s: &mut TcpStream) -> (u32, u32) { + let params = b"user\x00postgres\x00database\x00postgres\x00\x00"; + let mut msg = Vec::new(); + msg.extend_from_slice(&((8 + params.len()) as u32).to_be_bytes()); + msg.extend_from_slice(&PROTO_3.to_be_bytes()); + msg.extend_from_slice(params); + s.write_all(&msg).unwrap(); + + let (mut pid, mut secret) = (0u32, 0u32); + loop { + let (tag, body) = read_msg(s).expect("startup read"); + match tag { + b'R' => { + let kind = u32::from_be_bytes(body[..4].try_into().unwrap()); + assert_eq!( + kind, 0, + "expected trust auth (AuthenticationOk), got kind {kind}" + ); + } + b'K' => { + pid = u32::from_be_bytes(body[..4].try_into().unwrap()); + secret = u32::from_be_bytes(body[4..8].try_into().unwrap()); + } + b'Z' => return (pid, secret), + b'E' => panic!("startup error: {}", String::from_utf8_lossy(&body)), + _ => {} + } + } +} + +fn send_query(s: &mut TcpStream, sql: &str) { + let mut body = sql.as_bytes().to_vec(); + body.push(0); + let mut msg = vec![b'Q']; + msg.extend_from_slice(&((4 + body.len()) as u32).to_be_bytes()); + msg.extend_from_slice(&body); + s.write_all(&msg).unwrap(); +} + +/// Send a CancelRequest carrying (pid, secret) to `port` — a fresh connection, +/// exactly as libpq's PQcancel does. +fn send_cancel(port: u16, pid: u32, secret: u32) { + let mut s = TcpStream::connect(("127.0.0.1", port)).expect("connect cancel port"); + let mut msg = Vec::new(); + msg.extend_from_slice(&16u32.to_be_bytes()); + msg.extend_from_slice(&CANCEL_CODE.to_be_bytes()); + msg.extend_from_slice(&pid.to_be_bytes()); + msg.extend_from_slice(&secret.to_be_bytes()); + s.write_all(&msg).unwrap(); + let _ = s.read(&mut [0u8; 1]); // server closes after handling the cancel +} + +/// Start `SELECT pg_sleep(30)` on the worker at `query_port`, then fire a +/// CancelRequest for that session at `cancel_port`. Returns true iff the query +/// was cancelled within `wait` (i.e. the cancel was forwarded to the right worker). +fn misrouted_cancel_hits(query_port: u16, cancel_port: u16, wait: Duration) -> bool { + let mut conn = TcpStream::connect(("127.0.0.1", query_port)).expect("connect query port"); + let (pid, secret) = startup(&mut conn); + eprintln!(" worker gave BackendKeyData pid={pid} secret={secret:#010x}"); + send_query(&mut conn, "SELECT pg_sleep(30)"); + + // Give pgbouncer a moment to hand the query to a server before cancelling. + std::thread::sleep(Duration::from_millis(300)); + eprintln!(" sending CancelRequest to the OTHER worker (port {cancel_port})"); + send_cancel(cancel_port, pid, secret); + + conn.set_read_timeout(Some(wait)).unwrap(); + loop { + match read_msg(&mut conn) { + Ok((b'E', body)) => { + let text = String::from_utf8_lossy(&body); + let cancelled = text.contains("57014") || text.contains("canceling statement"); + eprintln!(" ErrorResponse: {}", text.replace('\0', " ")); + return cancelled; + } + Ok((b'C', _)) => { + eprintln!(" CommandComplete — pg_sleep ran to completion, cancel was dropped"); + return false; + } + Ok(_) => continue, + Err(e) => { + eprintln!(" no reply within {wait:?} ({e}) — cancel was dropped"); + return false; + } + } + } +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +#[test] +#[ignore = "requires Docker + beyond-pg-test:latest image"] +fn peering_forwards_misrouted_cancel() { + let _g = DOCKER.lock().unwrap_or_else(|e| e.into_inner()); + let Some(c) = Container::start() else { return }; + assert!( + c.wait_postgres_ready(Duration::from_secs(90)), + "postgres not ready\n{}", + c.logs() + ); + configure_and_start(&c, /* peering = */ true); + + let qport = c.host_port(6432); + let cport = c.host_port(6433); + assert!( + misrouted_cancel_hits(qport, cport, Duration::from_secs(10)), + "peering ON: a cancel routed to the wrong worker must be forwarded and cancel the query" + ); +} + +#[test] +#[ignore = "requires Docker + beyond-pg-test:latest image"] +fn no_peering_drops_misrouted_cancel() { + // Control: proves the test actually discriminates — without peering the same + // misrouted cancel is dropped, which is the behavior this feature fixes. + let _g = DOCKER.lock().unwrap_or_else(|e| e.into_inner()); + let Some(c) = Container::start() else { return }; + assert!( + c.wait_postgres_ready(Duration::from_secs(90)), + "postgres not ready\n{}", + c.logs() + ); + configure_and_start(&c, /* peering = */ false); + + let qport = c.host_port(6432); + let cport = c.host_port(6433); + assert!( + !misrouted_cancel_hits(qport, cport, Duration::from_secs(6)), + "peering OFF: a cancel routed to the wrong worker has no peer to forward to and must be dropped" + ); +} From e2bdf2ade6a3dcff4425e7e49953e7d9dcfd288a Mon Sep 17 00:00:00 2001 From: Jared Lunde Date: Sat, 11 Jul 2026 10:43:59 -0700 Subject: [PATCH 2/2] style: apply cargo fmt to pre-existing working-tree changes Formatting-only line rewrapping (no logic changes) in files that were dirty in the working tree; bundled in so the branch is fmt-clean. Co-Authored-By: Claude Opus 4.8 (1M context) --- beyond-pg-core/src/mmds.rs | 8 +++++--- beyond-pg-init/src/volumes.rs | 28 ++++++++++++++++++++++------ src/tls.rs | 4 ++-- src/vsock.rs | 3 +-- 4 files changed, 30 insertions(+), 13 deletions(-) diff --git a/beyond-pg-core/src/mmds.rs b/beyond-pg-core/src/mmds.rs index a406d42..0e5d706 100644 --- a/beyond-pg-core/src/mmds.rs +++ b/beyond-pg-core/src/mmds.rs @@ -57,7 +57,9 @@ pub enum MmdsError { InvalidPassword, #[error("BEYOND_PG_PRIMARY_CONNINFO is required when BEYOND_PG_TIER=replica")] MissingPrimaryConninfo, - #[error("BEYOND_PG_RECOVERY_TARGET_TIME requires BEYOND_PG_WAL_SINK (the WAL source to replay from)")] + #[error( + "BEYOND_PG_RECOVERY_TARGET_TIME requires BEYOND_PG_WAL_SINK (the WAL source to replay from)" + )] RecoveryTargetWithoutWalSource, #[error("MMDS metadata not available: {0}")] Unavailable(String), @@ -252,8 +254,8 @@ mod tests { .expect("should parse"); assert_eq!(cfg.replication_password.as_deref(), Some("replpw")); // Empty string is treated as absent (same as the other optional keys). - let cfg2 = parse(base_meta(&[("BEYOND_PG_REPLICATION_PASSWORD", "")])) - .expect("should parse"); + let cfg2 = + parse(base_meta(&[("BEYOND_PG_REPLICATION_PASSWORD", "")])).expect("should parse"); assert!(cfg2.replication_password.is_none()); } diff --git a/beyond-pg-init/src/volumes.rs b/beyond-pg-init/src/volumes.rs index 659be10..244220c 100644 --- a/beyond-pg-init/src/volumes.rs +++ b/beyond-pg-init/src/volumes.rs @@ -41,7 +41,10 @@ pub fn mount_from_mmds() { if attachments.is_empty() { return; } - eprintln!("[init] mounting {} data volume(s) from MMDS", attachments.len()); + eprintln!( + "[init] mounting {} data volume(s) from MMDS", + attachments.len() + ); mount_data_volumes(&attachments); } @@ -77,7 +80,10 @@ fn parse_attachments(val: &Value) -> Vec { path: v["path"].as_str()?.to_string(), fstype: v["fstype"].as_str().unwrap_or("ext4").to_string(), readonly: v["readonly"].as_bool().unwrap_or(false), - storage_class: v["storage_class"].as_str().unwrap_or("standard").to_string(), + storage_class: v["storage_class"] + .as_str() + .unwrap_or("standard") + .to_string(), }) }) .collect() @@ -100,7 +106,10 @@ fn mount_data_volumes(attachments: &[AttachmentMeta]) { fn mount_one(a: &AttachmentMeta) -> Result<(), String> { if !wait_for_device(&a.device) { - return Err(format!("device {} not present after {DEVICE_WAIT:?}", a.device)); + return Err(format!( + "device {} not present after {DEVICE_WAIT:?}", + a.device + )); } ensure_dir(&a.path)?; @@ -182,7 +191,10 @@ fn is_mounted(path: &str, device: &str) -> Result { /// Return the device currently mounted at `path`, or `None` if nothing is. fn mounted_device(path: &str) -> Result, String> { let info = read_mountinfo()?; - Ok(info.into_iter().find(|(mp, _)| mp == path).map(|(_, dev)| dev)) + Ok(info + .into_iter() + .find(|(mp, _)| mp == path) + .map(|(_, dev)| dev)) } fn read_mountinfo() -> Result, String> { @@ -312,9 +324,13 @@ mod tests { #[test] fn parse_mountinfo_extracts_mountpoint_and_source() { - let line = "36 35 8:1 / /var/lib/postgresql rw,noatime shared:1 - ext4 /dev/vdb rw,data=ordered"; + let line = + "36 35 8:1 / /var/lib/postgresql rw,noatime shared:1 - ext4 /dev/vdb rw,data=ordered"; let entries = parse_mountinfo(line); - assert_eq!(entries, vec![("/var/lib/postgresql".to_string(), "/dev/vdb".to_string())]); + assert_eq!( + entries, + vec![("/var/lib/postgresql".to_string(), "/dev/vdb".to_string())] + ); } #[test] diff --git a/src/tls.rs b/src/tls.rs index 25cd818..a3fed5e 100644 --- a/src/tls.rs +++ b/src/tls.rs @@ -258,8 +258,8 @@ fn generate_cert(cert_path: &Path, key_path: &Path) -> Result<(), TlsError> { // performance; v2 only appends the public key, which the cert already // carries). Falls back to rcgen's PEM if the key isn't the expected Ed25519 // shape (e.g. a future algorithm change). - let key_pem = ed25519_pkcs8_v1_pem(&key_pair.serialize_der()) - .unwrap_or_else(|| key_pair.serialize_pem()); + let key_pem = + ed25519_pkcs8_v1_pem(&key_pair.serialize_der()).unwrap_or_else(|| key_pair.serialize_pem()); crate::config::write_atomic_bytes_with_mode(key_path, key_pem.as_bytes(), 0o600)?; Ok(()) diff --git a/src/vsock.rs b/src/vsock.rs index f8e8fe4..bb53d18 100644 --- a/src/vsock.rs +++ b/src/vsock.rs @@ -123,8 +123,7 @@ mod tests { #[serde(with = "serde_bytes")] bytes: Vec, } - let outer: OuterIn = - rmp_serde::from_slice(&frame[5..]).expect("decode AppMessagePayload"); + let outer: OuterIn = rmp_serde::from_slice(&frame[5..]).expect("decode AppMessagePayload"); // 3. Inner workload envelope: [AppKind=46][msgpack_named(payload)]. let (kind, rest) = outer.bytes.split_first().expect("non-empty inner bytes");