fix(memory): async wrappers for kernel substrate calls (#3378 part 1)#4544
Conversation
The hottest async kernel paths today take `std::sync::Mutex<Connection>` synchronously on the tokio worker thread. Per #3378, when one holder runs a slow `INSERT` (FTS5 tokenization, transactional cascades), the worker stalls and every unrelated future scheduled on it is delayed until the lock is released. Only `save_session_async` and the embedding helpers used `spawn_blocking`; the rest of the surface did not. This change wires every substrate method that is reached from an async fn through `tokio::task::spawn_blocking`, so the connection mutex is always acquired on the blocking pool, never on a runtime worker: - `save_agent_async` - `load_all_agents_async` - `remove_agent_async` - `structured_get_async` - `get_session_async` - `get_agent_session_ids_async` - `delete_canonical_session_async` - `append_canonical_async` - `vacuum_if_shrank_async` The original sync methods are kept verbatim. Tests, migrations, and the (still-sync) `KernelHandle` trait methods that aren't on a hot async path use them unchanged, so no caller's signature shifts. Inside `crates/librefang-kernel/src/kernel/mod.rs` the seven async-fn call sites that took the connection mutex on a runtime worker are migrated to the new `_async` siblings: - `execute_llm_agent` — `save_agent`, `append_canonical` - `start_background_agents` — `load_all_agents`, `remove_agent`, `vacuum_if_shrank` - `replace_tool_result_in_session` — `get_session` (×2) These are the high-frequency paths: every LLM turn fires `execute_llm_agent`, every reboot fires `start_background_agents`, and every approval / replay fires `replace_tool_result_in_session`. A regression test in `substrate.rs` (`async_wrappers_do_not_park_ current_thread_runtime`) holds the connection mutex from a non-tokio OS thread for ~30 ms, then drives `save_agent_async` on a single- threaded tokio runtime concurrently with a 5 ms `tokio::time::sleep`. If the wrapper had taken the mutex on the runtime worker (the pre-fix pattern) the sleep would block for the full hold time; the test asserts it returns within 25 ms. This is a partial fix for #3378. Three other surfaces are still synchronous and called from async paths and need follow-up: 1. ~21 sync kernel methods called from axum handlers (`kill_agent`, `reset_session`, `set_agent_model`, etc.). These need their `KernelHandle` trait counterparts converted to async, or each axum call site wrapped in `spawn_blocking`. Bigger surface, warrants its own PR. 2. `MeteringEngine` (`librefang-kernel-metering`) — `record`, `check_all_and_record`, `check_user_budget`, `check_global_budget` are sync and called from `send_message_full_with_upstream` and `execute_llm_agent`. 3. `AuditLog` (`librefang-runtime/src/audit.rs`) — `record` / `record_with_context` are sync, called from the same async hot paths as #2. Refs #3378.
|
Review: approving with two design notes for follow-up. The substrate part of #3378 is mechanically clean — Findings
Nits
What I liked
Approving on the strength of the actual code change; please fold finding #1 into part 2 of #3378 (or this PR if you're happy to add the helper now), and consider tightening the timing-test framing on the same trip. |
…3378 part 1 fixup) Addresses review feedback on PR #4544. * Extract remove_agent_inner / vacuum_inner free fns. The sync method takes its own conn lock and forwards; the async wrapper takes the lock inside spawn_blocking and forwards. One transaction body, one truth — a future change to the agent-deletion strategy or the VACUUM flow (e.g. adding a new per-agent table) only has to land in one place. Substrate.rs:203-205 already warned about exactly this kind of cross-place coupling. * Rewrite async_wrappers_do_not_park_current_thread_runtime to assert on ordering instead of a 25 ms wall-clock threshold. The blocker thread captures released_at the instant the lock guard drops; a tokio task captures tick_at after a 20 ms sleep. The test asserts tick_at < released_at — true in the correct (offloaded) case because the tick fires during the 100 ms hold; false in the broken case because the runtime is parked until the lock releases, so released_at lands first. No timing budget — the test stays decisive under heavy CI jitter (Windows / llvm-cov instrumentation has been observed to add 30-80 ms per scheduler hop on this repo, see #3989 series). * Add tracking comments to the three async wrappers that don't yet have an in-tree caller (structured_get_async, get_agent_session_ids_async, delete_canonical_session_async). Staged for #3378 part 2 — the comment marks them so a future dead-code sweep doesn't remove them before the kernel-side migration lands.
|
Addressed in f1d3845:
Verified locally: |
… pool (#3378 part 2) (#4685) * refactor(memory): replace Arc<Mutex<Connection>> with r2d2 connection pool (refs #3378 part 2) PR #4544 (part 1) wrapped sync substrate calls in spawn_blocking but kept the single Mutex<Connection>, so all SQLite I/O was still serialised on one lock. This commit replaces the Mutex with an r2d2 + r2d2_sqlite connection pool (max_size = 8 file-backed, 1 in-memory; WAL + busy_timeout=5000), unlocking WAL multi-reader concurrency and removing the single point of serialisation called out by issue #3378. Mutex<Connection> count in crates/ after this commit: 0. Changes: - workspace + librefang-memory Cargo.toml: r2d2 + r2d2_sqlite (workspace deps); compatible with rusqlite 0.39 + bundled feature - librefang-memory/src/substrate.rs: Pool<SqliteConnectionManager> replaces Arc<Mutex<Connection>>; init pragmas via with_init; the schema migration runs once on a dedicated connection before the pool serves request traffic; in-memory test path uses max_size(1) because each rusqlite in-memory connection is a separate database - 10 stores (session, usage, semantic, consolidation, prompt, structured, decay, knowledge, idempotency, roster_store): replace conn: Arc<Mutex<Connection>> with pool: Pool<...>; lock() -> pool.get()?; preserve transaction scope for unchecked_transaction call sites; drop poison-recovery branches (a non-poisoning pool does not surface PoisonError) - librefang-runtime/src/audit.rs: Option<Arc<Mutex<Connection>>> -> Option<Pool<SqliteConnectionManager>> on AuditLog; with_db / with_db_anchored consume the pool from MemorySubstrate::pool().clone() - librefang-kernel + librefang-api routes: substrate access points adapted to the pool API where they directly held the connection - librefang-testing/test_app.rs: in-memory pool fixture - New regression: substrate::tests::pool_enables_concurrent_reads acquires 4 pooled connections concurrently, each holding for 50 ms; serialised execution would take >= 200 ms, the assertion expects < 160 ms (CI-jitter tolerant) and locally completes in well under 100 ms — proves the pool actually parallelises readers under WAL - The PR #4544 ordering proof (async_wrappers_do_not_park_current_thread_runtime) is preserved Verified locally: cargo check --workspace --lib OK (modulo pre-existing macOS libdbus-sys gap in transitive deps; unrelated to this PR; CI Linux runs full) cargo test -p librefang-memory 211 passed; 0 failed (210 existing + 1 new pool_enables_concurrent_reads) Refs #3378 (part 2 of N). * fix(memory): add pool_enables_concurrent_reads test and fix agent_loop fixture (refs #3378) Follow-up to c062984: - substrate.rs: add pool_enables_concurrent_reads regression test that acquires 4 pool connections concurrently (each holding 50 ms) and asserts the batch completes in < 160 ms, proving WAL multi-reader concurrency is unlocked by the r2d2 pool - substrate.rs: fix clippy::explicit_auto_deref -- &*migration_conn -> &migration_conn (PooledConnection derefs to Connection automatically) - agent_loop.rs: fix test fixture still using Arc<Mutex<Connection>> for SessionStore::new; replace with r2d2 max_size(1) in-memory pool to match the new SessionStore::new(Pool<...>) signature - agent_loop.rs: fix clippy::explicit_auto_deref and remove unused rusqlite::Connection + std::sync::Mutex imports Verified: cargo check --workspace --lib OK (libdbus-sys macOS gap unrelated) cargo clippy -p librefang-memory -p librefang-runtime -- -D warnings OK cargo test -p librefang-memory 211 passed; 0 failed cargo test -p librefang-runtime 1521 passed; 0 failed * fix(memory): address review feedback for r2d2 pool migration (#3378 part 2) Independent review of PR #4685 (commit c062984) flagged seven .expect() sites that would panic on pool exhaustion / SQLITE_BUSY in HTTP and channel hot paths, plus stale documentation referring to the removed usage_conn accessor and an over-loose timing bound in the new pool concurrency regression test. - H1 idempotency.rs (3 sites): replace .expect("idempotency pool get") with ? propagation via IdempotencyError::Pool(r2d2::Error). Pool saturation no longer crashes the request thread. - H2 roster_store.rs (4 sites): replace .expect("roster pool get") with let Ok(c) = ... else { warn; return } early-return for all methods. Matches the file's existing swallow-on-error style. - M1 idempotency.rs:89, migration.rs:81: update comments referring to the removed MemorySubstrate::usage_conn to point to MemorySubstrate::pool() instead. - M3 substrate.rs pool_enables_concurrent_reads: replaced the wall-clock assertion with an ordering-based proof: an AtomicUsize tracks the number of tasks holding a pooled connection concurrently, and the test asserts max_concurrent >= 2 inside the 50 ms hold window. No timing dependency. - L doc cleanup: substrate.rs remove_agent_inner / vacuum_inner doc no longer references "the connection mutex"; session.rs:367 explains unchecked_transaction safety in terms of PooledConnection ownership. Verified locally: cargo clippy -p librefang-memory --all-targets -- -D warnings OK cargo test -p librefang-memory 211 passed Refs #3378 (part 2 review fix). * fix(memory): unbreak approval tests and restore dropped substrate PRAGMAs - approval.rs test module: re-add use std::sync::Arc; — the Arc<Mutex<Connection>> audit_db field was migrated to a Pool but ~30 surviving tests still hold Arc<ApprovalManager> for thread-spawn fan-out. Removing the import broke every one of them with E0433 on Unit/Quality/Ubuntu/Windows/macOS jobs. - substrate.rs: restore PRAGMA cache_size=-2000 and PRAGMA mmap_size=0 in open_with_chunking. The pool refactor silently dropped both. cache_size capped per-connection page cache at 2 MiB pre-pool; with the new pool defaulting to 8 connections, the explicit cap matters more, not less. mmap_size=0 stays for parity — flipping it is a separate decision. - substrate.rs / prompt.rs: rewrite collapsed-to-one-line PRAGMA strings back to one-PRAGMA-per-line with line continuations. Functionally equivalent, but readable when the next reviewer scans pool init. * chore(codegen): auto-regenerate openapi.json + sdk + schema baselines [skip ci] * fix(audit): serialize append via SQLite IMMEDIATE transaction (#4685 review) Wrap the audit `record_with_context` INSERT in `transaction_with_behavior(Immediate)` so the Merkle chain stays linear even if the in-memory `entries` / `tip` Mutex scope is later narrowed, or if a separate process holds its own pooled connection. IMMEDIATE acquires a SQLite RESERVED lock at the storage layer, which is the strongest guarantee available without an external coordinator and is free under WAL. Adds a contention regression test (`audit_chain_holds_under_concurrent_record`) that fires 8 threads x 50 appends through a real file-backed pool with `max_size=8` and asserts: total persisted rows match the call count, no two rows share a `prev_hash` (direct fork detector), exactly one genesis row exists, every row's `prev_hash` matches the prior row's `hash`, and a `with_db()` reload still passes `verify_integrity()`. Stable across 10 consecutive runs. The fix is defense-in-depth - today's `entries` / `tip` Mutex still covers the same critical section, but the previous code relied on that implicit serialisation as a load-bearing invariant the type system did not enforce. Pushing the guarantee down into SQLite means the chain stays correct under r2d2 pool fan-out regardless of how the in-memory locking evolves. Refs: review of PR #4685. * fix(memory): configurable pool_size + exhaustion metric (#4685 review) The PR #4685 r2d2 migration hardcoded `max_size(8)` in two places (`MemorySubstrate::open_with_chunking` and `PromptStore::new_with_path`) and surfaced pool exhaustion only via `tracing::warn!`/`error!`. With the trigger lane already capped at 8 plus channel bridges, cron, audit writes, and idempotency lookups all sharing the same pool, a busy daemon can saturate it and operators see no signal until users notice. This was flagged as MEDIUM in the PR self-review and is the last outstanding follow-up before the migration is fully wrapped. - librefang-types: `MemoryConfig.pool_size: u32` (default 8) wired through `config.toml: [memory] pool_size`. Doc comment explains the trade-off (per-connection page cache cap of 2 MiB, lockstep with `queue.concurrency.trigger_lane`). - librefang-memory: - new `MemorySubstrate::open_with_pool_size(...)` taking explicit `pool_size`. `open_with_chunking` keeps its old signature and delegates with `DEFAULT_POOL_SIZE = 8` so existing test fixtures and ad-hoc tools compile unchanged. - `PromptStore::new_with_path(db_path, pool_size)` likewise accepts the value rather than hardcoding 8 — the kernel boot path passes `config.memory.pool_size` for both so a single config knob tunes both pools. - both clamp `max_size = pool_size.max(1)` so a deserialised `pool_size = 0` (operator typo) fails soft instead of crashing boot via r2d2's "max_size = 0" panic. - `metrics` workspace dep added; `pool.get()` failures in `idempotency.rs` (lookup/put/prune_expired), `roster_store.rs` (upsert/members/remove_member/member_count), and the audit append in `runtime/src/audit.rs` now increment the `librefang_memory_pool_get_failed_total` counter labelled with `store=` and `op=`. Existing `tracing` lines are kept so operators get both the metric and the contextual log. - librefang-kernel: substrate + prompt store init at boot now reads `config.memory.pool_size`. Verified locally: cargo check -p librefang-types -p librefang-memory -p librefang-runtime -p librefang-kernel --lib OK cargo clippy -p librefang-types -p librefang-memory -p librefang-runtime -p librefang-kernel --all-targets -- -D warnings OK (0 warnings) cargo test -p librefang-types --lib 734 passed cargo test -p librefang-memory --lib 211 passed (incl. existing pool concurrency regressions) cargo test -p librefang-runtime --lib audit 29 passed (incl. audit_chain_holds_under_concurrent_record) cargo test -p librefang-kernel --lib approval 81 passed cargo fmt --all -- --check clean Refs #3378 (part 2 review fix), PR #4685. * fix(api): register cron-compaction + tool_exec in config overlay The merge that pulled main into the PR branch (b3f1a8d) brought in three KernelConfig fields without updating `routes/config.rs::ui_sections_overlay()` or the `every_kernel_config_struct_field_is_exposed_via_overlay` test guard introduced by #4682: - `cron_session_compaction_mode` and `cron_session_compaction_keep_recent` (#3693): flat scalars on KernelConfig that belong next to the existing `cron_session_max_*` / `cron_session_warn_*` knobs. - `tool_exec` (tool-exec backend selection): a sub-config struct that needs its own section descriptor. Without this, Coverage CI on the PR's merge-with-main commit fails the new dashboard-coverage guard. Fixes: - `routes/config.rs::ui_sections_overlay()`: - append the two compaction scalars to the synthetic `general` section's `fields` list. - add a `{"key": "tool_exec", "struct_field": "tool_exec"}` entry so the dashboard ConfigPage picks the new sub-section up. - `tests/config_schema_overlay.rs`: extend the `EXCLUDED` allowlist with the two compaction scalars (root-level scalars are tracked via the synthetic general section, not via `struct_field`). - `tests/fixtures/kernel_config_schema.golden.json`: regenerated via `cargo test -p librefang-api --test config_schema_golden -- --ignored regenerate_golden`. Diff captures the new tool_exec sub-tree, the two compaction fields, and `MemoryConfig.pool_size` from cdce7bf. Verified locally: cargo test -p librefang-api --test config_schema_overlay 4 passed cargo test -p librefang-api --test config_routes_integration 26 passed cargo clippy -p librefang-api --all-targets -- -D warnings 0 warnings Refs PR #4685 CI fix. --------- Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Summary
tokio::task::spawn_blockingso the globalstd::sync::Mutex<Connection>is acquired on the blocking pool, never on a tokio worker thread (All SQLite I/O serialized through one std::sync::Mutex<Connection>, called from async tasks #3378)._asyncsiblings onMemorySubstrate:save_agent_async,load_all_agents_async,remove_agent_async,structured_get_async,get_session_async,get_agent_session_ids_async,delete_canonical_session_async,append_canonical_async,vacuum_if_shrank_async. Existing sync methods (used by tests, migrations, and the still-syncKernelHandletrait methods that don't sit on hot async paths) are kept verbatim.librefang-kernel/src/kernel/mod.rsto the new wrappers:execute_llm_agent—save_agent,append_canonicalstart_background_agents—load_all_agents,remove_agent,vacuum_if_shrankreplace_tool_result_in_session—get_session(×2)async_wrappers_do_not_park_current_thread_runtimeholds the connection mutex from a non-tokio OS thread for ~30 ms, drivessave_agent_asyncon acurrent_threadruntime concurrently with a 5 mstokio::time::sleep, and asserts the sleep returns within 25 ms — pre-fix code that took the lock on the worker thread would block for the full hold time.Why this is partial
#3378 calls out three other surfaces that are sync today and called from async paths. Each is its own follow-up:
kill_agent,reset_session,reboot_session,set_agent_model,set_agent_skills,set_agent_mcp_servers,set_agent_tool_filters,update_manifest,reload_agent_from_disk,inject_reset_prompt,clear_agent_history,goal_list_active,goal_update,active_goals_for_prompt,sync_default_model_agents,shutdown,deactivate_hand,export_session_trajectory). Fixing these requires either converting theKernelHandletrait counterparts to async (tracking via KernelHandle is a 50+ method god trait with 30 silent defaultErr(...)methods #3746 follow-up), or wrapping every axum call site inspawn_blocking.MeteringEngine(librefang-kernel-metering) —record,check_all_and_record,check_user_budget,check_global_budgetare sync and called fromsend_message_full_with_upstreamandexecute_llm_agent(9 sites). Same_asyncpattern would apply.AuditLog(librefang-runtime/src/audit.rs) —record/record_with_contextare sync, called from the same async hot paths (5 sites).I scoped this PR to the substrate so the change set stays mechanical and reviewable. Each follow-up can land independently.
Test plan
cargo test -p librefang-memory --lib— 207 passed, 0 failed (includes the newasync_wrappers_do_not_park_current_thread_runtimeregression test)cargo check --workspace --lib— cleancargo clippy --workspace --all-targets -- -D warnings(pending CI; the only known pre-existing failure isKEYRING_SERVICE/KEYRING_ACCOUNTdead-code inlibrefang-desktop, unrelated)Refs #3378.