systemprompt_database/repository/service/
endpoint.rs1use systemprompt_identifiers::{InstanceId, ServiceName};
17use systemprompt_manifest::profile::NodeRole;
18use systemprompt_manifest::services::ServiceStatus;
19
20use super::liveness::INSTANCE_DEAD_AFTER_SECS;
21use super::repo::ServiceRepository;
22use crate::error::DatabaseResult;
23
24#[derive(Debug, Clone, PartialEq, Eq)]
26pub struct ProcessOwner {
27 pub instance_id: InstanceId,
28 pub advertise_url: String,
29}
30
31impl ServiceRepository {
32 pub async fn advertise_endpoint(&self, advertise_url: &str) -> DatabaseResult<()> {
33 sqlx::query!(
34 r#"
35 INSERT INTO instance_endpoints (instance_id, advertise_url, heartbeat_at)
36 VALUES ($1, $2, CURRENT_TIMESTAMP)
37 ON CONFLICT (instance_id) DO UPDATE SET
38 advertise_url = EXCLUDED.advertise_url,
39 heartbeat_at = EXCLUDED.heartbeat_at
40 "#,
41 self.instance_id.as_str(),
42 advertise_url
43 )
44 .execute(&*self.write_pool)
45 .await?;
46 Ok(())
47 }
48
49 pub async fn find_process_owner(
50 &self,
51 service_name: &ServiceName,
52 ) -> DatabaseResult<Option<ProcessOwner>> {
53 let row = sqlx::query!(
54 r#"
55 SELECT s.instance_id AS "instance_id: InstanceId", e.advertise_url
56 FROM services s
57 JOIN instance_endpoints e ON e.instance_id = s.instance_id
58 JOIN instance_presence p ON p.instance_id = s.instance_id
59 WHERE s.name = $1
60 AND s.status = $2
61 AND s.instance_id <> $3
62 AND s.heartbeat_at
63 >= CURRENT_TIMESTAMP - make_interval(secs => $4::double precision)
64 AND e.heartbeat_at
65 >= CURRENT_TIMESTAMP - make_interval(secs => $4::double precision)
66 AND (
67 p.role <> $5
68 OR EXISTS (
69 SELECT 1 FROM instance_leases l
70 WHERE l.lease_key = s.module_name || ':' || s.name
71 AND l.owner_instance_id = s.instance_id
72 AND l.expires_at > CURRENT_TIMESTAMP
73 )
74 )
75 ORDER BY s.heartbeat_at DESC
76 LIMIT 1
77 "#,
78 service_name.as_str(),
79 ServiceStatus::Running.as_str(),
80 self.instance_id.as_str(),
81 INSTANCE_DEAD_AFTER_SECS as f64,
82 NodeRole::Worker.as_str()
83 )
84 .fetch_optional(&*self.write_pool)
85 .await?;
86 Ok(row.map(|row| ProcessOwner {
87 instance_id: row.instance_id,
88 advertise_url: row.advertise_url,
89 }))
90 }
91
92 pub async fn delete_stale_endpoints(&self, older_than_secs: i64) -> DatabaseResult<u64> {
93 let result = sqlx::query!(
94 r#"
95 DELETE FROM instance_endpoints
96 WHERE heartbeat_at < CURRENT_TIMESTAMP - make_interval(secs => $1::double precision)
97 "#,
98 older_than_secs as f64
99 )
100 .execute(&*self.write_pool)
101 .await?;
102 Ok(result.rows_affected())
103 }
104}