|
1 | | -use std::rc::Weak; |
| 1 | +use crate::connection::Connection; |
2 | 2 | use crate::database::Database; |
3 | | -use crate::pool::Pool; |
| 3 | +use crate::pool::{Pool, PoolConnection, PoolOptions}; |
4 | 4 | use crate::sync::AsyncOnceCell; |
| 5 | +use cfg_if::cfg_if; |
| 6 | +use pin_project_lite::pin_project; |
| 7 | +use std::future::Future; |
| 8 | +use std::ops::{Deref, DerefMut}; |
| 9 | +use std::pin::Pin; |
| 10 | +use std::task::{Context, Poll}; |
5 | 11 |
|
6 | 12 | pub struct TestMasterPool<DB: Database> { |
7 | 13 | inner: AsyncOnceCell<Inner<DB>>, |
8 | 14 | } |
9 | 15 |
|
| 16 | +pub struct TestMasterConnection<DB: Database> { |
| 17 | + conn: PoolConnection<DB>, |
| 18 | + |
| 19 | + #[cfg(feature = "_rt-tokio")] |
| 20 | + handle: tokio::runtime::Handle, |
| 21 | +} |
| 22 | + |
10 | 23 | struct Inner<DB: Database> { |
11 | 24 | pool: Pool<DB>, |
12 | 25 |
|
13 | 26 | #[cfg(feature = "_rt-tokio")] |
14 | | - |
| 27 | + handle: tokio::runtime::Handle, |
15 | 28 | } |
16 | 29 |
|
| 30 | +macro_rules! poll_with_handle( |
| 31 | + ($handle:expr, $fut:expr) => { |
| 32 | + PollWithHandle { |
| 33 | + fut: $fut, |
| 34 | + #[cfg(feature = "_rt-tokio")] |
| 35 | + handle: &$handle, |
| 36 | + #[cfg(not(feature = "_rt-tokio"))] |
| 37 | + _marker: std::marker::PhantomData, |
| 38 | + } |
| 39 | + } |
| 40 | +); |
17 | 41 |
|
18 | 42 | impl<DB: Database> TestMasterPool<DB> { |
| 43 | + pub const fn new() -> Self { |
| 44 | + TestMasterPool { |
| 45 | + inner: AsyncOnceCell::const_new(), |
| 46 | + } |
| 47 | + } |
| 48 | + |
| 49 | + pub async fn connect( |
| 50 | + &self, |
| 51 | + opts: &<DB::Connection as Connection>::Options, |
| 52 | + ) -> crate::Result<TestMasterConnection<DB>> { |
| 53 | + self.inner |
| 54 | + .get_or_try_init::<_, _, crate::Error>(|| { |
| 55 | + let opts = opts.clone(); |
| 56 | + |
| 57 | + async move { |
| 58 | + #[cfg(feature = "_rt-tokio")] |
| 59 | + let handle = spawn_test_runtime(); |
| 60 | + |
| 61 | + // Ensure this pool is linked to our master runtime so it can survive an individual |
| 62 | + // test runtime shutting down. |
| 63 | + let pool = poll_with_handle!( |
| 64 | + handle, |
| 65 | + PoolOptions::new() |
| 66 | + // Tests don't need a master connection for very long |
| 67 | + .max_connections(1) |
| 68 | + .test_before_acquire(false) |
| 69 | + .connect_with(opts) |
| 70 | + ) |
| 71 | + .await?; |
| 72 | + |
| 73 | + Ok(Inner { |
| 74 | + pool, |
| 75 | + #[cfg(feature = "_rt-tokio")] |
| 76 | + handle, |
| 77 | + }) |
| 78 | + } |
| 79 | + }) |
| 80 | + .await? |
| 81 | + .acquire() |
| 82 | + .await |
| 83 | + } |
| 84 | + |
| 85 | + /// # Panics |
| 86 | + /// If [`Self::connect()`] has not already completed successfully. |
| 87 | + pub async fn acquire(&self) -> crate::Result<TestMasterConnection<DB>> { |
| 88 | + self.inner |
| 89 | + .get() |
| 90 | + .expect("`TestMasterPool::connect()` has not been called") |
| 91 | + .acquire() |
| 92 | + .await |
| 93 | + } |
| 94 | +} |
| 95 | + |
| 96 | +impl<DB: Database> Deref for TestMasterConnection<DB> { |
| 97 | + type Target = PoolConnection<DB>; |
| 98 | + |
| 99 | + fn deref(&self) -> &Self::Target { |
| 100 | + &self.conn |
| 101 | + } |
| 102 | +} |
| 103 | + |
| 104 | +impl<DB: Database> DerefMut for TestMasterConnection<DB> { |
| 105 | + fn deref_mut(&mut self) -> &mut Self::Target { |
| 106 | + &mut self.conn |
| 107 | + } |
| 108 | +} |
| 109 | + |
| 110 | +impl<DB: Database> Drop for TestMasterConnection<DB> { |
| 111 | + fn drop(&mut self) { |
| 112 | + cfg_if!( |
| 113 | + if #[cfg(feature = "_rt-tokio")] { |
| 114 | + self.handle.spawn(self.conn.release()); |
| 115 | + } else { |
| 116 | + crate::rt::spawn(self.conn.release()); |
| 117 | + } |
| 118 | + ); |
| 119 | + } |
| 120 | +} |
| 121 | + |
| 122 | +impl<DB: Database> Inner<DB> { |
| 123 | + async fn acquire(&self) -> crate::Result<TestMasterConnection<DB>> { |
| 124 | + Ok(TestMasterConnection { |
| 125 | + // Ostensibly we only need to enter the runtime if the connection isn't already established |
| 126 | + conn: poll_with_handle!(self.handle, self.pool.acquire()).await?, |
| 127 | + #[cfg(feature = "_rt-tokio")] |
| 128 | + handle: self.handle.clone(), |
| 129 | + }) |
| 130 | + } |
| 131 | +} |
| 132 | + |
| 133 | +// It's likely not advisable to hold an `EnterGuard` across an `.await` point, |
| 134 | +// so we need to define an adapter that only enters the alternate runtime when it's polled. |
| 135 | +#[cfg(feature = "_rt-tokio")] |
| 136 | +pin_project! { |
| 137 | + struct PollWithHandle<'a, F> { |
| 138 | + #[pin] |
| 139 | + fut: F, |
| 140 | + handle: &'a tokio::runtime::Handle, |
| 141 | + } |
| 142 | +} |
| 143 | + |
| 144 | +#[cfg(not(feature = "_rt-tokio"))] |
| 145 | +pin_project! { |
| 146 | + struct PollWithHandle<'a, F> { |
| 147 | + #[pin] |
| 148 | + fut: F, |
| 149 | + _marker: std::marker::PhantomData<&'a ()>, |
| 150 | + } |
| 151 | +} |
| 152 | + |
| 153 | +impl<F: Future> Future for PollWithHandle<'_, F> { |
| 154 | + type Output = F::Output; |
| 155 | + |
| 156 | + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { |
| 157 | + let this = self.project(); |
| 158 | + |
| 159 | + #[cfg(feature = "_rt-tokio")] |
| 160 | + let _guard = this.handle.enter(); |
| 161 | + |
| 162 | + this.fut.poll(cx) |
| 163 | + } |
| 164 | +} |
| 165 | + |
| 166 | +#[cfg(feature = "_rt-tokio")] |
| 167 | +fn spawn_test_runtime() -> tokio::runtime::Handle { |
| 168 | + // Instead of forcing the `rt-multi-thread` feature on, |
| 169 | + // we just run a current-thread runtime in a background thread that we ourselves spawn. |
| 170 | + let rt = tokio::runtime::Builder::new_current_thread() |
| 171 | + .name("sqlx-test-master-pool") |
| 172 | + .enable_all() |
| 173 | + .build() |
| 174 | + .expect("failed to spawn master runtime"); |
| 175 | + |
| 176 | + let handle = rt.handle().clone(); |
| 177 | + |
| 178 | + std::thread::Builder::new() |
| 179 | + .name("sqlx-test-master-pool".into()) |
| 180 | + .spawn(move || rt.block_on(std::future::pending::<()>())) |
| 181 | + .expect("failed to spawn thread for master runtime"); |
19 | 182 |
|
| 183 | + handle |
20 | 184 | } |
0 commit comments