Pure style fix, no behavior change -- CI's clippy run (-D warnings)
rejected the nested if in get_runtime().
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Review on this PR found three more real bugs in the previous commit's
exit-lifecycle design, all traceable to the same root cause: using
Arc::strong_count as a proxy for "is anything still using the runtime"
is the wrong signal. It reflects who currently holds a reference, not
who's logically still in flight, and it let a spawned task become the
final, worker-thread-side owner of the Runtime -- which reproduced
unprompted in this branch's own test suite (a panic: "cannot drop a
runtime in a context where blocking is not allowed").
Specifically:
- A task holding an Arc for its own whole lifetime could end up being
the one to drop the Runtime when it completes, which panics if that
drop happens on one of the runtime's own worker threads.
- shutdown() bailed out as soon as it saw an empty slot, so an install
already in progress at that exact moment could finish afterward,
leaving a live, un-waited-for runtime behind as the interpreter moved
on into teardown.
- Once the bounded wait elapsed, shutdown() simply returned without
touching the runtime at all, recreating the original race for any
call slower than the grace period.
Replaces Arc-based tracking with an explicit OUTSTANDING counter that
every top-level spawn/spawn_blocking/block_on call increments before it
starts and decrements only once it has truly finished, via an
OutstandingGuard token that carries no reference to the runtime itself.
shutdown() waits for this to reach zero (not for Arc::strong_count to
look low) before ever touching the runtime, and forces the retirement
attempt anyway once its bound elapses instead of leaving the runtime
untouched.
Also narrows what the fork-child handler is allowed to do. ArcSwapOption
turned out not to be safe to call from there either: swap/compare_and_swap
reconcile reader "debts" internally (thread-local state, potentially an
allocation), none of which is safe in a forked child that may have
inherited another thread's lock mid-acquisition. The handler now only
bumps a bare GENERATION counter; get_runtime() compares its installed
runtime's generation against the live counter from ordinary (non-signal)
context and rebuilds on a mismatch, exactly the check the handler used
to perform directly.
Verified against the same concurrent reproducer as the previous commit:
50/50 clean runs (vs. hangs and, separately, an unprompted panic under
the prior design), plus a 30-second-join variant and a targeted test
forcing the install-vs-shutdown race directly.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
The previous commit's shutdown() freed the shared Tokio runtime as soon
as it looked unused (an AtomicPtr swap + Box::from_raw), but callers
only ever held a bare reference into it, with no way to tell shutdown
they were still using it. A concurrent caller could dereference memory
shutdown had already freed.
Switching to an ArcSwapOption<Runtime> so every caller holds an owned
Arc clone for as long as it's using the runtime fixes the
use-after-free, but isn't sufficient on its own, for two reasons:
- future_into_py spawns a task that, once running, spawns a second one
to do the real work and awaits its JoinHandle. If both nested calls
independently re-resolved "the current runtime" via get_runtime(), a
reclaim landing between them could bind the outer and inner task to
two different runtime instances, permanently stalling the outer task
if its own runtime was the one torn down while it awaited the inner
one.
- Runtime::shutdown_timeout gives spawned (non-blocking) tasks no
grace period at all: a task keeps running until it next yields, then
is dropped. An outermost caller that only held its Arc for the
instant it submitted a task let the runtime's reference count fall
back to baseline immediately, long before the task itself finished,
so shutdown could reclaim the runtime while a top-level task was
still genuinely in flight, abandoning it silently.
spawn/spawn_blocking now use Handle::try_current() to pin any nested
call to whatever runtime is already executing it, and the outermost
call moves its Arc clone into the spawned task itself so the reference
count stays elevated for the task's whole lifetime, not just its
submission.
Verified with two new targeted unit tests plus the existing concurrent
stress test, and empirically against a concurrent reproducer built
from this PR's own review feedback: 100+ clean runs with no hangs or
crashes, versus a consistent ~35-55% failure rate beforehand.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
The Python bindings route every async call through a single, lazily
created Tokio runtime, stored behind a raw pointer so it can survive
being reset (not dropped) across fork() -- dropping it there would
hang joining now-dead worker threads. That leak-on-fork path is
deliberate and documented, but normal process exit had no equivalent
handling at all: nothing ever tells the runtime to shut down, so its
worker threads keep running, completely uncoordinated with the
interpreter, until the process ends.
If a worker thread is mid-task exactly as Py_Finalize begins tearing
down interpreter state, it can panic on state that's already gone. That
panic happens on a background thread with no PyO3-wrapped call frame to
catch it -- and if it lands while another unwind is already in
progress, Rust aborts the whole process outright rather than merely
failing that one call. This reproduces as an intermittent SIGABRT on
otherwise-successful short-lived processes that open a connection, do a
handful of operations, and exit normally.
Give the runtime a coordinated, bounded shutdown at the one point it
was missing one: register a Python atexit callback that calls
runtime::shutdown(), which takes the runtime through the same atomic
slot the fork handler uses and calls shutdown_timeout() rather than a
bare drop -- bounded, so a wedged background task can never hang
process exit, but real in-flight work gets a genuine chance to finish
first. Runs with the GIL released (Python::detach), since the wait
would otherwise be able to deadlock against any in-flight task that
itself needs the GIL to complete.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
By default the read freshness provider was not included in the namespace
client, preventing the read freshness headers from being included in the
request. This prevents checkout_latest() from working as expected when
using the namespace client.
This fix ensures the provided is built into the client when the
namespace impl and properties are provided.
## Summary
PyTorch's `DataLoader` uses fork-based multiprocessing by default on
Linux, but threads do not survive `fork()`. LanceDB's Python bindings
drive async work through two threaded layers, both of which become inert
in a forked child:
- `BackgroundEventLoop` runs an asyncio loop on a Python
`threading.Thread`.
- `pyo3-async-runtimes::tokio` holds a global multi-threaded tokio
runtime whose worker threads also die on fork — and its runtime lives in
a `OnceLock` that cannot be replaced after first use.
As a result, any `Permutation` (or other async API) used inside a
fork-based `DataLoader` worker hangs indefinitely. This PR makes both
layers fork-safe so `Permutation` works as a `torch.utils.data.Dataset`
with `num_workers > 0`.
## Approach
### Rust — new `python/src/runtime.rs`
Mirrors the pattern used in [Lance's Python
bindings](https://github.com/lance-format/lance/blob/456198cd6f42be07f99617a6d7e39d6209cdf3cc/python/src/lib.rs#L139),
adapted for the async-bridge use case.
- `LanceRuntime` implements `pyo3_async_runtimes::generic::Runtime +
ContextExt`, backed by an `AtomicPtr<tokio::runtime::Runtime>` we own
(sidestepping `pyo3-async-runtimes`'s frozen `OnceLock` global).
- A `pthread_atfork(after_in_child)` handler nulls the pointer; the next
`spawn` rebuilds the runtime in the child. The previous runtime is
intentionally **leaked** — calling `Drop` would try to join now-dead
worker threads and hang.
- `runtime::future_into_py` is a drop-in for
`pyo3_async_runtimes::tokio::future_into_py`. All ~80 call sites in
`arrow.rs` / `connection.rs` / `permutation.rs` / `query.rs` /
`table.rs` are updated to route through it.
- `python/Cargo.toml` adds `libc = "0.2"` and the tokio
`rt-multi-thread` feature.
### Python — `lancedb/background_loop.py`
- Refactors `BackgroundEventLoop.__init__` to a reusable `_start()`
method.
- An `os.register_at_fork(after_in_child=…)` hook calls `LOOP._start()`
to give the singleton a fresh asyncio loop and thread **in place**. This
matters because the rest of the codebase imports `LOOP` via `from
.background_loop import LOOP` — rebinding the module attribute would
leave those references holding the dead loop.
### Python — `lancedb/__init__.py`
Removes the `__warn_on_fork` pre-fork warning (and the now-unused
`import warnings`). Fork is supported.
## Test plan
- [x] New `test_permutation_dataloader_fork_workers` in
`python/tests/test_torch.py`: runs a `Permutation` through
`torch.utils.data.DataLoader(num_workers=2,
multiprocessing_context="fork")` inside a spawn-isolated child with a
30s hang detector. **Pre-fix**: timed out at 36s. **Post-fix**: passes
in ~3.6s.
- [x] New `test_remote_connection_after_fork` in
`python/tests/test_remote_db.py`: forks a child that creates a fresh
`lancedb.connect(...)` against a mock HTTP server and calls
`table_names()`; passes in <1s, validates the runtime reset is
sufficient for fresh remote clients.
- [x] All 62 tests in `test_torch.py` + `test_permutation.py` pass.
- [x] All 35 tests in `test_remote_db.py` pass.
- [x] `test_table.py` (87) + `test_db.py` + `test_query.py` (157, minus
one unrelated `sentence_transformers` import skip) — 244 passing.
- [x] `cargo clippy -p lancedb-python --tests` clean.
- [x] `cargo fmt`, `ruff check`, `ruff format` all clean.
## Known limitation (follow-up)
This PR makes a **freshly-built** `lancedb.connect(...)` work in a
forked child. An **inherited** `Connection` from the parent still
carries an inherited `reqwest::Client` whose hyper connection pool
references socket FDs and TCP/TLS state shared with the parent — using
it from the child after fork is unsafe (especially with HTTP/1.1
keep-alive). The recommended pattern for fork-based `DataLoader` workers
that hit a remote DB is to construct a new connection inside the worker.
Auto-clearing inherited HTTP client pools on fork would require tracking
live `Connection` instances in `lancedb` core and is left for a
follow-up PR.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
---------
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
ChilePiquinandClaude Sonnet 5
Ryan Green
Weston PaceandClaude Opus 4.7