Skip to main content

systemprompt_database/repository/service/
maintenance.rs

1//! Liveness maintenance on the `services` registry: heartbeats, the stale-row
2//! reaper for this instance, and the cross-instance reap of dead replicas.
3//!
4//! [`ServiceRepository::reap_stale`] is the only stale-row reaper. It deletes
5//! this instance's `error` rows and `running` rows with no pid outright, then
6//! asks the injected [`ProcessLiveness`] about every remaining `running` pid
7//! and deletes the rows whose process is gone or is not a child this
8//! installation spawned.
9//!
10//! Copyright (c) systemprompt.io — Business Source License 1.1.
11//! See <https://systemprompt.io> for licensing details.
12
13use systemprompt_identifiers::ServiceName;
14use systemprompt_manifest::services::{ServiceModule, ServiceStatus};
15use systemprompt_traits::ProcessLiveness;
16
17use super::claim::INSTANCE_CLAIM_CLASS;
18use super::repo::ServiceRepository;
19use crate::error::{DatabaseResult, RepositoryError};
20
21/// One `running` row whose recorded pid turned out not to be ours.
22#[derive(Debug, Clone, PartialEq, Eq)]
23pub struct ReapedService {
24    pub name: ServiceName,
25    pub module_name: ServiceModule,
26    pub pid: u32,
27}
28
29/// What one [`ServiceRepository::reap_stale`] pass removed.
30#[derive(Debug, Clone, Default, PartialEq, Eq)]
31pub struct Reaped {
32    pub stale_rows: u64,
33    pub dead: Vec<ReapedService>,
34}
35
36impl Reaped {
37    #[must_use]
38    pub const fn total(&self) -> u64 {
39        self.stale_rows + self.dead.len() as u64
40    }
41}
42
43impl ServiceRepository {
44    pub async fn reap_stale(&self, liveness: &dyn ProcessLiveness) -> DatabaseResult<Reaped> {
45        let stale_rows = sqlx::query!(
46            r#"
47            DELETE FROM services
48            WHERE instance_id = $1
49              AND (status = $2 OR (status = $3 AND pid IS NULL))
50            "#,
51            self.instance_id.as_str(),
52            ServiceStatus::Error.as_str(),
53            ServiceStatus::Running.as_str()
54        )
55        .execute(&*self.write_pool)
56        .await?
57        .rows_affected();
58
59        let candidates = sqlx::query!(
60            r#"
61            SELECT name, module_name, pid AS "pid!"
62            FROM services
63            WHERE instance_id = $1 AND status = $2 AND pid IS NOT NULL
64            ORDER BY name
65            "#,
66            self.instance_id.as_str(),
67            ServiceStatus::Running.as_str()
68        )
69        .fetch_all(&*self.write_pool)
70        .await?;
71
72        let mut dead = Vec::new();
73        for row in candidates {
74            let name = ServiceName::new(row.name);
75            let module_name = row.module_name.parse::<ServiceModule>().map_err(|source| {
76                RepositoryError::Decode {
77                    context: "services.module_name".to_owned(),
78                    source: Box::new(source),
79                }
80            })?;
81            let Ok(pid) = u32::try_from(row.pid) else {
82                self.delete_service(&name).await?;
83                dead.push(ReapedService {
84                    name,
85                    module_name,
86                    pid: 0,
87                });
88                continue;
89            };
90            if liveness.owns(pid, &name).await {
91                continue;
92            }
93            self.delete_service(&name).await?;
94            dead.push(ReapedService {
95                name,
96                module_name,
97                pid,
98            });
99        }
100
101        Ok(Reaped { stale_rows, dead })
102    }
103
104    pub async fn touch_heartbeat(&self) -> DatabaseResult<u64> {
105        let result = sqlx::query!(
106            r#"UPDATE services SET heartbeat_at = CURRENT_TIMESTAMP WHERE instance_id = $1"#,
107            self.instance_id.as_str()
108        )
109        .execute(&*self.write_pool)
110        .await?;
111        Ok(result.rows_affected())
112    }
113
114    pub async fn delete_dead_instances(&self, older_than_secs: i64) -> DatabaseResult<u64> {
115        let result = sqlx::query!(
116            r#"
117            DELETE FROM services
118            WHERE heartbeat_at < CURRENT_TIMESTAMP - make_interval(secs => $1::double precision)
119              AND NOT EXISTS (
120                  SELECT 1 FROM pg_locks l
121                  WHERE l.locktype = 'advisory'
122                    AND l.granted
123                    AND l.classid = ($2::int)::oid
124                    AND l.objid = (hashtext(services.instance_id))::oid
125                    AND l.objsubid = 2
126              )
127            "#,
128            older_than_secs as f64,
129            INSTANCE_CLAIM_CLASS
130        )
131        .execute(&*self.write_pool)
132        .await?;
133        Ok(result.rows_affected())
134    }
135}