postjevsql.git / tests / support / postgres.rs
postgres.rsannotatedpostgres.rssource304 lines · 11.9 KB · raw

A throwaway Postgres per test, run from the devshell's binaries of the major buck built for (POSTJEVSQL_POSTGRES_BIN, POSTJEVSQL_PG_MAJOR): initdb into a temp dir, postgres as a child process, the built extension served in place. Same shape as a testcontainers module (builder, start, teardown on drop), without docker.

The library is found through dynamic_library_path on every major. The control and install script are found through extension_control_path from PG18. PG17 has no such setting and reads them only from its own sharedir, which it locates relative to the postgres it was started as: nixpkgs' relative-to-symlinks patch keeps it from resolving the link. So on PG17 the instance runs from a prefix of symlinks in its temp dir, whose share/postgresql/extension also links the built tree's files (measured on 17, 2026-09-26).

Teardown cannot leak: the postmaster gets PR_SET_PDEATHSIG, so the kernel sends it SIGQUIT (immediate shutdown) when the test thread that spawned it exits, even if the test process is SIGKILLed.

At most one instance per CPU runs at once, machine-wide: each holds an flock on one of available_parallelism() slot files in /tmp/postjevsql-test-slots-<uid> for its life. buck runs several test binaries at once, and each runs one test per core, each with its own server, so without the cap a full suite on 3 cores reached a load average of 28 (2026-09-26). The path is fixed rather than temp_dir() so every worktree's daemon and every test sandbox share the one set of slots.

A test that needs several instances at once takes their slots together with TestPostgres::start_all, all or none. Taking them one at a time is hold-and-wait: on 3 cores, three sidecar_cli tests each held a target's slot and waited for a sidecar's until buck killed them at 10 minutes (2026-09-27). So a second start on a thread that already holds a slot panics, naming start_all, instead of risking that. At least 2 slots, so a pair always fits.

It listens only on a Unix socket in its own temp dir, never on TCP. A port picked by binding 0 and releasing it can be taken by any other process before the postmaster binds it, and then the test either waits for a server that never starts or talks to another test's.

44use std::fs::File;
45use std::os::fd::AsRawFd;
46use std::os::unix::process::CommandExt;
47use std::collections::HashMap;
48use std::path::{Path, PathBuf};
49use std::process::{Child, Command, Stdio};
50use std::sync::Mutex;
51use std::thread::ThreadId;
52use std::time::{Duration, Instant};

