fix(gateway): offload SessionStore calls off the event loop via asyncio.to_thread#61842
Closed
kenyonxu wants to merge 2 commits into
Closed
fix(gateway): offload SessionStore calls off the event loop via asyncio.to_thread#61842kenyonxu wants to merge 2 commits into
kenyonxu wants to merge 2 commits into
Conversation
…cio.to_thread Every inbound message calls get_or_create_session which synchronously executes _is_session_ended_in_db → db.get_session → conn.execute on the asyncio event loop. On a ~1.4GB state.db, this blocks the loop for seconds to minutes, starving Discord heartbeats. Upstream NousResearch#55159 fixed the same pattern for self._session_db in gateway/run.py but missed SessionStore._db in gateway/session.py. This follows the exact same approach as NousResearch#55159: - session.py internals stay fully synchronous (zero changes) - Threading.Lock contract is preserved - All hot-path callers in run.py and slash_commands.py wrap calls with await asyncio.to_thread(self.session_store.method, ...) Affected: get_or_create_session, switch_session, update_session, load_transcript, rewrite_transcript, rewind_session, reset_session, set_model_override, _save — ~24 call sites across 2 files. # Conflicts: # gateway/slash_commands.py
…check (NousResearch#5) The sync _session_has_compression_in_flight sat on the message hot path and blocked the event loop twice: under session_store._lock during _ensure_loaded_locked (JSON read) and via db.get_compression_lock_holder (SQLite SELECT). Async-ify the method and offload both sources via asyncio.to_thread; await the call site in _handle_active_session_busy_message.
Collaborator
Related: #60980 (api_server SessionDB offload) and the merged event-loop-starvation cluster #60984 / #53603 / #52090. This PR closes a distinct hole -- |
Collaborator
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
PR #55159 added
AsyncSessionDBto offloadSessionDBcalls from the gateway event loop, and even included an AST guard to prevent new raw on-loop calls. ButSessionStoreingateway/session.pycreates its ownSessionDB()directly (line 949), bypassing theAsyncSessionDBfacade entirely. So everyget_or_create_session→_is_session_ended_in_db→db.get_sessionchain still runs synchronously on the event loop.On a large
state.db(~1.4 GB in production), this blocks the loop for seconds to minutes, starving Discord/Slack heartbeats and delaying session activation (#53297 — 15-30s delay for existing sessions).This PR follows the exact same approach as #55159:
SessionStoreinternals stay fully synchronous (zero changes), and all hot-path callers wrap calls withawait asyncio.to_thread(self.session_store.method, ...).Changes
Commit 1: Offload session store calls (
gateway/run.py,gateway/slash_commands.py)~24 call sites wrapped with
asyncio.to_thread:get_or_create_session,switch_session,update_sessionload_transcript,rewrite_transcript,rewind_sessionreset_session,set_model_override,_saveCommit 2: Async-ify compression-in-flight check (
gateway/run.py)_session_has_compression_in_flightblocked the event loop twice on the message hot path:session_store._lockduring_ensure_loaded_locked(JSON read)db.get_compression_lock_holder(SQLite SELECT)Both sources are now offloaded via
asyncio.to_thread. The method signature changed fromdeftoasync def; the call site at_handle_active_session_busy_messageaddsawait.Why not fix this inside SessionStore?
The
Threading.Lockcontract onSessionStoreis preserved — all internals remain synchronous and thread-safe. The fix is purely at the call boundary, matching the pattern established by #55159 for_session_db.Related
to_threadpattern as merged PR fix(gateway): offload channel directory session scans off the event loop #60984 (channel directory offload)Verification
tests/gateway/test_compression_in_flight_check.py— 7 tests (includes thread-offload verification)tests/gateway/test_session.py— 100 teststests/gateway/test_session_store_runtime_stale_guard.py— 11 teststests/gateway/test_compression_interrupt_demotion_56391.py— existing tests adapted