Coverage for src/local_deep_research/web/auth/connection_cleanup.py: 95%
128 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-22 21:20 +0000
1"""
2Automatic cleanup of idle database connections.
4Periodically closes database connections for users who have no active sessions
5and no active research, preventing resource leaks when users close their browser
6without logging out.
8Also periodically disposes all QueuePool engines to release accumulated WAL/SHM
9file handles. See ADR-0004 for why this is necessary with SQLCipher + WAL mode.
10"""
12import os
13import time
14from pathlib import Path
16from apscheduler.schedulers.background import BackgroundScheduler
17from loguru import logger
19from ...database.session_passwords import session_password_store
20from ...database.thread_local_session import (
21 cleanup_dead_threads,
22 clear_user_credentials,
23)
24from ...web.research_state import get_usernames_with_active_research
26# ---------------------------------------------------------------------------
27# File Descriptor Monitoring
28# ---------------------------------------------------------------------------
29# WHY: After days of idle operation in Docker, the app crashed with
30# OSError: [Errno 24] Too many open files
31# This monitoring logs the FD count every 5 minutes so we can correlate
32# FD growth with specific events and find leaks.
33#
34# WHAT IT LOGS:
35# - open_fds: total open file descriptors for the process
36# - pool_engines: number of per-user QueuePool engines
37# - pool_checked_out: connections currently checked out from QueuePool
38# - protected_users: users with active sessions
39#
40# HOW TO USE: grep "Resource monitor" in container logs. If open_fds
41# grows steadily over hours, something is leaking.
42# ---------------------------------------------------------------------------
44# Dispose all pool engines every 30 minutes to release WAL/SHM handles.
45# SQLCipher + WAL mode leaks handles when connections close out of order
46# (which QueuePool's pool_recycle causes). Periodic dispose() closes ALL
47# pooled connections at once, resetting the handle state cleanly.
48# The next DB operation transparently reopens a fresh connection.
49_DISPOSE_INTERVAL_SECONDS = 1800
50_last_dispose_time = 0.0
53def _pop_per_user_locks(username: str) -> None:
54 """Release safe-to-remove per-user lock-cache entries for ``username``.
56 The library-init, backup, queue-processor, and library-RAG modules each
57 maintain per-user lock registries. Plain locks cannot be safely removed
58 because a caller may already hold a reference before acquisition, so their
59 compatibility cleanup hooks retain stable identity. The library-RAG cache
60 uses tracked acquisition and can safely evict idle entries.
62 The research-start gate, queue admission view, library-init lock and backup
63 lock are bounded by the user population.
65 Lazy-imported here to keep this module's import graph shallow:
66 ``connection_cleanup`` runs at startup and shouldn't pull in the
67 queue / backup / library-init / library-RAG modules eagerly.
68 """
69 try:
70 from ...database.library_init import pop_user_init_lock
72 pop_user_init_lock(username)
73 except Exception:
74 # Surface at WARNING to match the sibling scheduler-unregister
75 # error handler in this same module (line ~111). A failure
76 # here means the lock-dict entry will accumulate on every
77 # subsequent close cycle for this user; we want it visible.
78 logger.warning(f"Failed to pop _user_init_locks for {username}")
80 try:
81 from ...database.backup.backup_service import pop_user_lock
83 pop_user_lock(username)
84 except Exception:
85 logger.warning(f"Failed to pop _user_locks for {username}")
87 try:
88 from ...web.queue.processor_v2 import queue_processor
90 queue_processor.pop_user_critical_lock(username)
91 except Exception:
92 logger.warning(f"Failed to pop _user_critical_locks for {username}")
94 # NOTE: the per-user research-start gate (_user_research_start_gates in
95 # web/research_state.py) is deliberately NOT popped here — it may be held
96 # across a multi-second rekey and removing a held gate would let a
97 # concurrent same-user check_and_start_research create a second gate
98 # instance and bypass the exclusion. See the docstring above.
100 try:
101 from ...research_library.services.library_rag_service import (
102 pop_faiss_locks_for_user,
103 )
105 pop_faiss_locks_for_user(username)
106 except Exception:
107 logger.warning(f"Failed to pop _faiss_write_locks for {username}")
110def _evict_collection_engines_cache(username: str) -> None:
111 """Best-effort drop of ``username``'s cached collection search engines.
113 The collection-engine cache in ``search_engines_config`` holds the user's
114 collection names/descriptions for up to its 30 s TTL (5 minutes under
115 stale-serve after a DB error). Like the lock entries above it is
116 per-user residue, so every teardown path calls this: logout, password
117 change, and the idle-connection sweep (#6609). Lazy-imported for the
118 same import-graph reason as ``_pop_per_user_locks``.
119 """
120 try:
121 from ...web_search_engines.search_engines_config import (
122 evict_collection_engines_cache,
123 )
125 evict_collection_engines_cache(username)
126 except Exception:
127 logger.warning(
128 f"Failed to evict cached collection engines for {username}"
129 )
132def _disconnect_all_user_sockets(username: str) -> None:
133 """Best-effort disconnect of ALL of ``username``'s live sockets.
135 Called from the idle-connection sweep, where the user has no active
136 session at all — so every one of their still-open sockets (authorised
137 once at handshake and never re-checked) should be severed. Lazy-imported
138 to keep this module's startup import graph shallow and to tolerate
139 non-web contexts (the socket server may not exist).
141 Retargeted from the deleted Flask ``SocketIOService`` (#5535) onto the
142 ASGI socket layer. Behaviour is equivalent: ``disconnect_user`` severs
143 every sid authenticated as this user, and the resulting ``disconnect``
144 handler drops their subscriptions. It schedules onto the main loop and
145 returns False rather than raising when no loop is running, which is
146 the non-web case main handled by catching ValueError.
147 """
148 try:
149 from ...web.services.socketio_asgi import disconnect_user
151 disconnect_user(username)
152 except Exception:
153 logger.warning(f"Failed to disconnect sockets for idle user {username}")
156def _count_open_fds() -> int:
157 """Count open file descriptors for the current process."""
158 proc_fd = Path("/proc/self/fd")
159 if proc_fd.is_dir():
160 try:
161 return len(list(proc_fd.iterdir()))
162 except OSError:
163 pass
164 import resource
166 soft_limit = resource.getrlimit(resource.RLIMIT_NOFILE)[0]
167 count = 0
168 for fd in range(soft_limit):
169 try:
170 os.fstat(fd)
171 count += 1
172 except OSError:
173 pass
174 return count
177def cleanup_idle_connections(session_manager, db_manager):
178 """Close db connections for users with no active sessions and no active research."""
179 # 1. Purge expired sessions first
180 session_manager.cleanup_expired_sessions()
182 # 2. Get protected usernames (active sessions OR active research)
183 active_usernames = session_manager.get_active_usernames()
184 researching_usernames = get_usernames_with_active_research()
185 protected = active_usernames | researching_usernames
187 # 3. Get usernames with open connections
188 connected_usernames = db_manager.get_connected_usernames()
190 # 4. Find idle candidates
191 candidates = connected_usernames - protected
193 # 5. Double-check before closing (narrows race window)
194 closed = 0
195 for username in candidates:
196 if session_manager.has_active_sessions_for(username):
197 logger.debug(
198 f"Skipped {username} (active session appeared since snapshot)"
199 )
200 continue # User logged in since snapshot
201 if username in get_usernames_with_active_research():
202 logger.debug(
203 f"Skipped {username} (active research appeared since snapshot)"
204 )
205 continue # Research started since snapshot
206 # Unregister news scheduler jobs (matches logout pattern in routes.py)
207 try:
208 from ...scheduler.background import (
209 get_background_job_scheduler,
210 )
212 sched = get_background_job_scheduler()
213 if sched.is_running:
214 sched.unregister_user(username)
215 except Exception:
216 logger.warning(
217 f"Failed to unregister scheduler for {username}",
218 )
219 try:
220 db_manager.close_user_database(username)
221 session_password_store.clear_all_for_user(username)
222 closed += 1
223 logger.debug(f"Closed idle connection for {username}")
224 except Exception:
225 logger.warning(f"Connection cleanup failed for {username}")
226 # Drop cached plaintext credentials on pooled worker threads, for the
227 # same reason logout does. This sweep is the teardown path for the
228 # MAJORITY of users — most close the tab rather than clicking logout —
229 # so omitting it left the SQLCipher master key in a process-global
230 # dict indefinitely after the server had already decided the user was
231 # gone and closed their database. Outside the try above for the same
232 # reason _pop_per_user_locks is: independent of engine teardown, and
233 # it matters most on the path where close raises.
234 clear_user_credentials(username)
235 # This user has no active session at all (that's why we're closing
236 # their DB), so tear down every one of their still-open sockets — a
237 # socket authorised at handshake is never re-checked and would
238 # otherwise keep receiving the user's events after the session lapsed.
239 _disconnect_all_user_sockets(username)
240 # Run lock-cache cleanup regardless of whether close succeeded.
241 # Stable plain-lock registries keep their identities; tracked caches
242 # can evict idle entries. This remains independent of engine teardown.
243 _pop_per_user_locks(username)
244 # The user's cached collection engines are teardown residue too,
245 # dropped outside the try for the same reason (#6609).
246 _evict_collection_engines_cache(username)
248 if closed:
249 logger.info(f"Connection cleanup: closed {closed} idle connection(s)")
250 logger.debug(
251 f"Connection cleanup: evaluated {len(candidates)} candidate(s), "
252 f"closed {closed}, protected {len(protected)} active user(s)"
253 )
255 # Sweep dead-thread sessions and credentials — safety net when neither
256 # HTTP requests nor the queue processor are triggering sweeps.
257 cleanup_dead_threads()
259 # --- Periodic pool dispose to release WAL/SHM handles ---
260 # SQLCipher + WAL mode accumulates file handles when QueuePool recycles
261 # connections out of open-order (ADR-0004). Periodically calling
262 # dispose() on all engines closes ALL pooled connections, releasing any
263 # leaked handles. The pool is transparently recreated on the next DB
264 # operation.
265 #
266 # Safe to run against engines with checked-out connections: SA 2.0
267 # `QueuePool.dispose` only drains idle queue entries and
268 # `Engine.dispose` calls `pool.recreate()`; a thread holding a
269 # checked-out connection keeps using it until return. SA docs are
270 # explicit — "Connections that are still checked out will not be
271 # closed". The post-login bulk write (_perform_post_login_tasks in
272 # web/auth/routes.py) is additionally protected by being a single
273 # atomic transaction, so any interruption (dispose, crash, OOM)
274 # rolls back cleanly without leaving partial state.
275 #
276 # Do not add a `checkedout() > 0` skip guard here without first
277 # reproducing a real torn-write against the actual SA source path:
278 # see PR #3487 discussion — the speculative skip introduces an
279 # unbounded "skip forever" risk on busy engines in exchange for
280 # preventing a failure mode that SA 2.0 does not produce.
281 global _last_dispose_time
282 now = time.monotonic()
283 if now - _last_dispose_time >= _DISPOSE_INTERVAL_SECONDS:
284 _last_dispose_time = now
285 # Delegate to the fenced pool-reset API: it takes each user's
286 # rekey lifecycle lock and the on-disk sidecar lock before
287 # checkpointing and disposing, so this periodic sweep can never
288 # race a live password rekey (checkpointing a mid-rekey WAL, or
289 # disposing the engine a rekey is about to evict anyway). Engines
290 # whose sidecar is contended are skipped, not torn down. Individual
291 # dispose failures are logged (type name only) inside the API.
292 try:
293 disposed = db_manager.reset_connection_pools()
294 except Exception as exc:
295 # The sweep must never take the scheduler down with it. Only
296 # the exception TYPE NAME is logged — the value can carry
297 # sensitive locals (DB paths, query fragments).
298 logger.warning(
299 f"Periodic pool dispose failed: {type(exc).__name__}"
300 )
301 disposed = 0
302 if disposed:
303 logger.info(
304 f"Pool dispose: reset {disposed} engine(s) to release "
305 f"WAL/SHM handles"
306 )
308 # --- FD monitoring ---
309 try:
310 fd_count = _count_open_fds()
311 pool_engine_count = len(db_manager.connections)
312 pool_checked_out = 0
313 with db_manager._connections_lock:
314 for engine in db_manager.connections.values():
315 try:
316 pool_checked_out += engine.pool.checkedout()
317 except Exception: # noqa: silent-exception
318 pass
319 logger.debug(
320 f"Resource monitor: open_fds={fd_count}, "
321 f"pool_engines={pool_engine_count}, "
322 f"pool_checked_out={pool_checked_out}, "
323 f"protected_users={len(protected)}"
324 )
325 if fd_count > 800:
326 logger.warning(
327 f"High FD count ({fd_count}) — approaching system limit. "
328 f"Check for resource leaks."
329 )
330 except Exception:
331 logger.debug("FD monitoring failed") # noqa: silent-exception
334def start_connection_cleanup_scheduler(
335 session_manager, db_manager, interval_seconds=300
336):
337 """Start APScheduler job for periodic connection cleanup.
339 Args:
340 session_manager: The SessionManager singleton.
341 db_manager: The DatabaseManager singleton.
342 interval_seconds: How often to run cleanup (default: 5 minutes).
344 Returns:
345 The BackgroundScheduler instance (for shutdown registration).
346 """
347 scheduler = BackgroundScheduler()
348 scheduler.add_job(
349 cleanup_idle_connections,
350 "interval",
351 seconds=interval_seconds,
352 args=[session_manager, db_manager],
353 id="cleanup_idle_connections",
354 jitter=30,
355 )
356 scheduler.start()
357 logger.info(
358 f"Connection cleanup scheduler started "
359 f"(interval={interval_seconds}s, jitter=30s)"
360 )
361 return scheduler