From cf0abb7a09501b785df763084ef0c233756a22ef Mon Sep 17 00:00:00 2001 From: autodev-bot Date: Tue, 21 Jul 2026 12:44:30 +0800 Subject: [PATCH 1/3] fix(searcher): bound thread pool to stop leak in Searcher retrieval paths Refs: #1273 `Searcher` used to open a fresh `ContextThreadPoolExecutor` on every `/search` request in four methods (`_retrieve_paths`, `_retrieve_from_long_term_and_user`, `_retrieve_from_tool_memory`, `_deduplicate_rawfile_results`). If any submitted task hung (slow Neo4j, unresponsive embedding API), the `with` block blocked in `shutdown(wait=True)` and the worker threads were never freed. Over hours/days, container thread count grew without bound (users reported 8,744 threads), at which point `/search` silently returned empty, `/chat` returned 503, and OpenBLAS could no longer `pthread_create`. Fix: - Pre-allocate four class-level shared `ContextThreadPoolExecutor` instances in `Searcher.__init__`, mirroring the original `max_workers` per site (5 / 3 / 2 / 10) with distinct `thread_name_prefix` for observability. Total ceiling per Searcher drops from unbounded to 20 worker threads regardless of QPS. - Rewrite the four call sites to submit to the shared executor without a `with` block, so the pool outlives the request. - Add `SEARCH_FUTURE_RESULT_TIMEOUT = 30.0` module constant and pass it to every `future.result()` and `as_completed()` in the four methods, catching timeouts with a warning so a hung sub-task cannot block a request forever. Kept four separate pools instead of one shared pool: `_retrieve_paths` workers call into `_retrieve_from_long_term_and_user` / `_retrieve_from_tool_memory` (nested submits). A single pool with `max_workers < fan-out` would deadlock; disjoint pools mirror the original semantics 1:1 and stay deadlock-free. Tests: 4 new regression tests locking in the fix (shared executor identity across calls, no naked `.result()`, survival under 5x-timeout slow sub-task), plus the 4 existing behavioural tests still pass. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../tree_text_memory/retrieve/searcher.py | 423 ++++++++++-------- tests/memories/textual/test_tree_searcher.py | 174 ++++++- 2 files changed, 418 insertions(+), 179 deletions(-) diff --git a/src/memos/memories/textual/tree_text_memory/retrieve/searcher.py b/src/memos/memories/textual/tree_text_memory/retrieve/searcher.py index cd27d92a1..d2f6b1725 100644 --- a/src/memos/memories/textual/tree_text_memory/retrieve/searcher.py +++ b/src/memos/memories/textual/tree_text_memory/retrieve/searcher.py @@ -37,6 +37,13 @@ logger = get_logger(__name__) KEYWORD_EXTRACT_TOP_K = 3 KEYWORD_ALLOW_POS = ("n", "nr", "nrt", "ns", "nt", "nz", "vn", "v", "t", "eng", "m") +# Timeout (seconds) applied to every future.result() inside the four search +# methods that used to open a fresh ContextThreadPoolExecutor per request +# (#1273). A per-future timeout is a safety net: because the executors are now +# shared class-level pools, hung sub-tasks cannot make the total thread count +# grow with QPS, but a hung task can still occupy one worker slot indefinitely. +# The timeout ensures the calling request does not block forever on such tasks. +SEARCH_FUTURE_RESULT_TIMEOUT = 30.0 COT_DICT = { "fine": {"en": COT_PROMPT, "zh": COT_PROMPT_ZH}, "fast": {"en": SIMPLE_COT_PROMPT, "zh": SIMPLE_COT_PROMPT_ZH}, @@ -76,6 +83,26 @@ def __init__( self.manual_close_internet = manual_close_internet self.tokenizer = tokenizer self._usage_executor = ContextThreadPoolExecutor(max_workers=4, thread_name_prefix="usage") + # Shared, class-level executors for search paths. Previously the + # four methods below opened a fresh ContextThreadPoolExecutor on + # every request; if any sub-task hung (slow Neo4j, unresponsive + # embedding API), `with ... as executor` would block in + # shutdown(wait=True) and leak threads, growing container thread + # count without bound (#1273). Using class-level pools bounds the + # total worker thread count for a Searcher instance to + # 5 + 3 + 2 + 10 = 20 regardless of QPS. + self._search_paths_executor = ContextThreadPoolExecutor( + max_workers=5, thread_name_prefix="search-paths" + ) + self._search_long_term_executor = ContextThreadPoolExecutor( + max_workers=3, thread_name_prefix="search-long-term" + ) + self._search_tool_mem_executor = ContextThreadPoolExecutor( + max_workers=2, thread_name_prefix="search-tool-mem" + ) + self._search_dedup_executor = ContextThreadPoolExecutor( + max_workers=10, thread_name_prefix="search-dedup" + ) def _maybe_rerank( self, @@ -386,10 +413,59 @@ def _retrieve_paths( } id_filter = {k: v for k, v in id_filter.items() if v is not None} - with ContextThreadPoolExecutor(max_workers=5) as executor: + # #1273: use the shared class-level executor. Do NOT wrap in `with` + # — the executor lives for the lifetime of the Searcher, not the + # request. + executor = self._search_paths_executor + tasks.append( + executor.submit( + self._retrieve_from_working_memory, + query, + parsed_goal, + query_embedding, + top_k, + memory_type, + search_filter, + search_priority, + user_name, + id_filter, + rerank=rerank, + ) + ) + tasks.append( + executor.submit( + self._retrieve_from_long_term_and_user, + query, + parsed_goal, + query_embedding, + top_k, + memory_type, + search_filter, + search_priority, + user_name, + id_filter, + mode=mode, + rerank=rerank, + ) + ) + tasks.append( + executor.submit( + self._retrieve_from_internet, + query, + parsed_goal, + query_embedding, + top_k, + info, + mode, + memory_type, + user_name, + rerank=rerank, + ) + ) + if self.use_fulltext: tasks.append( executor.submit( - self._retrieve_from_working_memory, + self._retrieve_from_keyword, query, parsed_goal, query_embedding, @@ -402,13 +478,14 @@ def _retrieve_paths( rerank=rerank, ) ) + if search_tool_memory: tasks.append( executor.submit( - self._retrieve_from_long_term_and_user, + self._retrieve_from_tool_memory, query, parsed_goal, query_embedding, - top_k, + tool_mem_top_k, memory_type, search_filter, search_priority, @@ -418,90 +495,50 @@ def _retrieve_paths( rerank=rerank, ) ) + if include_skill_memory: tasks.append( executor.submit( - self._retrieve_from_internet, + self._retrieve_from_skill_memory, query, parsed_goal, query_embedding, - top_k, - info, - mode, + skill_mem_top_k, memory_type, + search_filter, + search_priority, user_name, + id_filter, + mode=mode, rerank=rerank, ) ) - if self.use_fulltext: - tasks.append( - executor.submit( - self._retrieve_from_keyword, - query, - parsed_goal, - query_embedding, - top_k, - memory_type, - search_filter, - search_priority, - user_name, - id_filter, - rerank=rerank, - ) - ) - if search_tool_memory: - tasks.append( - executor.submit( - self._retrieve_from_tool_memory, - query, - parsed_goal, - query_embedding, - tool_mem_top_k, - memory_type, - search_filter, - search_priority, - user_name, - id_filter, - mode=mode, - rerank=rerank, - ) - ) - if include_skill_memory: - tasks.append( - executor.submit( - self._retrieve_from_skill_memory, - query, - parsed_goal, - query_embedding, - skill_mem_top_k, - memory_type, - search_filter, - search_priority, - user_name, - id_filter, - mode=mode, - rerank=rerank, - ) + if include_preference_memory: + tasks.append( + executor.submit( + self._retrieve_from_preference_memory, + query, + parsed_goal, + query_embedding, + pref_mem_top_k, + memory_type, + search_filter, + search_priority, + user_name, + id_filter, + mode=mode, + rerank=rerank, ) - if include_preference_memory: - tasks.append( - executor.submit( - self._retrieve_from_preference_memory, - query, - parsed_goal, - query_embedding, - pref_mem_top_k, - memory_type, - search_filter, - search_priority, - user_name, - id_filter, - mode=mode, - rerank=rerank, - ) + ) + results = [] + for t in tasks: + try: + results.extend(t.result(timeout=SEARCH_FUTURE_RESULT_TIMEOUT)) + except Exception as e: + logger.warning( + "[SEARCH] retrieval sub-task failed or timed out after %.1fs: %s", + SEARCH_FUTURE_RESULT_TIMEOUT, + e, ) - results = [] - for t in tasks: - results.extend(t.result()) logger.info(f"[SEARCH] Total raw results: {len(results)}") return results @@ -764,61 +801,69 @@ def _retrieve_from_long_term_and_user( else: cot_embeddings = query_embedding - with ContextThreadPoolExecutor(max_workers=3) as executor: - if memory_type in ["All", "AllSummaryMemory", "LongTermMemory"]: - tasks.append( - executor.submit( - self.graph_retriever.retrieve, - query=query, - parsed_goal=parsed_goal, - query_embedding=cot_embeddings, - top_k=top_k * 2, - memory_scope="LongTermMemory", - search_filter=search_filter, - search_priority=search_priority, - user_name=user_name, - id_filter=id_filter, - use_fast_graph=self.use_fast_graph, - ) + # #1273: shared class-level executor; do NOT wrap in `with`. + executor = self._search_long_term_executor + if memory_type in ["All", "AllSummaryMemory", "LongTermMemory"]: + tasks.append( + executor.submit( + self.graph_retriever.retrieve, + query=query, + parsed_goal=parsed_goal, + query_embedding=cot_embeddings, + top_k=top_k * 2, + memory_scope="LongTermMemory", + search_filter=search_filter, + search_priority=search_priority, + user_name=user_name, + id_filter=id_filter, + use_fast_graph=self.use_fast_graph, ) - if memory_type in ["All", "AllSummaryMemory", "UserMemory"]: - tasks.append( - executor.submit( - self.graph_retriever.retrieve, - query=query, - parsed_goal=parsed_goal, - query_embedding=cot_embeddings, - top_k=top_k * 2, - memory_scope="UserMemory", - search_filter=search_filter, - search_priority=search_priority, - user_name=user_name, - id_filter=id_filter, - use_fast_graph=self.use_fast_graph, - ) + ) + if memory_type in ["All", "AllSummaryMemory", "UserMemory"]: + tasks.append( + executor.submit( + self.graph_retriever.retrieve, + query=query, + parsed_goal=parsed_goal, + query_embedding=cot_embeddings, + top_k=top_k * 2, + memory_scope="UserMemory", + search_filter=search_filter, + search_priority=search_priority, + user_name=user_name, + id_filter=id_filter, + use_fast_graph=self.use_fast_graph, ) - if memory_type in ["RawFileMemory"]: - tasks.append( - executor.submit( - self.graph_retriever.retrieve, - query=query, - parsed_goal=parsed_goal, - query_embedding=cot_embeddings, - top_k=top_k * 2, - memory_scope="RawFileMemory", - search_filter=search_filter, - search_priority=search_priority, - user_name=user_name, - id_filter=id_filter, - use_fast_graph=self.use_fast_graph, - ) + ) + if memory_type in ["RawFileMemory"]: + tasks.append( + executor.submit( + self.graph_retriever.retrieve, + query=query, + parsed_goal=parsed_goal, + query_embedding=cot_embeddings, + top_k=top_k * 2, + memory_scope="RawFileMemory", + search_filter=search_filter, + search_priority=search_priority, + user_name=user_name, + id_filter=id_filter, + use_fast_graph=self.use_fast_graph, ) + ) - # Collect results from all tasks - for task in tasks: - results.extend(task.result()) - results = self._deduplicate_rawfile_results(results, user_name=user_name) - results = self._filter_intermediate_content(results) + # Collect results from all tasks + for task in tasks: + try: + results.extend(task.result(timeout=SEARCH_FUTURE_RESULT_TIMEOUT)) + except Exception as e: + logger.warning( + "[SEARCH] long-term/user sub-task failed or timed out after %.1fs: %s", + SEARCH_FUTURE_RESULT_TIMEOUT, + e, + ) + results = self._deduplicate_rawfile_results(results, user_name=user_name) + results = self._filter_intermediate_content(results) return self._maybe_rerank( rerank, @@ -928,47 +973,56 @@ def _retrieve_from_tool_memory( else: cot_embeddings = query_embedding - with ContextThreadPoolExecutor(max_workers=2) as executor: - if memory_type in ["All", "ToolSchemaMemory"]: - tasks.append( - executor.submit( - self.graph_retriever.retrieve, - query=query, - parsed_goal=parsed_goal, - query_embedding=cot_embeddings, - top_k=top_k * 2, - memory_scope="ToolSchemaMemory", - search_filter=search_filter, - search_priority=search_priority, - user_name=user_name, - id_filter=id_filter, - use_fast_graph=self.use_fast_graph, - ) + # #1273: shared class-level executor; do NOT wrap in `with`. + executor = self._search_tool_mem_executor + if memory_type in ["All", "ToolSchemaMemory"]: + tasks.append( + executor.submit( + self.graph_retriever.retrieve, + query=query, + parsed_goal=parsed_goal, + query_embedding=cot_embeddings, + top_k=top_k * 2, + memory_scope="ToolSchemaMemory", + search_filter=search_filter, + search_priority=search_priority, + user_name=user_name, + id_filter=id_filter, + use_fast_graph=self.use_fast_graph, ) - if memory_type in ["All", "ToolTrajectoryMemory"]: - tasks.append( - executor.submit( - self.graph_retriever.retrieve, - query=query, - parsed_goal=parsed_goal, - query_embedding=cot_embeddings, - top_k=top_k * 2, - memory_scope="ToolTrajectoryMemory", - search_filter=search_filter, - search_priority=search_priority, - user_name=user_name, - id_filter=id_filter, - use_fast_graph=self.use_fast_graph, - ) + ) + if memory_type in ["All", "ToolTrajectoryMemory"]: + tasks.append( + executor.submit( + self.graph_retriever.retrieve, + query=query, + parsed_goal=parsed_goal, + query_embedding=cot_embeddings, + top_k=top_k * 2, + memory_scope="ToolTrajectoryMemory", + search_filter=search_filter, + search_priority=search_priority, + user_name=user_name, + id_filter=id_filter, + use_fast_graph=self.use_fast_graph, ) + ) - # Collect results from all tasks - for task in tasks: - rsp = task.result() - if rsp and rsp[0].metadata.memory_type == "ToolSchemaMemory": - results["ToolSchemaMemory"].extend(rsp) - elif rsp and rsp[0].metadata.memory_type == "ToolTrajectoryMemory": - results["ToolTrajectoryMemory"].extend(rsp) + # Collect results from all tasks + for task in tasks: + try: + rsp = task.result(timeout=SEARCH_FUTURE_RESULT_TIMEOUT) + except Exception as e: + logger.warning( + "[SEARCH] tool-memory sub-task failed or timed out after %.1fs: %s", + SEARCH_FUTURE_RESULT_TIMEOUT, + e, + ) + continue + if rsp and rsp[0].metadata.memory_type == "ToolSchemaMemory": + results["ToolSchemaMemory"].extend(rsp) + elif rsp and rsp[0].metadata.memory_type == "ToolTrajectoryMemory": + results["ToolTrajectoryMemory"].extend(rsp) schema_reranked = self._maybe_rerank( rerank, @@ -1298,20 +1352,24 @@ def _deduplicate_rawfile_results(self, results, user_name: str | None = None): if not rawfile_items: return results - with ContextThreadPoolExecutor(max_workers=min(len(rawfile_items), 10)) as executor: - futures = [ - executor.submit( - self.graph_store.get_edges, - rawfile_item.id, - type="SUMMARY", - direction="OUTGOING", - user_name=user_name, - ) - for rawfile_item in rawfile_items - ] - for future in as_completed(futures): + # #1273: shared class-level executor; do NOT wrap in `with`. The + # pool is sized once (max_workers=10) instead of `min(len, 10)` per + # request to avoid rebuilding the pool. + executor = self._search_dedup_executor + futures = [ + executor.submit( + self.graph_store.get_edges, + rawfile_item.id, + type="SUMMARY", + direction="OUTGOING", + user_name=user_name, + ) + for rawfile_item in rawfile_items + ] + try: + for future in as_completed(futures, timeout=SEARCH_FUTURE_RESULT_TIMEOUT): try: - edges = future.result() + edges = future.result(timeout=SEARCH_FUTURE_RESULT_TIMEOUT) for edge in edges: summary_target_id = edge.get("to") if summary_target_id: @@ -1321,6 +1379,15 @@ def _deduplicate_rawfile_results(self, results, user_name: str | None = None): ) except Exception as e: logger.warning(f"[DEDUP] Failed to get summary target ids: {e}") + except Exception as e: + # as_completed(...) itself timed out — some futures never + # finished within SEARCH_FUTURE_RESULT_TIMEOUT. Continue with + # whatever ids we already collected. + logger.warning( + "[DEDUP] Timed out waiting for edge lookups after %.1fs: %s", + SEARCH_FUTURE_RESULT_TIMEOUT, + e, + ) filtered_results = [] for item in results: diff --git a/tests/memories/textual/test_tree_searcher.py b/tests/memories/textual/test_tree_searcher.py index b79958ca1..ffc3d236d 100644 --- a/tests/memories/textual/test_tree_searcher.py +++ b/tests/memories/textual/test_tree_searcher.py @@ -1,8 +1,12 @@ -from unittest.mock import MagicMock +import time + +from unittest.mock import MagicMock, patch import pytest +from memos.context.context import ContextThreadPoolExecutor from memos.memories.textual.item import TextualMemoryItem, TreeNodeTextualMemoryMetadata +from memos.memories.textual.tree_text_memory.retrieve import searcher as searcher_module from memos.memories.textual.tree_text_memory.retrieve.searcher import Searcher from memos.reranker.base import BaseReranker @@ -143,3 +147,171 @@ def test_searcher_respects_memory_type(mock_searcher): ) # WorkingMemory triggers only once path A assert mock_searcher.graph_retriever.retrieve.call_args[1]["memory_scope"] == "WorkingMemory" + + +# --------------------------------------------------------------------------- +# Bug #1273 regression: bounded thread pools +# --------------------------------------------------------------------------- + + +def _make_searcher(): + dispatcher_llm = MagicMock() + graph_store = MagicMock() + embedder = MagicMock() + reranker = MagicMock(spec=BaseReranker) + return Searcher(dispatcher_llm, graph_store, embedder, reranker) + + +def test_searcher_creates_shared_executors_in_init(): + """Regression for #1273: Searcher must pre-allocate class-level shared + ContextThreadPoolExecutor instances instead of constructing them per request.""" + s = _make_searcher() + for attr, expected_workers in [ + ("_search_paths_executor", 5), + ("_search_long_term_executor", 3), + ("_search_tool_mem_executor", 2), + ("_search_dedup_executor", 10), + ]: + assert hasattr(s, attr), f"Searcher must expose {attr}" + executor = getattr(s, attr) + assert isinstance(executor, ContextThreadPoolExecutor), ( + f"{attr} must be a ContextThreadPoolExecutor, got {type(executor)!r}" + ) + assert executor._max_workers == expected_workers, ( + f"{attr} max_workers={executor._max_workers}, expected {expected_workers}" + ) + + +def test_searcher_reuses_pool_across_retrieve_paths_calls(): + """Regression for #1273: repeated _retrieve_paths calls must NOT spawn a + fresh ContextThreadPoolExecutor each time. The class-level pool must be + reused, keeping thread count bounded regardless of QPS.""" + s = _make_searcher() + + # stub internals so _retrieve_paths returns immediately + s.task_goal_parser = MagicMock() + s.graph_retriever = MagicMock() + s.graph_retriever.retrieve.return_value = [] + s.reasoner = MagicMock() + + parsed_goal = MagicMock() + parsed_goal.memories = [] + parsed_goal.rephrased_query = None + s.task_goal_parser.parse.return_value = parsed_goal + s.embedder.embed.return_value = [[0.1] * 5] + + pool_ids_seen = [] + original_paths_executor = s._search_paths_executor + orig_submit = original_paths_executor.submit + + def tracking_submit(fn, *args, **kwargs): + pool_ids_seen.append(id(original_paths_executor)) + return orig_submit(fn, *args, **kwargs) + + with patch.object(original_paths_executor, "submit", side_effect=tracking_submit): + # First call + s._retrieve_paths( + query="q", + parsed_goal=parsed_goal, + query_embedding=[[0.1] * 5], + info={}, + top_k=1, + mode="fast", + memory_type="WorkingMemory", + ) + # Second call + s._retrieve_paths( + query="q", + parsed_goal=parsed_goal, + query_embedding=[[0.1] * 5], + info={}, + top_k=1, + mode="WorkingMemory", + memory_type="WorkingMemory", + ) + + assert pool_ids_seen, "expected at least one submit to the shared paths executor" + assert len(set(pool_ids_seen)) == 1, ( + "the paths executor must be a single shared instance across calls; " + f"saw distinct ids: {set(pool_ids_seen)}" + ) + # Executor must NOT have been shut down by the with-block + assert not original_paths_executor._shutdown, ( + "shared paths executor must remain open between requests; " + "presence of shutdown indicates per-request lifetime and bug #1273 regression" + ) + + +def test_searcher_future_result_calls_use_timeout(): + """Regression for #1273: future.result() calls in the four affected methods + MUST pass a timeout so a hung sub-task cannot block a request forever.""" + assert hasattr(searcher_module, "SEARCH_FUTURE_RESULT_TIMEOUT"), ( + "expected module-level SEARCH_FUTURE_RESULT_TIMEOUT constant" + ) + assert isinstance(searcher_module.SEARCH_FUTURE_RESULT_TIMEOUT, (int, float)) + assert searcher_module.SEARCH_FUTURE_RESULT_TIMEOUT > 0 + + import inspect + + src = inspect.getsource(searcher_module) + # No naked .result() with no arguments in the source of the four methods. + for method_name in [ + "_retrieve_paths", + "_retrieve_from_long_term_and_user", + "_retrieve_from_tool_memory", + "_deduplicate_rawfile_results", + ]: + method = getattr(Searcher, method_name) + method_src = inspect.getsource(method) + # Every .result(...) call in these methods must reference the constant + # (either directly, or via a helper that uses it). + for line in method_src.splitlines(): + stripped = line.strip() + if ".result()" in stripped and not stripped.startswith("#"): + pytest.fail( + f"{method_name}: naked future.result() found — must pass " + f"timeout=SEARCH_FUTURE_RESULT_TIMEOUT. line: {stripped!r}" + ) + assert "SEARCH_FUTURE_RESULT_TIMEOUT" in src, ( + "SEARCH_FUTURE_RESULT_TIMEOUT constant must be applied in searcher.py" + ) + + +def test_searcher_survives_slow_subtask(monkeypatch): + """Regression for #1273: if a submitted sub-task exceeds + SEARCH_FUTURE_RESULT_TIMEOUT, .search() must recover with a warning + instead of raising or hanging.""" + # shrink timeout so the test runs fast + monkeypatch.setattr(searcher_module, "SEARCH_FUTURE_RESULT_TIMEOUT", 0.1) + + s = _make_searcher() + s.task_goal_parser = MagicMock() + s.graph_retriever = MagicMock() + s.reasoner = MagicMock() + + parsed_goal = MagicMock() + parsed_goal.memories = [] + parsed_goal.rephrased_query = None + s.task_goal_parser.parse.return_value = parsed_goal + s.embedder.embed.return_value = [[0.1] * 5] + + def slow_retrieve(*_args, **_kwargs): + time.sleep(0.5) # 5x the timeout + return [make_item("late", 0.1)[0]] + + s.graph_retriever.retrieve.side_effect = slow_retrieve + s.reranker.rerank.return_value = [] + + start = time.time() + result = s.search( + query="anything", + top_k=1, + info={"user_id": "u"}, + mode="fast", + memory_type="WorkingMemory", + ) + elapsed = time.time() - start + + # Must return (not raise), and quickly (well below the sleep of 0.5s). + assert elapsed < 0.45, f"search took {elapsed:.3f}s, expected <0.45s once timeout={0.1}s fires" + assert isinstance(result, list) From 79e7d6a338af791fae2e77a0d809a00f6989fc8f Mon Sep 17 00:00:00 2001 From: autodev-bot Date: Tue, 21 Jul 2026 13:02:46 +0800 Subject: [PATCH 2/3] fix(searcher): address OCR round-1 findings on executor lifecycle and tests - drop redundant inner timeout on future.result() inside as_completed loop - add Searcher.close() + context-manager protocol to shut down the five shared executors so worker threads do not outlive the instance - tests: guard private _max_workers access, use valid mode='fine' in pool reuse test, hoist inspect import, replace string matching with AST-based detection of unbounded future.result() calls Co-Authored-By: Claude Opus 4.7 (1M context) --- .../tree_text_memory/retrieve/searcher.py | 31 ++++- tests/memories/textual/test_tree_searcher.py | 114 ++++++++++++++---- 2 files changed, 122 insertions(+), 23 deletions(-) diff --git a/src/memos/memories/textual/tree_text_memory/retrieve/searcher.py b/src/memos/memories/textual/tree_text_memory/retrieve/searcher.py index d2f6b1725..23131de45 100644 --- a/src/memos/memories/textual/tree_text_memory/retrieve/searcher.py +++ b/src/memos/memories/textual/tree_text_memory/retrieve/searcher.py @@ -104,6 +104,32 @@ def __init__( max_workers=10, thread_name_prefix="search-dedup" ) + def close(self) -> None: + """Shut down all shared thread-pool executors owned by this Searcher. + + The five class-level pools created in ``__init__`` (usage tracking + + the four search pools introduced for #1273) hold non-daemon worker + threads that would otherwise outlive the instance. Callers that + create short-lived ``Searcher`` objects should call ``close()`` when + done, or use the instance as a context manager + (``with Searcher(...) as searcher: ...``). ``shutdown(wait=False)`` + is used so a hung sub-task cannot block the caller; idempotent. + """ + for executor in ( + self._usage_executor, + self._search_paths_executor, + self._search_long_term_executor, + self._search_tool_mem_executor, + self._search_dedup_executor, + ): + executor.shutdown(wait=False) + + def __enter__(self) -> "Searcher": + return self + + def __exit__(self, exc_type, exc_val, exc_tb) -> None: + self.close() + def _maybe_rerank( self, enabled: bool, @@ -1369,7 +1395,10 @@ def _deduplicate_rawfile_results(self, results, user_name: str | None = None): try: for future in as_completed(futures, timeout=SEARCH_FUTURE_RESULT_TIMEOUT): try: - edges = future.result(timeout=SEARCH_FUTURE_RESULT_TIMEOUT) + # No inner timeout: a future yielded by as_completed() is + # already done, so result() returns immediately. Passing a + # second timeout here would be redundant and misleading. + edges = future.result() for edge in edges: summary_target_id = edge.get("to") if summary_target_id: diff --git a/tests/memories/textual/test_tree_searcher.py b/tests/memories/textual/test_tree_searcher.py index ffc3d236d..3db5224c5 100644 --- a/tests/memories/textual/test_tree_searcher.py +++ b/tests/memories/textual/test_tree_searcher.py @@ -1,3 +1,5 @@ +import ast +import inspect import time from unittest.mock import MagicMock, patch @@ -177,8 +179,16 @@ def test_searcher_creates_shared_executors_in_init(): assert isinstance(executor, ContextThreadPoolExecutor), ( f"{attr} must be a ContextThreadPoolExecutor, got {type(executor)!r}" ) - assert executor._max_workers == expected_workers, ( - f"{attr} max_workers={executor._max_workers}, expected {expected_workers}" + # NOTE: ThreadPoolExecutor does not expose max_workers publicly, so we + # inspect the private `_max_workers` attribute. This is stable on + # CPython 3.9-3.13 (the versions this project tests against); use + # getattr so a future rename fails with a clear assertion message + # instead of an AttributeError. + actual_workers = getattr(executor, "_max_workers", None) + assert actual_workers == expected_workers, ( + f"{attr} max_workers={actual_workers}, expected {expected_workers} " + "(None means the private _max_workers attribute is gone — " + "update this test for the current Python version)" ) @@ -219,14 +229,14 @@ def tracking_submit(fn, *args, **kwargs): mode="fast", memory_type="WorkingMemory", ) - # Second call + # Second call (distinct valid mode to also cover the 'fine' path) s._retrieve_paths( query="q", parsed_goal=parsed_goal, query_embedding=[[0.1] * 5], info={}, top_k=1, - mode="WorkingMemory", + mode="fine", memory_type="WorkingMemory", ) @@ -242,41 +252,101 @@ def tracking_submit(fn, *args, **kwargs): ) +def _searcher_method_ast(method_name: str) -> ast.FunctionDef: + """Return the AST node of a Searcher method, resolved from the module + source. Unlike ``inspect.getsource(getattr(Searcher, name))`` this is + immune to decorators such as ``@timed`` (which does not use + ``functools.wraps``, so the bound method's ``__code__`` points at the + decorator's wrapper, not the real method body).""" + tree = ast.parse(inspect.getsource(searcher_module)) + for node in ast.walk(tree): + if isinstance(node, ast.ClassDef) and node.name == "Searcher": + for item in node.body: + if isinstance(item, ast.FunctionDef) and item.name == method_name: + return item + raise AssertionError(f"method {method_name} not found on Searcher") + + +def _is_as_completed_call(node: ast.AST) -> bool: + return isinstance(node, ast.Call) and ( + (isinstance(node.func, ast.Name) and node.func.id == "as_completed") + or (isinstance(node.func, ast.Attribute) and node.func.attr == "as_completed") + ) + + def test_searcher_future_result_calls_use_timeout(): - """Regression for #1273: future.result() calls in the four affected methods - MUST pass a timeout so a hung sub-task cannot block a request forever.""" + """Regression for #1273: waiting on sub-task futures in the four affected + methods must be time-bounded so a hung sub-task cannot block a request + forever. A wait is bounded when either: + - ``future.result(timeout=...)`` is passed a timeout directly, or + - futures are collected via ``as_completed(..., timeout=...)`` — futures + yielded by ``as_completed`` are already done, so their ``.result()`` + correctly takes no timeout. + """ assert hasattr(searcher_module, "SEARCH_FUTURE_RESULT_TIMEOUT"), ( "expected module-level SEARCH_FUTURE_RESULT_TIMEOUT constant" ) assert isinstance(searcher_module.SEARCH_FUTURE_RESULT_TIMEOUT, (int, float)) assert searcher_module.SEARCH_FUTURE_RESULT_TIMEOUT > 0 - import inspect - - src = inspect.getsource(searcher_module) - # No naked .result() with no arguments in the source of the four methods. for method_name in [ "_retrieve_paths", "_retrieve_from_long_term_and_user", "_retrieve_from_tool_memory", "_deduplicate_rawfile_results", ]: - method = getattr(Searcher, method_name) - method_src = inspect.getsource(method) - # Every .result(...) call in these methods must reference the constant - # (either directly, or via a helper that uses it). - for line in method_src.splitlines(): - stripped = line.strip() - if ".result()" in stripped and not stripped.startswith("#"): - pytest.fail( - f"{method_name}: naked future.result() found — must pass " - f"timeout=SEARCH_FUTURE_RESULT_TIMEOUT. line: {stripped!r}" - ) - assert "SEARCH_FUTURE_RESULT_TIMEOUT" in src, ( + method_ast = _searcher_method_ast(method_name) + result_calls = [ + node + for node in ast.walk(method_ast) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr == "result" + ] + assert result_calls, f"{method_name}: expected at least one future .result() call" + + has_bounded_as_completed = any( + _is_as_completed_call(node) and any(kw.arg == "timeout" for kw in node.keywords) + for node in ast.walk(method_ast) + ) + for call in result_calls: + has_timeout_kwarg = any(kw.arg == "timeout" for kw in call.keywords) + assert has_timeout_kwarg or has_bounded_as_completed, ( + f"{method_name}: unbounded future .result() call at line " + f"{call.lineno} — pass timeout=SEARCH_FUTURE_RESULT_TIMEOUT or " + f"iterate futures via as_completed(..., timeout=...)" + ) + + assert "SEARCH_FUTURE_RESULT_TIMEOUT" in inspect.getsource(searcher_module), ( "SEARCH_FUTURE_RESULT_TIMEOUT constant must be applied in searcher.py" ) +def test_searcher_close_shuts_down_all_executors(): + """Searcher.close() (and the context-manager protocol) must shut down all + five shared executors so worker threads do not outlive the instance.""" + executor_attrs = [ + "_usage_executor", + "_search_paths_executor", + "_search_long_term_executor", + "_search_tool_mem_executor", + "_search_dedup_executor", + ] + + s = _make_searcher() + s.close() + for attr in executor_attrs: + assert getattr(s, attr)._shutdown, f"{attr} must be shut down after close()" + # close() must be idempotent + s.close() + + with _make_searcher() as s2: + for attr in executor_attrs: + assert not getattr(s2, attr)._shutdown + for attr in executor_attrs: + assert getattr(s2, attr)._shutdown, f"{attr} must be shut down on context exit" + + def test_searcher_survives_slow_subtask(monkeypatch): """Regression for #1273: if a submitted sub-task exceeds SEARCH_FUTURE_RESULT_TIMEOUT, .search() must recover with a warning From 4a885bf53af01cf8f619adcbc5fb890ac9bacd34 Mon Sep 17 00:00:00 2001 From: autodev-bot Date: Tue, 21 Jul 2026 13:14:28 +0800 Subject: [PATCH 3/3] test(searcher): drop private ThreadPoolExecutor attrs from #1273 tests Replace all _max_workers/_shutdown accesses with public-API behavioral checks per OCR round 2: - verify pool worker bounds by submitting N+1 blocking tasks and asserting exactly N run concurrently - verify live/shut-down state via submit() succeeding vs raising RuntimeError (documented Executor.shutdown() behavior) Co-Authored-By: Claude Opus 4.7 (1M context) --- tests/memories/textual/test_tree_searcher.py | 114 ++++++++++++++----- 1 file changed, 85 insertions(+), 29 deletions(-) diff --git a/tests/memories/textual/test_tree_searcher.py b/tests/memories/textual/test_tree_searcher.py index 3db5224c5..976385301 100644 --- a/tests/memories/textual/test_tree_searcher.py +++ b/tests/memories/textual/test_tree_searcher.py @@ -1,5 +1,6 @@ import ast import inspect +import threading import time from unittest.mock import MagicMock, patch @@ -164,32 +165,69 @@ def _make_searcher(): return Searcher(dispatcher_llm, graph_store, embedder, reranker) +def _assert_pool_concurrency(executor, expected_workers: int, attr: str) -> None: + """Verify the executor's worker bound using only the public + submit()/result() API: submit ``expected_workers + 1`` blocking tasks — + at least ``expected_workers`` must start concurrently, and the surplus + task must stay queued until a slot frees up. This avoids private, + undocumented CPython internals such as ``_max_workers``.""" + release = threading.Event() + lock = threading.Lock() + started: list[str] = [] + + def blocker() -> None: + with lock: + started.append(threading.current_thread().name) + release.wait(timeout=10.0) + + futures = [executor.submit(blocker) for _ in range(expected_workers + 1)] + try: + deadline = time.monotonic() + 5.0 + while time.monotonic() < deadline: + with lock: + if len(started) >= expected_workers: + break + time.sleep(0.01) + with lock: + count = len(started) + assert count >= expected_workers, ( + f"{attr}: only {count} of {expected_workers} expected workers started within 5s" + ) + # Give the surplus task a chance to (incorrectly) start running. + time.sleep(0.2) + with lock: + count = len(started) + assert count == expected_workers, ( + f"{attr}: {count} tasks ran concurrently, expected at most " + f"{expected_workers} (pool sized larger than configured)" + ) + finally: + release.set() + for f in futures: + f.result(timeout=10.0) + + def test_searcher_creates_shared_executors_in_init(): """Regression for #1273: Searcher must pre-allocate class-level shared - ContextThreadPoolExecutor instances instead of constructing them per request.""" + ContextThreadPoolExecutor instances instead of constructing them per + request. Worker bounds are verified behaviorally through the public + submit()/result() API — no private attributes involved.""" s = _make_searcher() - for attr, expected_workers in [ - ("_search_paths_executor", 5), - ("_search_long_term_executor", 3), - ("_search_tool_mem_executor", 2), - ("_search_dedup_executor", 10), - ]: - assert hasattr(s, attr), f"Searcher must expose {attr}" - executor = getattr(s, attr) - assert isinstance(executor, ContextThreadPoolExecutor), ( - f"{attr} must be a ContextThreadPoolExecutor, got {type(executor)!r}" - ) - # NOTE: ThreadPoolExecutor does not expose max_workers publicly, so we - # inspect the private `_max_workers` attribute. This is stable on - # CPython 3.9-3.13 (the versions this project tests against); use - # getattr so a future rename fails with a clear assertion message - # instead of an AttributeError. - actual_workers = getattr(executor, "_max_workers", None) - assert actual_workers == expected_workers, ( - f"{attr} max_workers={actual_workers}, expected {expected_workers} " - "(None means the private _max_workers attribute is gone — " - "update this test for the current Python version)" - ) + try: + for attr, expected_workers in [ + ("_search_paths_executor", 5), + ("_search_long_term_executor", 3), + ("_search_tool_mem_executor", 2), + ("_search_dedup_executor", 10), + ]: + assert hasattr(s, attr), f"Searcher must expose {attr}" + executor = getattr(s, attr) + assert isinstance(executor, ContextThreadPoolExecutor), ( + f"{attr} must be a ContextThreadPoolExecutor, got {type(executor)!r}" + ) + _assert_pool_concurrency(executor, expected_workers, attr) + finally: + s.close() def test_searcher_reuses_pool_across_retrieve_paths_calls(): @@ -245,10 +283,14 @@ def tracking_submit(fn, *args, **kwargs): "the paths executor must be a single shared instance across calls; " f"saw distinct ids: {set(pool_ids_seen)}" ) - # Executor must NOT have been shut down by the with-block - assert not original_paths_executor._shutdown, ( + # The executor must NOT have been shut down between requests. Verified + # behaviorally via the public API: a live executor accepts new work, + # while submit() after shutdown raises RuntimeError (documented + # concurrent.futures behavior) — no private attributes involved. + probe = original_paths_executor.submit(lambda: "alive") + assert probe.result(timeout=5.0) == "alive", ( "shared paths executor must remain open between requests; " - "presence of shutdown indicates per-request lifetime and bug #1273 regression" + "a shut-down executor indicates per-request lifetime and bug #1273 regression" ) @@ -322,6 +364,17 @@ def test_searcher_future_result_calls_use_timeout(): ) +def _assert_executor_shut_down(executor, attr: str, context: str) -> None: + """A shut-down executor rejects new work with RuntimeError — this is + documented public behavior of concurrent.futures.Executor.shutdown(), + so no private attributes are needed.""" + try: + executor.submit(lambda: None) + except RuntimeError: + return + pytest.fail(f"{attr} must reject new work after {context}") + + def test_searcher_close_shuts_down_all_executors(): """Searcher.close() (and the context-manager protocol) must shut down all five shared executors so worker threads do not outlive the instance.""" @@ -336,15 +389,18 @@ def test_searcher_close_shuts_down_all_executors(): s = _make_searcher() s.close() for attr in executor_attrs: - assert getattr(s, attr)._shutdown, f"{attr} must be shut down after close()" + _assert_executor_shut_down(getattr(s, attr), attr, "close()") # close() must be idempotent s.close() with _make_searcher() as s2: for attr in executor_attrs: - assert not getattr(s2, attr)._shutdown + # Executors must still accept work inside the context. + assert getattr(s2, attr).submit(lambda: "alive").result(timeout=5.0) == "alive", ( + f"{attr} must be usable before context exit" + ) for attr in executor_attrs: - assert getattr(s2, attr)._shutdown, f"{attr} must be shut down on context exit" + _assert_executor_shut_down(getattr(s2, attr), attr, "context exit") def test_searcher_survives_slow_subtask(monkeypatch):