systemprompt_database/repository/service/
repo.rs

1//! Async repository over the `services` registry table.
2//!
3//! Rows are keyed by `(instance_id, name)`: every replica registers, judges
4//! and reaps only the processes it spawned itself. The single cross-instance
5//! statement is [`ServiceRepository::delete_dead_instances`], which reaps rows
6//! whose heartbeat stopped, so a replica that vanished without cleanup is
7//! garbage-collected by the scheduler rather than by the next node to boot.
8//!
9//! Copyright (c) systemprompt.io — Business Source License 1.1.
10//! See <https://systemprompt.io> for licensing details.
11
12use std::sync::Arc;
13
14use sqlx::PgPool;
15use systemprompt_identifiers::InstanceId;
16
17use super::model::{CreateServiceInput, ServiceConfig};
18use crate::DbPool;
19use crate::error::DatabaseResult;
20
21#[derive(Debug, Clone)]
22pub struct ServiceRepository {
23    pool: Arc<PgPool>,
24    write_pool: Arc<PgPool>,
25    instance_id: InstanceId,
26}
27
28impl ServiceRepository {
29    pub fn new(db: &DbPool, instance_id: InstanceId) -> DatabaseResult<Self> {
30        let pool = db.pool_arc()?;
31        let write_pool = db.write_pool_arc()?;
32        Ok(Self {
33            pool,
34            write_pool,
35            instance_id,
36        })
37    }
38
39    pub const fn instance_id(&self) -> &InstanceId {
40        &self.instance_id
41    }
42
43    pub async fn find_service_by_name(&self, name: &str) -> DatabaseResult<Option<ServiceConfig>> {
44        let row = sqlx::query_as!(
45            ServiceConfig,
46            r#"
47            SELECT instance_id, name, module_name, status, pid, port, binary_mtime,
48                   heartbeat_at::text as "heartbeat_at!",
49                   created_at::text as "created_at!", updated_at::text as "updated_at!"
50            FROM services
51            WHERE instance_id = $1 AND name = $2
52            "#,
53            self.instance_id.as_str(),
54            name
55        )
56        .fetch_optional(&*self.pool)
57        .await?;
58        Ok(row)
59    }
60
61    pub async fn list_all_agent_service_names(&self) -> DatabaseResult<Vec<String>> {
62        let rows = sqlx::query!(
63            r#"SELECT name FROM services WHERE instance_id = $1 AND module_name = 'agent'"#,
64            self.instance_id.as_str()
65        )
66        .fetch_all(&*self.pool)
67        .await?;
68        Ok(rows.into_iter().map(|r| r.name).collect())
69    }
70
71    pub async fn list_mcp_services(&self) -> DatabaseResult<Vec<ServiceConfig>> {
72        self.list_services_by_type("mcp").await
73    }
74
75    pub async fn create_service(&self, input: CreateServiceInput<'_>) -> DatabaseResult<()> {
76        let port_i32 = i32::from(input.port);
77        sqlx::query!(
78            r#"
79            INSERT INTO services (instance_id, name, module_name, status, port, binary_mtime)
80            VALUES ($1, $2, $3, $4, $5, $6)
81            ON CONFLICT (instance_id, name) DO UPDATE SET
82              module_name = EXCLUDED.module_name,
83              status = EXCLUDED.status,
84              port = EXCLUDED.port,
85              binary_mtime = EXCLUDED.binary_mtime,
86              heartbeat_at = CURRENT_TIMESTAMP,
87              updated_at = CURRENT_TIMESTAMP
88            "#,
89            self.instance_id.as_str(),
90            input.name,
91            input.module_name,
92            input.status,
93            port_i32,
94            input.binary_mtime
95        )
96        .execute(&*self.write_pool)
97        .await?;
98        Ok(())
99    }
100
101    pub async fn update_service_status(
102        &self,
103        service_name: &str,
104        status: &str,
105    ) -> DatabaseResult<()> {
106        sqlx::query!(
107            r#"UPDATE services SET status = $1, updated_at = CURRENT_TIMESTAMP
108               WHERE instance_id = $2 AND name = $3"#,
109            status,
110            self.instance_id.as_str(),
111            service_name
112        )
113        .execute(&*self.write_pool)
114        .await?;
115        Ok(())
116    }
117
118    pub async fn delete_service(&self, service_name: &str) -> DatabaseResult<()> {
119        sqlx::query!(
120            r#"DELETE FROM services WHERE instance_id = $1 AND name = $2"#,
121            self.instance_id.as_str(),
122            service_name
123        )
124        .execute(&*self.write_pool)
125        .await?;
126        Ok(())
127    }
128
129    pub async fn update_service_pid(&self, service_name: &str, pid: i32) -> DatabaseResult<()> {
130        sqlx::query!(
131            r#"UPDATE services SET pid = $1, updated_at = CURRENT_TIMESTAMP
132               WHERE instance_id = $2 AND name = $3"#,
133            pid,
134            self.instance_id.as_str(),
135            service_name
136        )
137        .execute(&*self.write_pool)
138        .await?;
139        Ok(())
140    }
141
142    pub async fn clear_service_pid(&self, service_name: &str) -> DatabaseResult<()> {
143        sqlx::query!(
144            r#"UPDATE services SET pid = NULL, updated_at = CURRENT_TIMESTAMP
145               WHERE instance_id = $1 AND name = $2"#,
146            self.instance_id.as_str(),
147            service_name
148        )
149        .execute(&*self.write_pool)
150        .await?;
151        Ok(())
152    }
153
154    pub async fn list_all_running_services(&self) -> DatabaseResult<Vec<ServiceConfig>> {
155        let rows = sqlx::query_as!(
156            ServiceConfig,
157            r#"
158            SELECT instance_id, name, module_name, status, pid, port, binary_mtime,
159                   heartbeat_at::text as "heartbeat_at!",
160                   created_at::text as "created_at!", updated_at::text as "updated_at!"
161            FROM services
162            WHERE instance_id = $1 AND status = 'running'
163            ORDER BY name
164            "#,
165            self.instance_id.as_str()
166        )
167        .fetch_all(&*self.pool)
168        .await?;
169        Ok(rows)
170    }
171
172    pub async fn count_running_services(&self, module_name: &str) -> DatabaseResult<usize> {
173        let row = sqlx::query!(
174            r#"SELECT COUNT(*) as "count!" FROM services
175               WHERE instance_id = $1 AND module_name = $2 AND status = 'running'"#,
176            self.instance_id.as_str(),
177            module_name
178        )
179        .fetch_one(&*self.pool)
180        .await?;
181        Ok(usize::try_from(row.count).unwrap_or(0))
182    }
183
184    pub async fn mark_service_crashed(&self, service_name: &str) -> DatabaseResult<()> {
185        sqlx::query!(
186            r#"UPDATE services SET status = 'error', pid = NULL, updated_at = CURRENT_TIMESTAMP
187               WHERE instance_id = $1 AND name = $2"#,
188            self.instance_id.as_str(),
189            service_name
190        )
191        .execute(&*self.write_pool)
192        .await?;
193        Ok(())
194    }
195
196    pub async fn update_service_stopped(&self, service_name: &str) -> DatabaseResult<()> {
197        sqlx::query!(
198            r#"UPDATE services SET status = 'stopped', pid = NULL, updated_at = CURRENT_TIMESTAMP
199               WHERE instance_id = $1 AND name = $2"#,
200            self.instance_id.as_str(),
201            service_name
202        )
203        .execute(&*self.write_pool)
204        .await?;
205        Ok(())
206    }
207
208    pub async fn list_running_services_with_pid(&self) -> DatabaseResult<Vec<ServiceConfig>> {
209        self.list_all_running_services().await
210    }
211
212    pub async fn list_services_by_type(
213        &self,
214        module_name: &str,
215    ) -> DatabaseResult<Vec<ServiceConfig>> {
216        let rows = sqlx::query_as!(
217            ServiceConfig,
218            r#"
219            SELECT instance_id, name, module_name, status, pid, port, binary_mtime,
220                   heartbeat_at::text as "heartbeat_at!",
221                   created_at::text as "created_at!", updated_at::text as "updated_at!"
222            FROM services
223            WHERE instance_id = $1 AND module_name = $2
224            ORDER BY name
225            "#,
226            self.instance_id.as_str(),
227            module_name
228        )
229        .fetch_all(&*self.pool)
230        .await?;
231        Ok(rows)
232    }
233
234    pub async fn cleanup_stale_entries(&self) -> DatabaseResult<u64> {
235        let result = sqlx::query!(
236            r#"
237            DELETE FROM services
238            WHERE instance_id = $1
239              AND (status IN ('error', 'crashed')
240                   OR (status = 'running' AND pid IS NULL))
241            "#,
242            self.instance_id.as_str()
243        )
244        .execute(&*self.write_pool)
245        .await?;
246        Ok(result.rows_affected())
247    }
248
249    pub async fn touch_heartbeat(&self) -> DatabaseResult<u64> {
250        let result = sqlx::query!(
251            r#"UPDATE services SET heartbeat_at = CURRENT_TIMESTAMP WHERE instance_id = $1"#,
252            self.instance_id.as_str()
253        )
254        .execute(&*self.write_pool)
255        .await?;
256        Ok(result.rows_affected())
257    }
258
259    pub async fn delete_dead_instances(&self, older_than_secs: i64) -> DatabaseResult<u64> {
260        let result = sqlx::query!(
261            r#"
262            DELETE FROM services
263            WHERE heartbeat_at < CURRENT_TIMESTAMP - make_interval(secs => $1::double precision)
264            "#,
265            older_than_secs as f64
266        )
267        .execute(&*self.write_pool)
268        .await?;
269        Ok(result.rows_affected())
270    }
271}
systemprompt_database/repository/service/repo.rs

systemprompt_database/repository/service/
repo.rs
