1use super::{
11 ExecutionSummary, FireRecord, PutDefinition, Result, Rollover, ScheduleRow, StoreError,
12 WorkerRecord, SCHEMA,
13};
14use crate::model::{Event, ExecutionState, HistoryEvent};
15use orbit_ir::WorkflowIr;
16use postgres::Client;
17use std::sync::mpsc;
18use tokio_postgres_rustls::MakeRustlsConnect;
19
20type Job = Box<dyn FnOnce(&mut Client) + Send>;
21
22pub struct PgStore {
23 jobs: mpsc::Sender<Job>,
24 mode: PgMode,
25}
26
27#[derive(Debug, Clone, Copy, PartialEq, Eq)]
31pub enum PgMode {
32 ReadWrite,
33 ReadOnly,
36 Fenced(i64),
39 Shared,
45}
46
47pub(crate) const CLUSTER_SCHEMA: &str = r#"
49CREATE TABLE IF NOT EXISTS orbit_cluster (
50 id INT PRIMARY KEY,
51 term BIGINT NOT NULL,
52 leader_node TEXT,
53 leader_addr TEXT,
54 since_ms BIGINT NOT NULL
55);
56INSERT INTO orbit_cluster (id, term, since_ms) VALUES (1, 0, 0) ON CONFLICT DO NOTHING;
57"#;
58
59pub(crate) fn session_client(url: &str, application_name: &str) -> Result<Client> {
63 use std::str::FromStr;
64 use std::time::Duration;
65 let mut config = postgres::Config::from_str(url)?;
66 config
67 .application_name(application_name)
68 .connect_timeout(Duration::from_secs(5))
69 .keepalives(true)
70 .keepalives_idle(Duration::from_secs(5))
71 .tcp_user_timeout(Duration::from_secs(5));
72 Ok(config.connect(tls()?)?)
73}
74
75fn read_only() -> StoreError {
76 StoreError("this server is a standby: the store is read-only".into())
77}
78
79fn tls() -> Result<MakeRustlsConnect> {
84 let mut roots = rustls::RootCertStore::empty();
85 roots.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());
86 if let Ok(path) = std::env::var("ORBIT_DB_CA_CERT") {
87 let pem = std::fs::read(&path)
88 .map_err(|e| StoreError(format!("cannot read ORBIT_DB_CA_CERT {path}: {e}")))?;
89 for cert in rustls_pemfile::certs(&mut pem.as_slice()) {
90 let cert =
91 cert.map_err(|e| StoreError(format!("invalid certificate in {path}: {e}")))?;
92 roots
93 .add(cert)
94 .map_err(|e| StoreError(format!("invalid CA certificate in {path}: {e}")))?;
95 }
96 }
97 let provider = std::sync::Arc::new(rustls::crypto::aws_lc_rs::default_provider());
98 let config = rustls::ClientConfig::builder_with_provider(provider)
99 .with_safe_default_protocol_versions()
100 .map_err(|e| StoreError(format!("TLS configuration: {e}")))?
101 .with_root_certificates(roots)
102 .with_no_client_auth();
103 Ok(MakeRustlsConnect::new(config))
104}
105
106fn connect_client(url: &str) -> Result<Client> {
107 let mut client = Client::connect(url, tls()?)?;
108 client.batch_execute(SCHEMA)?;
109 let rows = client.query_one("SELECT COUNT(*) FROM schema_version", &[])?;
110 if rows.get::<_, i64>(0) == 0 {
111 client.execute("INSERT INTO schema_version (version) VALUES (1)", &[])?;
112 }
113 Ok(client)
114}
115
116type EventRow = (i64, &'static str, Option<String>, String, i64);
117
118fn bind<'a>(
121 id: &'a String,
122 projection: Option<&[&'a (dyn postgres::types::ToSql + Sync); 7]>,
123 chunk: &'a [EventRow],
124) -> Vec<&'a (dyn postgres::types::ToSql + Sync)> {
125 let mut v: Vec<&'a (dyn postgres::types::ToSql + Sync)> =
126 Vec::with_capacity(8 + chunk.len() * 5);
127 v.push(id);
128 if let Some(p) = projection {
129 v.extend(p.iter().copied());
130 }
131 for (seq, ty, step, payload, ts) in chunk {
132 v.extend([
133 seq as &(dyn postgres::types::ToSql + Sync),
134 ty,
135 step,
136 payload,
137 ts,
138 ]);
139 }
140 v
141}
142
143fn connect_client_serialized(url: &str) -> Result<Client> {
147 let mut client = Client::connect(url, tls()?)?;
148 let mut tx = client.transaction()?;
149 tx.execute("SELECT pg_advisory_xact_lock($1)", &[&BASE_SCHEMA_LOCK_KEY])?;
150 tx.batch_execute(SCHEMA)?;
151 if tx
152 .query_one("SELECT COUNT(*) FROM schema_version", &[])?
153 .get::<_, i64>(0)
154 == 0
155 {
156 tx.execute("INSERT INTO schema_version (version) VALUES (1)", &[])?;
157 }
158 tx.commit()?;
159 Ok(client)
160}
161
162const BASE_SCHEMA_LOCK_KEY: i64 = 0x6f72_6269_7473_6362;
164
165fn bad_payload(e: serde_json::Error) -> StoreError {
166 StoreError(format!("postgres: unreadable event payload: {e}"))
167}
168
169impl PgStore {
170 pub fn connect(url: &str) -> Result<PgStore> {
171 PgStore::connect_with(url, PgMode::ReadWrite)
172 }
173
174 pub fn connect_with(url: &str, mode: PgMode) -> Result<PgStore> {
176 let url = url.to_string();
177 let (jobs, rx) = mpsc::channel::<Job>();
178 let (ready_tx, ready_rx) = mpsc::channel::<Result<()>>();
179 std::thread::Builder::new()
180 .name("orbit-postgres".into())
181 .spawn(move || {
182 let open = |url: &str| match mode {
183 PgMode::Shared => connect_client_serialized(url),
184 _ => connect_client(url),
185 };
186 let mut client = match open(&url) {
187 Ok(c) => {
188 let _ = ready_tx.send(Ok(()));
189 c
190 }
191 Err(e) => {
192 let _ = ready_tx.send(Err(e));
193 return;
194 }
195 };
196 for job in rx {
197 if client.is_closed() {
198 match open(&url) {
199 Ok(c) => {
200 tracing::warn!("reconnected to PostgreSQL");
201 client = c;
202 }
203 Err(e) => tracing::error!(error = %e, "PostgreSQL reconnect failed"),
204 }
205 }
206 job(&mut client);
207 }
208 })
209 .map_err(|e| StoreError(format!("cannot start PostgreSQL thread: {e}")))?;
210 ready_rx
211 .recv()
212 .map_err(|_| StoreError("PostgreSQL thread exited during startup".into()))??;
213 Ok(PgStore { jobs, mode })
214 }
215
216 pub(crate) fn run<T: Send + 'static>(
217 &self,
218 f: impl FnOnce(&mut Client) -> Result<T> + Send + 'static,
219 ) -> Result<T> {
220 let (tx, rx) = mpsc::channel();
221 self.jobs
222 .send(Box::new(move |c| {
223 let _ = tx.send(f(c));
224 }))
225 .map_err(|_| StoreError("PostgreSQL connection thread has stopped".into()))?;
226 rx.recv()
227 .map_err(|_| StoreError("PostgreSQL connection thread has stopped".into()))?
228 }
229
230 pub fn mode(&self) -> PgMode {
231 self.mode
232 }
233
234 pub fn put_definition(&self, ir: &WorkflowIr, now: u64) -> Result<PutDefinition> {
235 if self.mode == PgMode::ReadOnly {
236 return Err(read_only());
237 }
238 let (name, version, hash) = (
239 ir.name.clone(),
240 ir.version.clone(),
241 ir.definition_hash.clone(),
242 );
243 let text = serde_json::to_string(ir).expect("IR serialises");
244 self.run(move |c| {
245 let existing = c.query_opt(
246 "SELECT definition_hash FROM workflow_definitions WHERE name = $1 AND version = $2",
247 &[&name, &version],
248 )?;
249 match existing.map(|r| r.get::<_, String>(0)) {
250 Some(h) if h == hash => Ok(PutDefinition::Unchanged),
251 Some(h) => Ok(PutDefinition::Conflict { existing_hash: h }),
252 None => {
253 let n = c.execute(
255 "INSERT INTO workflow_definitions (name, version, definition_hash, ir, created_at_ms)
256 VALUES ($1, $2, $3, $4, $5) ON CONFLICT DO NOTHING",
257 &[&name, &version, &hash, &text, &(now as i64)],
258 )?;
259 Ok(if n == 1 { PutDefinition::Created } else { PutDefinition::Unchanged })
260 }
261 }
262 })
263 }
264
265 pub fn definitions(&self) -> Result<Vec<WorkflowIr>> {
266 let rows = self.run(|c| {
267 Ok(c.query(
268 "SELECT ir FROM workflow_definitions ORDER BY name, created_at_ms",
269 &[],
270 )?
271 .into_iter()
272 .map(|r| r.get::<_, String>(0))
273 .collect::<Vec<_>>())
274 })?;
275 Ok(rows
276 .into_iter()
277 .filter_map(|t| match serde_json::from_str::<WorkflowIr>(&t) {
278 Ok(ir) => Some(ir),
279 Err(e) => {
280 tracing::error!(error = %e, "skipping unreadable workflow definition");
281 None
282 }
283 })
284 .collect())
285 }
286
287 pub fn execution_exists(&self, id: &str) -> Result<bool> {
288 let id = id.to_string();
289 self.run(move |c| {
290 Ok(c.query_opt(
291 "SELECT 1 FROM workflow_executions WHERE execution_id = $1",
292 &[&id],
293 )?
294 .is_some())
295 })
296 }
297
298 pub fn append(&mut self, state: &ExecutionState, events: &[HistoryEvent]) -> Result<()> {
299 let fence = match self.mode {
300 PgMode::ReadWrite => None,
301 PgMode::ReadOnly => return Err(read_only()),
302 PgMode::Fenced(term) => Some(term),
303 PgMode::Shared => {
304 return Err(StoreError(
305 "shared mode: history appends must carry the execution's ownership epoch"
306 .into(),
307 ))
308 }
309 };
310 let rows: Vec<EventRow> = events
311 .iter()
312 .map(|h| {
313 (
314 h.seq as i64,
315 h.event.name(),
316 h.event.step().map(str::to_string),
317 serde_json::to_string(&h.event).expect("event serialises"),
318 h.timestamp_ms as i64,
319 )
320 })
321 .collect();
322 let id = state.id.clone();
323 let projection = (
324 state.workflow.clone(),
325 state.version.clone(),
326 state.definition_hash.clone(),
327 state.status.as_str(),
328 state.created_at_ms as i64,
329 state.updated_at_ms as i64,
330 state.last_seq as i64,
331 );
332 self.run(move |c| {
333 type P<'a> = &'a (dyn postgres::types::ToSql + Sync);
334 let (workflow, version, hash, status, created, updated, last_seq) = &projection;
335 let statement = |chunk: &[EventRow],
340 with_projection: bool|
341 -> String {
342 let base = if with_projection { 8 } else { 1 };
343 let mut sql = String::from(
344 "INSERT INTO execution_events (execution_id, seq, event_type, step, payload, timestamp_ms) VALUES ",
345 );
346 for i in 0..chunk.len() {
347 let b = base + 1 + i * 5;
348 if i > 0 {
349 sql.push(',');
350 }
351 sql.push_str(&format!(
352 "($1, ${}, ${}, ${}, ${}, ${})",
353 b,
354 b + 1,
355 b + 2,
356 b + 3,
357 b + 4
358 ));
359 }
360 if !with_projection {
361 return sql;
362 }
363 format!(
364 "WITH ev AS ({sql} RETURNING 1) \
365 INSERT INTO workflow_executions (execution_id, workflow, version, definition_hash, status, created_at_ms, updated_at_ms, last_seq) \
366 VALUES ($1, $2, $3, $4, $5, $6, $7, $8) \
367 ON CONFLICT (execution_id) DO UPDATE SET status = EXCLUDED.status, updated_at_ms = EXCLUDED.updated_at_ms, last_seq = EXCLUDED.last_seq"
368 )
369 };
370 let proj: [P<'_>; 7] = [workflow, version, hash, status, created, updated, last_seq];
371 const CHUNK: usize = 1000;
372 if fence.is_none() && rows.len() <= CHUNK {
373 c.execute(&statement(&rows, true), &bind(&id, Some(&proj), &rows))?;
374 return Ok(());
375 }
376 let mut tx = c.transaction()?;
377 if let Some(term) = fence {
378 let current: Option<i64> = tx
381 .query_opt("SELECT term FROM orbit_cluster WHERE id = 1 FOR SHARE", &[])?
382 .map(|r| r.get(0));
383 if current != Some(term) {
384 return Err(StoreError(format!(
385 "fenced: this server led term {term}, but the cluster is at term {}; another server is the leader",
386 current.map_or("unknown".to_string(), |t| t.to_string())
387 )));
388 }
389 }
390 let chunks: Vec<_> = rows.chunks(CHUNK).collect();
391 for (i, chunk) in chunks.iter().enumerate() {
392 let last = i + 1 == chunks.len();
393 tx.execute(
394 &statement(chunk, last),
395 &bind(&id, last.then_some(&proj), chunk),
396 )?;
397 }
398 tx.commit()?;
399 Ok(())
400 })
401 }
402
403 pub fn history(&self, id: &str) -> Result<Vec<HistoryEvent>> {
404 let id = id.to_string();
405 let rows = self.run(move |c| {
406 Ok(c.query(
407 "SELECT seq, timestamp_ms, payload FROM execution_events WHERE execution_id = $1 ORDER BY seq",
408 &[&id],
409 )?
410 .into_iter()
411 .map(|r| (r.get::<_, i64>(0), r.get::<_, i64>(1), r.get::<_, String>(2)))
412 .collect::<Vec<_>>())
413 })?;
414 rows.into_iter()
415 .map(|(seq, ts, payload)| {
416 let event: Event = serde_json::from_str(&payload).map_err(bad_payload)?;
417 Ok(HistoryEvent {
418 seq: seq as u64,
419 timestamp_ms: ts as u64,
420 event,
421 })
422 })
423 .collect()
424 }
425
426 pub fn append_rollover(
427 &mut self,
428 state: &ExecutionState,
429 events: &[HistoryEvent],
430 rollover: Rollover<'_>,
431 ) -> Result<()> {
432 let fence = match self.mode {
433 PgMode::ReadWrite => None,
434 PgMode::ReadOnly => return Err(read_only()),
435 PgMode::Fenced(term) => Some(term),
436 PgMode::Shared => {
437 return Err(StoreError(
438 "shared mode: history appends must carry the execution's ownership epoch"
439 .into(),
440 ))
441 }
442 };
443 let tail = event_rows(rollover.events);
444 let rows = event_rows(events);
445 let run = rollover.run as i64;
446 let id = state.id.clone();
447 let projection = projection_of(state);
448 self.run(move |c| {
449 let mut tx = c.transaction()?;
450 if let Some(term) = fence {
451 check_term(&mut tx, term)?;
452 }
453 archive_run(&mut tx, &id, run, &tail)?;
454 insert_events(&mut tx, &id, &rows, &projection)?;
455 tx.commit()?;
456 Ok(())
457 })
458 }
459
460 pub fn run_history(&self, id: &str, run: u32) -> Result<Vec<HistoryEvent>> {
461 let id = id.to_string();
462 let rows = self.run(move |c| {
463 Ok(c.query(
464 "SELECT seq, timestamp_ms, payload FROM execution_run_events WHERE execution_id = $1 AND run = $2 ORDER BY seq",
465 &[&id, &(run as i64)],
466 )?
467 .into_iter()
468 .map(|r| (r.get::<_, i64>(0), r.get::<_, i64>(1), r.get::<_, String>(2)))
469 .collect::<Vec<_>>())
470 })?;
471 rows.into_iter()
472 .map(|(seq, ts, payload)| {
473 let event: Event = serde_json::from_str(&payload).map_err(bad_payload)?;
474 Ok(HistoryEvent {
475 seq: seq as u64,
476 timestamp_ms: ts as u64,
477 event,
478 })
479 })
480 .collect()
481 }
482
483 pub fn execution_status(&self, id: &str) -> Result<Option<String>> {
484 let id = id.to_string();
485 self.run(move |c| {
486 Ok(c.query_opt(
487 "SELECT status FROM workflow_executions WHERE execution_id = $1",
488 &[&id],
489 )?
490 .map(|r| r.get(0)))
491 })
492 }
493
494 pub fn active_execution_ids(&self) -> Result<Vec<String>> {
495 if matches!(self.mode, PgMode::ReadOnly | PgMode::Shared) {
496 return Ok(vec![]);
497 }
498 self.run(|c| {
499 Ok(c.query(
500 "SELECT execution_id FROM workflow_executions WHERE status IN ('RUNNING', 'COMPENSATING') ORDER BY created_at_ms",
501 &[],
502 )?
503 .into_iter()
504 .map(|r| r.get(0))
505 .collect())
506 })
507 }
508
509 pub fn list_executions(
510 &self,
511 status: Option<&str>,
512 workflow: Option<&str>,
513 limit: u32,
514 ) -> Result<Vec<ExecutionSummary>> {
515 let (status, workflow) = (status.map(str::to_string), workflow.map(str::to_string));
516 self.run(move |c| {
517 Ok(c.query(
518 "SELECT execution_id, workflow, version, status, created_at_ms, updated_at_ms FROM workflow_executions
519 WHERE ($1::TEXT IS NULL OR status = $1) AND ($2::TEXT IS NULL OR workflow = $2)
520 ORDER BY created_at_ms DESC LIMIT $3",
521 &[&status, &workflow, &(limit as i64)],
522 )?
523 .into_iter()
524 .map(|r| ExecutionSummary {
525 execution_id: r.get(0),
526 workflow: r.get(1),
527 version: r.get(2),
528 status: r.get(3),
529 created_at_ms: r.get::<_, i64>(4) as u64,
530 updated_at_ms: r.get::<_, i64>(5) as u64,
531 })
532 .collect())
533 })
534 }
535
536 pub fn count_by_status(&self) -> Result<Vec<(String, u64)>> {
537 self.run(|c| {
538 Ok(c.query(
539 "SELECT status, COUNT(*) FROM workflow_executions GROUP BY status",
540 &[],
541 )?
542 .into_iter()
543 .map(|r| (r.get(0), r.get::<_, i64>(1) as u64))
544 .collect())
545 })
546 }
547
548 pub fn upsert_worker(&self, w: &WorkerRecord) -> Result<()> {
549 if self.mode == PgMode::ReadOnly {
550 return Err(read_only());
551 }
552 let w = w.clone();
553 self.run(move |c| {
554 c.execute(
555 "INSERT INTO worker_registry (worker_id, runtime, sdk, contracts, registered_at_ms, last_seen_ms) VALUES ($1, $2, $3, $4, $5, $6)
556 ON CONFLICT (worker_id) DO UPDATE SET runtime = EXCLUDED.runtime, sdk = EXCLUDED.sdk, contracts = EXCLUDED.contracts,
557 registered_at_ms = EXCLUDED.registered_at_ms, last_seen_ms = EXCLUDED.last_seen_ms",
558 &[&w.worker_id, &w.runtime, &w.sdk, &w.contracts.to_string(), &(w.registered_at_ms as i64), &(w.last_seen_ms as i64)],
559 )?;
560 Ok(())
561 })
562 }
563
564 pub fn touch_worker(&self, worker_id: &str, now: u64) -> Result<()> {
565 if self.mode == PgMode::ReadOnly {
566 return Err(read_only());
567 }
568 let id = worker_id.to_string();
569 self.run(move |c| {
570 c.execute(
571 "UPDATE worker_registry SET last_seen_ms = $2 WHERE worker_id = $1",
572 &[&id, &(now as i64)],
573 )?;
574 Ok(())
575 })
576 }
577
578 pub fn workers(&self) -> Result<Vec<WorkerRecord>> {
579 self.run(|c| {
580 Ok(c.query(
581 "SELECT worker_id, runtime, sdk, contracts, registered_at_ms, last_seen_ms FROM worker_registry ORDER BY worker_id",
582 &[],
583 )?
584 .into_iter()
585 .map(|r| WorkerRecord {
586 worker_id: r.get(0),
587 runtime: r.get(1),
588 sdk: r.get(2),
589 contracts: serde_json::from_str(&r.get::<_, String>(3)).unwrap_or_default(),
590 registered_at_ms: r.get::<_, i64>(4) as u64,
591 last_seen_ms: r.get::<_, i64>(5) as u64,
592 })
593 .collect())
594 })
595 }
596}
597
598type Projection = (String, String, String, &'static str, i64, i64, i64);
599
600fn projection_of(state: &ExecutionState) -> Projection {
601 (
602 state.workflow.clone(),
603 state.version.clone(),
604 state.definition_hash.clone(),
605 state.status.as_str(),
606 state.created_at_ms as i64,
607 state.updated_at_ms as i64,
608 state.last_seq as i64,
609 )
610}
611
612fn event_rows(events: &[HistoryEvent]) -> Vec<EventRow> {
613 events
614 .iter()
615 .map(|h| {
616 (
617 h.seq as i64,
618 h.event.name(),
619 h.event.step().map(str::to_string),
620 serde_json::to_string(&h.event).expect("event serialises"),
621 h.timestamp_ms as i64,
622 )
623 })
624 .collect()
625}
626
627fn check_term(tx: &mut postgres::Transaction<'_>, term: i64) -> Result<()> {
629 let current: Option<i64> = tx
630 .query_opt("SELECT term FROM orbit_cluster WHERE id = 1 FOR SHARE", &[])?
631 .map(|r| r.get(0));
632 if current != Some(term) {
633 return Err(StoreError(format!(
634 "fenced: this server led term {term}, but the cluster is at term {}; another server is the leader",
635 current.map_or("unknown".to_string(), |t| t.to_string())
636 )));
637 }
638 Ok(())
639}
640
641fn archive_run(
644 tx: &mut postgres::Transaction<'_>,
645 id: &String,
646 run: i64,
647 tail: &[EventRow],
648) -> Result<()> {
649 tx.execute(
650 "INSERT INTO execution_run_events (execution_id, run, seq, event_type, step, payload, timestamp_ms)
651 SELECT execution_id, $2, seq, event_type, step, payload, timestamp_ms FROM execution_events WHERE execution_id = $1",
652 &[id, &run],
653 )?;
654 tx.execute(
655 "DELETE FROM execution_events WHERE execution_id = $1",
656 &[id],
657 )?;
658 let ins = tx.prepare(
659 "INSERT INTO execution_run_events (execution_id, run, seq, event_type, step, payload, timestamp_ms)
660 VALUES ($1, $2, $3, $4, $5, $6, $7)",
661 )?;
662 for (seq, ty, step, payload, ts) in tail {
663 tx.execute(&ins, &[id, &run, seq, ty, step, payload, ts])?;
664 }
665 Ok(())
666}
667
668fn insert_events(
671 tx: &mut postgres::Transaction<'_>,
672 id: &String,
673 rows: &[EventRow],
674 projection: &Projection,
675) -> Result<()> {
676 let ins = tx.prepare(
677 "INSERT INTO execution_events (execution_id, seq, event_type, step, payload, timestamp_ms)
678 VALUES ($1, $2, $3, $4, $5, $6)",
679 )?;
680 for (seq, ty, step, payload, ts) in rows {
681 tx.execute(&ins, &[id, seq, ty, step, payload, ts])?;
682 }
683 let (workflow, version, hash, status, created, updated, last_seq) = projection;
684 tx.execute(
685 "INSERT INTO workflow_executions (execution_id, workflow, version, definition_hash, status, created_at_ms, updated_at_ms, last_seq)
686 VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
687 ON CONFLICT (execution_id) DO UPDATE SET status = EXCLUDED.status, updated_at_ms = EXCLUDED.updated_at_ms, last_seq = EXCLUDED.last_seq",
688 &[id, workflow, version, hash, status, created, updated, last_seq],
689 )?;
690 Ok(())
691}
692
693const SCHEDULE_COLUMNS: &str = "schedule_id, spec, paused, next_fire_ms, due_ms, last_fire_ms, last_execution_id, fired_count, skipped_count, last_error, created_at_ms, updated_at_ms";
696
697fn schedule_row(r: &postgres::Row) -> ScheduleRow {
698 let opt = |v: Option<i64>| v.map(|x| x as u64);
699 ScheduleRow {
700 schedule_id: r.get(0),
701 spec: r.get(1),
702 paused: r.get::<_, i64>(2) != 0,
703 next_fire_ms: opt(r.get(3)),
704 due_ms: opt(r.get(4)),
705 last_fire_ms: opt(r.get(5)),
706 last_execution_id: r.get(6),
707 fired: r.get::<_, i64>(7) as u64,
708 skipped: r.get::<_, i64>(8) as u64,
709 last_error: r.get(9),
710 created_at_ms: r.get::<_, i64>(10) as u64,
711 updated_at_ms: r.get::<_, i64>(11) as u64,
712 }
713}
714
715fn opt_i64(v: Option<u64>) -> Option<i64> {
716 v.map(|x| x as i64)
717}
718
719impl PgStore {
720 fn writable(&self) -> Result<Option<i64>> {
721 match self.mode {
722 PgMode::ReadOnly => Err(read_only()),
723 PgMode::Fenced(term) => Ok(Some(term)),
724 _ => Ok(None),
725 }
726 }
727
728 pub fn insert_schedule(&self, row: &ScheduleRow) -> Result<bool> {
729 self.writable()?;
730 let r = row.clone();
731 self.run(move |c| {
732 let n = c.execute(
733 &format!(
734 "INSERT INTO orbit_schedules ({SCHEDULE_COLUMNS}) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12)
735 ON CONFLICT (schedule_id) DO NOTHING"
736 ),
737 &[
738 &r.schedule_id,
739 &r.spec,
740 &(r.paused as i64),
741 &opt_i64(r.next_fire_ms),
742 &opt_i64(r.due_ms),
743 &opt_i64(r.last_fire_ms),
744 &r.last_execution_id,
745 &(r.fired as i64),
746 &(r.skipped as i64),
747 &r.last_error,
748 &(r.created_at_ms as i64),
749 &(r.updated_at_ms as i64),
750 ],
751 )?;
752 Ok(n == 1)
753 })
754 }
755
756 pub fn update_schedule(&self, row: &ScheduleRow) -> Result<bool> {
757 self.writable()?;
758 let r = row.clone();
759 self.run(move |c| {
760 let n = c.execute(
761 "UPDATE orbit_schedules SET spec = $2, paused = $3, next_fire_ms = $4, due_ms = $5, last_error = $6,
762 claimed_by = NULL, claimed_until_ms = NULL, updated_at_ms = $7 WHERE schedule_id = $1",
763 &[
764 &r.schedule_id,
765 &r.spec,
766 &(r.paused as i64),
767 &opt_i64(r.next_fire_ms),
768 &opt_i64(r.due_ms),
769 &r.last_error,
770 &(r.updated_at_ms as i64),
771 ],
772 )?;
773 Ok(n == 1)
774 })
775 }
776
777 pub fn schedule(&self, id: &str) -> Result<Option<ScheduleRow>> {
778 let id = id.to_string();
779 self.run(move |c| {
780 Ok(c.query_opt(
781 &format!("SELECT {SCHEDULE_COLUMNS} FROM orbit_schedules WHERE schedule_id = $1"),
782 &[&id],
783 )?
784 .map(|r| schedule_row(&r)))
785 })
786 }
787
788 pub fn schedules(&self) -> Result<Vec<ScheduleRow>> {
789 self.run(|c| {
790 Ok(c.query(
791 &format!("SELECT {SCHEDULE_COLUMNS} FROM orbit_schedules ORDER BY schedule_id"),
792 &[],
793 )?
794 .iter()
795 .map(schedule_row)
796 .collect())
797 })
798 }
799
800 pub fn delete_schedule(&self, id: &str) -> Result<bool> {
801 self.writable()?;
802 let id = id.to_string();
803 self.run(move |c| {
804 Ok(c.execute("DELETE FROM orbit_schedules WHERE schedule_id = $1", &[&id])? == 1)
805 })
806 }
807
808 pub fn due_schedules(&self, now: u64) -> Result<Vec<ScheduleRow>> {
809 let now = now as i64;
810 self.run(move |c| {
811 Ok(c.query(
812 &format!(
813 "SELECT {SCHEDULE_COLUMNS} FROM orbit_schedules
814 WHERE paused = 0 AND due_ms IS NOT NULL AND due_ms <= $1
815 AND (claimed_until_ms IS NULL OR claimed_until_ms <= $1)
816 ORDER BY due_ms LIMIT 100"
817 ),
818 &[&now],
819 )?
820 .iter()
821 .map(schedule_row)
822 .collect())
823 })
824 }
825
826 pub fn claim_schedule(
827 &self,
828 id: &str,
829 next_fire_ms: u64,
830 node: &str,
831 now: u64,
832 until: u64,
833 ) -> Result<bool> {
834 let fence = self.writable()?;
835 let (id, node) = (id.to_string(), node.to_string());
836 self.run(move |c| {
837 let mut tx = c.transaction()?;
838 if let Some(term) = fence {
839 check_term(&mut tx, term)?;
841 }
842 let n = tx.execute(
843 "UPDATE orbit_schedules SET claimed_by = $3, claimed_until_ms = $5
844 WHERE schedule_id = $1 AND next_fire_ms = $2 AND paused = 0
845 AND (claimed_until_ms IS NULL OR claimed_until_ms <= $4)",
846 &[
847 &id,
848 &(next_fire_ms as i64),
849 &node,
850 &(now as i64),
851 &(until as i64),
852 ],
853 )?;
854 tx.commit()?;
855 Ok(n == 1)
856 })
857 }
858
859 pub fn finish_schedule_fire(
860 &self,
861 id: &str,
862 next_fire_ms: u64,
863 node: &str,
864 now: u64,
865 rec: &FireRecord,
866 ) -> Result<bool> {
867 let fence = self.writable()?;
868 let (id, node, rec) = (id.to_string(), node.to_string(), rec.clone());
869 self.run(move |c| {
870 let mut tx = c.transaction()?;
871 if let Some(term) = fence {
872 check_term(&mut tx, term)?;
873 }
874 let n = tx.execute(
875 "UPDATE orbit_schedules SET next_fire_ms = $4, due_ms = $5,
876 last_fire_ms = COALESCE($6, last_fire_ms), last_execution_id = COALESCE($7, last_execution_id),
877 fired_count = fired_count + $8, skipped_count = skipped_count + $9, last_error = $10,
878 claimed_by = NULL, claimed_until_ms = NULL, updated_at_ms = $11
879 WHERE schedule_id = $1 AND next_fire_ms = $2 AND claimed_by = $3",
880 &[
881 &id,
882 &(next_fire_ms as i64),
883 &node,
884 &opt_i64(rec.next_fire_ms),
885 &opt_i64(rec.due_ms),
886 &opt_i64(rec.last_fire_ms),
887 &rec.last_execution_id,
888 &(rec.fired as i64),
889 &(rec.skipped as i64),
890 &rec.last_error,
891 &(now as i64),
892 ],
893 )?;
894 tx.commit()?;
895 Ok(n == 1)
896 })
897 }
898
899 pub fn release_schedule_claims(&self) -> Result<u64> {
900 self.writable()?;
901 self.run(|c| {
902 Ok(c.execute(
903 "UPDATE orbit_schedules SET claimed_by = NULL, claimed_until_ms = NULL WHERE claimed_by IS NOT NULL",
904 &[],
905 )?)
906 })
907 }
908
909 pub fn release_schedule_claim(
910 &self,
911 id: &str,
912 node: &str,
913 retry_at: u64,
914 error: &str,
915 now: u64,
916 ) -> Result<()> {
917 self.writable()?;
918 let (id, node, error) = (id.to_string(), node.to_string(), error.to_string());
919 self.run(move |c| {
920 c.execute(
921 "UPDATE orbit_schedules SET claimed_by = NULL, claimed_until_ms = $3, last_error = $4, updated_at_ms = $5
922 WHERE schedule_id = $1 AND claimed_by = $2",
923 &[&id, &node, &(retry_at as i64), &error, &(now as i64)],
924 )?;
925 Ok(())
926 })
927 }
928
929 pub fn record_schedule_trigger(&self, id: &str, execution_id: &str, now: u64) -> Result<()> {
930 self.writable()?;
931 let (id, exec) = (id.to_string(), execution_id.to_string());
932 self.run(move |c| {
933 c.execute(
934 "UPDATE orbit_schedules SET last_execution_id = $2, fired_count = fired_count + 1, updated_at_ms = $3
935 WHERE schedule_id = $1",
936 &[&id, &exec, &(now as i64)],
937 )?;
938 Ok(())
939 })
940 }
941}
942
943pub(crate) const NOW_MS: &str = "(extract(epoch from clock_timestamp()) * 1000)::bigint";
949
950pub(crate) const SHARED_SCHEMA: &str = r#"
953CREATE TABLE IF NOT EXISTS orbit_nodes (
954 node_id TEXT PRIMARY KEY,
955 incarnation TEXT NOT NULL,
956 addr TEXT,
957 heartbeat_ms BIGINT NOT NULL,
958 started_ms BIGINT NOT NULL
959);
960CREATE TABLE IF NOT EXISTS orbit_owners (
961 execution_id TEXT PRIMARY KEY,
962 root_id TEXT NOT NULL,
963 node_id TEXT NOT NULL,
964 epoch BIGINT NOT NULL,
965 since_ms BIGINT NOT NULL
966);
967CREATE INDEX IF NOT EXISTS idx_orbit_owners_node ON orbit_owners(node_id);
968CREATE INDEX IF NOT EXISTS idx_orbit_owners_root ON orbit_owners(root_id);
969CREATE TABLE IF NOT EXISTS orbit_tasks (
970 id BIGSERIAL PRIMARY KEY,
971 execution_id TEXT NOT NULL,
972 unit_key TEXT NOT NULL,
973 attempt BIGINT NOT NULL,
974 offer BIGINT NOT NULL,
975 task_name TEXT NOT NULL,
976 contract_version TEXT NOT NULL,
977 language TEXT,
978 payload TEXT NOT NULL,
979 enqueued_ms BIGINT NOT NULL,
980 state TEXT NOT NULL,
981 token TEXT,
982 claimed_by TEXT,
983 claimed_inc TEXT,
984 claimed_session TEXT,
985 claimed_worker TEXT,
986 claimed_ms BIGINT,
987 updated_ms BIGINT NOT NULL,
988 UNIQUE (execution_id, unit_key, attempt, offer)
989);
990CREATE INDEX IF NOT EXISTS idx_orbit_tasks_queued ON orbit_tasks(task_name, enqueued_ms) WHERE state = 'queued';
991CREATE INDEX IF NOT EXISTS idx_orbit_tasks_token ON orbit_tasks(token);
992CREATE INDEX IF NOT EXISTS idx_orbit_tasks_session ON orbit_tasks(claimed_session);
993CREATE TABLE IF NOT EXISTS orbit_inbox (
994 seq BIGSERIAL PRIMARY KEY,
995 execution_id TEXT NOT NULL,
996 kind TEXT NOT NULL,
997 token TEXT,
998 payload TEXT NOT NULL,
999 created_ms BIGINT NOT NULL,
1000 done_ms BIGINT,
1001 result TEXT
1002);
1003CREATE INDEX IF NOT EXISTS idx_orbit_inbox_execution ON orbit_inbox(execution_id, seq);
1004CREATE TABLE IF NOT EXISTS orbit_sessions (
1005 session_id TEXT PRIMARY KEY,
1006 worker_id TEXT NOT NULL,
1007 runtime TEXT NOT NULL,
1008 contracts TEXT NOT NULL,
1009 node_id TEXT NOT NULL,
1010 registered_ms BIGINT NOT NULL,
1011 last_seen_ms BIGINT NOT NULL
1012);
1013CREATE INDEX IF NOT EXISTS idx_orbit_sessions_worker ON orbit_sessions(worker_id);
1014"#;
1015
1016#[derive(Debug, Clone)]
1018pub(crate) struct TaskOffer {
1019 pub unit_key: String,
1020 pub attempt: i64,
1021 pub offer: i64,
1025 pub task_name: String,
1026 pub contract_version: String,
1027 pub language: Option<String>,
1028 pub payload: String,
1030 pub enqueued_ms: i64,
1031}
1032
1033#[derive(Debug, Clone)]
1037pub(crate) struct UnitSync {
1038 pub unit_key: String,
1039 pub keep: Option<(i64, i64)>,
1040}
1041
1042#[derive(Debug, Clone, Default)]
1044pub(crate) struct SharedAppend {
1045 pub node_id: String,
1046 pub epoch: Option<i64>,
1049 pub parent: Option<String>,
1050 pub offers: Vec<TaskOffer>,
1051 pub syncs: Vec<UnitSync>,
1052 pub done_tokens: Vec<String>,
1054 pub consumed: Vec<(i64, Option<String>)>,
1057 pub finished_reply: Option<String>,
1060 pub rollover: Option<(u32, Vec<HistoryEvent>)>,
1064}
1065
1066impl PgStore {
1067 pub(crate) fn append_shared(
1072 &mut self,
1073 state: &ExecutionState,
1074 events: &[HistoryEvent],
1075 ext: SharedAppend,
1076 ) -> Result<()> {
1077 if self.mode != PgMode::Shared {
1078 return Err(StoreError("append_shared needs a shared-mode store".into()));
1079 }
1080 let rows: Vec<(i64, &'static str, Option<String>, String, i64)> = events
1081 .iter()
1082 .map(|h| {
1083 (
1084 h.seq as i64,
1085 h.event.name(),
1086 h.event.step().map(str::to_string),
1087 serde_json::to_string(&h.event).expect("event serialises"),
1088 h.timestamp_ms as i64,
1089 )
1090 })
1091 .collect();
1092 let id = state.id.clone();
1093 let projection = (
1094 state.workflow.clone(),
1095 state.version.clone(),
1096 state.definition_hash.clone(),
1097 state.status.as_str(),
1098 state.created_at_ms as i64,
1099 state.updated_at_ms as i64,
1100 state.last_seq as i64,
1101 );
1102 let now = state.updated_at_ms as i64;
1103 self.run(move |c| {
1104 let mut tx = c.transaction()?;
1105 let node = &ext.node_id;
1106 match ext.epoch {
1107 Some(epoch) => {
1108 let owner: Option<(String, i64)> = tx
1111 .query_opt(
1112 "SELECT node_id, epoch FROM orbit_owners WHERE execution_id = $1 FOR SHARE",
1113 &[&id],
1114 )?
1115 .map(|r| (r.get(0), r.get(1)));
1116 if owner.as_ref() != Some(&(node.clone(), epoch)) {
1117 return Err(StoreError(match owner {
1118 Some((n, e)) => format!(
1119 "fenced: execution `{id}` is owned by node `{n}` at epoch {e}; this node held epoch {epoch}"
1120 ),
1121 None => format!(
1122 "fenced: execution `{id}` has no owner any more; this node held epoch {epoch}"
1123 ),
1124 }));
1125 }
1126 }
1127 None => {
1128 let created = tx.execute(
1129 &format!(
1130 "INSERT INTO orbit_owners (execution_id, root_id, node_id, epoch, since_ms)
1131 VALUES ($1, COALESCE((SELECT root_id FROM orbit_owners WHERE execution_id = $2), $1), $3, 1, {NOW_MS})
1132 ON CONFLICT DO NOTHING"
1133 ),
1134 &[&id, &ext.parent, node],
1135 )?;
1136 if created == 0 {
1137 return Err(StoreError(format!(
1138 "fenced: execution `{id}` was created by another server at the same time"
1139 )));
1140 }
1141 }
1142 }
1143 if let Some((run, tail)) = &ext.rollover {
1144 archive_run(&mut tx, &id, *run as i64, &event_rows(tail))?;
1145 tx.execute("DELETE FROM orbit_tasks WHERE execution_id = $1", &[&id])?;
1146 }
1147 let ins = tx.prepare(
1148 "INSERT INTO execution_events (execution_id, seq, event_type, step, payload, timestamp_ms)
1149 VALUES ($1, $2, $3, $4, $5, $6)",
1150 )?;
1151 for (seq, ty, step, payload, ts) in &rows {
1152 tx.execute(&ins, &[&id, seq, ty, step, payload, ts])?;
1153 }
1154 let (workflow, version, hash, status, created, updated, last_seq) = projection;
1155 tx.execute(
1156 "INSERT INTO workflow_executions (execution_id, workflow, version, definition_hash, status, created_at_ms, updated_at_ms, last_seq)
1157 VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
1158 ON CONFLICT (execution_id) DO UPDATE SET status = EXCLUDED.status, updated_at_ms = EXCLUDED.updated_at_ms, last_seq = EXCLUDED.last_seq",
1159 &[&id, &workflow, &version, &hash, &status, &created, &updated, &last_seq],
1160 )?;
1161 if !ext.done_tokens.is_empty() {
1162 tx.execute(
1163 "UPDATE orbit_tasks SET state = 'done', updated_ms = $2 WHERE token = ANY($1) AND state IN ('claimed', 'cancelled')",
1164 &[&ext.done_tokens, &now],
1165 )?;
1166 }
1167 if !ext.offers.is_empty() {
1168 let ins = tx.prepare(
1169 "INSERT INTO orbit_tasks (execution_id, unit_key, attempt, offer, task_name, contract_version, language, payload, enqueued_ms, state, updated_ms)
1170 VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, 'queued', $10)
1171 ON CONFLICT (execution_id, unit_key, attempt, offer) DO NOTHING",
1172 )?;
1173 for o in &ext.offers {
1174 tx.execute(
1175 &ins,
1176 &[&id, &o.unit_key, &o.attempt, &o.offer, &o.task_name, &o.contract_version, &o.language, &o.payload, &o.enqueued_ms, &now],
1177 )?;
1178 }
1179 }
1180 if !ext.syncs.is_empty() {
1181 let upd = tx.prepare(
1182 "UPDATE orbit_tasks SET state = 'cancelled', updated_ms = $5
1183 WHERE execution_id = $1 AND unit_key = $2 AND state IN ('queued', 'claimed') AND NOT (attempt = $3 AND offer = $4)",
1184 )?;
1185 for s in &ext.syncs {
1186 let (a, o) = s.keep.unwrap_or((-1, -1));
1187 tx.execute(&upd, &[&id, &s.unit_key, &a, &o, &now])?;
1188 }
1189 }
1190 for (seq, reply) in &ext.consumed {
1191 match reply {
1192 None => tx.execute("DELETE FROM orbit_inbox WHERE seq = $1", &[seq])?,
1193 Some(r) => tx.execute(
1194 "UPDATE orbit_inbox SET done_ms = $2, result = $3 WHERE seq = $1",
1195 &[seq, &now, r],
1196 )?,
1197 };
1198 }
1199 if let Some(reply) = &ext.finished_reply {
1200 tx.execute(
1201 "UPDATE orbit_tasks SET state = 'cancelled', updated_ms = $2 WHERE execution_id = $1 AND state IN ('queued', 'claimed')",
1202 &[&id, &now],
1203 )?;
1204 tx.execute("DELETE FROM orbit_owners WHERE execution_id = $1", &[&id])?;
1205 tx.execute(
1206 "UPDATE orbit_inbox SET done_ms = $2, result = $3 WHERE execution_id = $1 AND done_ms IS NULL AND kind IN ('signal', 'cancel')",
1207 &[&id, &now, reply],
1208 )?;
1209 tx.execute(
1210 "DELETE FROM orbit_inbox WHERE execution_id = $1 AND done_ms IS NULL",
1211 &[&id],
1212 )?;
1213 tx.execute("SELECT pg_notify('orbit_done', $1)", &[&id])?;
1214 }
1215 if !ext.offers.is_empty() {
1216 tx.execute("SELECT pg_notify('orbit_tasks', '')", &[])?;
1217 }
1218 tx.commit()?;
1219 Ok(())
1220 })
1221 }
1222}