cairn_journal_postgres/
lib.rs

1//! PostgreSQL-backed [`Journal`] for Cairn.
2//!
3//! Appends are atomic compare-and-set on the run's next sequence number, so
4//! any number of processes and machines can drive, resume and approve the
5//! same runs: exactly one writer wins each position, the rest get
6//! [`JournalError::Conflict`] and can reload and retry.
7
8use std::time::Duration;
9
10use async_trait::async_trait;
11use cairn_runtime::{Event, Journal, JournalError, LeaseHolder};
12use sqlx::postgres::{PgPool, PgPoolOptions};
13
14const SCHEMA: &str = r#"
15CREATE TABLE IF NOT EXISTS cairn_events (
16    run_id  TEXT        NOT NULL,
17    seq     BIGINT      NOT NULL,
18    at      TIMESTAMPTZ NOT NULL,
19    kind    TEXT        NOT NULL,
20    event   JSONB       NOT NULL,
21    PRIMARY KEY (run_id, seq)
22);
23CREATE INDEX IF NOT EXISTS cairn_events_kind_at ON cairn_events (kind, at);
24CREATE TABLE IF NOT EXISTS cairn_leases (
25    run_id     TEXT        PRIMARY KEY,
26    owner      TEXT        NOT NULL,
27    expires_at TIMESTAMPTZ NOT NULL
28);
29"#;
30
31/// Arbitrary constant identifying Cairn's schema-migration advisory lock.
32const SCHEMA_LOCK: i64 = 0x6361_6972_6e00_0001;
33
34#[derive(Clone)]
35pub struct PostgresJournal {
36    pool: PgPool,
37}
38
39fn backend(e: sqlx::Error) -> JournalError {
40    JournalError::Backend(e.to_string())
41}
42
43impl PostgresJournal {
44    /// Connect and create the schema if needed.
45    pub async fn connect(url: &str) -> Result<Self, JournalError> {
46        let pool = PgPoolOptions::new().max_connections(16).connect(url).await.map_err(backend)?;
47        Self::from_pool(pool).await
48    }
49
50    /// Use an existing pool and create the schema if needed.
51    pub async fn from_pool(pool: PgPool) -> Result<Self, JournalError> {
52        // `CREATE ... IF NOT EXISTS` races when several replicas start at
53        // once; serialize schema setup behind a transaction-scoped lock.
54        let mut tx = pool.begin().await.map_err(backend)?;
55        sqlx::query("SELECT pg_advisory_xact_lock($1)").bind(SCHEMA_LOCK).execute(&mut *tx).await.map_err(backend)?;
56        sqlx::raw_sql(SCHEMA).execute(&mut *tx).await.map_err(backend)?;
57        tx.commit().await.map_err(backend)?;
58        Ok(Self { pool })
59    }
60
61    pub fn pool(&self) -> &PgPool {
62        &self.pool
63    }
64
65    /// Run ids that are suspended waiting for a human decision.
66    pub async fn runs_awaiting_approval(&self) -> Result<Vec<String>, JournalError> {
67        // A run is awaiting approval if it has an approval request with no
68        // matching decision and has not reached a terminal event.
69        sqlx::query_scalar(
70            r#"
71            SELECT DISTINCT r.run_id FROM cairn_events r
72            WHERE r.kind = 'approval_requested'
73              AND NOT EXISTS (
74                SELECT 1 FROM cairn_events d
75                WHERE d.run_id = r.run_id AND d.kind = 'approval_decided'
76                  AND d.event->'kind'->>'call_id' = r.event->'kind'->>'call_id')
77              AND NOT EXISTS (
78                SELECT 1 FROM cairn_events t
79                WHERE t.run_id = r.run_id
80                  AND t.kind IN ('run_completed', 'run_failed', 'budget_exceeded'))
81            ORDER BY r.run_id
82            "#,
83        )
84        .fetch_all(&self.pool)
85        .await
86        .map_err(backend)
87    }
88}
89
90#[async_trait]
91impl Journal for PostgresJournal {
92    async fn append(&self, event: &Event) -> Result<(), JournalError> {
93        let json = serde_json::to_value(event).map_err(|e| JournalError::Corrupt(e.to_string()))?;
94        let seq = i64::try_from(event.seq).map_err(|_| JournalError::Corrupt("seq out of range".into()))?;
95
96        // Insert only if `seq` is exactly the next position. A concurrent
97        // writer for the same position loses on the primary key.
98        let result = sqlx::query(
99            r#"
100            INSERT INTO cairn_events (run_id, seq, at, kind, event)
101            SELECT $1, $2, $3, $4, $5
102            WHERE (SELECT COALESCE(MAX(seq) + 1, 0) FROM cairn_events WHERE run_id = $1) = $2
103            "#,
104        )
105        .bind(&event.run_id)
106        .bind(seq)
107        .bind(event.at)
108        .bind(event.kind.name())
109        .bind(json)
110        .execute(&self.pool)
111        .await;
112
113        let inserted = match result {
114            Ok(r) => r.rows_affected() == 1,
115            Err(sqlx::Error::Database(e)) if e.is_unique_violation() => false,
116            Err(e) => return Err(backend(e)),
117        };
118        if inserted {
119            return Ok(());
120        }
121        let expected: i64 = sqlx::query_scalar("SELECT COALESCE(MAX(seq) + 1, 0) FROM cairn_events WHERE run_id = $1")
122            .bind(&event.run_id)
123            .fetch_one(&self.pool)
124            .await
125            .map_err(backend)?;
126        Err(JournalError::Conflict { run_id: event.run_id.clone(), expected: expected as u64, got: event.seq })
127    }
128
129    /// Leases use the database clock, so worker clock skew does not matter.
130    async fn acquire_lease(
131        &self,
132        run_id: &str,
133        owner: &str,
134        ttl: Duration,
135    ) -> Result<Result<(), LeaseHolder>, JournalError> {
136        let taken: Option<String> = sqlx::query_scalar(
137            r#"
138            INSERT INTO cairn_leases (run_id, owner, expires_at)
139            VALUES ($1, $2, now() + make_interval(secs => $3))
140            ON CONFLICT (run_id) DO UPDATE SET owner = EXCLUDED.owner, expires_at = EXCLUDED.expires_at
141            WHERE cairn_leases.expires_at < now() OR cairn_leases.owner = EXCLUDED.owner
142            RETURNING owner
143            "#,
144        )
145        .bind(run_id)
146        .bind(owner)
147        .bind(ttl.as_secs_f64())
148        .fetch_optional(&self.pool)
149        .await
150        .map_err(backend)?;
151        if taken.is_some() {
152            return Ok(Ok(()));
153        }
154        let holder: Option<(String, chrono::DateTime<chrono::Utc>)> =
155            sqlx::query_as("SELECT owner, expires_at FROM cairn_leases WHERE run_id = $1")
156                .bind(run_id)
157                .fetch_optional(&self.pool)
158                .await
159                .map_err(backend)?;
160        match holder {
161            Some((owner, expires_at)) => Ok(Err(LeaseHolder { owner, expires_at })),
162            // Released between the two statements: try again.
163            None => Box::pin(self.acquire_lease(run_id, owner, ttl)).await,
164        }
165    }
166
167    async fn renew_lease(&self, run_id: &str, owner: &str, ttl: Duration) -> Result<bool, JournalError> {
168        let updated = sqlx::query(
169            "UPDATE cairn_leases SET expires_at = now() + make_interval(secs => $3) WHERE run_id = $1 AND owner = $2",
170        )
171        .bind(run_id)
172        .bind(owner)
173        .bind(ttl.as_secs_f64())
174        .execute(&self.pool)
175        .await
176        .map_err(backend)?;
177        Ok(updated.rows_affected() == 1)
178    }
179
180    async fn release_lease(&self, run_id: &str, owner: &str) -> Result<(), JournalError> {
181        sqlx::query("DELETE FROM cairn_leases WHERE run_id = $1 AND owner = $2")
182            .bind(run_id)
183            .bind(owner)
184            .execute(&self.pool)
185            .await
186            .map_err(backend)?;
187        Ok(())
188    }
189
190    async fn load(&self, run_id: &str) -> Result<Vec<Event>, JournalError> {
191        let rows: Vec<serde_json::Value> =
192            sqlx::query_scalar("SELECT event FROM cairn_events WHERE run_id = $1 ORDER BY seq")
193                .bind(run_id)
194                .fetch_all(&self.pool)
195                .await
196                .map_err(backend)?;
197        rows.into_iter().map(|v| serde_json::from_value(v).map_err(|e| JournalError::Corrupt(e.to_string()))).collect()
198    }
199}
cairn_journal_postgres/lib.rs

cairn_journal_postgres/
lib.rs
