golem-rust 2.0.0-dev.7

// Copyright 2024-2026 Golem Cloud
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

mod compfn;

use std::fmt::{Debug, Display, Formatter};
use std::future::Future;
use std::pin::Pin;
use std::rc::Rc;

use crate::bindings::golem::api::host::{OplogIndex, get_oplog_index, set_oplog_index};
use crate::mark_atomic_operation;

pub use compfn::*;

/// A boxed local future. Used as the return type for transaction closures
/// to allow the returned future to borrow from the transaction reference.
///
/// Use [`boxed`] to construct this from an async block.
pub type LocalBoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + 'a>>;

/// Wraps an async block into a [`LocalBoxFuture`]. Use this when passing
/// closures to [`fallible_transaction`] or [`infallible_transaction`]:
///
/// ```ignore
/// fallible_transaction(|tx| boxed(async move {
///     tx.execute(op1, input).await?;
///     Ok(result)
/// })).await;
/// ```
#[inline]
pub fn boxed<'a, T>(fut: impl Future<Output = T> + 'a) -> LocalBoxFuture<'a, T> {
    Box::pin(fut)
}

/// Represents an atomic operation of the transaction which has a rollback action.
///
/// Implement this trait and use it within a `transaction` block.
/// Operations can also be constructed from closures using `operation` or `sync_operation`.
#[allow(async_fn_in_trait)]
pub trait Operation: Clone {
    type In: Clone;
    type Out: Clone;
    type Err: Clone;

    /// Executes the operation which may fail with a domain error
    async fn execute(&self, input: Self::In) -> Result<Self::Out, Self::Err>;

    /// Executes a compensation action for the operation.
    async fn compensate(&self, input: Self::In, result: Self::Out) -> Result<(), Self::Err>;
}

/// Constructs an `Operation` from two async closures: one for executing the operation,
/// and one for rolling it back. The rollback operation always sees the input and
/// the output of the operation.
///
/// This operation can run the compensation in both fallible and infallible transactions.
#[allow(clippy::type_complexity)]
pub fn operation<In: Clone, Out: Clone, Err: Clone>(
    execute_fn: impl Fn(In) -> Pin<Box<dyn Future<Output = Result<Out, Err>>>> + 'static,
    compensate_fn: impl Fn(In, Out) -> Pin<Box<dyn Future<Output = Result<(), Err>>>> + 'static,
) -> impl Operation<In = In, Out = Out, Err = Err> {
    FnOperation {
        execute_fn: Rc::new(execute_fn),
        compensate_fn: Rc::new(compensate_fn),
    }
}

/// Constructs an `Operation` from two synchronous closures: one for executing the operation,
/// and one for rolling it back. This is a convenience wrapper around `operation` for
/// cases where the execute and compensate logic is synchronous.
pub fn sync_operation<In: Clone + 'static, Out: Clone + 'static, Err: Clone + 'static>(
    execute_fn: impl Fn(In) -> Result<Out, Err> + 'static,
    compensate_fn: impl Fn(In, Out) -> Result<(), Err> + 'static,
) -> impl Operation<In = In, Out = Out, Err = Err> {
    let execute_fn = Rc::new(execute_fn);
    let compensate_fn = Rc::new(compensate_fn);
    operation(
        move |input| {
            let f = execute_fn.clone();
            Box::pin(async move { f(input) })
        },
        move |input, output| {
            let f = compensate_fn.clone();
            Box::pin(async move { f(input, output) })
        },
    )
}

#[allow(clippy::type_complexity)]
struct FnOperation<In, Out, Err> {
    execute_fn: Rc<dyn Fn(In) -> Pin<Box<dyn Future<Output = Result<Out, Err>>>>>,
    compensate_fn: Rc<dyn Fn(In, Out) -> Pin<Box<dyn Future<Output = Result<(), Err>>>>>,
}

