Skip to content

fix(db): lock sessions individually so SQL tabs execute concurrently - #264

Merged
feigeCode merged 2 commits into
mainfrom
fix/db-session-lock-per-session
Sep 21, 2026
Merged

feigeCode merged 2 commits into
mainfrom
fix/db-session-lock-per-session

Conversation

@feigeCode

@feigeCode feigeCode commented Sep 21, 2026

Copy link
Copy Markdown
Owner

Description

Problem. Two SQL tabs could not execute concurrently. While tab A was running a statement, tab B was blocked — even when it was connected to a different database — and metadata / table-data / config reads queued behind it as well.

Root cause. ConnectionManager stored every session in one global map behind a single RwLock. get_session_connection() took the write lock on that map and handed it to SessionConnectionGuard, which held it for the whole lifetime of a statement (conn.execute(..).await, conn.execute_streaming(..).await, …) — about 25 call sites followed this pattern. Since tokio::sync::RwLock is fair, the queued writer also blocked subsequent readers, so unrelated sessions serialized on the same lock. Lock granularity was the only issue: the driver layer has no global lock, and try_acquire_session already gives each session its own physical connection.

Fix — lock per session instead of per manager.

  • ConnectionSession.connection becomes Arc<AsyncMutex<Box<dyn DbConnection + Send + Sync>>>, so a statement only locks the connection it uses.
  • New SessionState behind its own StdMutex; bookkeeping reads no longer wait for an in-flight statement. All ConnectionSession methods now take &self.
  • The session map holds Arc<ConnectionSession>; get_session_connection() holds the map lock only to look the session up and claim it, then releases it. SessionConnectionGuard now owns that session's connection lock and no longer needs a lifetime parameter, so the existing call sites are unchanged.
  • Teardown paths (close_session, remove_all_sessions, cleanup_expired_sessions, release_session_internal) now detach the session under the map lock and close it outside it, so disconnecting one connection no longer stalls every other tab.
  • Preserved semantics: the same session stays exclusive (a physical connection is still serialized), busy/transactional sessions are never reclaimed, and cancel/close still wait for the running statement to finish.

Review follow-up (commit 24e87f8)

The first revision fixed the guard granularity but left the session-pool paths inconsistent. Three issues were raised in review; all three are fixed in 24e87f862, together with the lock-order invariant they depend on.

1. [P1] Reuse still waited on a connection lock while holding the session table lock. try_acquire_session called session.config().await for every candidate while holding sessions.write(), so one session busy with a long statement stalled the whole pool — the original bug, still reachable through create_session() (the follow-up path used by metadata, table data and plugin operations). The session table lock is now held only to pick and claim a candidate: the live identity check and the ping happen after it is released, a candidate that fails validation is skipped instead of stalling, and candidates are matched against a cached DatabaseIdentity (refreshed by verify_and_sync_database, the only writer of DbConnectionConfig::database, so it cannot drift).

2. [P2] A stale release could clear a newer claim. release_session_internal finishing after the next get_session_connection had already claimed the session marked a busy session as idle, and the idle sweep could then recycle the connection of a running statement (session not found). Claims are now reference counted, with a one-shot reservation: create_session / try_acquire_session reserve the session, the statement's get_session_connection consumes that reservation, so each statement adds exactly one claim and its release returns the counter to zero. Two concurrent create_session calls still never share a session, and an unconsumed reservation cannot pin a session forever (it is still collectable after the idle timeout).

3. [P2] disconnect() could run twice. release_session_internal called close() even when detach_session() returned None, racing a concurrent close_session / remove_all_sessions. Only the task that actually removes the session from the pool closes it now, and try_acquire_session discards a stale session through that same path.

Also included: get_session_connection re-checks that the session is still pooled after it took the connection, so a connection removed meanwhile is never handed out; list_sessions reads the cached identity instead of the live config; get_session_config still returns the live config but no longer waits while holding the session table lock. Lock order is now only "connection lock, then session table lock", and no path holds the session table lock while awaiting a connection lock.

Break Changes

None. Public signatures are unchanged; get_session_connection now returns SessionConnectionGuard without a lifetime parameter, which only relaxes the previous borrow.

How to Test

cargo test -p db --lib            # 1308 passed, incl. 43 manager tests
cargo test -p db                  # full suite, incl. real_postgres / real_sqlite
cargo test -p db_view --lib       # 717 passed
cargo check -p main               # OK
cargo clippy -p db --all-targets  # only the 3 pre-existing warnings in manager.rs
cargo run                         # not run (story tests need a real database / UI session)

Regression tests reproduce the reported bug and fail with a timeout before this change:

  • session_connection_guards_do_not_serialize_different_sessions — a second session must acquire its connection without waiting for the first session's statement, while the same session must remain exclusive.
  • executing_on_one_session_does_not_block_another_session — while session A is blocked inside execute, session B's execute_session must complete within the timeout.

Five more tests cover the three review issues listed above (each reproduced red first, then green):

  • try_acquire_session_does_not_block_the_pool_while_a_session_is_busy — while a session is busy with a statement, an unrelated connection must still be able to take its own session, and the busy session must not be handed out again.
  • overlapping_release_keeps_the_next_statement_claim and releasing_a_session_does_not_report_the_next_statement_as_idle — a release that finishes after the next statement claimed the session must not report it as idle.
  • concurrent_release_and_close_disconnect_the_session_once — release and close_session racing on one session must disconnect the physical connection exactly once.
  • a_statement_claims_and_releases_a_session_exactly_once — a claim/release cycle returns the session to idle, reusable, and eligible for idle cleanup.

