systemprompt_database/repository/service/
presence.rs

1//! Per-replica role presence: every replica records its `server.role` with
2//! each registry heartbeat, so the live replica count per role is readable
3//! cluster-wide, including for roles that register no process, and the live
4//! members of one role are listable for spreading work across them.
5//!
6//! Copyright (c) systemprompt.io — Business Source License 1.1.
7//! See <https://systemprompt.io> for licensing details.
8
9use systemprompt_identifiers::InstanceId;
10use systemprompt_manifest::profile::NodeRole;
11
12use super::liveness::INSTANCE_DEAD_AFTER_SECS;
13use super::repo::ServiceRepository;
14use crate::error::DatabaseResult;
15
16/// The number of live replicas currently heartbeating one role.
17#[derive(Debug, Clone, PartialEq, Eq)]
18pub struct RoleCount {
19    pub role: String,
20    pub instances: i64,
21}
22
23impl ServiceRepository {
24    pub async fn touch_presence(&self, role: NodeRole) -> DatabaseResult<()> {
25        sqlx::query!(
26            r#"
27            INSERT INTO instance_presence (instance_id, role, heartbeat_at)
28            VALUES ($1, $2, CURRENT_TIMESTAMP)
29            ON CONFLICT (instance_id) DO UPDATE SET
30                role = EXCLUDED.role,
31                heartbeat_at = EXCLUDED.heartbeat_at
32            "#,
33            self.instance_id.as_str(),
34            role.as_str()
35        )
36        .execute(&*self.write_pool)
37        .await?;
38        Ok(())
39    }
40
41    pub async fn count_live_instances_by_role(&self) -> DatabaseResult<Vec<RoleCount>> {
42        let rows = sqlx::query_as!(
43            RoleCount,
44            r#"
45            SELECT role, COUNT(*) AS "instances!"
46            FROM instance_presence
47            WHERE heartbeat_at
48                  >= CURRENT_TIMESTAMP - make_interval(secs => $1::double precision)
49            GROUP BY role
50            ORDER BY role
51            "#,
52            INSTANCE_DEAD_AFTER_SECS as f64
53        )
54        .fetch_all(&*self.pool)
55        .await?;
56        Ok(rows)
57    }
58
59    pub async fn list_reachable_instances(
60        &self,
61        role: NodeRole,
62        fresh_within_secs: i64,
63    ) -> DatabaseResult<Vec<InstanceId>> {
64        let rows = sqlx::query_scalar!(
65            r#"
66            SELECT p.instance_id AS "instance_id!: InstanceId"
67            FROM instance_presence p
68            JOIN instance_endpoints e ON e.instance_id = p.instance_id
69            WHERE p.role = $1
70              AND p.heartbeat_at
71                  >= CURRENT_TIMESTAMP - make_interval(secs => $2::double precision)
72              AND e.heartbeat_at
73                  >= CURRENT_TIMESTAMP - make_interval(secs => $2::double precision)
74            ORDER BY p.instance_id
75            "#,
76            role.as_str(),
77            fresh_within_secs as f64
78        )
79        .fetch_all(&*self.write_pool)
80        .await?;
81        Ok(rows)
82    }
83
84    pub async fn delete_stale_presence(&self, older_than_secs: i64) -> DatabaseResult<u64> {
85        let result = sqlx::query!(
86            r#"
87            DELETE FROM instance_presence
88            WHERE heartbeat_at < CURRENT_TIMESTAMP - make_interval(secs => $1::double precision)
89            "#,
90            older_than_secs as f64
91        )
92        .execute(&*self.write_pool)
93        .await?;
94        Ok(result.rows_affected())
95    }
96}
systemprompt_database/repository/service/presence.rs

systemprompt_database/repository/service/
presence.rs
