feat(bin): abort connections before dropping temp databases in parall…

…el run (#247) * feat(bin): abort connections before dropping temp databases in parallel run Signed-off-by: Bugen Zhao <[email protected]> * fmt Signed-off-by: Bugen Zhao <[email protected]> * refine comments Signed-off-by: Bugen Zhao <[email protected]> * use one from `tokio-util` Signed-off-by: Bugen Zhao <[email protected]> * bump version and add release notes Signed-off-by: Bugen Zhao <[email protected]> --------- Signed-off-by: Bugen Zhao <[email protected]>
risinglightdb · Jan 14, 2025 · f0cda70 · f0cda70
1 parent 3c4ee72
commit f0cda70
Show file tree

Hide file tree

Showing 5 changed files with 26 additions and 7 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## Unreleased
 
+## [0.26.3] - 2025-01-14
+
+* bin: when `--fail-fast` is enabled, abort all remaining connections before dropping temporary databases.
+
 ## [0.26.2] - 2025-01-08
 
 * bin: support `--fail-fast`, and add env vars `SLT_FAIL_FAST` and `SLT_KEEP_DB_ON_FAILURE`

diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -3,7 +3,7 @@ resolver = "2"
 members = ["sqllogictest", "sqllogictest-bin", "sqllogictest-engines", "tests"]
 
 [workspace.package]
-version = "0.26.2"
+version = "0.26.3"
 edition = "2021"
 homepage = "https://github.com/risinglightdb/sqllogictest-rs"
 keywords = ["sql", "database", "parser", "cli"]

diff --git a/sqllogictest-bin/Cargo.toml b/sqllogictest-bin/Cargo.toml
@@ -33,6 +33,7 @@ tokio = { version = "1", features = [
     "fs",
     "process",
 ] }
+tokio-util = { version = "0.7.12", features = ["rt"] }
 fs-err = "3.0.0"
 tracing-subscriber = { version = "0.3", features = ["env-filter"] }
 tracing = "0.1"
diff --git a/sqllogictest-bin/src/main.rs b/sqllogictest-bin/src/main.rs
@@ -20,6 +20,7 @@ use sqllogictest::{
     default_column_validator, default_normalizer, default_validator, update_record_with_output,
     AsyncDB, Injected, MakeConnection, Record, Runner,
 };
+use tokio_util::task::AbortOnDropHandle;
 
 #[derive(Default, Copy, Clone, Debug, PartialEq, Eq, ValueEnum)]
 #[must_use]
@@ -314,21 +315,21 @@ async fn run_parallel(
         }
     }
 
-    let mut stream = futures::stream::iter(create_databases.into_iter())
+    let mut stream = futures::stream::iter(create_databases)
         .map(|(db_name, filename)| {
             let mut config = config.clone();
             config.db.clone_from(&db_name);
             let file = filename.to_string_lossy().to_string();
             let engine = engine.clone();
             let labels = labels.to_vec();
             async move {
-                let (buf, res) = tokio::spawn(async move {
+                let (buf, res) = AbortOnDropHandle::new(tokio::spawn(async move {
                     let mut buf = vec![];
                     let res =
                         connect_and_run_test_file(&mut buf, filename, &engine, config, &labels)
                             .await;
                     (buf, res)
-                })
+                }))
                 .await
                 .unwrap();
                 (db_name, file, res, buf)
@@ -396,6 +397,10 @@ async fn run_parallel(
         start.elapsed().as_millis()
     );
 
+    // If `fail_fast`, there could be some ongoing cases (then active connections)
+    // in the stream. Abort them before dropping temporary databases.
+    drop(stream);
+
     for db_name in db_names {
         if keep_db_on_failure && failed_db.contains(&db_name) {
             eprintln!(