1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156

/// An auxiliary transaction that gets co-committed with the STM transaction,
/// unless the STM transaction is aborted or has to be retried, in which case
/// the auxiliary transaction is rolled back.
///
/// The auxiliary transaction can also signal that its is unable to be committed,
/// in which case the whole atomic transaction will be retried.
///
/// The database is not expected to return an error here, because of how `Transaction::commit` works.
/// If there is a failure that needs to be surfaced, at the moment the database would have to buffer
/// that error and return it on a subsequent operation, mapped to an `StmError::Abort`.
pub trait Aux {
    /// Commit the auxiliary transaction if the STM transaction did not detect any errors.
    /// The STM transaction is checked first, because committing the database involves IO
    /// and is expected to be slower.
    ///
    /// Return `false` if there are write conflicts in the persistent database itself,
    /// to cause a complete retry for the whole transaction.
    fn commit(self) -> bool;
    /// Rollback the auxiliary transaction if the STM transaction was aborted, or it's going to be retried.
    fn rollback(self);
}

/// Empty implementation for when we are not using any auxiliary transaction.
pub(crate) struct NoAux;

impl Aux for NoAux {
    fn commit(self) -> bool {
        true
    }
    fn rollback(self) {}
}

#[cfg(test)]
mod test {
    use std::{
        sync::{Arc, Mutex},
        thread,
    };

    use crate::auxtx::*;
    use crate::{
        abort, atomically, atomically_aux, atomically_or_err_aux, retry, test::TestError, TVar,
    };

    #[derive(Clone)]
    struct TestAuxDb {
        counter: Arc<Mutex<i32>>,
    }
    struct TestAuxTx<'a> {
        db: &'a TestAuxDb,
        counter: i32,
        finished: bool,
    }

    impl TestAuxDb {
        fn begin(&self) -> TestAuxTx {
            let guard = self.counter.lock().unwrap();
            TestAuxTx {
                db: self,
                counter: *guard,
                finished: false,
            }
        }
    }

    impl<'a> Aux for TestAuxTx<'a> {
        fn commit(mut self) -> bool {
            let mut guard = self.db.counter.lock().unwrap();
            *guard = self.counter;
            self.finished = true;
            true
        }
        fn rollback(mut self) {
            self.finished = true;
        }
    }

    impl Drop for TestAuxTx<'_> {
        fn drop(&mut self) {
            if !self.finished && !thread::panicking() {
                panic!("Transaction prematurely dropped. Must call `.commit()` or `.rollback()`.");
            }
        }
    }

    impl TestAuxDb {
        fn new() -> TestAuxDb {
            TestAuxDb {
                counter: Arc::new(Mutex::new(0)),
            }
        }

        fn counter(&self) -> i32 {
            *self.counter.lock().unwrap()
        }
    }

    #[tokio::test]
    async fn aux_commit_rollback() {
        let db = TestAuxDb::new();

        atomically_or_err_aux(
            || db.begin(),
            |atx| {
                atx.counter = 1;
                abort(TestError)?;
                Ok(())
            },
        )
        .await
        .expect_err("Should be aborted");

        assert_eq!(db.counter(), 0);

        atomically_aux(
            || db.begin(),
            |atx| {
                atx.counter = 1;
                Ok(())
            },
        )
        .await;

        assert_eq!(db.counter(), 1);

        let ta = TVar::new(42);
        let dbc = db.clone();
        let tac = ta.clone();
        let (sender, mut receiver) = tokio::sync::mpsc::unbounded_channel();
        let handle = tokio::spawn(async move {
            let _ = atomically_aux(
                || dbc.begin(),
                |atx| {
                    let a = tac.read()?;
                    atx.counter = atx.counter + *a;
                    if *a == 42 {
                        // Signal that we are entering the retry.
                        sender.send(()).unwrap();
                        retry()?;
                    }
                    Ok(())
                },
            )
            .await;
        });

        let _ = receiver.recv().await;
        assert_eq!(db.counter(), 1);

        // Writing a value to `ta` will trigger the retry.
        atomically(|| ta.write(10)).await;
        handle.await.unwrap();

        assert_eq!(db.counter(), 11);
    }
}