systemprompt_database/lifecycle/installation/extension/
lock.rs

1//! Session-pinned Postgres advisory lock serialising concurrent bootstraps.
2//!
3//! Copyright (c) systemprompt.io — Business Source License 1.1.
4//! See <https://systemprompt.io> for licensing details.
5
6use sqlx::Postgres;
7use sqlx::pool::PoolConnection;
8use systemprompt_extension::LoaderError;
9use tracing::{debug, warn};
10
11use crate::services::DatabaseProvider;
12
13pub const BOOTSTRAP_ADVISORY_LOCK_KEY: i64 = 0x73_70_72_6F_6D_70_74_01;
14
15/// Session-pinned advisory lock serialising concurrent bootstraps.
16///
17/// Only the acquiring Postgres session can release its advisory lock, so the
18/// guard pins that connection. Dropping the guard without `release` closes
19/// the session instead of returning it to the pool, so the lock never
20/// outlives a cancelled or panicking holder.
21pub struct BootstrapLockGuard {
22    conn: Option<PoolConnection<Postgres>>,
23}
24
25impl std::fmt::Debug for BootstrapLockGuard {
26    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
27        f.debug_struct("BootstrapLockGuard")
28            .field("key", &BOOTSTRAP_ADVISORY_LOCK_KEY)
29            .field("held", &self.conn.is_some())
30            .finish()
31    }
32}
33
34impl BootstrapLockGuard {
35    pub async fn acquire(db: &dyn DatabaseProvider) -> Result<Self, LoaderError> {
36        let mut conn = db.get_postgres_pool().acquire().await.map_err(|e| {
37            LoaderError::SchemaInstallationFailed {
38                extension: "database".to_owned(),
39                message: format!("Failed to acquire bootstrap lock connection: {e}"),
40            }
41        })?;
42
43        sqlx::query!("SELECT pg_advisory_lock($1)", BOOTSTRAP_ADVISORY_LOCK_KEY)
44            .execute(conn.as_mut())
45            .await
46            .map_err(|e| LoaderError::SchemaInstallationFailed {
47                extension: "database".to_owned(),
48                message: format!("Failed to acquire bootstrap advisory lock: {e}"),
49            })?;
50
51        debug!(
52            key = BOOTSTRAP_ADVISORY_LOCK_KEY,
53            "Acquired bootstrap advisory lock"
54        );
55
56        Ok(Self { conn: Some(conn) })
57    }
58
59    pub async fn release(mut self) {
60        let Some(mut conn) = self.conn.take() else {
61            return;
62        };
63        match sqlx::query_scalar!("SELECT pg_advisory_unlock($1)", BOOTSTRAP_ADVISORY_LOCK_KEY)
64            .fetch_one(conn.as_mut())
65            .await
66        {
67            Ok(Some(true)) => drop(conn),
68            Ok(released) => {
69                warn!(
70                    key = BOOTSTRAP_ADVISORY_LOCK_KEY,
71                    ?released,
72                    "Bootstrap advisory lock was not held by this session at release"
73                );
74                drop(conn);
75            },
76            Err(e) => {
77                warn!(
78                    error = %e,
79                    "Failed to release bootstrap advisory lock; closing its session instead of pooling it"
80                );
81                let session = conn.detach();
82                drop(session);
83            },
84        }
85    }
86}
87
88impl Drop for BootstrapLockGuard {
89    fn drop(&mut self) {
90        // Why: a connection returned to the pool keeps its session, and with it
91        // the advisory lock; detaching closes the session so a cancelled or
92        // panicking install cannot leave every other replica blocked.
93        if let Some(conn) = self.conn.take() {
94            warn!(
95                key = BOOTSTRAP_ADVISORY_LOCK_KEY,
96                "BootstrapLockGuard dropped without explicit release; closing its session"
97            );
98            let session = conn.detach();
99            drop(session);
100        }
101    }
102}
systemprompt_database/lifecycle/installation/extension/lock.rs

systemprompt_database/lifecycle/installation/extension/
lock.rs
