use std::collections::HashMap;
use std::collections::hash_map::Entry;
use std::future::Future;
use std::marker::PhantomData;
use std::pin::Pin;
use std::sync::{Arc, RwLock};
use serde::Serialize;
use serde::de::DeserializeOwned;
use crate::error::{DurableError, EncodedError, EngineOnly, Failure};
use crate::instance::ExecutorSlot;
use crate::serialization::{PayloadEncoding, SerializationFormat, decode_input, encode};
pub(crate) use crate::sysdb::types::QualifiedWorkflowName;
use crate::{DBOS, Error, Executor, Result};
pub(crate) type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
pub(crate) type ErasedWorkflow = Arc<
dyn Fn(
Option<String>,
PayloadEncoding,
) -> BoxFuture<'static, std::result::Result<Option<String>, Failure>>
+ Send
+ Sync,
>;
#[derive(Clone)]
pub(crate) struct RegisteredWorkflow {
pub(crate) body: ErasedWorkflow,
pub(crate) serialization: SerializationFormat,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct WorkflowOptions {
pub serialization: SerializationFormat,
}
#[derive(Default)]
pub(crate) struct Registry {
state: RwLock<State>,
}
#[derive(Default)]
struct State {
workflows: HashMap<QualifiedWorkflowName, RegisteredWorkflow>,
frozen: bool,
}
fn refuse_unregistrable(name: &QualifiedWorkflowName) -> Result<()> {
let detail = if name.name().is_empty() {
"a workflow needs a name"
} else if name.config_name().is_some() && name.class_name().is_none() {
"a configured instance needs a class"
} else {
return Ok(());
};
Err(Error::InvalidArgument {
operation: "register a workflow".into(),
detail: format!("{detail}, got `{name}`"),
})
}
pub(crate) type Snapshot = Arc<HashMap<QualifiedWorkflowName, RegisteredWorkflow>>;
impl Registry {
fn insert(&self, name: QualifiedWorkflowName, workflow: RegisteredWorkflow) -> Result<()> {
let mut state = self.write();
if state.frozen {
return Err(Error::AlreadyLaunched {
operation: "register_workflow".into(),
});
}
match state.workflows.entry(name) {
Entry::Occupied(taken) => Err(Error::AlreadyRegistered {
workflow: taken.key().clone(),
}),
Entry::Vacant(slot) => {
slot.insert(workflow);
Ok(())
}
}
}
pub(crate) fn snapshot(&self) -> Snapshot {
let mut state = self.write();
state.frozen = true;
Arc::new(state.workflows.clone())
}
pub(crate) fn thaw(&self) {
self.write().frozen = false;
}
fn write(&self) -> std::sync::RwLockWriteGuard<'_, State> {
self.state
.write()
.unwrap_or_else(std::sync::PoisonError::into_inner)
}
}
pub struct WorkflowRef<P, R, E = EngineOnly> {
executor: Arc<ExecutorSlot>,
name: Arc<QualifiedWorkflowName>,
serialization: SerializationFormat,
types: PhantomData<fn(P) -> (R, E)>,
}
impl<P, R, E> WorkflowRef<P, R, E> {
pub fn qualified_name(&self) -> &QualifiedWorkflowName {
&self.name
}
pub fn name(&self) -> &str {
self.name.name()
}
pub(crate) fn executor(&self, operation: &'static str) -> Result<Arc<Executor>> {
self.executor.get(operation)
}
pub(crate) fn serialization(&self) -> SerializationFormat {
self.serialization
}
}
impl<P, R, E> Clone for WorkflowRef<P, R, E> {
fn clone(&self) -> Self {
Self {
executor: Arc::clone(&self.executor),
name: Arc::clone(&self.name),
serialization: self.serialization,
types: PhantomData,
}
}
}
impl<P, R, E> std::fmt::Debug for WorkflowRef<P, R, E> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("WorkflowRef")
.field("name", &self.name.to_string())
.finish_non_exhaustive()
}
}
impl DBOS {
pub fn register_workflow<P, R, E, F, Fut>(
&self,
name: impl Into<QualifiedWorkflowName>,
workflow: F,
) -> Result<WorkflowRef<P, R, E>>
where
F: Fn(P) -> Fut + Send + Sync + 'static,
Fut: Future<Output = Result<R, E>> + Send + 'static,
P: Serialize + DeserializeOwned + Send + 'static,
R: Serialize + DeserializeOwned + Send + 'static,
E: DurableError,
{
self.register_workflow_with(name, WorkflowOptions::default(), workflow)
}
pub fn register_workflow_with<P, R, E, F, Fut>(
&self,
name: impl Into<QualifiedWorkflowName>,
options: WorkflowOptions,
workflow: F,
) -> Result<WorkflowRef<P, R, E>>
where
F: Fn(P) -> Fut + Send + Sync + 'static,
Fut: Future<Output = Result<R, E>> + Send + 'static,
P: Serialize + DeserializeOwned + Send + 'static,
R: Serialize + DeserializeOwned + Send + 'static,
E: DurableError,
{
let name = name.into();
refuse_unregistrable(&name)?;
let workflow = Arc::new(workflow);
let erased: ErasedWorkflow =
Arc::new(move |input: Option<String>, encoding: PayloadEncoding| {
let workflow = Arc::clone(&workflow);
Box::pin(async move {
let input = decode_input::<P, EngineOnly>(input.as_deref(), &encoding)
.map_err(Failure::Control)?;
match workflow(input).await {
Ok(output) => match encode::<R, E>(&output, &encoding, "result") {
Ok(encoded) => Ok(Some(encoded)),
Err(refused) => {
tracing::warn!(
error = %refused,
"the workflow's result could not be encoded; it is recorded \
as the workflow's error"
);
Err(Failure::Recorded(EncodedError::of(&refused, encoding)))
}
},
Err(error) => Err(match error.control() {
Some(control) => Failure::Control(control),
None => Failure::Recorded(EncodedError::of(&error, encoding)),
}),
}
})
});
self.registry().insert(
name.clone(),
RegisteredWorkflow {
body: erased,
serialization: options.serialization,
},
)?;
tracing::debug!(workflow = %name, "registered a workflow");
Ok(WorkflowRef {
executor: self.executor_slot(),
name: Arc::new(name),
serialization: options.serialization,
types: PhantomData,
})
}
}
#[cfg(test)]
mod tests {
use std::borrow::Cow;
use super::*;
use crate::Config;
use crate::sysdb::types::WorkflowRecord;
fn dbos() -> DBOS {
DBOS::new(Config::new("registry-test", "postgres://unused"))
}
async fn takes_nothing(_: ()) -> Result<String> {
Ok("nothing".to_owned())
}
async fn takes_one(n: u32) -> Result<u32> {
Ok(n * 2)
}
#[test]
fn registration_infers_the_argument_and_result_types() {
let dbos = dbos();
let nothing = dbos
.register_workflow("takes_nothing", takes_nothing)
.unwrap();
let one = dbos.register_workflow("takes_one", takes_one).unwrap();
assert_eq!(nothing.name(), "takes_nothing");
assert_eq!(one.name(), "takes_one");
dbos.register_workflow(
"closure",
|s: String| async move { Ok::<_, Error>(s.len()) },
)
.unwrap();
}
#[tokio::test]
async fn an_erased_workflow_round_trips_through_json() {
let dbos = dbos();
dbos.register_workflow("double", takes_one).unwrap();
let snapshot = dbos.registry().snapshot();
let erased = &snapshot
.get(&QualifiedWorkflowName::new("double"))
.expect("registered")
.body;
assert_eq!(
erased(Some("21".to_owned()), PayloadEncoding::Native)
.await
.unwrap(),
Some("42".to_owned())
);
}
#[tokio::test]
async fn a_zero_argument_workflow_is_called_with_no_input_at_all() {
let dbos = dbos();
dbos.register_workflow("nothing", takes_nothing).unwrap();
let snapshot = dbos.registry().snapshot();
let erased = &snapshot
.get(&QualifiedWorkflowName::new("nothing"))
.expect("registered")
.body;
assert_eq!(
erased(None, PayloadEncoding::Native).await.unwrap(),
Some("\"nothing\"".to_owned())
);
}
#[tokio::test]
async fn a_malformed_argument_is_reported_rather_than_panicking() {
let dbos = dbos();
dbos.register_workflow("double", takes_one).unwrap();
let snapshot = dbos.registry().snapshot();
let erased = &snapshot
.get(&QualifiedWorkflowName::new("double"))
.expect("registered")
.body;
let err = erased(Some("\"not a number\"".to_owned()), PayloadEncoding::Native)
.await
.unwrap_err();
assert!(
matches!(
err,
Failure::Control(Error::Deserialization {
what: Cow::Borrowed("argument"),
..
})
),
"{err:?}"
);
}
#[tokio::test]
async fn a_result_the_encoding_refuses_is_recorded_as_the_workflows_error() {
let dbos = dbos();
dbos.register_workflow("not_a_number", |()| async { Ok::<_, Error>(f64::NAN) })
.unwrap();
let snapshot = dbos.registry().snapshot();
let erased = &snapshot
.get(&QualifiedWorkflowName::new("not_a_number"))
.expect("registered")
.body;
let err = erased(None, PayloadEncoding::Portable).await.unwrap_err();
let Failure::Recorded(recorded) = err else {
panic!("expected a recorded failure, got {err:?}")
};
let read = recorded.decode::<EngineOnly>("wf-1");
assert!(
matches!(&read, Error::Serialization { what, .. } if what == "result"),
"{read:?}"
);
}
#[test]
fn a_snapshot_closes_the_registry_and_releasing_it_reopens() {
let dbos = dbos();
dbos.register_workflow("before", takes_nothing).unwrap();
let snapshot = dbos.registry().snapshot();
assert_eq!(snapshot.len(), 1);
let err = dbos.register_workflow("after", takes_nothing).unwrap_err();
assert!(
matches!(
err,
Error::AlreadyLaunched {
operation: Cow::Borrowed("register_workflow")
}
),
"{err}"
);
assert_eq!(
dbos.registry().snapshot().len(),
1,
"the refused registration is not in the map either"
);
dbos.registry().thaw();
dbos.register_workflow("after", takes_nothing)
.expect("releasing the snapshot reopens registration");
assert_eq!(dbos.registry().snapshot().len(), 2);
}
#[test]
fn a_key_naming_nothing_is_refused() {
let dbos = dbos();
for (key, part) in [
(QualifiedWorkflowName::new(""), "needs a name"),
(
QualifiedWorkflowName::instance("unclassed", "", "eu"),
"needs a class",
),
] {
match dbos.register_workflow(key, takes_one) {
Err(Error::InvalidArgument { detail, .. }) => {
assert!(detail.contains(part), "{detail}");
}
Err(other) => panic!("expected InvalidArgument, got {other}"),
Ok(_) => panic!("registered a key naming nothing"),
}
}
assert!(dbos.registry().snapshot().is_empty());
}
#[test]
fn one_identity_registers_once() {
let dbos = dbos();
let key = QualifiedWorkflowName::instance("same", "Checkout", "eu");
dbos.register_workflow(&key, takes_nothing).unwrap();
let err = dbos.register_workflow(&key, takes_nothing).unwrap_err();
match err {
Error::AlreadyRegistered { workflow: taken } => assert_eq!(taken, key),
other => panic!("expected AlreadyRegistered, got {other}"),
}
}
#[test]
fn a_name_and_an_instance_of_it_are_different_identities() {
let dbos = dbos();
dbos.register_workflow("shared", takes_nothing).unwrap();
dbos.register_workflow(
QualifiedWorkflowName::instance("shared", "Checkout", "eu"),
takes_nothing,
)
.unwrap();
assert_eq!(dbos.registry().snapshot().len(), 2);
}
#[test]
fn a_row_written_with_empty_strings_resolves_to_the_same_key_as_one_written_with_nulls() {
let free = QualifiedWorkflowName::new("checkout");
assert_eq!(
QualifiedWorkflowName::from_row("checkout", None, None),
free
);
assert_eq!(
QualifiedWorkflowName::from_row("checkout", Some(""), Some("")),
free
);
}
#[test]
fn a_name_converts_to_a_free_functions_key() {
let free = QualifiedWorkflowName::new("checkout");
let name = "checkout".to_owned();
assert_eq!(QualifiedWorkflowName::from("checkout"), free);
assert_eq!(QualifiedWorkflowName::from(&name), free);
assert_eq!(QualifiedWorkflowName::from(name), free);
let instance = QualifiedWorkflowName::instance("checkout", "Checkout", "eu");
assert_eq!(QualifiedWorkflowName::from(&instance), instance);
}
#[test]
fn a_record_yields_its_key_and_a_nameless_record_none() {
let record = |row: serde_json::Value| -> WorkflowRecord {
let mut fields = serde_json::json!({ "workflow_id": "w", "status": "PENDING" });
fields
.as_object_mut()
.unwrap()
.extend(row.as_object().unwrap().clone());
serde_json::from_value(fields).unwrap()
};
let instance = record(serde_json::json!({
"name": "checkout", "class_name": "Checkout", "config_name": "eu",
}));
assert_eq!(
instance.qualified_name(),
Some(QualifiedWorkflowName::instance(
"checkout", "Checkout", "eu"
))
);
let blank = record(serde_json::json!({
"name": "checkout", "class_name": "", "config_name": "",
}));
assert_eq!(
blank.qualified_name(),
Some(QualifiedWorkflowName::new("checkout"))
);
assert_eq!(record(serde_json::json!({})).qualified_name(), None);
assert_eq!(
record(serde_json::json!({ "name": "" })).qualified_name(),
None
);
}
#[test]
fn a_key_displays_as_its_slash_separated_triple() {
assert_eq!(
QualifiedWorkflowName::new("checkout").to_string(),
"checkout"
);
assert_eq!(
QualifiedWorkflowName::instance("checkout", "Checkout", "eu").to_string(),
"checkout/Checkout/eu"
);
}
}