systemprompt_database/repository/
lease.rs

1//! [`LeaseRepository`] over `instance_leases`: one owner per key, expiring
2//! unless the owner claims again.
3//!
4//! A claim takes the row when it is absent, expired, already ours, or held by
5//! a replica that is no longer alive — the same judgement as
6//! [`systemprompt_traits::InstanceRegistry::is_alive`], evaluated inside the
7//! claiming statement so two replicas racing for a dead owner's lease cannot
8//! both win it.
9//!
10//! Copyright (c) systemprompt.io — Business Source License 1.1.
11//! See <https://systemprompt.io> for licensing details.
12
13use std::sync::Arc;
14use std::time::Duration;
15
16use async_trait::async_trait;
17use chrono::{DateTime, Utc};
18use sqlx::PgPool;
19use systemprompt_identifiers::{InstanceId, LeaseKey};
20use systemprompt_traits::{
21    CoordinationError, CoordinationResult, Lease, LeaseRepository, RepositoryError,
22};
23
24use super::service::{INSTANCE_CLAIM_CLASS, INSTANCE_DEAD_AFTER_SECS};
25use crate::DbPool;
26
27const MAX_LEASE_TTL: Duration = Duration::from_hours(24);
28
29#[derive(Debug, Clone)]
30pub struct InstanceLeaseRepository {
31    write_pool: Arc<PgPool>,
32    instance_id: InstanceId,
33}
34
35struct LeaseRow {
36    lease_key: LeaseKey,
37    owner_instance_id: String,
38    acquired_at: DateTime<Utc>,
39    expires_at: DateTime<Utc>,
40}
41
42impl LeaseRow {
43    fn into_lease(self) -> Lease {
44        Lease {
45            key: self.lease_key,
46            owner: InstanceId::new(self.owner_instance_id),
47            acquired_at: self.acquired_at,
48            expires_at: self.expires_at,
49        }
50    }
51}
52
53impl InstanceLeaseRepository {
54    pub fn new(db: &DbPool, instance_id: InstanceId) -> Self {
55        Self {
56            write_pool: db.write_pool(),
57            instance_id,
58        }
59    }
60
61    pub const fn instance_id(&self) -> &InstanceId {
62        &self.instance_id
63    }
64}
65
66fn ttl_secs(ttl: Duration) -> CoordinationResult<f64> {
67    if ttl < Duration::from_secs(1) || ttl > MAX_LEASE_TTL {
68        return Err(CoordinationError::InvalidTtl { ttl });
69    }
70    Ok(ttl.as_secs_f64())
71}
72
73#[async_trait]
74impl LeaseRepository for InstanceLeaseRepository {
75    async fn claim(&self, key: &LeaseKey, ttl: Duration) -> CoordinationResult<Option<Lease>> {
76        let ttl = ttl_secs(ttl)?;
77        let row = sqlx::query_as!(
78            LeaseRow,
79            r#"
80            INSERT INTO instance_leases AS held
81                (lease_key, owner_instance_id, acquired_at, expires_at)
82            VALUES ($1, $2, CURRENT_TIMESTAMP,
83                    CURRENT_TIMESTAMP + make_interval(secs => $3::double precision))
84            ON CONFLICT (lease_key) DO UPDATE SET
85                owner_instance_id = EXCLUDED.owner_instance_id,
86                acquired_at = CASE
87                    WHEN held.owner_instance_id = EXCLUDED.owner_instance_id
88                        THEN held.acquired_at
89                    ELSE EXCLUDED.acquired_at
90                END,
91                expires_at = EXCLUDED.expires_at
92            WHERE held.owner_instance_id = EXCLUDED.owner_instance_id
93               OR held.expires_at <= CURRENT_TIMESTAMP
94               OR NOT (
95                   EXISTS (
96                       SELECT 1 FROM pg_locks l
97                       WHERE l.locktype = 'advisory'
98                         AND l.granted
99                         AND l.classid = ($4::int)::oid
100                         AND l.objid = (hashtext(held.owner_instance_id))::oid
101                         AND l.objsubid = 2
102                   )
103                   OR EXISTS (
104                       SELECT 1 FROM services s
105                       WHERE s.instance_id = held.owner_instance_id
106                         AND s.heartbeat_at
107                             >= CURRENT_TIMESTAMP - make_interval(secs => $5::double precision)
108                   )
109               )
110            RETURNING lease_key AS "lease_key!: LeaseKey", owner_instance_id AS "owner_instance_id!",
111                      acquired_at AS "acquired_at!", expires_at AS "expires_at!"
112            "#,
113            key.as_str(),
114            self.instance_id.as_str(),
115            ttl,
116            INSTANCE_CLAIM_CLASS,
117            INSTANCE_DEAD_AFTER_SECS as f64
118        )
119        .fetch_optional(&*self.write_pool)
120        .await
121        .map_err(RepositoryError::from)?;
122        Ok(row.map(LeaseRow::into_lease))
123    }
124
125    async fn release(&self, key: &LeaseKey) -> CoordinationResult<bool> {
126        let result = sqlx::query!(
127            "DELETE FROM instance_leases WHERE lease_key = $1 AND owner_instance_id = $2",
128            key.as_str(),
129            self.instance_id.as_str()
130        )
131        .execute(&*self.write_pool)
132        .await
133        .map_err(RepositoryError::from)?;
134        Ok(result.rows_affected() > 0)
135    }
136
137    async fn holder(&self, key: &LeaseKey) -> CoordinationResult<Option<Lease>> {
138        let row = sqlx::query_as!(
139            LeaseRow,
140            r#"
141            SELECT lease_key AS "lease_key: LeaseKey", owner_instance_id, acquired_at, expires_at
142            FROM instance_leases
143            WHERE lease_key = $1 AND expires_at > CURRENT_TIMESTAMP
144            "#,
145            key.as_str()
146        )
147        .fetch_optional(&*self.write_pool)
148        .await
149        .map_err(RepositoryError::from)?;
150        Ok(row.map(LeaseRow::into_lease))
151    }
152}
systemprompt_database/repository/lease.rs

systemprompt_database/repository/
lease.rs
