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 update_service_port(&self, service_name: &str, port: u16) -> DatabaseResult<()> {
143        let port_i32 = i32::from(port);
144        sqlx::query!(
145            r#"UPDATE services SET port = $1, updated_at = CURRENT_TIMESTAMP
146               WHERE instance_id = $2 AND name = $3"#,
147            port_i32,
148            self.instance_id.as_str(),
149            service_name
150        )
151        .execute(&*self.write_pool)
152        .await?;
153        Ok(())
154    }
155
156    pub async fn clear_service_pid(&self, service_name: &str) -> DatabaseResult<()> {
157        sqlx::query!(
158            r#"UPDATE services SET pid = NULL, updated_at = CURRENT_TIMESTAMP
159               WHERE instance_id = $1 AND name = $2"#,
160            self.instance_id.as_str(),
161            service_name
162        )
163        .execute(&*self.write_pool)
164        .await?;
165        Ok(())
166    }
167
168    pub async fn list_all_running_services(&self) -> DatabaseResult<Vec<ServiceConfig>> {
169        let rows = sqlx::query_as!(
170            ServiceConfig,
171            r#"
172            SELECT instance_id, name, module_name, status, pid, port, binary_mtime,
173                   heartbeat_at::text as "heartbeat_at!",
174                   created_at::text as "created_at!", updated_at::text as "updated_at!"
175            FROM services
176            WHERE instance_id = $1 AND status = 'running'
177            ORDER BY name
178            "#,
179            self.instance_id.as_str()
180        )
181        .fetch_all(&*self.pool)
182        .await?;
183        Ok(rows)
184    }
185
186    pub async fn count_running_services(&self, module_name: &str) -> DatabaseResult<usize> {
187        let row = sqlx::query!(
188            r#"SELECT COUNT(*) as "count!" FROM services
189               WHERE instance_id = $1 AND module_name = $2 AND status = 'running'"#,
190            self.instance_id.as_str(),
191            module_name
192        )
193        .fetch_one(&*self.pool)
194        .await?;
195        Ok(usize::try_from(row.count).unwrap_or(0))
196    }
197
198    pub async fn mark_service_crashed(&self, service_name: &str) -> DatabaseResult<()> {
199        sqlx::query!(
200            r#"UPDATE services SET status = 'error', pid = NULL, updated_at = CURRENT_TIMESTAMP
201               WHERE instance_id = $1 AND name = $2"#,
202            self.instance_id.as_str(),
203            service_name
204        )
205        .execute(&*self.write_pool)
206        .await?;
207        Ok(())
208    }
209
210    pub async fn update_service_stopped(&self, service_name: &str) -> DatabaseResult<()> {
211        sqlx::query!(
212            r#"UPDATE services SET status = 'stopped', pid = NULL, updated_at = CURRENT_TIMESTAMP
213               WHERE instance_id = $1 AND name = $2"#,
214            self.instance_id.as_str(),
215            service_name
216        )
217        .execute(&*self.write_pool)
218        .await?;
219        Ok(())
220    }
221
222    pub async fn list_running_services_with_pid(&self) -> DatabaseResult<Vec<ServiceConfig>> {
223        self.list_all_running_services().await
224    }
225
226    pub async fn list_services_by_type(
227        &self,
228        module_name: &str,
229    ) -> DatabaseResult<Vec<ServiceConfig>> {
230        let rows = sqlx::query_as!(
231            ServiceConfig,
232            r#"
233            SELECT instance_id, name, module_name, status, pid, port, binary_mtime,
234                   heartbeat_at::text as "heartbeat_at!",
235                   created_at::text as "created_at!", updated_at::text as "updated_at!"
236            FROM services
237            WHERE instance_id = $1 AND module_name = $2
238            ORDER BY name
239            "#,
240            self.instance_id.as_str(),
241            module_name
242        )
243        .fetch_all(&*self.pool)
244        .await?;
245        Ok(rows)
246    }
247
248    pub async fn cleanup_stale_entries(&self) -> DatabaseResult<u64> {
249        let result = sqlx::query!(
250            r#"
251            DELETE FROM services
252            WHERE instance_id = $1
253              AND (status IN ('error', 'crashed')
254                   OR (status = 'running' AND pid IS NULL))
255            "#,
256            self.instance_id.as_str()
257        )
258        .execute(&*self.write_pool)
259        .await?;
260        Ok(result.rows_affected())
261    }
262
263    pub async fn touch_heartbeat(&self) -> DatabaseResult<u64> {
264        let result = sqlx::query!(
265            r#"UPDATE services SET heartbeat_at = CURRENT_TIMESTAMP WHERE instance_id = $1"#,
266            self.instance_id.as_str()
267        )
268        .execute(&*self.write_pool)
269        .await?;
270        Ok(result.rows_affected())
271    }
272
273    pub async fn delete_dead_instances(&self, older_than_secs: i64) -> DatabaseResult<u64> {
274        let result = sqlx::query!(
275            r#"
276            DELETE FROM services
277            WHERE heartbeat_at < CURRENT_TIMESTAMP - make_interval(secs => $1::double precision)
278            "#,
279            older_than_secs as f64
280        )
281        .execute(&*self.write_pool)
282        .await?;
283        Ok(result.rows_affected())
284    }
285}
systemprompt_database/repository/service/repo.rs

systemprompt_database/repository/service/
repo.rs
