use std::sync::Arc;
use crate::adapter::WrappingDispenser;
use crate::adapter::{ExecutionError, OpDispenser, OpResult};
use crate::wrapper_registry::{WrapperName, WrapperRegistration, WrapperSubject};
pub const NAME: WrapperName = WrapperName::new("memo");
fn triggers(s: WrapperSubject) -> bool {
let Some(template) = s.op() else {
return false;
};
template
.params
.get("memo")
.map(|v| v.is_string() || v.is_object())
.unwrap_or(false)
}
fn describe_assignment(s: WrapperSubject) -> Option<String> {
let template = s.op()?;
let v = template.params.get("memo")?;
if let Some(s) = v.as_str() {
if s.is_empty() {
return None;
}
Some(format!("memo: \"{s}\" (before+after)"))
} else if let Some(obj) = v.as_object() {
let before = obj.get("before").and_then(|x| x.as_str());
let after = obj.get("after").and_then(|x| x.as_str());
match (before, after) {
(Some(b), Some(a)) => Some(format!("memo: before \"{b}\" / after \"{a}\"")),
(Some(b), None) => Some(format!("memo: before \"{b}\"")),
(None, Some(a)) => Some(format!("memo: after \"{a}\"")),
(None, None) => None,
}
} else {
None
}
}
inventory::submit! {
WrapperRegistration {
name: NAME,
owned_fields: &["memo"],
triggers,
requires_inner: &[super::traverse::NAME],
forbids_outer: &[],
mutually_exclusive_with: &[],
describe_assignment,
levels: &[crate::wrapper_registry::WrapperLevel::Op],
}
}
pub struct MemoDispenser {
inner: Arc<dyn OpDispenser>,
before_template: Option<String>,
after_template: Option<String>,
memo_state: Arc<arc_swap::ArcSwap<String>>,
}
impl MemoDispenser {
pub fn wrap(
inner: Arc<dyn OpDispenser>,
before_template: Option<String>,
after_template: Option<String>,
memo_state: Arc<arc_swap::ArcSwap<String>>,
) -> Arc<dyn OpDispenser> {
Arc::new(Self {
inner,
before_template,
after_template,
memo_state,
})
}
fn publish(&self, template: &str, wires: &dyn crate::wires::WireSource) {
match crate::wires::substitute_via_wires(template, wires) {
Ok(rendered) => {
self.memo_state.store(Arc::new(rendered));
}
Err(e) => {
crate::diag!(
crate::observer::LogLevel::Debug,
"memo: substitution failed for '{template}': {e}"
);
}
}
}
}
impl WrappingDispenser for MemoDispenser {}
impl OpDispenser for MemoDispenser {
fn execute<'a>(
&'a self,
cycle: u64,
ctx: &'a crate::fixture::ExecCtx<'a>,
) -> std::pin::Pin<
Box<dyn std::future::Future<Output = Result<OpResult, ExecutionError>> + Send + 'a>,
> {
Box::pin(async move {
if let Some(t) = &self.before_template {
self.publish(t, ctx.wires);
}
let result = self.inner.execute(cycle, ctx).await?;
if let Some(t) = &self.after_template {
self.publish(t, ctx.wires);
}
Ok(result)
})
}
fn inner_dispenser(&self) -> Option<&dyn OpDispenser> {
Some(self.inner.as_ref())
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::adapter::{AdapterError, ExecutionError, OpResult, ResultBody};
use crate::fixture::{ExecCtx, ResolvedPulls};
#[derive(Debug)]
struct ResultDispBody {
value: serde_json::Value,
count: u64,
}
impl ResultBody for ResultDispBody {
fn to_json(&self) -> serde_json::Value {
self.value.clone()
}
fn as_any(&self) -> &dyn std::any::Any {
self
}
fn element_count(&self) -> u64 {
self.count
}
}
struct FakeInner {
body: Option<ResultDispBody>,
error: Option<&'static str>,
}
impl OpDispenser for FakeInner {
fn execute<'a>(
&'a self,
_cycle: u64,
_ctx: &'a ExecCtx<'a>,
) -> std::pin::Pin<
Box<dyn std::future::Future<Output = Result<OpResult, ExecutionError>> + Send + 'a>,
> {
Box::pin(async move {
if let Some(msg) = self.error {
return Err(ExecutionError::Op(AdapterError {
error_name: "test".into(),
message: msg.into(),
retryable: false,
}));
}
Ok(OpResult {
body: self.body.as_ref().map(|b| {
Box::new(ResultDispBody {
value: b.value.clone(),
count: b.count,
}) as Box<dyn ResultBody>
}),
skipped: false,
})
})
}
}
fn empty_ctx() -> (crate::adapter::ResolvedFields, ResolvedPulls) {
let fields = crate::adapter::ResolvedFields::new(vec![], vec![]);
let pulls = ResolvedPulls::empty();
(fields, pulls)
}
#[tokio::test]
async fn memo_wrapper_publishes_before_and_after() {
let memo = Arc::new(arc_swap::ArcSwap::from_pointee(String::new()));
let inner = Arc::new(FakeInner {
body: None,
error: None,
});
let dispenser = MemoDispenser::wrap(
inner,
Some("before-state".into()),
Some("after-state".into()),
memo.clone(),
);
let (fields, pulls) = empty_ctx();
let ctx = ExecCtx::new(&fields, &pulls);
let _ = dispenser.execute(0, &ctx).await.expect("inner ok");
assert_eq!(memo.load().as_str(), "after-state");
}
#[tokio::test]
async fn memo_wrapper_only_before_when_after_unset() {
let memo = Arc::new(arc_swap::ArcSwap::from_pointee(String::new()));
let inner = Arc::new(FakeInner {
body: None,
error: None,
});
let dispenser = MemoDispenser::wrap(inner, Some("ready".into()), None, memo.clone());
let (fields, pulls) = empty_ctx();
let ctx = ExecCtx::new(&fields, &pulls);
let _ = dispenser.execute(0, &ctx).await.expect("inner ok");
assert_eq!(memo.load().as_str(), "ready");
}
#[tokio::test]
async fn memo_wrapper_does_not_run_after_on_inner_error() {
let memo = Arc::new(arc_swap::ArcSwap::from_pointee(String::new()));
let inner = Arc::new(FakeInner {
body: None,
error: Some("boom"),
});
let dispenser = MemoDispenser::wrap(
inner,
Some("attempting".into()),
Some("finished".into()),
memo.clone(),
);
let (fields, pulls) = empty_ctx();
let ctx = ExecCtx::new(&fields, &pulls);
let res = dispenser.execute(0, &ctx).await;
assert!(res.is_err());
assert_eq!(
memo.load().as_str(),
"attempting",
"after-template must not run on inner error"
);
}
}