use std::future::Future;
use std::time::Duration;
use serde::Deserialize;
use serde::Serialize;
use serde::de::DeserializeOwned;
use taquba::{Queue, WaitOutcome};
use crate::keys::hash_input;
use crate::memo::Memo;
use crate::{Result, Step, StepError, StepErrorKind, StepOutcome};
pub(crate) const OUTCOME_KEY: &str = "workflow.outcome";
#[derive(Debug, Clone, Serialize, Deserialize)]
pub(crate) struct OutcomeRecord {
pub(crate) input_hash: [u8; 32],
pub(crate) outcome: StoredOutcome,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub(crate) enum StoredOutcome {
Success {
output: Vec<u8>,
},
Failure {
kind: StoredErrorKind,
message: String,
},
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub(crate) enum StoredErrorKind {
Transient,
Permanent,
}
impl From<StepErrorKind> for StoredErrorKind {
fn from(kind: StepErrorKind) -> Self {
match kind {
StepErrorKind::Transient => Self::Transient,
StepErrorKind::Permanent => Self::Permanent,
}
}
}
pub(crate) async fn run_typed_step<I, O, F, Fut>(
step: &Step,
name: &str,
input: &[u8],
handler: F,
) -> std::result::Result<StepOutcome, StepError>
where
I: DeserializeOwned,
O: Serialize,
F: FnOnce(I) -> Fut,
Fut: Future<Output = std::result::Result<O, StepError>>,
{
run_recorded(step, async {
let input: I = rmp_serde::from_slice(input)
.map_err(|err| StepError::permanent(format!("invalid input for `{name}`: {err}")))?;
let output = handler(input).await?;
rmp_serde::to_vec_named(&output).map_err(|err| {
StepError::permanent(format!(
"`{name}` produced an output that failed to serialize: {err}"
))
})
})
.await
}
async fn run_recorded<F>(step: &Step, produce: F) -> std::result::Result<StepOutcome, StepError>
where
F: Future<Output = std::result::Result<Vec<u8>, StepError>>,
{
let input_hash = hash_input(&step.payload);
match produce.await {
Ok(bytes) => {
let record = OutcomeRecord {
input_hash,
outcome: StoredOutcome::Success {
output: bytes.clone(),
},
};
write_outcome(&step.run_memo, &record)
.await
.map_err(|err| StepError::transient(err.to_string()))?;
Ok(StepOutcome::Succeed { result: bytes })
}
Err(err) => {
let exhausted = step.attempts >= step.max_attempts;
if matches!(err.kind, StepErrorKind::Permanent) || exhausted {
let record = OutcomeRecord {
input_hash,
outcome: StoredOutcome::Failure {
kind: err.kind.into(),
message: err.message.clone(),
},
};
if let Err(write_err) = write_outcome(&step.run_memo, &record).await {
tracing::warn!(
run_id = %step.run_id,
"failed to persist the run's failure outcome: {write_err}"
);
}
}
Err(err)
}
}
}
pub(crate) async fn read_outcome(run_memo: &Memo) -> Result<Option<OutcomeRecord>> {
match run_memo.get(OUTCOME_KEY).await? {
None => Ok(None),
Some(bytes) => match rmp_serde::from_slice(&bytes) {
Ok(record) => Ok(Some(record)),
Err(err) => {
tracing::warn!(
run_id = %run_memo.run_id(),
error = %err,
"outcome record failed to decode; treated as absent",
);
Ok(None)
}
},
}
}
pub(crate) async fn write_outcome(run_memo: &Memo, record: &OutcomeRecord) -> Result<()> {
let bytes = rmp_serde::to_vec_named(record)?;
run_memo.put(OUTCOME_KEY, &bytes).await
}
pub(crate) enum Terminal {
Recorded(OutcomeRecord),
Unrecorded(Unrecorded),
}
pub(crate) enum Unrecorded {
Done,
Dead(Option<String>),
Cancelled,
NotFound,
}
pub(crate) async fn wait_terminal(
queue: &Queue,
run_memo: &Memo,
job_id: &str,
timeout: Duration,
) -> Result<Option<Terminal>> {
let unrecorded = match queue.wait_for_completion(job_id, timeout).await? {
WaitOutcome::TimedOut => return Ok(None),
WaitOutcome::Done(_) => Unrecorded::Done,
WaitOutcome::Dead(record) => Unrecorded::Dead(record.last_error),
WaitOutcome::Cancelled => Unrecorded::Cancelled,
WaitOutcome::NotFound => Unrecorded::NotFound,
};
Ok(Some(match read_outcome(run_memo).await? {
Some(record) => Terminal::Recorded(record),
None => Terminal::Unrecorded(unrecorded),
}))
}