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}"); + } +}