impl<In, Out, Err> Clone for FnOperation<In, Out, Err> {
    fn clone(&self) -> Self {
        Self {
            execute_fn: self.execute_fn.clone(),
            compensate_fn: self.compensate_fn.clone(),
        }
    }
}

impl<In: Clone, Out: Clone, Err: Clone> Operation for FnOperation<In, Out, Err> {
    type In = In;
    type Out = Out;
    type Err = Err;

    async fn execute(&self, input: In) -> Result<Out, Err> {
        (self.execute_fn)(input).await
    }

    async fn compensate(&self, input: In, result: Out) -> Result<(), Err> {
        (self.compensate_fn)(input, result).await
    }
}

/// The result of a transaction execution.
pub type TransactionResult<Out, Err> = Result<Out, TransactionFailure<Err>>;

/// The result of a transaction execution that failed.
#[derive(Debug)]
pub enum TransactionFailure<Err> {
    /// One of the operations failed with an error, and the transaction was fully rolled back.
    FailedAndRolledBackCompletely(Err),
    /// One of the operations failed with an error, and the transaction was partially rolled back
    /// because the compensation action of one of the operations also failed.
    FailedAndRolledBackPartially {
        failure: Err,
        compensation_failure: Err,
    },
}

impl<Err: Display> Display for TransactionFailure<Err> {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            TransactionFailure::FailedAndRolledBackCompletely(err) => {
                write!(
                    f,
                    "Transaction failed with {err} and rolled back completely."
                )
            }
            TransactionFailure::FailedAndRolledBackPartially {
                failure,
                compensation_failure,
            } => write!(
                f,
                "Transaction failed with {failure} and rolled back partially; compensation failed with: {compensation_failure}."
            ),
        }
    }
}

/// Fallible transaction execution. If any operation fails, all the already executed
/// successful operation's compensation actions are executed in reverse order and the transaction
/// returns with a failure.
pub async fn fallible_transaction<Out, Err>(
    f: impl for<'a> FnOnce(&'a mut FallibleTransaction<Err>) -> LocalBoxFuture<'a, Result<Out, Err>>,
) -> TransactionResult<Out, Err>
where
    Err: Clone + 'static,
{
    let mut transaction = FallibleTransaction::new();
    match f(&mut transaction).await {
        Ok(output) => Ok(output),
        Err(error) => Err(transaction.on_fail(error).await),
    }
}

/// Retry the transaction in case of failure. If any operation returns with a failure, all
/// the already executed successful operation's compensation actions are executed in reverse order
/// and the transaction gets retried, using Golem's active retry policy.
pub async fn infallible_transaction<Out>(
    f: impl for<'a> FnOnce(&'a mut InfallibleTransaction) -> LocalBoxFuture<'a, Out>,
) -> Out {
    let oplog_index = get_oplog_index();
    let _atomic_region = mark_atomic_operation();
    let mut transaction = InfallibleTransaction::new(oplog_index);
    f(&mut transaction).await
}

/// Same as `infallible_transaction`, but with strong rollback guarantees. The compensation actions
/// are guaranteed to be always executed before the transaction gets retried, even if it
/// fails due to a panic or an external executor failure.
pub async fn infallible_transaction_with_strong_rollback_guarantees<Out>(
    _f: impl for<'a> FnOnce(&'a mut InfallibleTransaction) -> LocalBoxFuture<'a, Out>,
) -> Out {
    unimplemented!()
}

/// A generic interface for defining transactions, where the transaction mode is
/// determined by the function's parameter (it can be `FallibleTransaction` or `InfallibleTransaction`).
///
/// This makes switching between different transaction guarantees easier, but is more constrained
/// than using the specific transaction functions where for retried transactions errors does
/// not have to be handled.
pub async fn transaction<Out: 'static, Err, T>(
    f: impl for<'a> FnOnce(&'a mut T) -> LocalBoxFuture<'a, Result<Out, Err>>,
) -> TransactionResult<Out, Err>
where
    T: Transaction<Err>,
{
    T::run(f).await
}

