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            .query_opt(
32                &format!(
33                    "INSERT INTO durable_object_host_leases (host_id, session_id, route, expires_at_ms) \
34                     VALUES ($1, $2, $3, {REGISTRY_NOW_MS} + $4) \
35                     ON CONFLICT (host_id) DO UPDATE \
36                     SET session_id = EXCLUDED.session_id, route = EXCLUDED.route, \
37                         expires_at_ms = EXCLUDED.expires_at_ms \
38                     WHERE durable_object_host_leases.session_id = EXCLUDED.session_id \
39                        OR durable_object_host_leases.expires_at_ms <= {REGISTRY_NOW_MS} \
40                     RETURNING expires_at_ms"
41                ),
42                &[
43                    &request.id.as_str(),
44                    &request.session_id,
45                    &request.route,
46                    &duration_ms,
47                ],
48            )
49            .await
50            .context("register PostgreSQL host lease")?
51            .context("host lease is held by a different active session")?;
52        Ok(HostLease {
53            id: request.id.clone(),
54            session_id: request.session_id.clone(),
55            route: request.route.clone(),
56            expires_at_ms: u64::try_from(row.get::<_, i64>(0))
57                .context("PostgreSQL host lease expiration is negative")?,
58        })
59    }
60
61    async fn unregister(&self, id: &HostId, session_id: &str) -> Result<()> {
62        self.database
63            .execute(
64                "DELETE FROM durable_object_host_leases WHERE host_id = $1 AND session_id = $2",
65                &[&id.as_str(), &session_id],
66            )
67            .await
68            .context("unregister PostgreSQL host lease")?;
69        Ok(())
70    }
71}
72
73#[async_trait]
74impl HostLeaseStore for PostgresHostLeaseStore {
75    async fn lease_status(&self, id: &HostId) -> Result<HostLeaseStatus> {
76        let row = self
77            .database
78            .query_one(
79                &format!(
80                    "SELECT {REGISTRY_NOW_MS} AS store_now_ms, lease.session_id, \
81                            lease.route, lease.expires_at_ms \
82                     FROM (SELECT 1) AS clock_row \
83                     LEFT JOIN durable_object_host_leases AS lease ON lease.host_id = $1"
84                ),
85                &[&id.as_str()],
86            )
87            .await
88            .context("read PostgreSQL host lease status")?;
89        let store_now_ms = u64::try_from(row.get::<_, i64>(0))
90            .context("PostgreSQL lease-store clock is before the Unix epoch")?;
91        let lease = row
92            .get::<_, Option<String>>(1)
93            .map(|session_id| {
94                Ok::<_, anyhow::Error>(HostLease {
95                    id: id.clone(),
96                    session_id,
97                    route: row
98                        .get::<_, Option<String>>(2)
99                        .context("PostgreSQL host lease row is missing its route")?,
100                    expires_at_ms: u64::try_from(
101                        row.get::<_, Option<i64>>(3)
102                            .context("PostgreSQL host lease row is missing its expiration")?,
103                    )
104                    .context("PostgreSQL host lease expiration is negative")?,
105                })
106            })
107            .transpose()?;
108        Ok(HostLeaseStatus {
109            lease,
110            store_now_ms,
111        })
112    }
113}
little_durable_objects/host_leases/postgres.rs

little_durable_objects/host_leases/
postgres.rs
