use super::{
Arc, BTreeMap, Budget, Checkpoint, CheckpointError, CheckpointErrorKind, CheckpointId,
ClaimedWorkflow, Deserialize, Duration, LeaseDuration, Mutex, MutexGuard, Reverse, Serialize,
SystemWorkflowClock, Usage, Value, WorkerId, WorkflowBudgetAuditCursor,
WorkflowBudgetAuditEvent, WorkflowBudgetAuditKind, WorkflowBudgetAuditLimit,
WorkflowBudgetAuditProjectionId, WorkflowBudgetAuditProjectionLease,
WorkflowBudgetForfeitReason, WorkflowBudgetReservationOutcome, WorkflowCancelOutcome,
WorkflowCheckpointHistoryLimit, WorkflowCheckpointPhase, WorkflowCheckpointRevision,
WorkflowClock, WorkflowDisposition, WorkflowForkCommand, WorkflowForkOutcome,
WorkflowInterruptRequest, WorkflowLease, WorkflowLineage, WorkflowSignal, WorkflowSignalId,
WorkflowSignalOutcome, WorkflowSignalRetention, WorkflowSignalSnapshot, WorkflowSignalState,
WorkflowStore, WorkflowStoreError, WorkflowStoreErrorKind, WorkflowStoreFuture, WorkflowTask,
WorkflowTaskSnapshot, WorkflowTaskStatus, WorkflowTenantBudgetPolicy,
WorkflowTenantBudgetSnapshot, WorkflowTenantId, WorkflowTenantListLimit, WorkflowTenantPolicy,
WorkflowWait, WorkflowWake, decode_revision, fork_checkpoint,
};
mod budget;
mod checkpoint;
mod signal;
mod task;
use signal::take_buffered_signal;
use task::{is_non_terminal, require_current_lease, require_tenant, workflow_not_found};
#[derive(Clone)]
pub struct InMemoryWorkflowStore {
tasks: Arc<Mutex<BTreeMap<CheckpointId, StoredTask>>>,
checkpoints: Arc<Mutex<StoredCheckpoints>>,
signals: Arc<Mutex<BTreeMap<WorkflowSignalId, StoredSignal>>>,
admission: Arc<Mutex<AdmissionState>>,
clock: Arc<dyn WorkflowClock>,
}
impl Default for InMemoryWorkflowStore {
fn default() -> Self {
Self::with_clock(Arc::new(SystemWorkflowClock))
}
}
impl InMemoryWorkflowStore {
pub fn new() -> Self {
Self::default()
}
pub fn with_clock(clock: Arc<dyn WorkflowClock>) -> Self {
Self {
tasks: Arc::new(Mutex::new(BTreeMap::new())),
checkpoints: Arc::new(Mutex::new(StoredCheckpoints::default())),
signals: Arc::new(Mutex::new(BTreeMap::new())),
admission: Arc::new(Mutex::new(AdmissionState::default())),
clock,
}
}
fn tasks(&self) -> MutexGuard<'_, BTreeMap<CheckpointId, StoredTask>> {
self.tasks
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
}
}
impl std::fmt::Debug for InMemoryWorkflowStore {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("InMemoryWorkflowStore")
.finish_non_exhaustive()
}
}
#[derive(Clone, Debug, Default, Deserialize, Serialize)]
struct AdmissionState {
tenants: BTreeMap<WorkflowTenantId, StoredTenant>,
budgets: BTreeMap<WorkflowTenantId, StoredTenantBudget>,
budget_audit_projections:
BTreeMap<(WorkflowTenantId, WorkflowBudgetAuditProjectionId), StoredBudgetAuditProjection>,
next_claim_sequence: u64,
}
#[derive(Clone, Copy, Debug, Deserialize, Serialize)]
struct StoredTenant {
policy: WorkflowTenantPolicy,
last_claim_sequence: u64,
}
#[derive(Clone, Debug, Deserialize, Serialize)]
struct StoredTenantBudget {
policy: WorkflowTenantBudgetPolicy,
window_started_at_ms: u64,
committed: Usage,
reserved: Usage,
reservations: BTreeMap<CheckpointId, StoredBudgetReservation>,
next_audit_sequence: u64,
audit_events: Vec<StoredBudgetAuditEvent>,
}
#[derive(Clone, Copy, Debug, Deserialize, Serialize)]
struct StoredBudgetReservation {
baseline: Usage,
amount: Usage,
reserved_at_ms: u64,
expires_at_ms: u64,
}
#[derive(Clone, Copy, Debug, Deserialize, Serialize)]
struct StoredBudgetAuditEvent {
cursor: WorkflowBudgetAuditCursor,
checkpoint_id: Option<CheckpointId>,
occurred_at_ms: u64,
kind: WorkflowBudgetAuditKind,
usage: Usage,
reservation_age_ms: Option<u64>,
limit: Budget,
committed: Usage,
reserved: Usage,
}
#[derive(Clone, Debug, Default, Deserialize, Serialize)]
struct StoredBudgetAuditProjection {
cursor: WorkflowBudgetAuditCursor,
owner: Option<WorkerId>,
fencing_token: u64,
expires_at_ms: Option<u64>,
}
#[derive(Clone, Debug, Deserialize, Serialize)]
struct StoredTask {
task: WorkflowTask,
state: StoredState,
attempts: u64,
fencing_token: u64,
wake: Option<WorkflowWake>,
lineage: Option<WorkflowLineage>,
created_at_ms: u64,
updated_at_ms: u64,
}
#[derive(Clone, Debug, Default, Deserialize, Serialize)]
struct StoredCheckpoints {
latest: BTreeMap<CheckpointId, Checkpoint>,
history: BTreeMap<(CheckpointId, u64), Checkpoint>,
}
#[derive(Clone, Debug, Deserialize, Serialize)]
enum StoredState {
Queued {
available_at_ms: u64,
},
Leased(WorkflowLease),
WaitingTimer {
wake_at_ms: u64,
},
WaitingSignal {
name: crate::WorkflowSignalName,
},
WaitingSignalOrTimeout {
name: crate::WorkflowSignalName,
wake_at_ms: u64,
},
WaitingInterrupt {
request: WorkflowInterruptRequest,
},
Completed,
Failed(String),
Cancelled,
}
#[derive(Clone, Debug, Deserialize, Serialize)]
struct StoredSignal {
tenant_id: WorkflowTenantId,
signal: WorkflowSignal,
consumed: bool,
dead_lettered: bool,
#[serde(default)]
compaction_protected: bool,
accepted_at_ms: u64,
}
const PERSISTENT_SNAPSHOT_VERSION: u32 = 1;
#[derive(Deserialize, Serialize)]
struct PersistentWorkflowSnapshot {
version: u32,
tasks: Vec<(CheckpointId, StoredTask)>,
checkpoints: Vec<(CheckpointId, Checkpoint)>,
checkpoint_history: Vec<(CheckpointId, u64, Checkpoint)>,
signals: Vec<(WorkflowSignalId, StoredSignal)>,
tenants: Vec<(WorkflowTenantId, StoredTenant)>,
budgets: Vec<(WorkflowTenantId, StoredTenantBudget)>,
budget_audit_projections: Vec<(
WorkflowTenantId,
WorkflowBudgetAuditProjectionId,
StoredBudgetAuditProjection,
)>,
next_claim_sequence: u64,
}
impl InMemoryWorkflowStore {
#[doc(hidden)]
pub fn export_persistent_snapshot(&self) -> Result<Vec<u8>, WorkflowStoreError> {
let tasks = self.tasks();
let checkpoints = self
.checkpoints
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let signals = self
.signals
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let admission = self
.admission
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let snapshot = PersistentWorkflowSnapshot {
version: PERSISTENT_SNAPSHOT_VERSION,
tasks: tasks.iter().map(|(id, task)| (*id, task.clone())).collect(),
checkpoints: checkpoints
.latest
.iter()
.map(|(id, checkpoint)| (*id, checkpoint.clone()))
.collect(),
checkpoint_history: checkpoints
.history
.iter()
.map(|((id, revision), checkpoint)| (*id, *revision, checkpoint.clone()))
.collect(),
signals: signals
.iter()
.map(|(id, signal)| (*id, signal.clone()))
.collect(),
tenants: admission
.tenants
.iter()
.map(|(id, tenant)| (id.clone(), *tenant))
.collect(),
budgets: admission
.budgets
.iter()
.map(|(id, budget)| (id.clone(), budget.clone()))
.collect(),
budget_audit_projections: admission
.budget_audit_projections
.iter()
.map(|((tenant_id, projection_id), projection)| {
(tenant_id.clone(), projection_id.clone(), projection.clone())
})
.collect(),
next_claim_sequence: admission.next_claim_sequence,
};
serde_json::to_vec(&snapshot).map_err(|error| {
WorkflowStoreError::new(
WorkflowStoreErrorKind::Storage,
format!("workflow snapshot encoding failed: {error}"),
)
})
}
#[doc(hidden)]
pub fn from_persistent_snapshot(
encoded: &[u8],
clock: Arc<dyn WorkflowClock>,
) -> Result<Self, WorkflowStoreError> {
let snapshot: PersistentWorkflowSnapshot =
serde_json::from_slice(encoded).map_err(|error| {
WorkflowStoreError::new(
WorkflowStoreErrorKind::Storage,
format!("workflow snapshot decoding failed: {error}"),
)
})?;
if snapshot.version != PERSISTENT_SNAPSHOT_VERSION {
return Err(WorkflowStoreError::new(
WorkflowStoreErrorKind::Storage,
format!("unsupported workflow snapshot version {}", snapshot.version),
));
}
let checkpoints = StoredCheckpoints {
latest: snapshot.checkpoints.into_iter().collect(),
history: snapshot
.checkpoint_history
.into_iter()
.map(|(id, revision, checkpoint)| ((id, revision), checkpoint))
.collect(),
};
let admission = AdmissionState {
tenants: snapshot.tenants.into_iter().collect(),
budgets: snapshot.budgets.into_iter().collect(),
budget_audit_projections: snapshot
.budget_audit_projections
.into_iter()
.map(|(tenant_id, projection_id, projection)| {
((tenant_id, projection_id), projection)
})
.collect(),
next_claim_sequence: snapshot.next_claim_sequence,
};
Ok(Self {
tasks: Arc::new(Mutex::new(snapshot.tasks.into_iter().collect())),
checkpoints: Arc::new(Mutex::new(checkpoints)),
signals: Arc::new(Mutex::new(snapshot.signals.into_iter().collect())),
admission: Arc::new(Mutex::new(admission)),
clock,
})
}
}
impl InMemoryWorkflowStore {
fn suspend(
&self,
stored: &mut StoredTask,
checkpoint_id: CheckpointId,
wait: WorkflowWait,
now: u64,
) -> StoredState {
match wait {
WorkflowWait::Timer { delay_ms } => {
stored.wake = None;
StoredState::WaitingTimer {
wake_at_ms: now.saturating_add(delay_ms),
}
}
WorkflowWait::Signal { name } => {
self.suspend_signal(stored, checkpoint_id, name, None, now)
}
WorkflowWait::SignalOrTimeout { name, timeout_ms } => {
self.suspend_signal(stored, checkpoint_id, name, Some(timeout_ms), now)
}
WorkflowWait::Interrupt { request } => {
let name = request.signal_name();
let mut signals = self
.signals
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if let Some(wake) = take_buffered_signal(&mut signals, checkpoint_id, &name) {
stored.wake = Some(wake);
StoredState::Queued {
available_at_ms: now,
}
} else {
stored.wake = None;
StoredState::WaitingInterrupt { request }
}
}
}
}
fn suspend_signal(
&self,
stored: &mut StoredTask,
checkpoint_id: CheckpointId,
name: crate::WorkflowSignalName,
timeout_ms: Option<u64>,
now: u64,
) -> StoredState {
let mut signals = self
.signals
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if let Some(wake) = take_buffered_signal(&mut signals, checkpoint_id, &name) {
stored.wake = Some(wake);
return StoredState::Queued {
available_at_ms: now,
};
}
stored.wake = None;
match timeout_ms {
Some(timeout_ms) => StoredState::WaitingSignalOrTimeout {
name,
wake_at_ms: now.saturating_add(timeout_ms),
},
None => StoredState::WaitingSignal { name },
}
}
}
impl WorkflowStore for InMemoryWorkflowStore {
fn current_time_ms(&self) -> WorkflowStoreFuture<'_, Result<u64, WorkflowStoreError>> {
Box::pin(async move { Ok(self.clock.now_ms()) })
}
fn set_tenant_policy(
&self,
tenant_id: WorkflowTenantId,
policy: WorkflowTenantPolicy,
) -> WorkflowStoreFuture<'_, Result<(), WorkflowStoreError>> {
self.set_tenant_policy_impl(tenant_id, policy)
}
fn set_tenant_budget_policy(
&self,
tenant_id: WorkflowTenantId,
policy: WorkflowTenantBudgetPolicy,
) -> WorkflowStoreFuture<'_, Result<(), WorkflowStoreError>> {
self.set_tenant_budget_policy_impl(tenant_id, policy)
}
fn list_tenant_budgets(
&self,
after: Option<WorkflowTenantId>,
limit: WorkflowTenantListLimit,
) -> WorkflowStoreFuture<'_, Result<Vec<WorkflowTenantId>, WorkflowStoreError>> {
self.list_tenant_budgets_impl(after, limit)
}
fn inspect_tenant_budget(
&self,
tenant_id: WorkflowTenantId,
) -> WorkflowStoreFuture<'_, Result<WorkflowTenantBudgetSnapshot, WorkflowStoreError>> {
self.inspect_tenant_budget_impl(tenant_id)
}
fn list_tenant_budget_audit(
&self,
tenant_id: WorkflowTenantId,
after: Option<WorkflowBudgetAuditCursor>,
limit: WorkflowBudgetAuditLimit,
) -> WorkflowStoreFuture<'_, Result<Vec<WorkflowBudgetAuditEvent>, WorkflowStoreError>> {
self.list_tenant_budget_audit_impl(tenant_id, after, limit)
}
fn compact_tenant_budget_audit(
&self,
tenant_id: WorkflowTenantId,
through: WorkflowBudgetAuditCursor,
) -> WorkflowStoreFuture<'_, Result<u64, WorkflowStoreError>> {
self.compact_tenant_budget_audit_impl(tenant_id, through)
}
fn load_or_create_tenant_budget_audit_projection(
&self,
tenant_id: WorkflowTenantId,
projection_id: WorkflowBudgetAuditProjectionId,
) -> WorkflowStoreFuture<'_, Result<WorkflowBudgetAuditCursor, WorkflowStoreError>> {
self.load_or_create_tenant_budget_audit_projection_impl(tenant_id, projection_id)
}
fn advance_tenant_budget_audit_projection(
&self,
tenant_id: WorkflowTenantId,
projection_id: WorkflowBudgetAuditProjectionId,
expected: WorkflowBudgetAuditCursor,
next: WorkflowBudgetAuditCursor,
) -> WorkflowStoreFuture<'_, Result<bool, WorkflowStoreError>> {
self.advance_tenant_budget_audit_projection_impl(tenant_id, projection_id, expected, next)
}
fn claim_tenant_budget_audit_projection(
&self,
tenant_id: WorkflowTenantId,
projection_id: WorkflowBudgetAuditProjectionId,
owner: WorkerId,
lease: LeaseDuration,
) -> WorkflowStoreFuture<
'_,
Result<Option<WorkflowBudgetAuditProjectionLease>, WorkflowStoreError>,
> {
self.claim_tenant_budget_audit_projection_impl(tenant_id, projection_id, owner, lease)
}
fn heartbeat_tenant_budget_audit_projection(
&self,
lease: WorkflowBudgetAuditProjectionLease,
extension: LeaseDuration,
) -> WorkflowStoreFuture<'_, Result<WorkflowBudgetAuditProjectionLease, WorkflowStoreError>>
{
self.heartbeat_tenant_budget_audit_projection_impl(lease, extension)
}
fn advance_tenant_budget_audit_projection_lease(
&self,
lease: WorkflowBudgetAuditProjectionLease,
next: WorkflowBudgetAuditCursor,
) -> WorkflowStoreFuture<'_, Result<WorkflowBudgetAuditProjectionLease, WorkflowStoreError>>
{
self.advance_tenant_budget_audit_projection_lease_impl(lease, next)
}
fn release_tenant_budget_audit_projection(
&self,
lease: WorkflowBudgetAuditProjectionLease,
) -> WorkflowStoreFuture<'_, Result<(), WorkflowStoreError>> {
self.release_tenant_budget_audit_projection_impl(lease)
}
fn reserve_budget(
&self,
lease: WorkflowLease,
workflow_limit: Budget,
baseline: Usage,
) -> WorkflowStoreFuture<'_, Result<WorkflowBudgetReservationOutcome, WorkflowStoreError>> {
self.reserve_budget_impl(lease, workflow_limit, baseline)
}
fn settle_budget(
&self,
lease: WorkflowLease,
cumulative: Usage,
) -> WorkflowStoreFuture<'_, Result<(), WorkflowStoreError>> {
self.settle_budget_impl(lease, cumulative)
}
fn enqueue(
&self,
task: WorkflowTask,
) -> WorkflowStoreFuture<'_, Result<(), WorkflowStoreError>> {
self.enqueue_impl(task)
}
fn claim(
&self,
worker: WorkerId,
lease: LeaseDuration,
) -> WorkflowStoreFuture<'_, Result<Option<ClaimedWorkflow>, WorkflowStoreError>> {
self.claim_impl(worker, lease)
}
fn heartbeat(
&self,
lease: WorkflowLease,
extension: LeaseDuration,
) -> WorkflowStoreFuture<'_, Result<WorkflowLease, WorkflowStoreError>> {
self.heartbeat_impl(lease, extension)
}
fn finish(
&self,
lease: WorkflowLease,
disposition: WorkflowDisposition,
) -> WorkflowStoreFuture<'_, Result<(), WorkflowStoreError>> {
self.finish_impl(lease, disposition)
}
fn publish_signal(
&self,
tenant_id: WorkflowTenantId,
signal: WorkflowSignal,
) -> WorkflowStoreFuture<'_, Result<WorkflowSignalOutcome, WorkflowStoreError>> {
self.publish_signal_impl(tenant_id, signal, false)
}
fn publish_control_signal(
&self,
tenant_id: WorkflowTenantId,
signal: WorkflowSignal,
) -> WorkflowStoreFuture<'_, Result<WorkflowSignalOutcome, WorkflowStoreError>> {
self.publish_signal_impl(tenant_id, signal, true)
}
fn cancel(
&self,
tenant_id: WorkflowTenantId,
checkpoint_id: CheckpointId,
) -> WorkflowStoreFuture<'_, Result<WorkflowCancelOutcome, WorkflowStoreError>> {
self.cancel_impl(tenant_id, checkpoint_id)
}
fn inspect_signal(
&self,
tenant_id: WorkflowTenantId,
signal_id: WorkflowSignalId,
) -> WorkflowStoreFuture<'_, Result<WorkflowSignalSnapshot, WorkflowStoreError>> {
self.inspect_signal_impl(tenant_id, signal_id)
}
fn load_signal_payload(
&self,
tenant_id: WorkflowTenantId,
signal_id: WorkflowSignalId,
) -> WorkflowStoreFuture<'_, Result<Value, WorkflowStoreError>> {
Box::pin(async move {
let signals = self
.signals
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let stored = signals.get(&signal_id).ok_or_else(|| {
WorkflowStoreError::new(
WorkflowStoreErrorKind::NotFound,
"workflow signal does not exist",
)
})?;
if stored.tenant_id != tenant_id {
return Err(WorkflowStoreError::new(
WorkflowStoreErrorKind::TenantMismatch,
"workflow tenant does not own signal",
));
}
Ok(stored.signal.payload.clone())
})
}
fn compact_signals(
&self,
tenant_id: WorkflowTenantId,
retention: WorkflowSignalRetention,
) -> WorkflowStoreFuture<'_, Result<u64, WorkflowStoreError>> {
self.compact_signals_impl(tenant_id, retention)
}
fn inspect(
&self,
tenant_id: WorkflowTenantId,
checkpoint_id: CheckpointId,
) -> WorkflowStoreFuture<'_, Result<WorkflowTaskSnapshot, WorkflowStoreError>> {
self.inspect_impl(tenant_id, checkpoint_id)
}
fn load_task_input(
&self,
tenant_id: WorkflowTenantId,
checkpoint_id: CheckpointId,
) -> WorkflowStoreFuture<'_, Result<Value, WorkflowStoreError>> {
Box::pin(async move {
let tasks = self.tasks();
let stored = tasks.get(&checkpoint_id).ok_or_else(|| {
WorkflowStoreError::new(WorkflowStoreErrorKind::NotFound, "workflow not found")
})?;
if stored.task.tenant_id != tenant_id {
return Err(WorkflowStoreError::new(
WorkflowStoreErrorKind::TenantMismatch,
"workflow tenant does not own task",
));
}
Ok(stored.task.input.clone())
})
}
fn list_checkpoint_history(
&self,
tenant_id: WorkflowTenantId,
checkpoint_id: CheckpointId,
after_revision: Option<u64>,
limit: WorkflowCheckpointHistoryLimit,
) -> WorkflowStoreFuture<'_, Result<Vec<WorkflowCheckpointRevision>, WorkflowStoreError>> {
self.list_checkpoint_history_impl(tenant_id, checkpoint_id, after_revision, limit)
}
fn load_checkpoint_revision(
&self,
tenant_id: WorkflowTenantId,
checkpoint_id: CheckpointId,
revision: u64,
) -> WorkflowStoreFuture<'_, Result<WorkflowCheckpointRevision, WorkflowStoreError>> {
self.load_checkpoint_revision_impl(tenant_id, checkpoint_id, revision)
}
fn fork_workflow(
&self,
tenant_id: WorkflowTenantId,
command: WorkflowForkCommand,
) -> WorkflowStoreFuture<'_, Result<WorkflowForkOutcome, WorkflowStoreError>> {
self.fork_workflow_impl(tenant_id, command)
}
fn load_checkpoint(
&self,
lease: WorkflowLease,
) -> WorkflowStoreFuture<'_, Result<Checkpoint, CheckpointError>> {
self.load_checkpoint_impl(lease)
}
fn compare_and_swap_checkpoint(
&self,
lease: WorkflowLease,
checkpoint: Checkpoint,
expected_revision: Option<u64>,
) -> WorkflowStoreFuture<'_, Result<(), CheckpointError>> {
self.compare_and_swap_checkpoint_impl(lease, checkpoint, expected_revision)
}
}