systemprompt_database/repository/service/
claim.rs

1//! Boot-time claim on this process's replica identity.
2//!
3//! Every `services` row is keyed by `(instance_id, name)`, and each replica
4//! reconciles (and reaps) only its own rows. Two live processes sharing one
5//! instance id would therefore evict each other's MCP registrations. A
6//! session-scoped Postgres advisory lock keyed on the instance id makes the
7//! identity exclusive for the process lifetime: a second process with the
8//! same id is refused at boot, and the claim vanishes with the session on a
9//! crash, so no stale claim ever needs reaping.
10//!
11//! A crash only frees the claim once the server notices the session is gone.
12//! A replaced host (a Fly machine redeployed under the same `HOSTNAME`) never
13//! sends a FIN, and with the default two-hour keepalive the server would hold
14//! the dead session — and the identity — long after the process died. The
15//! claim session therefore asks the server to probe its client every few
16//! seconds, so a vanished holder is torn down within
17//! [`CLAIM_DEAD_PEER_SECS`].
18//!
19//! The session is detached from the session pool (`database_session_url`
20//! when set), never from a pool that may sit behind a transaction-mode
21//! pooler, which would hand the locked backend to another client.
22//!
23//! Copyright (c) systemprompt.io — Business Source License 1.1.
24//! See <https://systemprompt.io> for licensing details.
25
26use sqlx::{Connection, PgConnection};
27use systemprompt_identifiers::InstanceId;
28use systemprompt_traits::RepositoryError;
29use tracing::warn;
30
31use super::repo::ServiceRepository;
32
33// Why: Postgres advisory locks share one key space per database; the
34// two-key form with a dedicated class keeps instance claims disjoint from
35// the single-key `hashtext(job_name)` job locks.
36pub const INSTANCE_CLAIM_CLASS: i32 = 0x5350_0001;
37
38pub const CLAIM_DEAD_PEER_SECS: u64 = 25;
39
40/// Failure to claim this process's replica identity.
41#[derive(Debug, thiserror::Error)]
42pub enum InstanceClaimError {
43    #[error(
44        "instance id '{instance_id}' is already claimed by a live process; give each replica its \
45         own server.instance_id or leave it unset so HOSTNAME is used"
46    )]
47    Claimed { instance_id: InstanceId },
48
49    #[error("could not claim instance id: {0}")]
50    Repository(#[from] RepositoryError),
51}
52
53impl From<sqlx::Error> for InstanceClaimError {
54    fn from(err: sqlx::Error) -> Self {
55        Self::Repository(RepositoryError::from(err))
56    }
57}
58
59/// Exclusive claim on an instance id, held on a dedicated session.
60pub struct InstanceClaim {
61    session: Option<PgConnection>,
62    instance_id: InstanceId,
63}
64
65impl InstanceClaim {
66    pub const fn instance_id(&self) -> &InstanceId {
67        &self.instance_id
68    }
69
70    pub async fn release(mut self) {
71        if let Some(mut session) = self.session.take() {
72            if let Err(e) = sqlx::query_scalar!(
73                "SELECT pg_advisory_unlock($1, hashtext($2))",
74                INSTANCE_CLAIM_CLASS,
75                self.instance_id.as_str()
76            )
77            .fetch_one(&mut session)
78            .await
79            {
80                warn!(
81                    instance_id = %self.instance_id,
82                    error = %e,
83                    "Failed to release instance claim; closing its session"
84                );
85            }
86            if let Err(e) = session.close().await {
87                warn!(
88                    instance_id = %self.instance_id,
89                    error = %e,
90                    "Failed to close instance claim session"
91                );
92            }
93        }
94    }
95}
96
97impl std::fmt::Debug for InstanceClaim {
98    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
99        f.debug_struct("InstanceClaim")
100            .field("instance_id", &self.instance_id)
101            .field("held", &self.session.is_some())
102            .finish()
103    }
104}
105
106impl ServiceRepository {
107    pub async fn claim_instance(&self) -> Result<InstanceClaim, InstanceClaimError> {
108        let mut session = self.session_pool.acquire().await?.detach();
109        // Why: idle 10 s + 3 probes 5 s apart bounds a dead peer at
110        // CLAIM_DEAD_PEER_SECS; the server-side settings are what free a
111        // session whose client host disappeared without closing it.
112        sqlx::query_scalar!(
113            r#"SELECT set_config('tcp_keepalives_idle', '10', false)
114                   || set_config('tcp_keepalives_interval', '5', false)
115                   || set_config('tcp_keepalives_count', '3', false) AS "applied!""#
116        )
117        .fetch_one(&mut session)
118        .await?;
119        let acquired = sqlx::query_scalar!(
120            r#"SELECT pg_try_advisory_lock($1, hashtext($2)) AS "acquired!""#,
121            INSTANCE_CLAIM_CLASS,
122            self.instance_id.as_str()
123        )
124        .fetch_one(&mut session)
125        .await?;
126
127        if !acquired {
128            if let Err(e) = session.close().await {
129                warn!(
130                    instance_id = %self.instance_id,
131                    error = %e,
132                    "Failed to close refused claim session"
133                );
134            }
135            return Err(InstanceClaimError::Claimed {
136                instance_id: self.instance_id.clone(),
137            });
138        }
139
140        Ok(InstanceClaim {
141            session: Some(session),
142            instance_id: self.instance_id.clone(),
143        })
144    }
145}
systemprompt_database/repository/service/claim.rs

systemprompt_database/repository/service/
claim.rs