54use tempfile::TempDir;
55
56pub struct TestPostgres {
57    ext: Option<PathBuf>,
58    settings: Vec<(String, String)>,
59    env: Vec<(String, String)>,
60}
61
62pub struct Instance {
63    child: Child,
64    dir: TempDir,
65    // Released after `Drop` has stopped the server.
66    _slot: Slot,
67}
68
69impl Default for TestPostgres {
70    fn default() -> Self {
71        Self::new()
72    }
73}
74
75impl TestPostgres {
76    pub fn new() -> Self {
77        TestPostgres {
78            ext: None,
79            settings: vec![("fsync".into(), "off".into())],
80            env: Vec::new(),
81        }
82    }

Serves the built extension tree without installing it anywhere.

85    pub fn with_extension(mut self, ext_dir: &Path) -> Self {
86        let ext = ext_dir.canonicalize().expect("extension tree exists");
87        let path = ext.to_str().expect("utf-8 path").to_owned();
88        let controls = has_extension_control_path();
89        self.ext = Some(ext);
90        let this = self.setting("dynamic_library_path", &format!("{path}/lib:$libdir"));
91        if controls {
92            this.setting("extension_control_path", &format!("{path}/share:$system"))
93        } else {
94            this
95        }
96    }

98    pub fn setting(mut self, name: &str, value: &str) -> Self {
99        self.settings.push((name.into(), value.into()));
100        self
101    }
102
103    pub fn env_var(mut self, name: &str, value: &str) -> Self {
104        self.env.push((name.into(), value.into()));
105        self
106    }

One instance. A test that needs two at once uses start_all.

109    pub async fn start(self) -> Instance {
110        let [instance] = Self::start_all([self]).await;
111        instance
112    }

Several instances whose slots are taken together, all or none.

115    pub async fn start_all<const N: usize>(configs: [TestPostgres; N]) -> [Instance; N] {
116        let mut slots = acquire_slots(N).await.into_iter();
117        let mut instances = Vec::with_capacity(N);
118        for config in configs {
119            instances.push(config.launch(slots.next().expect("one slot per instance")).await);
120        }
121        let Ok(instances) = instances.try_into() else { unreachable!("one instance per config") };
122        instances
123    }

125    async fn launch(self, slot: Slot) -> Instance {
126        let dir = tempfile::Builder::new().prefix("postjevsql-pg").tempdir().expect("temp dir");
127        let bin = match &self.ext {
128            Some(ext) if !has_extension_control_path() => linked_prefix(dir.path(), ext),
129            _ => PathBuf::from(env("POSTJEVSQL_POSTGRES_BIN")),
130        };
131        let data = dir.path().join("data");
132
133        let status = Command::new(bin.join("initdb"))
134            .args(["--username=postgres", "--auth=trust", "--encoding=UTF8", "--locale=C", "--no-sync"])
135            .arg("-D")
136            .arg(&data)
137            .stdout(Stdio::null())
138            .status()
139            .expect("initdb runs");
140        assert!(status.success(), "initdb failed");
141
142        let mut postgres = Command::new(bin.join("postgres"));
143        postgres
144            .arg("-D")
145            .arg(&data)
146            .args(["-c", "listen_addresses="])
147            .arg("-c")
148            .arg(format!("unix_socket_directories={}", dir.path().display()))
149            // Only what the test sets: a developer's real key or libpq
150            // settings must never reach the instance under test.
151            .env_clear()
152            .env("PATH", std::env::var_os("PATH").unwrap_or_default())
153            .envs(self.env)
154            .stdout(Stdio::null());
155        for (name, value) in &self.settings {
156            postgres.arg("-c").arg(format!("{name}={value}"));
157        }
158        // SAFETY: prctl is async-signal-safe and touches no Rust state.
159        unsafe {
160            postgres.pre_exec(|| {
161                if libc::prctl(libc::PR_SET_PDEATHSIG, libc::SIGQUIT) != 0 {
162                    return Err(std::io::Error::last_os_error());
163                }
164                Ok(())
165            });
166        }
167        let child = postgres.spawn().expect("postgres starts");
168        let instance = Instance { child, dir, _slot: slot };
169        instance.wait_ready().await;
170        instance
171    }
172}
173
174impl Instance {
175    pub fn conn_str(&self) -> String {
176        format!("host={} user=postgres dbname=postgres", self.dir.path().display())
177    }
178
179    async fn wait_ready(&self) {
180        let deadline = Instant::now() + Duration::from_secs(30);
181        loop {
182            match tokio_postgres::connect(&self.conn_str(), tokio_postgres::NoTls).await {
183                Ok(_) => return,
184                Err(e) if Instant::now() > deadline => panic!("postgres never became ready: {e}"),
185                Err(_) => tokio::time::sleep(Duration::from_millis(20)).await,
186            }
187        }
188    }
189}
190
191impl Drop for Instance {
192    fn drop(&mut self) {
193        // SIGQUIT: immediate shutdown, which also stops the backends.
194        // SAFETY: the pid is our own un-reaped child, so it is not reused.
195        unsafe { libc::kill(self.child.id() as libc::pid_t, libc::SIGQUIT) };
196        let _ = self.child.wait();
197    }
198}

A locked slot file, counted against the thread that took it. Closing the file releases the lock, and the kernel closes it if the process dies. The file is close-on-exec, so the server never inherits it.

203struct Slot {
204    _file: File,
205    thread: ThreadId,
206}

Slots held per thread, so a second acquisition on one thread is caught.

209static HELD: Mutex<Option<HashMap<ThreadId, usize>>> = Mutex::new(None);

