1mod compfn;
16
17use std::fmt::{Debug, Display, Formatter};
18use std::future::Future;
19use std::pin::Pin;
20use std::rc::Rc;
21
22pub use compfn::*;
23
24pub type LocalBoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + 'a>>;
29
30#[inline]
40pub fn boxed<'a, T>(fut: impl Future<Output = T> + 'a) -> LocalBoxFuture<'a, T> {
41 Box::pin(fut)
42}
43
44#[allow(async_fn_in_trait)]
49pub trait Operation: Clone {
50 type In: Clone;
51 type Out: Clone;
52 type Err: Clone;
53
54 async fn execute(&self, input: Self::In) -> Result<Self::Out, Self::Err>;
56
57 async fn compensate(&self, input: Self::In, result: Self::Out) -> Result<(), Self::Err>;
59}
60
61#[allow(clippy::type_complexity)]
67pub fn operation<In: Clone, Out: Clone, Err: Clone>(
68 execute_fn: impl Fn(In) -> Pin<Box<dyn Future<Output = Result<Out, Err>>>> + 'static,
69 compensate_fn: impl Fn(In, Out) -> Pin<Box<dyn Future<Output = Result<(), Err>>>> + 'static,
70) -> impl Operation<In = In, Out = Out, Err = Err> {
71 FnOperation {
72 execute_fn: Rc::new(execute_fn),
73 compensate_fn: Rc::new(compensate_fn),
74 }
75}
76
77pub fn sync_operation<In: Clone + 'static, Out: Clone + 'static, Err: Clone + 'static>(
81 execute_fn: impl Fn(In) -> Result<Out, Err> + 'static,
82 compensate_fn: impl Fn(In, Out) -> Result<(), Err> + 'static,
83) -> impl Operation<In = In, Out = Out, Err = Err> {
84 let execute_fn = Rc::new(execute_fn);
85 let compensate_fn = Rc::new(compensate_fn);
86 operation(
87 move |input| {
88 let f = execute_fn.clone();
89 Box::pin(async move { f(input) })
90 },
91 move |input, output| {
92 let f = compensate_fn.clone();
93 Box::pin(async move { f(input, output) })
94 },
95 )
96}
97
98#[allow(clippy::type_complexity)]
99struct FnOperation<In, Out, Err> {
100 execute_fn: Rc<dyn Fn(In) -> Pin<Box<dyn Future<Output = Result<Out, Err>>>>>,
101 compensate_fn: Rc<dyn Fn(In, Out) -> Pin<Box<dyn Future<Output = Result<(), Err>>>>>,
102}
103
104impl<In, Out, Err> Clone for FnOperation<In, Out, Err> {
105 fn clone(&self) -> Self {
106 Self {
107 execute_fn: self.execute_fn.clone(),
108 compensate_fn: self.compensate_fn.clone(),
109 }
110 }
111}
112
113impl<In: Clone, Out: Clone, Err: Clone> Operation for FnOperation<In, Out, Err> {
114 type In = In;
115 type Out = Out;
116 type Err = Err;
117
118 async fn execute(&self, input: In) -> Result<Out, Err> {
119 (self.execute_fn)(input).await
120 }
121
122 async fn compensate(&self, input: In, result: Out) -> Result<(), Err> {
123 (self.compensate_fn)(input, result).await
124 }
125}
126
127pub type TransactionResult<Out, Err> = Result<Out, TransactionFailure<Err>>;
129
130#[derive(Debug)]
132pub enum TransactionFailure<Err> {
133 RetryExhausted(Err),
135 RetryPolicyError { failure: Err, message: String },
137 FailedAndRolledBackCompletely(Err),
139 FailedAndRolledBackPartially {
142 failure: Err,
143 compensation_failure: Err,
144 },
145}
146
147impl<Err: Display> Display for TransactionFailure<Err> {
148 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
149 match self {
150 TransactionFailure::RetryExhausted(err) => {
151 write!(f, "Transaction retry policy exhausted: {err}")
152 }
153 TransactionFailure::RetryPolicyError { failure, message } => write!(
154 f,
155 "Transaction failed with {failure}; retry policy error: {message}"
156 ),
157 TransactionFailure::FailedAndRolledBackCompletely(err) => {
158 write!(
159 f,
160 "Transaction failed with {err} and rolled back completely."
161 )
162 }
163 TransactionFailure::FailedAndRolledBackPartially {
164 failure,
165 compensation_failure,
166 } => write!(
167 f,
168 "Transaction failed with {failure} and rolled back partially; compensation failed with: {compensation_failure}."
169 ),
170 }
171 }
172}
173
174pub async fn fallible_transaction<Out, Err>(
178 f: impl for<'a> FnOnce(&'a mut FallibleTransaction<Err>) -> LocalBoxFuture<'a, Result<Out, Err>>,
179) -> TransactionResult<Out, Err>
180where
181 Err: Clone + 'static,
182{
183 let mut transaction = FallibleTransaction::new();
184 match f(&mut transaction).await {
185 Ok(output) => Ok(output),
186 Err(error) => Err(transaction.on_fail(error).await),
187 }
188}
189
190pub async fn infallible_transaction<Out, Err: Clone + 'static>(
197 mut f: impl for<'a> FnMut(&'a mut FallibleTransaction<Err>) -> LocalBoxFuture<'a, Result<Out, Err>>,
198) -> TransactionResult<Out, Err> {
199 let started = saga_host::now();
200 let mut state = Vec::new();
201 let mut attempt = 0_i64;
202 loop {
203 let mut transaction = FallibleTransaction::new();
204 match f(&mut transaction).await {
205 Ok(output) => return Ok(output),
206 Err(error) => {
207 let failure = transaction.on_fail(error).await;
208 let TransactionFailure::FailedAndRolledBackCompletely(error) = failure else {
209 return Err(failure);
210 };
211 attempt = attempt.saturating_add(1);
212 match saga_next(&state, saga_host::now().saturating_sub(started), attempt) {
213 Ok(decision) => {
214 state = decision.state;
215 match decision.delay {
216 Some(delay) => saga_host::wait(delay).await,
217 None => return Err(TransactionFailure::RetryExhausted(error)),
218 }
219 }
220 Err(message) => {
221 return Err(TransactionFailure::RetryPolicyError {
222 failure: error,
223 message,
224 });
225 }
226 }
227 }
228 }
229 }
230}
231
232fn saga_next(
233 state: &[u8],
234 elapsed: u64,
235 attempt: i64,
236) -> Result<crate::bindings::golem::api::retry::RetryDecision, String> {
237 use crate::bindings::golem::api::retry::PredicateValue;
238 let noun = saga_host::noun()?;
239 let properties = vec![
240 ("attempt".to_string(), PredicateValue::Integer(attempt)),
241 (
242 "error-kind".to_string(),
243 PredicateValue::Text("operation-error".to_string()),
244 ),
245 ];
246 let policy = saga_host::resolve("saga", &noun, &properties)?;
247 saga_host::next(policy.as_ref(), state, elapsed, "saga", &noun, &properties)
248}
249
250#[cfg(test)]
251mod saga_host {
252 use crate::bindings::golem::api::retry::{PredicateValue, RetryDecision, RetryPolicy};
253 use std::cell::Cell;
254 thread_local! {
255 pub static CLOCK: Cell<u64> = const { Cell::new(0) };
256 pub static ELAPSED: Cell<u64> = const { Cell::new(0) };
257 pub static POLICY_ERROR: Cell<bool> = const { Cell::new(false) };
258 pub static RESOLVE_ERROR: Cell<bool> = const { Cell::new(false) };
259 }
260 pub fn now() -> u64 {
261 CLOCK.get()
262 }
263 pub async fn wait(delay: u64) {
264 CLOCK.set(CLOCK.get() + delay);
265 }
266 pub fn noun() -> Result<String, String> {
267 Ok("urn:worker:component/agent".into())
268 }
269 pub fn resolve(
270 verb: &str,
271 noun: &str,
272 properties: &[(String, PredicateValue)],
273 ) -> Result<Option<RetryPolicy>, String> {
274 assert_eq!(verb, "saga");
275 assert_eq!(noun, "urn:worker:component/agent");
276 assert_eq!(properties[1].0, "error-kind");
277 assert!(
278 matches!(&properties[1].1, PredicateValue::Text(value) if value == "operation-error")
279 );
280 if RESOLVE_ERROR.replace(false) {
281 return Err("invalid resolver predicate".into());
282 }
283 Ok(None)
284 }
285 pub fn next(
286 policy: Option<&RetryPolicy>,
287 state: &[u8],
288 elapsed: u64,
289 verb: &str,
290 noun: &str,
291 properties: &[(String, PredicateValue)],
292 ) -> Result<RetryDecision, String> {
293 assert!(policy.is_none());
294 ELAPSED.set(elapsed);
295 resolve(verb, noun, properties)?;
296 assert_eq!(properties[0].0, "attempt");
297 let PredicateValue::Integer(attempt) = properties[0].1 else {
298 panic!("missing attempt")
299 };
300 if POLICY_ERROR.replace(false) {
301 return Err("invalid policy".into());
302 }
303 let expected = if attempt == 1 {
304 vec![]
305 } else {
306 vec![attempt as u8 - 1]
307 };
308 assert_eq!(state, expected);
309 assert!(elapsed >= (attempt as u64 - 1) * 10);
310 Ok(RetryDecision {
311 state: vec![attempt as u8],
312 delay: (attempt < 3).then_some(10),
313 })
314 }
315}
316
317#[cfg(not(test))]
318mod saga_host {
319 pub use crate::bindings::golem::api::retry::next_retry_delay as next;
320 pub use crate::retry::resolve_retry_policy as resolve;
321 pub fn now() -> u64 {
322 crate::wasip3::clocks::monotonic_clock::now()
323 }
324 pub async fn wait(delay: u64) {
325 crate::wasip3::clocks::monotonic_clock::wait_for(delay).await;
326 }
327 pub fn noun() -> Result<String, String> {
328 let metadata = crate::get_self_metadata().map_err(|e| format!("{e:?}"))?;
329 let id = metadata.agent_id;
330 Ok(format!(
331 "urn:worker:{}/{}",
332 id.component_id.uuid, id.agent_id
333 ))
334 }
335}
336
337pub async fn transaction<Out: 'static, Err, T>(
339 f: impl for<'a> FnOnce(&'a mut T) -> LocalBoxFuture<'a, Result<Out, Err>>,
340) -> TransactionResult<Out, Err>
341where
342 T: Transaction<Err>,
343{
344 T::run(f).await
345}
346
347#[allow(clippy::type_complexity)]
349struct CompensationAction<Err> {
350 action: Box<dyn FnOnce() -> Pin<Box<dyn Future<Output = Result<(), Err>>>>>,
351}
352
353impl<Err> CompensationAction<Err> {
354 pub async fn execute(self) -> Result<(), Err> {
355 (self.action)().await
356 }
357}
358
359pub struct FallibleTransaction<Err> {
366 compensations: Vec<CompensationAction<Err>>,
367}
368
369impl<Err: Clone + 'static> FallibleTransaction<Err> {
370 fn new() -> Self {
371 Self {
372 compensations: Vec::new(),
373 }
374 }
375
376 pub async fn execute<OpIn: Clone + 'static, OpOut: Clone + 'static>(
377 &mut self,
378 operation: impl Operation<In = OpIn, Out = OpOut, Err = Err> + 'static,
379 input: OpIn,
380 ) -> Result<OpOut, Err> {
381 let result = operation.execute(input.clone()).await;
382 if let Ok(output) = &result {
383 let cloned_op = operation.clone();
384 let cloned_out = output.clone();
385 self.compensations.push(CompensationAction {
386 action: Box::new(move || {
387 Box::pin(async move {
388 cloned_op
389 .compensate(input.clone(), cloned_out.clone())
390 .await
391 })
392 }),
393 });
394 }
395 result
396 }
397
398 async fn on_fail(&mut self, failure: Err) -> TransactionFailure<Err> {
399 for compensation_action in self.compensations.drain(..).rev() {
400 if let Err(compensation_failure) = compensation_action.execute().await {
401 return TransactionFailure::FailedAndRolledBackPartially {
402 failure,
403 compensation_failure,
404 };
405 }
406 }
407 TransactionFailure::FailedAndRolledBackCompletely(failure)
408 }
409}
410
411#[allow(async_fn_in_trait)]
415pub trait Transaction<Err> {
416 async fn execute<OpIn: Clone + 'static, OpOut: Clone + 'static>(
417 &mut self,
418 operation: impl Operation<In = OpIn, Out = OpOut, Err = Err> + 'static,
419 input: OpIn,
420 ) -> Result<OpOut, Err>;
421
422 async fn fail(&mut self, error: Err) -> Result<(), Err>;
423
424 async fn run<Out: 'static>(
425 f: impl for<'a> FnOnce(&'a mut Self) -> LocalBoxFuture<'a, Result<Out, Err>>,
426 ) -> TransactionResult<Out, Err>;
427}
428
429impl<Err: Clone + 'static> Transaction<Err> for FallibleTransaction<Err> {
430 async fn execute<OpIn: Clone + 'static, OpOut: Clone + 'static>(
431 &mut self,
432 operation: impl Operation<In = OpIn, Out = OpOut, Err = Err> + 'static,
433 input: OpIn,
434 ) -> Result<OpOut, Err> {
435 FallibleTransaction::execute(self, operation, input).await
436 }
437
438 async fn fail(&mut self, error: Err) -> Result<(), Err> {
439 Err(error)
440 }
441
442 async fn run<Out: 'static>(
443 f: impl for<'a> FnOnce(&'a mut Self) -> LocalBoxFuture<'a, Result<Out, Err>>,
444 ) -> TransactionResult<Out, Err> {
445 fallible_transaction(f).await
446 }
447}
448
449#[cfg(test)]
450mod tests {
451 use std::cell::RefCell;
452 use std::rc::Rc;
453 use test_r::test;
454
455 use crate::{boxed, fallible_transaction, infallible_transaction, sync_operation};
456
457 #[test]
458 async fn bounded_retries_wait_and_exhaust() {
459 super::saga_host::CLOCK.set(0);
460 let op = sync_operation(
461 |()| {
462 super::saga_host::CLOCK.set(super::saga_host::CLOCK.get() + 2);
463 Ok::<_, &str>(())
464 },
465 |(), ()| {
466 super::saga_host::CLOCK.set(super::saga_host::CLOCK.get() + 3);
467 Ok(())
468 },
469 );
470 let mut attempts = 0;
471 let result = infallible_transaction(|tx| {
472 attempts += 1;
473 let op = op.clone();
474 boxed(async move {
475 tx.execute(op, ()).await?;
476 Err::<(), _>("operation")
477 })
478 })
479 .await;
480 assert!(matches!(
481 result,
482 Err(crate::TransactionFailure::RetryExhausted("operation"))
483 ));
484 assert_eq!(attempts, 3);
485 assert_eq!(super::saga_host::CLOCK.get(), 35);
486 assert_eq!(super::saga_host::ELAPSED.get(), 35);
487 }
488
489 #[test]
490 async fn resolver_error_preserves_original_failure() {
491 super::saga_host::RESOLVE_ERROR.set(true);
492 let mut attempts = 0;
493 let result = infallible_transaction(|_| {
494 attempts += 1;
495 boxed(async { Err::<(), _>("original") })
496 })
497 .await;
498 assert!(matches!(
499 result,
500 Err(crate::TransactionFailure::RetryPolicyError { failure: "original", message })
501 if message == "invalid resolver predicate"
502 ));
503 assert_eq!(attempts, 1);
504 }
505
506 #[test]
507 async fn invalid_policy_never_retries() {
508 super::saga_host::POLICY_ERROR.set(true);
509 let mut attempts = 0;
510 let result = infallible_transaction(|_| {
511 attempts += 1;
512 boxed(async { Err::<(), _>("operation") })
513 })
514 .await;
515 assert!(matches!(
516 result,
517 Err(crate::TransactionFailure::RetryPolicyError {
518 failure: "operation",
519 ..
520 })
521 ));
522 assert_eq!(attempts, 1);
523 }
524
525 #[test]
526 async fn compensation_order_and_partial_failure() {
527 let log = Rc::new(RefCell::new(Vec::new()));
528 let compensation_log = log.clone();
529 let op = sync_operation(
530 |input: u8| Ok::<_, &str>(input),
531 move |input, output| {
532 assert_eq!(input, output);
533 compensation_log.borrow_mut().push(input);
534 if input == 2 {
535 Err("compensation")
536 } else {
537 Ok(())
538 }
539 },
540 );
541 let result = fallible_transaction(|tx| {
542 boxed(async move {
543 tx.execute(op.clone(), 1).await?;
544 tx.execute(op.clone(), 2).await?;
545 tx.execute(op, 3).await?;
546 Err::<(), _>("body")
547 })
548 })
549 .await;
550 assert!(matches!(
551 result,
552 Err(crate::TransactionFailure::FailedAndRolledBackPartially {
553 failure: "body",
554 compensation_failure: "compensation"
555 })
556 ));
557 assert_eq!(*log.borrow(), [3, 2]);
558 }
559
560 #[test]
561 async fn infallible_stops_on_compensation_failure() {
562 let op = sync_operation(|()| Ok::<_, &str>(()), |(), ()| Err("compensation"));
563 let mut attempts = 0;
564 let result = infallible_transaction(|tx| {
565 attempts += 1;
566 assert_eq!(attempts, 1);
567 let op = op.clone();
568 boxed(async move {
569 tx.execute(op, ()).await?;
570 Err::<(), _>("body")
571 })
572 })
573 .await;
574 assert!(matches!(
575 result,
576 Err(crate::TransactionFailure::FailedAndRolledBackPartially {
577 failure: "body",
578 compensation_failure: "compensation"
579 })
580 ));
581 }
582
583 #[test]
584 async fn fallible_compensates_successful_operations_only() {
585 let log = Rc::new(RefCell::new(Vec::new()));
586
587 let log1 = log.clone();
588 let log2 = log.clone();
589 let log3 = log.clone();
590 let log4 = log.clone();
591
592 let op1 = sync_operation(
593 move |input: String| {
594 log1.borrow_mut().push(format!("op1 execute {input}"));
595 Ok(())
596 },
597 move |input: String, _| {
598 log2.borrow_mut().push(format!("op1 rollback {input}"));
599 Ok(())
600 },
601 );
602
603 let op2 = sync_operation(
604 move |_: ()| {
605 log3.clone().borrow_mut().push("op2 execute".to_string());
606 Err::<(), &str>("op2 error")
607 },
608 move |_: (), _| {
609 log4.clone().borrow_mut().push("op2 rollback".to_string());
610 Ok(())
611 },
612 );
613
614 let result = fallible_transaction(|tx| {
615 boxed(async move {
616 println!("First we execute op1");
617 tx.execute(op1, "hello".to_string()).await?;
618 println!("Then execute op2");
619 tx.execute(op2, ()).await?;
620 println!("Finally compute a result");
621 Ok(11)
622 })
623 })
624 .await;
625
626 assert!(matches!(
627 result,
628 Err(crate::TransactionFailure::FailedAndRolledBackCompletely(
629 "op2 error"
630 ))
631 ));
632 assert_eq!(
633 *log.borrow(),
634 ["op1 execute hello", "op2 execute", "op1 rollback hello"]
635 );
636 }
637
638 #[test]
639 async fn infallible_retries_with_fresh_compensations() {
640 let log = Rc::new(RefCell::new(Vec::new()));
641
642 let log1 = log.clone();
643 let log2 = log.clone();
644 let log3 = log.clone();
645 let log4 = log.clone();
646
647 let op1 = sync_operation(
648 move |input: String| {
649 log1.borrow_mut().push(format!("op1 execute {input}"));
650 Ok::<(), &str>(())
651 },
652 move |input: String, _| {
653 log2.borrow_mut().push(format!("op1 rollback {input}"));
654 Ok(())
655 },
656 );
657
658 let attempts = Rc::new(std::cell::Cell::new(0));
659 let count = attempts.clone();
660 let op2 = sync_operation(
661 move |_: ()| {
662 log3.clone().borrow_mut().push("op2 execute".to_string());
663 count.set(count.get() + 1);
664 if count.get() < 3 {
665 Err("op2 error")
666 } else {
667 Ok(())
668 }
669 },
670 move |_: (), r| {
671 log4.clone()
672 .borrow_mut()
673 .push(format!("op2 rollback {r:?}"));
674 Ok(())
675 },
676 );
677
678 let result = infallible_transaction(|tx| {
679 let op1 = op1.clone();
680 let op2 = op2.clone();
681 boxed(async move {
682 println!("First we execute op1");
683 tx.execute(op1.clone(), "hello".to_string()).await?;
684 tx.execute(op1, "world".to_string()).await?;
685 println!("Then execute op2");
686 tx.execute(op2, ()).await?;
687 println!("Finally compute a result");
688 Ok(11)
689 })
690 })
691 .await;
692
693 assert_eq!(result.unwrap(), 11);
694 assert_eq!(attempts.get(), 3);
695 assert_eq!(
696 *log.borrow(),
697 [
698 "op1 execute hello",
699 "op1 execute world",
700 "op2 execute",
701 "op1 rollback world",
702 "op1 rollback hello",
703 "op1 execute hello",
704 "op1 execute world",
705 "op2 execute",
706 "op1 rollback world",
707 "op1 rollback hello",
708 "op1 execute hello",
709 "op1 execute world",
710 "op2 execute"
711 ]
712 );
713 }
714}
715
716#[cfg(test)]
717#[cfg(feature = "macro")]
718mod macro_tests {
719 use crate::{boxed, fallible_transaction, infallible_transaction};
720 use golem_rust_macro::golem_operation;
721 use test_r::test;
722
723 mod golem_rust {
724 pub use crate::*;
725 }
726
727 #[golem_operation(compensation=test_compensation)]
728 fn test_operation(input1: u64, input2: f32) -> Result<bool, String> {
729 println!("Op input: {input1}, {input2}");
730 Ok(true)
731 }
732
733 fn test_compensation(_: bool, input1: u64, input2: f32) -> Result<(), String> {
734 println!("Compensation input: {input1}, {input2}");
735 Ok(())
736 }
737
738 #[golem_operation(compensation=test_compensation_2)]
739 fn test_operation_2(input1: u64, input2: f32) -> Result<bool, String> {
740 println!("Op input: {input1}, {input2}");
741 Ok(true)
742 }
743
744 fn test_compensation_2(result: bool) -> Result<(), String> {
745 println!("Compensation for operation result {result:?}");
746 Ok(())
747 }
748
749 #[golem_operation(compensation=test_compensation_3)]
750 fn test_operation_3(input: String) -> Result<(), String> {
751 println!("Op input: {input}");
752 Ok(())
753 }
754
755 fn test_compensation_3() -> Result<(), String> {
756 println!("Compensation for operation, not using any input");
757 Ok(())
758 }
759
760 #[golem_operation(compensation=test_compensation_4)]
761 fn test_operation_4(input: u64) -> Result<(), String> {
762 println!("Op input: {input}");
763 Ok(())
764 }
765
766 fn test_compensation_4(_: (), input: u64) -> Result<(), String> {
767 println!("Compensation for operation with single input {input}");
768 Ok(())
769 }
770
771 #[test]
773 #[ignore]
774 async fn tx_test_1() {
775 let result = fallible_transaction(|tx| {
776 boxed(async move {
777 println!("Executing the annotated function as an operation directly");
778 tx.test_operation(1, 0.1).await?;
779 tx.test_operation_2(1, 0.1).await?;
780 tx.test_operation_3("test".to_string()).await?;
781 tx.test_operation_4(1).await?;
782
783 Ok(11)
784 })
785 })
786 .await;
787
788 println!("{result:?}");
789 }
790
791 #[test]
793 #[ignore]
794 async fn tx_test_2() {
795 let result = infallible_transaction(|tx| {
796 boxed(async move {
797 println!("Executing the annotated function as an operation directly");
798 tx.test_operation(1, 0.1).await?;
799 Ok(11)
800 })
801 })
802 .await;
803
804 println!("{result:?}");
805 }
806}