systemprompt_database/repository/
advisory_lock.rs

1//! [`AdvisoryLocks`] for extensions: session-scoped Postgres advisory locks
2//! taken on the session pool.
3//!
4//! An advisory lock belongs to the backend session that took it, so the
5//! connection must stay pinned until release and must never come from a pool
6//! behind a transaction-mode pooler, which would hand the locked backend to
7//! another client between statements. Keys use the two-key form under
8//! [`EXTENSION_LOCK_CLASS`], disjoint from the scheduler's single-key job
9//! locks and from the instance claims.
10//!
11//! Copyright (c) systemprompt.io — Business Source License 1.1.
12//! See <https://systemprompt.io> for licensing details.
13
14use std::sync::Arc;
15use std::time::Duration;
16
17use async_trait::async_trait;
18use sqlx::pool::PoolConnection;
19use sqlx::{PgPool, Postgres};
20use systemprompt_identifiers::AdvisoryLockKey;
21use systemprompt_traits::{
22    AdvisoryLockGuard, AdvisoryLocks, CoordinationResult, HeldAdvisoryLock, RepositoryError,
23};
24use tokio::time::Instant;
25use tracing::warn;
26
27use crate::DbPool;
28
29// Why: Postgres advisory locks share one key space per database; a class of
30// its own keeps extension locks from colliding with job or instance locks.
31pub const EXTENSION_LOCK_CLASS: i32 = 0x5350_0002;
32
33const LOCK_POLL_INTERVAL: Duration = Duration::from_millis(250);
34
35#[derive(Debug, Clone)]
36pub struct SessionAdvisoryLocks {
37    session_pool: Arc<PgPool>,
38}
39
40impl SessionAdvisoryLocks {
41    pub fn new(db: &DbPool) -> Self {
42        Self {
43            session_pool: db.session_pool(),
44        }
45    }
46}
47
48#[async_trait]
49impl AdvisoryLocks for SessionAdvisoryLocks {
50    async fn try_acquire(
51        &self,
52        key: &AdvisoryLockKey,
53    ) -> CoordinationResult<Option<AdvisoryLockGuard>> {
54        let mut conn = self
55            .session_pool
56            .acquire()
57            .await
58            .map_err(RepositoryError::from)?;
59        let acquired = sqlx::query_scalar!(
60            r#"SELECT pg_try_advisory_lock($1, hashtext($2)) AS "acquired!""#,
61            EXTENSION_LOCK_CLASS,
62            key.as_str()
63        )
64        .fetch_one(conn.as_mut())
65        .await
66        .map_err(RepositoryError::from)?;
67        if !acquired {
68            return Ok(None);
69        }
70        Ok(Some(Box::new(SessionAdvisoryLock {
71            conn: Some(conn),
72            key: key.clone(),
73        })))
74    }
75
76    async fn acquire_within(
77        &self,
78        key: &AdvisoryLockKey,
79        wait: Duration,
80    ) -> CoordinationResult<Option<AdvisoryLockGuard>> {
81        let deadline = Instant::now() + wait;
82        loop {
83            if let Some(guard) = self.try_acquire(key).await? {
84                return Ok(Some(guard));
85            }
86            let now = Instant::now();
87            if now >= deadline {
88                return Ok(None);
89            }
90            tokio::time::sleep(LOCK_POLL_INTERVAL.min(deadline - now)).await;
91        }
92    }
93}
94
95pub struct SessionAdvisoryLock {
96    conn: Option<PoolConnection<Postgres>>,
97    key: AdvisoryLockKey,
98}
99
100impl std::fmt::Debug for SessionAdvisoryLock {
101    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
102        f.debug_struct("SessionAdvisoryLock")
103            .field("key", &self.key)
104            .field("held", &self.conn.is_some())
105            .finish()
106    }
107}
108
109impl Drop for SessionAdvisoryLock {
110    fn drop(&mut self) {
111        // Why: a connection returned to the pool keeps its session, and with it
112        // the lock; detaching closes the session so the lock is freed now.
113        if let Some(conn) = self.conn.take() {
114            warn!(
115                key = %self.key,
116                "advisory lock dropped without explicit release; closing its session"
117            );
118            let session = conn.detach();
119            drop(session);
120        }
121    }
122}
123
124#[async_trait]
125impl HeldAdvisoryLock for SessionAdvisoryLock {
126    fn key(&self) -> &AdvisoryLockKey {
127        &self.key
128    }
129
130    async fn release(mut self: Box<Self>) {
131        let Some(mut conn) = self.conn.take() else {
132            return;
133        };
134        let unlocked = sqlx::query_scalar!(
135            r#"SELECT pg_advisory_unlock($1, hashtext($2)) AS "unlocked!""#,
136            EXTENSION_LOCK_CLASS,
137            self.key.as_str()
138        )
139        .fetch_one(conn.as_mut())
140        .await;
141        match unlocked {
142            Ok(true) => {},
143            Ok(false) => {
144                warn!(key = %self.key, "advisory lock was not held at release; closing its session");
145                let session = conn.detach();
146                drop(session);
147            },
148            Err(error) => {
149                warn!(key = %self.key, %error, "advisory lock release failed; closing its session");
150                let session = conn.detach();
151                drop(session);
152            },
153        }
154    }
155}
systemprompt_database/repository/advisory_lock.rs

systemprompt_database/repository/
advisory_lock.rs