Manual check: open two SQL tabs on the same connection, start a long-running query in tab A, then execute a query in tab B. B returns immediately instead of queueing until A finishes.

AI Assistance

This refactor was drafted with an AI assistant (Finch) and reviewed line by line afterwards. Every command listed above was executed locally on this branch; the failing-before / passing-after behaviour of the regression tests was verified manually.

Checklist

  • I have read the CONTRIBUTING document and followed the guidelines.
  • Reviewed the changes in this PR and confirmed AI generated code (If any) is accurate.
  • Passed cargo run for story tests related to the changes.
  • Tested macOS, Windows and Linux platforms performance (if the change is platform-specific)

Note: the story tests (cargo run) were not run — the verification above is test-suite level. Platform testing was done on macOS only; the change is platform-independent Rust code in the db layer, but Windows/Linux runs are still pending.

问题:A 页签执行时 B 页签无法执行,即使 B 连的是另一台库。根因是
ConnectionManager 的整张会话表共用一把 RwLock:get_session_connection()
取得写锁后交给 SessionConnectionGuard,并一直持有到整条语句执行结束。
于是任一语句执行期间,所有页签的取连接、表数据查询、元数据与配置读取
全部排队(tokio RwLock 公平锁,排队的写请求还会连带阻塞后续读请求)。

修复:
- ConnectionSession.connection 改为 Arc<AsyncMutex<Box<dyn DbConnection>>>,
  每个会话持有自己的连接锁,语句只锁自己那条连接。
- 新增 SessionState(in_use/last_active)放入独立 StdMutex,状态查询不再
  等待执行中的语句;ConnectionSession 各方法改收 &self。
- 会话表改存 Arc<ConnectionSession>;get_session_connection 仅在查表与标记
  in_use 期间持有会话表锁,随后立即释放,guard 只持有该会话的连接锁
  (不再带生命周期,调用点无需改动)。
- close_session / remove_all_sessions / cleanup_expired_sessions /
  release_session_internal 改为锁内摘出会话、锁外关闭,避免拆连接时再次
  造成全局停顿。
- 保留原有语义:同一 session 仍互斥(物理连接串行)、被占用的会话不会被
  清理、取消与关闭仍等待执行结束;锁顺序统一为「会话表 → 连接锁」,
  不存在反向持锁路径。

验证:
- 新增两个并发回归测试(修复前均以 200ms/500ms 超时失败):
  session_connection_guards_do_not_serialize_different_sessions 断言不同
  session 取连接不互相等待且同一 session 仍互斥;
  executing_on_one_session_does_not_block_another_session 断言 A 阻塞执行时
  B 的 execute_session 仍能在超时内完成。
- cargo test -p db --lib 1313 passed / 0 failed(含 manager 38 个测试)。
- cargo test -p db 全量通过(含 real_postgres / real_sqlite 集成测试)。
- cargo test -p db_view --lib 717 passed。
- cargo check -p main 通过;cargo clippy -p db --all-targets 在 manager.rs
  无新增告警。
按 review 意见修复上一版(a044b63d3)留下的三处问题:

1. try_acquire_session 在持有会话表写锁时逐条 await 候选会话的 config()
   (原 manager.rs:633),一个正在跑长语句的会话会连带堵住其它页签与该连接的
   元数据路径 —— 本次要修的串行问题在复用/建会话路径上并没有修干净。
   现在会话表锁内只做「候选筛选 + 占用」,连接锁与 ping 一律在锁外;候选身份
   改用缓存的 DatabaseIdentity(由 verify_and_sync_database 同步刷新),校验
   失败或会话已不在池内就换下一个候选,不再持锁等待。

2. release_session_internal 用 bool 覆盖 in_use,释放与下一条语句的占用交错时
   会把正在执行的会话标成空闲,随后被 idle 清理回收(表现为 session not found)。
   改为引用计数占用 + 预留位:create_session/try_acquire_session 只做一次预留,
   由紧接着取连接的语句消费,保证一条语句恰好占用一次、释放一次,且两个并发
   create_session 不会共用同一会话。

3. release_session_internal 即使 detach 未成功也会 close_session,与并发的
   close_session/remove_all_sessions 重复 disconnect。现在只有真正把会话从池中
   摘出的调用方才负责断开;try_acquire_session 丢弃失效会话也走同一条路径。

附带:get_session_connection 在取得连接后复查会话是否仍在池内,避免把已摘除并
关闭的连接交出去;list_sessions 改读缓存身份,不再等待忙碌会话的连接锁;
get_session_config 保持读取实时配置,但不再在持有会话表锁时等待。

验证(macOS 本地,crates/db edition 2021):
- cargo test -p db --lib:1308 passed / 0 failed
- cargo test -p db:全量通过(tests/real_postgres.rs 4、tests/real_sqlite.rs 4,
  doc-test 1 ignored)
- cargo test -p db_view --lib:717 passed / 0 failed
- cargo check -p main:通过
- cargo clippy -p db --all-targets:manager.rs 仅剩 3 条既有告警(行号后移)
- rustfmt --edition 2021 crates/db/src/manager.rs

新增 5 个回归测试,均先复现失败再修复通过:过期释放覆盖新占用、复用扫描堵住
其它连接、并发 release/close 只 disconnect 一次、一次语句只占用一次占用单位、
复用扫描遇到忙碌会话立即返回。未跑 story 测试(需真实数据库与交互环境),
未在 Windows/Linux 上验证。
@feigeCode
feigeCode merged commit 46f4c67 into main Sep 21, 2026
8 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant