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

[sundapeng]

Purpose

Linked issue: #872

A filtered information_schema query through the Python binding never returns:

SELECT column_name FROM paimon.information_schema.columns WHERE table_schema = 'db';

The plan has a RepartitionExec, so the metadata enumeration runs on a worker of the process runtime. The enumeration calls CatalogProvider::schema() on that worker. block_on_with_runtime then blocks the worker with a spawned thread and join(), in the middle of a poll.

Two things are lost with the worker:

• Its LIFO slot. Tokio does not steal from the LIFO slot (1.51 added stealing, 1.52.2 reverted it), so the task in it never runs. Here it is the hyper connection task of the REST client. It was woken when the previous table().await dropped its response body, and a push into an empty LIFO slot wakes no other worker.
• The I/O driver, when this worker was the last one parked on it. The other workers sleep on a condvar and nothing wakes them. No I/O and no timers run after that, so even the 30s timeout of the REST client does not fire.

The helper thread waits for a request that is never written to the socket. gdb on the hung process shows lifo_slot: Some(..) with the task in NOTIFIED state and had_driver: Yes on the blocked worker. ss shows one established connection with empty queues.

The number of workers does not matter. It hangs 10/10 on 2, 4 and 8 cores. With one core, a plan that spawns tasks hangs the same way, for example UNION ALL over two information_schema scans. Turning off the LIFO slot in the same binary makes it finish 10/10.

In the Python binding no runtime is entered on the calling thread, so runtime() is the process runtime. DataFusion and the helper thread use the same runtime, there is only one in this path.

Brief change log

• Mark the threads of the process runtime in on_thread_start. On these threads block_on_with_runtime uses block_in_place. It moves the LIFO slot to the run queue and hands the worker core to another thread before blocking.
• All other callers keep the thread-based path. block_in_place panics inside a LocalSet and on a current-thread runtime, and Tokio has no API to ask whether it is allowed. This is the case raised in the review of fix(datafusion): preserve active runtime in blocking callbacks #873.
• Declare the rt-multi-thread Tokio feature. Runtime::new() already needed it and only got it through feature unification.

This does not close #872. A worker of a runtime owned by the caller still takes the thread-based path. #881 removes the blocking call from the discovery callbacks. This change fixes the bridge itself, so it also covers the callbacks that #881 leaves on it. The two only touch the same line in Cargo.toml.

Tests

• runtime::tests::blocking_on_a_process_runtime_worker_keeps_its_queued_tasks_running: a task spawns a second task, which lands in the LIFO slot, then blocks on its result. It fails after 30s without the fix, on any number of cores.
• runtime::tests::callers_outside_the_process_runtime_keep_the_thread_based_path: a LocalSet on a multi-thread runtime (the root future and a spawn_local task) and a current-thread runtime. With the approach of fix(datafusion): preserve active runtime in blocking callbacks #873 it panics with can call blocking only when running on the multi-threaded runtime.
• tests/catalog_callbacks_on_runtime_workers.rs: a REST catalog queried from a plain thread, the way the Python binding does it, with the filtered query and a UNION ALL query. Without the fix it fails on the filtered query with 8 cores and on the UNION ALL query with 1 core.
• cargo fmt --all -- --check
• cargo clippy --locked -p paimon-datafusion --all-targets -- -D warnings
• cargo test --locked -p paimon-datafusion: 801 passed. 39 tests in table::tests and tests/read_tables.rs need the shared test warehouse and fail locally with TableNotExist. The same 39 fail on main.

API and Format

No.

Documentation

No.

[JingsongLi]

+1