Skip to main content

orbit_runtime/store/
postgres.rs

1//! PostgreSQL backend.
2//!
3//! The `postgres` crate drives its own async runtime internally and must not be
4//! called from inside another one, while the engine runs synchronously inside
5//! the HTTP server's tokio runtime. So every statement runs on one dedicated
6//! connection thread, fed through a channel. Calls stay synchronous for the
7//! engine, the connection is reconnected automatically if it drops, and each
8//! `append` is a single transaction, exactly like the SQLite backend.
9
10use super::{ExecutionSummary, PutDefinition, Result, StoreError, WorkerRecord, SCHEMA};
11use crate::model::{Event, ExecutionState, HistoryEvent};
12use orbit_ir::WorkflowIr;
13use postgres::Client;
14use std::sync::mpsc;
15use tokio_postgres_rustls::MakeRustlsConnect;
16
17type Job = Box<dyn FnOnce(&mut Client) + Send>;
18
19pub struct PgStore {
20    jobs: mpsc::Sender<Job>,
21    mode: PgMode,
22}
23
24/// How a [`PgStore`] may write. Single-node servers use `ReadWrite`; the
25/// cluster mode (`orbit server --cluster`, see `crate::cluster`) gives a
26/// standby a `ReadOnly` store and the leader a `Fenced` one.
27#[derive(Debug, Clone, Copy, PartialEq, Eq)]
28pub enum PgMode {
29    ReadWrite,
30    /// Reads only: writes fail and no execution counts as active, so an
31    /// engine opened on it recovers nothing and never appends.
32    ReadOnly,
33    /// Every history append verifies, in its own transaction, that the
34    /// cluster term is still this one, so a deposed leader cannot write.
35    Fenced(i64),
36    /// Shared mode (`--cluster-mode shared`, see `crate::shared`): every node
37    /// writes, but only to the executions it owns. Appends go through
38    /// [`PgStore::append_shared`], which checks the execution's ownership
39    /// epoch; no execution counts as active at open, because a node loads
40    /// the executions it owns or takes over explicitly.
41    Shared,
42}
43
44/// Table holding the cluster's leadership term (created by `crate::cluster`).
45pub(crate) const CLUSTER_SCHEMA: &str = r#"
46CREATE TABLE IF NOT EXISTS orbit_cluster (
47    id INT PRIMARY KEY,
48    term BIGINT NOT NULL,
49    leader_node TEXT,
50    leader_addr TEXT,
51    since_ms BIGINT NOT NULL
52);
53INSERT INTO orbit_cluster (id, term, since_ms) VALUES (1, 0, 0) ON CONFLICT DO NOTHING;
54"#;
55
56/// A plain session for the cluster's leader-election connection: same TLS
57/// rules as the store, plus timeouts so that a dead network path is noticed
58/// within seconds rather than minutes.
59pub(crate) fn session_client(url: &str, application_name: &str) -> Result<Client> {
60    use std::str::FromStr;
61    use std::time::Duration;
62    let mut config = postgres::Config::from_str(url)?;
63    config
64        .application_name(application_name)
65        .connect_timeout(Duration::from_secs(5))
66        .keepalives(true)
67        .keepalives_idle(Duration::from_secs(5))
68        .tcp_user_timeout(Duration::from_secs(5));
69    Ok(config.connect(tls()?)?)
70}
71
72fn read_only() -> StoreError {
73    StoreError("this server is a standby: the store is read-only".into())
74}
75
76/// TLS for PostgreSQL connections, following the URL's `sslmode`: `prefer`
77/// (the default) uses TLS when the server offers it, `require` insists on it,
78/// `disable` never uses it. Certificates are verified against the public web
79/// PKI roots plus, if set, the PEM CA bundle in `ORBIT_DB_CA_CERT`.
80fn tls() -> Result<MakeRustlsConnect> {
81    let mut roots = rustls::RootCertStore::empty();
82    roots.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());
83    if let Ok(path) = std::env::var("ORBIT_DB_CA_CERT") {
84        let pem = std::fs::read(&path)
85            .map_err(|e| StoreError(format!("cannot read ORBIT_DB_CA_CERT {path}: {e}")))?;
86        for cert in rustls_pemfile::certs(&mut pem.as_slice()) {
87            let cert =
88                cert.map_err(|e| StoreError(format!("invalid certificate in {path}: {e}")))?;
89            roots
90                .add(cert)
91                .map_err(|e| StoreError(format!("invalid CA certificate in {path}: {e}")))?;
92        }
93    }
94    let provider = std::sync::Arc::new(rustls::crypto::aws_lc_rs::default_provider());
95    let config = rustls::ClientConfig::builder_with_provider(provider)
96        .with_safe_default_protocol_versions()
97        .map_err(|e| StoreError(format!("TLS configuration: {e}")))?
98        .with_root_certificates(roots)
99        .with_no_client_auth();
100    Ok(MakeRustlsConnect::new(config))
101}
102
103fn connect_client(url: &str) -> Result<Client> {
104    let mut client = Client::connect(url, tls()?)?;
105    client.batch_execute(SCHEMA)?;
106    let rows = client.query_one("SELECT COUNT(*) FROM schema_version", &[])?;
107    if rows.get::<_, i64>(0) == 0 {
108        client.execute("INSERT INTO schema_version (version) VALUES (1)", &[])?;
109    }
110    Ok(client)
111}
112
113type EventRow = (i64, &'static str, Option<String>, String, i64);
114
115/// Bind parameters for `append`'s statement: the execution id, the projection
116/// (when this statement writes it), then five per event.
117fn bind<'a>(
118    id: &'a String,
119    projection: Option<&[&'a (dyn postgres::types::ToSql + Sync); 7]>,
120    chunk: &'a [EventRow],
121) -> Vec<&'a (dyn postgres::types::ToSql + Sync)> {
122    let mut v: Vec<&'a (dyn postgres::types::ToSql + Sync)> =
123        Vec::with_capacity(8 + chunk.len() * 5);
124    v.push(id);
125    if let Some(p) = projection {
126        v.extend(p.iter().copied());
127    }
128    for (seq, ty, step, payload, ts) in chunk {
129        v.extend([
130            seq as &(dyn postgres::types::ToSql + Sync),
131            ty,
132            step,
133            payload,
134            ts,
135        ]);
136    }
137    v
138}
139
140/// Like [`connect_client`], with the schema created under an advisory lock:
141/// shared-mode servers usually start together, and concurrent
142/// `CREATE TABLE IF NOT EXISTS` statements can fail.
143fn connect_client_serialized(url: &str) -> Result<Client> {
144    let mut client = Client::connect(url, tls()?)?;
145    let mut tx = client.transaction()?;
146    tx.execute("SELECT pg_advisory_xact_lock($1)", &[&BASE_SCHEMA_LOCK_KEY])?;
147    tx.batch_execute(SCHEMA)?;
148    if tx
149        .query_one("SELECT COUNT(*) FROM schema_version", &[])?
150        .get::<_, i64>(0)
151        == 0
152    {
153        tx.execute("INSERT INTO schema_version (version) VALUES (1)", &[])?;
154    }
155    tx.commit()?;
156    Ok(client)
157}
158
159/// Advisory lock key serialising creation of the base schema (shared mode).
160const BASE_SCHEMA_LOCK_KEY: i64 = 0x6f72_6269_7473_6362;
161
162fn bad_payload(e: serde_json::Error) -> StoreError {
163    StoreError(format!("postgres: unreadable event payload: {e}"))
164}
165
166impl PgStore {
167    pub fn connect(url: &str) -> Result<PgStore> {
168        PgStore::connect_with(url, PgMode::ReadWrite)
169    }
170
171    /// Connect with a write mode (see [`PgMode`]).
172    pub fn connect_with(url: &str, mode: PgMode) -> Result<PgStore> {
173        let url = url.to_string();
174        let (jobs, rx) = mpsc::channel::<Job>();
175        let (ready_tx, ready_rx) = mpsc::channel::<Result<()>>();
176        std::thread::Builder::new()
177            .name("orbit-postgres".into())
178            .spawn(move || {
179                let open = |url: &str| match mode {
180                    PgMode::Shared => connect_client_serialized(url),
181                    _ => connect_client(url),
182                };
183                let mut client = match open(&url) {
184                    Ok(c) => {
185                        let _ = ready_tx.send(Ok(()));
186                        c
187                    }
188                    Err(e) => {
189                        let _ = ready_tx.send(Err(e));
190                        return;
191                    }
192                };
193                for job in rx {
194                    if client.is_closed() {
195                        match open(&url) {
196                            Ok(c) => {
197                                tracing::warn!("reconnected to PostgreSQL");
198                                client = c;
199                            }
200                            Err(e) => tracing::error!(error = %e, "PostgreSQL reconnect failed"),
201                        }
202                    }
203                    job(&mut client);
204                }
205            })
206            .map_err(|e| StoreError(format!("cannot start PostgreSQL thread: {e}")))?;
207        ready_rx
208            .recv()
209            .map_err(|_| StoreError("PostgreSQL thread exited during startup".into()))??;
210        Ok(PgStore { jobs, mode })
211    }
212
213    pub(crate) fn run<T: Send + 'static>(
214        &self,
215        f: impl FnOnce(&mut Client) -> Result<T> + Send + 'static,
216    ) -> Result<T> {
217        let (tx, rx) = mpsc::channel();
218        self.jobs
219            .send(Box::new(move |c| {
220                let _ = tx.send(f(c));
221            }))
222            .map_err(|_| StoreError("PostgreSQL connection thread has stopped".into()))?;
223        rx.recv()
224            .map_err(|_| StoreError("PostgreSQL connection thread has stopped".into()))?
225    }
226
227    pub fn mode(&self) -> PgMode {
228        self.mode
229    }
230
231    pub fn put_definition(&self, ir: &WorkflowIr, now: u64) -> Result<PutDefinition> {
232        if self.mode == PgMode::ReadOnly {
233            return Err(read_only());
234        }
235        let (name, version, hash) = (
236            ir.name.clone(),
237            ir.version.clone(),
238            ir.definition_hash.clone(),
239        );
240        let text = serde_json::to_string(ir).expect("IR serialises");
241        self.run(move |c| {
242            let existing = c.query_opt(
243                "SELECT definition_hash FROM workflow_definitions WHERE name = $1 AND version = $2",
244                &[&name, &version],
245            )?;
246            match existing.map(|r| r.get::<_, String>(0)) {
247                Some(h) if h == hash => Ok(PutDefinition::Unchanged),
248                Some(h) => Ok(PutDefinition::Conflict { existing_hash: h }),
249                None => {
250                    // ON CONFLICT DO NOTHING closes the race with a concurrent registration.
251                    let n = c.execute(
252                        "INSERT INTO workflow_definitions (name, version, definition_hash, ir, created_at_ms)
253                         VALUES ($1, $2, $3, $4, $5) ON CONFLICT DO NOTHING",
254                        &[&name, &version, &hash, &text, &(now as i64)],
255                    )?;
256                    Ok(if n == 1 { PutDefinition::Created } else { PutDefinition::Unchanged })
257                }
258            }
259        })
260    }
261
262    pub fn definitions(&self) -> Result<Vec<WorkflowIr>> {
263        let rows = self.run(|c| {
264            Ok(c.query(
265                "SELECT ir FROM workflow_definitions ORDER BY name, created_at_ms",
266                &[],
267            )?
268            .into_iter()
269            .map(|r| r.get::<_, String>(0))
270            .collect::<Vec<_>>())
271        })?;
272        Ok(rows
273            .into_iter()
274            .filter_map(|t| match serde_json::from_str::<WorkflowIr>(&t) {
275                Ok(ir) => Some(ir),
276                Err(e) => {
277                    tracing::error!(error = %e, "skipping unreadable workflow definition");
278                    None
279                }
280            })
281            .collect())
282    }
283
284    pub fn execution_exists(&self, id: &str) -> Result<bool> {
285        let id = id.to_string();
286        self.run(move |c| {
287            Ok(c.query_opt(
288                "SELECT 1 FROM workflow_executions WHERE execution_id = $1",
289                &[&id],
290            )?
291            .is_some())
292        })
293    }
294
295    pub fn append(&mut self, state: &ExecutionState, events: &[HistoryEvent]) -> Result<()> {
296        let fence = match self.mode {
297            PgMode::ReadWrite => None,
298            PgMode::ReadOnly => return Err(read_only()),
299            PgMode::Fenced(term) => Some(term),
300            PgMode::Shared => {
301                return Err(StoreError(
302                    "shared mode: history appends must carry the execution's ownership epoch"
303                        .into(),
304                ))
305            }
306        };
307        let rows: Vec<EventRow> = events
308            .iter()
309            .map(|h| {
310                (
311                    h.seq as i64,
312                    h.event.name(),
313                    h.event.step().map(str::to_string),
314                    serde_json::to_string(&h.event).expect("event serialises"),
315                    h.timestamp_ms as i64,
316                )
317            })
318            .collect();
319        let id = state.id.clone();
320        let projection = (
321            state.workflow.clone(),
322            state.version.clone(),
323            state.definition_hash.clone(),
324            state.status.as_str(),
325            state.created_at_ms as i64,
326            state.updated_at_ms as i64,
327            state.last_seq as i64,
328        );
329        self.run(move |c| {
330            type P<'a> = &'a (dyn postgres::types::ToSql + Sync);
331            let (workflow, version, hash, status, created, updated, last_seq) = &projection;
332            // Events (in chunks that stay well under PostgreSQL's 65535 bind
333            // parameters) and the execution projection, written by one
334            // statement: a data-modifying CTE. Unfenced, that statement is the
335            // whole transaction, so a typical commit is one round trip.
336            let statement = |chunk: &[EventRow],
337                             with_projection: bool|
338             -> String {
339                let base = if with_projection { 8 } else { 1 };
340                let mut sql = String::from(
341                    "INSERT INTO execution_events (execution_id, seq, event_type, step, payload, timestamp_ms) VALUES ",
342                );
343                for i in 0..chunk.len() {
344                    let b = base + 1 + i * 5;
345                    if i > 0 {
346                        sql.push(',');
347                    }
348                    sql.push_str(&format!(
349                        "($1, ${}, ${}, ${}, ${}, ${})",
350                        b,
351                        b + 1,
352                        b + 2,
353                        b + 3,
354                        b + 4
355                    ));
356                }
357                if !with_projection {
358                    return sql;
359                }
360                format!(
361                    "WITH ev AS ({sql} RETURNING 1) \
362                     INSERT INTO workflow_executions (execution_id, workflow, version, definition_hash, status, created_at_ms, updated_at_ms, last_seq) \
363                     VALUES ($1, $2, $3, $4, $5, $6, $7, $8) \
364                     ON CONFLICT (execution_id) DO UPDATE SET status = EXCLUDED.status, updated_at_ms = EXCLUDED.updated_at_ms, last_seq = EXCLUDED.last_seq"
365                )
366            };
367            let proj: [P<'_>; 7] = [workflow, version, hash, status, created, updated, last_seq];
368            const CHUNK: usize = 1000;
369            if fence.is_none() && rows.len() <= CHUNK {
370                c.execute(&statement(&rows, true), &bind(&id, Some(&proj), &rows))?;
371                return Ok(());
372            }
373            let mut tx = c.transaction()?;
374            if let Some(term) = fence {
375                // FOR SHARE: a new leader's term increment waits for this
376                // transaction, and this read sees any increment committed first.
377                let current: Option<i64> = tx
378                    .query_opt("SELECT term FROM orbit_cluster WHERE id = 1 FOR SHARE", &[])?
379                    .map(|r| r.get(0));
380                if current != Some(term) {
381                    return Err(StoreError(format!(
382                        "fenced: this server led term {term}, but the cluster is at term {}; another server is the leader",
383                        current.map_or("unknown".to_string(), |t| t.to_string())
384                    )));
385                }
386            }
387            let chunks: Vec<_> = rows.chunks(CHUNK).collect();
388            for (i, chunk) in chunks.iter().enumerate() {
389                let last = i + 1 == chunks.len();
390                tx.execute(
391                    &statement(chunk, last),
392                    &bind(&id, last.then_some(&proj), chunk),
393                )?;
394            }
395            tx.commit()?;
396            Ok(())
397        })
398    }
399
400    pub fn history(&self, id: &str) -> Result<Vec<HistoryEvent>> {
401        let id = id.to_string();
402        let rows = self.run(move |c| {
403            Ok(c.query(
404                "SELECT seq, timestamp_ms, payload FROM execution_events WHERE execution_id = $1 ORDER BY seq",
405                &[&id],
406            )?
407            .into_iter()
408            .map(|r| (r.get::<_, i64>(0), r.get::<_, i64>(1), r.get::<_, String>(2)))
409            .collect::<Vec<_>>())
410        })?;
411        rows.into_iter()
412            .map(|(seq, ts, payload)| {
413                let event: Event = serde_json::from_str(&payload).map_err(bad_payload)?;
414                Ok(HistoryEvent {
415                    seq: seq as u64,
416                    timestamp_ms: ts as u64,
417                    event,
418                })
419            })
420            .collect()
421    }
422
423    pub fn active_execution_ids(&self) -> Result<Vec<String>> {
424        if matches!(self.mode, PgMode::ReadOnly | PgMode::Shared) {
425            return Ok(vec![]);
426        }
427        self.run(|c| {
428            Ok(c.query(
429                "SELECT execution_id FROM workflow_executions WHERE status IN ('RUNNING', 'COMPENSATING') ORDER BY created_at_ms",
430                &[],
431            )?
432            .into_iter()
433            .map(|r| r.get(0))
434            .collect())
435        })
436    }
437
438    pub fn list_executions(
439        &self,
440        status: Option<&str>,
441        workflow: Option<&str>,
442        limit: u32,
443    ) -> Result<Vec<ExecutionSummary>> {
444        let (status, workflow) = (status.map(str::to_string), workflow.map(str::to_string));
445        self.run(move |c| {
446            Ok(c.query(
447                "SELECT execution_id, workflow, version, status, created_at_ms, updated_at_ms FROM workflow_executions
448                 WHERE ($1::TEXT IS NULL OR status = $1) AND ($2::TEXT IS NULL OR workflow = $2)
449                 ORDER BY created_at_ms DESC LIMIT $3",
450                &[&status, &workflow, &(limit as i64)],
451            )?
452            .into_iter()
453            .map(|r| ExecutionSummary {
454                execution_id: r.get(0),
455                workflow: r.get(1),
456                version: r.get(2),
457                status: r.get(3),
458                created_at_ms: r.get::<_, i64>(4) as u64,
459                updated_at_ms: r.get::<_, i64>(5) as u64,
460            })
461            .collect())
462        })
463    }
464
465    pub fn count_by_status(&self) -> Result<Vec<(String, u64)>> {
466        self.run(|c| {
467            Ok(c.query(
468                "SELECT status, COUNT(*) FROM workflow_executions GROUP BY status",
469                &[],
470            )?
471            .into_iter()
472            .map(|r| (r.get(0), r.get::<_, i64>(1) as u64))
473            .collect())
474        })
475    }
476
477    pub fn upsert_worker(&self, w: &WorkerRecord) -> Result<()> {
478        if self.mode == PgMode::ReadOnly {
479            return Err(read_only());
480        }
481        let w = w.clone();
482        self.run(move |c| {
483            c.execute(
484                "INSERT INTO worker_registry (worker_id, runtime, sdk, contracts, registered_at_ms, last_seen_ms) VALUES ($1, $2, $3, $4, $5, $6)
485                 ON CONFLICT (worker_id) DO UPDATE SET runtime = EXCLUDED.runtime, sdk = EXCLUDED.sdk, contracts = EXCLUDED.contracts,
486                 registered_at_ms = EXCLUDED.registered_at_ms, last_seen_ms = EXCLUDED.last_seen_ms",
487                &[&w.worker_id, &w.runtime, &w.sdk, &w.contracts.to_string(), &(w.registered_at_ms as i64), &(w.last_seen_ms as i64)],
488            )?;
489            Ok(())
490        })
491    }
492
493    pub fn touch_worker(&self, worker_id: &str, now: u64) -> Result<()> {
494        if self.mode == PgMode::ReadOnly {
495            return Err(read_only());
496        }
497        let id = worker_id.to_string();
498        self.run(move |c| {
499            c.execute(
500                "UPDATE worker_registry SET last_seen_ms = $2 WHERE worker_id = $1",
501                &[&id, &(now as i64)],
502            )?;
503            Ok(())
504        })
505    }
506
507    pub fn workers(&self) -> Result<Vec<WorkerRecord>> {
508        self.run(|c| {
509            Ok(c.query(
510                "SELECT worker_id, runtime, sdk, contracts, registered_at_ms, last_seen_ms FROM worker_registry ORDER BY worker_id",
511                &[],
512            )?
513            .into_iter()
514            .map(|r| WorkerRecord {
515                worker_id: r.get(0),
516                runtime: r.get(1),
517                sdk: r.get(2),
518                contracts: serde_json::from_str(&r.get::<_, String>(3)).unwrap_or_default(),
519                registered_at_ms: r.get::<_, i64>(4) as u64,
520                last_seen_ms: r.get::<_, i64>(5) as u64,
521            })
522            .collect())
523        })
524    }
525}
526
527// ------------------------------------------------------------- shared mode
528
529/// Database time in milliseconds. Shared mode measures node heartbeats and
530/// worker session liveness with it, so clock differences between servers do
531/// not decide who is alive.
532pub(crate) const NOW_MS: &str = "(extract(epoch from clock_timestamp()) * 1000)::bigint";
533
534/// Tables of shared mode (created by `crate::shared`). See
535/// `docs/src/runtime/storage.md` for what each column means.
536pub(crate) const SHARED_SCHEMA: &str = r#"
537CREATE TABLE IF NOT EXISTS orbit_nodes (
538    node_id TEXT PRIMARY KEY,
539    incarnation TEXT NOT NULL,
540    addr TEXT,
541    heartbeat_ms BIGINT NOT NULL,
542    started_ms BIGINT NOT NULL
543);
544CREATE TABLE IF NOT EXISTS orbit_owners (
545    execution_id TEXT PRIMARY KEY,
546    root_id TEXT NOT NULL,
547    node_id TEXT NOT NULL,
548    epoch BIGINT NOT NULL,
549    since_ms BIGINT NOT NULL
550);
551CREATE INDEX IF NOT EXISTS idx_orbit_owners_node ON orbit_owners(node_id);
552CREATE INDEX IF NOT EXISTS idx_orbit_owners_root ON orbit_owners(root_id);
553CREATE TABLE IF NOT EXISTS orbit_tasks (
554    id BIGSERIAL PRIMARY KEY,
555    execution_id TEXT NOT NULL,
556    unit_key TEXT NOT NULL,
557    attempt BIGINT NOT NULL,
558    offer BIGINT NOT NULL,
559    task_name TEXT NOT NULL,
560    contract_version TEXT NOT NULL,
561    language TEXT,
562    payload TEXT NOT NULL,
563    enqueued_ms BIGINT NOT NULL,
564    state TEXT NOT NULL,
565    token TEXT,
566    claimed_by TEXT,
567    claimed_inc TEXT,
568    claimed_session TEXT,
569    claimed_worker TEXT,
570    claimed_ms BIGINT,
571    updated_ms BIGINT NOT NULL,
572    UNIQUE (execution_id, unit_key, attempt, offer)
573);
574CREATE INDEX IF NOT EXISTS idx_orbit_tasks_queued ON orbit_tasks(task_name, enqueued_ms) WHERE state = 'queued';
575CREATE INDEX IF NOT EXISTS idx_orbit_tasks_token ON orbit_tasks(token);
576CREATE INDEX IF NOT EXISTS idx_orbit_tasks_session ON orbit_tasks(claimed_session);
577CREATE TABLE IF NOT EXISTS orbit_inbox (
578    seq BIGSERIAL PRIMARY KEY,
579    execution_id TEXT NOT NULL,
580    kind TEXT NOT NULL,
581    token TEXT,
582    payload TEXT NOT NULL,
583    created_ms BIGINT NOT NULL,
584    done_ms BIGINT,
585    result TEXT
586);
587CREATE INDEX IF NOT EXISTS idx_orbit_inbox_execution ON orbit_inbox(execution_id, seq);
588CREATE TABLE IF NOT EXISTS orbit_sessions (
589    session_id TEXT PRIMARY KEY,
590    worker_id TEXT NOT NULL,
591    runtime TEXT NOT NULL,
592    contracts TEXT NOT NULL,
593    node_id TEXT NOT NULL,
594    registered_ms BIGINT NOT NULL,
595    last_seen_ms BIGINT NOT NULL
596);
597CREATE INDEX IF NOT EXISTS idx_orbit_sessions_worker ON orbit_sessions(worker_id);
598"#;
599
600/// A worker task to put on the shared queue (one row per offer of a unit).
601#[derive(Debug, Clone)]
602pub(crate) struct TaskOffer {
603    pub unit_key: String,
604    pub attempt: i64,
605    /// The unit's dispatch count when it was scheduled: together with the
606    /// attempt it identifies this offer, so a claim of an older offer of the
607    /// same unit is recognised as stale.
608    pub offer: i64,
609    pub task_name: String,
610    pub contract_version: String,
611    pub language: Option<String>,
612    /// The task assignment without token and deadline (JSON).
613    pub payload: String,
614    pub enqueued_ms: i64,
615}
616
617/// Queue rows of a unit to withdraw: every queued or claimed row of the unit
618/// except the offer in `keep` (attempt, offer) becomes `cancelled`, which
619/// tells the worker holding it to stop at its next heartbeat.
620#[derive(Debug, Clone)]
621pub(crate) struct UnitSync {
622    pub unit_key: String,
623    pub keep: Option<(i64, i64)>,
624}
625
626/// Everything shared mode writes in the same transaction as a history append.
627#[derive(Debug, Clone, Default)]
628pub(crate) struct SharedAppend {
629    pub node_id: String,
630    /// The ownership epoch this node holds for the execution; `None` for a
631    /// new execution, whose ownership row is created at epoch 1.
632    pub epoch: Option<i64>,
633    pub parent: Option<String>,
634    pub offers: Vec<TaskOffer>,
635    pub syncs: Vec<UnitSync>,
636    /// Tokens whose result arrived: their queue rows become `done`.
637    pub done_tokens: Vec<String>,
638    /// Inbox entries applied by this append: (seq, reply). Entries without a
639    /// reply are deleted; the others are marked done with the reply.
640    pub consumed: Vec<(i64, Option<String>)>,
641    /// Set when the execution reaches a terminal status: the reply recorded
642    /// for commands still waiting in its inbox.
643    pub finished_reply: Option<String>,
644}
645
646impl PgStore {
647    /// Shared-mode append: checks that this node still owns the execution at
648    /// its epoch (`SELECT ... FOR SHARE`, the same pattern as the cluster
649    /// term), then writes the events, the projection and the queue and inbox
650    /// changes in one transaction.
651    pub(crate) fn append_shared(
652        &mut self,
653        state: &ExecutionState,
654        events: &[HistoryEvent],
655        ext: SharedAppend,
656    ) -> Result<()> {
657        if self.mode != PgMode::Shared {
658            return Err(StoreError("append_shared needs a shared-mode store".into()));
659        }
660        let rows: Vec<(i64, &'static str, Option<String>, String, i64)> = events
661            .iter()
662            .map(|h| {
663                (
664                    h.seq as i64,
665                    h.event.name(),
666                    h.event.step().map(str::to_string),
667                    serde_json::to_string(&h.event).expect("event serialises"),
668                    h.timestamp_ms as i64,
669                )
670            })
671            .collect();
672        let id = state.id.clone();
673        let projection = (
674            state.workflow.clone(),
675            state.version.clone(),
676            state.definition_hash.clone(),
677            state.status.as_str(),
678            state.created_at_ms as i64,
679            state.updated_at_ms as i64,
680            state.last_seq as i64,
681        );
682        let now = state.updated_at_ms as i64;
683        self.run(move |c| {
684            let mut tx = c.transaction()?;
685            let node = &ext.node_id;
686            match ext.epoch {
687                Some(epoch) => {
688                    // FOR SHARE: a takeover's epoch increment waits for this
689                    // transaction, and this read sees any increment committed first.
690                    let owner: Option<(String, i64)> = tx
691                        .query_opt(
692                            "SELECT node_id, epoch FROM orbit_owners WHERE execution_id = $1 FOR SHARE",
693                            &[&id],
694                        )?
695                        .map(|r| (r.get(0), r.get(1)));
696                    if owner.as_ref() != Some(&(node.clone(), epoch)) {
697                        return Err(StoreError(match owner {
698                            Some((n, e)) => format!(
699                                "fenced: execution `{id}` is owned by node `{n}` at epoch {e}; this node held epoch {epoch}"
700                            ),
701                            None => format!(
702                                "fenced: execution `{id}` has no owner any more; this node held epoch {epoch}"
703                            ),
704                        }));
705                    }
706                }
707                None => {
708                    let created = tx.execute(
709                        &format!(
710                            "INSERT INTO orbit_owners (execution_id, root_id, node_id, epoch, since_ms)
711                             VALUES ($1, COALESCE((SELECT root_id FROM orbit_owners WHERE execution_id = $2), $1), $3, 1, {NOW_MS})
712                             ON CONFLICT DO NOTHING"
713                        ),
714                        &[&id, &ext.parent, node],
715                    )?;
716                    if created == 0 {
717                        return Err(StoreError(format!(
718                            "fenced: execution `{id}` was created by another server at the same time"
719                        )));
720                    }
721                }
722            }
723            let ins = tx.prepare(
724                "INSERT INTO execution_events (execution_id, seq, event_type, step, payload, timestamp_ms)
725                 VALUES ($1, $2, $3, $4, $5, $6)",
726            )?;
727            for (seq, ty, step, payload, ts) in &rows {
728                tx.execute(&ins, &[&id, seq, ty, step, payload, ts])?;
729            }
730            let (workflow, version, hash, status, created, updated, last_seq) = projection;
731            tx.execute(
732                "INSERT INTO workflow_executions (execution_id, workflow, version, definition_hash, status, created_at_ms, updated_at_ms, last_seq)
733                 VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
734                 ON CONFLICT (execution_id) DO UPDATE SET status = EXCLUDED.status, updated_at_ms = EXCLUDED.updated_at_ms, last_seq = EXCLUDED.last_seq",
735                &[&id, &workflow, &version, &hash, &status, &created, &updated, &last_seq],
736            )?;
737            if !ext.done_tokens.is_empty() {
738                tx.execute(
739                    "UPDATE orbit_tasks SET state = 'done', updated_ms = $2 WHERE token = ANY($1) AND state IN ('claimed', 'cancelled')",
740                    &[&ext.done_tokens, &now],
741                )?;
742            }
743            if !ext.offers.is_empty() {
744                let ins = tx.prepare(
745                    "INSERT INTO orbit_tasks (execution_id, unit_key, attempt, offer, task_name, contract_version, language, payload, enqueued_ms, state, updated_ms)
746                     VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, 'queued', $10)
747                     ON CONFLICT (execution_id, unit_key, attempt, offer) DO NOTHING",
748                )?;
749                for o in &ext.offers {
750                    tx.execute(
751                        &ins,
752                        &[&id, &o.unit_key, &o.attempt, &o.offer, &o.task_name, &o.contract_version, &o.language, &o.payload, &o.enqueued_ms, &now],
753                    )?;
754                }
755            }
756            if !ext.syncs.is_empty() {
757                let upd = tx.prepare(
758                    "UPDATE orbit_tasks SET state = 'cancelled', updated_ms = $5
759                     WHERE execution_id = $1 AND unit_key = $2 AND state IN ('queued', 'claimed') AND NOT (attempt = $3 AND offer = $4)",
760                )?;
761                for s in &ext.syncs {
762                    let (a, o) = s.keep.unwrap_or((-1, -1));
763                    tx.execute(&upd, &[&id, &s.unit_key, &a, &o, &now])?;
764                }
765            }
766            for (seq, reply) in &ext.consumed {
767                match reply {
768                    None => tx.execute("DELETE FROM orbit_inbox WHERE seq = $1", &[seq])?,
769                    Some(r) => tx.execute(
770                        "UPDATE orbit_inbox SET done_ms = $2, result = $3 WHERE seq = $1",
771                        &[seq, &now, r],
772                    )?,
773                };
774            }
775            if let Some(reply) = &ext.finished_reply {
776                tx.execute(
777                    "UPDATE orbit_tasks SET state = 'cancelled', updated_ms = $2 WHERE execution_id = $1 AND state IN ('queued', 'claimed')",
778                    &[&id, &now],
779                )?;
780                tx.execute("DELETE FROM orbit_owners WHERE execution_id = $1", &[&id])?;
781                tx.execute(
782                    "UPDATE orbit_inbox SET done_ms = $2, result = $3 WHERE execution_id = $1 AND done_ms IS NULL AND kind IN ('signal', 'cancel')",
783                    &[&id, &now, reply],
784                )?;
785                tx.execute(
786                    "DELETE FROM orbit_inbox WHERE execution_id = $1 AND done_ms IS NULL",
787                    &[&id],
788                )?;
789                tx.execute("SELECT pg_notify('orbit_done', $1)", &[&id])?;
790            }
791            if !ext.offers.is_empty() {
792                tx.execute("SELECT pg_notify('orbit_tasks', '')", &[])?;
793            }
794            tx.commit()?;
795            Ok(())
796        })
797    }
798}