Skip to main content

datasources/
unit_of_work_transaction.rs

1use crate::connection::Connection;
2use crate::datasource_error::DatasourceError;
3use crate::i_unit_of_work::IUnitOfWork;
4use crate::result::Result;
5use crate::transaction::Transaction;
6use async_trait::async_trait;
7use std::collections::HashSet;
8use std::sync::{Arc, LazyLock, Mutex};
9
10static OPEN: LazyLock<Mutex<HashSet<usize>>> = LazyLock::new(|| Mutex::new(HashSet::new()));
11
12enum Phase {
13    New,
14    Joined,
15    Owner(Arc<dyn Transaction>),
16}
17
18pub struct UnitOfWorkTransaction {
19    connection: Arc<dyn Connection>,
20    phase: Phase,
21}
22
23impl UnitOfWorkTransaction {
24    pub fn new(connection: Arc<dyn Connection>) -> Self {
25        Self {
26            connection,
27            phase: Phase::New,
28        }
29    }
30
31    fn release(&self) {
32        OPEN.lock()
33            .expect("transactions")
34            .remove(&data_ptr(&self.connection));
35    }
36}
37
38#[async_trait]
39impl IUnitOfWork for UnitOfWorkTransaction {
40    async fn begin(&mut self) -> Result<()> {
41        if !matches!(self.phase, Phase::New) {
42            return Err(DatasourceError::message("Unit of work already begun"));
43        }
44        let key = data_ptr(&self.connection);
45        let owner = {
46            let mut open = OPEN.lock().expect("transactions");
47            open.insert(key)
48        };
49        if !owner {
50            self.phase = Phase::Joined;
51            return Ok(());
52        }
53        match self.connection.start_transaction(None).await {
54            Ok(txn) => {
55                self.phase = Phase::Owner(txn);
56                Ok(())
57            }
58            Err(error) => {
59                OPEN.lock().expect("transactions").remove(&key);
60                Err(error)
61            }
62        }
63    }
64
65    async fn commit(&mut self) -> Result<()> {
66        match std::mem::replace(&mut self.phase, Phase::Joined) {
67            Phase::New => {
68                self.phase = Phase::New;
69                Err(DatasourceError::message("Unit of work has not begun"))
70            }
71            Phase::Joined => Ok(()),
72            Phase::Owner(txn) => match txn.commit().await {
73                Ok(()) => {
74                    self.release();
75                    Ok(())
76                }
77                Err(error) => {
78                    self.phase = Phase::Owner(txn);
79                    Err(error)
80                }
81            },
82        }
83    }
84
85    async fn abort(&mut self) -> Result<()> {
86        match std::mem::replace(&mut self.phase, Phase::Joined) {
87            Phase::New => {
88                self.phase = Phase::New;
89                Err(DatasourceError::message("Unit of work has not begun"))
90            }
91            Phase::Joined => Ok(()),
92            Phase::Owner(txn) => {
93                let outcome = txn.abort().await;
94                self.release();
95                outcome
96            }
97        }
98    }
99}
100
101fn data_ptr(connection: &Arc<dyn Connection>) -> usize {
102    &**connection as *const dyn Connection as *const () as usize
103}