mod listener;
mod notifier;
use std::collections::{HashMap, HashSet};
use async_trait::async_trait;
use sqlx::postgres::{PgDatabaseError, PgPoolOptions, PgSeverity};
use sqlx::{AssertSqlSafe, PgPool, Row};
use self::listener::Listener;
use self::notifier::Notifier;
use super::PARTITIONED_DEQUEUE_SWEEP_CAP;
use super::migrations::{self, quote_identifier};
use super::notify::{
EVENTS_CHANNEL, Registry, STREAMS_CHANNEL, event_key, message_key, stream_key,
};
use super::retry::{RetryPolicy, with_retry};
use std::sync::{Arc, OnceLock};
use std::time::Duration;
use super::types::step_names;
use super::types::{
AppScope, ApplicationRowCounts, AttributeUpdate, AwaitedOutcome, AwaitedStream, Change,
DatabaseTime, Debounce, DebounceHolder, DebounceRequest, DequeueStop, Dequeued, EncodedValue,
EventRecord, Fork, ForkOptions, ForkPoint, GetEventCaller, IdleTransactionTimeout,
InitWorkflowCaller, Message, NewQueue, NewSchedule, NewWorkflow, NotificationRecord,
OnExistingQueue, Outcome, PollingConcurrency, QueueRecord, QueueUpdate, RateLimit,
RenameBatching, RenameFrom, RewindOptions, ScheduleFilter, ScheduleRecord, ScheduleStatus,
ScheduleUpdate, StepCaller, StepCodec, StepRecord, StepTiming, StreamRead, StreamRecord,
Timestamp, UpdatedQueue, UpsertedQueue, VersionInfo, WakeTime, WorkflowDelay, WorkflowFilter,
WorkflowIdReusePolicy, WorkflowOwner, WorkflowRecord, WorkflowStatus, WrittenBy,
decode_step_value, duration_from_ms, duration_from_secs, duration_to_ms, encode_step_value,
is_valid_application_name,
};
use super::{
BackendError, BackendErrorKind, DEFAULT_SCHEMA, Error, INTERNAL_QUEUE, NULL_TOPIC,
OutcomeWrite, PORTABLE_JSON, STREAM_CLOSED, SystemDatabase, WorkflowInitResult,
is_stream_closed,
};
impl From<sqlx::Error> for Error {
fn from(e: sqlx::Error) -> Self {
let sqlstate = e
.as_database_error()
.and_then(|d| d.code())
.map(|c| c.into_owned());
Error::Backend(BackendError {
kind: classify(&e, sqlstate.as_deref()),
message: e.to_string(),
sqlstate,
})
}
}
fn is_unique_violation(error: &sqlx::Error) -> bool {
error
.as_database_error()
.and_then(|e| e.code())
.is_some_and(|code| code == "23505")
}
fn is_foreign_key_violation(error: &sqlx::Error) -> bool {
error
.as_database_error()
.and_then(|e| e.code())
.is_some_and(|code| code == "23503")
}
fn empty_to_none(value: Option<&str>) -> Option<&str> {
value.filter(|v| !v.is_empty())
}
fn classify(error: &sqlx::Error, sqlstate: Option<&str>) -> BackendErrorKind {
if let Some(code) = sqlstate {
let database = error.as_database_error();
let severity = database
.and_then(|e| e.try_downcast_ref::<PgDatabaseError>())
.map(PgDatabaseError::severity);
return classify_sqlstate(code, severity, database.map(|e| e.message()));
}
match error {
sqlx::Error::Io(_)
| sqlx::Error::Tls(_)
| sqlx::Error::PoolTimedOut
| sqlx::Error::WorkerCrashed => BackendErrorKind::Connection,
sqlx::Error::PoolClosed => BackendErrorKind::Permanent,
_ => BackendErrorKind::Permanent,
}
}
const COCKROACH_SESSION_CANCELED: &str = "query execution canceled";
fn classify_sqlstate(
code: &str,
severity: Option<PgSeverity>,
message: Option<&str>,
) -> BackendErrorKind {
if matches!(code, "40001" | "40P01") {
return BackendErrorKind::Conflict;
}
if code == "57P04" {
return BackendErrorKind::Permanent;
}
if code == "57014" {
return match message {
Some(message) if message.trim() == COCKROACH_SESSION_CANCELED => {
BackendErrorKind::Connection
}
_ => BackendErrorKind::Permanent,
};
}
if code == "08P01" {
return match severity {
Some(PgSeverity::Fatal | PgSeverity::Panic) => BackendErrorKind::Connection,
_ => BackendErrorKind::Permanent,
};
}
if code == "25P03" {
return BackendErrorKind::Connection;
}
match &code[..2.min(code.len())] {
"08" | "53" | "57" => BackendErrorKind::Connection,
"XX" => match message {
Some(message) if is_transport_failure(message) => BackendErrorKind::Connection,
_ => BackendErrorKind::Permanent,
},
_ => BackendErrorKind::Permanent,
}
}
fn is_transport_failure(message: &str) -> bool {
const NEEDLES: [&str; 5] = [
"i/o timeout",
"broken pipe",
"connection reset by peer",
"connection refused",
"use of closed network connection",
];
let message = message.to_ascii_lowercase();
NEEDLES.iter().any(|needle| message.contains(needle))
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Settings<'a> {
pub schema: &'a str,
pub retry: RetryPolicy,
pub executor_id: Option<&'a str>,
pub application_name: Option<&'a str>,
pub polling_concurrency: PollingConcurrency,
pub notification_coalesce: Option<Duration>,
}
impl Default for Settings<'_> {
fn default() -> Self {
Self {
schema: DEFAULT_SCHEMA,
retry: RetryPolicy::default(),
executor_id: None,
application_name: None,
polling_concurrency: PollingConcurrency::Default,
notification_coalesce: None,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Config<'a> {
pub url: &'a str,
pub max_connections: u32,
pub use_listen_notify: bool,
pub idle_transaction_timeout: IdleTransactionTimeout,
pub migrate: bool,
pub settings: Settings<'a>,
}
impl<'a> Config<'a> {
pub fn new(url: &'a str) -> Self {
Self {
url,
max_connections: 10,
use_listen_notify: true,
idle_transaction_timeout: IdleTransactionTimeout::Default,
migrate: true,
settings: Settings::default(),
}
}
}
fn options_set_idle_transaction_timeout(url: &str) -> bool {
url.parse::<sqlx::postgres::PgConnectOptions>()
.ok()
.and_then(|options| {
options
.get_options()
.map(|options| option_names(options).any(|name| name == IDLE_TIMEOUT_PARAMETER))
})
.unwrap_or(false)
}
const IDLE_TIMEOUT_PARAMETER: &str = "idle_in_transaction_session_timeout";
fn option_names(options: &str) -> impl Iterator<Item = String> + '_ {
let mut words = Vec::new();
let mut word = String::new();
let mut chars = options.chars();
while let Some(c) = chars.next() {
match c {
'\\' => word.extend(chars.next()),
c if c.is_whitespace() => {
if !word.is_empty() {
words.push(std::mem::take(&mut word));
}
}
c => word.push(c),
}
}
if !word.is_empty() {
words.push(word);
}
let mut words = words.into_iter();
std::iter::from_fn(move || {
loop {
let word = words.next()?;
let assignment = if word == "-c" {
words.next()?
} else if let Some(rest) = word.strip_prefix("--") {
rest.to_owned()
} else if let Some(rest) = word.strip_prefix("-c") {
rest.to_owned()
} else {
continue;
};
if let Some((name, _)) = assignment.split_once('=') {
return Some(name.to_ascii_lowercase().replace('-', "_"));
}
}
})
}
pub struct PostgresSystemDatabase {
pool: PgPool,
tables: Tables,
retry: RetryPolicy,
executor_id: Option<String>,
application_name: Option<String>,
polling: tokio::sync::Semaphore,
notify: Arc<Registry>,
listener: Arc<Listener>,
listener_task: std::sync::Mutex<Option<crate::rt::JoinHandle<()>>>,
notifier: Arc<Notifier>,
notifier_task: std::sync::Mutex<Option<crate::rt::JoinHandle<()>>>,
}
impl PostgresSystemDatabase {
pub async fn connect(config: &Config<'_>) -> Result<Self, Error> {
config.idle_transaction_timeout.validate()?;
if config.migrate {
ensure_database_exists(config.url).await?;
}
let mut pool = PgPoolOptions::new().max_connections(config.max_connections);
if let Some(millis) = config.idle_transaction_timeout.millis()
&& !options_set_idle_transaction_timeout(config.url)
{
pool = pool.after_connect(move |conn, _| {
Box::pin(async move {
sqlx::query(AssertSqlSafe(format!(
"SET idle_in_transaction_session_timeout = {millis}"
)))
.execute(conn)
.await?;
Ok(())
})
});
}
let pool = pool.connect(config.url).await?;
let prepared = if config.migrate {
migrations::runner::run(&pool, config.settings.schema, config.use_listen_notify)
.await
.map(|_| ())
} else {
migrations::runner::verify(&pool, config.settings.schema).await
};
prepared.map_err(|e| {
Error::Backend(BackendError {
message: e.to_string(),
sqlstate: None,
kind: BackendErrorKind::Permanent,
})
})?;
let notify: Arc<Registry> = Arc::default();
let handle = Self {
polling: tokio::sync::Semaphore::new(
config
.settings
.polling_concurrency
.permits(pool.options().get_max_connections()),
),
listener: Arc::new(Listener::new(pool.clone(), Arc::clone(¬ify))),
notifier: Arc::new(Notifier::new(
pool.clone(),
Arc::clone(¬ify),
config.settings.notification_coalesce,
)),
pool,
notify,
listener_task: std::sync::Mutex::default(),
notifier_task: std::sync::Mutex::default(),
tables: Tables::new(config.settings.schema),
retry: config.settings.retry,
executor_id: config.settings.executor_id.map(str::to_owned),
application_name: config.settings.application_name.map(str::to_owned),
};
if config.use_listen_notify {
handle.start_notifications().await;
}
Ok(handle)
}
pub async fn start_notifications(&self) -> bool {
match migrations::runner::detect_dialect(&self.pool).await {
Ok(migrations::Dialect::Cockroach) => {
tracing::debug!(
"not starting a notification listener: CockroachDB has no LISTEN/NOTIFY, so \
waits re-query on their own interval, which is what delivers there"
);
return false;
}
Ok(migrations::Dialect::Postgres) => {}
Err(e) => {
tracing::warn!(error = %e, "could not start a notification listener");
return false;
}
}
let mut listening = self.listener_task.lock().expect("listener lock");
if listening.as_ref().is_none_or(|task| task.is_finished()) {
*listening = Some(crate::rt::spawn(Arc::clone(&self.listener).run()));
}
let mut notifying = self.notifier_task.lock().expect("notifier lock");
if notifying.as_ref().is_none_or(|task| task.is_finished()) {
self.notifier.enable();
*notifying = Some(crate::rt::spawn(Arc::clone(&self.notifier).run()));
}
true
}
pub fn is_pushing(&self) -> bool {
self.notifier.is_pushing()
}
pub fn is_delivering(&self) -> bool {
self.listener.is_delivering()
}
pub fn from_pool(pool: PgPool, settings: &Settings<'_>) -> Self {
let notify: Arc<Registry> = Arc::default();
Self {
polling: tokio::sync::Semaphore::new(
settings
.polling_concurrency
.permits(pool.options().get_max_connections()),
),
listener: Arc::new(Listener::new(pool.clone(), Arc::clone(¬ify))),
notifier: Arc::new(Notifier::new(
pool.clone(),
Arc::clone(¬ify),
settings.notification_coalesce,
)),
pool,
notify,
listener_task: std::sync::Mutex::default(),
notifier_task: std::sync::Mutex::default(),
tables: Tables::new(settings.schema),
retry: settings.retry,
executor_id: settings.executor_id.map(str::to_owned),
application_name: settings.application_name.map(str::to_owned),
}
}
}
impl Drop for PostgresSystemDatabase {
fn drop(&mut self) {
self.notifier.stop();
let task = self.listener_task.lock().expect("listener lock").take();
if let Some(task) = task {
task.abort();
}
}
}
fn rename_source_predicate(source: RenameFrom<'_>, param: usize) -> String {
match source {
RenameFrom::Application(_) => format!("application_name = ${param}"),
RenameFrom::ApplicationAndUnclaimed(_) => {
format!("(application_name = ${param} OR application_name IS NULL)")
}
RenameFrom::Unclaimed => format!("(${param}::text IS NULL AND application_name IS NULL)"),
}
}
#[derive(Debug, Clone, Copy)]
enum StartSteps<'a> {
Given(&'a [i32]),
Resolve {
aggregate: &'a str,
named: Option<&'a str>,
},
}
enum StepConn {
Pooled(sqlx::pool::PoolConnection<sqlx::Postgres>),
Transactional(Box<sqlx::Transaction<'static, sqlx::Postgres>>),
}
impl std::ops::Deref for StepConn {
type Target = sqlx::PgConnection;
fn deref(&self) -> &Self::Target {
match self {
StepConn::Pooled(conn) => conn,
StepConn::Transactional(tx) => tx,
}
}
}
impl std::ops::DerefMut for StepConn {
fn deref_mut(&mut self) -> &mut Self::Target {
match self {
StepConn::Pooled(conn) => conn,
StepConn::Transactional(tx) => tx,
}
}
}
fn replayed_output<T: serde::de::DeserializeOwned>(
step: &StepRecord,
caller: StepCaller<'_>,
step_name: &str,
) -> Result<T, Error> {
let StepCaller {
workflow_id,
step_id,
..
} = caller;
tracing::debug!(workflow_id, step_id, step_name, "replaying a step");
if let Some(recorded) = step.error.as_deref() {
return Err(decode_step_value(
step, recorded, caller, step_name, "an error",
)?);
}
let recorded = step.output.as_deref().ok_or_else(|| {
Error::Malformed(format!(
"workflow {workflow_id} step {step_id} ({step_name}) has no recorded output"
))
})?;
decode_step_value(step, recorded, caller, step_name, "an output")
}
fn refuse_self_target(caller_id: &str, targets: &[&str], action: &str) -> Result<(), Error> {
if targets.contains(&caller_id) {
return Err(Error::InvalidInput {
field: "workflow_ids".into(),
detail: format!("workflow {caller_id} cannot {action} itself"),
});
}
Ok(())
}
impl PostgresSystemDatabase {
async fn root_containing(
&self,
conn: &mut sqlx::PgConnection,
caller_id: &str,
roots: &[&str],
) -> Result<String, Error> {
let workflow_table = self.tables.workflow_status.as_str();
let mut seen = std::collections::HashSet::from([caller_id.to_owned()]);
let mut current = caller_id.to_owned();
loop {
let parent: Option<Option<String>> = sqlx::query_scalar(AssertSqlSafe(format!(
"SELECT parent_workflow_id FROM {workflow_table} WHERE workflow_uuid = $1"
)))
.bind(¤t)
.fetch_optional(&mut *conn)
.await?;
match parent.flatten() {
Some(parent) if roots.contains(&parent.as_str()) => return Ok(parent),
Some(parent) if seen.insert(parent.clone()) => current = parent,
_ => return Ok(roots.join(", ")),
}
}
}
async fn run_transactional_step<T, F, Fut>(
&self,
caller: Option<StepCaller<'_>>,
step_name: &str,
started_at: Timestamp,
work: F,
) -> Result<T, Error>
where
T: serde::Serialize + serde::de::DeserializeOwned,
F: Fn(sqlx::Transaction<'static, sqlx::Postgres>) -> Fut + Send + Sync,
Fut: Future<Output = Result<(sqlx::Transaction<'static, sqlx::Postgres>, T), Error>> + Send,
{
let mut tx = self.pool.begin().await?;
if let Some(caller) = caller {
self.check_owner_on(&mut tx, caller.owner()).await?;
if let Some(step) = self
.check_step_on(&mut tx, caller.workflow_id, caller.step_id, step_name)
.await?
{
tx.commit().await?;
return replayed_output(&step, caller, step_name);
}
}
let (mut tx, value) = match work(tx).await {
Ok(done) => done,
Err(error) => {
return match caller {
Some(caller) => Err(self
.record_step_failure(caller, step_name, error, started_at)
.await),
None => Err(error),
};
}
};
if let Some(caller) = caller {
self.record_step_output(&mut tx, caller, step_name, &value, started_at)
.await?;
}
tx.commit().await?;
Ok(value)
}
async fn record_child_workflow_on(
&self,
conn: &mut sqlx::PgConnection,
parent_workflow_id: &str,
child_workflow_id: &str,
step_id: i32,
step_name: &str,
started_at: Option<Timestamp>,
) -> Result<(), Error> {
if child_workflow_id.is_empty() {
return Err(Error::InvalidInput {
field: "child_workflow_id".into(),
detail: "must not be empty".to_owned(),
});
}
let steps_table = self.tables.operation_outputs.as_str();
let completed_at = started_at.map(|_| Timestamp::now());
let application_name = self.application_name.as_deref();
let stored: Option<Option<String>> = sqlx::query_scalar(AssertSqlSafe(format!(
"INSERT INTO {steps_table} (workflow_uuid, function_id, function_name, \
child_workflow_id, started_at_epoch_ms, completed_at_epoch_ms, \
application_name) \
VALUES ($1, $2, $3, $4, $5, $6, $7) \
ON CONFLICT (workflow_uuid, function_id) DO UPDATE \
SET child_workflow_id = {steps_table}.child_workflow_id \
RETURNING child_workflow_id"
)))
.bind(parent_workflow_id)
.bind(step_id)
.bind(step_name)
.bind(child_workflow_id)
.bind(started_at.map(Timestamp::as_epoch_ms))
.bind(completed_at.map(Timestamp::as_epoch_ms))
.bind(application_name)
.fetch_optional(&mut *conn)
.await?;
if let Some(stored) = stored
&& stored.as_deref() != Some(child_workflow_id)
{
return Err(Error::StepNondeterminism {
workflow_id: parent_workflow_id.to_owned(),
step_id,
});
}
Ok(())
}
async fn record_step_output<T: serde::Serialize>(
&self,
conn: &mut sqlx::PgConnection,
caller: StepCaller<'_>,
step_name: &str,
value: &T,
started_at: Timestamp,
) -> Result<(), Error> {
let recorded = encode_step_value(value, caller.codec, "an output")?;
self.record_step_on(
conn,
caller.workflow_id,
caller.step_id,
step_name,
Outcome::Output(Some(&recorded.value)),
recorded.serialization.as_deref(),
Some(StepTiming {
started_at,
completed_at: Timestamp::now(),
}),
None,
)
.await
}
async fn record_step_failure(
&self,
caller: StepCaller<'_>,
step_name: &str,
error: Error,
started_at: Timestamp,
) -> Error {
if !error.should_record() {
return error;
}
let StepCaller {
workflow_id,
step_id,
codec,
..
} = caller;
let recorded = match encode_step_value(&error, codec, "an error") {
Ok(recorded) => recorded,
Err(encode_error) => {
tracing::warn!(
workflow_id,
step_id,
step_name,
error = %error,
encode_error = %encode_error,
"the step failed, but its error could not be encoded to record it"
);
return encode_error;
}
};
let record = async {
let mut tx = self.pool.begin().await?;
self.check_owner_on(&mut tx, caller.owner()).await?;
self.record_step_on(
&mut tx,
workflow_id,
step_id,
step_name,
Outcome::Error(&recorded.value),
recorded.serialization.as_deref(),
Some(StepTiming {
started_at,
completed_at: Timestamp::now(),
}),
None,
)
.await?;
tx.commit().await?;
Ok::<(), Error>(())
};
match record.await {
Ok(()) => {
tracing::debug!(
workflow_id,
step_id,
step_name,
error = %error,
"the step failed; its error is recorded"
);
error
}
Err(record_error) => record_error,
}
}
async fn run_single_statement_step<T, F, Fut>(
&self,
caller: Option<StepCaller<'_>>,
step_name: &str,
started_at: Timestamp,
work: F,
) -> Result<T, Error>
where
T: serde::Serialize + serde::de::DeserializeOwned,
F: Fn(StepConn) -> Fut + Send + Sync,
Fut: Future<Output = Result<(StepConn, T), Error>> + Send,
{
let Some(caller) = caller else {
let (_conn, value) = work(StepConn::Pooled(self.pool.acquire().await?)).await?;
return Ok(value);
};
let mut conn = StepConn::Transactional(Box::new(self.pool.begin().await?));
self.check_owner_on(&mut conn, caller.owner()).await?;
if let Some(step) = self
.check_step_on(&mut conn, caller.workflow_id, caller.step_id, step_name)
.await?
{
return replayed_output(&step, caller, step_name);
}
let (mut conn, value) = match work(conn).await {
Ok(done) => done,
Err(error) => {
return Err(self
.record_step_failure(caller, step_name, error, started_at)
.await);
}
};
self.record_step_output(&mut conn, caller, step_name, &value, started_at)
.await?;
match conn {
StepConn::Transactional(tx) => (*tx).commit().await?,
StepConn::Pooled(_) => {}
}
Ok(value)
}
async fn fork_on(
&self,
mut tx: sqlx::Transaction<'static, sqlx::Postgres>,
original_ids: &[&str],
forked_ids: &[String],
steps: StartSteps<'_>,
options: &ForkOptions<'_>,
) -> Result<sqlx::Transaction<'static, sqlx::Postgres>, Error> {
let resolved: Vec<i32>;
let start_steps: &[i32] = match steps {
StartSteps::Given(given) => given,
StartSteps::Resolve { aggregate, named } => {
resolved = self
.resolve_fork_points_on(&mut tx, original_ids, aggregate, named)
.await?;
&resolved
}
};
let copies_anything = start_steps.iter().any(|&step| step > 0);
let queue_name = options.queue_name.unwrap_or(INTERNAL_QUEUE);
let timeout_ms = options.timeout_ms()?;
let (replace_from, replace_to): (Vec<&str>, Vec<&str>) =
options.replacement_children.iter().copied().unzip();
let (replace_from, replace_to) = (&replace_from[..], &replace_to[..]);
let workflow_table = self.tables.workflow_status.as_str();
let input_table = self.tables.workflow_input.as_str();
let steps_table = self.tables.operation_outputs.as_str();
let events_table = self.tables.workflow_events.as_str();
let history_table = self.tables.workflow_events_history.as_str();
let streams_table = self.tables.streams.as_str();
let application_name = self.application_name.as_deref();
let found: Vec<String> = sqlx::query_scalar(AssertSqlSafe(format!(
"SELECT workflow_uuid FROM {workflow_table} WHERE workflow_uuid = ANY($1)"
)))
.bind(original_ids)
.fetch_all(&mut *tx)
.await?;
let missing: Vec<String> = original_ids
.iter()
.filter(|id| !found.iter().any(|f| f == *id))
.map(|id| (*id).to_owned())
.collect();
if !missing.is_empty() {
return Err(Error::NonExistentWorkflow {
workflow_ids: missing,
});
}
sqlx::query(AssertSqlSafe(format!(
"INSERT INTO {workflow_table} (workflow_uuid, status, name, class_name, config_name, \
application_version, application_id, authenticated_user, authenticated_roles, \
assumed_role, serialization, request, queue_name, \
queue_partition_key, forked_from, attributes, workflow_timeout_ms, \
application_name) \
SELECT m.fork_id, 'ENQUEUED', w.name, w.class_name, w.config_name, \
COALESCE($4, w.application_version), w.application_id, w.authenticated_user, \
w.authenticated_roles, w.assumed_role, w.serialization, w.request, \
$5, $6, w.workflow_uuid, w.attributes, $7, \
COALESCE(w.application_name, $8) \
FROM unnest($1::text[], $2::text[], $3::int4[]) \
AS m(original_id, fork_id, start_step) \
JOIN {workflow_table} w ON w.workflow_uuid = m.original_id"
)))
.bind(original_ids)
.bind(forked_ids)
.bind(start_steps)
.bind(options.application_version)
.bind(queue_name)
.bind(options.queue_partition_key)
.bind(timeout_ms)
.bind(application_name)
.execute(&mut *tx)
.await?;
sqlx::query(AssertSqlSafe(format!(
"INSERT INTO {input_table} (workflow_uuid, inputs) \
SELECT m.fork_id, COALESCE(i.inputs, w.inputs) \
FROM unnest($1::text[], $2::text[]) AS m(original_id, fork_id) \
JOIN {workflow_table} w ON w.workflow_uuid = m.original_id \
LEFT JOIN {input_table} i ON i.workflow_uuid = m.original_id \
ON CONFLICT (workflow_uuid) DO UPDATE SET inputs = EXCLUDED.inputs"
)))
.bind(original_ids)
.bind(forked_ids)
.execute(&mut *tx)
.await?;
sqlx::query(AssertSqlSafe(format!(
"UPDATE {workflow_table} SET was_forked_from = TRUE WHERE workflow_uuid = ANY($1)"
)))
.bind(original_ids)
.execute(&mut *tx)
.await?;
if copies_anything {
sqlx::query(AssertSqlSafe(format!(
"INSERT INTO {steps_table} (workflow_uuid, function_id, output, error, \
serialization, function_name, child_workflow_id, started_at_epoch_ms, \
completed_at_epoch_ms, application_name) \
SELECT m.fork_id, o.function_id, o.output, o.error, o.serialization, \
o.function_name, \
COALESCE(r.replacement, o.child_workflow_id), \
o.started_at_epoch_ms, o.completed_at_epoch_ms, \
COALESCE(w.application_name, $6) \
FROM unnest($1::text[], $2::text[], $3::int4[]) \
AS m(original_id, fork_id, start_step) \
JOIN {steps_table} o \
ON o.workflow_uuid = m.original_id AND o.function_id < m.start_step \
JOIN {workflow_table} w ON w.workflow_uuid = m.original_id \
LEFT JOIN unnest($4::text[], $5::text[]) AS r(original, replacement) \
ON r.original = o.child_workflow_id"
)))
.bind(original_ids)
.bind(forked_ids)
.bind(start_steps)
.bind(replace_from)
.bind(replace_to)
.bind(application_name)
.execute(&mut *tx)
.await?;
sqlx::query(AssertSqlSafe(format!(
"INSERT INTO {history_table} (workflow_uuid, function_id, key, value, serialization) \
SELECT m.fork_id, h.function_id, h.key, h.value, h.serialization \
FROM unnest($1::text[], $2::text[], $3::int4[]) \
AS m(original_id, fork_id, start_step) \
JOIN {history_table} h \
ON h.workflow_uuid = m.original_id AND h.function_id < m.start_step"
)))
.bind(original_ids)
.bind(forked_ids)
.bind(start_steps)
.execute(&mut *tx)
.await?;
sqlx::query(AssertSqlSafe(format!(
"INSERT INTO {events_table} (workflow_uuid, key, value, serialization) \
SELECT fork_id, key, value, serialization FROM ( \
SELECT m.fork_id, h.key, h.value, h.serialization, \
row_number() OVER (PARTITION BY m.fork_id, h.key \
ORDER BY h.function_id DESC) AS rn \
FROM unnest($1::text[], $2::text[], $3::int4[]) \
AS m(original_id, fork_id, start_step) \
JOIN {history_table} h \
ON h.workflow_uuid = m.original_id AND h.function_id < m.start_step \
) latest WHERE rn = 1"
)))
.bind(original_ids)
.bind(forked_ids)
.bind(start_steps)
.execute(&mut *tx)
.await?;
sqlx::query(AssertSqlSafe(format!(
"INSERT INTO {streams_table} (workflow_uuid, function_id, key, value, \
serialization, \"offset\") \
SELECT m.fork_id, s.function_id, s.key, s.value, s.serialization, s.\"offset\" \
FROM unnest($1::text[], $2::text[], $3::int4[]) \
AS m(original_id, fork_id, start_step) \
JOIN {streams_table} s \
ON s.workflow_uuid = m.original_id AND s.function_id < m.start_step"
)))
.bind(original_ids)
.bind(forked_ids)
.bind(start_steps)
.execute(&mut *tx)
.await?;
}
tracing::info!(
count = original_ids.len(),
queue_name,
"forked workflows onto the queue"
);
Ok(tx)
}
async fn fork_batch(
&self,
original_ids: &[&str],
forked_ids: &[String],
steps: StartSteps<'_>,
options: &ForkOptions<'_>,
caller: Option<StepCaller<'_>>,
) -> Result<Vec<String>, Error> {
let started_at = Timestamp::now();
with_retry(&self.retry, "fork_workflows", move || async move {
self.run_transactional_step(
caller,
step_names::FORK_WORKFLOW,
started_at,
|tx| async move {
let tx = self
.fork_on(tx, original_ids, forked_ids, steps, options)
.await?;
Ok((tx, forked_ids.to_vec()))
},
)
.await
})
.await
}
async fn upsert_schedule_on(
&self,
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
schedule: &NewSchedule<'_>,
schedule_id: &str,
) -> Result<(), Error> {
let schedules_table = self.tables.workflow_schedules.as_str();
let application_name = schedule
.application_name
.or(self.application_name.as_deref());
let owner = resolve_owning_application(
tx,
schedules_table,
"schedule_name",
schedule.schedule_name,
application_name,
"Schedule",
)
.await?;
sqlx::query(AssertSqlSafe(format!(
"INSERT INTO {schedules_table} \
(schedule_id, schedule_name, workflow_name, workflow_class_name, schedule, status, \
context, last_fired_at, automatic_backfill, cron_timezone, queue_name, \
application_name) \
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12) \
ON CONFLICT (schedule_name) DO UPDATE SET \
workflow_name = EXCLUDED.workflow_name, \
workflow_class_name = EXCLUDED.workflow_class_name, \
schedule = EXCLUDED.schedule, \
context = EXCLUDED.context, \
automatic_backfill = EXCLUDED.automatic_backfill, \
cron_timezone = EXCLUDED.cron_timezone, \
queue_name = EXCLUDED.queue_name, \
application_name = COALESCE({schedules_table}.application_name, \
EXCLUDED.application_name)"
)))
.bind(schedule_id)
.bind(schedule.schedule_name)
.bind(schedule.workflow_name)
.bind(schedule.workflow_class_name)
.bind(schedule.schedule)
.bind(schedule.status.as_str())
.bind(schedule.context)
.bind(schedule.last_fired_at.map(Timestamp::to_iso8601))
.bind(schedule.automatic_backfill)
.bind(schedule.cron_timezone)
.bind(schedule.queue_name)
.bind(owner.as_deref())
.execute(&mut **tx)
.await?;
resolve_owning_application(
tx,
schedules_table,
"schedule_name",
schedule.schedule_name,
application_name,
"Schedule",
)
.await?;
tracing::debug!(
schedule_name = schedule.schedule_name,
"registered or updated a schedule"
);
Ok(())
}
async fn is_latest_application_version(
&self,
conn: &mut sqlx::PgConnection,
application_version: &str,
) -> Result<bool, Error> {
let versions_table = self.tables.application_versions.as_str();
let latest: Option<String> =
latest_version_query(versions_table, self.application_name.as_deref())
.fetch_optional(&mut *conn)
.await?
.map(|row| row.try_get("version_name"))
.transpose()?;
Ok(latest.is_none_or(|name| name == application_version))
}
}
fn version_predicate(is_latest: bool, param: usize) -> String {
if is_latest {
format!("(application_version = ${param} OR application_version IS NULL)")
} else {
format!("application_version = ${param}")
}
}
struct Tables {
workflow_status: String,
workflow_input: String,
workflow_output: String,
operation_outputs: String,
application_versions: String,
notifications: String,
workflow_events: String,
workflow_events_history: String,
streams: String,
queues: String,
workflow_schedules: String,
}
impl Tables {
fn new(schema: &str) -> Self {
let schema = quote_identifier(schema);
Self {
workflow_status: format!("{schema}.{}", quote_identifier("workflow_status")),
workflow_input: format!("{schema}.{}", quote_identifier("workflow_input")),
workflow_output: format!("{schema}.{}", quote_identifier("workflow_output")),
operation_outputs: format!("{schema}.{}", quote_identifier("operation_outputs")),
application_versions: format!("{schema}.{}", quote_identifier("application_versions")),
notifications: format!("{schema}.{}", quote_identifier("notifications")),
workflow_events: format!("{schema}.{}", quote_identifier("workflow_events")),
workflow_events_history: format!(
"{schema}.{}",
quote_identifier("workflow_events_history")
),
streams: format!("{schema}.{}", quote_identifier("streams")),
queues: format!("{schema}.{}", quote_identifier("queues")),
workflow_schedules: format!("{schema}.{}", quote_identifier("workflow_schedules")),
}
}
}
async fn ensure_database_exists(url: &str) -> Result<(), Error> {
let Some((maintenance_url, database)) = split_database(url) else {
return Ok(());
};
let admin = match PgPoolOptions::new()
.max_connections(1)
.connect(&maintenance_url)
.await
{
Ok(pool) => pool,
Err(_) => return Ok(()),
};
let exists = sqlx::query("SELECT 1 FROM pg_database WHERE datname = $1")
.bind(&database)
.fetch_optional(&admin)
.await?
.is_some();
if !exists {
let quoted = quote_identifier(&database);
match sqlx::raw_sql(AssertSqlSafe(format!("CREATE DATABASE {quoted}")))
.execute(&admin)
.await
{
Ok(_) => tracing::info!(database = %database, "created the system database"),
Err(e) if e.as_database_error().and_then(|d| d.code()).as_deref() == Some("42P04") => {}
Err(e) => return Err(e.into()),
}
}
admin.close().await;
Ok(())
}
fn split_database(url: &str) -> Option<(String, String)> {
let (prefix, tail) = url.split_once("://")?;
let (authority, rest) = tail.split_once('/')?;
let (database, query) = match rest.split_once('?') {
Some((db, q)) => (db, Some(q)),
None => (rest, None),
};
if database.is_empty() {
return None;
}
let mut maintenance = format!("{prefix}://{authority}/postgres");
if let Some(q) = query {
maintenance.push('?');
maintenance.push_str(q);
}
Some((maintenance, database.to_owned()))
}
const WORKFLOW_COLUMNS: &str = "workflow_uuid, status, name, class_name, config_name, \
serialization, executor_id, application_version, recovery_attempts, \
queue_name, created_at, updated_at, started_at_epoch_ms, completed_at, forked_from, \
parent_workflow_id, was_forked_from, owner_xid, application_id, authenticated_user, \
authenticated_roles, assumed_role, request, deduplication_id, priority, \
queue_partition_key, rate_limited, schedule_name, workflow_timeout_ms, \
workflow_deadline_epoch_ms, delay_until_epoch_ms, debounce_deadline_epoch_ms, \
is_debounced, application_name, attributes::text AS attributes";
fn workflow_payloads(tables: &Tables, load_input: bool, load_output: bool) -> String {
let inputs = if load_input {
payload_column(tables, PayloadColumn::Inputs)
} else {
"NULL::text AS inputs".to_owned()
};
let outcome = if load_output {
format!(
"{}, {}",
payload_column(tables, PayloadColumn::Output),
payload_column(tables, PayloadColumn::Error)
)
} else {
"NULL::text AS output, NULL::text AS error".to_owned()
};
format!("{inputs}, {outcome}")
}
#[derive(Clone, Copy)]
enum PayloadColumn {
Inputs,
Output,
Error,
}
fn payload_column(tables: &Tables, column: PayloadColumn) -> String {
let status = &tables.workflow_status;
let (table, name) = match column {
PayloadColumn::Inputs => (&tables.workflow_input, "inputs"),
PayloadColumn::Output => (&tables.workflow_output, "output"),
PayloadColumn::Error => (&tables.workflow_output, "error"),
};
format!(
"COALESCE((SELECT p.{name} FROM {table} p WHERE p.workflow_uuid = {status}.workflow_uuid), \
{status}.{name}) AS {name}"
)
}
fn encode_roles(roles: &[&str]) -> Option<String> {
(!roles.is_empty()).then(|| serde_json::to_string(roles).expect("a list of strings is JSON"))
}
fn decode_roles(stored: Option<String>) -> Result<Vec<String>, Error> {
let Some(json) = stored else {
return Ok(Vec::new());
};
serde_json::from_str(&json).map_err(|e| {
Error::Malformed(format!(
"authenticated_roles is not a JSON array of strings: {e}"
))
})
}
#[derive(Debug, Clone, Copy)]
enum SleepKind<'a> {
Durable(&'a dyn StepCodec),
Deadline,
}
const STREAM_OFFSET_ATTEMPTS: u32 = 16;
const NOW_MS_SQL: &str = "(EXTRACT(epoch FROM now()) * 1000)::bigint";
const CLOCK_MS_SQL: &str = "(EXTRACT(epoch FROM clock_timestamp()) * 1000)::bigint";
fn now_ms_plus_sql(operand: &str) -> String {
format!(
"CASE WHEN ({operand}) IS NULL THEN NULL \
ELSE LEAST(GREATEST({NOW_MS_SQL}::numeric + ({operand})::numeric, \
({min})::numeric), ({max})::numeric)::bigint END",
min = i64::MIN,
max = i64::MAX,
)
}
fn timed_out_set_sql() -> String {
format!(
"status = 'CANCELLED', owner_xid = NULL, queue_name = NULL, deduplication_id = NULL, \
started_at_epoch_ms = NULL, updated_at = {NOW_MS_SQL}, completed_at = {NOW_MS_SQL}"
)
}
fn timed_out_where_sql() -> String {
format!(
"status IN ({UNSETTLED}) AND workflow_deadline_epoch_ms IS NOT NULL \
AND workflow_deadline_epoch_ms <= {NOW_MS_SQL}"
)
}
fn timed_out_sweep_sql(workflow_table: &str) -> String {
format!(
"UPDATE {workflow_table} SET {set} \
WHERE workflow_uuid IN ( \
SELECT workflow_uuid FROM {workflow_table} \
WHERE {timed_out} \
AND ($1::text IS NULL OR application_name = $1 OR application_name IS NULL) \
ORDER BY workflow_deadline_epoch_ms \
LIMIT $2 \
FOR UPDATE SKIP LOCKED) \
RETURNING workflow_uuid",
set = timed_out_set_sql(),
timed_out = timed_out_where_sql(),
)
}
const UNSETTLED: &str = "'PENDING', 'ENQUEUED', 'DELAYED'";
const QUEUE_COLUMNS: &str = "name, concurrency, worker_concurrency, rate_limit_max, \
rate_limit_period_sec, partition_concurrency, partition_worker_concurrency, \
partition_rate_limit_max, partition_rate_limit_period_sec, polling_interval_sec, \
application_name";
const QUEUE_PENDING_PREDICATE: &str = "queue_name = $2 AND status = 'PENDING' \
AND ($1::text IS NULL OR application_name = $1 OR application_name IS NULL)";
fn bind_new_queue<'q>(
query: sqlx::query::Query<'q, sqlx::Postgres, sqlx::postgres::PgArguments>,
queue: &NewQueue<'q>,
owner: Option<&'q str>,
) -> sqlx::query::Query<'q, sqlx::Postgres, sqlx::postgres::PgArguments> {
query
.bind(queue.name)
.bind(queue.concurrency)
.bind(queue.worker_concurrency)
.bind(queue.rate_limit.map(|r| r.limit))
.bind(queue.rate_limit.map(|r| r.period.as_secs_f64()))
.bind(true)
.bind(queue.is_partitioned())
.bind(queue.partition_concurrency)
.bind(queue.partition_worker_concurrency)
.bind(queue.partition_rate_limit.map(|r| r.limit))
.bind(queue.partition_rate_limit.map(|r| r.period.as_secs_f64()))
.bind(queue.polling_interval.as_secs_f64())
.bind(owner)
}
const SCHEDULE_COLUMNS: &str = "schedule_id, schedule_name, workflow_name, workflow_class_name, \
schedule, status, context, last_fired_at, automatic_backfill, cron_timezone, queue_name, \
application_name";
fn escape_like(value: &str) -> String {
value
.replace('\\', "\\\\")
.replace('%', "\\%")
.replace('_', "\\_")
}
fn schedule_from_row(row: &sqlx::postgres::PgRow) -> Result<ScheduleRecord, Error> {
let stored: String = row.try_get("status")?;
let status = ScheduleStatus::parse(&stored).ok_or_else(|| {
Error::Malformed(format!(
"schedule status {stored:?} is not one this build knows"
))
})?;
let stored: Option<String> = row.try_get("last_fired_at")?;
let last_fired_at = stored
.map(|stored| {
Timestamp::parse_iso8601(&stored).ok_or_else(|| {
Error::Malformed(format!(
"last_fired_at {stored:?} is not an ISO-8601 instant"
))
})
})
.transpose()?;
Ok(ScheduleRecord {
schedule_id: row.try_get("schedule_id")?,
schedule_name: row.try_get("schedule_name")?,
workflow_name: row.try_get("workflow_name")?,
workflow_class_name: row.try_get("workflow_class_name")?,
schedule: row.try_get("schedule")?,
status,
context: row.try_get("context")?,
last_fired_at,
automatic_backfill: row.try_get("automatic_backfill")?,
cron_timezone: row.try_get("cron_timezone")?,
queue_name: row.try_get("queue_name")?,
application_name: row.try_get("application_name")?,
})
}
fn queue_from_row(row: &sqlx::postgres::PgRow) -> Result<QueueRecord, Error> {
let period = |column: &str, value: Option<f64>| match value {
None => Ok(None),
Some(secs) => duration_from_secs(secs)
.map(Some)
.ok_or_else(|| Error::Malformed(format!("{column} is not a duration: {secs}"))),
};
let rate_limit = |max_column: &str, period_column: &str| -> Result<Option<RateLimit>, Error> {
Ok(
match (
row.try_get::<Option<i32>, _>(max_column)?,
period(period_column, row.try_get(period_column)?)?,
) {
(Some(limit), Some(period)) => Some(RateLimit { limit, period }),
_ => None,
},
)
};
Ok(QueueRecord {
name: row.try_get("name")?,
concurrency: row.try_get("concurrency")?,
worker_concurrency: row.try_get("worker_concurrency")?,
rate_limit: rate_limit("rate_limit_max", "rate_limit_period_sec")?,
partition_concurrency: row.try_get("partition_concurrency")?,
partition_worker_concurrency: row.try_get("partition_worker_concurrency")?,
partition_rate_limit: rate_limit(
"partition_rate_limit_max",
"partition_rate_limit_period_sec",
)?,
polling_interval: period("polling_interval_sec", row.try_get("polling_interval_sec")?)?
.ok_or_else(|| Error::Malformed("polling_interval_sec is null".to_owned()))?,
application_name: row.try_get("application_name")?,
})
}
const VERSION_COLUMNS: &str =
"version_id, version_name, version_timestamp, created_at, application_name";
fn latest_version_query<'q>(
versions_table: &str,
application_name: Option<&'q str>,
) -> sqlx::query::Query<'q, sqlx::Postgres, sqlx::postgres::PgArguments> {
let query = sqlx::query(AssertSqlSafe(latest_version_sql(
versions_table,
application_name.is_some(),
)));
match application_name {
Some(application_name) => query.bind(application_name),
None => query,
}
}
fn latest_version_sql(versions_table: &str, named: bool) -> String {
if !named {
return format!(
"SELECT {VERSION_COLUMNS} FROM {versions_table} ORDER BY version_timestamp DESC LIMIT 1"
);
}
format!(
"(SELECT {VERSION_COLUMNS} FROM {versions_table} \
WHERE application_name = $1 \
ORDER BY application_name, version_timestamp DESC LIMIT 1) \
UNION ALL \
(SELECT {VERSION_COLUMNS} FROM {versions_table} \
WHERE application_name IS NULL \
ORDER BY application_name, version_timestamp DESC LIMIT 1) \
ORDER BY version_timestamp DESC LIMIT 1"
)
}
async fn resolve_owning_application(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
table: &str,
key_column: &str,
name: &str,
claimant: Option<&str>,
kind: &'static str,
) -> Result<Option<String>, Error> {
let holder: Option<Option<String>> = sqlx::query_scalar(AssertSqlSafe(format!(
"SELECT application_name FROM {table} WHERE {key_column} = $1"
)))
.bind(name)
.fetch_optional(&mut **tx)
.await?;
let Some(Some(holder)) = holder else {
return Ok(claimant.map(str::to_owned));
};
match claimant {
None => Ok(Some(holder)),
Some(claimant) if claimant == holder => Ok(Some(holder)),
Some(claimant) => Err(Error::RegisteredByAnother {
kind: kind.into(),
name: name.to_owned(),
holder,
claimant: Some(claimant.to_owned()),
}),
}
}
fn decode_wake_time(
output: Option<&str>,
workflow_id: &str,
step_id: i32,
) -> Result<Timestamp, Error> {
let recorded = output.ok_or_else(|| {
Error::Malformed(format!(
"workflow {workflow_id} step {step_id} is a sleep with no recorded wake time"
))
})?;
recorded
.trim()
.parse::<i64>()
.map(Timestamp::from_epoch_ms)
.map_err(|e| {
Error::Malformed(format!(
"workflow {workflow_id} step {step_id} wake time {recorded:?} is not epoch milliseconds: {e}"
))
})
}
fn replayed_sleep(
step: &StepRecord,
kind: SleepKind<'_>,
workflow_id: &str,
step_id: i32,
owner_xid: &str,
) -> Result<Timestamp, Error> {
if let (SleepKind::Durable(codec), Some(recorded)) = (kind, step.error.as_deref()) {
let caller = StepCaller {
workflow_id,
step_id,
codec,
owner_xid,
};
return Err(decode_step_value(
step,
recorded,
caller,
step_names::SLEEP,
"an error",
)?);
}
decode_wake_time(step.output.as_deref(), workflow_id, step_id)
}
fn version_from_row(row: &sqlx::postgres::PgRow) -> Result<VersionInfo, Error> {
Ok(VersionInfo {
version_id: row.try_get("version_id")?,
version_name: row.try_get("version_name")?,
version_timestamp: Timestamp::from_epoch_ms(row.try_get("version_timestamp")?),
created_at: Timestamp::from_epoch_ms(row.try_get("created_at")?),
app_name: row.try_get("application_name")?,
})
}
fn parse_workflow_status(text: &str) -> Result<WorkflowStatus, Error> {
WorkflowStatus::parse(text)
.ok_or_else(|| Error::Malformed(format!("unknown workflow status {text:?}")))
}
fn workflow_from_row(row: &sqlx::postgres::PgRow) -> Result<WorkflowRecord, Error> {
let status = parse_workflow_status(&row.try_get::<String, _>("status")?)?;
Ok(WorkflowRecord {
workflow_id: row.try_get("workflow_uuid")?,
status,
name: row.try_get("name")?,
class_name: row.try_get("class_name")?,
config_name: row.try_get("config_name")?,
input: row.try_get("inputs")?,
output: row.try_get("output")?,
error: row.try_get("error")?,
serialization: row.try_get("serialization")?,
executor_id: row.try_get("executor_id")?,
app_version: row.try_get("application_version")?,
recovery_attempts: row
.try_get::<Option<i64>, _>("recovery_attempts")?
.unwrap_or(0),
queue_name: row.try_get("queue_name")?,
created_at: Timestamp::from_epoch_ms(row.try_get("created_at")?),
updated_at: Timestamp::from_epoch_ms(row.try_get("updated_at")?),
started_at: row
.try_get::<Option<i64>, _>("started_at_epoch_ms")?
.map(Timestamp::from_epoch_ms),
completed_at: row
.try_get::<Option<i64>, _>("completed_at")?
.map(Timestamp::from_epoch_ms),
forked_from: row.try_get("forked_from")?,
parent_workflow_id: row.try_get("parent_workflow_id")?,
was_forked_from: row
.try_get::<Option<bool>, _>("was_forked_from")?
.unwrap_or(false),
owner_xid: row.try_get("owner_xid")?,
app_name: row.try_get("application_name")?,
app_id: row.try_get("application_id")?,
authenticated_user: row.try_get("authenticated_user")?,
authenticated_roles: decode_roles(
row.try_get::<Option<String>, _>("authenticated_roles")?,
)?,
assumed_role: row.try_get("assumed_role")?,
request: row.try_get("request")?,
deduplication_id: row.try_get("deduplication_id")?,
priority: row.try_get::<Option<i32>, _>("priority")?.unwrap_or(0),
queue_partition_key: row.try_get("queue_partition_key")?,
rate_limited: row
.try_get::<Option<bool>, _>("rate_limited")?
.unwrap_or(false),
schedule_name: row.try_get("schedule_name")?,
timeout: row
.try_get::<Option<i64>, _>("workflow_timeout_ms")?
.and_then(duration_from_ms),
deadline: row
.try_get::<Option<i64>, _>("workflow_deadline_epoch_ms")?
.map(Timestamp::from_epoch_ms),
delay_until: row
.try_get::<Option<i64>, _>("delay_until_epoch_ms")?
.map(Timestamp::from_epoch_ms),
debounce_deadline: row
.try_get::<Option<i64>, _>("debounce_deadline_epoch_ms")?
.map(Timestamp::from_epoch_ms),
is_debounced: row
.try_get::<Option<bool>, _>("is_debounced")?
.unwrap_or(false),
attributes: row.try_get("attributes")?,
})
}
impl PostgresSystemDatabase {
async fn check_step_on(
&self,
conn: &mut sqlx::PgConnection,
workflow_id: &str,
step_id: i32,
step_name: &str,
) -> Result<Option<StepRecord>, Error> {
let workflow_table = &self.tables.workflow_status;
let steps_table = &self.tables.operation_outputs;
let (workflow_table, steps_table) = (workflow_table.as_str(), steps_table.as_str());
let step_columns = qualified_step_columns("o");
{
let row = sqlx::query(AssertSqlSafe(format!(
"SELECT s.status, {step_columns} \
FROM {workflow_table} s \
LEFT JOIN {steps_table} o \
ON o.workflow_uuid = s.workflow_uuid AND o.function_id = $2 \
WHERE s.workflow_uuid = $1"
)))
.bind(workflow_id)
.bind(step_id)
.fetch_optional(&mut *conn)
.await?;
let Some(row) = row else {
return Err(Error::NonExistentWorkflow {
workflow_ids: vec![workflow_id.to_owned()],
});
};
let status: String = row.try_get("status")?;
if status == WorkflowStatus::Cancelled.as_str() {
return Err(Error::WorkflowCancelled {
workflow_id: workflow_id.to_owned(),
});
}
if row.try_get::<Option<i32>, _>("function_id")?.is_none() {
return Ok(None);
}
let step = step_from_row(&row, workflow_id)?;
if step.step_name != step_name {
return Err(Error::UnexpectedStep {
workflow_id: workflow_id.to_owned(),
step_id,
expected: step_name.to_owned(),
recorded: step.step_name,
});
}
Ok(Some(step))
}
}
async fn check_owner_on(
&self,
conn: &mut sqlx::PgConnection,
owner: WorkflowOwner<'_>,
) -> Result<DatabaseTime, Error> {
let workflow_table = self.tables.workflow_status.as_str();
let current: Option<(Option<String>, String, i64)> =
sqlx::query_as(AssertSqlSafe(format!(
"SELECT owner_xid, status, {CLOCK_MS_SQL} FROM ( \
SELECT owner_xid, status FROM {workflow_table} \
WHERE workflow_uuid = $1 FOR NO KEY UPDATE \
) locked"
)))
.bind(owner.workflow_id)
.fetch_optional(&mut *conn)
.await?;
match current {
Some((_, status, _)) if status == WorkflowStatus::Cancelled.as_str() => {
Err(Error::WorkflowCancelled {
workflow_id: owner.workflow_id.to_owned(),
})
}
Some((Some(owner_xid), _, db_now)) if owner_xid == owner.owner_xid => {
Ok(DatabaseTime::read(db_now))
}
_ => Err(Error::WorkflowConflict {
workflow_id: owner.workflow_id.to_owned(),
}),
}
}
#[allow(clippy::too_many_arguments)]
async fn record_step_on(
&self,
conn: &mut sqlx::PgConnection,
workflow_id: &str,
step_id: i32,
step_name: &str,
outcome: Outcome<'_>,
serialization: Option<&str>,
timing: Option<StepTiming>,
child_workflow_id: Option<&str>,
) -> Result<(), Error> {
if workflow_id.is_empty() {
return Err(Error::InvalidInput {
field: "workflow_id".into(),
detail: "must not be empty".to_owned(),
});
}
if step_id < 0 {
return Err(Error::InvalidInput {
field: "step_id".into(),
detail: "must not be negative".to_owned(),
});
}
let workflow_table = &self.tables.workflow_status;
let steps_table = &self.tables.operation_outputs;
let (workflow_table, steps_table) = (workflow_table.as_str(), steps_table.as_str());
let executor_id = self.executor_id.as_deref();
let application_name = self.application_name.as_deref();
let (output, error) = outcome.columns();
{
let stored: Option<Option<i64>> = sqlx::query_scalar(AssertSqlSafe(format!(
"INSERT INTO {steps_table} (workflow_uuid, function_id, function_name, output, \
error, serialization, started_at_epoch_ms, completed_at_epoch_ms, \
application_name, child_workflow_id) \
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10) \
ON CONFLICT (workflow_uuid, function_id) DO UPDATE \
SET completed_at_epoch_ms = {steps_table}.completed_at_epoch_ms \
RETURNING completed_at_epoch_ms"
)))
.bind(workflow_id)
.bind(step_id)
.bind(step_name)
.bind(output)
.bind(error)
.bind(serialization)
.bind(timing.map(|t| t.started_at.as_epoch_ms()))
.bind(timing.map(|t| t.completed_at.as_epoch_ms()))
.bind(application_name)
.bind(child_workflow_id)
.fetch_optional(&mut *conn)
.await?;
let ours = stored.flatten() == timing.map(|t| t.completed_at.as_epoch_ms());
if !ours {
tracing::warn!(
workflow_id,
step_id,
step_name,
"step was already recorded by another execution"
);
return Err(Error::StepNondeterminism {
workflow_id: workflow_id.to_owned(),
step_id,
});
}
if let Some(executor_id) = executor_id {
sqlx::query(AssertSqlSafe(format!(
"UPDATE {workflow_table} SET executor_id = $2 \
WHERE workflow_uuid = $1 AND executor_id IS DISTINCT FROM $2"
)))
.bind(workflow_id)
.bind(executor_id)
.execute(&mut *conn)
.await?;
}
Ok(())
}
}
async fn cancel_batch<S>(
&self,
conn: &mut sqlx::PgConnection,
workflow_ids: &[S],
) -> Result<Vec<String>, Error>
where
S: AsRef<str> + Sync,
{
let workflow_table = &self.tables.workflow_status;
let ids: Vec<&str> = workflow_ids.iter().map(AsRef::as_ref).collect();
let cancelled: Vec<String> = sqlx::query_scalar(AssertSqlSafe(format!(
"UPDATE {workflow_table} SET status = 'CANCELLED', queue_name = NULL, \
deduplication_id = NULL, started_at_epoch_ms = NULL, owner_xid = NULL, \
updated_at = {NOW_MS_SQL}, completed_at = {NOW_MS_SQL} \
WHERE workflow_uuid = ANY($1) \
AND status NOT IN ('SUCCESS', 'ERROR', 'CANCELLED') \
RETURNING workflow_uuid"
)))
.bind(ids)
.fetch_all(&mut *conn)
.await?;
Ok(cancelled)
}
async fn direct_children<S>(
&self,
conn: &mut sqlx::PgConnection,
workflow_ids: &[S],
) -> Result<Vec<String>, Error>
where
S: AsRef<str> + Sync,
{
let workflow_table = &self.tables.workflow_status;
let ids: Vec<&str> = workflow_ids.iter().map(AsRef::as_ref).collect();
let children: Vec<String> = sqlx::query_scalar(AssertSqlSafe(format!(
"SELECT workflow_uuid FROM {workflow_table} WHERE parent_workflow_id = ANY($1)"
)))
.bind(ids)
.fetch_all(&mut *conn)
.await?;
Ok(children)
}
}
const STEP_COLUMNS: &str = "function_id, function_name, \
child_workflow_id, serialization, started_at_epoch_ms, completed_at_epoch_ms";
fn step_payloads(load_output: bool) -> &'static str {
if load_output {
"output, error"
} else {
"NULL::text AS output, NULL::text AS error"
}
}
fn step_listing_query(
steps_table: &str,
workflow_id: &str,
load_output: bool,
limit: Option<i64>,
offset: Option<i64>,
) -> sqlx::QueryBuilder<sqlx::Postgres> {
let mut q = sqlx::QueryBuilder::<sqlx::Postgres>::new("SELECT ");
q.push(STEP_COLUMNS)
.push(", ")
.push(step_payloads(load_output))
.push(" FROM ")
.push(steps_table)
.push(" WHERE workflow_uuid = ")
.push_bind(workflow_id)
.push(" ORDER BY function_id");
if let Some(limit) = limit {
q.push(" LIMIT ").push_bind(limit);
}
if let Some(offset) = offset {
q.push(" OFFSET ").push_bind(offset);
}
q
}
fn qualified_step_columns(alias: &str) -> String {
format!("{STEP_COLUMNS}, {}", step_payloads(true))
.split(',')
.map(|c| format!("{alias}.{}", c.trim()))
.collect::<Vec<_>>()
.join(", ")
}
fn step_from_row(row: &sqlx::postgres::PgRow, workflow_id: &str) -> Result<StepRecord, Error> {
Ok(StepRecord {
workflow_id: workflow_id.to_owned(),
step_id: row.try_get("function_id")?,
step_name: row.try_get("function_name")?,
output: row.try_get("output")?,
error: row.try_get("error")?,
child_workflow_id: row.try_get("child_workflow_id")?,
serialization: row.try_get("serialization")?,
started_at: row
.try_get::<Option<i64>, _>("started_at_epoch_ms")?
.map(Timestamp::from_epoch_ms),
completed_at: row
.try_get::<Option<i64>, _>("completed_at_epoch_ms")?
.map(Timestamp::from_epoch_ms),
})
}
impl PostgresSystemDatabase {
async fn checkpoint_sleep(
&self,
workflow_id: &str,
step_id: i32,
duration: Duration,
kind: SleepKind<'_>,
owner_xid: &str,
) -> Result<WakeTime, Error> {
let pool = &self.pool;
let first: OnceLock<(Timestamp, Result<Timestamp, Error>)> = OnceLock::new();
let first = &first;
with_retry(&self.retry, "checkpoint_sleep", move || async move {
let mut conn = pool.begin().await?;
let clock = self
.check_owner_on(
&mut conn,
WorkflowOwner {
workflow_id,
owner_xid,
},
)
.await?;
let (started_at, wake_at) = first.get_or_init(|| {
let started_at = clock.now();
let wake_at = match kind {
SleepKind::Durable(_) => started_at.try_add("duration", duration),
SleepKind::Deadline => Ok(started_at
.checked_add(duration)
.unwrap_or(Timestamp::from_epoch_ms(i64::MAX))),
};
(started_at, wake_at)
});
let started_at = *started_at;
let answer = |at: Timestamp| WakeTime { at, clock };
if let Some(step) = self
.check_step_on(&mut conn, workflow_id, step_id, step_names::SLEEP)
.await?
{
tracing::debug!(workflow_id, step_id, "replaying sleep");
return replayed_sleep(&step, kind, workflow_id, step_id, owner_xid).map(answer);
}
tracing::debug!(
workflow_id,
step_id,
duration_ms = duration.as_millis(),
"running sleep"
);
let encoded: EncodedValue;
let (outcome, completed_at) = match (wake_at, kind) {
(Ok(wake_at), kind) => {
encoded = EncodedValue {
value: wake_at.as_epoch_ms().to_string(),
serialization: Some(PORTABLE_JSON.to_owned()),
};
let completed_at = match kind {
SleepKind::Durable(_) => *wake_at,
SleepKind::Deadline => started_at,
};
(Outcome::Output(Some(&encoded.value)), completed_at)
}
(Err(refusal), SleepKind::Durable(codec)) => {
encoded = encode_step_value(refusal, codec, "an error").inspect_err(
|encode_error| {
tracing::warn!(
workflow_id,
step_id,
error = %refusal,
encode_error = %encode_error,
"the sleep was refused, but the refusal could not be encoded to \
record it"
);
},
)?;
(Outcome::Error(&encoded.value), started_at)
}
(Err(refusal), SleepKind::Deadline) => return Err(refusal.clone()),
};
match self
.record_step_on(
&mut conn,
workflow_id,
step_id,
step_names::SLEEP,
outcome,
encoded.serialization.as_deref(),
Some(StepTiming {
started_at,
completed_at,
}),
None,
)
.await
{
Ok(()) => {
conn.commit().await?;
wake_at.clone().map(answer)
}
Err(Error::StepNondeterminism { .. }) => {
let step = self
.check_step_on(&mut conn, workflow_id, step_id, step_names::SLEEP)
.await?
.ok_or_else(|| {
Error::Malformed(
"sleep reported as recorded but cannot be read back".to_owned(),
)
})?;
replayed_sleep(&step, kind, workflow_id, step_id, owner_xid).map(answer)
}
Err(e) => Err(e),
}
})
.await
}
}
impl PostgresSystemDatabase {
async fn resolve_fork_points_on(
&self,
conn: &mut sqlx::PgConnection,
workflow_ids: &[&str],
aggregate: &str,
named: Option<&str>,
) -> Result<Vec<i32>, Error> {
let steps_table = self.tables.operation_outputs.as_str();
let filter = if named.is_some() {
" AND function_name = $2"
} else {
""
};
let sql = format!(
"SELECT workflow_uuid, {aggregate} AS start_step FROM {steps_table} \
WHERE workflow_uuid = ANY($1){filter} GROUP BY workflow_uuid"
);
let mut query = sqlx::query_as(AssertSqlSafe(sql)).bind(workflow_ids);
if let Some(name) = named {
query = query.bind(name);
}
let rows: Vec<(String, i32)> = query.fetch_all(&mut *conn).await?;
let resolved: HashMap<&str, i32> =
rows.iter().map(|(id, step)| (id.as_str(), *step)).collect();
let mut start_steps = Vec::with_capacity(workflow_ids.len());
let mut missing = Vec::new();
for id in workflow_ids {
match resolved.get(id) {
Some(&step) => start_steps.push(step),
None => missing.push((*id).to_owned()),
}
}
if !missing.is_empty() {
return Err(Error::NoForkPoint {
workflow_ids: missing,
step_name: named.map(str::to_owned),
});
}
Ok(start_steps)
}
}
impl PostgresSystemDatabase {
async fn descendant_forks(
conn: &mut sqlx::PgConnection,
table: &str,
roots: &[&str],
) -> Result<HashMap<String, Vec<String>>, Error> {
let mut children: HashMap<String, Vec<String>> = HashMap::new();
let mut seen: HashSet<String> = roots.iter().map(|r| (*r).to_owned()).collect();
let mut frontier: Vec<String> = seen.iter().cloned().collect();
while !frontier.is_empty() {
let borrowed: Vec<&str> = frontier.iter().map(String::as_str).collect();
let rows: Vec<(String, String)> = sqlx::query_as(AssertSqlSafe(format!(
"SELECT workflow_uuid, forked_from FROM {table} \
WHERE forked_from = ANY($1) AND forked_from IS NOT NULL"
)))
.bind(&borrowed)
.fetch_all(&mut *conn)
.await?;
let mut next = Vec::new();
for (forked_id, forked_from) in rows {
children
.entry(forked_from)
.or_default()
.push(forked_id.clone());
if seen.insert(forked_id.clone()) {
next.push(forked_id);
}
}
frontier = next;
}
let mut descendants = HashMap::new();
for root in roots {
if descendants.contains_key(*root) {
continue;
}
let mut found: HashSet<&str> = HashSet::new();
let mut stack: Vec<&str> = children
.get(*root)
.map(|c| c.iter().map(String::as_str).collect())
.unwrap_or_default();
while let Some(node) = stack.pop() {
if node != *root
&& found.insert(node)
&& let Some(grandchildren) = children.get(node)
{
stack.extend(grandchildren.iter().map(String::as_str));
}
}
let mut found: Vec<String> = found.into_iter().map(str::to_owned).collect();
found.sort();
descendants.insert((*root).to_owned(), found);
}
Ok(descendants)
}
}
impl PostgresSystemDatabase {
async fn rename_application_in_batches(
&self,
table: &str,
source: RenameFrom<'_>,
new_name: &str,
batching: RenameBatching,
) -> Result<u64, Error> {
let key_column = "workflow_uuid";
let update_predicate = rename_source_predicate(source, 2);
let select_predicate = rename_source_predicate(source, 1);
let update_predicate = update_predicate.as_str();
let select_predicate = select_predicate.as_str();
let renamed_from = source.application();
let pool = &self.pool;
let RenameBatching::Batched(batch_size) = batching else {
return with_retry(&self.retry, "rename_rows", move || async move {
let moved = sqlx::query(AssertSqlSafe(format!(
"UPDATE {table} SET application_name = $1 WHERE {update_predicate}"
)))
.bind(new_name)
.bind(renamed_from)
.execute(pool)
.await?
.rows_affected();
Ok(moved)
})
.await;
};
let select_bound = format!(
"SELECT DISTINCT {key_column} FROM {table} \
WHERE {select_predicate} AND {key_column} > COALESCE($2, '') \
ORDER BY {key_column} LIMIT 1 OFFSET {}",
batch_size - 1
);
let update_range = format!(
"UPDATE {table} SET application_name = $1 \
WHERE {update_predicate} \
AND {key_column} > COALESCE($3, '') AND {key_column} <= $4"
);
let update_rest =
format!("UPDATE {table} SET application_name = $1 WHERE {update_predicate}");
let select_bound = select_bound.as_str();
let update_range = update_range.as_str();
let update_rest = update_rest.as_str();
let mut total = 0;
let mut watermark: Option<String> = None;
loop {
let above = watermark.clone();
let (moved, upper) = with_retry(&self.retry, "rename_row_batch", || {
let above = above.clone();
async move {
let upper: Option<String> = sqlx::query_scalar(AssertSqlSafe(select_bound))
.bind(renamed_from)
.bind(above.as_deref())
.fetch_optional(pool)
.await?
.flatten();
let moved = match &upper {
Some(upper) => sqlx::query(AssertSqlSafe(update_range))
.bind(new_name)
.bind(renamed_from)
.bind(above.as_deref())
.bind(upper.as_str()),
None => sqlx::query(AssertSqlSafe(update_rest))
.bind(new_name)
.bind(renamed_from),
}
.execute(pool)
.await?
.rows_affected();
Ok((moved, upper))
}
})
.await?;
total += moved;
match upper {
None => return Ok(total),
Some(upper) => watermark = Some(upper),
}
}
}
}
impl PostgresSystemDatabase {
async fn is_rewound_child(
&self,
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
caller: Option<InitWorkflowCaller<'_>>,
row: &sqlx::postgres::PgRow,
workflow_id: &str,
) -> Result<bool, Error> {
let Some(caller) = caller else {
return Ok(false);
};
let stored_parent: Option<String> = row.try_get("parent_workflow_id")?;
if stored_parent.as_deref() != Some(caller.parent_workflow_id) {
return Ok(false);
}
let steps_table = self.tables.operation_outputs.as_str();
let recorded: bool = sqlx::query_scalar(AssertSqlSafe(format!(
"SELECT EXISTS (SELECT 1 FROM {steps_table} \
WHERE workflow_uuid = $1 AND child_workflow_id = $2 AND function_id <> $3)"
)))
.bind(caller.parent_workflow_id)
.bind(workflow_id)
.bind(caller.step_id)
.fetch_one(&mut **tx)
.await?;
Ok(!recorded)
}
async fn init_workflow_as(
&self,
workflow: &NewWorkflow<'_>,
caller: Option<InitWorkflowCaller<'_>>,
reuse_policy: WorkflowIdReusePolicy,
creator_xid: &str,
) -> Result<WorkflowInitResult, Error> {
workflow.validate()?;
let workflow_table = &self.tables.workflow_status;
let initial_status = workflow.initial_status();
let queued = matches!(
initial_status,
WorkflowStatus::Enqueued | WorkflowStatus::Delayed
);
let initial_attempts = i64::from(!queued);
let to_ms = |field, duration: Option<Duration>| {
duration.map(|d| duration_to_ms(field, d)).transpose()
};
let timeout_ms = to_ms("timeout", workflow.timeout)?;
let delay_ms = to_ms("delay", workflow.delay)?;
let debounce_timeout_ms = to_ms("debounce_timeout", workflow.debounce_timeout)?;
let (workflow_table, input_table, pool) = (
workflow_table.as_str(),
self.tables.workflow_input.as_str(),
&self.pool,
);
with_retry(&self.retry, "init_workflow", move || async move {
let mut tx = pool.begin().await?;
if let Some(caller) = caller {
self.check_owner_on(
&mut tx,
WorkflowOwner {
workflow_id: caller.parent_workflow_id,
owner_xid: caller.owner_xid,
},
)
.await?;
}
let deadline_from_timeout = now_ms_plus_sql("$18::bigint");
let delay_until = now_ms_plus_sql("$10::bigint");
let debounce_deadline = now_ms_plus_sql("$25::bigint");
let row = sqlx::query(AssertSqlSafe(format!(
"INSERT INTO {workflow_table} (workflow_uuid, status, \
name, class_name, config_name, \
queue_name, deduplication_id, priority, queue_partition_key, delay_until_epoch_ms, \
authenticated_user, assumed_role, authenticated_roles, \
executor_id, application_version, application_id, \
created_at, updated_at, recovery_attempts, \
workflow_timeout_ms, workflow_deadline_epoch_ms, \
parent_workflow_id, owner_xid, serialization, attributes, schedule_name, \
debounce_deadline_epoch_ms, is_debounced, application_name, creator_xid) \
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, \
CASE WHEN $10::bigint IS NULL THEN NULL \
WHEN $25::bigint IS NULL THEN {delay_until} \
ELSE LEAST({delay_until}, {debounce_deadline}) END, \
$11, $12, $13, $14, $15, $16, \
{NOW_MS_SQL}, {NOW_MS_SQL}, $17, $18, \
COALESCE($19, CASE WHEN $2 = 'PENDING' THEN {deadline_from_timeout} END), \
$20, CASE WHEN $2 = 'PENDING' THEN $21 END, $22, $23::jsonb, $24, \
{debounce_deadline}, $26, $27, $21) \
ON CONFLICT (workflow_uuid) DO UPDATE SET \
creator_xid = {workflow_table}.creator_xid \
RETURNING recovery_attempts, status, name, class_name, config_name, queue_name, \
deduplication_id, workflow_deadline_epoch_ms, creator_xid, serialization, \
parent_workflow_id, \
created_at = {NOW_MS_SQL} AS inserted, \
{CLOCK_MS_SQL} AS db_now"
)))
.bind(workflow.workflow_id)
.bind(initial_status.as_str())
.bind(workflow.name)
.bind(workflow.class_name)
.bind(workflow.config_name)
.bind(workflow.queue_name)
.bind(workflow.deduplication_id)
.bind(workflow.priority)
.bind(workflow.queue_partition_key)
.bind(delay_ms)
.bind(empty_to_none(workflow.authenticated_user))
.bind(empty_to_none(workflow.assumed_role))
.bind(encode_roles(&workflow.authenticated_roles))
.bind(workflow.executor_id)
.bind(workflow.application_version)
.bind(workflow.application_id)
.bind(initial_attempts)
.bind(timeout_ms)
.bind(workflow.deadline.map(Timestamp::as_epoch_ms))
.bind(caller.map(|caller| caller.parent_workflow_id))
.bind(creator_xid)
.bind(workflow.serialization)
.bind(workflow.attributes)
.bind(workflow.schedule_name)
.bind(debounce_timeout_ms)
.bind(workflow.is_debounced)
.bind(
workflow
.application_name
.or(self.application_name.as_deref()),
)
.fetch_one(&mut *tx)
.await
.map_err(|e| match Error::from(e) {
Error::Backend(b)
if b.sqlstate.as_deref() == Some("23505")
&& let Some(deduplication_id) = workflow.deduplication_id
&& let Some(queue_name) = workflow.queue_name =>
{
Error::QueueDeduplicated {
workflow_id: workflow.workflow_id.to_owned(),
queue_name: queue_name.to_owned(),
deduplication_id: deduplication_id.to_owned(),
}
}
other => other,
})?;
let clock = DatabaseTime::read(row.try_get("db_now")?);
let recovery_attempts: i64 = row.try_get("recovery_attempts")?;
let status = parse_workflow_status(&row.try_get::<String, _>("status")?)?;
let stored_creator: Option<String> = row.try_get("creator_xid")?;
let deadline: Option<i64> = row.try_get("workflow_deadline_epoch_ms")?;
let serialization: Option<String> = row.try_get("serialization")?;
if reuse_policy == WorkflowIdReusePolicy::Reject
&& stored_creator.as_deref() != Some(creator_xid)
&& !self
.is_rewound_child(&mut tx, caller, &row, workflow.workflow_id)
.await?
{
let name: Option<String> = row.try_get("name")?;
return Err(Error::WorkflowIdInUse {
workflow_id: workflow.workflow_id.to_owned(),
status,
name: name.filter(|name| !name.is_empty()),
});
}
for (field, stored, offered) in [
(
"function name",
row.try_get::<Option<String>, _>("name")?,
&workflow.name,
),
("class name", row.try_get("class_name")?, &workflow.class_name),
("config name", row.try_get("config_name")?, &workflow.config_name),
] {
if stored.as_deref().filter(|stored| !stored.is_empty()) != offered.as_deref() {
return Err(Error::ConflictingWorkflow {
workflow_id: workflow.workflow_id.to_owned(),
detail: format!(
"existing {field} is {stored:?}, but {offered:?} was provided"
),
});
}
}
let stored_queue: Option<String> = row.try_get("queue_name")?;
if stored_queue.as_deref() != workflow.queue_name {
tracing::warn!(
workflow_id = %workflow.workflow_id,
stored = ?stored_queue,
provided = ?workflow.queue_name,
"workflow already exists on a different queue; the stored queue is kept"
);
}
let should_execute = stored_creator.as_deref() == Some(creator_xid);
let inserted: bool = row.try_get("inserted")?;
if workflow.is_debounced && !should_execute {
let stored_deduplication_id: Option<String> = row.try_get("deduplication_id")?;
if let Some(deduplication_id) = workflow.deduplication_id
&& let Some(queue_name) = workflow.queue_name
&& stored_deduplication_id.as_deref() == Some(deduplication_id)
&& stored_queue.as_deref() == Some(queue_name)
{
return Err(Error::QueueDeduplicated {
workflow_id: workflow.workflow_id.to_owned(),
queue_name: queue_name.to_owned(),
deduplication_id: deduplication_id.to_owned(),
});
}
return Err(Error::ConflictingWorkflow {
workflow_id: workflow.workflow_id.to_owned(),
detail: "a debounce enqueues a new workflow, but a workflow with this id \
already exists"
.to_owned(),
});
}
if !should_execute {
tracing::debug!(
workflow_id = workflow.workflow_id,
"another caller created this workflow; recorded but not run"
);
}
if should_execute && inserted {
sqlx::query(AssertSqlSafe(format!(
"INSERT INTO {input_table} (workflow_uuid, inputs) VALUES ($1, $2) \
ON CONFLICT (workflow_uuid) DO UPDATE SET inputs = EXCLUDED.inputs"
)))
.bind(workflow.workflow_id)
.bind(workflow.input)
.execute(&mut *tx)
.await?;
}
if let Some(caller) = caller {
self.record_child_workflow_on(
&mut tx,
caller.parent_workflow_id,
workflow.workflow_id,
caller.step_id,
caller.step_name,
Some(caller.started_at),
)
.await?;
}
tx.commit().await?;
Ok(WorkflowInitResult {
status,
recovery_attempts,
deadline: deadline.map(Timestamp::from_epoch_ms),
clock,
serialization,
should_execute,
creator_xid: creator_xid.to_owned(),
})
})
.await
}
}
#[async_trait]
impl SystemDatabase for PostgresSystemDatabase {
async fn init_workflow(
&self,
workflow: &NewWorkflow,
caller: Option<InitWorkflowCaller<'_>>,
reuse_policy: WorkflowIdReusePolicy,
) -> Result<WorkflowInitResult, Error> {
let creator_xid = uuid::Uuid::new_v4().to_string();
self.init_workflow_as(workflow, caller, reuse_policy, &creator_xid)
.await
}
async fn database_time(&self) -> Result<DatabaseTime, Error> {
let pool = &self.pool;
with_retry(&self.retry, "database_time", move || async move {
let epoch_ms: i64 = sqlx::query_scalar(AssertSqlSafe(format!("SELECT {CLOCK_MS_SQL}")))
.fetch_one(pool)
.await?;
Ok(DatabaseTime::read(epoch_ms))
})
.await
}
async fn get_workflow(&self, workflow_id: &str) -> Result<Option<WorkflowRecord>, Error> {
let workflow_table = &self.tables.workflow_status;
let select = format!(
"SELECT {WORKFLOW_COLUMNS}, {} FROM {workflow_table} WHERE workflow_uuid = $1",
workflow_payloads(&self.tables, true, true)
);
let (select, pool) = (select.as_str(), &self.pool);
with_retry(&self.retry, "get_workflow", move || async move {
let row = sqlx::query(AssertSqlSafe(select))
.bind(workflow_id)
.fetch_optional(pool)
.await?;
row.as_ref().map(workflow_from_row).transpose()
})
.await
}
async fn list_workflows(
&self,
filter: &WorkflowFilter<'_>,
caller: Option<StepCaller<'_>>,
) -> Result<Vec<WorkflowRecord>, Error> {
let workflow_table = &self.tables.workflow_status;
let status: Vec<&str> = filter
.status
.iter()
.copied()
.map(WorkflowStatus::as_str)
.collect();
let prefixes: Vec<String> = filter
.workflow_id_prefixes
.iter()
.map(|p| {
format!(
"{}%",
p.replace('\\', r"\\")
.replace('%', r"\%")
.replace('_', r"\_")
)
})
.collect();
let workflow_table = workflow_table.as_str();
let (status, prefixes) = (status.as_slice(), prefixes.as_slice());
let application_name = self.application_name.as_deref();
let started_at = Timestamp::now();
with_retry(&self.retry, "list_workflows", move || async move {
self.run_single_statement_step(
caller,
step_names::LIST_WORKFLOWS,
started_at,
|mut conn| async move {
let mut q = sqlx::QueryBuilder::<sqlx::Postgres>::new("SELECT ");
q.push(WORKFLOW_COLUMNS).push(", ").push(workflow_payloads(
&self.tables,
filter.load_input,
filter.load_output,
));
q.push(" FROM ").push(workflow_table);
let mut first = true;
let mut clause = |q: &mut sqlx::QueryBuilder<sqlx::Postgres>, sql: &str| {
q.push(if first { " WHERE " } else { " AND " });
first = false;
q.push(sql);
};
macro_rules! any_of {
($values:expr, $column:literal) => {
if !$values.is_empty() {
clause(&mut q, concat!($column, " = ANY("));
q.push_bind($values).push(")");
}
};
}
any_of!(&filter.workflow_ids[..], "workflow_uuid");
match &filter.app_scope {
AppScope::Any => {}
AppScope::Named(names) if names.is_empty() => {}
AppScope::Named(names) => {
clause(&mut q, "(application_name = ANY(");
q.push_bind(&names[..])
.push(") OR application_name IS NULL)");
}
AppScope::Unset if !filter.workflow_ids.is_empty() => {}
AppScope::Unset => {
if let Some(name) = application_name {
clause(&mut q, "(application_name = ");
q.push_bind(name).push(" OR application_name IS NULL)");
}
}
}
any_of!(&filter.names[..], "name");
any_of!(&filter.class_names[..], "class_name");
any_of!(&filter.config_names[..], "config_name");
any_of!(status, "status");
any_of!(&filter.app_versions[..], "application_version");
any_of!(&filter.executor_ids[..], "executor_id");
any_of!(&filter.authenticated_users[..], "authenticated_user");
any_of!(&filter.queue_names[..], "queue_name");
any_of!(&filter.schedule_names[..], "schedule_name");
any_of!(&filter.deduplication_ids[..], "deduplication_id");
any_of!(&filter.parent_workflow_ids[..], "parent_workflow_id");
any_of!(&filter.forked_from[..], "forked_from");
if !prefixes.is_empty() {
clause(&mut q, "workflow_uuid LIKE ANY(");
q.push_bind(prefixes).push(")");
}
macro_rules! compare {
($value:expr, $sql:literal) => {
if let Some(v) = $value {
clause(&mut q, $sql);
q.push_bind(v.as_epoch_ms());
}
};
}
compare!(filter.created_after, "created_at >= ");
compare!(filter.created_before, "created_at <= ");
compare!(filter.completed_after, "completed_at >= ");
compare!(filter.completed_before, "completed_at <= ");
compare!(filter.started_after, "started_at_epoch_ms >= ");
compare!(filter.started_before, "started_at_epoch_ms <= ");
macro_rules! flag {
($value:expr, $sql:literal) => {
if let Some(v) = $value {
clause(&mut q, $sql);
q.push_bind(v);
}
};
}
flag!(filter.was_forked_from, "was_forked_from = ");
flag!(filter.is_debounced, "is_debounced = ");
if filter.queues_only {
clause(&mut q, "queue_name IS NOT NULL");
}
if let Some(is_fork) = filter.is_fork {
clause(
&mut q,
if is_fork {
"forked_from IS NOT NULL"
} else {
"forked_from IS NULL"
},
);
}
if let Some(has_parent) = filter.has_parent {
clause(
&mut q,
if has_parent {
"parent_workflow_id IS NOT NULL"
} else {
"parent_workflow_id IS NULL"
},
);
}
if let Some(attributes) = &filter.attributes {
clause(&mut q, "attributes @> ");
q.push_bind(attributes).push("::jsonb");
}
q.push(if filter.sort_desc {
" ORDER BY created_at DESC"
} else {
" ORDER BY created_at ASC"
});
if let Some(limit) = filter.limit {
q.push(" LIMIT ").push_bind(limit);
}
if let Some(offset) = filter.offset {
q.push(" OFFSET ").push_bind(offset);
}
let rows = q.build().fetch_all(&mut *conn).await?;
let workflows: Vec<WorkflowRecord> = rows
.iter()
.map(workflow_from_row)
.collect::<Result<_, _>>()?;
Ok((conn, workflows))
},
)
.await
})
.await
}
async fn get_workflow_children(&self, workflow_id: &str) -> Result<Vec<String>, Error> {
with_retry(&self.retry, "get_workflow_children", move || async move {
let mut conn = self.pool.acquire().await?;
let mut seen = std::collections::HashSet::new();
let mut descendants: Vec<String> = Vec::new();
let mut absorb = |into: &mut Vec<String>, children: Vec<String>| {
for child in children {
if child != workflow_id && seen.insert(child.clone()) {
into.push(child);
}
}
};
absorb(
&mut descendants,
self.direct_children(&mut conn, &[workflow_id]).await?,
);
let mut start = 0;
while start < descendants.len() {
let end = descendants.len();
let children = self
.direct_children(&mut conn, &descendants[start..end])
.await?;
start = end;
absorb(&mut descendants, children);
}
Ok(descendants)
})
.await
}
async fn record_workflow_outcome(
&self,
workflow_id: &str,
outcome: Outcome<'_>,
owner_xid: &str,
) -> Result<OutcomeWrite, Error> {
let (workflow_table, output_table, pool) = (
self.tables.workflow_status.as_str(),
self.tables.workflow_output.as_str(),
&self.pool,
);
let (output, error) = outcome.columns();
with_retry(&self.retry, "record_workflow_outcome", move || async move {
let mut tx = pool.begin().await?;
let updated = sqlx::query(AssertSqlSafe(format!(
"UPDATE {workflow_table} SET status = $2, output = NULL, error = NULL, \
updated_at = {NOW_MS_SQL}, completed_at = {NOW_MS_SQL}, \
deduplication_id = NULL, owner_xid = NULL \
WHERE workflow_uuid = $1 AND status = 'PENDING' \
AND owner_xid = $3"
)))
.bind(workflow_id)
.bind(outcome.status().as_str())
.bind(owner_xid)
.execute(&mut *tx)
.await?
.rows_affected();
if updated == 0 {
tracing::debug!(
workflow_id,
"outcome not recorded; the workflow is no longer this run's to finish"
);
return Ok(OutcomeWrite::AlreadyFinished);
}
sqlx::query(AssertSqlSafe(format!(
"INSERT INTO {output_table} (workflow_uuid, output, error) VALUES ($1, $2, $3) \
ON CONFLICT (workflow_uuid) DO UPDATE \
SET output = EXCLUDED.output, error = EXCLUDED.error"
)))
.bind(workflow_id)
.bind(output)
.bind(error)
.execute(&mut *tx)
.await?;
tx.commit().await?;
Ok(OutcomeWrite::Recorded)
})
.await
}
async fn await_workflow_result(
&self,
workflow_id: &str,
poll_interval: Duration,
fail_if_missing: bool,
) -> Result<AwaitedOutcome, Error> {
let workflow_table = &self.tables.workflow_status;
let (workflow_table, pool) = (workflow_table.as_str(), &self.pool);
let select = format!(
"SELECT status, {}, {}, serialization, recovery_attempts \
FROM {workflow_table} WHERE workflow_uuid = $1",
payload_column(&self.tables, PayloadColumn::Output),
payload_column(&self.tables, PayloadColumn::Error),
);
let select = &select;
let polling = &self.polling;
loop {
let row = with_retry(&self.retry, "await_workflow_result", move || async move {
let _permit = polling
.acquire()
.await
.expect("the polling limiter is never closed");
Ok(sqlx::query(AssertSqlSafe(select.clone()))
.bind(workflow_id)
.fetch_optional(pool)
.await?)
})
.await?;
match row {
Some(row) => {
let status = parse_workflow_status(&row.try_get::<String, _>("status")?)?;
let settled = match status {
WorkflowStatus::Success => Some(AwaitedOutcome::Succeeded {
output: row.try_get("output")?,
serialization: row.try_get("serialization")?,
}),
WorkflowStatus::Error => Some(AwaitedOutcome::Failed {
error: row.try_get::<Option<String>, _>("error")?.ok_or_else(|| {
Error::Malformed(format!(
"workflow {workflow_id} failed with no error recorded"
))
})?,
serialization: row.try_get("serialization")?,
}),
WorkflowStatus::Cancelled => Some(AwaitedOutcome::Cancelled),
WorkflowStatus::MaxRecoveryAttemptsExceeded => {
Some(AwaitedOutcome::Parked {
recovery_attempts: row
.try_get::<Option<i64>, _>("recovery_attempts")?
.unwrap_or_default(),
})
}
WorkflowStatus::Pending
| WorkflowStatus::Enqueued
| WorkflowStatus::Delayed => None,
};
if let Some(settled) = settled {
tracing::debug!(workflow_id, status = status.as_str(), "workflow settled");
return Ok(settled);
}
}
None if fail_if_missing => {
return Err(Error::NonExistentWorkflow {
workflow_ids: vec![workflow_id.to_owned()],
});
}
None => {}
}
crate::rt::sleep(poll_interval).await;
}
}
async fn await_first_workflow_id(
&self,
workflow_ids: &[&str],
poll_interval: Duration,
) -> Result<String, Error> {
if workflow_ids.is_empty() {
return Err(Error::InvalidInput {
field: "workflow_ids".into(),
detail: "must name at least one workflow to wait for".to_owned(),
});
}
let workflow_table = self.tables.workflow_status.as_str();
let (pool, polling) = (&self.pool, &self.polling);
let select = format!(
"SELECT workflow_uuid FROM {workflow_table} \
WHERE workflow_uuid = ANY($1) AND status NOT IN ({UNSETTLED}) \
LIMIT 1"
);
let select = &select;
let ids: Vec<String> = workflow_ids.iter().map(|id| (*id).to_owned()).collect();
let ids = &ids;
loop {
let row = with_retry(&self.retry, "await_first_workflow_id", move || async move {
let _permit = polling
.acquire()
.await
.expect("the polling limiter is never closed");
Ok(sqlx::query(AssertSqlSafe(select.clone()))
.bind(ids)
.fetch_optional(pool)
.await?)
})
.await?;
if let Some(row) = row {
let winner: String = row.try_get("workflow_uuid")?;
tracing::debug!(workflow_id = winner, "the first of the awaited set settled");
return Ok(winner);
}
crate::rt::sleep(poll_interval).await;
}
}
async fn await_workflow_ids(
&self,
workflow_ids: &[&str],
poll_interval: Duration,
) -> Result<(), Error> {
if workflow_ids.is_empty() {
return Ok(());
}
let workflow_table = self.tables.workflow_status.as_str();
let (pool, polling) = (&self.pool, &self.polling);
let select = format!(
"SELECT workflow_uuid FROM {workflow_table} \
WHERE workflow_uuid = ANY($1) AND status NOT IN ({UNSETTLED})"
);
let select = &select;
let mut outstanding: HashSet<String> =
workflow_ids.iter().map(|id| (*id).to_owned()).collect();
while !outstanding.is_empty() {
let ids: Vec<&str> = outstanding.iter().map(String::as_str).collect();
let ids = &ids;
let settled = with_retry(&self.retry, "await_workflow_ids", move || async move {
let _permit = polling
.acquire()
.await
.expect("the polling limiter is never closed");
Ok(sqlx::query(AssertSqlSafe(select.clone()))
.bind(ids)
.fetch_all(pool)
.await?)
})
.await?;
for row in &settled {
let id: String = row.try_get("workflow_uuid")?;
outstanding.remove(&id);
}
if outstanding.is_empty() {
break;
}
crate::rt::sleep(poll_interval).await;
}
tracing::debug!("every awaited workflow settled");
Ok(())
}
async fn set_workflow_delay(
&self,
workflow_id: &str,
delay: WorkflowDelay,
caller: Option<StepCaller<'_>>,
) -> Result<(), Error> {
let workflow_table = &self.tables.workflow_status;
let delay = match delay {
WorkflowDelay::For(duration) => {
duration_to_ms("delay", duration).map(|ms| (None, Some(ms)))
}
WorkflowDelay::Until(instant) => Ok((Some(instant.as_epoch_ms()), None)),
};
let delay = &delay;
let delay_until_from_delay = now_ms_plus_sql("$3::bigint");
let delay_until_from_delay = delay_until_from_delay.as_str();
let workflow_table = workflow_table.as_str();
let started_at = Timestamp::now();
with_retry(&self.retry, "set_workflow_delay", move || async move {
self.run_single_statement_step(
caller,
step_names::SET_WORKFLOW_DELAY,
started_at,
|mut conn| async move {
let (until_ms, delay_ms) = delay.clone()?;
sqlx::query(AssertSqlSafe(format!(
"UPDATE {workflow_table} \
SET delay_until_epoch_ms = COALESCE($2::bigint, {delay_until_from_delay}), \
updated_at = {NOW_MS_SQL} \
WHERE workflow_uuid = $1 AND status = 'DELAYED'"
)))
.bind(workflow_id)
.bind(until_ms)
.bind(delay_ms)
.execute(&mut *conn)
.await?;
Ok((conn, ()))
},
)
.await
})
.await
}
async fn update_workflow_attributes(
&self,
workflow_id: &str,
update: AttributeUpdate<'_>,
caller: Option<StepCaller<'_>>,
) -> Result<(), Error> {
update.validate()?;
let workflow_table = self.tables.workflow_status.as_str();
let started_at = Timestamp::now();
let (set, attributes) = match update {
AttributeUpdate::Replace(attributes) => ("$2::jsonb", attributes),
AttributeUpdate::Merge(attributes) => (
"CASE WHEN jsonb_typeof(attributes) = 'object' \
THEN attributes ELSE '{}'::jsonb END || $2::jsonb",
Some(attributes),
),
};
with_retry(
&self.retry,
"update_workflow_attributes",
move || async move {
self.run_single_statement_step(
caller,
step_names::UPDATE_WORKFLOW_ATTRIBUTES,
started_at,
|mut conn| async move {
let updated = sqlx::query(AssertSqlSafe(format!(
"UPDATE {workflow_table} \
SET attributes = {set}, updated_at = {NOW_MS_SQL} \
WHERE workflow_uuid = $1"
)))
.bind(workflow_id)
.bind(attributes)
.execute(&mut *conn)
.await?;
if updated.rows_affected() == 0 {
return Err(Error::NonExistentWorkflow {
workflow_ids: vec![workflow_id.to_owned()],
});
}
Ok((conn, ()))
},
)
.await
},
)
.await
}
async fn reenqueue_for_recovery(
&self,
executor_ids: &[&str],
application_version: &str,
recovery_queue: &str,
) -> Result<Vec<String>, Error> {
if executor_ids.is_empty() {
return Ok(Vec::new());
}
let workflow_table = &self.tables.workflow_status;
let (workflow_table, pool) = (workflow_table.as_str(), &self.pool);
let application_name = self.application_name.as_deref();
with_retry(&self.retry, "reenqueue_for_recovery", move || async move {
let ids: Vec<String> = sqlx::query_scalar(AssertSqlSafe(format!(
"UPDATE {workflow_table} \
SET status = 'ENQUEUED', \
started_at_epoch_ms = NULL, \
owner_xid = NULL, \
updated_at = {NOW_MS_SQL}, \
queue_name = COALESCE(NULLIF(queue_name, ''), $1) \
WHERE status = 'PENDING' \
AND executor_id = ANY($2) \
AND application_version = $3 \
AND ($4::text IS NULL \
OR application_name = $4 \
OR application_name IS NULL) \
RETURNING workflow_uuid"
)))
.bind(recovery_queue)
.bind(executor_ids)
.bind(application_version)
.bind(application_name)
.fetch_all(pool)
.await?;
Ok(ids)
})
.await
}
async fn dead_letter_workflows(
&self,
workflow_ids: &[&str],
min_recovery_attempts: i64,
) -> Result<Vec<String>, Error> {
if workflow_ids.is_empty() {
return Ok(Vec::new());
}
let workflow_table = &self.tables.workflow_status;
let (workflow_table, pool) = (workflow_table.as_str(), &self.pool);
with_retry(&self.retry, "dead_letter_workflows", move || async move {
let ids: Vec<String> = sqlx::query_scalar(AssertSqlSafe(format!(
"UPDATE {workflow_table} \
SET status = 'MAX_RECOVERY_ATTEMPTS_EXCEEDED', owner_xid = NULL, \
deduplication_id = NULL, started_at_epoch_ms = NULL, queue_name = NULL, \
updated_at = {NOW_MS_SQL}, completed_at = {NOW_MS_SQL} \
WHERE workflow_uuid = ANY($1) AND status = 'PENDING' \
AND recovery_attempts >= $2 \
RETURNING workflow_uuid"
)))
.bind(workflow_ids)
.bind(min_recovery_attempts)
.fetch_all(pool)
.await?;
Ok(ids)
})
.await
}
async fn cancel_if_timed_out(&self, workflow_id: &str, owner_xid: &str) -> Result<bool, Error> {
let workflow_table = &self.tables.workflow_status;
let (workflow_table, pool) = (workflow_table.as_str(), &self.pool);
with_retry(&self.retry, "cancel_if_timed_out", move || async move {
let moved = sqlx::query(AssertSqlSafe(format!(
"UPDATE {workflow_table} SET {set} \
WHERE workflow_uuid = $1 AND owner_xid = $2 AND {timed_out}",
set = timed_out_set_sql(),
timed_out = timed_out_where_sql(),
)))
.bind(workflow_id)
.bind(owner_xid)
.execute(pool)
.await?
.rows_affected();
Ok(moved > 0)
})
.await
}
async fn cancel_timed_out_workflows(&self, limit: i64) -> Result<Vec<String>, Error> {
let workflow_table = &self.tables.workflow_status;
let (workflow_table, pool) = (workflow_table.as_str(), &self.pool);
let application_name = self.application_name.as_deref();
with_retry(
&self.retry,
"cancel_timed_out_workflows",
move || async move {
let cancelled: Vec<String> =
sqlx::query_scalar(AssertSqlSafe(timed_out_sweep_sql(workflow_table)))
.bind(application_name)
.bind(limit)
.fetch_all(pool)
.await?;
Ok(cancelled)
},
)
.await
}
async fn transition_delayed_workflows(&self) -> Result<u64, Error> {
let workflow_table = &self.tables.workflow_status;
let (workflow_table, pool) = (workflow_table.as_str(), &self.pool);
let application_name = self.application_name.as_deref();
with_retry(
&self.retry,
"transition_delayed_workflows",
move || async move {
let moved = sqlx::query(AssertSqlSafe(format!(
"UPDATE {workflow_table} SET status = 'ENQUEUED', updated_at = {NOW_MS_SQL}, \
deduplication_id = CASE WHEN is_debounced THEN NULL \
ELSE deduplication_id END \
WHERE status = 'DELAYED' AND delay_until_epoch_ms <= {NOW_MS_SQL} \
AND ($1::text IS NULL \
OR application_name = $1 \
OR application_name IS NULL)"
)))
.bind(application_name)
.execute(pool)
.await?
.rows_affected();
if moved > 0 {
tracing::debug!(moved, "released delayed workflows");
}
Ok(moved)
},
)
.await
}
async fn cancel_workflows(
&self,
workflow_ids: &[&str],
cancel_children: bool,
from_workflow: Option<&str>,
caller: Option<StepCaller<'_>>,
) -> Result<Vec<String>, Error> {
if workflow_ids.is_empty() && caller.is_none() {
return Ok(Vec::new());
}
let started_at = Timestamp::now();
let cancelled = with_retry(&self.retry, "cancel_workflows", move || async move {
self.run_transactional_step(
caller,
step_names::CANCEL_WORKFLOW,
started_at,
|mut tx| async move {
if let Some(caller_id) = from_workflow {
refuse_self_target(caller_id, workflow_ids, "cancel")?;
}
let mut cancelled = self.cancel_batch(&mut tx, workflow_ids).await?;
if !cancel_children {
return Ok((tx, cancelled));
}
let mut seen: std::collections::HashSet<String> =
workflow_ids.iter().map(|id| (*id).to_owned()).collect();
let mut frontier = self.direct_children(&mut tx, workflow_ids).await?;
frontier.retain(|c| seen.insert(c.clone()));
while !frontier.is_empty() {
if let Some(caller_id) = from_workflow
&& frontier.iter().any(|id| id == caller_id)
{
let root = self
.root_containing(&mut tx, caller_id, workflow_ids)
.await?;
return Err(Error::InvalidInput {
field: "workflow_ids".into(),
detail: format!(
"workflow {caller_id} cannot cancel the tree of {root}, which \
contains it"
),
});
}
cancelled.extend(self.cancel_batch(&mut tx, &frontier).await?);
let children = self.direct_children(&mut tx, &frontier).await?;
frontier = children
.into_iter()
.filter(|c| seen.insert(c.clone()))
.collect();
}
Ok((tx, cancelled))
},
)
.await
})
.await?;
tracing::debug!(
requested = workflow_ids.len(),
cancelled = cancelled.len(),
cancel_children,
"cancelled workflows"
);
Ok(cancelled)
}
async fn resume_workflows(
&self,
workflow_ids: &[&str],
queue_name: Option<&str>,
from_workflow: Option<&str>,
caller: Option<StepCaller<'_>>,
) -> Result<Vec<String>, Error> {
if workflow_ids.is_empty() && caller.is_none() {
return Ok(Vec::new());
}
let workflow_table = self.tables.workflow_status.as_str();
let queue = queue_name.unwrap_or(INTERNAL_QUEUE);
let started_at = Timestamp::now();
with_retry(&self.retry, "resume_workflows", move || async move {
self.run_transactional_step(
caller,
step_names::RESUME_WORKFLOW,
started_at,
|mut tx| async move {
if let Some(caller_id) = from_workflow {
refuse_self_target(caller_id, workflow_ids, "resume")?;
}
let existing: Vec<String> = sqlx::query_scalar(AssertSqlSafe(format!(
"SELECT workflow_uuid FROM {workflow_table} WHERE workflow_uuid = ANY($1)"
)))
.bind(workflow_ids)
.fetch_all(&mut *tx)
.await?;
let missing: Vec<String> = workflow_ids
.iter()
.filter(|id| !existing.iter().any(|e| e == *id))
.map(|id| (*id).to_owned())
.collect();
if !missing.is_empty() {
return Err(Error::NonExistentWorkflow {
workflow_ids: missing,
});
}
let resumed: Vec<String> = sqlx::query_scalar(AssertSqlSafe(format!(
"UPDATE {workflow_table} SET status = 'ENQUEUED', queue_name = $2, \
recovery_attempts = 0, workflow_deadline_epoch_ms = NULL, owner_xid = NULL, \
deduplication_id = NULL, started_at_epoch_ms = NULL, completed_at = NULL, \
updated_at = {NOW_MS_SQL} \
WHERE workflow_uuid = ANY($1) AND status NOT IN ('SUCCESS', 'ERROR') \
RETURNING workflow_uuid"
)))
.bind(workflow_ids)
.bind(queue)
.fetch_all(&mut *tx)
.await?;
tracing::debug!(
requested = workflow_ids.len(),
resumed = resumed.len(),
queue,
"resumed workflows"
);
Ok((tx, resumed))
},
)
.await
})
.await
}
async fn delete_workflows(
&self,
workflow_ids: &[&str],
delete_children: bool,
caller: Option<StepCaller<'_>>,
) -> Result<u64, Error> {
if workflow_ids.is_empty() && caller.is_none() {
return Ok(0);
}
let mut children: Vec<String> = Vec::new();
if delete_children {
for id in workflow_ids {
children.extend(self.get_workflow_children(id).await?);
}
}
let mut targets: Vec<&str> = workflow_ids.to_vec();
targets.extend(children.iter().map(String::as_str));
targets.sort_unstable();
targets.dedup();
if let Some(StepCaller {
workflow_id: caller_id,
..
}) = caller
&& targets.contains(&caller_id)
{
return Err(Error::InvalidInput {
field: "workflow_ids".into(),
detail: format!(
"workflow {caller_id} cannot delete itself: the step checkpoint for the delete \
is written in the same transaction and would outlive the row it references"
),
});
}
let tables = &self.tables;
let targets = targets.as_slice();
let started_at = Timestamp::now();
with_retry(&self.retry, "delete_workflows", move || async move {
self.run_transactional_step(
caller,
step_names::DELETE_WORKFLOW,
started_at,
|mut tx| async move {
for table in [
&tables.workflow_input,
&tables.workflow_output,
&tables.operation_outputs,
] {
sqlx::query(AssertSqlSafe(format!(
"DELETE FROM {table} WHERE workflow_uuid = ANY($1)"
)))
.bind(targets)
.execute(&mut *tx)
.await?;
}
let workflow_table = &tables.workflow_status;
let deleted = sqlx::query(AssertSqlSafe(format!(
"DELETE FROM {workflow_table} WHERE workflow_uuid = ANY($1)"
)))
.bind(targets)
.execute(&mut *tx)
.await?
.rows_affected();
tracing::debug!(deleted, targets = targets.len(), "deleted workflows");
Ok((tx, deleted))
},
)
.await
})
.await
}
async fn fork_workflows(
&self,
forks: &[Fork<'_>],
options: &ForkOptions<'_>,
caller: Option<StepCaller<'_>>,
) -> Result<Vec<String>, Error> {
if forks.is_empty() && caller.is_none() {
return Ok(Vec::new());
}
options.validate()?;
for fork in forks {
fork.validate()?;
}
let forked_ids: Vec<String> = forks
.iter()
.map(|f| {
f.forked_id
.map_or_else(|| uuid::Uuid::new_v4().to_string(), str::to_owned)
})
.collect();
let original_ids: Vec<&str> = forks.iter().map(|f| f.original_id).collect();
let start_steps: Vec<i32> = forks.iter().map(|f| f.start_step).collect();
self.fork_batch(
&original_ids,
&forked_ids,
StartSteps::Given(&start_steps),
options,
caller,
)
.await
}
async fn fork_from(
&self,
workflow_ids: &[&str],
point: ForkPoint<'_>,
options: &ForkOptions<'_>,
caller: Option<StepCaller<'_>>,
) -> Result<Vec<String>, Error> {
if workflow_ids.is_empty() && caller.is_none() {
return Ok(Vec::new());
}
const LAST_STEP: &str = "MAX(function_id)";
const LAST_FAILURE: &str =
"COALESCE(MAX(function_id) FILTER (WHERE error IS NOT NULL), MAX(function_id))";
options.validate()?;
let given_step = match point {
ForkPoint::Step(step) => step,
_ => 0,
};
for &original_id in workflow_ids {
Fork {
original_id,
forked_id: None,
start_step: given_step,
}
.validate()?;
}
let forked_ids: Vec<String> = workflow_ids
.iter()
.map(|_| uuid::Uuid::new_v4().to_string())
.collect();
let given = vec![given_step; workflow_ids.len()];
let steps = match point {
ForkPoint::Step(_) => StartSteps::Given(&given),
ForkPoint::LastFailure => StartSteps::Resolve {
aggregate: LAST_FAILURE,
named: None,
},
ForkPoint::LastStep => StartSteps::Resolve {
aggregate: LAST_STEP,
named: None,
},
ForkPoint::StepNamed(name) => StartSteps::Resolve {
aggregate: LAST_STEP,
named: Some(name),
},
};
self.fork_batch(workflow_ids, &forked_ids, steps, options, caller)
.await
}
async fn rewind_workflow(
&self,
workflow_id: &str,
start_step: i32,
options: &RewindOptions<'_>,
caller: Option<StepCaller<'_>>,
) -> Result<(), Error> {
if start_step < 0 {
return Err(Error::InvalidInput {
field: "start_step".into(),
detail: format!("must not be negative, got {start_step}"),
});
}
options.validate()?;
let tables = &self.tables;
let queue = options.queue_name.unwrap_or(INTERNAL_QUEUE);
let started_at = Timestamp::now();
with_retry(&self.retry, "rewind_workflow", move || async move {
self.run_transactional_step(
caller,
step_names::REWIND_WORKFLOW,
started_at,
|mut tx| async move {
let workflow_table = tables.workflow_status.as_str();
let events_table = tables.workflow_events.as_str();
let history_table = tables.workflow_events_history.as_str();
let streams_table = tables.streams.as_str();
let steps_table = tables.operation_outputs.as_str();
let notifications_table = tables.notifications.as_str();
let output_table = tables.workflow_output.as_str();
let status: Option<String> = sqlx::query_scalar(AssertSqlSafe(format!(
"SELECT status FROM {workflow_table} WHERE workflow_uuid = $1"
)))
.bind(workflow_id)
.fetch_optional(&mut *tx)
.await?;
let Some(status) = status else {
return Err(Error::NonExistentWorkflow {
workflow_ids: vec![workflow_id.to_owned()],
});
};
let status = parse_workflow_status(&status)?;
if status.is_active() {
return Err(Error::WorkflowNotRewindable {
workflow_id: workflow_id.to_owned(),
status: status.as_str().to_owned(),
});
}
let published_past_cut = format!(
"SELECT key FROM {history_table} \
WHERE workflow_uuid = $1 AND function_id >= $2"
);
sqlx::query(AssertSqlSafe(format!(
"DELETE FROM {events_table} \
WHERE workflow_uuid = $1 AND key IN ({published_past_cut})"
)))
.bind(workflow_id)
.bind(start_step)
.execute(&mut *tx)
.await?;
sqlx::query(AssertSqlSafe(format!(
"INSERT INTO {events_table} (workflow_uuid, key, value, serialization) \
SELECT workflow_uuid, key, value, serialization FROM ( \
SELECT h.workflow_uuid, h.key, h.value, h.serialization, \
row_number() OVER (PARTITION BY h.key \
ORDER BY h.function_id DESC) AS rn \
FROM {history_table} h \
WHERE h.workflow_uuid = $1 AND h.function_id < $2 \
AND h.key IN ({published_past_cut}) \
) latest WHERE rn = 1"
)))
.bind(workflow_id)
.bind(start_step)
.execute(&mut *tx)
.await?;
sqlx::query(AssertSqlSafe(format!(
"DELETE FROM {streams_table} \
WHERE workflow_uuid = $1 AND function_id >= $2 \
AND value = $3 AND serialization = $4"
)))
.bind(workflow_id)
.bind(start_step)
.bind(STREAM_CLOSED)
.bind(PORTABLE_JSON)
.execute(&mut *tx)
.await?;
for table in [steps_table, history_table] {
sqlx::query(AssertSqlSafe(format!(
"DELETE FROM {table} WHERE workflow_uuid = $1 AND function_id >= $2"
)))
.bind(workflow_id)
.bind(start_step)
.execute(&mut *tx)
.await?;
}
sqlx::query(AssertSqlSafe(format!(
"DELETE FROM {notifications_table} \
WHERE destination_uuid = $1 \
AND (consumed_by_function_id >= $2 OR consumed = FALSE)"
)))
.bind(workflow_id)
.bind(start_step)
.execute(&mut *tx)
.await?;
sqlx::query(AssertSqlSafe(format!(
"DELETE FROM {output_table} WHERE workflow_uuid = $1"
)))
.bind(workflow_id)
.execute(&mut *tx)
.await?;
let updated = sqlx::query(AssertSqlSafe(format!(
"UPDATE {workflow_table} SET status = 'ENQUEUED', queue_name = $3, \
queue_partition_key = $4, \
application_version = COALESCE($5, application_version), \
recovery_attempts = 0, owner_xid = NULL, \
workflow_deadline_epoch_ms = NULL, deduplication_id = NULL, \
started_at_epoch_ms = NULL, completed_at = NULL, \
output = NULL, error = NULL, updated_at = {NOW_MS_SQL} \
WHERE workflow_uuid = $1 AND status = $2"
)))
.bind(workflow_id)
.bind(status.as_str())
.bind(queue)
.bind(options.queue_partition_key)
.bind(options.application_version)
.execute(&mut *tx)
.await?
.rows_affected();
if updated != 1 {
return Err(Error::RewindInterrupted {
workflow_id: workflow_id.to_owned(),
status: status.as_str().to_owned(),
});
}
tracing::debug!(workflow_id, start_step, queue, "rewound workflow");
Ok((tx, ()))
},
)
.await
})
.await
}
async fn send_message(
&self,
message: &Message<'_>,
serialization: Option<&str>,
caller: Option<(&str, i32)>,
send_to_forks: bool,
owner: Option<WorkflowOwner<'_>>,
) -> Result<(), Error> {
self.deliver(
std::slice::from_ref(message),
serialization,
caller,
step_names::SEND,
send_to_forks,
owner,
)
.await
}
async fn send_messages(
&self,
messages: &[Message<'_>],
serialization: Option<&str>,
caller: Option<(&str, i32)>,
send_to_forks: bool,
owner: Option<WorkflowOwner<'_>>,
) -> Result<(), Error> {
self.deliver(
messages,
serialization,
caller,
step_names::SEND_BULK,
send_to_forks,
owner,
)
.await
}
async fn recv(
&self,
workflow_id: &str,
step_id: i32,
timeout_step_id: i32,
topic: Option<&str>,
timeout: Duration,
owner_xid: &str,
) -> Result<Option<EncodedValue>, Error> {
let notifications_table = self.tables.notifications.as_str();
let (pool, polling) = (&self.pool, &self.polling);
let stored_topic = topic.unwrap_or(NULL_TOPIC);
let probe = format!(
"SELECT 1 FROM {notifications_table} \
WHERE destination_uuid = $1 AND topic = $2 AND consumed = FALSE LIMIT 1"
);
let consume = format!(
"UPDATE {notifications_table} SET consumed = TRUE, consumed_by_function_id = $3 \
WHERE message_uuid = ( \
SELECT message_uuid FROM {notifications_table} \
WHERE destination_uuid = $1 AND topic = $2 AND consumed = FALSE \
ORDER BY created_at_epoch_ms ASC LIMIT 1 \
) AND consumed = FALSE \
RETURNING message, serialization"
);
let (probe, consume) = (&probe, &consume);
if let Some(step) = self
.check_step(workflow_id, step_id, step_names::RECV)
.await?
{
tracing::debug!(
workflow_id,
step_id = step_id,
topic = stored_topic,
"replaying recv"
);
return Ok(step.output.map(|value| EncodedValue {
value,
serialization: step.serialization,
}));
}
let mut subscription = self.notify.subscribe(message_key(workflow_id, topic));
let deadline = self
.checkpoint_sleep(
workflow_id,
timeout_step_id,
timeout,
SleepKind::Deadline,
owner_xid,
)
.await?;
let until = deadline.arrives();
let started_at = deadline.clock.now();
enum Take {
Answered(Option<EncodedValue>),
Missed,
}
loop {
let has_message = with_retry(&self.retry, "recv", move || async move {
let _permit = polling
.acquire()
.await
.expect("the polling limiter is never closed");
Ok(sqlx::query(AssertSqlSafe(probe.clone()))
.bind(workflow_id)
.bind(stored_topic)
.fetch_optional(pool)
.await?
.is_some())
})
.await?;
let remaining = until.saturating_duration_since(crate::rt::now());
if has_message || remaining.is_zero() {
let timing = StepTiming {
started_at,
completed_at: deadline.clock.now(),
};
let take = with_retry(&self.retry, "recv", move || async move {
let mut tx = pool.begin().await?;
self.check_owner_on(
&mut tx,
WorkflowOwner {
workflow_id,
owner_xid,
},
)
.await?;
if let Some(step) = self
.check_step_on(&mut tx, workflow_id, step_id, step_names::RECV)
.await?
{
tx.commit().await?;
tracing::debug!(workflow_id, step_id = step_id, "replaying recv");
return Ok(Take::Answered(step.output.map(|value| EncodedValue {
value,
serialization: step.serialization,
})));
}
let row = sqlx::query(AssertSqlSafe(consume.clone()))
.bind(workflow_id)
.bind(stored_topic)
.bind(step_id)
.fetch_optional(&mut *tx)
.await?;
let message = match row {
Some(row) => Some(EncodedValue {
value: row.try_get("message")?,
serialization: row.try_get("serialization")?,
}),
None => None,
};
if message.is_none() {
let time_left = until > crate::rt::now();
if time_left {
tx.rollback().await?;
tracing::debug!(
workflow_id,
step_id = step_id,
topic = stored_topic,
"recv found its message taken; waiting on"
);
return Ok(Take::Missed);
}
}
self.record_step_on(
&mut tx,
workflow_id,
step_id,
step_names::RECV,
Outcome::Output(message.as_ref().map(|m| m.value.as_str())),
message.as_ref().and_then(|m| m.serialization.as_deref()),
Some(timing),
None,
)
.await?;
tx.commit().await?;
tracing::debug!(
workflow_id,
step_id = step_id,
topic = stored_topic,
received = message.is_some(),
"recv finished"
);
Ok(Take::Answered(message))
})
.await?;
if let Take::Answered(answer) = take {
return Ok(answer);
}
continue;
}
let _ = crate::rt::timeout(
remaining.min(self.listener.poll_interval()),
subscription.notified(),
)
.await;
}
}
async fn write_stream(
&self,
workflow_id: &str,
step_id: i32,
key: &str,
value: &str,
serialization: Option<&str>,
written_by: WrittenBy,
owner_xid: &str,
) -> Result<(), Error> {
let step_name = if is_stream_closed(value, serialization) {
step_names::CLOSE_STREAM
} else {
step_names::WRITE_STREAM
};
let timing = StepTiming {
started_at: Timestamp::now(),
completed_at: Timestamp::now(),
};
let streams_table = self.tables.streams.as_str();
let pool = &self.pool;
let insert = format!(
"INSERT INTO {streams_table} \
(workflow_uuid, function_id, key, value, serialization, \"offset\") \
SELECT $1, $2, $3, $4, $5, COALESCE( \
(SELECT MAX(\"offset\") FROM {streams_table} \
WHERE workflow_uuid = $1 AND key = $3), -1) + 1"
);
let insert = &insert;
with_retry(&self.retry, "write_stream", move || async move {
for attempt in 0..STREAM_OFFSET_ATTEMPTS {
let mut tx = pool.begin().await?;
self.check_owner_on(
&mut tx,
WorkflowOwner {
workflow_id,
owner_xid,
},
)
.await?;
if written_by == WrittenBy::Workflow
&& self
.check_step_on(&mut tx, workflow_id, step_id, step_name)
.await?
.is_some()
{
tracing::debug!(workflow_id, step_id, key, "replaying stream write");
return Ok(());
}
match sqlx::query(AssertSqlSafe(insert.clone()))
.bind(workflow_id)
.bind(step_id)
.bind(key)
.bind(value)
.bind(serialization)
.execute(&mut *tx)
.await
{
Ok(_) => {}
Err(e) if is_unique_violation(&e) => {
tracing::debug!(
workflow_id,
key,
attempt,
"stream offset taken; recomputing"
);
continue;
}
Err(e) if is_foreign_key_violation(&e) => {
return Err(Error::NonExistentWorkflow {
workflow_ids: vec![workflow_id.to_owned()],
});
}
Err(e) => return Err(e.into()),
}
if written_by == WrittenBy::Workflow {
self.record_step_on(
&mut tx,
workflow_id,
step_id,
step_name,
Outcome::Output(None),
None,
Some(timing),
None,
)
.await?;
}
tx.commit().await?;
self.notifier.signal(STREAMS_CHANNEL, workflow_id, key);
return Ok(());
}
Err(Error::Backend(BackendError {
message: format!(
"stream {key} on workflow {workflow_id} lost \
{STREAM_OFFSET_ATTEMPTS} offset races"
),
sqlstate: None,
kind: BackendErrorKind::Conflict,
}))
})
.await
}
async fn close_stream(
&self,
workflow_id: &str,
step_id: i32,
key: &str,
written_by: WrittenBy,
owner_xid: &str,
) -> Result<(), Error> {
self.write_stream(
workflow_id,
step_id,
key,
STREAM_CLOSED,
Some(PORTABLE_JSON),
written_by,
owner_xid,
)
.await
}
#[cfg(test)]
fn is_pushing(&self) -> bool {
PostgresSystemDatabase::is_pushing(self)
}
#[cfg(test)]
fn is_delivering(&self) -> bool {
PostgresSystemDatabase::is_delivering(self)
}
async fn close(&self) {
self.notifier.stop();
let task = self.notifier_task.lock().expect("notifier lock").take();
if let Some(task) = task {
let _ = task.await;
}
self.pool.close().await;
let task = self.listener_task.lock().expect("listener lock").take();
if let Some(task) = task {
let _ = task.await;
}
}
async fn check_step(
&self,
workflow_id: &str,
step_id: i32,
step_name: &str,
) -> Result<Option<StepRecord>, Error> {
let pool = &self.pool;
with_retry(&self.retry, "check_step", move || async move {
let mut conn = pool.acquire().await?;
self.check_step_on(&mut conn, workflow_id, step_id, step_name)
.await
})
.await
}
async fn check_step_at_or_after(
&self,
workflow_id: &str,
step_id: i32,
) -> Result<Option<(i32, String)>, Error> {
let workflow_table = &self.tables.workflow_status;
let steps_table = &self.tables.operation_outputs;
let (workflow_table, steps_table, pool) =
(workflow_table.as_str(), steps_table.as_str(), &self.pool);
with_retry(&self.retry, "check_step_at_or_after", move || async move {
let row = sqlx::query(AssertSqlSafe(format!(
"SELECT s.status, o.function_id, o.function_name \
FROM {workflow_table} s \
LEFT JOIN ( \
SELECT function_id, function_name FROM {steps_table} \
WHERE workflow_uuid = $1 AND function_id >= $2 \
ORDER BY function_id LIMIT 1 \
) o ON TRUE \
WHERE s.workflow_uuid = $1"
)))
.bind(workflow_id)
.bind(step_id)
.fetch_optional(pool)
.await?;
let Some(row) = row else {
return Err(Error::NonExistentWorkflow {
workflow_ids: vec![workflow_id.to_owned()],
});
};
let status: String = row.try_get("status")?;
if status == WorkflowStatus::Cancelled.as_str() {
return Err(Error::WorkflowCancelled {
workflow_id: workflow_id.to_owned(),
});
}
let Some(found) = row.try_get::<Option<i32>, _>("function_id")? else {
return Ok(None);
};
Ok(Some((found, row.try_get("function_name")?)))
})
.await
}
async fn record_step(
&self,
workflow_id: &str,
step_id: i32,
step_name: &str,
outcome: Outcome<'_>,
serialization: Option<&str>,
timing: Option<StepTiming>,
owner_xid: &str,
) -> Result<(), Error> {
let pool = &self.pool;
with_retry(&self.retry, "record_step", move || async move {
let mut tx = pool.begin().await?;
self.check_owner_on(
&mut tx,
WorkflowOwner {
workflow_id,
owner_xid,
},
)
.await?;
self.record_step_on(
&mut tx,
workflow_id,
step_id,
step_name,
outcome,
serialization,
timing,
None,
)
.await?;
tx.commit().await?;
Ok(())
})
.await
}
#[allow(clippy::too_many_arguments)]
async fn record_child_result(
&self,
parent_workflow_id: &str,
step_id: i32,
child_workflow_id: &str,
outcome: Outcome<'_>,
serialization: Option<&str>,
timing: Option<StepTiming>,
owner_xid: &str,
) -> Result<(), Error> {
let pool = &self.pool;
with_retry(&self.retry, "record_child_result", move || async move {
let mut tx = pool.begin().await?;
self.check_owner_on(
&mut tx,
WorkflowOwner {
workflow_id: parent_workflow_id,
owner_xid,
},
)
.await?;
self.record_step_on(
&mut tx,
parent_workflow_id,
step_id,
step_names::GET_RESULT,
outcome,
serialization,
timing,
Some(child_workflow_id),
)
.await?;
tx.commit().await?;
Ok(())
})
.await
}
async fn list_workflow_steps(
&self,
workflow_id: &str,
load_output: bool,
limit: Option<i64>,
offset: Option<i64>,
caller: Option<StepCaller<'_>>,
) -> Result<Vec<StepRecord>, Error> {
let steps_table = self.tables.operation_outputs.as_str();
let started_at = Timestamp::now();
with_retry(&self.retry, "list_workflow_steps", move || async move {
self.run_single_statement_step(
caller,
step_names::LIST_WORKFLOW_STEPS,
started_at,
|mut conn| async move {
let mut q =
step_listing_query(steps_table, workflow_id, load_output, limit, offset);
let rows = q.build().fetch_all(&mut *conn).await?;
let steps: Vec<StepRecord> = rows
.iter()
.map(|r| step_from_row(r, workflow_id))
.collect::<Result<_, _>>()?;
Ok((conn, steps))
},
)
.await
})
.await
}
async fn record_sleep(
&self,
caller: StepCaller<'_>,
duration: Duration,
) -> Result<WakeTime, Error> {
self.checkpoint_sleep(
caller.workflow_id,
caller.step_id,
duration,
SleepKind::Durable(caller.codec),
caller.owner_xid,
)
.await
}
async fn set_event(
&self,
workflow_id: &str,
step_id: i32,
key: &str,
value: &str,
serialization: Option<&str>,
owner_xid: &str,
) -> Result<(), Error> {
let events = &self.tables.workflow_events;
let history = &self.tables.workflow_events_history;
let (events, history, pool) = (events.as_str(), history.as_str(), &self.pool);
let timing = StepTiming {
started_at: Timestamp::now(),
completed_at: Timestamp::now(),
};
with_retry(&self.retry, "set_event", move || async move {
let mut tx = pool.begin().await?;
self.check_owner_on(
&mut tx,
WorkflowOwner {
workflow_id,
owner_xid,
},
)
.await?;
if self
.check_step_on(&mut tx, workflow_id, step_id, step_names::SET_EVENT)
.await?
.is_some()
{
tracing::debug!(workflow_id, step_id, key, "replaying set_event");
return Ok(());
}
tracing::debug!(workflow_id, step_id, key, "running set_event");
sqlx::query(AssertSqlSafe(format!(
"INSERT INTO {events} (workflow_uuid, key, value, serialization) \
VALUES ($1, $2, $3, $4) \
ON CONFLICT (workflow_uuid, key) DO UPDATE \
SET value = EXCLUDED.value, serialization = EXCLUDED.serialization"
)))
.bind(workflow_id)
.bind(key)
.bind(value)
.bind(serialization)
.execute(&mut *tx)
.await?;
sqlx::query(AssertSqlSafe(format!(
"INSERT INTO {history} (workflow_uuid, function_id, key, value, serialization) \
VALUES ($1, $2, $3, $4, $5) \
ON CONFLICT (workflow_uuid, key, function_id) DO UPDATE \
SET value = EXCLUDED.value, serialization = EXCLUDED.serialization"
)))
.bind(workflow_id)
.bind(step_id)
.bind(key)
.bind(value)
.bind(serialization)
.execute(&mut *tx)
.await?;
self.record_step_on(
&mut tx,
workflow_id,
step_id,
step_names::SET_EVENT,
Outcome::Output(None),
None,
Some(timing),
None,
)
.await?;
tx.commit().await?;
self.notifier.signal(EVENTS_CHANNEL, workflow_id, key);
Ok(())
})
.await
}
async fn get_event(
&self,
workflow_id: &str,
key: &str,
timeout: Duration,
caller: Option<GetEventCaller<'_>>,
) -> Result<Option<EncodedValue>, Error> {
let events_table = self.tables.workflow_events.as_str();
let (pool, polling) = (&self.pool, &self.polling);
let select = format!(
"SELECT value, serialization FROM {events_table} WHERE workflow_uuid = $1 AND key = $2"
);
let select = &select;
if let Some(caller) = caller
&& let Some(step) = self
.check_step(caller.workflow_id, caller.step_id, step_names::GET_EVENT)
.await?
{
tracing::debug!(
workflow_id = caller.workflow_id,
step_id = caller.step_id,
key,
"replaying get_event"
);
return Ok(step.output.map(|value| EncodedValue {
value,
serialization: step.serialization,
}));
}
let mut subscription = self.notify.subscribe(event_key(workflow_id, key));
let (until, clock) = match caller {
Some(caller) => {
let deadline = self
.checkpoint_sleep(
caller.workflow_id,
caller.timeout_step_id,
timeout,
SleepKind::Deadline,
caller.owner_xid,
)
.await?;
(Some(deadline.arrives()), Some(deadline.clock))
}
None => (crate::rt::now().checked_add(timeout), None),
};
let started_at = clock.map_or_else(Timestamp::now, DatabaseTime::now);
let found = loop {
let row = with_retry(&self.retry, "get_event", move || async move {
let _permit = polling
.acquire()
.await
.expect("the polling limiter is never closed");
Ok(sqlx::query(AssertSqlSafe(select.clone()))
.bind(workflow_id)
.bind(key)
.fetch_optional(pool)
.await?)
})
.await?;
if let Some(row) = row {
break Some(EncodedValue {
value: row.try_get("value")?,
serialization: row.try_get("serialization")?,
});
}
let remaining = until.map_or(Duration::MAX, |until| {
until.saturating_duration_since(crate::rt::now())
});
if remaining.is_zero() {
break None;
}
let _ = crate::rt::timeout(
remaining.min(self.listener.poll_interval()),
subscription.notified(),
)
.await;
};
drop(subscription);
let Some(caller) = caller else {
tracing::debug!(
workflow_id,
key,
found = found.is_some(),
"get_event finished"
);
return Ok(found);
};
let timing = StepTiming {
started_at,
completed_at: clock.map_or_else(Timestamp::now, DatabaseTime::now),
};
let found = found.as_ref();
with_retry(&self.retry, "get_event", move || async move {
let mut tx = pool.begin().await?;
self.check_owner_on(
&mut tx,
WorkflowOwner {
workflow_id: caller.workflow_id,
owner_xid: caller.owner_xid,
},
)
.await?;
if let Some(step) = self
.check_step_on(
&mut tx,
caller.workflow_id,
caller.step_id,
step_names::GET_EVENT,
)
.await?
{
tx.commit().await?;
tracing::debug!(
workflow_id = caller.workflow_id,
step_id = caller.step_id,
key,
"replaying get_event"
);
return Ok(step.output.map(|value| EncodedValue {
value,
serialization: step.serialization,
}));
}
self.record_step_on(
&mut tx,
caller.workflow_id,
caller.step_id,
step_names::GET_EVENT,
Outcome::Output(found.map(|v| v.value.as_str())),
found.and_then(|v| v.serialization.as_deref()),
Some(timing),
None,
)
.await?;
tx.commit().await?;
tracing::debug!(
workflow_id = caller.workflow_id,
step_id = caller.step_id,
key,
found = found.is_some(),
"get_event finished"
);
Ok(found.cloned())
})
.await
}
async fn get_all_notifications(
&self,
workflow_id: &str,
) -> Result<Vec<NotificationRecord>, Error> {
let notifications_table = &self.tables.notifications;
let (notifications_table, pool) = (notifications_table.as_str(), &self.pool);
with_retry(&self.retry, "get_all_notifications", move || async move {
let rows = sqlx::query(AssertSqlSafe(format!(
"SELECT message_uuid, topic, message, serialization, created_at_epoch_ms, \
consumed \
FROM {notifications_table} WHERE destination_uuid = $1 ORDER BY created_at_epoch_ms"
)))
.bind(workflow_id)
.fetch_all(pool)
.await?;
rows.iter()
.map(|row| {
Ok(NotificationRecord {
message_uuid: row.try_get("message_uuid")?,
topic: row.try_get("topic")?,
message: row.try_get("message")?,
serialization: row.try_get("serialization")?,
created_at: Timestamp::from_epoch_ms(row.try_get("created_at_epoch_ms")?),
consumed: row.try_get("consumed")?,
})
})
.collect()
})
.await
}
async fn get_all_events(&self, workflow_id: &str) -> Result<Vec<EventRecord>, Error> {
let events_table = &self.tables.workflow_events;
let (events_table, pool) = (events_table.as_str(), &self.pool);
with_retry(&self.retry, "get_all_events", move || async move {
let rows = sqlx::query(AssertSqlSafe(format!(
"SELECT key, value, serialization FROM {events_table} \
WHERE workflow_uuid = $1 ORDER BY key"
)))
.bind(workflow_id)
.fetch_all(pool)
.await?;
rows.iter()
.map(|row| {
Ok(EventRecord {
key: row.try_get("key")?,
value: row.try_get("value")?,
serialization: row.try_get("serialization")?,
})
})
.collect()
})
.await
}
async fn read_stream_values(
&self,
workflow_id: &str,
key: &str,
offset: i32,
limit: i32,
) -> Result<StreamRead, Error> {
let workflow_table = self.tables.workflow_status.as_str();
let streams_table = self.tables.streams.as_str();
let (pool, polling) = (&self.pool, &self.polling);
let select = format!(
"SELECT w.status AS status, s.value AS value, s.serialization AS serialization, \
s.\"offset\" AS stream_offset \
FROM {workflow_table} w \
LEFT OUTER JOIN {streams_table} s \
ON s.workflow_uuid = w.workflow_uuid AND s.key = $2 \
AND s.\"offset\" >= $3 AND s.\"offset\" < $4 \
WHERE w.workflow_uuid = $1 \
ORDER BY s.\"offset\""
);
let select = &select;
if limit < 1 {
return Err(Error::InvalidInput {
field: "limit".into(),
detail: format!("must be at least 1, got {limit}"),
});
}
let end = i64::from(offset) + i64::from(limit);
with_retry(&self.retry, "read_stream_values", move || async move {
let _permit = polling
.acquire()
.await
.expect("the polling limiter is never closed");
let rows = sqlx::query(AssertSqlSafe(select.clone()))
.bind(workflow_id)
.bind(key)
.bind(offset)
.bind(end)
.fetch_all(pool)
.await?;
let Some(first) = rows.first() else {
return Err(Error::NonExistentWorkflow {
workflow_ids: vec![workflow_id.to_owned()],
});
};
let status = parse_workflow_status(&first.try_get::<String, _>("status")?)?;
let mut values = Vec::with_capacity(rows.len());
for (expected, row) in (i64::from(offset)..).zip(&rows) {
match row.try_get::<Option<i32>, _>("stream_offset")? {
Some(at) if i64::from(at) == expected => values.push(EncodedValue {
value: row.try_get("value")?,
serialization: row.try_get("serialization")?,
}),
_ => break,
}
}
Ok(StreamRead { status, values })
})
.await
}
async fn await_stream_values(
&self,
workflow_id: &str,
key: &str,
offset: i32,
limit: i32,
deadline: Option<Timestamp>,
polling_interval: Duration,
) -> Result<AwaitedStream, Error> {
let mut subscription = self.notify.subscribe(stream_key(workflow_id, key));
let mut final_look = false;
loop {
subscription.clear();
let read = self
.read_stream_values(workflow_id, key, offset, limit)
.await?;
if !read.values.is_empty() {
return Ok(AwaitedStream::Values(read.values));
}
if final_look {
return Ok(AwaitedStream::Ended);
}
if !read.status.is_active() {
final_look = true;
continue;
}
let mut wait = polling_interval;
if let Some(deadline) = deadline {
let remaining = deadline
.duration_since(Timestamp::now())
.unwrap_or(Duration::ZERO);
if remaining.is_zero() {
return Ok(AwaitedStream::TimedOut);
}
wait = wait.min(remaining);
}
let _ = crate::rt::timeout(wait, subscription.notified()).await;
}
}
async fn get_all_stream_entries(&self, workflow_id: &str) -> Result<Vec<StreamRecord>, Error> {
let streams_table = &self.tables.streams;
let (streams_table, pool) = (streams_table.as_str(), &self.pool);
with_retry(&self.retry, "get_all_stream_entries", move || async move {
let rows = sqlx::query(AssertSqlSafe(format!(
"SELECT key, \"offset\", value, serialization, function_id FROM {streams_table} \
WHERE workflow_uuid = $1 ORDER BY key, \"offset\""
)))
.bind(workflow_id)
.fetch_all(pool)
.await?;
rows.iter()
.map(|row| {
Ok(StreamRecord {
key: row.try_get("key")?,
offset: row.try_get("offset")?,
value: row.try_get("value")?,
serialization: row.try_get("serialization")?,
step_id: row.try_get("function_id")?,
})
})
.collect()
})
.await
}
async fn create_application_version(
&self,
version_name: &str,
application_name: Option<&str>,
) -> Result<(), Error> {
let versions_table = &self.tables.application_versions;
let version_id = uuid::Uuid::new_v4().to_string();
let (versions_table, pool, version_id) =
(versions_table.as_str(), &self.pool, version_id.as_str());
let application_name = application_name.or(self.application_name.as_deref());
with_retry(
&self.retry,
"create_application_version",
move || async move {
let mut tx = pool.begin().await?;
let claimed = sqlx::query(AssertSqlSafe(format!(
"UPDATE {versions_table} SET application_name = $1 \
WHERE version_name = $2 AND application_name IS NULL"
)))
.bind(application_name)
.bind(version_name)
.execute(&mut *tx)
.await?
.rows_affected();
if claimed == 0 {
sqlx::query(AssertSqlSafe(format!(
"INSERT INTO {versions_table} (version_id, version_name, application_name) \
VALUES ($1, $2, $3) ON CONFLICT DO NOTHING"
)))
.bind(version_id)
.bind(version_name)
.bind(application_name)
.execute(&mut *tx)
.await?;
}
resolve_owning_application(
&mut tx,
versions_table,
"version_name",
version_name,
application_name,
"Application version",
)
.await?;
tx.commit().await?;
Ok(())
},
)
.await
}
async fn list_application_versions(&self) -> Result<Vec<VersionInfo>, Error> {
let versions_table = &self.tables.application_versions;
let (versions_table, pool) = (versions_table.as_str(), &self.pool);
let application_name = self.application_name.as_deref();
with_retry(
&self.retry,
"list_application_versions",
move || async move {
let rows = sqlx::query(AssertSqlSafe(format!(
"SELECT {VERSION_COLUMNS} FROM {versions_table} \
WHERE ($1::text IS NULL \
OR application_name = $1 \
OR application_name IS NULL) \
ORDER BY version_timestamp DESC"
)))
.bind(application_name)
.fetch_all(pool)
.await?;
rows.iter().map(version_from_row).collect()
},
)
.await
}
async fn get_latest_application_version(
&self,
application_name: Option<&str>,
) -> Result<Option<VersionInfo>, Error> {
let versions_table = &self.tables.application_versions;
let (versions_table, pool) = (versions_table.as_str(), &self.pool);
let application_name = application_name.or(self.application_name.as_deref());
with_retry(
&self.retry,
"get_latest_application_version",
move || async move {
let row = latest_version_query(versions_table, application_name)
.fetch_optional(pool)
.await?;
row.as_ref().map(version_from_row).transpose()
},
)
.await
}
async fn promote_application_version(
&self,
version_name: &str,
application_name: Option<&str>,
) -> Result<(), Error> {
let versions_table = &self.tables.application_versions;
let (versions_table, pool) = (versions_table.as_str(), &self.pool);
let application_name = application_name.or(self.application_name.as_deref());
with_retry(
&self.retry,
"promote_application_version",
move || async move {
let mut tx = pool.begin().await?;
let owner = resolve_owning_application(
&mut tx,
versions_table,
"version_name",
version_name,
application_name,
"Application version",
)
.await?;
let probe = |filter: &str| {
format!(
"COALESCE((SELECT version_timestamp FROM {versions_table} {filter} \
ORDER BY application_name, version_timestamp DESC LIMIT 1), 0)"
)
};
let highest = if application_name.is_some() {
format!(
"GREATEST({}, {})",
probe("WHERE application_name = $2"),
probe("WHERE application_name IS NULL")
)
} else {
format!(
"COALESCE((SELECT version_timestamp FROM {versions_table} \
ORDER BY version_timestamp DESC LIMIT 1), 0)"
)
};
sqlx::query(AssertSqlSafe(format!(
"UPDATE {versions_table} SET \
version_timestamp = \
GREATEST({CLOCK_MS_SQL}, LEAST({highest}, 9223372036854775806) + 1), \
application_name = $2 \
WHERE version_name = $1 \
AND (application_name IS NULL OR application_name = $2)"
)))
.bind(version_name)
.bind(owner.as_deref())
.execute(&mut *tx)
.await?;
tx.commit().await?;
Ok(())
},
)
.await
}
async fn upsert_queue(
&self,
queue: &NewQueue<'_>,
on_existing: OnExistingQueue,
) -> Result<UpsertedQueue, Error> {
let queues_table = self.tables.queues.as_str();
let pool = &self.pool;
let application_name = queue.application_name.or(self.application_name.as_deref());
with_retry(&self.retry, "upsert_queue", move || async move {
let mut tx = pool.begin().await?;
let owner = resolve_owning_application(
&mut tx,
queues_table,
"name",
queue.name,
application_name,
"Queue",
)
.await?;
let before = loop {
let inserted = bind_new_queue(
sqlx::query(AssertSqlSafe(format!(
"INSERT INTO {queues_table} \
(name, concurrency, worker_concurrency, rate_limit_max, \
rate_limit_period_sec, priority_enabled, partition_queue, \
partition_concurrency, partition_worker_concurrency, \
partition_rate_limit_max, partition_rate_limit_period_sec, \
polling_interval_sec, updated_at, application_name) \
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, {NOW_MS_SQL}, $13) \
ON CONFLICT (name) DO NOTHING \
RETURNING {QUEUE_COLUMNS}"
))),
queue,
owner.as_deref(),
)
.fetch_optional(&mut *tx)
.await?;
if let Some(row) = inserted {
let after = queue_from_row(&row)?;
tx.commit().await?;
return Ok(UpsertedQueue {
before: None,
after,
});
}
let found = sqlx::query(AssertSqlSafe(format!(
"SELECT {QUEUE_COLUMNS} FROM {queues_table} WHERE name = $1 FOR UPDATE"
)))
.bind(queue.name)
.fetch_optional(&mut *tx)
.await?;
if let Some(row) = found {
break queue_from_row(&row)?;
}
};
let owner = resolve_owning_application(
&mut tx,
queues_table,
"name",
queue.name,
application_name,
"Queue",
)
.await?;
let after = match on_existing {
OnExistingQueue::Leave => before.clone(),
OnExistingQueue::Update => {
let row = bind_new_queue(
sqlx::query(AssertSqlSafe(format!(
"UPDATE {queues_table} SET \
concurrency = $2, worker_concurrency = $3, rate_limit_max = $4, \
rate_limit_period_sec = $5, priority_enabled = $6, \
partition_queue = $7, partition_concurrency = $8, \
partition_worker_concurrency = $9, partition_rate_limit_max = $10, \
partition_rate_limit_period_sec = $11, polling_interval_sec = $12, \
updated_at = {NOW_MS_SQL}, \
application_name = COALESCE(application_name, $13) \
WHERE name = $1 \
RETURNING {QUEUE_COLUMNS}"
))),
queue,
owner.as_deref(),
)
.fetch_one(&mut *tx)
.await?;
queue_from_row(&row)?
}
};
tx.commit().await?;
Ok(UpsertedQueue {
before: Some(before),
after,
})
})
.await
}
#[allow(clippy::too_many_arguments)]
async fn start_queued_workflows(
&self,
queue: &QueueRecord,
executor_id: &str,
application_version: &str,
owner_xid: &str,
partition_key: Option<&str>,
local_running_count: i64,
partition_local_running_count: i64,
) -> Result<Dequeued, Error> {
if partition_key == Some("") {
return Err(Error::InvalidInput {
field: "partition_key".into(),
detail: "must be absent rather than empty".to_owned(),
});
}
let workflow_table = self.tables.workflow_status.as_str();
let pool = &self.pool;
let application_name = self.application_name.as_deref();
let owner_predicate =
"($1::text IS NULL OR application_name = $1 OR application_name IS NULL)";
let partition_predicate = "($2::text IS NULL OR queue_partition_key = $2)";
let has_shared_budget = queue.concurrency.is_some()
|| queue.partition_concurrency.is_some()
|| queue.rate_limit.is_some()
|| queue.partition_rate_limit.is_some();
let budget_spans_partitions =
partition_key.is_some() && (queue.concurrency.is_some() || queue.rate_limit.is_some());
let rate_limited = queue.rate_limit.is_some() || queue.partition_rate_limit.is_some();
let period_ms = |limit: Option<RateLimit>| {
limit.map_or(0, |limit| {
i64::try_from(limit.period.as_millis()).unwrap_or(i64::MAX)
})
};
let rate_limit_period_ms = period_ms(queue.rate_limit);
let partition_rate_limit_period_ms = period_ms(queue.partition_rate_limit);
let retry = self.retry.without_conflict_replay();
with_retry(&retry, "start_queued_workflows", move || async move {
let mut tx = pool.begin().await?;
if has_shared_budget {
let isolation = if budget_spans_partitions {
"SERIALIZABLE"
} else {
"REPEATABLE READ"
};
sqlx::raw_sql(AssertSqlSafe(format!(
"SET TRANSACTION ISOLATION LEVEL {isolation}"
)))
.execute(&mut *tx)
.await?;
}
let narrow = |current: Option<i64>, available: i64| -> Option<i64> {
Some(current.map_or(available, |taken: i64| taken.min(available)))
};
let mut max_tasks: Option<i64> = None;
let stopped = |stop: DequeueStop| Dequeued {
workflow_ids: Vec::new(),
stopped_by: Some(stop),
};
if let Some(worker_concurrency) = queue.worker_concurrency {
max_tasks = narrow(
max_tasks,
(i64::from(worker_concurrency) - local_running_count).max(0),
);
if max_tasks == Some(0) {
tx.commit().await?;
return Ok(stopped(DequeueStop::WorkerConcurrency));
}
}
if let Some(worker_concurrency) = queue.partition_worker_concurrency
&& partition_key.is_some()
{
max_tasks = narrow(
max_tasks,
(i64::from(worker_concurrency) - partition_local_running_count).max(0),
);
if max_tasks == Some(0) {
tx.commit().await?;
return Ok(stopped(DequeueStop::PartitionWorkerConcurrency));
}
}
if let Some(limit) = queue.rate_limit {
let recent_starts: i64 = sqlx::query_scalar(AssertSqlSafe(format!(
"SELECT count(*) FROM {workflow_table} \
WHERE queue_name = $2 AND rate_limited = TRUE \
AND status NOT IN ('ENQUEUED', 'DELAYED') \
AND started_at_epoch_ms > {NOW_MS_SQL} - $3 \
AND {owner_predicate}"
)))
.bind(application_name)
.bind(&queue.name)
.bind(rate_limit_period_ms)
.fetch_one(&mut *tx)
.await?;
max_tasks = narrow(max_tasks, (i64::from(limit.limit) - recent_starts).max(0));
if max_tasks == Some(0) {
tx.commit().await?;
return Ok(stopped(DequeueStop::RateLimit));
}
}
if let (Some(limit), Some(key)) = (queue.partition_rate_limit, partition_key) {
let partition_starts: i64 = sqlx::query_scalar(AssertSqlSafe(format!(
"SELECT count(*) FROM {workflow_table} \
WHERE queue_name = $3 AND rate_limited = TRUE \
AND status NOT IN ('ENQUEUED', 'DELAYED') \
AND started_at_epoch_ms > {NOW_MS_SQL} - $4 \
AND queue_partition_key = $2 \
AND {owner_predicate}"
)))
.bind(application_name)
.bind(key)
.bind(&queue.name)
.bind(partition_rate_limit_period_ms)
.fetch_one(&mut *tx)
.await?;
max_tasks = narrow(
max_tasks,
(i64::from(limit.limit) - partition_starts).max(0),
);
if max_tasks == Some(0) {
tx.commit().await?;
return Ok(stopped(DequeueStop::PartitionRateLimit));
}
}
if let Some(concurrency) = queue.concurrency {
let running: i64 = sqlx::query_scalar(AssertSqlSafe(format!(
"SELECT count(*) FROM {workflow_table} WHERE {QUEUE_PENDING_PREDICATE}"
)))
.bind(application_name)
.bind(&queue.name)
.fetch_one(&mut *tx)
.await?;
if concurrency > 0 && running > i64::from(concurrency) {
tracing::warn!(
queue = %queue.name,
running,
concurrency,
"pending workflows exceed the queue's global concurrency limit"
);
}
max_tasks = narrow(max_tasks, (i64::from(concurrency) - running).max(0));
if max_tasks == Some(0) {
tx.commit().await?;
return Ok(stopped(DequeueStop::Concurrency));
}
}
if let (Some(concurrency), Some(key)) = (queue.partition_concurrency, partition_key) {
let running: i64 = sqlx::query_scalar(AssertSqlSafe(format!(
"SELECT count(*) FROM {workflow_table} \
WHERE queue_name = $3 AND status = 'PENDING' \
AND queue_partition_key = $2 AND {owner_predicate}"
)))
.bind(application_name)
.bind(key)
.bind(&queue.name)
.fetch_one(&mut *tx)
.await?;
if running > i64::from(concurrency) {
tracing::warn!(
queue = %queue.name,
partition = %key,
running,
concurrency,
"pending workflows exceed the partition's concurrency limit"
);
}
max_tasks = narrow(max_tasks, (i64::from(concurrency) - running).max(0));
if max_tasks == Some(0) {
tx.commit().await?;
return Ok(stopped(DequeueStop::PartitionConcurrency));
}
}
let is_latest = self
.is_latest_application_version(&mut tx, application_version)
.await?;
let version_predicate = version_predicate(is_latest, 3);
let lock = if has_shared_budget {
"FOR UPDATE NOWAIT"
} else {
"FOR UPDATE SKIP LOCKED"
};
let limit = match max_tasks {
Some(n) => format!(" LIMIT {n}"),
None => String::new(),
};
let candidates: Vec<String> = sqlx::query_scalar(AssertSqlSafe(format!(
"SELECT workflow_uuid FROM {workflow_table} \
WHERE queue_name = $4 AND status = 'ENQUEUED' \
AND {version_predicate} AND {owner_predicate} AND {partition_predicate} \
ORDER BY priority ASC, created_at ASC{limit} {lock}"
)))
.bind(application_name)
.bind(partition_key)
.bind(application_version)
.bind(&queue.name)
.fetch_all(&mut *tx)
.await?;
let started: Vec<String> = if candidates.is_empty() {
Vec::new()
} else {
let deadline_from_timeout = now_ms_plus_sql("workflow_timeout_ms");
let flipped: HashSet<String> = sqlx::query_scalar(AssertSqlSafe(format!(
"UPDATE {workflow_table} \
SET status = 'PENDING', executor_id = $2, application_version = $3, \
owner_xid = $6, \
started_at_epoch_ms = {NOW_MS_SQL}, rate_limited = $4, \
updated_at = {NOW_MS_SQL}, \
application_name = COALESCE(application_name, $1), \
recovery_attempts = recovery_attempts + 1, \
workflow_deadline_epoch_ms = CASE \
WHEN workflow_timeout_ms IS NOT NULL \
AND workflow_deadline_epoch_ms IS NULL \
THEN {deadline_from_timeout} \
ELSE workflow_deadline_epoch_ms \
END \
WHERE workflow_uuid = ANY($5::text[]) AND status = 'ENQUEUED' \
AND {owner_predicate} \
RETURNING workflow_uuid"
)))
.bind(application_name)
.bind(executor_id)
.bind(application_version)
.bind(rate_limited)
.bind(candidates.as_slice())
.bind(owner_xid)
.fetch_all(&mut *tx)
.await?
.into_iter()
.collect();
candidates
.into_iter()
.filter(|id| flipped.contains(id))
.collect()
};
tx.commit().await?;
if !started.is_empty() {
tracing::debug!(queue = %queue.name, started = started.len(), "dequeued");
}
Ok(Dequeued {
workflow_ids: started,
stopped_by: None,
})
})
.await
}
async fn count_pending_queue_workflows(&self, queue_name: &str) -> Result<i64, Error> {
let workflow_table = self.tables.workflow_status.as_str();
let pool = &self.pool;
let application_name = self.application_name.as_deref();
with_retry(
&self.retry,
"count_pending_queue_workflows",
move || async move {
let pending: i64 = sqlx::query_scalar(AssertSqlSafe(format!(
"SELECT count(*) FROM {workflow_table} WHERE {QUEUE_PENDING_PREDICATE}"
)))
.bind(application_name)
.bind(queue_name)
.fetch_one(pool)
.await?;
Ok(pending)
},
)
.await
}
async fn get_queue_partitions(&self, queue_name: &str) -> Result<Vec<String>, Error> {
let workflow_table = self.tables.workflow_status.as_str();
let pool = &self.pool;
let application_name = self.application_name.as_deref();
let eligible_predicate = "queue_name = $1 AND status = 'ENQUEUED' \
AND ($2::text IS NULL OR application_name = $2 \
OR application_name IS NULL)";
with_retry(&self.retry, "get_queue_partitions", move || async move {
let partitions: Vec<String> = sqlx::query_scalar(AssertSqlSafe(format!(
"WITH RECURSIVE partitions AS ( \
(SELECT MIN(queue_partition_key) AS pk FROM {workflow_table} \
WHERE {eligible_predicate} AND queue_partition_key IS NOT NULL) \
UNION ALL \
(SELECT (SELECT MIN(queue_partition_key) FROM {workflow_table} \
WHERE {eligible_predicate} AND queue_partition_key > partitions.pk) \
FROM partitions WHERE partitions.pk IS NOT NULL) \
) \
SELECT pk FROM partitions WHERE pk IS NOT NULL"
)))
.bind(queue_name)
.bind(application_name)
.fetch_all(pool)
.await?;
Ok(partitions)
})
.await
}
async fn start_queued_partitioned_workflows(
&self,
queue: &QueueRecord,
executor_id: &str,
application_version: &str,
owner_xid: &str,
max_tasks: Option<i64>,
) -> Result<Vec<String>, Error> {
if queue.partition_concurrency != Some(1)
|| queue.concurrency.is_some()
|| queue.rate_limit.is_some()
|| queue.partition_rate_limit.is_some()
{
return Err(Error::InvalidInput {
field: "queue".into(),
detail: format!(
"a partitioned sweep needs partition concurrency 1 and no other limit, but \
{:?} has partition_concurrency={:?} concurrency={:?} rate_limited={} \
partition_rate_limited={}",
queue.name,
queue.partition_concurrency,
queue.concurrency,
queue.rate_limit.is_some(),
queue.partition_rate_limit.is_some(),
),
});
}
if max_tasks == Some(0) {
return Ok(Vec::new());
}
let cap = i64::from(PARTITIONED_DEQUEUE_SWEEP_CAP);
let sweep_limit = max_tasks.map_or(cap, |budget| budget.min(cap));
let sweep_order = if sweep_limit < cap {
"random()"
} else {
"partitions.pk ASC"
};
let workflow_table = self.tables.workflow_status.as_str();
let pool = &self.pool;
let application_name = self.application_name.as_deref();
let retry = self.retry.without_conflict_replay();
with_retry(
&retry,
"start_queued_partitioned_workflows",
move || async move {
let mut tx = pool.begin().await?;
let is_latest = self
.is_latest_application_version(&mut tx, application_version)
.await?;
let version_predicate = version_predicate(is_latest, 3);
let eligible_predicate = "queue_name = $1 AND status = 'ENQUEUED' \
AND ($2::text IS NULL OR application_name = $2 \
OR application_name IS NULL)";
let candidates: Vec<String> = sqlx::query_scalar(AssertSqlSafe(format!(
"WITH RECURSIVE partitions AS ( \
(SELECT MIN(queue_partition_key) AS pk FROM {workflow_table} \
WHERE {eligible_predicate} AND queue_partition_key IS NOT NULL) \
UNION ALL \
(SELECT (SELECT MIN(queue_partition_key) FROM {workflow_table} \
WHERE {eligible_predicate} AND queue_partition_key > partitions.pk) \
FROM partitions WHERE partitions.pk IS NOT NULL) \
) \
, chosen AS ( \
SELECT partitions.pk FROM partitions \
WHERE partitions.pk IS NOT NULL \
AND NOT EXISTS ( \
SELECT 1 FROM {workflow_table} \
WHERE queue_name = $1 AND status = 'PENDING' \
AND queue_partition_key IS NOT NULL \
AND queue_partition_key = partitions.pk \
) \
ORDER BY {sweep_order} LIMIT $4 \
) \
SELECT head.workflow_uuid FROM chosen \
JOIN LATERAL ( \
SELECT workflow_uuid FROM {workflow_table} \
WHERE {eligible_predicate} AND queue_partition_key = chosen.pk \
AND {version_predicate} \
ORDER BY priority ASC, created_at ASC, workflow_uuid ASC LIMIT 1 \
) head ON TRUE \
ORDER BY chosen.pk ASC"
)))
.bind(&queue.name)
.bind(application_name)
.bind(application_version)
.bind(sweep_limit)
.fetch_all(&mut *tx)
.await?;
if candidates.is_empty() {
tx.commit().await?;
return Ok(Vec::new());
}
let claim_predicate = format!(
"workflow_uuid = ANY($1::text[]) AND status = 'ENQUEUED' \
AND queue_name = $2 AND queue_partition_key IS NOT NULL \
AND {version_predicate} \
AND ($4::text IS NULL OR application_name = $4 \
OR application_name IS NULL)"
);
let locked: Vec<String> = sqlx::query_scalar(AssertSqlSafe(format!(
"SELECT workflow_uuid FROM {workflow_table} WHERE {claim_predicate} \
FOR UPDATE SKIP LOCKED"
)))
.bind(&candidates)
.bind(&queue.name)
.bind(application_version)
.bind(application_name)
.fetch_all(&mut *tx)
.await?;
let locked: HashSet<&str> = locked.iter().map(String::as_str).collect();
let claiming: Vec<&str> = candidates
.iter()
.map(String::as_str)
.filter(|id| locked.contains(id))
.collect();
if claiming.is_empty() {
tx.commit().await?;
return Ok(Vec::new());
}
let deadline_from_timeout = now_ms_plus_sql("workflow_timeout_ms");
let flipped: Vec<String> = sqlx::query_scalar(AssertSqlSafe(format!(
"UPDATE {workflow_table} \
SET status = 'PENDING', executor_id = $5, application_version = $3, \
owner_xid = $6, \
started_at_epoch_ms = {NOW_MS_SQL}, rate_limited = FALSE, \
updated_at = {NOW_MS_SQL}, \
application_name = COALESCE(application_name, $4), \
recovery_attempts = recovery_attempts + 1, \
workflow_deadline_epoch_ms = CASE \
WHEN workflow_timeout_ms IS NOT NULL \
AND workflow_deadline_epoch_ms IS NULL \
THEN {deadline_from_timeout} \
ELSE workflow_deadline_epoch_ms \
END \
WHERE {claim_predicate} RETURNING workflow_uuid"
)))
.bind(&claiming)
.bind(&queue.name)
.bind(application_version)
.bind(application_name)
.bind(executor_id)
.bind(owner_xid)
.fetch_all(&mut *tx)
.await?;
tx.commit().await?;
let flipped: HashSet<&str> = flipped.iter().map(String::as_str).collect();
let started: Vec<String> = claiming
.into_iter()
.filter(|id| flipped.contains(id))
.map(str::to_owned)
.collect();
if !started.is_empty() {
tracing::debug!(
queue = %queue.name,
started = started.len(),
"dequeued a partition sweep"
);
}
Ok(started)
},
)
.await
}
async fn get_deduplication_key_holder(
&self,
queue_name: &str,
deduplication_id: &str,
) -> Result<Option<String>, Error> {
let workflow_table = self.tables.workflow_status.as_str();
let pool = &self.pool;
with_retry(
&self.retry,
"get_deduplication_key_holder",
move || async move {
let holder: Option<String> = sqlx::query_scalar(AssertSqlSafe(format!(
"SELECT workflow_uuid FROM {workflow_table} \
WHERE queue_name = $1 AND deduplication_id = $2"
)))
.bind(queue_name)
.bind(deduplication_id)
.fetch_optional(pool)
.await?;
tracing::debug!(
queue_name,
deduplication_id,
?holder,
"read the key's holder"
);
Ok(holder)
},
)
.await
}
async fn get_queue(&self, name: &str) -> Result<Option<QueueRecord>, Error> {
let queues_table = self.tables.queues.as_str();
let pool = &self.pool;
with_retry(&self.retry, "get_queue", move || async move {
let row = sqlx::query(AssertSqlSafe(format!(
"SELECT {QUEUE_COLUMNS} FROM {queues_table} WHERE name = $1"
)))
.bind(name)
.fetch_optional(pool)
.await?;
row.as_ref().map(queue_from_row).transpose()
})
.await
}
async fn list_queues(&self, applications: &AppScope<'_>) -> Result<Vec<QueueRecord>, Error> {
let queues_table = self.tables.queues.as_str();
let pool = &self.pool;
let application_name = self.application_name.as_deref();
with_retry(&self.retry, "list_queues", move || async move {
let mut q = sqlx::QueryBuilder::<sqlx::Postgres>::new("SELECT ");
q.push(QUEUE_COLUMNS).push(" FROM ").push(queues_table);
match applications {
AppScope::Any => {}
AppScope::Named(names) if names.is_empty() => {}
AppScope::Named(names) => {
q.push(" WHERE (application_name = ANY(")
.push_bind(&names[..])
.push(") OR application_name IS NULL)");
}
AppScope::Unset => {
if let Some(name) = application_name {
q.push(" WHERE (application_name = ")
.push_bind(name)
.push(" OR application_name IS NULL)");
}
}
}
q.push(" ORDER BY name");
let rows = q.build().fetch_all(pool).await?;
rows.iter().map(queue_from_row).collect()
})
.await
}
async fn update_queue(
&self,
name: &str,
update: &QueueUpdate,
validate: &(
dyn for<'r, 's> Fn(&'r QueueRecord, &'s QueueRecord) -> Result<(), Error> + Send + Sync
),
) -> Result<UpdatedQueue, Error> {
let queues_table = self.tables.queues.as_str();
let pool = &self.pool;
with_retry(&self.retry, "update_queue", move || async move {
let mut tx = pool.begin().await?;
let row = sqlx::query(AssertSqlSafe(format!(
"SELECT {QUEUE_COLUMNS} FROM {queues_table} WHERE name = $1 FOR UPDATE"
)))
.bind(name)
.fetch_optional(&mut *tx)
.await?;
let Some(stored) = row.as_ref().map(queue_from_row).transpose()? else {
return Err(Error::NotRegistered {
kind: "Queue".into(),
name: name.to_owned(),
});
};
if update.is_empty() {
tx.commit().await?;
return Ok(UpdatedQueue {
before: stored.clone(),
after: stored,
});
}
let merged = update.apply_to(&stored);
validate(&stored, &merged)?;
let mut q = sqlx::QueryBuilder::<sqlx::Postgres>::new("UPDATE ");
q.push(queues_table).push(" SET ");
let mut set = q.separated(", ");
macro_rules! assign {
($field:expr, $column:literal) => {
if let Some(value) = $field {
set.push(concat!($column, " = "));
set.push_bind_unseparated(value);
}
};
}
assign!(update.concurrency.set(), "concurrency");
assign!(update.worker_concurrency.set(), "worker_concurrency");
set.push("priority_enabled = TRUE");
if !(update.partition_concurrency.is_leave()
&& update.partition_worker_concurrency.is_leave()
&& update.partition_rate_limit.is_leave())
{
set.push("partition_queue = ");
set.push_bind_unseparated(merged.is_partitioned());
}
assign!(update.partition_concurrency.set(), "partition_concurrency");
assign!(
update.partition_worker_concurrency.set(),
"partition_worker_concurrency"
);
assign!(
update.polling_interval.set().map(|d| d.as_secs_f64()),
"polling_interval_sec"
);
if let Some(limit) = update.rate_limit.set() {
set.push("rate_limit_max = ");
set.push_bind_unseparated(limit.map(|l| l.limit));
set.push("rate_limit_period_sec = ");
set.push_bind_unseparated(limit.map(|l| l.period.as_secs_f64()));
}
if let Some(limit) = update.partition_rate_limit.set() {
set.push("partition_rate_limit_max = ");
set.push_bind_unseparated(limit.map(|l| l.limit));
set.push("partition_rate_limit_period_sec = ");
set.push_bind_unseparated(limit.map(|l| l.period.as_secs_f64()));
}
set.push("updated_at = ").push_unseparated(NOW_MS_SQL);
q.push(" WHERE name = ").push_bind(name);
q.push(" RETURNING ").push(QUEUE_COLUMNS);
let row = q.build().fetch_one(&mut *tx).await?;
let written = queue_from_row(&row)?;
tx.commit().await?;
Ok(UpdatedQueue {
before: stored,
after: written,
})
})
.await
}
async fn debounce_delayed_workflow(
&self,
request: &DebounceRequest<'_>,
caller: Option<StepCaller<'_>>,
) -> Result<Debounce, Error> {
request.validate()?;
let workflow_table = self.tables.workflow_status.as_str();
let input_table = self.tables.workflow_input.as_str();
let application_name = request
.application_name
.or(self.application_name.as_deref());
let period_ms = duration_to_ms("period", request.period)?;
let delay_until = now_ms_plus_sql("$4::bigint");
let delay_until = delay_until.as_str();
let started_at = Timestamp::now();
with_retry(
&self.retry,
"debounce_delayed_workflow",
move || async move {
self.run_transactional_step(
caller,
step_names::DEBOUNCE,
started_at,
|mut tx| async move {
let bounced: Option<String> = sqlx::query_scalar(AssertSqlSafe(format!(
"UPDATE {workflow_table} \
SET delay_until_epoch_ms = CASE \
WHEN debounce_deadline_epoch_ms IS NOT NULL \
AND debounce_deadline_epoch_ms < {delay_until} \
THEN debounce_deadline_epoch_ms \
ELSE {delay_until} \
END, \
serialization = $5, \
updated_at = {NOW_MS_SQL}, \
application_name = COALESCE(application_name, $6) \
WHERE name = $1 AND queue_name = $2 AND deduplication_id = $3 \
AND NULLIF(class_name, '') IS NOT DISTINCT FROM $7 \
AND NULLIF(config_name, '') IS NOT DISTINCT FROM $8 \
AND status = 'DELAYED' AND is_debounced = TRUE \
AND ($6::text IS NULL OR application_name = $6 \
OR application_name IS NULL) \
RETURNING workflow_uuid"
)))
.bind(request.workflow.name())
.bind(request.queue_name)
.bind(request.deduplication_id)
.bind(period_ms)
.bind(request.serialization)
.bind(application_name)
.bind(request.workflow.class_name())
.bind(request.workflow.config_name())
.fetch_optional(&mut *tx)
.await?;
if let Some(workflow_id) = bounced {
sqlx::query(AssertSqlSafe(format!(
"INSERT INTO {input_table} (workflow_uuid, inputs) VALUES ($1, $2) \
ON CONFLICT (workflow_uuid) DO UPDATE SET inputs = EXCLUDED.inputs"
)))
.bind(&workflow_id)
.bind(request.inputs)
.execute(&mut *tx)
.await?;
return Ok((tx, Debounce::Bounced { workflow_id }));
}
type HolderRow = (
String,
bool,
Option<String>,
Option<String>,
Option<String>,
Option<String>,
);
let holder: Option<HolderRow> = sqlx::query_as(AssertSqlSafe(format!(
"SELECT workflow_uuid, is_debounced, name, class_name, config_name, \
application_name \
FROM {workflow_table} \
WHERE queue_name = $1 AND deduplication_id = $2"
)))
.bind(request.queue_name)
.bind(request.deduplication_id)
.fetch_optional(&mut *tx)
.await?;
let outcome = match holder {
None => Debounce::Unheld,
Some((
workflow_id,
is_debounced,
workflow_name,
class_name,
config_name,
application_name,
)) => Debounce::Held(DebounceHolder {
workflow_id,
is_debounced,
workflow_name,
class_name,
config_name,
application_name,
}),
};
Ok((tx, outcome))
},
)
.await
},
)
.await
}
async fn delete_queue(&self, name: &str) -> Result<(), Error> {
let queues_table = self.tables.queues.as_str();
let pool = &self.pool;
with_retry(&self.retry, "delete_queue", move || async move {
sqlx::query(AssertSqlSafe(format!(
"DELETE FROM {queues_table} WHERE name = $1"
)))
.bind(name)
.execute(pool)
.await?;
Ok(())
})
.await
}
async fn create_schedule(
&self,
schedule: &NewSchedule<'_>,
caller: Option<StepCaller<'_>>,
) -> Result<(), Error> {
let schedules_table = self.tables.workflow_schedules.as_str();
let application_name = schedule
.application_name
.or(self.application_name.as_deref());
let schedule_id = schedule
.schedule_id
.map_or_else(|| uuid::Uuid::new_v4().to_string(), str::to_owned);
let schedule_id = schedule_id.as_str();
let started_at = Timestamp::now();
with_retry(&self.retry, "create_schedule", move || async move {
self.run_transactional_step(
caller,
step_names::CREATE_SCHEDULE,
started_at,
|mut tx| async move {
let owner = resolve_owning_application(
&mut tx,
schedules_table,
"schedule_name",
schedule.schedule_name,
application_name,
"Schedule",
)
.await?;
let inserted = sqlx::query(AssertSqlSafe(format!(
"INSERT INTO {schedules_table} \
(schedule_id, schedule_name, workflow_name, workflow_class_name, \
schedule, status, context, last_fired_at, automatic_backfill, \
cron_timezone, queue_name, application_name) \
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12)"
)))
.bind(schedule_id)
.bind(schedule.schedule_name)
.bind(schedule.workflow_name)
.bind(schedule.workflow_class_name)
.bind(schedule.schedule)
.bind(schedule.status.as_str())
.bind(schedule.context)
.bind(schedule.last_fired_at.map(Timestamp::to_iso8601))
.bind(schedule.automatic_backfill)
.bind(schedule.cron_timezone)
.bind(schedule.queue_name)
.bind(owner.as_deref())
.execute(&mut *tx)
.await;
match inserted {
Ok(_) => {
tracing::debug!(
schedule_name = schedule.schedule_name,
schedule_id,
"registered a schedule"
);
Ok((tx, ()))
}
Err(error) if is_unique_violation(&error) => {
let kind = match error.as_database_error().and_then(|e| e.constraint())
{
Some(c) if c.ends_with("_pkey") => "Schedule id",
_ => "Schedule",
};
Err(Error::AlreadyRegistered {
kind: kind.into(),
name: match kind {
"Schedule id" => schedule_id.to_owned(),
_ => schedule.schedule_name.to_owned(),
},
})
}
Err(error) => Err(error.into()),
}
},
)
.await
})
.await
}
async fn upsert_schedule(
&self,
schedule: &NewSchedule<'_>,
caller: Option<StepCaller<'_>>,
) -> Result<(), Error> {
let schedule_id = schedule
.schedule_id
.map_or_else(|| uuid::Uuid::new_v4().to_string(), str::to_owned);
let schedule_id = schedule_id.as_str();
let started_at = Timestamp::now();
with_retry(&self.retry, "upsert_schedule", move || async move {
self.run_transactional_step(
caller,
step_names::UPSERT_SCHEDULE,
started_at,
|mut tx| async move {
self.upsert_schedule_on(&mut tx, schedule, schedule_id)
.await?;
Ok((tx, ()))
},
)
.await
})
.await
}
async fn get_schedule(
&self,
name: &str,
caller: Option<StepCaller<'_>>,
) -> Result<Option<ScheduleRecord>, Error> {
let schedules_table = self.tables.workflow_schedules.as_str();
let started_at = Timestamp::now();
with_retry(&self.retry, "get_schedule", move || async move {
self.run_transactional_step(
caller,
step_names::GET_SCHEDULE,
started_at,
|mut tx| async move {
let row = sqlx::query(AssertSqlSafe(format!(
"SELECT {SCHEDULE_COLUMNS} FROM {schedules_table} WHERE schedule_name = $1"
)))
.bind(name)
.fetch_optional(&mut *tx)
.await?;
let schedule = row.as_ref().map(schedule_from_row).transpose()?;
Ok((tx, schedule))
},
)
.await
})
.await
}
async fn apply_schedules(&self, schedules: &[NewSchedule<'_>]) -> Result<(), Error> {
let pool = &self.pool;
let ids: Vec<String> = schedules
.iter()
.map(|s| {
s.schedule_id
.map_or_else(|| uuid::Uuid::new_v4().to_string(), str::to_owned)
})
.collect();
let ids = ids.as_slice();
with_retry(&self.retry, "apply_schedules", move || async move {
let mut tx = pool.begin().await?;
for (schedule, schedule_id) in schedules.iter().zip(ids) {
self.upsert_schedule_on(&mut tx, schedule, schedule_id)
.await?;
}
tx.commit().await?;
tracing::debug!(count = schedules.len(), "applied schedules");
Ok(())
})
.await
}
async fn list_schedules(
&self,
filter: &ScheduleFilter<'_>,
caller: Option<StepCaller<'_>>,
) -> Result<Vec<ScheduleRecord>, Error> {
let schedules_table = self.tables.workflow_schedules.as_str();
let application_name = self.application_name.as_deref();
let statuses: Vec<&str> = filter.statuses.iter().map(|s| s.as_str()).collect();
let statuses = statuses.as_slice();
let started_at = Timestamp::now();
with_retry(&self.retry, "list_schedules", move || async move {
self.run_transactional_step(
caller,
step_names::LIST_SCHEDULES,
started_at,
|mut tx| async move {
let mut q = sqlx::QueryBuilder::<sqlx::Postgres>::new("SELECT ");
q.push(SCHEDULE_COLUMNS)
.push(" FROM ")
.push(schedules_table)
.push(" WHERE TRUE");
if !statuses.is_empty() {
q.push(" AND status = ANY(").push_bind(statuses).push(")");
}
if !filter.workflow_names.is_empty() {
q.push(" AND workflow_name = ANY(")
.push_bind(&filter.workflow_names[..])
.push(")");
}
if !filter.schedule_name_prefixes.is_empty() {
q.push(" AND (");
for (i, prefix) in filter.schedule_name_prefixes.iter().enumerate() {
if i > 0 {
q.push(" OR ");
}
q.push("schedule_name LIKE ")
.push_bind(format!("{}%", escape_like(prefix)));
}
q.push(")");
}
match &filter.app_scope {
AppScope::Any => {}
AppScope::Named(names) if names.is_empty() => {}
AppScope::Named(names) => {
q.push(" AND (application_name = ANY(")
.push_bind(&names[..])
.push(") OR application_name IS NULL)");
}
AppScope::Unset => {
if let Some(name) = application_name {
q.push(" AND (application_name = ")
.push_bind(name)
.push(" OR application_name IS NULL)");
}
}
}
q.push(" ORDER BY schedule_name");
let rows = q.build().fetch_all(&mut *tx).await?;
let schedules: Vec<ScheduleRecord> = rows
.iter()
.map(schedule_from_row)
.collect::<Result<_, _>>()?;
Ok((tx, schedules))
},
)
.await
})
.await
}
async fn update_schedule(
&self,
name: &str,
update: &ScheduleUpdate<'_>,
caller: Option<StepCaller<'_>>,
) -> Result<(), Error> {
let schedules_table = self.tables.workflow_schedules.as_str();
let started_at = Timestamp::now();
with_retry(&self.retry, "update_schedule", move || async move {
self.run_transactional_step(
caller,
step_names::UPDATE_SCHEDULE,
started_at,
|mut tx| async move {
let changed = if update.is_empty() {
sqlx::query_scalar::<_, i32>(AssertSqlSafe(format!(
"SELECT 1::int4 FROM {schedules_table} WHERE schedule_name = $1"
)))
.bind(name)
.fetch_optional(&mut *tx)
.await?
.is_some()
} else {
let mut q = sqlx::QueryBuilder::<sqlx::Postgres>::new("UPDATE ");
q.push(schedules_table).push(" SET ");
let mut separated = q.separated(", ");
if let Change::Set(schedule) = update.schedule {
separated
.push("schedule = ")
.push_bind_unseparated(schedule);
}
if let Change::Set(context) = update.context {
separated.push("context = ").push_bind_unseparated(context);
}
if let Change::Set(backfill) = update.automatic_backfill {
separated
.push("automatic_backfill = ")
.push_bind_unseparated(backfill);
}
if let Change::Set(timezone) = update.cron_timezone {
separated
.push("cron_timezone = ")
.push_bind_unseparated(timezone);
}
if let Change::Set(queue_name) = update.queue_name {
separated
.push("queue_name = ")
.push_bind_unseparated(queue_name);
}
q.push(" WHERE schedule_name = ").push_bind(name);
q.build().execute(&mut *tx).await?.rows_affected() > 0
};
if !changed {
return Err(Error::NotRegistered {
kind: "Schedule".into(),
name: name.to_owned(),
});
}
Ok((tx, ()))
},
)
.await
})
.await
}
async fn set_schedule_status(
&self,
name: &str,
status: ScheduleStatus,
caller: Option<StepCaller<'_>>,
) -> Result<(), Error> {
let schedules_table = self.tables.workflow_schedules.as_str();
let step_name = match status {
ScheduleStatus::Active => step_names::RESUME_SCHEDULE,
ScheduleStatus::Paused => step_names::PAUSE_SCHEDULE,
};
let started_at = Timestamp::now();
with_retry(&self.retry, "set_schedule_status", move || async move {
self.run_transactional_step(caller, step_name, started_at, |mut tx| async move {
let updated = sqlx::query(AssertSqlSafe(format!(
"UPDATE {schedules_table} SET status = $2 WHERE schedule_name = $1"
)))
.bind(name)
.bind(status.as_str())
.execute(&mut *tx)
.await?
.rows_affected();
if updated == 0 {
return Err(Error::NotRegistered {
kind: "Schedule".into(),
name: name.to_owned(),
});
}
tracing::debug!(
schedule_name = name,
status = status.as_str(),
"set a schedule's status"
);
Ok((tx, ()))
})
.await
})
.await
}
async fn update_schedule_last_fired_at(
&self,
name: &str,
last_fired_at: Timestamp,
) -> Result<(), Error> {
let schedules_table = self.tables.workflow_schedules.as_str();
let pool = &self.pool;
let last_fired_at = last_fired_at.to_iso8601();
let last_fired_at = last_fired_at.as_str();
with_retry(
&self.retry,
"update_schedule_last_fired_at",
move || async move {
sqlx::query(AssertSqlSafe(format!(
"UPDATE {schedules_table} SET last_fired_at = $2 WHERE schedule_name = $1"
)))
.bind(name)
.bind(last_fired_at)
.execute(pool)
.await?;
Ok(())
},
)
.await
}
async fn delete_schedule(
&self,
name: &str,
caller: Option<StepCaller<'_>>,
) -> Result<(), Error> {
let schedules_table = self.tables.workflow_schedules.as_str();
let started_at = Timestamp::now();
with_retry(&self.retry, "delete_schedule", move || async move {
self.run_transactional_step(
caller,
step_names::DELETE_SCHEDULE,
started_at,
|mut tx| async move {
sqlx::query(AssertSqlSafe(format!(
"DELETE FROM {schedules_table} WHERE schedule_name = $1"
)))
.bind(name)
.execute(&mut *tx)
.await?;
Ok((tx, ()))
},
)
.await
})
.await
}
async fn rename_application(
&self,
source: RenameFrom<'_>,
new_name: &str,
batching: RenameBatching,
) -> Result<ApplicationRowCounts, Error> {
if !is_valid_application_name(new_name) {
return Err(Error::InvalidInput {
field: "new_name".into(),
detail: "must be 3 to 30 characters of lowercase letters, digits, dashes and \
underscores"
.to_owned(),
});
}
if source.application() == Some(new_name) {
return Err(Error::InvalidInput {
field: "new_name".into(),
detail: format!("{new_name:?} already holds that name"),
});
}
if let RenameBatching::Batched(0) = batching {
return Err(Error::InvalidInput {
field: "batching".into(),
detail: "a batch must hold at least one workflow".to_owned(),
});
}
let workflow_table = self.tables.workflow_status.as_str();
let steps_table = self.tables.operation_outputs.as_str();
let queues_table = self.tables.queues.as_str();
let schedules_table = self.tables.workflow_schedules.as_str();
let versions_table = self.tables.application_versions.as_str();
let pool = &self.pool;
let predicate = rename_source_predicate(source, 2);
let (predicate, renamed_from) = (predicate.as_str(), source.application());
let counts = with_retry(&self.retry, "rename_application", move || async move {
let mut tx = pool.begin().await?;
let mut moved = [0u64; 4];
for (i, table) in [
queues_table,
schedules_table,
versions_table,
workflow_table,
]
.into_iter()
.enumerate()
{
let in_flight = if table == workflow_table {
" AND status IN ('PENDING', 'ENQUEUED', 'DELAYED')"
} else {
""
};
moved[i] = sqlx::query(AssertSqlSafe(format!(
"UPDATE {table} SET application_name = $1 WHERE {predicate}{in_flight}"
)))
.bind(new_name)
.bind(renamed_from)
.execute(&mut *tx)
.await?
.rows_affected();
}
tx.commit().await?;
Ok(moved)
})
.await?;
let terminal = self
.rename_application_in_batches(workflow_table, source, new_name, batching)
.await?;
let step_rows = self
.rename_application_in_batches(steps_table, source, new_name, batching)
.await?;
let counts = ApplicationRowCounts {
queues: counts[0],
schedules: counts[1],
versions: counts[2],
workflows: counts[3] + terminal,
steps: step_rows,
};
tracing::debug!(
queues = counts.queues,
schedules = counts.schedules,
versions = counts.versions,
workflows = counts.workflows,
steps = counts.steps,
"renamed application"
);
Ok(counts)
}
async fn record_child_workflow(
&self,
parent_workflow_id: &str,
child_workflow_id: &str,
step_id: i32,
step_name: &str,
started_at: Option<Timestamp>,
owner_xid: &str,
) -> Result<(), Error> {
let pool = &self.pool;
with_retry(&self.retry, "record_child_workflow", move || async move {
let mut tx = pool.begin().await?;
self.check_owner_on(
&mut tx,
WorkflowOwner {
workflow_id: parent_workflow_id,
owner_xid,
},
)
.await?;
self.record_child_workflow_on(
&mut tx,
parent_workflow_id,
child_workflow_id,
step_id,
step_name,
started_at,
)
.await?;
tx.commit().await?;
Ok(())
})
.await
}
}
impl PostgresSystemDatabase {
async fn deliver(
&self,
messages: &[Message<'_>],
serialization: Option<&str>,
caller: Option<(&str, i32)>,
step_name: &'static str,
send_to_forks: bool,
owner: Option<WorkflowOwner<'_>>,
) -> Result<(), Error> {
if messages.is_empty() && caller.is_none() {
return Ok(());
}
let mut keys = HashSet::with_capacity(messages.len());
for message in messages {
if let Some(key) = message.idempotency_key {
if key.is_empty() {
return Err(Error::InvalidInput {
field: "idempotency_key".into(),
detail: "must be absent rather than empty".to_owned(),
});
}
if !keys.insert(key) {
return Err(Error::InvalidInput {
field: "idempotency_key".into(),
detail: format!("{key} is used by more than one message"),
});
}
}
}
let fallback_ids: Vec<String> = messages
.iter()
.map(|_| uuid::Uuid::new_v4().to_string())
.collect();
let timing = StepTiming {
started_at: Timestamp::now(),
completed_at: Timestamp::now(),
};
let notifications_table = self.tables.notifications.as_str();
let workflow_table = self.tables.workflow_status.as_str();
let pool = &self.pool;
let fallback_ids = &fallback_ids;
with_retry(&self.retry, "send_messages", move || async move {
let mut tx = pool.begin().await?;
if let Some(owner) = owner {
self.check_owner_on(&mut tx, owner).await?;
}
if let Some((workflow_id, step_id)) = caller
&& self
.check_step_on(&mut tx, workflow_id, step_id, step_name)
.await?
.is_some()
{
tracing::debug!(
workflow_id,
step_id,
count = messages.len(),
"replaying send"
);
return Ok(());
}
let forks_of = if send_to_forks {
let roots: Vec<&str> = messages.iter().map(|m| m.destination_id).collect();
Self::descendant_forks(&mut tx, workflow_table, &roots).await?
} else {
HashMap::new()
};
let mut destination_ids = Vec::new();
let mut topics = Vec::new();
let mut payloads = Vec::new();
let mut message_ids = Vec::new();
for (message, fallback) in messages.iter().zip(fallback_ids) {
let forks = forks_of
.get(message.destination_id)
.map(Vec::as_slice)
.unwrap_or_default();
for destination in
std::iter::once(message.destination_id).chain(forks.iter().map(String::as_str))
{
destination_ids.push(destination);
topics.push(message.topic.unwrap_or(NULL_TOPIC));
payloads.push(message.message);
message_ids.push(match message.idempotency_key {
Some(key) => format!("{key}::{destination}"),
None => format!("{fallback}::{destination}"),
});
}
}
let sent = if destination_ids.is_empty() {
Ok(Default::default())
} else {
sqlx::query(AssertSqlSafe(format!(
"INSERT INTO {notifications_table} \
(destination_uuid, topic, message, message_uuid, serialization) \
SELECT * FROM unnest($1::text[], $2::text[], $3::text[], $4::text[]) \
AS m(destination_uuid, topic, message, message_uuid), \
(SELECT $5::text) AS s(serialization) \
ON CONFLICT (message_uuid) DO NOTHING"
)))
.bind(&destination_ids)
.bind(&topics)
.bind(&payloads)
.bind(&message_ids)
.bind(serialization)
.execute(&mut *tx)
.await
};
match sent {
Ok(result) => {
tracing::debug!(
requested = messages.len(),
delivered = result.rows_affected(),
send_to_forks,
"sent messages"
);
}
Err(e) if is_foreign_key_violation(&e) => {
let mut missing: Vec<String> = messages
.iter()
.map(|m| m.destination_id.to_owned())
.collect();
missing.sort();
missing.dedup();
return Err(Error::NonExistentWorkflow {
workflow_ids: missing,
});
}
Err(e) => return Err(e.into()),
}
if let Some((workflow_id, step_id)) = caller {
self.record_step_on(
&mut tx,
workflow_id,
step_id,
step_name,
Outcome::Output(None),
None,
Some(timing),
None,
)
.await?;
}
tx.commit().await?;
Ok(())
})
.await
}
}
#[cfg(test)]
mod tests {
use super::{
AssertSqlSafe, NOW_MS_SQL, PgSeverity, PollingConcurrency, PostgresSystemDatabase,
RetryPolicy, Settings, classify_sqlstate, is_transport_failure, latest_version_sql,
split_database, step_listing_query, timed_out_sweep_sql, with_retry,
};
use crate::sysdb::types::{
EncodedValue, NewWorkflow, Outcome, StepCaller, StepCodec, Timestamp,
WorkflowIdReusePolicy, WorkflowStatus,
};
use crate::sysdb::{BackendError, BackendErrorKind, Error, SystemDatabase};
use std::sync::atomic::{AtomicU32, Ordering};
use tokio::sync::Semaphore;
#[test]
fn only_serialization_failures_and_deadlocks_are_conflicts() {
for code in ["40001", "40P01"] {
assert_eq!(
classify_sqlstate(code, Some(PgSeverity::Error), None),
BackendErrorKind::Conflict,
"{code}"
);
}
for code in ["40000", "40002", "40003"] {
assert_eq!(
classify_sqlstate(code, Some(PgSeverity::Error), None),
BackendErrorKind::Permanent,
"{code}"
);
}
}
#[test]
fn the_connection_class_exceptions_are_permanent() {
let cancellations = [
Some("canceling statement due to statement timeout"),
Some("canceling statement due to user request"),
Some("query execution canceled due to statement timeout"),
None,
];
for severity in [Some(PgSeverity::Error), Some(PgSeverity::Fatal), None] {
for message in cancellations {
assert_eq!(
classify_sqlstate("57014", severity, message),
BackendErrorKind::Permanent,
"{severity:?} {message:?}"
);
}
assert_eq!(
classify_sqlstate("57P04", severity, None),
BackendErrorKind::Permanent,
"{severity:?}"
);
}
}
#[test]
fn a_cockroach_session_cancel_is_a_connection_failure() {
assert_eq!(
classify_sqlstate(
"57014",
Some(PgSeverity::Error),
Some("query execution canceled")
),
BackendErrorKind::Connection
);
}
#[test]
fn a_protocol_violation_is_split_by_severity() {
for severity in [Some(PgSeverity::Error), Some(PgSeverity::Log), None] {
assert_eq!(
classify_sqlstate("08P01", severity, None),
BackendErrorKind::Permanent,
"{severity:?}"
);
}
for severity in [PgSeverity::Fatal, PgSeverity::Panic] {
assert_eq!(
classify_sqlstate("08P01", Some(severity), Some("query_wait_timeout")),
BackendErrorKind::Connection,
"{severity:?}"
);
}
}
#[test]
fn connection_classes_match_by_prefix() {
for code in ["08000", "08006", "08007", "53300", "57P01", "57P05"] {
assert_eq!(
classify_sqlstate(code, Some(PgSeverity::Error), None),
BackendErrorKind::Connection,
"{code}"
);
}
assert_eq!(
classify_sqlstate(
"XXUUU",
Some(PgSeverity::Error),
Some("read tcp 10.0.0.1:26257->10.0.0.2:5432: i/o timeout")
),
BackendErrorKind::Connection
);
assert_eq!(
classify_sqlstate("XX000", Some(PgSeverity::Error), Some("internal error")),
BackendErrorKind::Permanent
);
assert_eq!(
classify_sqlstate("23505", Some(PgSeverity::Error), None),
BackendErrorKind::Permanent
);
}
#[test]
fn a_step_listing_without_outputs_reads_no_payload_column() {
let without = step_listing_query("steps", "wf", false, Some(10), Some(5))
.sql()
.as_str()
.to_owned();
let rest = without
.replace("NULL::text AS output", "")
.replace("NULL::text AS error", "");
assert!(
!rest.contains("output") && !rest.contains("error"),
"{without}"
);
assert!(without.contains(" LIMIT $2 OFFSET $3"), "{without}");
let with = step_listing_query("steps", "wf", true, None, None)
.sql()
.as_str()
.to_owned();
assert!(with.contains(", output, error FROM steps"), "{with}");
}
#[test]
fn a_dead_socket_reads_as_a_transport_failure() {
assert!(is_transport_failure(
"read tcp 172.17.0.3:26257->172.17.0.1:37026: i/o timeout"
));
assert!(is_transport_failure(
"write tcp 172.17.0.3:26257->172.17.0.1:37026: broken pipe"
));
assert!(is_transport_failure(
"read tcp 10.0.0.2:26257->10.0.0.9:5432: connection reset by peer"
));
assert!(is_transport_failure("use of closed network connection"));
}
#[test]
fn transport_text_matches_whatever_its_case() {
assert!(is_transport_failure("read tcp: I/O timeout"));
}
#[test]
fn a_real_internal_error_is_not_a_transport_failure() {
assert!(!is_transport_failure(
"internal error: expected LHS of assignment to be a variable"
));
assert!(!is_transport_failure("index corrupted: duplicate key"));
assert!(!is_transport_failure(""));
}
#[test]
fn the_default_polling_cap_is_half_the_pool() {
assert_eq!(PollingConcurrency::Default.permits(10), 5);
assert_eq!(PollingConcurrency::Default.permits(21), 10);
}
#[test]
fn a_pool_too_small_to_halve_still_admits_one_poll() {
assert_eq!(PollingConcurrency::Default.permits(1), 1);
assert_eq!(PollingConcurrency::Default.permits(0), 1);
}
#[test]
fn a_configured_polling_cap_is_taken_as_given() {
let at_most = |n| PollingConcurrency::AtMost(std::num::NonZeroU32::new(n).unwrap());
assert_eq!(at_most(3).permits(10), 3);
assert_eq!(at_most(100).permits(10), 100);
}
#[test]
fn an_unlimited_polling_cap_never_runs_out() {
assert_eq!(
PollingConcurrency::Unlimited.permits(10),
Semaphore::MAX_PERMITS
);
}
#[test]
fn splits_a_url_into_its_maintenance_form_and_database() {
assert_eq!(
split_database("postgresql://u:p@host:5432/dbos_sys"),
Some((
"postgresql://u:p@host:5432/postgres".to_owned(),
"dbos_sys".to_owned()
)),
);
}
#[test]
fn query_parameters_are_kept_on_the_maintenance_url() {
assert_eq!(
split_database("postgresql://h/db?sslmode=require"),
Some((
"postgresql://h/postgres?sslmode=require".to_owned(),
"db".to_owned()
)),
);
}
#[test]
fn a_url_without_a_database_has_nothing_to_split() {
assert_eq!(split_database("postgresql://host:5432/"), None);
assert_eq!(split_database("postgresql://host:5432"), None);
}
const STEP: &str = "test.transactionalStep";
const UNIT_TAG: &str = "unit_json";
#[derive(Debug)]
struct UnitCodec;
impl StepCodec for UnitCodec {
fn encode(&self, value: serde_json::Value) -> Result<EncodedValue, Error> {
Ok(EncodedValue {
value: value.to_string(),
serialization: Some(UNIT_TAG.to_owned()),
})
}
fn decode(
&self,
value: &str,
serialization: Option<&str>,
) -> Result<serde_json::Value, Error> {
match serialization {
Some(UNIT_TAG) => {
serde_json::from_str(value).map_err(|e| Error::Malformed(e.to_string()))
}
other => Err(Error::UnknownSerialization {
name: other.unwrap_or_default().to_owned(),
}),
}
}
}
fn unit_caller() -> Option<StepCaller<'static>> {
Some(StepCaller {
workflow_id: "wf-caller",
step_id: 0,
codec: &UnitCodec,
owner_xid: crate::context::TEST_OWNER,
})
}
async fn caller() -> (PostgresSystemDatabase, dbos_test_support::TestDatabase) {
let db = dbos_test_support::test_database().await;
let sys = PostgresSystemDatabase::from_pool(db.pool().await, &Settings::default());
sys.init_workflow(
&NewWorkflow {
name: Some("checkout"),
..NewWorkflow::new("wf-caller")
},
None,
WorkflowIdReusePolicy::ReturnExisting,
)
.await
.expect("init_workflow failed");
crate::context::own_for_test(&db, "wf-caller").await;
(sys, db)
}
async fn touch(
sys: &PostgresSystemDatabase,
tx: &mut sqlx::Transaction<'static, sqlx::Postgres>,
) -> Result<(), Error> {
let workflow_table = sys.tables.workflow_status.as_str();
sqlx::query(AssertSqlSafe(format!(
"UPDATE {workflow_table} SET name = 'touched' WHERE workflow_uuid = 'wf-caller'"
)))
.execute(&mut **tx)
.await?;
Ok(())
}
async fn name_of_caller(sys: &PostgresSystemDatabase) -> Option<String> {
sys.get_workflow("wf-caller")
.await
.unwrap()
.expect("wf-caller exists")
.name
}
#[tokio::test]
async fn a_malformed_statement_is_refused_not_retried() {
let db = dbos_test_support::test_database().await;
let pool = &db.pool().await;
let outcome = crate::rt::timeout(
std::time::Duration::from_secs(10),
with_retry(&RetryPolicy::default(), "malformed", move || async move {
let sum: i32 = sqlx::query_scalar("SELECT $1::INT4 + $2::INT4")
.bind(1_i32)
.fetch_one(pool)
.await?;
Ok(sum)
}),
)
.await
.expect("a malformed statement returns instead of being retried");
match outcome {
Err(Error::Backend(backend)) => {
assert_eq!(backend.kind, BackendErrorKind::Permanent, "{backend:?}");
if db.backend() == dbos_test_support::Backend::Postgres {
assert_eq!(backend.sqlstate.as_deref(), Some("08P01"), "{backend:?}");
}
}
other => panic!("expected a backend error, got {other:?}"),
}
}
#[tokio::test]
async fn a_statement_timeout_is_refused_not_retried() {
let db = dbos_test_support::test_database().await;
let pool = &db.pool().await;
let outcome = crate::rt::timeout(
std::time::Duration::from_secs(10),
with_retry(
&RetryPolicy::default(),
"statement_timeout",
move || async move {
let mut tx = pool.begin().await?;
sqlx::query("SET LOCAL statement_timeout = '50ms'")
.execute(&mut *tx)
.await?;
sqlx::query("SELECT pg_sleep(5)").execute(&mut *tx).await?;
tx.commit().await?;
Ok(())
},
),
)
.await
.expect("a cancelled statement returns instead of being retried");
match outcome {
Err(Error::Backend(backend)) => {
assert_eq!(backend.kind, BackendErrorKind::Permanent, "{backend:?}");
assert_eq!(backend.sqlstate.as_deref(), Some("57014"), "{backend:?}");
}
other => panic!("expected a backend error, got {other:?}"),
}
}
#[tokio::test]
async fn a_backend_failure_in_the_work_records_nothing() {
let (sys, _db) = caller().await;
let sys = &sys;
let outcome = sys
.run_transactional_step(unit_caller(), STEP, Timestamp::now(), |mut tx| async move {
touch(sys, &mut tx).await?;
let quotient: i64 = sqlx::query_scalar("SELECT 1::BIGINT / 0::BIGINT")
.fetch_one(&mut *tx)
.await?;
Ok((tx, quotient))
})
.await;
match outcome {
Err(Error::Backend(backend)) => {
assert_eq!(backend.kind, BackendErrorKind::Permanent, "{backend:?}")
}
other => panic!("expected a backend error, got {other:?}"),
}
assert_eq!(
sys.check_step("wf-caller", 0, STEP).await.unwrap(),
None,
"a backend failure is never recorded"
);
assert_eq!(name_of_caller(sys).await.as_deref(), Some("checkout"));
}
#[tokio::test]
async fn a_conflict_in_the_work_reruns_to_one_checkpoint() {
let (sys, _db) = caller().await;
let sys = &sys;
let attempts = &AtomicU32::new(0);
let started_at = Timestamp::now();
let run = || {
with_retry(&sys.retry, STEP, move || async move {
sys.run_transactional_step(unit_caller(), STEP, started_at, |mut tx| async move {
let attempt = attempts.fetch_add(1, Ordering::SeqCst);
touch(sys, &mut tx).await?;
if attempt == 0 {
return Err(Error::Backend(BackendError {
message: "could not serialize access".to_owned(),
sqlstate: Some("40001".to_owned()),
kind: BackendErrorKind::Conflict,
}));
}
Ok((tx, attempt))
})
.await
})
};
assert_eq!(run().await.unwrap(), 1);
assert_eq!(attempts.load(Ordering::SeqCst), 2);
let step = sys
.check_step("wf-caller", 0, STEP)
.await
.unwrap()
.expect("the success is recorded");
assert_eq!(step.output.as_deref(), Some("1"), "{step:?}");
assert_eq!(
step.serialization.as_deref(),
Some(UNIT_TAG),
"tagged by the caller's codec: {step:?}"
);
assert!(step.error.is_none(), "{step:?}");
assert_eq!(name_of_caller(sys).await.as_deref(), Some("touched"));
assert_eq!(
run().await.unwrap(),
1,
"the replay returns the recorded output"
);
assert_eq!(
attempts.load(Ordering::SeqCst),
2,
"the replay runs no work"
);
}
#[tokio::test]
async fn a_recorded_refusal_goes_through_the_callers_codec() {
let (sys, _db) = caller().await;
let sys = &sys;
let refusal = || Error::InvalidInput {
field: "workflow_ids".into(),
detail: "refused".to_owned(),
};
let run = || {
sys.run_transactional_step(unit_caller(), STEP, Timestamp::now(), |_tx| async move {
Err::<(sqlx::Transaction<'static, sqlx::Postgres>, i64), _>(refusal())
})
};
assert_eq!(run().await.unwrap_err(), refusal());
let step = sys
.check_step("wf-caller", 0, STEP)
.await
.unwrap()
.expect("the refusal is recorded");
assert_eq!(step.serialization.as_deref(), Some(UNIT_TAG), "{step:?}");
assert_eq!(
run().await.unwrap_err(),
refusal(),
"the replay raises the recorded refusal"
);
}
#[tokio::test]
async fn a_checkpoint_under_an_unreadable_tag_is_refused() {
let (sys, _db) = caller().await;
let sys = &sys;
sys.record_step(
"wf-caller",
0,
STEP,
Outcome::Output(Some("1")),
Some("pickle"),
None,
crate::context::TEST_OWNER,
)
.await
.unwrap();
let ran = &AtomicU32::new(0);
let outcome = sys
.run_transactional_step(unit_caller(), STEP, Timestamp::now(), |tx| async move {
ran.fetch_add(1, Ordering::SeqCst);
Ok((tx, 2_i64))
})
.await;
assert_eq!(
outcome,
Err(Error::UnknownSerialization {
name: "pickle".to_owned()
})
);
assert_eq!(ran.load(Ordering::SeqCst), 0, "the replay runs no work");
}
#[tokio::test]
async fn the_timeout_sweep_plans_through_the_deadline_index() {
let db = dbos_test_support::test_database().await;
if db.backend() == dbos_test_support::Backend::Cockroach {
return; }
let pool = db.pool().await;
let sys = PostgresSystemDatabase::from_pool(pool.clone(), &Settings::default());
let table = sys.tables.workflow_status.as_str();
sqlx::query(AssertSqlSafe(format!(
"INSERT INTO {table} (workflow_uuid, name, status, workflow_deadline_epoch_ms) \
SELECT 'wf-plan-' || g, 'plan_probe', \
CASE WHEN g % 20 = 0 THEN 'PENDING' ELSE 'SUCCESS' END, \
CASE WHEN g % 20 = 0 THEN {NOW_MS_SQL} + 3600000 + g \
ELSE {NOW_MS_SQL} - 60000 - g END \
FROM generate_series(1, 5000) g ORDER BY md5(g::text)"
)))
.execute(&pool)
.await
.expect("insert failed");
sqlx::query(AssertSqlSafe(format!("ANALYZE {table}")))
.execute(&pool)
.await
.unwrap();
let mut custom = Vec::new();
for application_name in [None, Some("plan-app")] {
let plan: Vec<String> = sqlx::query_scalar(AssertSqlSafe(format!(
"EXPLAIN {}",
timed_out_sweep_sql(table)
)))
.bind(application_name)
.bind(10_i64)
.fetch_all(&pool)
.await
.expect("EXPLAIN failed");
custom.push((format!("custom, application {application_name:?}"), plan));
}
let generic: Vec<String> = sqlx::raw_sql(AssertSqlSafe(format!(
"EXPLAIN (GENERIC_PLAN) {}",
timed_out_sweep_sql(table)
)))
.fetch_all(&pool)
.await
.expect("EXPLAIN (GENERIC_PLAN) failed")
.iter()
.map(|row| sqlx::Row::get::<String, _>(row, 0))
.collect();
let generic = ("generic".to_owned(), generic);
for (kind, plan) in custom.iter().chain([&generic]) {
let plan = plan.join("\n");
assert!(
plan.contains("idx_workflow_status_deadline"),
"{kind}: {plan}"
);
assert!(
plan.lines().any(|line| line.contains("Index Cond")
&& line.contains("workflow_deadline_epoch_ms")),
"{kind}: the deadline is an index condition: {plan}"
);
}
for (kind, plan) in &custom {
let plan = plan.join("\n");
assert!(!plan.contains("Seq Scan"), "{kind}: {plan}");
}
}
fn versions_handle(
pool: &sqlx::PgPool,
application_name: Option<&str>,
) -> PostgresSystemDatabase {
PostgresSystemDatabase::from_pool(
pool.clone(),
&Settings {
application_name,
..Settings::default()
},
)
}
async fn register_versions(sys: &PostgresSystemDatabase, pool: &sqlx::PgPool, per_owner: i64) {
let table = &sys.tables.application_versions;
sqlx::query(AssertSqlSafe(format!(
"INSERT INTO {table} \
(version_id, version_name, version_timestamp, created_at, application_name) \
SELECT g::text, 'version-' || g::text, g, g, \
CASE WHEN g <= $1 THEN 'ours' WHEN g <= 2 * $1 THEN NULL ELSE 'peer' END \
FROM generate_series(1, 3 * $1) g"
)))
.bind(per_owner)
.execute(pool)
.await
.expect("insert failed");
}
#[tokio::test]
async fn each_application_sees_its_own_and_unclaimed_versions() {
let db = dbos_test_support::test_database().await;
let pool = db.pool().await;
register_versions(&versions_handle(&pool, None), &pool, 10).await;
for (owner, expected) in [
(Some("ours"), "version-20"),
(Some("absent"), "version-20"),
(Some("peer"), "version-30"),
(None, "version-30"),
] {
let sys = versions_handle(&pool, owner);
let latest = sys
.get_latest_application_version(None)
.await
.expect("lookup failed")
.expect("a version is registered");
assert_eq!(latest.version_name, expected, "{owner:?}");
let mut conn = pool.acquire().await.unwrap();
assert!(
sys.is_latest_application_version(&mut conn, expected)
.await
.unwrap(),
"{owner:?}: the dequeue agrees"
);
assert!(
!sys.is_latest_application_version(&mut conn, "version-1")
.await
.unwrap(),
"{owner:?}"
);
}
}
#[tokio::test]
async fn the_latest_version_is_found_without_walking_the_history() {
let db = dbos_test_support::test_database().await;
if db.backend() == dbos_test_support::Backend::Cockroach {
return; }
let pool = db.pool().await;
let unnamed = versions_handle(&pool, None);
let table = &unnamed.tables.application_versions;
register_versions(&unnamed, &pool, 10_000).await;
sqlx::raw_sql(AssertSqlSafe(format!("VACUUM ANALYZE {table}")))
.execute(&pool)
.await
.unwrap();
for (owner, expected) in [
(Some("ours"), "version-20000"),
(Some("absent"), "version-20000"),
(Some("peer"), "version-30000"),
(None, "version-30000"),
] {
let latest = versions_handle(&pool, owner)
.get_latest_application_version(None)
.await
.expect("lookup failed")
.expect("a version is registered");
assert_eq!(latest.version_name, expected, "{owner:?}");
let explain = sqlx::query(AssertSqlSafe(format!(
"EXPLAIN (ANALYZE, BUFFERS, FORMAT JSON) {}",
latest_version_sql(table, owner.is_some())
)));
let explain = match owner {
Some(owner) => explain.bind(owner),
None => explain,
};
let row = explain.fetch_one(&pool).await.expect("EXPLAIN failed");
let plan = sqlx::Row::try_get_raw(&row, 0).unwrap().as_str().unwrap();
let plan: serde_json::Value = serde_json::from_str(plan).unwrap();
let plan = &plan[0]["Plan"];
let blocks = plan["Shared Hit Blocks"].as_u64().unwrap()
+ plan["Shared Read Blocks"].as_u64().unwrap();
assert!(blocks < 32, "{owner:?}: {blocks} blocks: {plan:#}");
fn nodes<'a>(node: &'a serde_json::Value, into: &mut Vec<&'a serde_json::Value>) {
into.push(node);
for child in node["Plans"].as_array().into_iter().flatten() {
nodes(child, into);
}
}
let mut all = Vec::new();
nodes(plan, &mut all);
let scans: Vec<_> = all
.into_iter()
.filter(|node| node["Node Type"].as_str().unwrap().contains("Scan"))
.collect();
assert!(!scans.is_empty(), "{owner:?}: {plan:#}");
for scan in scans {
assert!(
scan["Node Type"].as_str().unwrap().contains("Index"),
"{owner:?}: {plan:#}"
);
let rows =
scan["Actual Rows"].as_f64().unwrap() * scan["Actual Loops"].as_f64().unwrap();
assert!(rows <= 1.0, "{owner:?}: {plan:#}");
assert_eq!(
scan["Rows Removed by Filter"].as_f64().unwrap_or(0.0),
0.0,
"{owner:?}: {plan:#}"
);
}
}
}
#[tokio::test]
async fn promotion_and_re_registration_keep_a_versions_identity() {
let db = dbos_test_support::test_database().await;
let pool = db.pool().await;
let sys = versions_handle(&pool, Some("ours"));
register_versions(&sys, &pool, 10).await;
sys.promote_application_version("version-1", None)
.await
.expect("promotion failed");
let promoted = sys
.get_latest_application_version(None)
.await
.unwrap()
.expect("a version is registered");
assert_eq!(promoted.version_id, "1");
assert_eq!(promoted.version_name, "version-1");
assert_eq!(promoted.created_at, Timestamp::from_epoch_ms(1));
assert_eq!(promoted.app_name.as_deref(), Some("ours"));
sys.create_application_version("version-1", None)
.await
.expect("re-registration failed");
assert_eq!(
sys.get_latest_application_version(None).await.unwrap(),
Some(promoted)
);
let peer = sys
.get_latest_application_version(Some("peer"))
.await
.unwrap()
.expect("the peer has versions");
assert_eq!(peer.version_name, "version-30");
}
#[tokio::test]
async fn an_application_seeing_no_versions_has_no_latest() {
let db = dbos_test_support::test_database().await;
let pool = db.pool().await;
let sys = versions_handle(&pool, Some("absent"));
register_versions(&sys, &pool, 10).await;
let table = &sys.tables.application_versions;
sqlx::query(AssertSqlSafe(format!(
"DELETE FROM {table} WHERE application_name IS NULL"
)))
.execute(&pool)
.await
.unwrap();
assert_eq!(
sys.get_latest_application_version(None).await.unwrap(),
None
);
let mut conn = pool.acquire().await.unwrap();
assert!(
sys.is_latest_application_version(&mut conn, "anything")
.await
.unwrap()
);
}
#[test]
fn an_idle_transaction_timeout_is_retried_as_a_lost_connection() {
assert_eq!(
classify_sqlstate("25P03", Some(PgSeverity::Fatal), None),
BackendErrorKind::Connection
);
}
#[tokio::test]
async fn connect_sets_the_idle_transaction_timeout() {
use super::IdleTransactionTimeout::{After, Default, ServerSetting};
async fn show(pool: &sqlx::PgPool) -> String {
sqlx::query_scalar("SHOW idle_in_transaction_session_timeout")
.fetch_one(pool)
.await
.unwrap()
}
async fn connected(url: &str, timeout: super::IdleTransactionTimeout) -> String {
let sys = PostgresSystemDatabase::connect(&super::Config {
idle_transaction_timeout: timeout,
migrate: false,
..super::Config::new(url)
})
.await
.expect("connect failed");
show(&sys.pool).await
}
let db = dbos_test_support::test_database().await;
let with = |query: &str| {
let separator = if db.url().contains('?') { '&' } else { '?' };
format!("{}{separator}{query}", db.url())
};
let set_directly = async |millis: u64| {
let url = with(&format!(
"options=-c%20idle_in_transaction_session_timeout%3D{millis}"
));
show(&sqlx::PgPool::connect(&url).await.unwrap()).await
};
let ms = std::time::Duration::from_millis;
let server = show(&db.pool().await).await;
assert_eq!(
connected(db.url(), Default).await,
set_directly(60_000).await
);
assert_eq!(
connected(db.url(), After(ms(300))).await,
set_directly(300).await
);
assert_eq!(
connected(db.url(), After(std::time::Duration::from_micros(500))).await,
set_directly(1).await,
"under a millisecond is a millisecond"
);
assert_eq!(connected(db.url(), ServerSetting).await, server);
let own = set_directly(1234).await;
for options in [
"options=-c%20idle_in_transaction_session_timeout%3D1234",
"options=-c%20IDLE%5FIN%5FTRANSACTION%5FSESSION%5FTIMEOUT%3D1234",
"options=--idle-in-transaction-session-timeout%3D1234",
] {
assert_eq!(connected(&with(options), Default).await, own, "{options}");
}
for query in [
"idle_in_transaction_session_timeout=1234",
"options=-c%20application_name%3Didle_in_transaction_session_timeout",
] {
assert_eq!(
connected(&with(query), After(ms(300))).await,
set_directly(300).await,
"{query}"
);
}
let supplied = PostgresSystemDatabase::from_pool(db.pool().await, &Settings::default());
assert_eq!(
show(&supplied.pool).await,
server,
"a pool the caller built is left alone"
);
}
#[test]
fn option_names_are_read_as_the_server_reads_them() {
let names = |options| super::option_names(options).collect::<Vec<_>>();
assert_eq!(
names("-c a=1 -cB=2 --c-d=3 -c e\\ f=4"),
["a", "b", "c_d", "e f"]
);
assert_eq!(
names("-c application_name=idle_in_transaction_session_timeout"),
["application_name"]
);
assert!(names("-B 100 stray -c").is_empty());
}
#[tokio::test]
async fn connect_ends_transactions_left_idle() {
let db = dbos_test_support::test_database().await;
if db.backend() == dbos_test_support::Backend::Cockroach {
return;
}
let sys = PostgresSystemDatabase::connect(&super::Config {
idle_transaction_timeout: super::IdleTransactionTimeout::After(
std::time::Duration::from_millis(300),
),
migrate: false,
..super::Config::new(db.url())
})
.await
.expect("connect failed");
sys.init_workflow(
&NewWorkflow::new("wf-idle"),
None,
WorkflowIdReusePolicy::ReturnExisting,
)
.await
.unwrap();
let table = sys.tables.workflow_status.as_str();
let mut idle = sys.pool.begin().await.unwrap();
sqlx::query(AssertSqlSafe(format!(
"SELECT owner_xid FROM {table} WHERE workflow_uuid = 'wf-idle' FOR NO KEY UPDATE"
)))
.execute(&mut *idle)
.await
.unwrap();
crate::rt::sleep(std::time::Duration::from_secs(1)).await;
let other = db.pool().await;
crate::rt::timeout(
std::time::Duration::from_secs(5),
sqlx::query(AssertSqlSafe(format!(
"UPDATE {table} SET owner_xid = NULL WHERE workflow_uuid = 'wf-idle'"
)))
.execute(&other),
)
.await
.expect("the idle transaction still holds its lock")
.unwrap();
assert!(
sqlx::query("SELECT 1").execute(&mut *idle).await.is_err(),
"the idle transaction's session was ended"
);
}
#[tokio::test]
async fn a_retried_insert_keeps_the_input_a_bounce_wrote_since() {
use crate::sysdb::PORTABLE_JSON;
use crate::sysdb::types::{Debounce, DebounceRequest, QualifiedWorkflowName};
use std::time::Duration;
let db = dbos_test_support::test_database().await;
let sys = PostgresSystemDatabase::from_pool(db.pool().await, &Settings::default());
let hour = Duration::from_secs(3600);
let wf = NewWorkflow {
name: Some("digest"),
queue_name: Some("orders"),
deduplication_id: Some("digest-key"),
is_debounced: true,
delay: Some(hour),
input: Some("\"first\""),
serialization: Some("rust_serde"),
..NewWorkflow::new("wf-debounced")
};
let first = sys
.init_workflow_as(&wf, None, WorkflowIdReusePolicy::ReturnExisting, "creator")
.await
.expect("the insert failed");
assert!(first.should_execute);
let key = QualifiedWorkflowName::new("digest");
let bounce = DebounceRequest {
workflow: &key,
queue_name: "orders",
deduplication_id: "digest-key",
period: hour,
inputs: Some("\"later\""),
serialization: Some(PORTABLE_JSON),
application_name: None,
};
assert_eq!(
sys.debounce_delayed_workflow(&bounce, None).await.unwrap(),
Debounce::Bounced {
workflow_id: "wf-debounced".to_owned()
}
);
let retried = sys
.init_workflow_as(&wf, None, WorkflowIdReusePolicy::ReturnExisting, "creator")
.await
.expect("the retry failed");
assert!(retried.should_execute, "the retry recognises its own row");
let row = sys.get_workflow("wf-debounced").await.unwrap().unwrap();
assert_eq!(
row.input.as_deref(),
Some("\"later\""),
"the bounce's input"
);
assert_eq!(row.serialization.as_deref(), Some(PORTABLE_JSON));
}
#[tokio::test]
async fn a_rejecting_insert_refuses_another_callers_workflow_but_not_its_own_retry() {
let db = dbos_test_support::test_database().await;
let pool = db.pool().await;
let sys = PostgresSystemDatabase::from_pool(pool.clone(), &Settings::default());
let first = NewWorkflow {
name: Some("charge"),
input: Some("\"first\""),
serialization: Some("rust_serde"),
..NewWorkflow::new("wf-taken")
};
let second = NewWorkflow {
input: Some("\"second\""),
..first.clone()
};
let reject = WorkflowIdReusePolicy::Reject;
let created = sys
.init_workflow_as(&first, None, reject, "creator")
.await
.expect("the insert failed");
assert!(created.should_execute);
let retried = sys
.init_workflow_as(&first, None, reject, "creator")
.await
.expect("a retry of the insert that created the row is not a reuse");
assert!(retried.should_execute, "the retry recognises its own row");
let joined = sys
.init_workflow_as(
&second,
None,
WorkflowIdReusePolicy::ReturnExisting,
"joiner",
)
.await
.expect("the join failed");
assert!(!joined.should_execute, "a joiner does not run the workflow");
let before = sys.get_workflow("wf-taken").await.unwrap().unwrap();
for status in [
WorkflowStatus::Pending,
WorkflowStatus::Success,
WorkflowStatus::Error,
] {
sqlx::query("UPDATE dbos.workflow_status SET status = $1 WHERE workflow_uuid = $2")
.bind(status.as_str())
.bind("wf-taken")
.execute(&pool)
.await
.unwrap();
let refused = sys
.init_workflow_as(&second, None, reject, "rejecter")
.await;
assert_eq!(
refused.map(|_| ()),
Err(Error::WorkflowIdInUse {
workflow_id: "wf-taken".to_owned(),
status,
name: Some("charge".to_owned()),
}),
);
}
let after = sys.get_workflow("wf-taken").await.unwrap().unwrap();
sqlx::query("UPDATE dbos.workflow_status SET name = '' WHERE workflow_uuid = 'wf-taken'")
.execute(&pool)
.await
.unwrap();
let nameless = NewWorkflow {
name: None,
..second.clone()
};
assert!(matches!(
sys.init_workflow_as(&nameless, None, reject, "rejecter")
.await,
Err(Error::WorkflowIdInUse { name: None, .. })
));
assert_eq!(
after.input.as_deref(),
Some("\"first\""),
"the input is the first's"
);
assert_eq!(
after.updated_at, before.updated_at,
"a refusal writes nothing"
);
}
#[tokio::test]
async fn a_rejecting_child_start_adopts_only_its_parents_unrecorded_child() {
use crate::sysdb::types::{InitWorkflowCaller, WorkflowInitResult};
let db = dbos_test_support::test_database().await;
let pool = db.pool().await;
let sys = PostgresSystemDatabase::from_pool(pool.clone(), &Settings::default());
let reject = WorkflowIdReusePolicy::Reject;
let mut owners = std::collections::HashMap::new();
for parent in ["parent-a", "parent-b"] {
let created = sys
.init_workflow(
&NewWorkflow::new(parent),
None,
WorkflowIdReusePolicy::ReturnExisting,
)
.await
.unwrap();
owners.insert(parent, created.creator_xid);
}
let child = NewWorkflow {
name: Some("child"),
..NewWorkflow::new("the-child")
};
let at = |parent_workflow_id: &'static str, step_id| {
Some(InitWorkflowCaller {
parent_workflow_id,
owner_xid: &owners[parent_workflow_id],
step_id,
step_name: "child",
started_at: Timestamp::now(),
})
};
let in_use = |result: Result<WorkflowInitResult, Error>| {
matches!(result, Err(Error::WorkflowIdInUse { .. }))
};
sys.init_workflow_as(&child, at("parent-a", 1), reject, "first")
.await
.expect("a fresh id is not in use");
assert!(
in_use(
sys.init_workflow_as(&child, at("parent-a", 3), reject, "again")
.await
),
"the parent still records the child, so another start of the id is a reuse"
);
assert!(
in_use(
sys.init_workflow_as(&child, at("parent-b", 1), reject, "other")
.await
),
"another parent's start is a reuse"
);
sqlx::query("DELETE FROM dbos.operation_outputs WHERE workflow_uuid = 'parent-a'")
.execute(&pool)
.await
.unwrap();
let adopted = sys
.init_workflow_as(&child, at("parent-a", 1), reject, "replay")
.await
.expect("the replayed parent adopts its own child");
assert!(!adopted.should_execute, "adopted, not run again");
let retried = sys
.init_workflow_as(&child, at("parent-a", 1), reject, "replay")
.await
.expect("a retried adoption adopts again");
assert!(!retried.should_execute);
assert!(
in_use(
sys.init_workflow_as(&child, at("parent-b", 1), reject, "other")
.await
),
"the adoption is the parent's alone"
);
}
}