From 508bc33c159a7bb375a92d1d5d821fabfec22b86 Mon Sep 17 00:00:00 2001 From: Vlad Poenaru Date: Wed, 11 Nov 2026 08:34:12 -0800 Subject: [PATCH] index: receive inserter batches off the tokio workers The database inserters of the `--git` range, tree, `--commits` and lore pipelines hold a runtime worker while they wait for their next batch. They are tokio tasks that wait in a blocking `std::sync::mpsc` `recv()`, under a std `Mutex` all of them share. With `--db-threads` at or above the worker count, every worker but the one running an insert sits in `recv()` or in the mutex, and all of lance's async work for the run shares that one worker. On the default runtime that takes `--db-threads` at or above the CPU count. With the 32-worker cap in #112 it takes at most 32, so hosts with more CPUs reach it at `--db-threads 32`. How much slower such a run gets has not been measured. Add `recv_shared()`, which runs the locked receive on tokio's blocking pool, and use it at all four sites. A waiting inserter now occupies a blocking-pool thread instead of a worker; the pool's default limit is 512 threads. `recv_shared()` uses `spawn_blocking` rather than `block_in_place` because `block_in_place` panics on the current-thread runtime that `#[tokio::test]` gives tests. Each site handles a poisoned receiver lock as before: the range, tree and `--commits` inserters panic, as their `unwrap()` did, and the lore inserter logs and leaves its loop. The tuple receive in `git_range.rs` that runs on a std thread is unchanged. This conflicts with #113, which makes the lore inserter keep receiving until a batch holds `batch_size` emails. Whichever lands second resolves it by moving that loop into the closure, so the whole merge runs on the blocking pool: recv_shared(&result_rx_clone, move |rx| { rx.recv().map(|mut emails| { while emails.len() < batch_size { match rx.recv() { Ok(more) => emails.extend(more), Err(_) => break, } } emails }) }) `recv_shared_does_not_hold_workers` parks two waiters on a live channel in a two-worker runtime and checks that a newly spawned task still runs. With the waiters calling `lock().unwrap().recv()` directly, the same test fails after its 10 s timeout. --- src/git_range.rs | 16 ++++----- src/indexer.rs | 94 ++++++++++++++++++++++++++++++++++++++++++------ 2 files changed, 90 insertions(+), 20 deletions(-) diff --git a/src/git_range.rs b/src/git_range.rs index 8612688..49f3978 100644 --- a/src/git_range.rs +++ b/src/git_range.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: MIT OR Apache-2.0 use crate::git::walk_tree_at_commit_with_repo; -use crate::indexer::{list_shas_in_range, process_commits_pipeline}; +use crate::indexer::{list_shas_in_range, process_commits_pipeline, recv_shared}; use crate::types::{ArgumentFunction, GlobalVariable}; use crate::{ DatabaseManager, DispatchSite, FunctionInfo, GitFileEntry, GitFileManifestEntry, @@ -902,10 +902,9 @@ async fn process_git_tuples_streaming(config: StreamingConfig) -> Result { @@ -1165,10 +1164,9 @@ async fn process_git_tree_streaming(config: TreeStreamingConfig) -> Result { diff --git a/src/indexer.rs b/src/indexer.rs index 2752ea1..0ba4753 100644 --- a/src/indexer.rs +++ b/src/indexer.rs @@ -17,6 +17,31 @@ use tracing::{debug, error, info, warn}; use crate::git::resolve_to_commit; use crate::{DatabaseManager, GitCommitInfo}; +/// Receive from a std channel shared by several inserter tasks without +/// holding a tokio worker while waiting. +/// +/// `f` runs on the blocking pool with the receiver locked, so a task waiting +/// for a batch, or for another task to release the receiver, occupies a +/// blocking-pool thread instead of a runtime worker. Returns `None` if the +/// receiver's lock is poisoned, or if the runtime shut down before `f` ran. +/// A panic in `f` is resumed in the caller. Dropping the returned future +/// does not stop `f`; anything it received is lost. +pub(crate) async fn recv_shared( + rx: &Arc>>, + f: impl FnOnce(&mpsc::Receiver) -> R + Send + 'static, +) -> Option +where + T: Send + 'static, + R: Send + 'static, +{ + let rx = Arc::clone(rx); + match tokio::task::spawn_blocking(move || rx.lock().ok().map(|rx| f(&rx))).await { + Ok(result) => result, + Err(e) if e.is_panic() => std::panic::resume_unwind(e.into_panic()), + Err(_) => None, + } +} + /// Perform periodic optimization check during indexing operations /// Returns true if optimization was performed, false otherwise pub async fn check_and_optimize_if_needed( @@ -737,10 +762,9 @@ pub async fn process_commits_pipeline( let handle = tokio::spawn(async move { loop { // Get next batch from shared receiver - let batch = { - let rx = result_rx_clone.lock().unwrap(); - rx.recv() - }; + let batch = recv_shared(&result_rx_clone, |rx| rx.recv()) + .await + .expect("commit inserter receiver lock poisoned"); match batch { Ok(batch) => { @@ -958,13 +982,11 @@ pub async fn process_lore_commits_pipeline( let handle = tokio::spawn(async move { loop { // Get next batch from shared receiver - let batch = { - match result_rx_clone.lock() { - Ok(rx) => rx.recv(), - Err(e) => { - error!("Inserter {} failed to lock receiver: {}", inserter_id, e); - break; - } + let batch = match recv_shared(&result_rx_clone, |rx| rx.recv()).await { + Some(batch) => batch, + None => { + error!("Inserter {} failed to lock receiver", inserter_id); + break; } }; @@ -1141,6 +1163,56 @@ pub async fn index_git_commits( mod tests { use super::*; + /// Two inserter-style tasks waiting on a shared receiver whose sender + /// stays alive must leave the runtime's two workers free for other work. + /// The test body runs on the `block_on` thread, not on a worker, and + /// waits with std primitives, so it still fails (rather than hangs) if + /// the waiters do occupy both workers. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn recv_shared_does_not_hold_workers() { + let (tx, rx) = mpsc::channel::(); + let rx = Arc::new(std::sync::Mutex::new(rx)); + let entered = Arc::new(AtomicUsize::new(0)); + + let waiters: Vec<_> = (0..2) + .map(|_| { + let rx = Arc::clone(&rx); + let entered = Arc::clone(&entered); + tokio::spawn(async move { + entered.fetch_add(1, Ordering::SeqCst); + recv_shared(&rx, |rx| rx.recv()).await + }) + }) + .collect(); + + // Each waiter increments `entered` on a worker right before it + // starts waiting. Once both have, waiters that block their workers + // would now hold both. + let deadline = std::time::Instant::now() + std::time::Duration::from_secs(10); + while entered.load(Ordering::SeqCst) < 2 { + assert!( + std::time::Instant::now() < deadline, + "waiters never started" + ); + std::thread::sleep(std::time::Duration::from_millis(1)); + } + + let (done_tx, done_rx) = mpsc::channel(); + tokio::spawn(async move { + let _ = done_tx.send(7); + }); + assert_eq!( + done_rx.recv_timeout(std::time::Duration::from_secs(10)), + Ok(7), + "a task spawned while two receivers wait never ran" + ); + + drop(tx); + for waiter in waiters { + assert_eq!(waiter.await.unwrap(), Some(Err(mpsc::RecvError))); + } + } + /// Helper: run a git command in a directory, panic on failure. fn git(repo: &std::path::Path, args: &[&str]) -> String { let out = std::process::Command::new("git")