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//! Copyright (c) systemprompt.io — Business Source License 1.1.
20//! See <https://systemprompt.io> for licensing details.
21
22use sqlx::{Connection, PgConnection};
23use systemprompt_identifiers::InstanceId;
24use systemprompt_traits::RepositoryError;
25use tracing::warn;
26
27use super::repo::ServiceRepository;
28
29// Why: Postgres advisory locks share one key space per database; the
30// two-key form with a dedicated class keeps instance claims disjoint from
31// the single-key `hashtext(job_name)` job locks.
32pub const INSTANCE_CLAIM_CLASS: i32 = 0x5350_0001;
33
34pub const CLAIM_DEAD_PEER_SECS: u64 = 25;
35
36/// Failure to claim this process's replica identity.
37#[derive(Debug, thiserror::Error)]
38pub enum InstanceClaimError {
39    #[error(
40        "instance id '{instance_id}' is already claimed by a live process; give each replica its \
41         own server.instance_id or leave it unset so HOSTNAME is used"
42    )]
43    Claimed { instance_id: InstanceId },
44
45    #[error("could not claim instance id: {0}")]
46    Repository(#[from] RepositoryError),
47}
48
49impl From<sqlx::Error> for InstanceClaimError {
50    fn from(err: sqlx::Error) -> Self {
51        Self::Repository(RepositoryError::from(err))
52    }
53}
54
55/// Exclusive claim on an instance id, held on a dedicated session.
56pub struct InstanceClaim {
57    session: Option<PgConnection>,
58    instance_id: InstanceId,
59}
60
61impl InstanceClaim {
62    pub const fn instance_id(&self) -> &InstanceId {
63        &self.instance_id
64    }
65
66    pub async fn release(mut self) {
67        if let Some(mut session) = self.session.take() {
68            if let Err(e) = sqlx::query_scalar!(
69                "SELECT pg_advisory_unlock($1, hashtext($2))",
70                INSTANCE_CLAIM_CLASS,
71                self.instance_id.as_str()
72            )
73            .fetch_one(&mut session)
74            .await
75            {
76                warn!(
77                    instance_id = %self.instance_id,
78                    error = %e,
79                    "Failed to release instance claim; closing its session"
80                );
81            }
82            if let Err(e) = session.close().await {
83                warn!(
84                    instance_id = %self.instance_id,
85                    error = %e,
86                    "Failed to close instance claim session"
87                );
88            }
89        }
90    }
91}
92
93impl std::fmt::Debug for InstanceClaim {
94    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
95        f.debug_struct("InstanceClaim")
96            .field("instance_id", &self.instance_id)
97            .field("held", &self.session.is_some())
98            .finish()
99    }
100}
101
102impl ServiceRepository {
103    pub async fn claim_instance(&self) -> Result<InstanceClaim, InstanceClaimError> {
104        let mut session = self.write_pool.acquire().await?.detach();
105        // Why: idle 10 s + 3 probes 5 s apart bounds a dead peer at
106        // CLAIM_DEAD_PEER_SECS; the server-side settings are what free a
107        // session whose client host disappeared without closing it.
108        sqlx::query_scalar!(
109            r#"SELECT set_config('tcp_keepalives_idle', '10', false)
110                   || set_config('tcp_keepalives_interval', '5', false)
111                   || set_config('tcp_keepalives_count', '3', false) AS "applied!""#
112        )
113        .fetch_one(&mut session)
114        .await?;
115        let acquired = sqlx::query_scalar!(
116            r#"SELECT pg_try_advisory_lock($1, hashtext($2)) AS "acquired!""#,
117            INSTANCE_CLAIM_CLASS,
118            self.instance_id.as_str()
119        )
120        .fetch_one(&mut session)
121        .await?;
122
123        if !acquired {
124            if let Err(e) = session.close().await {
125                warn!(
126                    instance_id = %self.instance_id,
127                    error = %e,
128                    "Failed to close refused claim session"
129                );
130            }
131            return Err(InstanceClaimError::Claimed {
132                instance_id: self.instance_id.clone(),
133            });
134        }
135
136        Ok(InstanceClaim {
137            session: Some(session),
138            instance_id: self.instance_id.clone(),
139        })
140    }
141}
systemprompt_database/repository/service/claim.rs

systemprompt_database/repository/service/
claim.rs
