use aion_core::{Payload, RunId, WorkflowId};
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use crate::StoreError;
pub const DEFAULT_OUTBOX_ROUTE: &str = aion_core::DEFAULT_TASK_QUEUE;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum OutboxStatus {
Pending,
Claimed,
Done,
Failed,
Cancelled,
}
impl OutboxStatus {
#[must_use]
pub fn as_str(self) -> &'static str {
match self {
Self::Pending => "pending",
Self::Claimed => "claimed",
Self::Done => "done",
Self::Failed => "failed",
Self::Cancelled => "cancelled",
}
}
pub fn parse_token(value: &str) -> Result<Self, StoreError> {
match value {
"pending" => Ok(Self::Pending),
"claimed" => Ok(Self::Claimed),
"done" => Ok(Self::Done),
"failed" => Ok(Self::Failed),
"cancelled" => Ok(Self::Cancelled),
other => Err(StoreError::Serialization(format!(
"unknown outbox status: {other}"
))),
}
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ClaimScope {
pub namespace: String,
pub task_queue: String,
pub node: Option<String>,
}
impl ClaimScope {
#[must_use]
pub fn new(namespace: impl Into<String>, task_queue: impl Into<String>) -> Self {
Self {
namespace: namespace.into(),
task_queue: task_queue.into(),
node: None,
}
}
#[must_use]
pub fn with_node(mut self, node: Option<String>) -> Self {
self.node = node;
self
}
#[must_use]
pub fn admits(&self, row: &OutboxRow) -> bool {
row.namespace == self.namespace
&& row.task_queue == self.task_queue
&& match (&self.node, &row.node) {
(_, None) => true,
(Some(scope_node), Some(row_node)) => scope_node == row_node,
(None, Some(_)) => false,
}
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct OutboxRow {
pub dispatch_key: String,
pub workflow_id: WorkflowId,
pub ordinal: u64,
pub run_id: Option<RunId>,
pub namespace: String,
pub task_queue: String,
pub node: Option<String>,
pub activity_type: String,
pub input: Payload,
pub status: OutboxStatus,
pub attempt: u32,
pub visible_after: DateTime<Utc>,
pub claimed_at: Option<DateTime<Utc>>,
}
impl OutboxRow {
#[must_use]
pub fn dispatch_key_for(workflow_id: &WorkflowId, ordinal: u64) -> String {
format!("{workflow_id}:{ordinal}")
}
#[must_use]
pub fn pending(
workflow_id: WorkflowId,
ordinal: u64,
activity_type: String,
input: Payload,
now: DateTime<Utc>,
) -> Self {
let dispatch_key = Self::dispatch_key_for(&workflow_id, ordinal);
Self {
dispatch_key,
workflow_id,
ordinal,
run_id: None,
namespace: String::from(DEFAULT_OUTBOX_ROUTE),
task_queue: String::from(DEFAULT_OUTBOX_ROUTE),
node: None,
activity_type,
input,
status: OutboxStatus::Pending,
attempt: 0,
visible_after: now,
claimed_at: None,
}
}
#[must_use]
pub fn with_run_id(mut self, run_id: Option<RunId>) -> Self {
self.run_id = run_id;
self
}
#[must_use]
pub fn with_namespace(mut self, namespace: impl Into<String>) -> Self {
self.namespace = namespace.into();
self
}
#[must_use]
pub fn with_task_queue(mut self, task_queue: impl Into<String>) -> Self {
self.task_queue = task_queue.into();
self
}
#[must_use]
pub fn with_node(mut self, node: Option<String>) -> Self {
self.node = node;
self
}
}
#[async_trait]
pub trait OutboxStore: Send + Sync + 'static {
async fn append_outbox_batch(&self, rows: &[OutboxRow]) -> Result<(), StoreError>;
async fn claim_outbox_rows(&self, limit: u32) -> Result<Vec<OutboxRow>, StoreError>;
async fn claim_outbox_rows_excluding(
&self,
limit: u32,
held: &std::collections::HashSet<WorkflowId>,
) -> Result<Vec<OutboxRow>, StoreError> {
let _ = held;
self.claim_outbox_rows(limit).await
}
async fn claim_outbox_rows_scoped(
&self,
scope: &ClaimScope,
limit: u32,
) -> Result<Vec<OutboxRow>, StoreError>;
async fn claim_outbox_rows_scoped_excluding(
&self,
scope: &ClaimScope,
limit: u32,
held: &std::collections::HashSet<WorkflowId>,
) -> Result<Vec<OutboxRow>, StoreError> {
let _ = held;
self.claim_outbox_rows_scoped(scope, limit).await
}
async fn list_stale_claimed_outbox_rows(
&self,
older_than: DateTime<Utc>,
limit: u32,
) -> Result<Vec<OutboxRow>, StoreError> {
let _ = (older_than, limit);
Err(StoreError::Backend(String::from(
"this outbox store does not support the stale-claim liveness probe; \
refusing to report an empty stale set (override OutboxStore::list_stale_claimed_outbox_rows)",
)))
}
async fn list_unsettled_outbox_workflow_ids(&self) -> Result<Vec<WorkflowId>, StoreError> {
Err(StoreError::Backend(String::from(
"this outbox store does not support unsettled-workflow enumeration; \
refusing to report an empty set (override OutboxStore::list_unsettled_outbox_workflow_ids)",
)))
}
async fn cancel_outbox_rows_for_workflow(
&self,
workflow_id: &WorkflowId,
) -> Result<Vec<String>, StoreError> {
let _ = workflow_id;
Err(StoreError::Backend(String::from(
"this outbox store does not support workflow-terminal settlement; \
refusing to no-op a settle (override OutboxStore::cancel_outbox_rows_for_workflow)",
)))
}
async fn rearm_stale_claimed_outbox_rows(
&self,
older_than: DateTime<Utc>,
visible_after: DateTime<Utc>,
limit: u32,
) -> Result<Vec<OutboxRow>, StoreError>;
async fn complete_outbox_row(&self, dispatch_key: &str) -> Result<(), StoreError>;
async fn retry_outbox_row(
&self,
dispatch_key: &str,
next_attempt: u32,
visible_after: DateTime<Utc>,
) -> Result<(), StoreError>;
async fn fail_outbox_row(&self, dispatch_key: &str) -> Result<(), StoreError>;
async fn count_inflight_outbox_rows(&self, namespace: &str) -> Result<u64, StoreError>;
async fn count_claimed_outbox_rows(&self, namespace: &str) -> Result<u64, StoreError>;
async fn count_claimed_outbox_rows_by_namespace(
&self,
namespaces: &[&str],
) -> Result<std::collections::BTreeMap<String, u64>, StoreError> {
let mut counts = std::collections::BTreeMap::new();
for namespace in namespaces {
let count = self.count_claimed_outbox_rows(namespace).await?;
counts.insert((*namespace).to_owned(), count);
}
Ok(counts)
}
async fn pending_outbox_routes(&self) -> Result<Vec<ClaimScope>, StoreError>;
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use aion_core::{ContentType, Payload, WorkflowId};
use chrono::Utc;
use super::{ClaimScope, OutboxRow, OutboxStatus, OutboxStore};
#[test]
fn outbox_store_is_object_safe() {
let _: Option<Arc<dyn OutboxStore>> = None;
}
fn row(namespace: &str, task_queue: &str, node: Option<&str>) -> OutboxRow {
OutboxRow::pending(
WorkflowId::new_v4(),
0,
String::from("charge"),
Payload::new(ContentType::Json, b"{}".to_vec()),
Utc::now(),
)
.with_namespace(namespace)
.with_task_queue(task_queue)
.with_node(node.map(ToOwned::to_owned))
}
#[test]
fn scope_admits_matching_namespace_task_queue_and_unpinned_or_matching_node() {
let scope = ClaimScope::new("remote", "gpu").with_node(Some("box-7".to_owned()));
assert!(scope.admits(&row("remote", "gpu", Some("box-7"))));
assert!(scope.admits(&row("remote", "gpu", None)));
}
#[test]
fn scope_rejects_other_namespace_task_queue_or_pinned_to_other_node() {
let scope = ClaimScope::new("remote", "gpu").with_node(Some("box-7".to_owned()));
assert!(!scope.admits(&row("default", "gpu", None)));
assert!(!scope.admits(&row("remote", "cpu", None)));
assert!(!scope.admits(&row("remote", "gpu", Some("box-9"))));
}
#[test]
fn node_less_scope_admits_only_unpinned_rows() {
let scope = ClaimScope::new("remote", "gpu");
assert!(scope.admits(&row("remote", "gpu", None)));
assert!(!scope.admits(&row("remote", "gpu", Some("box-7"))));
}
#[test]
fn status_tokens_round_trip() -> Result<(), crate::StoreError> {
for status in [
OutboxStatus::Pending,
OutboxStatus::Claimed,
OutboxStatus::Done,
OutboxStatus::Failed,
OutboxStatus::Cancelled,
] {
let parsed = OutboxStatus::parse_token(status.as_str())?;
assert_eq!(parsed, status);
}
Ok(())
}
#[test]
fn unknown_status_token_is_rejected() {
assert!(OutboxStatus::parse_token("nope").is_err());
}
#[test]
fn dispatch_key_is_workflow_id_colon_ordinal() {
let workflow_id = aion_core::WorkflowId::new_v4();
let key = OutboxRow::dispatch_key_for(&workflow_id, 7);
assert_eq!(key, format!("{workflow_id}:7"));
}
}