feat(bin): add --shutdown-timeout to set a timeout for shutting down the database connections after a test file is finished (#265)

Author: BugenZhao

.github/workflows/ci.yml

@@ -68,7 +68,7 @@ jobs:
         uses: actions-rs/cargo@v1
         with:
           command: install
-          args: cargo-semver-checks --version ^0.39 --locked
+          args: cargo-semver-checks --version ^0.41 --locked
       - name: Check semver
         run: |
          cargo semver-checks check-release -p sqllogictest

CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## Unreleased
 
+## [0.28.3] - 2025-05-16
+
+* bin: Add `--shutdown-timeout` to set a timeout for shutting down the database connections after a test file is finished. By default, this is unspecified, meaning to wait forever.
+
 ## [0.28.2] - 2025-05-07
 
 * engines: Enhance graceful shutdown by canceling ongoing queries when shutting down postgres connections. This improves the behavior when receiving a `Ctrl-C` signal by ensuring that any long-running queries are properly canceled before the connection is closed.

Cargo.lock

@@ -3,7 +3,7 @@ resolver = "2"
 members = ["sqllogictest", "sqllogictest-bin", "sqllogictest-engines", "tests"]
 
 [workspace.package]
-version = "0.28.2"
+version = "0.28.3"
 edition = "2021"
 homepage = "https://github.com/risinglightdb/sqllogictest-rs"
 keywords = ["sql", "database", "parser", "cli"]

Cargo.toml

@@ -132,6 +132,11 @@ struct Opt {
     /// Total number of partitions for test sharding. More details in `partition_id`.
     #[clap(long, env = PARTITION_COUNT_ENV_KEY)]
     partition_count: Option<u64>,
+
+    /// Timeout in seconds for shutting down the connections to the database after a
+    /// test file is finished. By default, this is unspecified, meaning to wait forever.
+    #[clap(long = "shutdown-timeout", env = "SLT_SHUTDOWN_TIMEOUT")]
+    shutdown_timeout_secs: Option<u64>,
 }
 
 /// Connection configuration.
@@ -241,6 +246,7 @@ pub async fn main() -> Result<()> {
         labels,
         partition_count,
         partition_id,
+        shutdown_timeout_secs,
     } = Opt::from_arg_matches(&matches)
         .map_err(|err| err.exit())
         .unwrap();
@@ -341,6 +347,14 @@ pub async fn main() -> Result<()> {
         }
     });
 
+    let run_config = RunConfig {
+        labels,
+        junit: junit.clone(),
+        fail_fast,
+        cancel,
+        shutdown_timeout: shutdown_timeout_secs.map(Duration::from_secs),
+    };
+
     let result = if let Some(jobs) = jobs {
         run_parallel(
             jobs,
@@ -349,24 +363,11 @@ pub async fn main() -> Result<()> {
             all_files,
             &engine,
             config,
-            &labels,
-            junit.clone(),
-            fail_fast,
-            cancel,
+            run_config,
         )
         .await
     } else {
-        run_serial(
-            &mut test_suite,
-            all_files,
-            &engine,
-            config,
-            &labels,
-            junit.clone(),
-            fail_fast,
-            cancel,
-        )
-        .await
+        run_serial(&mut test_suite, all_files, &engine, config, run_config).await
     };
 
     report.add_test_suite(test_suite);
@@ -378,18 +379,28 @@ pub async fn main() -> Result<()> {
     result
 }
 
-#[allow(clippy::too_many_arguments)]
+struct RunConfig {
+    labels: Vec<String>,
+    junit: Option<String>,
+    fail_fast: bool,
+    cancel: CancellationToken,
+    shutdown_timeout: Option<Duration>,
+}
+
 async fn run_parallel(
     jobs: usize,
     keep_db_on_failure: bool,
     test_suite: &mut TestSuite,
     files: Vec<PathBuf>,
     engine: &EngineConfig,
     config: DBConfig,
-    labels: &[String],
-    junit: Option<String>,
-    fail_fast: bool,
-    cancel: CancellationToken,
+    RunConfig {
+        labels,
+        junit,
+        fail_fast,
+        cancel,
+        shutdown_timeout,
+    }: RunConfig,
 ) -> Result<()> {
     let mut create_databases = BTreeMap::new();
     let mut test_cases = BTreeSet::new();
@@ -427,7 +438,7 @@ async fn run_parallel(
             config.db.clone_from(&db_name);
             let file = filename.to_string_lossy().to_string();
             let engine = engine.clone();
-            let labels = labels.to_vec();
+            let labels = labels.clone();
             let cancel = cancel.clone();
             async move {
                 let res = AbortOnDropHandle::new(tokio::spawn(async move {
@@ -438,6 +449,7 @@ async fn run_parallel(
                         config,
                         &labels,
                         cancel,
+                        shutdown_timeout,
                     )
                     .await
                 }))
@@ -529,10 +541,13 @@ async fn run_serial(
     files: Vec<PathBuf>,
     engine: &EngineConfig,
     config: DBConfig,
-    labels: &[String],
-    junit: Option<String>,
-    fail_fast: bool,
-    cancel: CancellationToken,
+    RunConfig {
+        labels,
+        junit,
+        fail_fast,
+        cancel,
+        shutdown_timeout,
+    }: RunConfig,
 ) -> Result<()> {
     let mut failed_cases = vec![];
     let mut connection_refused = false;
@@ -545,8 +560,9 @@ async fn run_serial(
             file,
             engine,
             config.clone(),
-            labels,
+            &labels,
             cancel.clone(),
+            shutdown_timeout,
         )
         .await;
         stdout().flush()?;
@@ -694,6 +710,7 @@ async fn connect_and_run_test_file(
     config: DBConfig,
     labels: &[String],
     cancel: CancellationToken,
+    shutdown_timeout: Option<Duration>,
 ) -> RunResult {
     struct OutputGuard<O: Output>(O);
     impl<O: Output> Drop for OutputGuard<O> {
@@ -745,7 +762,7 @@ async fn connect_and_run_test_file(
             .unwrap();
             RunResult::Cancelled
         }
-        result = run_test_file(&mut out.0, &mut runner, filename) => {
+        result = run_test_file(&mut out.0, &mut runner, filename.clone()) => {
             if let Err(err) = &result {
                 writeln!(
                     out.0,
@@ -760,7 +777,21 @@ async fn connect_and_run_test_file(
     };
 
     drop(out); // flush the output before shutting down the runner
-    runner.shutdown_async().await;
+
+    match shutdown_timeout {
+        None => runner.shutdown_async().await,
+        Some(timeout) => {
+            if tokio::time::timeout(timeout, runner.shutdown_async())
+                .await
+                .is_err()
+            {
+                eprintln!(
+                    "shutting down connection to database for test {} timed out",
+                    filename.display()
+                )
+            }
+        }
+    }
 
     result
 }