From 0c72caf9e45706967af531f050220fcc64a9e1a1 Mon Sep 17 00:00:00 2001 From: Dapeng Sun Date: Mon, 21 Sep 2026 08:33:04 +0800 Subject: [PATCH] fix(datafusion): hand off the worker before blocking in catalog callbacks A filtered information_schema query through the Python binding never returns. The plan has a RepartitionExec, so the metadata enumeration runs on a worker of the process runtime, and block_on_with_runtime blocks that worker with a spawned thread and join() in the middle of a poll. Two things are lost with the worker. Its LIFO slot, which Tokio does not steal from, holds the hyper connection task that was woken when the previous table() call dropped its response body. And when the worker was the last one parked on the I/O driver, nobody takes the driver over, so no I/O and no timers run. The helper thread waits for a request that is never written to the socket. Mark the threads of the process runtime and use block_in_place on them. It moves the LIFO slot to the run queue and hands the worker core to another thread before blocking. Every other caller keeps the thread-based path, because block_in_place panics inside a LocalSet and on a current-thread runtime, and Tokio has no way to ask whether it is allowed. --- crates/integrations/datafusion/Cargo.toml | 2 +- crates/integrations/datafusion/src/runtime.rs | 80 +++++++++++-- .../catalog_callbacks_on_runtime_workers.rs | 112 ++++++++++++++++++ 3 files changed, 184 insertions(+), 10 deletions(-) create mode 100644 crates/integrations/datafusion/tests/catalog_callbacks_on_runtime_workers.rs diff --git a/crates/integrations/datafusion/Cargo.toml b/crates/integrations/datafusion/Cargo.toml index f7072fc36..a2f78ae27 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 aed51093a..6fc9bf0d9 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 = 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 = const { Cell::new(false) }; +} + +fn build_process_runtime() -> std::io::Result { + 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 000000000..84ac5e688 --- /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}"); + } +}