Switch ParquetObjectReader runtime tests to not depend on tokio_unsta…

…ble anymore

parquet/Cargo.toml

@@ -28,10 +28,6 @@ readme = "README.md"
 edition = { workspace = true }
 rust-version = "1.70.0"
 
-# we need this for the arrow::async_reader::test_runtime_is_used test
-[lints.rust]
-unexpected_cfgs = { level = "warn", check-cfg = ["cfg(tokio_unstable)"] }
-
 [target.'cfg(target_arch = "wasm32")'.dependencies]
 ahash = { version = "0.8", default-features = false, features = ["compile-time-rng"] }
 
@@ -85,7 +81,7 @@ lz4_flex = { version = "0.11", default-features = false, features = ["std", "fra
 zstd = { version = "0.13", default-features = false }
 serde_json = { version = "1.0", features = ["std"], default-features = false }
 arrow = { workspace = true, features = ["ipc", "test_utils", "prettyprint", "json"] }
-tokio = { version = "1.0", default-features = false, features = ["macros", "rt", "io-util", "fs"] }
+tokio = { version = "1.0", default-features = false, features = ["macros", "rt-multi-thread", "io-util", "fs"] }
 rand = { version = "0.8", default-features = false, features = ["std", "std_rng"] }
 object_store = { version = "0.11.0", default-features = false, features = ["azure"] }
 

parquet/src/arrow/async_reader/store.rs

@@ -173,11 +173,18 @@ impl AsyncFileReader for ParquetObjectReader {
 
 #[cfg(test)]
 mod tests {
-    use std::sync::Arc;
+    use std::{
+        convert::Infallible,
+        sync::{
+            atomic::{AtomicUsize, Ordering},
+            Arc,
+        },
+    };
 
     use futures::TryStreamExt;
 
     use arrow::util::test_util::parquet_test_data;
+    use futures::FutureExt;
     use object_store::local::LocalFileSystem;
     use object_store::path::Path;
     use object_store::{ObjectMeta, ObjectStore};
@@ -233,14 +240,24 @@ mod tests {
     #[tokio::test]
     // We need to mark this with the `target_has_atomic` because the spawned_tasks_count() fn is
     // only available for that cfg
-    #[cfg(all(target_has_atomic = "64", tokio_unstable))]
     async fn test_runtime_is_used() {
-        let rt = tokio::runtime::Builder::new_current_thread()
+        let num_actions = Arc::new(AtomicUsize::new(0));
+
+        let (a1, a2) = (num_actions.clone(), num_actions.clone());
+        let rt = tokio::runtime::Builder::new_multi_thread()
+            .on_thread_park(move || {
+                a1.fetch_add(1, Ordering::Relaxed);
+            })
+            .on_thread_unpark(move || {
+                a2.fetch_add(1, Ordering::Relaxed);
+            })
             .build()
             .unwrap();
 
         let (meta, store) = get_meta_store().await;
 
+        let initial_actions = num_actions.load(Ordering::Relaxed);
+
         let reader = ParquetObjectReader::new(store, meta).with_runtime(rt.handle().clone());
 
         let builder = ParquetRecordBatchStreamBuilder::new(reader).await.unwrap();
@@ -250,13 +267,33 @@ mod tests {
         assert_eq!(batches.len(), 1);
         assert_eq!(batches[0].num_rows(), 8);
 
-        // According to tokio documentation for the `spawned_tasks_count` method, this number
-        // starts at 0 when the runtime is created. So this check should actually verify what we
-        // want.
-        assert!(rt.metrics().spawned_tasks_count() > 0);
+        assert!(num_actions.load(Ordering::Relaxed) - initial_actions > 0);
 
         // Runtimes have to be dropped in blocking contexts, so we need to move this one to a new
         // blocking thread to drop it.
         tokio::runtime::Handle::current().spawn_blocking(move || drop(rt));
     }
+
+    #[tokio::test]
+    async fn test_runtime_thread_id_different() {
+        let rt = tokio::runtime::Builder::new_multi_thread()
+            .worker_threads(1)
+            .build()
+            .unwrap();
+
+        let (meta, store) = get_meta_store().await;
+
+        let reader = ParquetObjectReader::new(store, meta).with_runtime(rt.handle().clone());
+
+        let current_id = std::thread::current().id();
+
+        let other_id = reader
+            .spawn(|_, _| async move { Ok::<_, Infallible>(std::thread::current().id()) }.boxed())
+            .await
+            .unwrap();
+
+        assert_ne!(current_id, other_id);
+
+        tokio::runtime::Handle::current().spawn_blocking(move || drop(rt));
+    }
 }

parquet/src/errors.rs

@@ -107,6 +107,13 @@ impl From<str::Utf8Error> for ParquetError {
     }
 }
 
+#[cfg(test)]
+impl From<std::convert::Infallible> for ParquetError {
+    fn from(value: std::convert::Infallible) -> Self {
+        match value {}
+    }
+}
+
 #[cfg(feature = "arrow")]
 impl From<ArrowError> for ParquetError {
     fn from(e: ArrowError) -> ParquetError {