Skip to main content

golem_rust/
checkpoint.rs

1// Copyright 2024-2026 Golem Cloud
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use crate::bindings::golem::api::host::{OplogIndex, get_oplog_index, set_oplog_index};
16
17/// A checkpoint that captures the current oplog index and can revert execution back to it.
18pub struct Checkpoint {
19    oplog_index: OplogIndex,
20}
21
22impl Default for Checkpoint {
23    fn default() -> Self {
24        Self::new()
25    }
26}
27
28impl Checkpoint {
29    /// Creates a new checkpoint at the current oplog index.
30    pub fn new() -> Self {
31        Self {
32            oplog_index: get_oplog_index(),
33        }
34    }
35
36    /// Reverts execution back to the checkpoint's oplog index.
37    pub fn revert(&self) -> ! {
38        set_oplog_index(self.oplog_index);
39        unreachable!()
40    }
41
42    /// Runs the given function, reverting to the checkpoint on error.
43    pub fn run_or_revert<T, E>(&self, f: impl FnOnce() -> Result<T, E>) -> T {
44        match f() {
45            Ok(value) => value,
46            Err(_) => self.revert(),
47        }
48    }
49
50    /// Runs the given async function, reverting to the checkpoint on error.
51    pub async fn run_or_revert_async<T, E, F: std::future::Future<Output = Result<T, E>>>(
52        &self,
53        f: impl FnOnce() -> F,
54    ) -> T {
55        match f().await {
56            Ok(value) => value,
57            Err(_) => self.revert(),
58        }
59    }
60
61    /// Reverts to the checkpoint if the condition is false.
62    pub fn assert_or_revert(&self, condition: bool) {
63        if !condition {
64            self.revert();
65        }
66    }
67}
68
69/// Extension trait for `Result` that provides checkpoint-based revert on error.
70pub trait CheckpointResultExt<T, E> {
71    /// Returns the `Ok` value, or reverts to the checkpoint on `Err`.
72    fn unwrap_or_revert(self, checkpoint: &Checkpoint) -> T;
73
74    /// Returns the `Ok` value, or logs the message and reverts to the checkpoint on `Err`.
75    fn expect_or_revert(self, checkpoint: &Checkpoint, msg: &str) -> T;
76}
77
78impl<T, E> CheckpointResultExt<T, E> for Result<T, E> {
79    fn unwrap_or_revert(self, checkpoint: &Checkpoint) -> T {
80        match self {
81            Ok(value) => value,
82            Err(_) => checkpoint.revert(),
83        }
84    }
85
86    fn expect_or_revert(self, checkpoint: &Checkpoint, msg: &str) -> T {
87        match self {
88            Ok(value) => value,
89            Err(_) => {
90                eprintln!("{msg}");
91                checkpoint.revert();
92            }
93        }
94    }
95}
96
97/// Extension trait for `Option` that provides checkpoint-based revert on `None`.
98pub trait CheckpointOptionExt<T> {
99    /// Returns the `Some` value, or reverts to the checkpoint on `None`.
100    fn unwrap_or_revert(self, checkpoint: &Checkpoint) -> T;
101
102    /// Returns the `Some` value, or logs the message and reverts to the checkpoint on `None`.
103    fn expect_or_revert(self, checkpoint: &Checkpoint, msg: &str) -> T;
104}
105
106impl<T> CheckpointOptionExt<T> for Option<T> {
107    fn unwrap_or_revert(self, checkpoint: &Checkpoint) -> T {
108        match self {
109            Some(value) => value,
110            None => checkpoint.revert(),
111        }
112    }
113
114    fn expect_or_revert(self, checkpoint: &Checkpoint, msg: &str) -> T {
115        match self {
116            Some(value) => value,
117            None => {
118                eprintln!("{msg}");
119                checkpoint.revert();
120            }
121        }
122    }
123}
124
125/// Creates a checkpoint, runs the given function, and reverts on error.
126pub fn with_checkpoint<T, E>(f: impl FnOnce(&Checkpoint) -> Result<T, E>) -> T {
127    let checkpoint = Checkpoint::new();
128    match f(&checkpoint) {
129        Ok(value) => value,
130        Err(_) => checkpoint.revert(),
131    }
132}
133
134/// Creates a checkpoint, runs the given async function, and reverts on error.
135pub async fn with_checkpoint_async<T, E, F: std::future::Future<Output = Result<T, E>>>(
136    f: impl FnOnce(&Checkpoint) -> F,
137) -> T {
138    let checkpoint = Checkpoint::new();
139    match f(&checkpoint).await {
140        Ok(value) => value,
141        Err(_) => checkpoint.revert(),
142    }
143}