This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/paimon-rust.git
The following commit(s) were added to refs/heads/main by this push:
new 07c0fb45 fix(datafusion): hand off the worker before blocking in
catalog callbacks (#891)
07c0fb45 is described below
commit 07c0fb454658d3d1a83c6198ff0783f29e75c53a
Author: Dapeng Sun(孙大鹏) <[email protected]>
AuthorDate: Mon Sep 21 13:40:30 2026 +0800
fix(datafusion): hand off the worker before blocking in catalog callbacks
(#891)
---
crates/integrations/datafusion/Cargo.toml | 2 +-
crates/integrations/datafusion/src/runtime.rs | 80 +++++++++++++--
.../tests/catalog_callbacks_on_runtime_workers.rs | 112 +++++++++++++++++++++
3 files changed, 184 insertions(+), 10 deletions(-)
diff --git a/crates/integrations/datafusion/Cargo.toml
b/crates/integrations/datafusion/Cargo.toml
index f7072fc3..a2f78ae2 100644
--- a/crates/integrations/datafusion/Cargo.toml
+++ b/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"] }
diff --git a/crates/integrations/datafusion/src/runtime.rs
b/crates/integrations/datafusion/src/runtime.rs
index aed51093..6fc9bf0d 100644
--- a/crates/integrations/datafusion/src/runtime.rs
+++ b/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()
diff --git
a/crates/integrations/datafusion/tests/catalog_callbacks_on_runtime_workers.rs
b/crates/integrations/datafusion/tests/catalog_callbacks_on_runtime_workers.rs
new file mode 100644
index 00000000..84ac5e68
--- /dev/null
+++
b/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}");
+ }
+}
