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
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
//! Async lifecycle API for [`TestCluster`].
//!
//! Hosts the `async-api` feature surface: asynchronous startup on the
//! caller's Tokio runtime and explicit asynchronous shutdown.
use tracing::info_span;
use super::{
TestCluster,
guard::ClusterGuard,
handle::ClusterHandle,
runtime_mode::ClusterRuntime,
shutdown,
startup::{cache_config_from_bootstrap, start_postgres_async},
};
use crate::{
bootstrap_for_tests,
env::ScopedEnv,
error::BootstrapResult,
observability::LOG_TARGET,
};
enum StopAsyncPath {
WorkerManaged,
InProcess(Box<postgresql_embedded::PostgreSQL>),
Noop,
}
impl TestCluster {
/// Boots a `PostgreSQL` instance asynchronously for use in `#[tokio::test]` contexts.
///
/// Unlike [`TestCluster::new`], this constructor does not create its own Tokio runtime.
/// Instead, it runs on the caller's async runtime, making it safe to call from within
/// `#[tokio::test]` and other async contexts.
///
/// **Important:** Clusters created with `start_async()` should be shut down explicitly
/// using [`stop_async()`](Self::stop_async). The `Drop` implementation will attempt
/// best-effort cleanup but may not succeed if the runtime is no longer available.
///
/// # Errors
///
/// Returns an error if the bootstrap configuration cannot be prepared or if starting
/// the embedded cluster fails.
///
/// # Examples
///
/// ```no_run
/// use pg_embedded_setup_unpriv::TestCluster;
///
/// #[tokio::test]
/// async fn test_with_embedded_postgres() -> pg_embedded_setup_unpriv::BootstrapResult<()> {
/// let cluster = TestCluster::start_async().await?;
/// let url = cluster.settings().url("my_database");
/// // ... async database operations ...
/// cluster.stop_async().await?;
/// Ok(())
/// }
/// ```
pub async fn start_async() -> BootstrapResult<Self> {
let (handle, guard) = Self::start_async_split().await?;
Ok(Self { handle, guard })
}
/// Boots a `PostgreSQL` instance asynchronously and returns a separate handle and guard.
///
/// This is the async equivalent of [`new_split()`](Self::new_split).
///
/// # Returns
///
/// A tuple of:
/// - [`ClusterHandle`]: `Send + Sync` handle for accessing the cluster
/// - [`ClusterGuard`]: `!Send` guard managing shutdown and environment
///
/// # Errors
///
/// Returns an error if the bootstrap configuration cannot be prepared or if
/// starting the embedded cluster fails.
pub async fn start_async_split() -> BootstrapResult<(ClusterHandle, ClusterGuard)> {
use tracing::Instrument;
let span = info_span!(target: LOG_TARGET, "test_cluster", async_mode = true);
// Sync bootstrap preparation (no await needed).
// Resolve cache directory BEFORE applying test environment.
// Otherwise, the test sandbox's XDG_CACHE_HOME would be used.
let initial_bootstrap = bootstrap_for_tests()?;
let cache_config = cache_config_from_bootstrap(&initial_bootstrap);
let env_vars = initial_bootstrap.environment.to_env();
let env_guard = ScopedEnv::apply(&env_vars);
// Async postgres startup, instrumented with the span.
// Box::pin to avoid large future on the stack.
let outcome = Box::pin(start_postgres_async(
initial_bootstrap,
&env_vars,
&cache_config,
))
.instrument(span.clone())
.await?;
let handle = ClusterHandle::new(outcome.bootstrap.clone());
let guard = ClusterGuard {
runtime: ClusterRuntime::Async,
postgres: outcome.postgres,
bootstrap: outcome.bootstrap,
is_managed_via_worker: outcome.is_managed_via_worker,
env_vars,
worker_guard: None,
_env_guard: env_guard,
_cluster_span: span,
};
Ok((handle, guard))
}
fn stop_async_path(&mut self) -> StopAsyncPath {
if self.guard.is_managed_via_worker {
StopAsyncPath::WorkerManaged
} else if let Some(postgres) = self.guard.postgres.take() {
StopAsyncPath::InProcess(Box::new(postgres))
} else {
StopAsyncPath::Noop
}
}
/// Explicitly shuts down an async cluster.
///
/// This method should be called for clusters created with [`start_async()`](Self::start_async)
/// to ensure proper cleanup. It consumes `self` to prevent the `Drop` implementation from
/// attempting duplicate shutdown.
///
/// For worker-managed clusters (root privileges), the worker subprocess is invoked
/// synchronously via `spawn_blocking`.
///
/// # Errors
///
/// Returns an error if the shutdown operation fails. The cluster resources are released
/// regardless of whether shutdown succeeds.
///
/// # Examples
///
/// ```no_run
/// use pg_embedded_setup_unpriv::TestCluster;
///
/// #[tokio::test]
/// async fn test_explicit_shutdown() -> pg_embedded_setup_unpriv::BootstrapResult<()> {
/// let cluster = TestCluster::start_async().await?;
/// // ... use cluster ...
/// cluster.stop_async().await?;
/// Ok(())
/// }
/// ```
pub async fn stop_async(mut self) -> BootstrapResult<()> {
let context = shutdown::stop_context(self.handle.settings());
shutdown::log_async_stop(&context, self.guard.is_managed_via_worker);
match self.stop_async_path() {
StopAsyncPath::WorkerManaged => {
shutdown::stop_worker_managed_async(
&self.guard.bootstrap,
&self.guard.env_vars,
&context,
)
.await
}
StopAsyncPath::InProcess(postgres) => {
let cleanup = shutdown::InProcessCleanup {
cleanup_mode: self.guard.bootstrap.cleanup_mode,
settings: &self.guard.bootstrap.settings,
context: &context,
};
shutdown::stop_in_process_async(
*postgres,
self.guard.bootstrap.shutdown_timeout,
cleanup,
)
.await
}
StopAsyncPath::Noop => Ok(()),
}
}
}