From 140e3bfe5eab074948f41c5c14af514e4dc14b5c Mon Sep 17 00:00:00 2001
From: Bugen Zhao <i@bugenzhao.com>
Date: Tue, 14 Jan 2025 14:31:33 +0800
Subject: [PATCH] use one from `tokio-util`

Signed-off-by: Bugen Zhao <i@bugenzhao.com>
---
 Cargo.lock                   |  9 +++++++++
 sqllogictest-bin/Cargo.toml  |  1 +
 sqllogictest-bin/src/main.rs | 29 ++---------------------------
 3 files changed, 12 insertions(+), 27 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 61011e5..221f501 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -779,6 +779,12 @@ dependencies = [
  "ahash",
 ]
 
+[[package]]
+name = "hashbrown"
+version = "0.14.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1"
+
 [[package]]
 name = "hashbrown"
 version = "0.15.2"
@@ -1923,6 +1929,7 @@ dependencies = [
  "sqllogictest",
  "sqllogictest-engines",
  "tokio",
+ "tokio-util",
  "tracing",
  "tracing-subscriber",
 ]
@@ -2212,6 +2219,8 @@ dependencies = [
  "bytes",
  "futures-core",
  "futures-sink",
+ "futures-util",
+ "hashbrown 0.14.5",
  "pin-project-lite",
  "tokio",
 ]
diff --git a/sqllogictest-bin/Cargo.toml b/sqllogictest-bin/Cargo.toml
index f8aa4b4..6587463 100644
--- 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
index 50ec5d1..2c6dd46 100644
--- a/sqllogictest-bin/src/main.rs
+++ b/sqllogictest-bin/src/main.rs
@@ -20,7 +20,7 @@ use sqllogictest::{
     default_column_validator, default_normalizer, default_validator, update_record_with_output,
     AsyncDB, Injected, MakeConnection, Record, Runner,
 };
-use utils::AbortOnDropHandle;
+use tokio_util::task::AbortOnDropHandle;
 
 #[derive(Default, Copy, Clone, Debug, PartialEq, Eq, ValueEnum)]
 #[must_use]
@@ -323,7 +323,7 @@ async fn run_parallel(
             let engine = engine.clone();
             let labels = labels.to_vec();
             async move {
-                let (buf, res) = AbortOnDropHandle(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)
@@ -835,28 +835,3 @@ async fn update_record<M: MakeConnection>(
 
     Ok(())
 }
-
-mod utils {
-    use std::future::Future;
-    use std::pin::Pin;
-    use std::task::{Context, Poll};
-
-    use tokio::task::{JoinError, JoinHandle};
-
-    /// A wrapper around a [`tokio::task::JoinHandle`], which aborts the task when it is dropped.
-    pub struct AbortOnDropHandle<T>(pub JoinHandle<T>);
-
-    impl<T> Drop for AbortOnDropHandle<T> {
-        fn drop(&mut self) {
-            self.0.abort();
-        }
-    }
-
-    impl<T> Future for AbortOnDropHandle<T> {
-        type Output = Result<T, JoinError>;
-
-        fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
-            Pin::new(&mut self.0).poll(cx)
-        }
-    }
-}