Skip to content

Commit 681cfb7

Browse files
committed
pool: fix panic when using callbacks
add regression test
1 parent 339e058 commit 681cfb7

File tree

2 files changed

+136
-2
lines changed

2 files changed

+136
-2
lines changed

sqlx-core/src/pool/connection.rs

+4-2
Original file line numberDiff line numberDiff line change
@@ -325,8 +325,10 @@ impl<DB: Database> Floating<DB, Idle<DB>> {
325325
let now = Instant::now();
326326

327327
PoolConnectionMetadata {
328-
age: self.created_at.duration_since(now),
329-
idle_for: self.idle_since.duration_since(now),
328+
// NOTE: the receiver is the later `Instant` and the arg is the earlier
329+
// https://github.com/launchbadge/sqlx/issues/1912
330+
age: now.saturating_duration_since(self.created_at),
331+
idle_for: now.saturating_duration_since(self.idle_since),
330332
}
331333
}
332334
}

tests/any/pool.rs

+132
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,6 @@
11
use sqlx::any::AnyPoolOptions;
2+
use sqlx::Executor;
3+
use std::sync::atomic::AtomicI32;
24
use std::sync::{
35
atomic::{AtomicUsize, Ordering},
46
Arc,
@@ -64,3 +66,133 @@ async fn pool_should_be_returned_failed_transactions() -> anyhow::Result<()> {
6466

6567
Ok(())
6668
}
69+
70+
#[sqlx_macros::test]
71+
async fn test_pool_callbacks() -> anyhow::Result<()> {
72+
sqlx_test::setup_if_needed();
73+
74+
#[derive(sqlx::FromRow, Debug, PartialEq, Eq)]
75+
struct ConnStats {
76+
id: i32,
77+
before_acquire_calls: i32,
78+
after_release_calls: i32,
79+
}
80+
81+
let current_id = AtomicI32::new(0);
82+
83+
let pool = AnyPoolOptions::new()
84+
.max_connections(1)
85+
.acquire_timeout(Duration::from_secs(5))
86+
.after_connect(move |conn, meta| {
87+
assert_eq!(meta.age, Duration::ZERO);
88+
assert_eq!(meta.idle_for, Duration::ZERO);
89+
90+
let id = current_id.fetch_add(1, Ordering::AcqRel);
91+
92+
Box::pin(async move {
93+
let statement = format!(
94+
// language=SQL
95+
r#"
96+
CREATE TEMPORARY TABLE conn_stats(
97+
id int primary key,
98+
before_acquire_calls int default 0,
99+
after_release_calls int default 0
100+
);
101+
INSERT INTO conn_stats(id) VALUES ({});
102+
"#,
103+
// Until we have generalized bind parameters
104+
id
105+
);
106+
107+
conn.execute(&statement[..]).await?;
108+
Ok(())
109+
})
110+
})
111+
.before_acquire(|conn, meta| {
112+
// `age` and `idle_for` should both be nonzero
113+
assert_ne!(meta.age, Duration::ZERO);
114+
assert_ne!(meta.idle_for, Duration::ZERO);
115+
116+
Box::pin(async move {
117+
// MySQL and MariaDB don't support UPDATE ... RETURNING
118+
sqlx::query(
119+
r#"
120+
UPDATE conn_stats
121+
SET before_acquire_calls = before_acquire_calls + 1
122+
"#,
123+
)
124+
.execute(&mut *conn)
125+
.await?;
126+
127+
let stats: ConnStats = sqlx::query_as("SELECT * FROM conn_stats")
128+
.fetch_one(conn)
129+
.await?;
130+
131+
// For even IDs, cap by the number of before_acquire calls.
132+
// Ignore the check for odd IDs.
133+
Ok((stats.id & 1) == 1 || stats.before_acquire_calls < 3)
134+
})
135+
})
136+
.after_release(|conn, meta| {
137+
// `age` should be nonzero but `idle_for` should be zero.
138+
assert_ne!(meta.age, Duration::ZERO);
139+
assert_eq!(meta.idle_for, Duration::ZERO);
140+
141+
Box::pin(async move {
142+
sqlx::query(
143+
r#"
144+
UPDATE conn_stats
145+
SET after_release_calls = after_release_calls + 1
146+
"#,
147+
)
148+
.execute(&mut *conn)
149+
.await?;
150+
151+
let stats: ConnStats = sqlx::query_as("SELECT * FROM conn_stats")
152+
.fetch_one(conn)
153+
.await?;
154+
155+
// For odd IDs, cap by the number of before_release calls.
156+
// Ignore the check for even IDs.
157+
Ok((stats.id & 1) == 0 || stats.after_release_calls < 4)
158+
})
159+
})
160+
// Don't establish a connection yet.
161+
.connect_lazy(&dotenv::var("DATABASE_URL")?)?;
162+
163+
// Expected pattern of (id, before_acquire_calls, after_release_calls)
164+
let pattern = [
165+
// The connection pool starts empty.
166+
(0, 0, 0),
167+
(0, 1, 1),
168+
(0, 2, 2),
169+
(1, 0, 0),
170+
(1, 1, 1),
171+
(1, 2, 2),
172+
// We should expect one more `acquire` because the ID is odd
173+
(1, 3, 3),
174+
(2, 0, 0),
175+
(2, 1, 1),
176+
(2, 2, 2),
177+
(3, 0, 0),
178+
];
179+
180+
for (id, before_acquire_calls, after_release_calls) in pattern {
181+
let conn_stats: ConnStats = sqlx::query_as("SELECT * FROM conn_stats")
182+
.fetch_one(&pool)
183+
.await?;
184+
185+
assert_eq!(
186+
conn_stats,
187+
ConnStats {
188+
id,
189+
before_acquire_calls,
190+
after_release_calls
191+
}
192+
);
193+
}
194+
195+
pool.close().await;
196+
197+
Ok(())
198+
}

0 commit comments

Comments
 (0)