rust_query/
async_db.rs

1use std::{
2    future,
3    sync::{Arc, Mutex},
4    task::{Poll, Waker},
5};
6
7use crate::{Database, Transaction, migrate::Schema};
8
9/// This is an async wrapper for [Database].
10///
11/// Transaction methods on this type are identical to [Database], but they spawn
12/// an async task and return a future to await the transaction completion.
13/// Because of this difference, the transaction closures must be `'static`.
14///
15/// You can easily implement [DatabaseAsync] manually using `tokio::task::spawn_blocking`,
16/// but this wrapper is a little bit more efficient while also being runtime agnostic.
17pub struct DatabaseAsync<S> {
18    inner: Arc<Database<S>>,
19}
20
21impl<S> Clone for DatabaseAsync<S> {
22    fn clone(&self) -> Self {
23        Self {
24            inner: self.inner.clone(),
25        }
26    }
27}
28
29impl<S: 'static + Send + Sync + Schema> DatabaseAsync<S> {
30    /// Create an async wrapper for the [Database].
31    ///
32    /// The database is wrapped in an [Arc] as it needs to be shared with any thread
33    /// executing a transaction. These threads can live longer than the future that
34    /// started the transaction.
35    ///
36    /// By accepting an [Arc], you can keep your own clone of the [Arc] and use
37    /// the database synchronously and asynchronously at the same time!
38    pub fn new(db: Arc<Database<S>>) -> Self {
39        DatabaseAsync { inner: db }
40    }
41
42    #[doc = include_str!("database/transaction.md")]
43    pub async fn transaction<R: 'static + Send>(
44        &self,
45        f: impl 'static + Send + FnOnce(&'static Transaction<S>) -> R,
46    ) -> R {
47        let db = self.inner.clone();
48        async_run(move || db.transaction_local(f))
49            .await
50            .unwrap_or_else(|e| e.to_panic())
51    }
52
53    #[doc = include_str!("database/transaction_mut.md")]
54    pub async fn transaction_mut<O: 'static + Send, E: 'static + Send>(
55        &self,
56        f: impl 'static + Send + FnOnce(&'static mut Transaction<S>) -> Result<O, E>,
57    ) -> Result<O, E> {
58        let db = self.inner.clone();
59        async_run(move || db.transaction_mut_local(f))
60            .await
61            .unwrap_or_else(|e| e.to_panic())
62    }
63
64    #[doc = include_str!("database/transaction_mut_ok.md")]
65    pub async fn transaction_mut_ok<R: 'static + Send>(
66        &self,
67        f: impl 'static + Send + FnOnce(&'static mut Transaction<S>) -> R,
68    ) -> R {
69        self.transaction_mut(|txn| Ok::<R, std::convert::Infallible>(f(txn)))
70            .await
71            .unwrap()
72    }
73}
74
75async fn async_run<R: 'static + Send>(f: impl 'static + Send + FnOnce() -> R) -> R {
76    pub struct WakeOnDrop {
77        waker: Mutex<Waker>,
78    }
79
80    impl Drop for WakeOnDrop {
81        #[cfg_attr(false, mutants::skip)] // mutating this will make the test hang
82        fn drop(&mut self) {
83            self.waker.lock().unwrap().wake_by_ref();
84        }
85    }
86
87    // Initally we use a noop waker, because we will override it anyway.
88    let wake_on_drop = Arc::new(WakeOnDrop {
89        waker: Mutex::new(Waker::noop().clone()),
90    });
91    let weak = Arc::downgrade(&wake_on_drop);
92
93    let handle = std::thread::spawn(move || {
94        // waker will be called when thread finishes, even with panic.
95        let _wake_on_drop = wake_on_drop;
96        f()
97    });
98
99    // asynchonously wait for the thread to finish
100    future::poll_fn(|cx| {
101        if let Some(wake_on_drop) = weak.upgrade() {
102            wake_on_drop.waker.lock().unwrap().clone_from(cx.waker());
103            Poll::Pending
104        } else {
105            Poll::Ready(())
106        }
107    })
108    .await;
109
110    // we know that the thread is finished, so we block on it
111    match handle.join() {
112        Ok(val) => val,
113        Err(err) => std::panic::resume_unwind(err),
114    }
115}
rust_query/async_db.rs

rust_query/
async_db.rs