211fn held(thread: ThreadId, delta: isize) -> usize {
212    let mut map = HELD.lock().unwrap_or_else(|e| e.into_inner());
213    let count = map.get_or_insert_with(HashMap::new).entry(thread).or_default();
214    *count = count.checked_add_signed(delta).expect("slot count");
215    *count
216}
217
218impl Drop for Slot {
219    fn drop(&mut self) {
220        held(self.thread, -1);
221    }
222}

Waits until n slots are free at once and returns them locked. Never holds some while waiting for the rest.

226async fn acquire_slots(n: usize) -> Vec<Slot> {
227    let thread = std::thread::current().id();
228    assert!(
229        held(thread, 0) == 0,
230        "this thread already holds a Postgres slot; start every instance a test needs \
231         together with TestPostgres::start_all, or waiting for a second can deadlock"
232    );
233    // SAFETY: getuid cannot fail and touches no Rust state.
234    let dir = PathBuf::from(format!("/tmp/postjevsql-test-slots-{}", unsafe { libc::getuid() }));
235    std::fs::create_dir_all(&dir).expect("slot dir");
236    let slots = std::thread::available_parallelism().map_or(2, |n| n.get().max(2));
237    assert!(n <= slots, "{n} instances at once, but only {slots} slots");
238    loop {
239        let mut files = Vec::with_capacity(n);
240        for i in 0..slots {
241            let file = File::options()
242                .create(true)
243                .truncate(false)
244                .write(true)
245                .open(dir.join(i.to_string()))
246                .expect("slot file");
247            // SAFETY: the fd is open and owned by `file`.
248            if unsafe { libc::flock(file.as_raw_fd(), libc::LOCK_EX | libc::LOCK_NB) } == 0 {
249                files.push(file);
250                if files.len() == n {
251                    held(thread, n as isize);
252                    return files.into_iter().map(|_file| Slot { _file, thread }).collect();
253                }
254            }
255        }
256        // Too few free: release what was taken before waiting.
257        drop(files);
258        tokio::time::sleep(Duration::from_millis(50)).await;
259    }
260}

extension_control_path arrived in PG18.

263fn has_extension_control_path() -> bool {
264    let major: u32 = env("POSTJEVSQL_PG_MAJOR").parse().expect("POSTJEVSQL_PG_MAJOR is a number");
265    major >= 18
266}

A prefix of symlinks to the server's installation, with the built tree's control and SQL files added to its extension directory. Returns its bin.

271fn linked_prefix(dir: &Path, ext: &Path) -> PathBuf {
272    let real = PathBuf::from(env("POSTJEVSQL_POSTGRES_BIN"));
273    let real = real.parent().expect("bin has a parent");
274    let prefix = dir.join("prefix");
275    link_tree(real, &prefix, &[Path::new("share/postgresql/extension")]);
276    let extension = prefix.join("share/postgresql/extension");
277    for entry in std::fs::read_dir(ext.join("share/extension")).expect("built extension share") {
278        let entry = entry.expect("dir entry");
279        std::os::unix::fs::symlink(entry.path(), extension.join(entry.file_name())).expect("link extension file");
280    }
281    prefix.join("bin")
282}

Mirrors from at to: a symlink per entry, except that each of open (relative to from) and its ancestors is a real directory whose entries are linked in turn.

287fn link_tree(from: &Path, to: &Path, open: &[&Path]) {
288    std::fs::create_dir_all(to).expect("prefix dir");
289    for entry in std::fs::read_dir(from).expect("installation dir") {
290        let entry = entry.expect("dir entry");
291        let name = PathBuf::from(entry.file_name());
292        let inner: Vec<&Path> = open.iter().filter_map(|p| p.strip_prefix(&name).ok()).collect();
293        if inner.is_empty() {
294            std::os::unix::fs::symlink(entry.path(), to.join(&name)).expect("link installation entry");
295        } else {
296            let inner: Vec<&Path> = inner.into_iter().filter(|p| !p.as_os_str().is_empty()).collect();
297            link_tree(&entry.path(), &to.join(&name), &inner);
298        }
299    }
300}

302fn env(name: &str) -> String {
303    std::env::var(name).unwrap_or_else(|_| panic!("{name} is unset; run under buck2 test"))
304}