/// Helper struct for coupling compensation action and the result of the operation.
#[allow(clippy::type_complexity)]
struct CompensationAction<Err> {
    action: Box<dyn FnOnce() -> Pin<Box<dyn Future<Output = Result<(), Err>>>>>,
}

impl<Err> CompensationAction<Err> {
    pub async fn execute(self) -> Result<(), Err> {
        (self.action)().await
    }
}

/// FallibleTransaction is a sequence of operations that are executed in a way that if any of the
/// operations fails all the already performed operation's compensation actions got executed in
/// reverse order.
///
/// In case of fatal errors (panic) and external executor failures it does not perform the
/// compensation actions and the whole transaction gets retried.
pub struct FallibleTransaction<Err> {
    compensations: Vec<CompensationAction<Err>>,
}

impl<Err: Clone + 'static> FallibleTransaction<Err> {
    fn new() -> Self {
        Self {
            compensations: Vec::new(),
        }
    }

    pub async fn execute<OpIn: Clone + 'static, OpOut: Clone + 'static>(
        &mut self,
        operation: impl Operation<In = OpIn, Out = OpOut, Err = Err> + 'static,
        input: OpIn,
    ) -> Result<OpOut, Err> {
        let result = operation.execute(input.clone()).await;
        if let Ok(output) = &result {
            let cloned_op = operation.clone();
            let cloned_out = output.clone();
            self.compensations.push(CompensationAction {
                action: Box::new(move || {
                    Box::pin(async move {
                        cloned_op
                            .compensate(input.clone(), cloned_out.clone())
                            .await
                    })
                }),
            });
        }
        result
    }

    async fn on_fail(&mut self, failure: Err) -> TransactionFailure<Err> {
        for compensation_action in self.compensations.drain(..).rev() {
            if let Err(compensation_failure) = compensation_action.execute().await {
                return TransactionFailure::FailedAndRolledBackPartially {
                    failure,
                    compensation_failure,
                };
            }
        }
        TransactionFailure::FailedAndRolledBackCompletely(failure)
    }
}

/// InfallibleTransaction is a sequence of operations that are executed in a way that if any of the
/// operations or the underlying Golem executor fails, the whole transaction is going to
/// be retried.
///
/// In addition to that, **user level failures** (represented by the `Result::Err` value
/// of an operation) lead to performing the compensation actions of each already performed operation
/// in reverse order.
///
/// Fatal errors (panic) and external executor failures are currently cannot perform the
/// rollback actions.
pub struct InfallibleTransaction {
    begin_oplog_index: OplogIndex,
    compensations: Vec<CompensationAction<()>>,
}

impl InfallibleTransaction {
    fn new(begin_oplog_index: OplogIndex) -> Self {
        Self {
            begin_oplog_index,
            compensations: Vec::new(),
        }
    }

    pub async fn execute<
        OpIn: Clone + 'static,
        OpOut: Clone + 'static,
        OpErr: Debug + Clone + 'static,
    >(
        &mut self,
        operation: impl Operation<In = OpIn, Out = OpOut, Err = OpErr> + 'static,
        input: OpIn,
    ) -> OpOut {
        match operation.execute(input.clone()).await {
            Ok(output) => {
                let cloned_op = operation.clone();
                let cloned_out = output.clone();
                self.compensations.push(CompensationAction {
                    action: Box::new(move || {
                        Box::pin(async move {
                            cloned_op
                                .compensate(input.clone(), cloned_out.clone())
                                .await
                                .expect("Compensation action failed");
                            Ok(())
                        })
                    }),
                });
                output
            }
            Err(_) => {
                self.retry().await;
                unreachable!()
            }
        }
    }

    /// Stop executing the transaction and retry from the beginning, after executing the compensation actions
    pub async fn retry(&mut self) {
        for compensation_action in self.compensations.drain(..).rev() {
            let _ = compensation_action.execute().await;
        }
        set_oplog_index(self.begin_oplog_index);
    }
}

