little_durable_objects/host_leases/
postgres.rs

1use anyhow::{Context, Result};
2use async_trait::async_trait;
3
4use super::{HostLease, HostLeaseRegistry, HostLeaseRequest, HostLeaseStatus, HostLeaseStore};
5use crate::{host::HostId, postgres::PostgresDatabase};
6
7const REGISTRY_NOW_MS: &str = "(extract(epoch FROM clock_timestamp()) * 1000)::bigint";
8
9pub struct PostgresHostLeaseStore {
10    database: PostgresDatabase,
11}
12
13impl PostgresHostLeaseStore {
14    pub async fn connect(url: &str) -> Result<Self> {
15        Ok(Self::from_database(PostgresDatabase::connect(url).await?))
16    }
17
18    pub(crate) fn from_database(database: PostgresDatabase) -> Self {
19        Self { database }
20    }
21}
22
23#[async_trait]
24impl HostLeaseRegistry for PostgresHostLeaseStore {
25    async fn register(&self, request: &HostLeaseRequest) -> Result<HostLease> {
26        request.validate_duration()?;
27        let duration_ms = i64::try_from(request.duration_ms)
28            .expect("validated host lease duration must fit PostgreSQL BIGINT");
29        let row = self
30            .database
31            .client()
32            .query_opt(
33                &format!(
34                    "INSERT INTO durable_object_host_leases (host_id, session_id, route, expires_at_ms) \
35                     VALUES ($1, $2, $3, {REGISTRY_NOW_MS} + $4) \
36                     ON CONFLICT (host_id) DO UPDATE \
37                     SET session_id = EXCLUDED.session_id, route = EXCLUDED.route, \
38                         expires_at_ms = EXCLUDED.expires_at_ms \
39                     WHERE durable_object_host_leases.session_id = EXCLUDED.session_id \
40                        OR durable_object_host_leases.expires_at_ms <= {REGISTRY_NOW_MS} \
41                     RETURNING expires_at_ms"
42                ),
43                &[
44                    &request.id.as_str(),
45                    &request.session_id,
46                    &request.route,
47                    &duration_ms,
48                ],
49            )
50            .await
51            .context("register PostgreSQL host lease")?
52            .context("host lease is held by a different active session")?;
53        Ok(HostLease {
54            id: request.id.clone(),
55            session_id: request.session_id.clone(),
56            route: request.route.clone(),
57            expires_at_ms: u64::try_from(row.get::<_, i64>(0))
58                .context("PostgreSQL host lease expiration is negative")?,
59        })
60    }
61
62    async fn unregister(&self, id: &HostId, session_id: &str) -> Result<()> {
63        self.database
64            .client()
65            .execute(
66                "DELETE FROM durable_object_host_leases WHERE host_id = $1 AND session_id = $2",
67                &[&id.as_str(), &session_id],
68            )
69            .await
70            .context("unregister PostgreSQL host lease")?;
71        Ok(())
72    }
73}
74
75#[async_trait]
76impl HostLeaseStore for PostgresHostLeaseStore {
77    async fn lease_status(&self, id: &HostId) -> Result<HostLeaseStatus> {
78        let row = self
79            .database
80            .client()
81            .query_one(
82                &format!(
83                    "SELECT {REGISTRY_NOW_MS} AS store_now_ms, lease.session_id, \
84                            lease.route, lease.expires_at_ms \
85                     FROM (SELECT 1) AS clock_row \
86                     LEFT JOIN durable_object_host_leases AS lease ON lease.host_id = $1"
87                ),
88                &[&id.as_str()],
89            )
90            .await
91            .context("read PostgreSQL host lease status")?;
92        let store_now_ms = u64::try_from(row.get::<_, i64>(0))
93            .context("PostgreSQL lease-store clock is before the Unix epoch")?;
94        let lease = row
95            .get::<_, Option<String>>(1)
96            .map(|session_id| {
97                Ok::<_, anyhow::Error>(HostLease {
98                    id: id.clone(),
99                    session_id,
100                    route: row
101                        .get::<_, Option<String>>(2)
102                        .context("PostgreSQL host lease row is missing its route")?,
103                    expires_at_ms: u64::try_from(
104                        row.get::<_, Option<i64>>(3)
105                            .context("PostgreSQL host lease row is missing its expiration")?,
106                    )
107                    .context("PostgreSQL host lease expiration is negative")?,
108                })
109            })
110            .transpose()?;
111        Ok(HostLeaseStatus {
112            lease,
113            store_now_ms,
114        })
115    }
116}
little_durable_objects/host_leases/postgres.rs

little_durable_objects/host_leases/
postgres.rs
