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
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
//! Pooled connection wrapper: struct, accessors, RLS cleanup, transaction control,
//! COPY export, pipeline, LISTEN/NOTIFY delegation, and Drop.
use super::churn::{decrement_active_count_saturating, pool_churn_record_destroy};
use super::lifecycle::{PgPoolInner, execute_simple_with_timeout};
use crate::driver::{PgConnection, PgError, PgResult};
use std::sync::Arc;
use std::sync::atomic::Ordering;
use std::time::Instant;
/// A pooled connection with creation timestamp for idle tracking.
pub(super) struct PooledConn {
pub(super) conn: PgConnection,
pub(super) created_at: Instant,
pub(super) last_used: Instant,
}
/// A pooled connection handle.
///
/// Use [`PooledConnection::release`] for deterministic reset+return behavior.
/// If dropped without `release()`, the pool performs best-effort bounded async
/// cleanup; on any uncertainty it destroys the connection (fail-closed).
pub struct PooledConnection {
pub(super) conn: Option<PgConnection>,
pub(super) pool: Arc<PgPoolInner>,
pub(super) rls_dirty: bool,
pub(super) created_at: Instant,
}
impl PooledConnection {
/// Get a reference to the underlying connection, returning an error
/// if the connection has already been released.
pub(super) fn conn_ref(&self) -> PgResult<&PgConnection> {
self.conn
.as_ref()
.ok_or_else(|| PgError::Connection("Connection already released back to pool".into()))
}
/// Get a mutable reference to the underlying connection, returning an error
/// if the connection has already been released.
pub(super) fn conn_mut(&mut self) -> PgResult<&mut PgConnection> {
self.conn
.as_mut()
.ok_or_else(|| PgError::Connection("Connection already released back to pool".into()))
}
/// Get a shared reference to the underlying connection.
///
/// Returns an error if the connection has already been released.
pub fn get(&self) -> PgResult<&PgConnection> {
self.conn_ref()
}
/// Get a mutable reference to the underlying connection.
///
/// Returns an error if the connection has already been released.
pub fn get_mut(&mut self) -> PgResult<&mut PgConnection> {
self.conn_mut()
}
/// Get a token to cancel the currently running query.
pub fn cancel_token(&self) -> PgResult<crate::driver::CancelToken> {
let conn = self.conn_ref()?;
let (process_id, secret_key_bytes) = conn.get_cancel_key_bytes();
Ok(crate::driver::CancelToken {
host: self.pool.config.host.clone(),
port: self.pool.config.port,
process_id,
secret_key_bytes: secret_key_bytes.to_vec(),
})
}
/// Deterministic connection cleanup and pool return.
///
/// This is the **correct** way to return a connection to the pool.
/// COMMITs the transaction (which auto-resets transaction-local RLS
/// session variables) and returns the connection to the pool with
/// prepared statement caches intact.
///
/// If cleanup fails, the connection is destroyed (not returned to pool).
///
/// # Usage
/// ```ignore
/// let mut conn = pool.acquire_with_rls(ctx).await?;
/// let result = conn.fetch_all_cached(&cmd).await;
/// conn.release().await; // COMMIT + return to pool
/// result
/// ```
pub async fn release(mut self) {
if let Some(mut conn) = self.conn.take() {
if conn.is_io_desynced() {
tracing::warn!(
host = %self.pool.config.host,
port = self.pool.config.port,
user = %self.pool.config.user,
db = %self.pool.config.database,
"pool_release_desynced: dropping connection due to prior I/O/protocol desync"
);
decrement_active_count_saturating(&self.pool.active_count);
self.pool.semaphore.add_permits(1);
pool_churn_record_destroy(&self.pool.config, "release_desynced");
return;
}
// COMMIT the transaction opened by acquire_with_rls.
// Transaction-local set_config values auto-reset on COMMIT,
// so no explicit RLS cleanup is needed.
// Prepared statements survive — they are NOT transaction-scoped.
let reset_timeout = self.pool.config.connect_timeout;
if let Err(e) = execute_simple_with_timeout(
&mut conn,
crate::driver::rls::reset_sql(),
reset_timeout,
"pool release reset/COMMIT",
)
.await
{
tracing::error!(
host = %self.pool.config.host,
port = self.pool.config.port,
user = %self.pool.config.user,
db = %self.pool.config.database,
timeout_ms = reset_timeout.as_millis() as u64,
error = %e,
"pool_release_failed: reset/COMMIT failed; dropping connection to prevent state leak"
);
decrement_active_count_saturating(&self.pool.active_count);
self.pool.semaphore.add_permits(1);
pool_churn_record_destroy(&self.pool.config, "release_reset_failed");
return; // Connection destroyed — not returned to pool
}
self.pool.return_connection(conn, self.created_at).await;
}
}
// ==================== TRANSACTION CONTROL ====================
/// Begin an explicit transaction on this pooled connection.
///
/// Use this when you need multi-statement atomicity beyond the
/// implicit transaction created by `acquire_with_rls()`.
///
/// # Example
/// ```ignore
/// let mut conn = pool.acquire_with_rls(ctx).await?;
/// conn.begin().await?;
/// conn.execute(&insert1).await?;
/// conn.execute(&insert2).await?;
/// conn.commit().await?;
/// conn.release().await;
/// ```
pub async fn begin(&mut self) -> PgResult<()> {
self.conn_mut()?.begin_transaction().await
}
/// Commit the current transaction.
/// Makes all changes since `begin()` permanent.
pub async fn commit(&mut self) -> PgResult<()> {
self.conn_mut()?.commit().await
}
/// Rollback the current transaction.
/// Discards all changes since `begin()`.
pub async fn rollback(&mut self) -> PgResult<()> {
self.conn_mut()?.rollback().await
}
/// Create a named savepoint within the current transaction.
/// Use `rollback_to()` to return to this savepoint.
pub async fn savepoint(&mut self, name: &str) -> PgResult<()> {
self.conn_mut()?.savepoint(name).await
}
/// Rollback to a previously created savepoint.
/// Discards changes since the savepoint, but keeps the transaction open.
pub async fn rollback_to(&mut self, name: &str) -> PgResult<()> {
self.conn_mut()?.rollback_to(name).await
}
/// Release a savepoint (free resources).
/// After release, the savepoint cannot be rolled back to.
pub async fn release_savepoint(&mut self, name: &str) -> PgResult<()> {
self.conn_mut()?.release_savepoint(name).await
}
/// Execute multiple QAIL commands in a single PG pipeline round-trip.
///
/// Sends all queries as Parse+Bind+Execute in one write, receives all
/// responses in one read. Returns raw column data per query per row.
///
/// This is the fastest path for batch operations — amortizes TCP
/// overhead across N queries into a single syscall pair.
pub async fn pipeline_execute_rows_ast(
&mut self,
cmds: &[qail_core::ast::Qail],
) -> PgResult<Vec<Vec<Vec<Option<Vec<u8>>>>>> {
let conn = self.conn_mut()?;
conn.pipeline_execute_rows_ast(cmds).await
}
/// Run `EXPLAIN (FORMAT JSON)` on a Qail command and return cost estimates.
///
/// Uses `simple_query` under the hood — no additional round-trips beyond
/// the single EXPLAIN statement. Returns `None` if parsing fails or
/// the EXPLAIN output is unexpected.
pub async fn explain_estimate(
&mut self,
cmd: &qail_core::ast::Qail,
) -> PgResult<Option<crate::driver::explain::ExplainEstimate>> {
use qail_core::transpiler::ToSql;
let sql = cmd.to_sql();
let explain_sql = format!("EXPLAIN (FORMAT JSON) {}", sql);
let rows = self.conn_mut()?.simple_query(&explain_sql).await?;
// PostgreSQL returns the JSON plan as a single text column across one or more rows
let mut json_output = String::new();
for row in &rows {
if let Some(Some(val)) = row.columns.first()
&& let Ok(text) = std::str::from_utf8(val)
{
json_output.push_str(text);
}
}
Ok(crate::driver::explain::parse_explain_json(&json_output))
}
// ─── LISTEN / NOTIFY delegation ─────────────────────────────────
/// Subscribe to a PostgreSQL notification channel.
///
/// Delegates to [`PgConnection::listen`].
pub async fn listen(&mut self, channel: &str) -> PgResult<()> {
self.conn_mut()?.listen(channel).await
}
/// Unsubscribe from a PostgreSQL notification channel.
///
/// Delegates to [`PgConnection::unlisten`].
pub async fn unlisten(&mut self, channel: &str) -> PgResult<()> {
self.conn_mut()?.unlisten(channel).await
}
/// Unsubscribe from all notification channels.
///
/// Delegates to [`PgConnection::unlisten_all`].
pub async fn unlisten_all(&mut self) -> PgResult<()> {
self.conn_mut()?.unlisten_all().await
}
/// Wait for the next notification, blocking until one arrives.
///
/// Delegates to [`PgConnection::recv_notification`].
/// Useful for dedicated LISTEN connections in background tasks.
pub async fn recv_notification(
&mut self,
) -> PgResult<crate::driver::notification::Notification> {
self.conn_mut()?.recv_notification().await
}
}
impl Drop for PooledConnection {
fn drop(&mut self) {
if let Some(mut conn) = self.conn.take() {
// Safety net: connection was NOT released via `release()`.
// Best-effort strategy:
// 1) If connection is already desynced, destroy immediately.
// 2) Else, queue bounded async reset+return cleanup.
// 3) If cleanup queue/runtime unavailable, destroy.
//
// This preserves security (fail-closed) while reducing churn under
// accidental early-returns in handler code.
tracing::warn!(
host = %self.pool.config.host,
port = self.pool.config.port,
user = %self.pool.config.user,
db = %self.pool.config.database,
rls_dirty = self.rls_dirty,
"pool_connection_leaked: dropped without release()"
);
if conn.is_io_desynced() {
tracing::warn!(
host = %self.pool.config.host,
port = self.pool.config.port,
user = %self.pool.config.user,
db = %self.pool.config.database,
"pool_connection_leaked_desynced: destroying immediately"
);
decrement_active_count_saturating(&self.pool.active_count);
self.pool.semaphore.add_permits(1);
pool_churn_record_destroy(&self.pool.config, "dropped_without_release_desynced");
return;
}
let mut inflight = self.pool.leaked_cleanup_inflight.load(Ordering::Relaxed);
let max_inflight = self.pool.config.leaked_cleanup_queue;
loop {
if inflight >= max_inflight {
tracing::warn!(
host = %self.pool.config.host,
port = self.pool.config.port,
user = %self.pool.config.user,
db = %self.pool.config.database,
max_inflight,
"pool_connection_leaked_cleanup_queue_full: destroying connection"
);
decrement_active_count_saturating(&self.pool.active_count);
self.pool.semaphore.add_permits(1);
pool_churn_record_destroy(
&self.pool.config,
"dropped_without_release_cleanup_queue_full",
);
return;
}
match self.pool.leaked_cleanup_inflight.compare_exchange_weak(
inflight,
inflight + 1,
Ordering::AcqRel,
Ordering::Relaxed,
) {
Ok(_) => break,
Err(actual) => inflight = actual,
}
}
let pool = self.pool.clone();
let created_at = self.created_at;
let reset_timeout = pool.config.connect_timeout;
match tokio::runtime::Handle::try_current() {
Ok(handle) => {
handle.spawn(async move {
let cleanup_ok = execute_simple_with_timeout(
&mut conn,
crate::driver::rls::reset_sql(),
reset_timeout,
"pool leaked cleanup reset/COMMIT",
)
.await
.is_ok();
if cleanup_ok && !conn.is_io_desynced() {
pool.return_connection(conn, created_at).await;
} else {
tracing::warn!(
host = %pool.config.host,
port = pool.config.port,
user = %pool.config.user,
db = %pool.config.database,
timeout_ms = reset_timeout.as_millis() as u64,
"pool_connection_leaked_cleanup_failed: destroying connection"
);
decrement_active_count_saturating(&pool.active_count);
pool.semaphore.add_permits(1);
pool_churn_record_destroy(
&pool.config,
"dropped_without_release_cleanup_failed",
);
}
pool.leaked_cleanup_inflight.fetch_sub(1, Ordering::AcqRel);
});
}
Err(_) => {
pool.leaked_cleanup_inflight.fetch_sub(1, Ordering::AcqRel);
tracing::warn!(
host = %pool.config.host,
port = pool.config.port,
user = %pool.config.user,
db = %pool.config.database,
"pool_connection_leaked_no_runtime: destroying connection"
);
decrement_active_count_saturating(&pool.active_count);
pool.semaphore.add_permits(1);
pool_churn_record_destroy(&pool.config, "dropped_without_release_no_runtime");
}
}
}
}
}