/// A unified interface for the different types of transactions. Using it can make the code
/// easier to switch between different transactional guarantees but is more constrained in
/// terms of error types.
#[allow(async_fn_in_trait)]
pub trait Transaction<Err> {
    async fn execute<OpIn: Clone + 'static, OpOut: Clone + 'static>(
        &mut self,
        operation: impl Operation<In = OpIn, Out = OpOut, Err = Err> + 'static,
        input: OpIn,
    ) -> Result<OpOut, Err>;

    async fn fail(&mut self, error: Err) -> Result<(), Err>;

    async fn run<Out: 'static>(
        f: impl for<'a> FnOnce(&'a mut Self) -> LocalBoxFuture<'a, Result<Out, Err>>,
    ) -> TransactionResult<Out, Err>;
}

impl<Err: Clone + 'static> Transaction<Err> for FallibleTransaction<Err> {
    async fn execute<OpIn: Clone + 'static, OpOut: Clone + 'static>(
        &mut self,
        operation: impl Operation<In = OpIn, Out = OpOut, Err = Err> + 'static,
        input: OpIn,
    ) -> Result<OpOut, Err> {
        FallibleTransaction::execute(self, operation, input).await
    }

    async fn fail(&mut self, error: Err) -> Result<(), Err> {
        Err(error)
    }

    async fn run<Out: 'static>(
        f: impl for<'a> FnOnce(&'a mut Self) -> LocalBoxFuture<'a, Result<Out, Err>>,
    ) -> TransactionResult<Out, Err> {
        fallible_transaction(f).await
    }
}

impl<Err: Debug + Clone + 'static> Transaction<Err> for InfallibleTransaction {
    async fn execute<OpIn: Clone + 'static, OpOut: Clone + 'static>(
        &mut self,
        operation: impl Operation<In = OpIn, Out = OpOut, Err = Err> + 'static,
        input: OpIn,
    ) -> Result<OpOut, Err> {
        Ok(InfallibleTransaction::execute(self, operation, input).await)
    }

    async fn fail(&mut self, error: Err) -> Result<(), Err> {
        InfallibleTransaction::retry(self).await;
        Err(error)
    }

    async fn run<Out: 'static>(
        f: impl for<'a> FnOnce(&'a mut Self) -> LocalBoxFuture<'a, Result<Out, Err>>,
    ) -> TransactionResult<Out, Err> {
        Ok(infallible_transaction(|tx| -> LocalBoxFuture<'_, Out> {
            let fut = f(tx);
            Box::pin(async { fut.await.unwrap() })
        })
        .await)
    }
}

#[cfg(test)]
mod tests {
    use std::cell::RefCell;
    use std::rc::Rc;
    use test_r::test;

    use crate::{boxed, fallible_transaction, infallible_transaction, sync_operation};

    // Not a real test, just verifying that the code compiles
    #[test]
    #[ignore]
    async fn tx_test_1() {
        let log = Rc::new(RefCell::new(Vec::new()));

        let log1 = log.clone();
        let log2 = log.clone();
        let log3 = log.clone();
        let log4 = log.clone();

        let op1 = sync_operation(
            move |input: String| {
                log1.borrow_mut().push(format!("op1 execute {input}"));
                Ok(())
            },
            move |input: String, _| {
                log2.borrow_mut().push(format!("op1 rollback {input}"));
                Ok(())
            },
        );

        let op2 = sync_operation(
            move |_: ()| {
                log3.clone().borrow_mut().push("op2 execute".to_string());
                Err::<(), &str>("op2 error")
            },
            move |_: (), _| {
                log4.clone().borrow_mut().push("op2 rollback".to_string());
                Ok(())
            },
        );

        let result = fallible_transaction(|tx| {
            boxed(async move {
                println!("First we execute op1");
                tx.execute(op1, "hello".to_string()).await?;
                println!("Then execute op2");
                tx.execute(op2, ()).await?;
                println!("Finally compute a result");
                Ok(11)
            })
        })
        .await;

        println!("{log:?}");
        println!("{result:?}");
    }

