use chrono::{Duration, Utc};
use sayiir_core::codec;
use sayiir_core::snapshot::{ExecutionPosition, WorkflowSnapshot, WorkflowSnapshotState};
use sayiir_core::task_claim::{AvailableTask, TaskClaim};
use sayiir_persistence::{BackendError, SnapshotStore, TaskClaimStore};
use sqlx::Row;
use crate::backend::PostgresBackend;
use crate::error::PgError;
const ELIGIBILITY_PREDICATE: &str = "NOT EXISTS (
SELECT 1 FROM sayiir_workflow_claims c
WHERE c.instance_id = s.instance_id
AND (c.expires_at IS NULL OR c.expires_at > now())
)
AND NOT EXISTS (
SELECT 1 FROM sayiir_workflow_signals sig
WHERE sig.instance_id = s.instance_id
)";
impl<C> TaskClaimStore for PostgresBackend<C>
where
C: codec::SnapshotCodec,
{
#[tracing::instrument(
name = "db.claim_task",
skip(self),
fields(db.system = "postgresql"),
err(level = tracing::Level::ERROR),
)]
async fn claim_task(
&self,
instance_id: &str,
task_id: &sayiir_core::TaskId,
worker_id: &str,
ttl: Option<Duration>,
) -> Result<Option<TaskClaim>, BackendError> {
let expires_at = ttl.and_then(|d| Utc::now().checked_add_signed(d));
let query = "
WITH probe AS (
SELECT instance_id, current_task_id, status
FROM sayiir_workflow_snapshots
WHERE instance_id = $1
),
lock AS (
SELECT pg_try_advisory_xact_lock(hashtextextended($1::text, 0)) AS got
),
upsert AS (
INSERT INTO sayiir_workflow_claims
(instance_id, task_id, worker_id, expires_at)
SELECT probe.instance_id, probe.current_task_id, $3, $4
FROM probe, lock
WHERE lock.got
AND probe.status = 'InProgress'
AND probe.current_task_id = $2
AND NOT EXISTS (
SELECT 1 FROM sayiir_workflow_claims c
WHERE c.instance_id = probe.instance_id
AND (c.expires_at IS NULL OR c.expires_at > now())
)
-- Re-execution race protection lives in the
-- worker''s validate_task_preconditions path, not
-- here: gating claim_task on a completed row in
-- workflow_tasks would permanently lock out
-- recovery when a worker dies BETWEEN
-- save_task_result (which writes the workflow_tasks
-- row) and save_snapshot (which advances
-- current_task_id). After that partial-completion,
-- every subsequent claim_task would reject and the
-- workflow would stall forever. The sleeping-giants
-- storm hits exactly this pattern under load.
-- Gate ON CONFLICT DO UPDATE on the existing claim
-- being expired. Without this, two transactions that
-- both pass the SELECT-side NOT EXISTS check (e.g.
-- because the prior claim's row just got deleted by
-- release, or because both see the same expired row)
-- can both INSERT, and ON CONFLICT silently picks the
-- last writer — leaving the earlier worker thinking
-- it owns a claim that the row actually attributes
-- to a different worker. The WHERE here turns the
-- upsert into 'steal expired only'; concurrent
-- inserts on a still-live claim now produce zero
-- upsert rows for the loser (caller sees Ok(None)
-- and retries), instead of misreporting success.
ON CONFLICT (instance_id) DO UPDATE
SET task_id = EXCLUDED.task_id,
worker_id = EXCLUDED.worker_id,
claimed_at = now(),
expires_at = EXCLUDED.expires_at
WHERE sayiir_workflow_claims.expires_at IS NOT NULL
AND sayiir_workflow_claims.expires_at <= now()
RETURNING
FLOOR(EXTRACT(EPOCH FROM claimed_at))::BIGINT AS claimed_epoch,
FLOOR(EXTRACT(EPOCH FROM expires_at))::BIGINT AS expires_epoch
)
SELECT claimed_epoch, expires_epoch, TRUE AS snapshot_exists
FROM upsert
UNION ALL
SELECT
NULL::bigint, NULL::bigint,
EXISTS (SELECT 1 FROM probe)
WHERE NOT EXISTS (SELECT 1 FROM upsert)
";
let row = sqlx::query(query)
.bind(instance_id)
.bind(task_id.as_bytes().as_slice())
.bind(worker_id)
.bind(expires_at)
.fetch_one(&self.pool)
.await
.map_err(PgError)?;
if let Some(claimed_epoch) = row.get::<Option<i64>, _>("claimed_epoch") {
return Ok(Some(TaskClaim {
instance_id: std::sync::Arc::from(instance_id),
task_id: *task_id,
worker_id: worker_id.to_string(),
claimed_at: claimed_epoch.cast_unsigned(),
expires_at: row
.get::<Option<i64>, _>("expires_epoch")
.map(i64::cast_unsigned),
}));
}
if !row.get::<bool, _>("snapshot_exists") {
return Err(BackendError::NotFound(format!("{instance_id}:{task_id}")));
}
Ok(None)
}
#[tracing::instrument(
name = "db.release_task_claim",
skip(self),
fields(db.system = "postgresql"),
err(level = tracing::Level::ERROR),
)]
async fn release_task_claim(
&self,
instance_id: &str,
task_id: &sayiir_core::TaskId,
worker_id: &str,
) -> Result<(), BackendError> {
let row = sqlx::query(
"WITH del AS (
DELETE FROM sayiir_workflow_claims
WHERE instance_id = $1 AND worker_id = $2 AND task_id = $3
RETURNING 1
),
probe AS (
SELECT worker_id, task_id
FROM sayiir_workflow_claims
WHERE instance_id = $1 AND NOT EXISTS (SELECT 1 FROM del)
)
SELECT
EXISTS (SELECT 1 FROM del) AS done,
(SELECT worker_id FROM probe) AS owner,
(SELECT task_id FROM probe) AS current_task_id",
)
.bind(instance_id)
.bind(worker_id)
.bind(task_id.as_bytes().as_slice())
.fetch_one(&self.pool)
.await
.map_err(PgError)?;
disambiguate_writeback(&row, instance_id, task_id, worker_id)
}
#[tracing::instrument(
name = "db.extend_task_claim",
skip(self),
fields(db.system = "postgresql"),
err(level = tracing::Level::ERROR),
)]
async fn extend_task_claim(
&self,
instance_id: &str,
task_id: &sayiir_core::TaskId,
worker_id: &str,
additional_duration: Duration,
) -> Result<(), BackendError> {
let row = sqlx::query(
"WITH upd AS (
UPDATE sayiir_workflow_claims
SET expires_at = expires_at + $3
WHERE instance_id = $1
AND worker_id = $2
AND task_id = $4
AND expires_at IS NOT NULL
RETURNING 1
),
probe AS (
SELECT worker_id, task_id
FROM sayiir_workflow_claims
WHERE instance_id = $1 AND NOT EXISTS (SELECT 1 FROM upd)
)
SELECT
EXISTS (SELECT 1 FROM upd) AS done,
(SELECT worker_id FROM probe) AS owner,
(SELECT task_id FROM probe) AS current_task_id",
)
.bind(instance_id)
.bind(worker_id)
.bind(additional_duration)
.bind(task_id.as_bytes().as_slice())
.fetch_one(&self.pool)
.await
.map_err(PgError)?;
disambiguate_writeback(&row, instance_id, task_id, worker_id)
}
#[tracing::instrument(
name = "db.find_available_tasks",
skip(self),
fields(db.system = "postgresql"),
err(level = tracing::Level::ERROR),
)]
#[allow(clippy::cast_precision_loss, clippy::too_many_lines)]
async fn find_available_tasks(
&self,
worker_id: &str,
limit: usize,
aging_interval: Duration,
worker_tags: &[String],
) -> Result<Vec<AvailableTask>, BackendError> {
let aging_secs = (aging_interval.num_milliseconds() as f64 / 1000.0).max(1.0);
let worker_tags_vec: Vec<&str> = worker_tags.iter().map(String::as_str).collect();
let tag_filter = if worker_tags.is_empty() {
""
} else {
"AND s.task_tags <@ $3"
};
let query = format!(
"SELECT r.instance_id, h.data, r.trace_parent
FROM (
SELECT s.instance_id, s.history_version, s.trace_parent
FROM sayiir_workflow_snapshots s
WHERE s.status = 'InProgress'
AND (
s.position_kind = 'AtTask'
OR (s.position_kind = 'AtDelay'
AND s.delay_wake_at IS NOT NULL
AND s.delay_wake_at <= now())
)
AND {ELIGIBILITY_PREDICATE}
{tag_filter}
ORDER BY
(s.task_priority - EXTRACT(EPOCH FROM (now() - s.updated_at)) / $2) ASC,
s.updated_at ASC
LIMIT $1
FOR UPDATE OF s SKIP LOCKED
) r
JOIN sayiir_workflow_snapshot_history h
ON h.instance_id = r.instance_id AND h.version = r.history_version"
);
let mut q = sqlx::query(&query)
.bind(i64::try_from(limit).unwrap_or(i64::MAX))
.bind(aging_secs);
if !worker_tags.is_empty() {
q = q.bind(&worker_tags_vec);
}
let rows = q.fetch_all(&self.pool).await.map_err(PgError)?;
let mut decoded: Vec<WorkflowSnapshot> = Vec::with_capacity(rows.len());
for row in &rows {
let raw: &[u8] = row.get("data");
let mut snapshot = self.decode(raw)?;
snapshot.trace_parent = row.get("trace_parent");
decoded.push(snapshot);
}
let instance_id_strs: Vec<&str> = decoded.iter().map(|s| s.instance_id.as_ref()).collect();
let mut outputs_by_instance =
crate::history::fetch_task_outputs_batched(&self.pool, &instance_id_strs).await?;
for snapshot in &mut decoded {
if let Some(outputs) = outputs_by_instance.remove(snapshot.instance_id.as_ref()) {
snapshot.hydrate_task_outputs(outputs);
}
}
let mut available: Vec<(bool, AvailableTask)> = Vec::with_capacity(decoded.len());
for mut snapshot in decoded {
let task_id_to_dispatch: sayiir_core::TaskId = match &snapshot.state {
WorkflowSnapshotState::InProgress {
position:
ExecutionPosition::AtDelay {
wake_at,
next_task_id,
delay_id,
..
},
..
} if Utc::now() >= *wake_at => {
let next_id_opt = *next_task_id;
let delay_id_owned = *delay_id;
if let Some(next_id) = next_id_opt {
snapshot.update_position(ExecutionPosition::AtTask { task_id: next_id });
self.save_snapshot(&mut snapshot).await?;
next_id
} else {
let output = snapshot
.get_task_result_bytes(&delay_id_owned)
.unwrap_or_default();
snapshot.mark_completed(output);
self.save_snapshot(&mut snapshot).await?;
continue;
}
}
WorkflowSnapshotState::InProgress {
position: ExecutionPosition::AtTask { task_id },
completed_tasks,
..
} => {
if completed_tasks.contains_key(task_id) {
continue;
}
let task_id_copy = *task_id;
if let Some(rs) = snapshot.task_retries.get(&task_id_copy)
&& Utc::now() < rs.next_retry_at
{
continue;
}
task_id_copy
}
_ => continue,
};
let bias = snapshot.has_failed_on_worker(&task_id_to_dispatch, worker_id);
if let Some(task) = build_available_task(snapshot, task_id_to_dispatch) {
available.push((bias, task));
}
if available.len() >= limit {
break;
}
}
available.sort_by_key(|(bias, _)| *bias);
tracing::debug!(count = available.len(), "available tasks found");
Ok(available.into_iter().map(|(_, task)| task).collect())
}
async fn wait_for_wakeup(
&self,
timeout: std::time::Duration,
) -> Result<Option<sayiir_persistence::TaskWakeupHint>, BackendError> {
Ok(self.wakeup.wait(timeout).await)
}
#[tracing::instrument(
name = "db.find_hinted_task",
skip(self, hint),
fields(
db.system = "postgresql",
instance_id = %hint.instance_id,
task_id = %sayiir_core::TaskId::from_bytes(hint.task_id),
),
err(level = tracing::Level::ERROR),
)]
async fn find_hinted_task(
&self,
hint: &sayiir_persistence::TaskWakeupHint,
) -> Result<Option<AvailableTask>, BackendError> {
let query = format!(
"SELECT
s.trace_parent,
h.data,
(SELECT array_agg(t.task_id ORDER BY t.task_id)
FROM sayiir_workflow_tasks t
WHERE t.instance_id = s.instance_id
AND t.status = 'completed'
AND t.output IS NOT NULL) AS task_ids,
(SELECT array_agg(t.output ORDER BY t.task_id)
FROM sayiir_workflow_tasks t
WHERE t.instance_id = s.instance_id
AND t.status = 'completed'
AND t.output IS NOT NULL) AS task_outputs
FROM sayiir_workflow_snapshots s
JOIN sayiir_workflow_snapshot_history h
ON h.instance_id = s.instance_id AND h.version = s.history_version
WHERE s.instance_id = $1
AND s.current_task_id = $2
AND s.status = 'InProgress'
AND {ELIGIBILITY_PREDICATE}"
);
let row = sqlx::query(&query)
.bind(&hint.instance_id)
.bind(hint.task_id.as_slice())
.fetch_optional(&self.pool)
.await
.map_err(PgError)?;
let Some(row) = row else { return Ok(None) };
let data: Vec<u8> = row.get("data");
let mut snapshot = self.decode(&data)?;
snapshot.trace_parent = row.get("trace_parent");
let task_ids: Option<Vec<Vec<u8>>> = row.get("task_ids");
let task_outputs: Option<Vec<Vec<u8>>> = row.get("task_outputs");
if let (Some(ids), Some(outs)) = (task_ids, task_outputs) {
let mut outputs = Vec::with_capacity(ids.len());
for (id_bytes, out_bytes) in ids.into_iter().zip(outs.into_iter()) {
let tid = sayiir_core::TaskId::from_slice(&id_bytes).map_err(|_| {
BackendError::Backend(format!(
"sayiir_workflow_tasks.task_id for instance {} is {} bytes (expected 32)",
hint.instance_id,
id_bytes.len()
))
})?;
outputs.push((tid, bytes::Bytes::from(out_bytes)));
}
snapshot.hydrate_task_outputs(outputs);
}
let hint_task_id = sayiir_core::TaskId::from_bytes(hint.task_id);
let matches = match &snapshot.state {
WorkflowSnapshotState::InProgress {
position: ExecutionPosition::AtTask { task_id },
completed_tasks,
..
} => *task_id == hint_task_id && !completed_tasks.contains_key(task_id),
_ => false,
};
if !matches {
return Ok(None);
}
Ok(build_available_task(snapshot, hint_task_id))
}
}
fn disambiguate_writeback(
row: &sqlx::postgres::PgRow,
instance_id: &str,
task_id: &sayiir_core::TaskId,
worker_id: &str,
) -> Result<(), BackendError> {
if row.get::<bool, _>("done") {
return Ok(());
}
let owner: Option<String> = row.get("owner");
let Some(owner) = owner else {
return Err(BackendError::NotFound(format!("{instance_id}:{task_id}")));
};
if owner != worker_id {
return Err(BackendError::Backend(format!(
"Claim owned by different worker: {owner}"
)));
}
let current: Option<Vec<u8>> = row.get("current_task_id");
match current {
Some(bytes) if bytes.as_slice() == task_id.as_bytes().as_slice() => Ok(()),
_ => Err(BackendError::NotFound(format!("{instance_id}:{task_id}"))),
}
}
fn build_available_task(
snapshot: WorkflowSnapshot,
task_id: sayiir_core::TaskId,
) -> Option<AvailableTask> {
let completed_empty = match &snapshot.state {
sayiir_core::snapshot::WorkflowSnapshotState::InProgress {
completed_tasks, ..
} => completed_tasks.is_empty(),
_ => return None,
};
let input = if completed_empty {
snapshot.initial_input_bytes()
} else {
snapshot.get_last_task_output()
};
input.map(|input_bytes| AvailableTask {
instance_id: snapshot.instance_id.clone(),
task_id,
input: input_bytes,
workflow_definition_hash: snapshot.definition_hash,
trace_parent: snapshot.trace_parent.clone(),
snapshot: std::sync::Arc::new(snapshot),
})
}