use crate::defs::{AsQuery, AsSignal, AsUpdate, AsWorkflow, Make, Op};
use crate::error::{Error, Result};
use crate::failure::{Failure, IntoFailure};
use crate::limits::{Acquire, Priority};
use crate::runtime::{Host, LocalFuture, ParentClose, RunOutcome};
use fusee_core::record::{Lsn, PAUSED, Payload, RunId, Seq, Status, policy, value_kind};
use fusee_core::run::{AwaitStart, Outcome, RunInfo, StepStart};
use rustc_hash::FxHashMap;
use serde::Serialize;
use serde::de::DeserializeOwned;
use std::cell::{Cell, RefCell};
use std::collections::BTreeSet;
use std::future::Future;
use std::pin::Pin;
use std::rc::Rc;
use std::sync::{Arc, Mutex, MutexGuard};
use std::task::{Context, Poll, Wake, Waker};
use std::time::Duration;
pub type RetryIf = Arc<dyn Fn(&Failure, u32) -> bool + Send + Sync>;
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum Exhausted {
#[default]
Fail,
Pause,
}
#[derive(Clone)]
pub struct RetryPolicy {
pub attempts: u32,
pub initial: Duration,
pub factor: f64,
pub max: Duration,
pub jitter: bool,
pub max_elapsed: Option<Duration>,
pub non_retryable: Vec<String>,
pub retry_if: Option<RetryIf>,
pub on_exhausted: Option<Exhausted>,
}
impl std::fmt::Debug for RetryPolicy {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("RetryPolicy")
.field("attempts", &self.attempts)
.field("initial", &self.initial)
.field("factor", &self.factor)
.field("max", &self.max)
.field("jitter", &self.jitter)
.field("max_elapsed", &self.max_elapsed)
.field("non_retryable", &self.non_retryable)
.field("retry_if", &self.retry_if.as_ref().map(|_| "fn"))
.field("on_exhausted", &self.on_exhausted)
.finish()
}
}
impl PartialEq for RetryPolicy {
fn eq(&self, o: &Self) -> bool {
self.attempts == o.attempts
&& self.initial == o.initial
&& self.factor == o.factor
&& self.max == o.max
&& self.jitter == o.jitter
&& self.max_elapsed == o.max_elapsed
&& self.non_retryable == o.non_retryable
&& self.on_exhausted == o.on_exhausted
&& match (&self.retry_if, &o.retry_if) {
(None, None) => true,
(Some(a), Some(b)) => Arc::ptr_eq(a, b),
_ => false,
}
}
}
impl RetryPolicy {
pub fn none() -> Self {
Self {
attempts: 1,
..Self::default()
}
}
pub fn exponential(attempts: u32, initial: Duration, max: Duration) -> Self {
Self {
attempts,
initial,
max,
..Self::default()
}
}
pub fn attempts(attempts: u32) -> Self {
Self {
attempts,
..Self::default()
}
}
pub fn no_jitter(mut self) -> Self {
self.jitter = false;
self
}
pub fn max_elapsed(mut self, limit: Duration) -> Self {
self.max_elapsed = Some(limit);
self
}
pub fn non_retryable(mut self, kind: impl Into<String>) -> Self {
self.non_retryable.push(kind.into());
self
}
pub fn retry_if(mut self, f: impl Fn(&Failure, u32) -> bool + Send + Sync + 'static) -> Self {
self.retry_if = Some(Arc::new(f));
self
}
pub fn on_exhausted(mut self, then: Exhausted) -> Self {
self.on_exhausted = Some(then);
self
}
pub fn allows(&self, failure: &Failure, attempt: u32) -> bool {
attempt < self.attempts && self.retryable(failure, attempt)
}
pub fn retryable(&self, failure: &Failure, attempt: u32) -> bool {
failure.is_retryable()
&& !self.non_retryable.iter().any(|k| failure.has_kind(k))
&& self.retry_if.as_ref().is_none_or(|f| f(failure, attempt))
}
pub fn delay(&self, failed_attempt: u32) -> Duration {
let exp = self.factor.powi(failed_attempt.saturating_sub(1) as i32);
let d = self.initial.as_secs_f64() * exp;
Duration::from_secs_f64(d.min(self.max.as_secs_f64()))
}
}
impl Default for RetryPolicy {
fn default() -> Self {
Self {
attempts: 3,
initial: Duration::from_millis(100),
factor: 2.0,
max: Duration::from_secs(10),
jitter: true,
max_elapsed: None,
non_retryable: Vec::new(),
retry_if: None,
on_exhausted: None,
}
}
}
enum Next {
Retry(u64),
Exhausted,
Final,
}
#[derive(Clone, Debug, Default, PartialEq)]
pub struct StepOptions {
pub retry: RetryPolicy,
pub timeout: Option<Duration>,
pub heartbeat_timeout: Option<Duration>,
pub limit: Option<String>,
pub priority: Priority,
pub read_only: bool,
pub own_value: bool,
}
impl StepOptions {
pub fn own_value(mut self) -> Self {
self.own_value = true;
self
}
pub fn read_only(mut self) -> Self {
self.read_only = true;
self
}
pub fn retry(mut self, retry: RetryPolicy) -> Self {
self.retry = retry;
self
}
pub fn timeout(mut self, timeout: Duration) -> Self {
self.timeout = Some(timeout);
self
}
pub fn heartbeat_timeout(mut self, timeout: Duration) -> Self {
self.heartbeat_timeout = Some(timeout);
self
}
pub fn limit(mut self, name: &str) -> Self {
self.limit = Some(name.to_owned());
self
}
pub fn limit_by(mut self, name: &str, key: &str) -> Self {
self.limit = Some(format!("{name}/{key}"));
self
}
pub fn priority(mut self, level: i32) -> Self {
self.priority.level = level;
self
}
pub fn fairness(mut self, key: &str) -> Self {
self.priority.fairness_key = key.to_owned();
self
}
pub fn fairness_weight(mut self, weight: u32) -> Self {
self.priority.fairness_weight = weight;
self
}
}
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct WaitOptions {
pub timeout: Option<Duration>,
}
impl<M: Make<StepOptions>> Op<StepOptions, M> {
pub fn own_value(mut self) -> Self {
self.options_mut().own_value = true;
self
}
pub fn retry(mut self, retry: RetryPolicy) -> Self {
self.options_mut().retry = retry;
self
}
pub fn timeout(mut self, timeout: Duration) -> Self {
self.options_mut().timeout = Some(timeout);
self
}
pub fn heartbeat_timeout(mut self, timeout: Duration) -> Self {
self.options_mut().heartbeat_timeout = Some(timeout);
self
}
pub fn read_only(mut self) -> Self {
self.options_mut().read_only = true;
self
}
pub fn limit(mut self, name: &str) -> Self {
self.options_mut().limit = Some(name.to_owned());
self
}
pub fn limit_by(mut self, name: &str, key: &str) -> Self {
self.options_mut().limit = Some(format!("{name}/{key}"));
self
}
pub fn priority(mut self, level: i32) -> Self {
self.options_mut().priority.level = level;
self
}
pub fn fairness(mut self, key: &str) -> Self {
self.options_mut().priority.fairness_key = key.to_owned();
self
}
pub fn fairness_weight(mut self, weight: u32) -> Self {
self.options_mut().priority.fairness_weight = weight;
self
}
}
impl<M: Make<WaitOptions>> Op<WaitOptions, M> {
pub fn timeout(mut self, timeout: Duration) -> Self {
self.options_mut().timeout = Some(timeout);
self
}
}
impl<M: Make<ChildOptions>> Op<ChildOptions, M> {
pub fn parent_close(mut self, close: ParentClose) -> Self {
self.options_mut().close = close;
self
}
pub fn timeout(mut self, timeout: Duration) -> Self {
self.options_mut().timeout = Some(timeout);
self
}
}
#[derive(Clone, Debug, Default, PartialEq)]
pub struct ChildOptions {
pub close: ParentClose,
pub timeout: Option<Duration>,
}
pub type QueryFn = Rc<dyn Fn(Payload) -> std::result::Result<Payload, String>>;
pub type ValidateFn = Rc<dyn Fn(&Payload) -> std::result::Result<(), String>>;
pub type UpdateFn = Rc<dyn Fn(Payload) -> LocalFuture>;
#[derive(Default)]
pub struct Handlers {
queries: RefCell<FxHashMap<String, QueryFn>>,
updates: RefCell<FxHashMap<String, (ValidateFn, UpdateFn)>>,
}
impl Handlers {
pub fn query(&self, name: &str) -> Option<QueryFn> {
self.queries.borrow().get(name).cloned()
}
pub fn update(&self, name: &str) -> Option<(ValidateFn, UpdateFn)> {
self.updates.borrow().get(name).cloned()
}
pub(crate) fn clear(&self) {
let queries = std::mem::take(&mut *self.queries.borrow_mut());
let updates = std::mem::take(&mut *self.updates.borrow_mut());
drop(queries);
drop(updates);
}
}
#[derive(Clone)]
pub struct Ctx {
host: Rc<dyn Host>,
run: RunId,
s: Rc<CtxState>,
handlers: Rc<Handlers>,
outside: Rc<Cell<u32>>,
}
struct CtxState {
info: RunInfo,
current: Cell<Option<Seq>>,
beats: RefCell<FxHashMap<Seq, u64>>,
compensations: RefCell<Vec<(String, Compensation)>>,
}
type Compensation = Rc<dyn Fn() -> Pin<Box<dyn Future<Output = std::result::Result<(), Failure>>>>>;
fn payload_of(bytes: &[u8]) -> Payload {
thread_local! {
static NULL: Payload = Arc::from(&b"null"[..]);
}
if *bytes == *b"null"
&& let Ok(null) = NULL.try_with(Arc::clone)
{
return null;
}
Arc::from(bytes)
}
pub(crate) fn millis(d: Duration) -> u64 {
u64::try_from(d.as_millis()).unwrap_or(u64::MAX)
}
pub(crate) fn after(now_ms: u64, d: Duration) -> u64 {
now_ms.saturating_add(millis(d))
}
pub(crate) fn encode<T: Serialize + ?Sized>(value: &T) -> Result<Payload> {
thread_local! {
static SCRATCH: RefCell<Vec<u8>> = const { RefCell::new(Vec::new()) };
}
const KEEP: usize = 64 << 10;
let kept = SCRATCH.try_with(|scratch| {
let mut buf = scratch.try_borrow_mut().ok()?;
buf.clear();
let r = serde_json::to_writer(&mut *buf, value).map(|()| payload_of(&buf));
if buf.capacity() > KEEP {
*buf = Vec::new();
}
Some(r)
});
match kept {
Ok(Some(r)) => Ok(r?),
_ => Ok(Arc::from(serde_json::to_vec(value)?)),
}
}
pub(crate) fn decode<T: DeserializeOwned>(bytes: &[u8]) -> Result<T> {
if bytes.is_empty() {
return Ok(serde_json::from_slice(b"null")?);
}
Ok(serde_json::from_slice(bytes)?)
}
fn text(p: &[u8]) -> String {
String::from_utf8_lossy(p).into_owned()
}
enum Attempted<T> {
Settled(Result<T>),
Payload(Payload),
}
fn unencodable(e: Error) -> Failure {
Failure::fatal(format!(
"the step's result cannot be serialized as JSON ({e}); return a value serde_json can encode"
))
.kind("Codec")
}
fn jitter_fraction(run: RunId, seq: Seq, attempt: u32, now: u64) -> f64 {
let mut x =
run.rotate_left(32) ^ (u64::from(seq) << 8) ^ u64::from(attempt) ^ now.rotate_left(17);
x = x.wrapping_add(0x9E37_79B9_7F4A_7C15);
x = (x ^ (x >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
x = (x ^ (x >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
x ^= x >> 31;
(x >> 11) as f64 / (1u64 << 53) as f64
}
struct Delivery {
host: Rc<dyn Host>,
run: RunId,
seq: Seq,
done: bool,
}
impl Future for Delivery {
type Output = Outcome;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Outcome> {
match self.host.take_resolution(self.run, self.seq, cx.waker()) {
Some(o) => {
self.done = true;
Poll::Ready(o)
}
None => Poll::Pending,
}
}
}
impl Drop for Delivery {
fn drop(&mut self) {
if !self.done {
self.host.forget(self.run, self.seq);
}
}
}
struct Gate {
host: Rc<dyn Host>,
run: RunId,
seq: Seq,
gate: Lsn,
not_before: u64,
timer: Option<u64>,
}
impl Gate {
fn new(host: Rc<dyn Host>, run: RunId, seq: Seq, gate: Lsn, not_before: u64) -> Self {
Self {
host,
run,
seq,
gate,
not_before,
timer: None,
}
}
}
impl Future for Gate {
type Output = ();
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
if self.not_before == PAUSED {
if self.host.step_paused(self.run, self.seq) {
return Poll::Pending;
}
self.not_before = 0;
}
if self.not_before > 0 && self.host.now_ms() < self.not_before {
if self.timer.is_none() {
self.timer = Some(self.host.wake_at(self.not_before, cx.waker()));
}
return Poll::Pending;
}
if self
.host
.may_run(self.run, self.gate, self.not_before, cx.waker())
{
Poll::Ready(())
} else {
Poll::Pending
}
}
}
impl Drop for Gate {
fn drop(&mut self) {
if let Some(id) = self.timer {
self.host.cancel_timer(id);
}
}
}
struct Until {
host: Rc<dyn Host>,
at: u64,
timer: Option<u64>,
}
impl Future for Until {
type Output = ();
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
if self.host.now_ms() >= self.at {
return Poll::Ready(());
}
if self.timer.is_none() {
self.timer = Some(self.host.wake_at(self.at, cx.waker()));
}
Poll::Pending
}
}
impl Drop for Until {
fn drop(&mut self) {
if let Some(id) = self.timer {
self.host.cancel_timer(id);
}
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Completion {
pub token: String,
pub url: Option<String>,
}
thread_local! {
pub(crate) static CALLBACK_BASE: RefCell<Option<Arc<std::sync::RwLock<Option<String>>>>> =
const { RefCell::new(None) };
}
impl Completion {
fn new(token: String) -> Self {
let base = CALLBACK_BASE.with(|b| {
b.borrow()
.as_ref()
.and_then(|b| b.read().ok().and_then(|u| u.clone()))
});
let url = base.map(|b| {
format!(
"{}/v1/complete/{}",
b.trim_end_matches('/'),
crate::http::encode_segment(&token)
)
});
Self { token, url }
}
}
pub struct WorkflowLog<'a> {
ctx: &'a Ctx,
}
impl WorkflowLog<'_> {
pub fn debug(&self, message: impl std::fmt::Display) {
self.emit(crate::log::Level::Debug, message);
}
pub fn info(&self, message: impl std::fmt::Display) {
self.emit(crate::log::Level::Info, message);
}
pub fn warn(&self, message: impl std::fmt::Display) {
self.emit(crate::log::Level::Warn, message);
}
pub fn error(&self, message: impl std::fmt::Display) {
self.emit(crate::log::Level::Error, message);
}
fn emit(&self, level: crate::log::Level, message: impl std::fmt::Display) {
let c = self.ctx;
let seq = c.s.current.get();
if seq.is_none() && c.host.replaying(c.run) {
return;
}
let step = seq.and_then(|s| c.host.op_name(c.run, s));
let attempt = seq.map(|s| c.host.step_attempt(c.run, s));
let message = message.to_string();
crate::log::emit(&crate::log::Record {
level,
source: "workflow",
message: &message,
workflow_id: Some(&c.s.info.id),
workflow_type: Some(&c.s.info.wtype),
run: Some(c.run),
step: step.as_deref(),
attempt,
});
}
}
struct Keyed<'a, F> {
fut: Pin<&'a mut F>,
current: &'a Cell<Option<Seq>>,
outside: &'a Cell<u32>,
seq: Seq,
pending: bool,
}
impl<F: Future> Future for Keyed<'_, F> {
type Output = F::Output;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<F::Output> {
let prev = self.current.replace(Some(self.seq));
let r = self.fut.as_mut().poll(cx);
self.current.set(prev);
let pending = r.is_pending();
if pending != self.pending {
self.pending = pending;
let n = self.outside.get();
self.outside
.set(if pending { n + 1 } else { n.saturating_sub(1) });
}
r
}
}
impl<F> Drop for Keyed<'_, F> {
fn drop(&mut self) {
if self.pending {
self.outside.set(self.outside.get().saturating_sub(1));
}
}
}
struct BeatWatch {
host: Rc<dyn Host>,
state: Rc<CtxState>,
seq: Seq,
timeout_ms: u64,
armed: Option<(u64, u64)>,
}
impl Future for BeatWatch {
type Output = ();
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
let last = self
.state
.beats
.borrow()
.get(&self.seq)
.copied()
.unwrap_or(0);
let at = last + self.timeout_ms;
if self.host.now_ms() >= at {
return Poll::Ready(());
}
if self.armed.is_none_or(|(armed_for, _)| armed_for != at) {
if let Some((_, id)) = self.armed.take() {
self.host.cancel_timer(id);
}
let id = self.host.wake_at(at, cx.waker());
self.armed = Some((at, id));
}
Poll::Pending
}
}
impl Drop for BeatWatch {
fn drop(&mut self) {
if let Some((_, id)) = self.armed {
self.host.cancel_timer(id);
}
}
}
struct Permit {
host: Rc<dyn Host>,
limit: String,
ticket: u64,
}
impl Drop for Permit {
fn drop(&mut self) {
self.host.release(&self.limit, self.ticket);
}
}
struct AcquirePermit {
permit: Option<Permit>,
priority: Priority,
run: RunId,
}
impl Future for AcquirePermit {
type Output = Permit;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Permit> {
let Some(p) = self.permit.as_ref() else {
return Poll::Pending;
};
match p.host.acquire_in(
self.run,
&p.limit,
p.ticket,
&self.priority,
Some(cx.waker()),
) {
Acquire::Granted => match self.permit.take() {
Some(p) => Poll::Ready(p),
None => Poll::Pending,
},
Acquire::Wait { .. } => Poll::Pending,
}
}
}
impl Ctx {
#[doc(hidden)]
pub fn input_bytes<T>(&self, f: impl FnOnce(&[u8]) -> T) -> T {
f(&self.s.info.input)
}
pub(crate) fn new(
host: Rc<dyn Host>,
run: RunId,
info: RunInfo,
handlers: Rc<Handlers>,
outside: Rc<Cell<u32>>,
) -> Self {
Self {
host,
run,
s: Rc::new(CtxState {
info,
current: Cell::new(None),
beats: RefCell::new(FxHashMap::default()),
compensations: RefCell::new(Vec::new()),
}),
handlers,
outside,
}
}
pub fn workflow_id(&self) -> &str {
&self.s.info.id
}
pub fn workflow_type(&self) -> &str {
&self.s.info.wtype
}
pub fn version(&self) -> u32 {
self.s.info.version
}
pub fn run_id(&self) -> RunId {
self.run
}
pub fn started_ms(&self) -> u64 {
self.s.info.started_ms
}
pub fn traceparent(&self) -> Option<&str> {
self.s
.info
.trace
.as_deref()
.and_then(|t| std::str::from_utf8(t).ok())
}
pub fn key(&self) -> String {
let seq = self.s.current.get().unwrap_or(Seq::MAX);
self.host.key(self.run, seq)
}
pub fn compensate<F, Fut, E>(&self, name: &str, f: F)
where
F: Fn() -> Fut + 'static,
Fut: Future<Output = std::result::Result<(), E>> + 'static,
E: IntoFailure + 'static,
{
let f: Compensation = Rc::new(move || {
let fut = f();
Box::pin(async move { fut.await.map_err(IntoFailure::into_failure) })
});
self.s.compensations.borrow_mut().push((name.to_owned(), f));
}
pub(crate) async fn run_compensations(&self) {
loop {
let next = self.s.compensations.borrow_mut().pop();
let Some((name, f)) = next else { break };
let r = self
.step(&name, || {
let f = f.clone();
async move { f().await }
})
.await;
if let Err(e) = r {
self.log()
.error(format!("compensation '{name}' failed: {e}"));
}
}
}
pub(crate) fn has_compensations(&self) -> bool {
!self.s.compensations.borrow().is_empty()
}
pub fn log(&self) -> WorkflowLog<'_> {
WorkflowLog { ctx: self }
}
fn outcome_to<T: DeserializeOwned>(&self, name: &str, o: Outcome) -> Result<T> {
match o.status {
Status::Ok => decode(&o.payload),
Status::Failed => Err(Error::Step {
name: name.to_owned(),
failure: Failure::from_payload(&o.payload),
}),
Status::Cancelled => Err(Error::Cancelled(text(&o.payload))),
Status::TimedOut => Err(Error::TimedOut),
Status::Terminated | Status::ContinuedAsNew | Status::Reset => Err(Error::Terminated),
}
}
fn next_attempt(
&self,
policy: &RetryPolicy,
seq: Seq,
failure: &Failure,
attempt: u32,
first: &mut Option<u64>,
) -> Next {
if !policy.retryable(failure, attempt) {
return Next::Final;
}
match self.retry_at(policy, seq, attempt, first) {
Some(at) => Next::Retry(at),
None => Next::Exhausted,
}
}
fn pauses(&self, policy: &RetryPolicy) -> bool {
policy
.on_exhausted
.unwrap_or_else(|| self.host.on_exhausted())
== Exhausted::Pause
}
fn retry_at(
&self,
policy: &RetryPolicy,
seq: Seq,
attempt: u32,
first: &mut Option<u64>,
) -> Option<u64> {
if attempt >= policy.attempts {
return None;
}
let now = self.host.now_ms();
let first = *first.get_or_insert(now);
let mut delay = millis(policy.delay(attempt));
if policy.jitter && delay > 1 {
let half = delay / 2;
delay = half
+ (jitter_fraction(self.run, seq, attempt, now) * (delay - half) as f64) as u64;
}
let at = now.saturating_add(delay);
if policy
.max_elapsed
.is_some_and(|limit| at.saturating_sub(first) > millis(limit))
{
return None;
}
Some(at)
}
fn settle<T: Serialize + DeserializeOwned>(
&self,
seq: Seq,
opts: &StepOptions,
v: T,
kept: &Cell<Option<T>>,
) -> std::result::Result<Attempted<T>, Failure> {
thread_local! {
static BUF: RefCell<Vec<u8>> = const { RefCell::new(Vec::new()) };
}
let mut buf = BUF
.try_with(|b| {
b.try_borrow_mut()
.map(|mut b| std::mem::take(&mut *b))
.unwrap_or_default()
})
.unwrap_or_default();
buf.clear();
let r = match serde_json::to_writer(&mut buf, &v) {
Err(e) => Err(unencodable(e.into())),
Ok(()) if self.host.step_done_inline(self.run, seq, &buf) => Ok(Attempted::Settled(
if opts.own_value { Ok(v) } else { decode(&buf) },
)),
Ok(()) => {
if opts.own_value {
kept.set(Some(v));
}
Ok(Attempted::Payload(payload_of(&buf)))
}
};
if buf.capacity() <= 64 << 10 {
let _ = BUF.try_with(|b| {
if let Ok(mut b) = b.try_borrow_mut() {
*b = buf;
}
});
}
r
}
async fn attempt<T, E, Fut>(
&self,
seq: Seq,
opts: &StepOptions,
fut: Fut,
kept: &Cell<Option<T>>,
) -> std::result::Result<Attempted<T>, Failure>
where
T: Serialize + DeserializeOwned,
E: IntoFailure,
Fut: Future<Output = std::result::Result<T, E>>,
{
if opts.limit.is_none() && opts.timeout.is_none() && opts.heartbeat_timeout.is_none() {
let fut = std::pin::pin!(fut);
let keyed = Keyed {
fut,
current: &self.s.current,
outside: &self.outside,
seq,
pending: false,
};
return match keyed.await {
Ok(v) => self.settle(seq, opts, v, kept),
Err(e) => Err(e.into_failure()),
};
}
let _permit = match &opts.limit {
Some(limit) => Some(
AcquirePermit {
permit: Some(Permit {
host: self.host.clone(),
limit: limit.clone(),
ticket: self.host.limit_ticket(),
}),
priority: opts.priority.clone(),
run: self.run,
}
.await,
),
None => None,
};
let watched = opts.heartbeat_timeout.is_some();
if watched {
self.s.beats.borrow_mut().insert(seq, self.host.now_ms());
}
let fut = std::pin::pin!(fut);
let keyed = Keyed {
fut,
current: &self.s.current,
outside: &self.outside,
seq,
pending: false,
};
let deadline = opts.timeout.map(|t| after(self.host.now_ms(), t));
let watch = opts.heartbeat_timeout.map(|t| BeatWatch {
host: self.host.clone(),
state: self.s.clone(),
seq,
timeout_ms: millis(t),
armed: None,
});
let raced = Race3 {
fut: keyed,
until: deadline.map(|at| Until {
host: self.host.clone(),
at,
timer: None,
}),
watch,
}
.await;
if watched {
self.s.beats.borrow_mut().remove(&seq);
}
match raced {
Raced::Done(Ok(v)) => self.settle(seq, opts, v, kept),
Raced::Done(Err(e)) => Err(e.into_failure()),
Raced::TimedOut => Err(Failure::new(format!(
"step timed out after {} ms",
opts.timeout.map(|t| t.as_millis()).unwrap_or(0)
))
.kind("TimedOut")),
Raced::Silent => Err(Failure::new(format!(
"step sent no heartbeat for {} ms",
opts.heartbeat_timeout.map(|t| t.as_millis()).unwrap_or(0)
))
.kind("TimedOut")),
}
}
pub fn step<'a, T, E, F, Fut>(
&'a self,
name: &'a str,
f: F,
) -> Op<
StepOptions,
impl Make<StepOptions, Output = Result<T>, Fut = impl Future<Output = Result<T>> + 'a> + 'a,
>
where
T: Serialize + DeserializeOwned + 'a,
E: IntoFailure + 'a,
F: FnMut() -> Fut + 'a,
Fut: Future<Output = std::result::Result<T, E>> + 'a,
{
Op::new(StepOptions::default(), move |opts| {
self.run_step(name, opts, f)
})
}
async fn run_step<T, E, F, Fut>(&self, name: &str, opts: StepOptions, mut f: F) -> Result<T>
where
T: Serialize + DeserializeOwned,
E: IntoFailure,
F: FnMut() -> Fut,
Fut: Future<Output = std::result::Result<T, E>>,
{
let (seq, start, mut gate) = self.host.begin_step(self.run, name)?;
let kept: Cell<Option<T>> = Cell::new(None);
let gated = !opts.read_only;
if !gated {
gate = 0;
}
let mut delivery = Delivery {
host: self.host.clone(),
run: self.run,
seq,
done: false,
};
if let StepStart::Run {
attempt,
not_before_ms,
} = start
{
let mut attempt = attempt;
let mut not_before = not_before_ms;
let mut first = None;
loop {
Gate::new(self.host.clone(), self.run, seq, gate, not_before).await;
if not_before == PAUSED {
if !self.host.step_pending(self.run, seq) {
break;
}
attempt = 1;
first = None;
}
#[cfg(feature = "testing")]
let mocked = crate::testing::mocked(&self.s.info.wtype, name, attempt);
#[cfg(not(feature = "testing"))]
let mocked: Option<std::result::Result<Payload, Failure>> = None;
let result = match mocked {
Some(r) => r.map(Attempted::Payload),
None => self.attempt(seq, &opts, async { f().await }, &kept).await,
};
match result {
Ok(Attempted::Settled(r)) => {
delivery.done = true;
return r;
}
Ok(Attempted::Payload(payload)) => {
self.host.step_done(self.run, seq, Status::Ok, payload);
break;
}
Err(failure) => {
let next =
self.next_attempt(&opts.retry, seq, &failure, attempt, &mut first);
let at = match next {
Next::Retry(at) => Some(at),
Next::Exhausted if self.pauses(&opts.retry) => Some(PAUSED),
_ => None,
};
match at {
Some(at) => {
let retry_gate = self.host.step_retry(
self.run,
seq,
at,
failure.to_payload(),
)?;
gate = if gated { retry_gate } else { 0 };
not_before = at;
attempt += 1;
}
None => {
self.host.step_done(
self.run,
seq,
Status::Failed,
failure.to_payload(),
);
break;
}
}
}
}
}
}
let outcome = (&mut delivery).await;
if outcome.status == Status::Ok
&& let Some(v) = kept.take()
{
return Ok(v);
}
self.outcome_to(name, outcome)
}
pub fn await_completion<'a, T, E, F, Fut>(
&'a self,
name: &'a str,
send: F,
) -> Op<
StepOptions,
impl Make<StepOptions, Output = Result<T>, Fut = impl Future<Output = Result<T>> + 'a> + 'a,
>
where
T: DeserializeOwned + 'a,
E: IntoFailure + 'a,
F: FnMut(Completion) -> Fut + 'a,
Fut: Future<Output = std::result::Result<(), E>> + 'a,
{
Op::new(StepOptions::default(), move |opts| {
self.run_await_completion(name, opts, send)
})
}
async fn run_await_completion<T, E, F, Fut>(
&self,
name: &str,
opts: StepOptions,
mut send: F,
) -> Result<T>
where
T: DeserializeOwned,
E: IntoFailure,
F: FnMut(Completion) -> Fut,
Fut: Future<Output = std::result::Result<(), E>>,
{
let (seq, start, mut gate) = self.host.begin_await(self.run, name)?;
let mut delivery = Delivery {
host: self.host.clone(),
run: self.run,
seq,
done: false,
};
if let AwaitStart::Run {
attempt,
not_before_ms,
} = start
{
let mut attempt = attempt;
let mut not_before = not_before_ms;
let mut first = None;
loop {
Gate::new(self.host.clone(), self.run, seq, gate, not_before).await;
let token = self.host.completion_token(self.run, seq)?;
let completion = Completion::new(token);
let result = self
.attempt(
seq,
&opts,
async { send(completion).await.map(|()| true) },
&Cell::new(None),
)
.await;
match result {
Ok(_) => {
self.host.awaiting(self.run, seq)?;
break;
}
Err(failure) => {
let at = match self.next_attempt(
&opts.retry,
seq,
&failure,
attempt,
&mut first,
) {
Next::Retry(at) => Some(at),
Next::Exhausted | Next::Final => None,
};
match at {
Some(at) => {
gate = self.host.step_retry(
self.run,
seq,
at,
failure.to_payload(),
)?;
not_before = at;
attempt += 1;
}
None => {
self.host.step_done(
self.run,
seq,
Status::Failed,
failure.to_payload(),
);
break;
}
}
}
}
}
}
let outcome = (&mut delivery).await;
self.outcome_to(name, outcome)
}
pub fn heartbeat<T: Serialize>(&self, progress: &T) -> Result<()> {
let Some(seq) = self.s.current.get() else {
return Err(Error::Rejected(
"ctx.heartbeat() was called outside a step's function: call it inside ctx.step(name, || ...)".into(),
));
};
if let Some(last) = self.s.beats.borrow_mut().get_mut(&seq) {
*last = self.host.now_ms();
}
self.host.heartbeat(self.run, seq, encode(progress)?)
}
pub fn stream(&self, name: &str) -> Result<crate::stream::StreamWriter> {
crate::stream::StreamWriter::open(self, name)
}
pub(crate) fn host_ref(&self) -> &Rc<dyn Host> {
&self.host
}
pub(crate) fn running_step(&self) -> Option<Seq> {
self.s.current.get()
}
pub fn progress<T: DeserializeOwned>(&self) -> Result<Option<T>> {
let Some(seq) = self.s.current.get() else {
return Ok(None);
};
match self.host.progress(self.run, seq) {
Some(p) => decode(&p).map(Some),
None => Ok(None),
}
}
pub fn on_query<A, R, F>(&self, query: impl AsQuery<A, R>, f: F)
where
A: DeserializeOwned,
R: Serialize,
F: Fn(A) -> R + 'static,
{
let name = query.name();
let erased: QueryFn = Rc::new(move |arg: Payload| {
let a: A = decode(&arg).map_err(|e| e.to_string())?;
encode(&f(a)).map_err(|e| e.to_string())
});
self.handlers
.queries
.borrow_mut()
.insert(name.to_owned(), erased);
}
pub fn on_update<A, R, V, H, Fut>(&self, update: impl AsUpdate<A, R>, validate: V, handle: H)
where
A: DeserializeOwned + 'static,
R: Serialize + 'static,
V: Fn(&A) -> std::result::Result<(), String> + 'static,
H: Fn(A) -> Fut + 'static,
Fut: Future<Output = Result<R>> + 'static,
{
let name = update.name();
let v: ValidateFn = Rc::new(move |arg: &Payload| {
let a: A = decode(arg).map_err(|e| e.to_string())?;
validate(&a)
});
let handle = Rc::new(handle);
let h: UpdateFn = Rc::new(move |arg: Payload| {
let handle = handle.clone();
Box::pin(async move {
let a: A = decode(&arg).map_err(crate::engine::outcome_of)?;
let r = handle(a).await.map_err(crate::engine::outcome_of)?;
encode(&r).map_err(crate::engine::outcome_of)
}) as Pin<Box<dyn Future<Output = RunOutcome>>>
});
self.handlers
.updates
.borrow_mut()
.insert(name.to_owned(), (v, h));
}
pub(crate) fn on_update_raw(&self, name: &str, validate: ValidateFn, handle: UpdateFn) {
self.handlers
.updates
.borrow_mut()
.insert(name.to_owned(), (validate, handle));
}
pub(crate) fn on_query_raw(&self, name: &str, f: QueryFn) {
self.handlers
.queries
.borrow_mut()
.insert(name.to_owned(), f);
}
pub(crate) fn delivery(&self, seq: Seq) -> impl Future<Output = Outcome> + '_ {
Delivery {
host: self.host.clone(),
run: self.run,
seq,
done: false,
}
}
pub fn object(&self, wtype: &str, key: &str) -> crate::object::ObjectCall<'_> {
crate::object::ObjectCall {
ctx: self,
wtype: wtype.to_owned(),
key: key.to_owned(),
}
}
pub async fn continue_as_new<I: Serialize, T>(&self, input: &I) -> Result<T> {
self.host.continue_as_new(self.run, encode(input)?, None)?;
std::future::pending().await
}
pub fn journal_size(&self) -> u32 {
self.host.journal_size(self.run)
}
pub fn patched(&self, id: &str) -> Result<bool> {
self.host.patched(self.run, id, false)
}
pub fn deprecate_patch(&self, id: &str) -> Result<()> {
self.host.patched(self.run, id, true).map(|_| ())
}
pub fn set_attributes<T: Serialize>(&self, attributes: &T) -> Result<()> {
let v = serde_json::to_value(attributes)?;
if !v.is_object() {
return Err(Error::Codec("search attributes are a JSON object".into()));
}
self.host
.set_attributes(self.run, Arc::from(serde_json::to_vec(&v)?))
}
pub fn cancel_workflow(&self, target: &str, reason: &str) -> Result<()> {
self.host.cancel_workflow(self.run, target, reason)
}
pub async fn sleep(&self, name: &str, duration: Duration) -> Result<()> {
let at = after(self.host.now_ms(), duration);
self.sleep_until(name, at).await
}
pub async fn sleep_until(&self, name: &str, at_ms: u64) -> Result<()> {
let seq = self.host.begin_sleep(self.run, name, at_ms)?;
let o = Delivery {
host: self.host.clone(),
run: self.run,
seq,
done: false,
}
.await;
match o.status {
Status::Ok => Ok(()),
Status::Cancelled => Err(Error::Cancelled(text(&o.payload))),
_ => Err(Error::Failed(Failure::from_payload(&o.payload))),
}
}
pub fn wait_for<'a, T: DeserializeOwned + 'a>(
&'a self,
signal: impl AsSignal<T> + 'a,
) -> Op<
WaitOptions,
impl Make<
WaitOptions,
Output = Result<Option<T>>,
Fut = impl Future<Output = Result<Option<T>>> + 'a,
> + 'a,
> {
Op::new(WaitOptions::default(), move |o: WaitOptions| async move {
self.run_wait_for(signal.name(), o.timeout).await
})
}
async fn run_wait_for<T: DeserializeOwned>(
&self,
name: &str,
timeout: Option<Duration>,
) -> Result<Option<T>> {
let deadline = timeout.map(|t| after(self.host.now_ms(), t)).unwrap_or(0);
let seq = self.host.begin_wait(self.run, name, deadline)?;
let o = Delivery {
host: self.host.clone(),
run: self.run,
seq,
done: false,
}
.await;
match o.status {
Status::Ok => decode(&o.payload).map(Some),
Status::TimedOut => Ok(None),
Status::Cancelled => Err(Error::Cancelled(text(&o.payload))),
_ => Err(Error::Failed(Failure::from_payload(&o.payload))),
}
}
pub fn child<'a, I: Serialize + 'a, O: DeserializeOwned + 'a>(
&'a self,
workflow: impl AsWorkflow<I, O> + 'a,
id: &'a str,
input: &'a I,
) -> Op<
ChildOptions,
impl Make<ChildOptions, Output = Result<O>, Fut = impl Future<Output = Result<O>> + 'a> + 'a,
> {
Op::new(ChildOptions::default(), move |o| async move {
self.run_child(workflow.name(), id, input, o).await
})
}
pub(crate) async fn run_child<I: Serialize, O: DeserializeOwned>(
&self,
wtype: &str,
id: &str,
input: &I,
opts: ChildOptions,
) -> Result<O> {
let input = encode(input)?;
let deadline = opts
.timeout
.map(|t| after(self.host.now_ms(), t))
.unwrap_or(0);
let Some(seq) =
self.host
.begin_child(self.run, wtype, id, opts.close.bits(), deadline, input)?
else {
return Err(Error::failed("child did not register"));
};
let o = Delivery {
host: self.host.clone(),
run: self.run,
seq,
done: false,
}
.await;
match o.status {
Status::Ok => decode(&o.payload),
Status::Cancelled => Err(Error::Cancelled(text(&o.payload))),
Status::TimedOut => Err(Error::TimedOut),
Status::Terminated | Status::ContinuedAsNew | Status::Reset => Err(Error::Terminated),
Status::Failed => Err(Error::Child {
id: id.to_owned(),
failure: Failure::from_payload(&o.payload),
}),
}
}
pub fn start_child<'a, I: Serialize + 'a>(
&'a self,
workflow: impl crate::defs::NamesWorkflow<I> + 'a,
id: &'a str,
input: &'a I,
) -> Op<
ChildOptions,
impl Make<ChildOptions, Output = Result<()>, Fut = impl Future<Output = Result<()>> + 'a> + 'a,
> {
Op::new(ChildOptions::default(), move |o| async move {
self.run_start_child(workflow.name(), id, input, o)
})
}
pub(crate) fn run_start_child<I: Serialize>(
&self,
wtype: &str,
id: &str,
input: &I,
opts: ChildOptions,
) -> Result<()> {
let input = encode(input)?;
let deadline = opts
.timeout
.map(|t| after(self.host.now_ms(), t))
.unwrap_or(0);
self.host.begin_child(
self.run,
wtype,
id,
opts.close.bits() | policy::DETACHED,
deadline,
input,
)?;
Ok(())
}
pub fn signal_workflow<T: Serialize>(
&self,
id: &str,
signal: impl AsSignal<T>,
value: &T,
) -> Result<()> {
self.host.send(self.run, id, signal.name(), encode(value)?)
}
pub fn now_ms(&self) -> Result<u64> {
let v = self.host.value(self.run, value_kind::NOW, &mut |now, _| {
Arc::from(&now.to_le_bytes()[..])
})?;
let mut b = [0u8; 8];
if v.len() == 8 {
b.copy_from_slice(&v);
}
Ok(u64::from_le_bytes(b))
}
pub fn random(&self) -> Result<u64> {
let v = self.host.value(self.run, value_kind::RANDOM, &mut |_, r| {
Arc::from(&r.to_le_bytes()[..])
})?;
let mut b = [0u8; 8];
if v.len() == 8 {
b.copy_from_slice(&v);
}
Ok(u64::from_le_bytes(b))
}
pub fn uuid(&self) -> Result<String> {
let v = self.host.value(self.run, value_kind::UUID, &mut |now, r| {
let mut b = [0u8; 16];
b[..8].copy_from_slice(&r.to_le_bytes());
let mix = r.rotate_left(29) ^ now.wrapping_mul(0x9e37_79b9_7f4a_7c15);
b[8..].copy_from_slice(&mix.to_le_bytes());
b[6] = (b[6] & 0x0f) | 0x40;
b[8] = (b[8] & 0x3f) | 0x80;
Arc::from(&b[..])
})?;
let h: String = v.iter().map(|b| format!("{b:02x}")).collect();
if h.len() != 32 {
return Err(Error::Codec("recorded uuid has the wrong length".into()));
}
Ok(format!(
"{}-{}-{}-{}-{}",
&h[0..8],
&h[8..12],
&h[12..16],
&h[16..20],
&h[20..32]
))
}
pub async fn all<F: Future>(&self, futures: impl IntoIterator<Item = F>) -> Vec<F::Output> {
let items: Vec<_> = futures
.into_iter()
.map(|f| JoinItem::Pending(Box::pin(f)))
.collect();
JoinAll {
pending: items.len(),
branches: Branches::new(self.host.clone(), self.run, items.len()),
items,
}
.await
}
pub async fn race<F: Future>(
&self,
futures: impl IntoIterator<Item = F>,
) -> (usize, F::Output) {
let items: Vec<_> = futures.into_iter().map(|f| Box::pin(f)).collect();
Race {
branches: Branches::new(self.host.clone(), self.run, items.len()),
items,
}
.await
}
}
enum Raced<T> {
Done(T),
TimedOut,
Silent,
}
struct Race3<F> {
fut: F,
until: Option<Until>,
watch: Option<BeatWatch>,
}
impl<F: Future + Unpin> Future for Race3<F> {
type Output = Raced<F::Output>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
if let Poll::Ready(v) = Pin::new(&mut self.fut).poll(cx) {
return Poll::Ready(Raced::Done(v));
}
if let Some(u) = self.until.as_mut()
&& Pin::new(u).poll(cx).is_ready()
{
return Poll::Ready(Raced::TimedOut);
}
if let Some(w) = self.watch.as_mut()
&& Pin::new(w).poll(cx).is_ready()
{
return Poll::Ready(Raced::Silent);
}
Poll::Pending
}
}
fn lock<T>(m: &Mutex<T>) -> MutexGuard<'_, T> {
m.lock().unwrap_or_else(|e| e.into_inner())
}
#[derive(Default)]
struct Woken {
branches: Mutex<BTreeSet<usize>>,
parent: Mutex<Option<Waker>>,
}
struct BranchWaker {
woken: Arc<Woken>,
index: usize,
}
impl Wake for BranchWaker {
fn wake(self: Arc<Self>) {
self.wake_by_ref();
}
fn wake_by_ref(self: &Arc<Self>) {
lock(&self.woken.branches).insert(self.index);
let parent = lock(&self.woken.parent).clone();
if let Some(p) = parent {
p.wake();
}
}
}
struct Branches {
host: Rc<dyn Host>,
run: RunId,
woken: Arc<Woken>,
wakers: Vec<Waker>,
seen: Option<u64>,
}
impl Branches {
fn new(host: Rc<dyn Host>, run: RunId, n: usize) -> Self {
let woken = Arc::new(Woken::default());
let wakers = (0..n)
.map(|index| {
Waker::from(Arc::new(BranchWaker {
woken: woken.clone(),
index,
}))
})
.collect();
Self {
host,
run,
woken,
wakers,
seen: None,
}
}
fn begin(&mut self, cx: &Context<'_>) -> bool {
{
let mut parent = lock(&self.woken.parent);
if !parent.as_ref().is_some_and(|p| p.will_wake(cx.waker())) {
*parent = Some(cx.waker().clone());
}
}
let now = self.host.broadcasts(self.run);
let all = self.seen != Some(now);
self.seen = Some(now);
if all {
lock(&self.woken.branches).clear();
}
all
}
fn next(&self, all: bool, from: usize) -> Option<usize> {
if all {
return (from < self.wakers.len()).then_some(from);
}
let mut woken = lock(&self.woken.branches);
let i = woken.range(from..).next().copied()?;
woken.remove(&i);
Some(i)
}
}
enum JoinItem<F: Future> {
Pending(Pin<Box<F>>),
Done(Option<F::Output>),
}
struct JoinAll<F: Future> {
items: Vec<JoinItem<F>>,
branches: Branches,
pending: usize,
}
impl<F: Future> Unpin for JoinAll<F> {}
impl<F: Future> Future for JoinAll<F> {
type Output = Vec<F::Output>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = self.get_mut();
let all = this.branches.begin(cx);
let held = (this.pending > 1).then(|| this.branches.host.hold_inline());
let mut from = 0;
while let Some(i) = this.branches.next(all, from) {
from = i + 1;
let item = &mut this.items[i];
if let JoinItem::Pending(f) = item
&& let Poll::Ready(v) = f
.as_mut()
.poll(&mut Context::from_waker(&this.branches.wakers[i]))
{
*item = JoinItem::Done(Some(v));
this.pending -= 1;
}
}
if let Some(held) = held {
this.branches.host.resume_inline(held);
}
if this.pending > 0 {
return Poll::Pending;
}
Poll::Ready(
this.items
.iter_mut()
.filter_map(|i| match i {
JoinItem::Done(v) => v.take(),
JoinItem::Pending(_) => None,
})
.collect(),
)
}
}
struct Race<F> {
items: Vec<Pin<Box<F>>>,
branches: Branches,
}
impl<F: Future> Future for Race<F> {
type Output = (usize, F::Output);
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = self.get_mut();
let all = this.branches.begin(cx);
let mut from = 0;
while let Some(i) = this.branches.next(all, from) {
from = i + 1;
let mut branch = Context::from_waker(&this.branches.wakers[i]);
if let Poll::Ready(v) = this.items[i].as_mut().poll(&mut branch) {
return Poll::Ready((i, v));
}
}
Poll::Pending
}
}