fix(datafusion): hand off the worker before blocking in catalog callbacks

crates/integrations/datafusion/Cargo.toml

@@ -43,7 +43,7 @@ paimon = { workspace = true }
 futures = "0.3"
 serde = { version = "1", features = ["derive"] }
 serde_json = "1"
-tokio = { workspace = true, features = ["rt", "time", "fs"] }
+tokio = { workspace = true, features = ["rt", "rt-multi-thread", "time", "fs"] }
 lexical-write-float = "1.0.6"
 uuid = { version = "1", features = ["v4"] }
 

crates/integrations/datafusion/src/runtime.rs

@@ -15,11 +15,12 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use std::cell::Cell;
 use std::future::Future;
 use std::sync::atomic::{AtomicPtr, Ordering};
 use std::sync::OnceLock;
 
-use tokio::runtime::{Handle, Runtime};
+use tokio::runtime::{Builder, Handle, Runtime};
 
 struct ProcessRuntime {
     pid: u32,
@@ -28,6 +29,18 @@ struct ProcessRuntime {
 
 static RUNTIME: AtomicPtr<ProcessRuntime> = AtomicPtr::new(std::ptr::null_mut());
 
+thread_local! {
+    // Set on every thread the process runtime starts, where `block_in_place` is allowed.
+    static ON_PROCESS_RUNTIME_THREAD: Cell<bool> = const { Cell::new(false) };
+}
+
+fn build_process_runtime() -> std::io::Result<Runtime> {
+    Builder::new_multi_thread()
+        .enable_all()
+        .on_thread_start(|| ON_PROCESS_RUNTIME_THREAD.with(|flag| flag.set(true)))
+        .build()
+}
+
 fn global_runtime() -> &'static Runtime {
     let pid = std::process::id();
     let mut current = RUNTIME.load(Ordering::Acquire);
@@ -38,7 +51,7 @@ fn global_runtime() -> &'static Runtime {
             let state = unsafe { &*current };
             if state.pid == pid {
                 return state.runtime.get_or_init(|| {
-                    Runtime::new().expect(
+                    build_process_runtime().expect(
                         "failed to build global tokio runtime for paimon datafusion integration",
                     )
                 });
@@ -103,14 +116,19 @@ where
     F: Future + Send + 'static,
     F::Output: Send + 'static,
 {
-    if Handle::try_current().is_ok() {
-        let handle = global_runtime().handle().clone();
-        std::thread::spawn(move || handle.block_on(future))
-            .join()
-            .expect(panic_error)
-    } else {
-        global_runtime().block_on(future)
+    if Handle::try_current().is_err() {
+        return global_runtime().block_on(future);
+    }
+    if ON_PROCESS_RUNTIME_THREAD.with(Cell::get) {
+        // A worker blocked mid-poll strands its LIFO slot and the I/O driver it last parked on.
+        // `block_in_place` hands both to another thread first.
+        return tokio::task::block_in_place(|| global_runtime().block_on(future));
     }
+    // Threads of other runtimes: `block_in_place` panics in a `LocalSet` or a current-thread runtime.
+    let handle = global_runtime().handle().clone();
+    std::thread::spawn(move || handle.block_on(future))
+        .join()
+        .expect(panic_error)
 }
 
 #[cfg(test)]
@@ -142,6 +160,50 @@ mod tests {
         });
     }
 
+    #[test]
+    fn blocking_on_a_process_runtime_worker_keeps_its_queued_tasks_running() {
+        let (done_tx, done_rx) = std::sync::mpsc::channel();
+        global_runtime().spawn(async move {
+            let (value_tx, value_rx) = tokio::sync::oneshot::channel();
+            // Spawned from a worker, this lands in that worker's LIFO slot, which cannot be stolen.
+            tokio::spawn(async move {
+                let _ = value_tx.send(7);
+            });
+            let value = block_on_with_runtime(
+                async move { value_rx.await.unwrap() },
+                "blocking runtime test panicked",
+            );
+            done_tx.send(value).unwrap();
+        });
+        assert_eq!(
+            done_rx.recv_timeout(std::time::Duration::from_secs(30)),
+            Ok(7),
+            "the blocked worker stranded the task it had just queued"
+        );
+    }
+
+    #[test]
+    fn callers_outside_the_process_runtime_keep_the_thread_based_path() {
+        let block = || block_on_with_runtime(async { 7 }, "blocking runtime test panicked");
+
+        // `block_in_place` would panic in both of these.
+        let multi_thread = tokio::runtime::Builder::new_multi_thread()
+            .worker_threads(2)
+            .build()
+            .unwrap();
+        let local = tokio::task::LocalSet::new();
+        let in_local_set = local.block_on(&multi_thread, async move {
+            let spawned = tokio::task::spawn_local(async move { block() });
+            (block(), spawned.await.unwrap())
+        });
+        assert_eq!(in_local_set, (7, 7));
+
+        let current_thread = tokio::runtime::Builder::new_current_thread()
+            .build()
+            .unwrap();
+        assert_eq!(current_thread.block_on(async move { block() }), 7);
+    }
+
     #[test]
     fn entered_runtime_is_preserved() {
         let local = tokio::runtime::Builder::new_current_thread()

crates/integrations/datafusion/tests/catalog_callbacks_on_runtime_workers.rs

@@ -0,0 +1,112 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! `information_schema` queries whose plan spawns tasks, so the synchronous catalog callbacks
+//! run on a worker of the process runtime. Driven like the Python binding, from a plain thread.
+
+#[path = "../../../paimon/tests/mock_server.rs"]
+mod mock_server;
+
+use std::collections::HashMap;
+use std::sync::{mpsc, Arc};
+use std::time::Duration;
+
+use arrow_array::RecordBatch;
+use paimon::api::ConfigResponse;
+use paimon::catalog::RESTCatalog;
+use paimon::spec::{DataType, IntType, Schema};
+use paimon::{CatalogOptions, Options};
+use paimon_datafusion::SQLContext;
+
+use mock_server::start_mock_server;
+
+const WAREHOUSE: &str = "test_warehouse";
+const DATABASES: [&str; 2] = ["db_a", "db_b"];
+
+/// `(query, rows)`. Every table has two columns and every database has one table.
+const QUERIES: [(&str, usize); 2] = [
+    // With two or more target partitions the filter sits on a `RepartitionExec`.
+    (
+        "SELECT column_name FROM paimon.information_schema.columns WHERE table_schema = 'db_a'",
+        2,
+    ),
+    // Two output partitions are merged by spawned tasks whatever the number of cores.
+    (
+        "SELECT column_name FROM paimon.information_schema.columns \
+         UNION ALL SELECT column_name FROM paimon.information_schema.columns",
+        8,
+    ),
+];
+
+#[test]
+fn information_schema_queries_finish_when_the_plan_spawns_tasks() {
+    // The catalog service gets its own runtime, so the runtime under test cannot starve it.
+    let server_runtime = tokio::runtime::Builder::new_multi_thread()
+        .worker_threads(2)
+        .enable_all()
+        .build()
+        .unwrap();
+    let temp_dir = tempfile::tempdir().unwrap();
+    let server = server_runtime.block_on(start_mock_server(
+        WAREHOUSE.to_string(),
+        temp_dir.path().to_string_lossy().into_owned(),
+        ConfigResponse::new(HashMap::from([(
+            CatalogOptions::PREFIX.to_string(),
+            "mock-test".to_string(),
+        )])),
+        DATABASES.iter().map(|name| name.to_string()).collect(),
+    ));
+    for database in DATABASES {
+        let schema = Schema::builder()
+            .column("id", DataType::Int(IntType::new()))
+            .column("v", DataType::Int(IntType::new()))
+            .build()
+            .unwrap();
+        let path = format!("file://{}/{database}.db/t", temp_dir.path().display());
+        server.add_table_with_schema(database, "t", schema, &path);
+    }
+    let url = server.url().unwrap();
+
+    let (rows_tx, rows_rx) = mpsc::channel();
+    // No runtime is entered on this thread, so `runtime()` is the process runtime.
+    std::thread::spawn(move || {
+        paimon_datafusion::runtime::runtime().block_on(async move {
+            let mut options = Options::new();
+            options.set(CatalogOptions::URI, url);
+            options.set(CatalogOptions::WAREHOUSE, WAREHOUSE);
+            options.set(CatalogOptions::TOKEN_PROVIDER, "bear");
+            options.set(CatalogOptions::TOKEN, "test-token");
+            let catalog = Arc::new(RESTCatalog::new(options, true).await.unwrap());
+            let mut context = SQLContext::new();
+            context.register_catalog("paimon", catalog).await.unwrap();
+
+            for (query, _) in QUERIES {
+                let batches = context.sql(query).await.unwrap().collect().await.unwrap();
+                let rows: usize = batches.iter().map(RecordBatch::num_rows).sum();
+                rows_tx.send(rows).unwrap();
+            }
+        });
+    });
+
+    for (query, expected_rows) in QUERIES {
+        // A timer of the blocked runtime would never fire, so the deadline lives on this thread.
+        let rows = rows_rx
+            .recv_timeout(Duration::from_secs(60))
+            .unwrap_or_else(|_| panic!("a catalog callback blocked a runtime worker: {query}"));
+        assert_eq!(rows, expected_rows, "{query}");
+    }
+}