fix: ensure fresh test db's aren't accidentally deleted by do_cleanup (#2454)

If the first test thread is a bit slow after it acquires the
`DO_CLEANUP` permit, it can accidentally delete a fresh test db created
by another thread right before that other thread has a chance to open
its connection.

This fix simply records the current timestamp _before_ we acquire the
`DO_CLEANUP` permit and only cleans up test db's created before then.
This commit is contained in:
Philip Kannegaard Hayes 2023-05-04 13:23:04 -07:00 committed by GitHub
parent cbf8dd37e9
commit 1227d5d168
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 44 additions and 13 deletions

View file

@ -2,7 +2,7 @@ use std::fmt::Write;
use std::ops::Deref;
use std::str::FromStr;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Duration;
use std::time::{Duration, SystemTime};
use futures_core::future::BoxFuture;
@ -60,7 +60,12 @@ impl TestSupport for MySql {
let url = dotenvy::var("DATABASE_URL").expect("DATABASE_URL must be set");
let mut conn = MySqlConnection::connect(&url).await?;
let num_deleted = do_cleanup(&mut conn).await?;
let now = SystemTime::now()
.duration_since(SystemTime::UNIX_EPOCH)
.unwrap();
let num_deleted = do_cleanup(&mut conn, now).await?;
let _ = conn.close().await;
Ok(Some(num_deleted))
})
@ -123,9 +128,16 @@ async fn test_context(args: &TestArgs) -> Result<TestContext<MySql>, Error> {
)
.await?;
// Record the current time _before_ we acquire the `DO_CLEANUP` permit. This
// prevents the first test thread from accidentally deleting new test dbs
// created by other test threads if we're a bit slow.
let now = SystemTime::now()
.duration_since(SystemTime::UNIX_EPOCH)
.unwrap();
// Only run cleanup if the test binary just started.
if DO_CLEANUP.swap(false, Ordering::SeqCst) {
do_cleanup(&mut conn).await?;
do_cleanup(&mut conn, now).await?;
}
query("insert into _sqlx_test_databases(test_path) values (?)")
@ -163,10 +175,12 @@ async fn test_context(args: &TestArgs) -> Result<TestContext<MySql>, Error> {
})
}
async fn do_cleanup(conn: &mut MySqlConnection) -> Result<usize, Error> {
async fn do_cleanup(conn: &mut MySqlConnection, created_before: Duration) -> Result<usize, Error> {
let delete_db_ids: Vec<u64> = query_scalar(
"select db_id from _sqlx_test_databases where created_at < current_timestamp()",
"select db_id from _sqlx_test_databases \
where created_at < (cast(from_unixtime($1) as timestamp))",
)
.bind(&created_before.as_secs())
.fetch_all(&mut *conn)
.await?;

View file

@ -2,7 +2,7 @@ use std::fmt::Write;
use std::ops::Deref;
use std::str::FromStr;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Duration;
use std::time::{Duration, SystemTime};
use futures_core::future::BoxFuture;
@ -57,7 +57,12 @@ impl TestSupport for Postgres {
let url = dotenvy::var("DATABASE_URL").expect("DATABASE_URL must be set");
let mut conn = PgConnection::connect(&url).await?;
let num_deleted = do_cleanup(&mut conn).await?;
let now = SystemTime::now()
.duration_since(SystemTime::UNIX_EPOCH)
.unwrap();
let num_deleted = do_cleanup(&mut conn, now).await?;
let _ = conn.close().await;
Ok(Some(num_deleted))
})
@ -133,9 +138,16 @@ async fn test_context(args: &TestArgs) -> Result<TestContext<Postgres>, Error> {
)
.await?;
// Record the current time _before_ we acquire the `DO_CLEANUP` permit. This
// prevents the first test thread from accidentally deleting new test dbs
// created by other test threads if we're a bit slow.
let now = SystemTime::now()
.duration_since(SystemTime::UNIX_EPOCH)
.unwrap();
// Only run cleanup if the test binary just started.
if DO_CLEANUP.swap(false, Ordering::SeqCst) {
do_cleanup(&mut conn).await?;
do_cleanup(&mut conn, now).await?;
}
let new_db_name: String = query_scalar(
@ -170,11 +182,16 @@ async fn test_context(args: &TestArgs) -> Result<TestContext<Postgres>, Error> {
})
}
async fn do_cleanup(conn: &mut PgConnection) -> Result<usize, Error> {
let delete_db_names: Vec<String> =
query_scalar("select db_name from _sqlx_test.databases where created_at < now()")
.fetch_all(&mut *conn)
.await?;
async fn do_cleanup(conn: &mut PgConnection, created_before: Duration) -> Result<usize, Error> {
let created_before = i64::try_from(created_before.as_secs()).unwrap();
let delete_db_names: Vec<String> = query_scalar(
"select db_name from _sqlx_test.databases \
where created_at < (to_timestamp($1) at time zone 'UTC')",
)
.bind(&created_before)
.fetch_all(&mut *conn)
.await?;
if delete_db_names.is_empty() {
return Ok(0);