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

sofie_testing/
database.rs
