sofie_testing/
database.rs

1use std::sync::Mutex;
2use std::sync::atomic::{AtomicU32, Ordering};
3use std::thread::JoinHandle;
4
5use sofie_core::unpoisoned;
6use tokio::sync::{OnceCell, oneshot};
7use tokio_postgres::{Client, NoTls};
8
9const SERVER_URL_VARIABLE: &str = "SOFIE_TEST_DATABASE_URL";
10const DEFAULT_SERVER_URL: &str = "postgres://sofie:sofie@127.0.0.1:25432/sofie";
11const NAME_PREFIX: &str = "sofie_test_";
12/// Held per process for its lifetime; a database whose process no longer holds it was left behind and is dropped.
13const LIVENESS_CLASS: i32 = 0x6677_7465;
14
15static NEXT_DATABASE: AtomicU32 = AtomicU32::new(0);
16static LIVENESS: OnceCell<()> = OnceCell::const_new();
17static PENDING_DROPS: Mutex<Vec<JoinHandle<()>>> = Mutex::new(Vec::new());
18
19pub struct TestDatabase {
20    name: String,
21    server_url: String,
22    url: String,
23}
24
25impl TestDatabase {
26    pub async fn create() -> Self {
27        let server_url =
28            std::env::var(SERVER_URL_VARIABLE).unwrap_or_else(|_| DEFAULT_SERVER_URL.to_owned());
29        LIVENESS
30            .get_or_init(|| hold_liveness(server_url.clone()))
31            .await;
32        let name = format!(
33            "{NAME_PREFIX}{}_{}",
34            std::process::id(),
35            NEXT_DATABASE.fetch_add(1, Ordering::Relaxed)
36        );
37        run_on_server(
38            &server_url,
39            &[drop_statement(&name), format!("CREATE DATABASE {name}")],
40        )
41        .await;
42        Self {
43            url: with_database(&server_url, &name),
44            name,
45            server_url,
46        }
47    }
48
49    #[must_use]
50    pub fn url(&self) -> &str {
51        &self.url
52    }
53}
54
55/// Never blocks: Postgres waits for a half-open connection of the test's runtime, which runs only once the runtime does.
56impl Drop for TestDatabase {
57    fn drop(&mut self) {
58        let statement = drop_statement(&self.name);
59        let server_url = self.server_url.clone();
60        let dropping = std::thread::spawn(move || {
61            on_own_runtime(run_on_server(&server_url, &[statement]));
62        });
63        unpoisoned::lock(&PENDING_DROPS).push(dropping);
64    }
65}
66
67/// Called outside a runtime once connections are closed; a drop still running at exit is left to the next test process.
68pub fn finish_teardown() {
69    let pending = std::mem::take(&mut *unpoisoned::lock(&PENDING_DROPS));
70    for dropping in pending {
71        if let Err(panic) = dropping.join() {
72            std::panic::resume_unwind(panic);
73        }
74    }
75}
76
77fn on_own_runtime(work: impl Future<Output = ()>) {
78    tokio::runtime::Builder::new_current_thread()
79        .enable_all()
80        .build()
81        .expect("building a runtime for the test server")
82        .block_on(work);
83}
84
85async fn hold_liveness(server_url: String) {
86    let (locked, held) = oneshot::channel();
87    std::thread::spawn(move || {
88        on_own_runtime(async move {
89            let (client, connection) = connect(&server_url).await;
90            let connection = tokio::spawn(connection);
91            client
92                .execute(
93                    "SELECT pg_advisory_lock($1, $2)",
94                    &[&LIVENESS_CLASS, &own_process()],
95                )
96                .await
97                .expect("taking the test process's liveness lock");
98            let _ = locked.send(());
99            drop_abandoned(&client).await;
100            let _ = connection.await;
101        });
102    });
103    held.await
104        .expect("the liveness connection takes its lock or panics");
105}
106
107/// Dropped by the process that takes the dead one's lock, so two processes never drop the same one.
108async fn drop_abandoned(client: &Client) {
109    let databases = client
110        .query(
111            "SELECT datname FROM pg_database WHERE datname LIKE $1",
112            &[&format!("{NAME_PREFIX}%")],
113        )
114        .await
115        .expect("listing the test databases");
116    let mut dead = Vec::new();
117    for row in databases {
118        let name: String = row.get(0);
119        let Some(process) = process_of(&name) else {
120            continue;
121        };
122        if process != own_process() && !dead.contains(&process) && takes_lock(client, process).await
123        {
124            dead.push(process);
125        }
126        if dead.contains(&process) {
127            client
128                .batch_execute(&drop_statement(&name))
129                .await
130                .unwrap_or_else(|error| panic!("dropping abandoned database {name}: {error}"));
131        }
132    }
133}
134
135async fn takes_lock(client: &Client, process: i32) -> bool {
136    client
137        .query_one(
138            "SELECT pg_try_advisory_lock($1, $2)",
139            &[&LIVENESS_CLASS, &process],
140        )
141        .await
142        .expect("trying a test process's liveness lock")
143        .get(0)
144}
145
146fn drop_statement(name: &str) -> String {
147    format!("DROP DATABASE IF EXISTS {name} WITH (FORCE)")
148}
149
150fn own_process() -> i32 {
151    i32::try_from(std::process::id()).expect("a process id fits an i32")
152}
153
154fn process_of(name: &str) -> Option<i32> {
155    name.strip_prefix(NAME_PREFIX)?
156        .split_once('_')?
157        .0
158        .parse()
159        .ok()
160}
161
162async fn connect(
163    server_url: &str,
164) -> (
165    Client,
166    tokio_postgres::Connection<tokio_postgres::Socket, tokio_postgres::tls::NoTlsStream>,
167) {
168    tokio_postgres::connect(server_url, NoTls)
169        .await
170        .expect("connecting to the test server; is `docker compose up -d --wait` running?")
171}
172
173// Databases cannot be created or dropped in the implicit transaction of a multi-statement script.
174async fn run_on_server(server_url: &str, statements: &[String]) {
175    let (client, connection) = connect(server_url).await;
176    let connection = tokio::spawn(connection);
177    for sql in statements {
178        client
179            .batch_execute(sql)
180            .await
181            .unwrap_or_else(|error| panic!("running `{sql}`: {error}"));
182    }
183    drop(client);
184    connection
185        .await
186        .expect("joining the connection task")
187        .expect("closing the connection");
188}
189
190fn with_database(server_url: &str, name: &str) -> String {
191    let (base, options) = match server_url.split_once('?') {
192        Some((base, options)) => (base, Some(options)),
193        None => (server_url, None),
194    };
195    let host_start = base.find("://").map_or(0, |index| index + 3);
196    let path_start = base[host_start..]
197        .find('/')
198        .map_or(base.len(), |index| host_start + index);
199    let mut url = format!("{}/{name}", &base[..path_start]);
200    if let Some(options) = options {
201        url.push('?');
202        url.push_str(options);
203    }
204    url
205}
206
207#[cfg(test)]
208mod tests {
209    use super::{process_of, with_database};
210
211    #[test]
212    fn replaces_only_the_database_of_the_server_url() {
213        for (server_url, expected) in [
214            (
215                "postgres://sofie:sofie@127.0.0.1:25432/sofie",
216                "postgres://sofie:sofie@127.0.0.1:25432/t",
217            ),
218            ("postgres://db.test", "postgres://db.test/t"),
219            (
220                "postgres://a@db.test/x?sslmode=disable",
221                "postgres://a@db.test/t?sslmode=disable",
222            ),
223        ] {
224            assert_eq!(with_database(server_url, "t"), expected, "{server_url}");
225        }
226    }
227
228    #[test]
229    fn a_test_database_is_named_after_its_process() {
230        assert_eq!(process_of("sofie_test_4242_7"), Some(4242));
231        assert_eq!(process_of("sofie_test_x_7"), None);
232        assert_eq!(process_of("sofie"), None);
233    }
234}
sofie_testing/database.rs

sofie_testing/
database.rs