    // Not a real test, just verifying that the code compiles
    #[test]
    #[ignore]
    async fn tx_test_2() {
        let log = Rc::new(RefCell::new(Vec::new()));

        let log1 = log.clone();
        let log2 = log.clone();
        let log3 = log.clone();
        let log4 = log.clone();

        let op1 = sync_operation(
            move |input: String| {
                log1.borrow_mut().push(format!("op1 execute {input}"));
                Ok::<(), ()>(())
            },
            move |input: String, _| {
                log2.borrow_mut().push(format!("op1 rollback {input}"));
                Ok(())
            },
        );

        let op2 = sync_operation(
            move |_: ()| {
                log3.clone().borrow_mut().push("op2 execute".to_string());
                Err::<(), &str>("op2 error")
            },
            move |_: (), r| {
                log4.clone()
                    .borrow_mut()
                    .push(format!("op2 rollback {r:?}"));
                Ok(())
            },
        );

        let result = infallible_transaction(|tx| {
            boxed(async move {
                println!("First we execute op1");
                tx.execute(op1, "hello".to_string()).await;
                println!("Then execute op2");
                tx.execute(op2, ()).await;
                println!("Finally compute a result");
                11
            })
        })
        .await;

        println!("{log:?}");
        println!("{result:?}");
    }
}

#[cfg(test)]
#[cfg(feature = "macro")]
mod macro_tests {
    use crate::{boxed, fallible_transaction, infallible_transaction};
    use golem_rust_macro::golem_operation;
    use test_r::test;

    mod golem_rust {
        pub use crate::*;
    }

    #[golem_operation(compensation=test_compensation)]
    fn test_operation(input1: u64, input2: f32) -> Result<bool, String> {
        println!("Op input: {input1}, {input2}");
        Ok(true)
    }

    fn test_compensation(_: bool, input1: u64, input2: f32) -> Result<(), String> {
        println!("Compensation input: {input1}, {input2}");
        Ok(())
    }

    #[golem_operation(compensation=test_compensation_2)]
    fn test_operation_2(input1: u64, input2: f32) -> Result<bool, String> {
        println!("Op input: {input1}, {input2}");
        Ok(true)
    }

    fn test_compensation_2(result: bool) -> Result<(), String> {
        println!("Compensation for operation result {result:?}");
        Ok(())
    }

    #[golem_operation(compensation=test_compensation_3)]
    fn test_operation_3(input: String) -> Result<(), String> {
        println!("Op input: {input}");
        Ok(())
    }

    fn test_compensation_3() -> Result<(), String> {
        println!("Compensation for operation, not using any input");
        Ok(())
    }

    #[golem_operation(compensation=test_compensation_4)]
    fn test_operation_4(input: u64) -> Result<(), String> {
        println!("Op input: {input}");
        Ok(())
    }

    fn test_compensation_4(_: (), input: u64) -> Result<(), String> {
        println!("Compensation for operation with single input {input}");
        Ok(())
    }

    // Not a real test, just verifying that the code compiles
    #[test]
    #[ignore]
    async fn tx_test_1() {
        let result = fallible_transaction(|tx| {
            boxed(async move {
                println!("Executing the annotated function as an operation directly");
                tx.test_operation(1, 0.1).await?;
                tx.test_operation_2(1, 0.1).await?;
                tx.test_operation_3("test".to_string()).await?;
                tx.test_operation_4(1).await?;

                Ok(11)
            })
        })
        .await;

        println!("{result:?}");
    }

    // Not a real test, just verifying that the code compiles
    #[test]
    #[ignore]
    async fn tx_test_2() {
        let result = infallible_transaction(|tx| {
            boxed(async move {
                println!("Executing the annotated function as an operation directly");
                let _ = tx.test_operation(1, 0.1).await;
                11
            })
        })
        .await;

        println!("{result:?}");
    }
}