perf(core-py): share one Tokio runtime across session contexts

[ttw225]

Problem

Each context previously added about 18 runtime worker threads, producing
an approximately 1 + 18 × contexts native-thread count on this machine.
With the shared runtime, the count stays flat:

Measured on an 18-core macOS machine:

contexts	before	after
1	19	19
4	73	19
8	145	19

Changes

• Share one process-wide Tokio runtime across all session contexts.
• Rebuild the runtime lazily when first used in a forked child.
• Preserve runtime creation errors as Python exceptions.
• Add a POSIX-only fork regression test covering contexts created both
  before and after the fork.
• Add a parent-side timeout so deadlock regressions fail instead of
  hanging CI.

A forked child inherits the runtime handle but not its worker threads.
The inherited handle is deliberately leaked because dropping it may
wait for workers that were not copied by fork().

The Python API is unchanged. Runtime usage outside wren-core-py is
not affected.

Related work

#2485 releases the GIL around Rust execution while retaining the
per-context call_lock.

#2495 fixes a separate catalog-list race exposed by concurrent
transforms. It is now merged and complements this PR: #2495 isolates
catalog state, while this PR bounds Tokio worker growth.

#2504 tracks the remaining Python-side catalog documentation,
integration coverage, and eventual call_lock removal. This PR does
not change catalog semantics or close that issue.

Verification

• cargo test — 37 passed
• cargo test --no-default-features — 37 passed
• cargo clippy --all-targets --all-features -- -D warnings
• release extension built successfully with maturin
• Python test suite — 35 passed
• concurrent-call and fork tests passed across 10 repeated rounds
• native thread count remained flat at 19 for up to 8 contexts

Summary by CodeRabbit

• Bug Fixes

  • Improved session reliability after process forking by ensuring the async runtime state is correctly handled for the child process.
  • Confirmed SQL planning/execution paths remain functional before and after a fork, avoiding potential stale-state hangs.

• Tests

  • Added a regression test that validates session behavior across fork() on POSIX systems, including deadline-based polling, timeout handling, and exit-status reporting.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro

Run ID: 999e2bb5-282b-44b8-beef-0ce4028502c7

📥 Commits

Reviewing files that changed from the base of the PR and between db627e8 and 8a6ba50.

📒 Files selected for processing (2)

• core/wren-core-py/src/context.rs
• core/wren-core-py/tests/test_modeling_core.py

🚧 Files skipped from review as they are similar to previous changes (2)

• core/wren-core-py/tests/test_modeling_core.py
• core/wren-core-py/src/context.rs

---

Walkthrough

PySessionContext now uses a PID-aware shared Tokio runtime instead of storing one runtime per context. Blocking operations and registration paths use it, while a POSIX regression test verifies runtime recreation after fork().

Changes

Runtime migration

Layer / File(s)	Summary
Shared runtime and context construction core/wren-core-py/src/context.rs	Adds PID-aware shared runtime creation and removes per-context runtime storage from PySessionContext initialization.
Blocking API operations core/wren-core-py/src/context.rs	Routes SQL transformation, function lookup, queries, table registration, and explain operations through the shared runtime while retaining call locking.
MDL registration and fork regression coverage core/wren-core-py/src/context.rs, core/wren-core-py/tests/test_modeling_core.py	Updates MDL and remote-function registration runtime usage and adds a POSIX fork regression test.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Sequence Diagram(s)

sequenceDiagram
  participant SessionContext
  participant shared_runtime
  participant TokioRuntime
  participant ForkedProcess
  SessionContext->>shared_runtime: request runtime for transform_sql
  shared_runtime->>TokioRuntime: execute blocking SQL transformation
  TokioRuntime-->>SessionContext: return transformed SQL
  SessionContext->>ForkedProcess: fork process
  ForkedProcess->>shared_runtime: request runtime after fork
  shared_runtime->>TokioRuntime: recreate runtime for child PID
  TokioRuntime-->>ForkedProcess: return working runtime
  ForkedProcess->>SessionContext: transform SQL in child

Loading

Possibly related PRs

• Canner/WrenAI#2485: Refactors overlapping PySessionContext blocking-call and runtime handling.

Poem

A bunny hops where runtimes share,
Forked little processes find fresh air.
SQL transforms, then queries run,
New runtimes bloom beneath the sun.
Locks stay snug, the tests all cheer—
“No hanging hops in here!”

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title accurately summarizes the main change: sharing one Tokio runtime across SessionContext instances for performance.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

🧹 Nitpick comments (1)

core/wren-core-py/src/context.rs (1)

130-152: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Defer fetching the shared runtime.

The runtime variable is fetched at the beginning of the new method but remains unused in the None branch (lines 132-152) since register_function_by_data_source fetches its own runtime reference internally. To slightly reduce initialization overhead, consider deferring the shared_runtime()? fetch to where it is actually needed inside the Ok branch of AnalyzedWrenMDL::analyze.

♻️ Proposed refactor

Remove the early fetch:

     ) -> PyResult<Self> {
-        let runtime = shared_runtime()?;
-
         let Some(mdl_base64) = mdl_base64 else {

And move it inside the Ok block around line 219:

             match AnalyzedWrenMDL::analyze(
                 manifest,
                 Arc::clone(&properties_ref),
                 mdl::context::Mode::Unparse,
             ) {
                 Ok(analyzed_mdl) => {
                     let analyzed_mdl = Arc::new(analyzed_mdl);
+                    let runtime = shared_runtime()?;
                     let unparser_ctx = py
                         .detach(|| {
                             runtime.block_on(apply_wren_on_ctx(

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@core/wren-core-py/src/context.rs` around lines 130 - 152, Defer the shared
runtime lookup in the `new` method: remove the early `shared_runtime()?`
assignment before the `mdl_base64` branch, then fetch it only in the
`Some`/successful analyzed-MDL path where `runtime` is required. Keep the `None`
branch unchanged and preserve existing error propagation.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@core/wren-core-py/src/context.rs`:
- Around line 130-152: Defer the shared runtime lookup in the `new` method:
remove the early `shared_runtime()?` assignment before the `mdl_base64` branch,
then fetch it only in the `Some`/successful analyzed-MDL path where `runtime` is
required. Keep the `None` branch unchanged and preserve existing error
propagation.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro

Run ID: ee2f7b58-a99d-4157-9fc9-f7389dd02b24

📥 Commits

Reviewing files that changed from the base of the PR and between 283d089 and db627e8.

📒 Files selected for processing (2)

• core/wren-core-py/src/context.rs
• core/wren-core-py/tests/test_modeling_core.py

[ttw225]

🧹 Nitpick comments (1)

core/wren-core-py/src/context.rs (1)> 130-152: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Defer fetching the shared runtime.
The runtime variable is fetched at the beginning of the new method but remains unused in the None branch (lines 132-152) since register_function_by_data_source fetches its own runtime reference internally. To slightly reduce initialization overhead, consider deferring the shared_runtime()? fetch to where it is actually needed inside the Ok branch of AnalyzedWrenMDL::analyze.

Skipped. Although the local runtime binding is only used by the analyzed-MDL path, the no-manifest path may still require the shared runtime through register_function_by_data_source.

Keeping the eager lookup also preserves fail-fast runtime creation errors during Python context construction. Deferring it would only avoid initialization for a subset of paths, such as BigQuery, while changing when those errors are reported.

[goldmedal]

Thanks @ttw225, nice catch 👍 , thanks for working on the concurrent case.