wehub-resource-sync
298 lines
11 KiB
Python
298 lines
11 KiB
Python
"""Tests for the TaskRuntime terminal-state dict leak fix.
|
|
|
|
Verifies that, at terminal state, the four short-lived tracking dicts
|
|
(``_tasks``, ``_running_by_session``, ``_pending_by_session``,
|
|
``_last_envelope_by_session``) drop the task / session_key, while
|
|
``_session_locks`` is intentionally retained to prevent split-brain on
|
|
rapid re-enqueue. Also covers exception-path cleanup and a 10 000-task
|
|
tracemalloc-bounded soak.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import gc
|
|
import tracemalloc
|
|
from collections.abc import Awaitable, Callable
|
|
from typing import Any
|
|
from unittest.mock import MagicMock
|
|
|
|
import pytest
|
|
|
|
from opensquilla.gateway import task_runtime
|
|
from opensquilla.gateway.routing import RouteEnvelope, SourceKind
|
|
from opensquilla.gateway.task_runtime import TaskRuntime
|
|
from opensquilla.session.models import AgentTaskRecord
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _make_envelope(session_key: str = "agent-1::sess-1") -> RouteEnvelope:
|
|
return RouteEnvelope(
|
|
source_kind=SourceKind.WEB,
|
|
source_name="test",
|
|
agent_id="agent-1",
|
|
session_key=session_key,
|
|
input_provenance={"kind": "test"},
|
|
)
|
|
|
|
|
|
def _make_storage() -> Any:
|
|
"""Minimal storage mock."""
|
|
storage = MagicMock()
|
|
task_db: dict[str, AgentTaskRecord] = {}
|
|
|
|
async def create(record: AgentTaskRecord) -> None:
|
|
task_db[record.task_id] = record
|
|
|
|
async def update(task_id: str, **kwargs: Any) -> None:
|
|
rec = task_db.get(task_id)
|
|
if rec is None:
|
|
return
|
|
for k, v in kwargs.items():
|
|
if hasattr(rec, k):
|
|
object.__setattr__(rec, k, v)
|
|
|
|
async def get(task_id: str) -> AgentTaskRecord | None:
|
|
return task_db.get(task_id)
|
|
|
|
async def list_tasks(**_: Any) -> list[AgentTaskRecord]:
|
|
return list(task_db.values())
|
|
|
|
storage.create_agent_task = create
|
|
storage.update_agent_task = update
|
|
storage.get_agent_task = get
|
|
storage.list_agent_tasks = list_tasks
|
|
return storage
|
|
|
|
|
|
def _make_runtime(
|
|
turn_handler: Callable[..., Awaitable[Any]] | None = None,
|
|
max_concurrency: int = 4,
|
|
max_pending_per_session: int | None = 64,
|
|
) -> TaskRuntime:
|
|
async def _default_handler(_run: Any) -> None:
|
|
pass
|
|
|
|
return TaskRuntime(
|
|
storage=_make_storage(),
|
|
turn_handler=turn_handler or _default_handler,
|
|
max_concurrency=max_concurrency,
|
|
max_pending_per_session=max_pending_per_session,
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# terminal_clears_all_dicts
|
|
# ---------------------------------------------------------------------------
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_terminal_clears_all_dicts() -> None:
|
|
"""After a task succeeds, tracking dicts (except _session_locks) must not contain its key.
|
|
|
|
``_session_locks`` is intentionally NOT cleaned at terminal to prevent
|
|
split-brain under concurrent enqueue. All other dicts are cleaned.
|
|
"""
|
|
rt = _make_runtime()
|
|
env = _make_envelope("agent-1::sess-a")
|
|
handle = await rt.enqueue(env, "hello")
|
|
await rt.wait(handle.task_id, timeout=2.0)
|
|
|
|
sk = env.session_key
|
|
assert handle.task_id not in rt._tasks
|
|
assert sk not in rt._running_by_session
|
|
assert sk not in rt._pending_by_session
|
|
# _session_locks is intentionally retained: never pop while _execute may
|
|
# still hold the lock; prevents split-brain on rapid re-enqueue.
|
|
assert sk not in rt._last_envelope_by_session
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# cancel_clears_dicts
|
|
# ---------------------------------------------------------------------------
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancel_clears_dicts() -> None:
|
|
"""After a task is cancelled, all five tracking dicts must not contain its key."""
|
|
started = asyncio.Event()
|
|
blocker = asyncio.Event()
|
|
|
|
async def _blocking_handler(_run: Any) -> None:
|
|
started.set()
|
|
await blocker.wait() # blocks until test cancels
|
|
|
|
rt = _make_runtime(turn_handler=_blocking_handler)
|
|
env = _make_envelope("agent-1::sess-b")
|
|
handle = await rt.enqueue(env, "hello")
|
|
|
|
# Wait for the handler to actually start, then cancel.
|
|
await asyncio.wait_for(started.wait(), timeout=2.0)
|
|
await rt.cancel(task_id=handle.task_id)
|
|
await rt.wait(handle.task_id, timeout=2.0)
|
|
|
|
sk = env.session_key
|
|
assert handle.task_id not in rt._tasks
|
|
assert sk not in rt._running_by_session
|
|
assert sk not in rt._pending_by_session
|
|
# _session_locks is intentionally retained.
|
|
assert sk not in rt._last_envelope_by_session
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# session_lock_kept_during_pending
|
|
# ---------------------------------------------------------------------------
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_session_lock_kept_during_pending() -> None:
|
|
"""_session_locks must NOT be removed while another task is still pending."""
|
|
first_started = asyncio.Event()
|
|
first_release = asyncio.Event()
|
|
|
|
async def _slow_handler(_run: Any) -> None:
|
|
first_started.set()
|
|
await first_release.wait()
|
|
|
|
# Only 1 concurrency slot so the second task stays pending.
|
|
rt = _make_runtime(turn_handler=_slow_handler, max_concurrency=1)
|
|
env = _make_envelope("agent-1::sess-c")
|
|
|
|
handle1 = await rt.enqueue(env, "first")
|
|
await asyncio.wait_for(first_started.wait(), timeout=2.0)
|
|
|
|
# Enqueue second task — it will be QUEUED (pending) while first is running.
|
|
handle2 = await rt.enqueue(env, "second")
|
|
|
|
sk = env.session_key
|
|
# Session lock must exist because there is still a pending task.
|
|
assert sk in rt._session_locks
|
|
|
|
# Now let the first task finish.
|
|
first_release.set()
|
|
await rt.wait(handle1.task_id, timeout=2.0)
|
|
|
|
# The lock should still exist because the second task is still alive.
|
|
assert sk in rt._session_locks
|
|
|
|
# Wait for second task to finish.
|
|
await rt.wait(handle2.task_id, timeout=2.0)
|
|
|
|
# _session_locks is intentionally retained after all tasks complete;
|
|
# do not assert its absence here.
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# exception path cleans up
|
|
# ---------------------------------------------------------------------------
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_exception_path_clears_dicts() -> None:
|
|
"""Even when the turn handler raises, cleanup must run for 4 tracking dicts.
|
|
|
|
``_session_locks`` is intentionally NOT cleared on terminal: retaining
|
|
the lock prevents split-brain when a new enqueue races with _execute's
|
|
post-terminal cleanup. All other 4 dicts (``_tasks``,
|
|
``_running_by_session``, ``_pending_by_session``,
|
|
``_last_envelope_by_session``) must be cleaned up.
|
|
"""
|
|
|
|
async def _failing_handler(_run: Any) -> None:
|
|
raise RuntimeError("deliberate failure")
|
|
|
|
rt = _make_runtime(turn_handler=_failing_handler)
|
|
env = _make_envelope("agent-1::sess-d")
|
|
handle = await rt.enqueue(env, "hello")
|
|
await rt.wait(handle.task_id, timeout=2.0)
|
|
|
|
sk = env.session_key
|
|
assert handle.task_id not in rt._tasks
|
|
assert sk not in rt._running_by_session
|
|
# _session_locks is intentionally retained after terminal: prevents
|
|
# split-brain on rapid re-enqueue; lock is cheap and bounded per session_key.
|
|
assert sk not in rt._pending_by_session
|
|
assert sk not in rt._last_envelope_by_session
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# no_leak_under_load (tracemalloc quantitative)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_no_leak_under_load(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""10 000 tasks, each <=50 ms; dict sizes after GC must be within ±2 of baseline."""
|
|
num_tasks = 10_000
|
|
session_count = 50 # rotate sessions to mimic real load
|
|
monkeypatch.setattr(task_runtime, "_emit_metric", lambda *_args, **_kwargs: None)
|
|
|
|
async def _instant_handler(_run: Any) -> None:
|
|
pass # returns immediately — well under 50 ms
|
|
|
|
rt = _make_runtime(
|
|
turn_handler=_instant_handler,
|
|
max_concurrency=32,
|
|
max_pending_per_session=None,
|
|
)
|
|
|
|
# --- baseline snapshot (before any tasks) ---
|
|
gc.collect()
|
|
tracemalloc.start()
|
|
snap_before = tracemalloc.take_snapshot()
|
|
|
|
baseline_tasks = len(rt._tasks)
|
|
baseline_pending = len(rt._pending_by_session)
|
|
baseline_running = len(rt._running_by_session)
|
|
baseline_envelope = len(rt._last_envelope_by_session)
|
|
|
|
# --- run 10 000 tasks ---
|
|
handles = []
|
|
for i in range(num_tasks):
|
|
sk = f"agent-1::sess-load-{i % session_count}"
|
|
env = _make_envelope(sk)
|
|
h = await rt.enqueue(env, f"msg-{i}")
|
|
handles.append(h)
|
|
|
|
# Wait for all to complete.
|
|
await asyncio.gather(*(rt.wait(h.task_id, timeout=30.0) for h in handles))
|
|
|
|
# --- post-GC snapshot ---
|
|
gc.collect()
|
|
snap_after = tracemalloc.take_snapshot()
|
|
tracemalloc.stop()
|
|
|
|
after_tasks = len(rt._tasks)
|
|
after_locks = len(rt._session_locks)
|
|
after_pending = len(rt._pending_by_session)
|
|
after_running = len(rt._running_by_session)
|
|
after_envelope = len(rt._last_envelope_by_session)
|
|
|
|
tolerance = 2
|
|
assert abs(after_tasks - baseline_tasks) <= tolerance, (
|
|
f"_tasks leaked: baseline={baseline_tasks}, after={after_tasks}"
|
|
)
|
|
# _session_locks is intentionally NOT cleaned at terminal to prevent
|
|
# split-brain on rapid re-enqueue. The dict grows by # unique session_keys
|
|
# (capped at session_count=50 here), not by # tasks. We verify it is bounded
|
|
# by session_count rather than by num_tasks.
|
|
assert after_locks <= session_count + tolerance, (
|
|
f"_session_locks grew beyond unique session count: {after_locks} > {session_count}"
|
|
)
|
|
assert abs(after_pending - baseline_pending) <= tolerance, (
|
|
f"_pending_by_session leaked: baseline={baseline_pending}, after={after_pending}"
|
|
)
|
|
assert abs(after_running - baseline_running) <= tolerance, (
|
|
f"_running_by_session leaked: baseline={baseline_running}, after={after_running}"
|
|
)
|
|
assert abs(after_envelope - baseline_envelope) <= tolerance, (
|
|
f"_last_envelope_by_session leaked: baseline={baseline_envelope}, after={after_envelope}"
|
|
)
|
|
|
|
# Confirm memory allocation delta is reasonable (no catastrophic growth).
|
|
# Informational only — the dict-size assertions above are authoritative.
|
|
# 10 000 asyncio tasks create significant transient allocation for
|
|
# Task/Future/Event objects; allow up to 200 MB of incidental growth.
|
|
top_stats = snap_after.compare_to(snap_before, "lineno")
|
|
total_added = sum(s.size_diff for s in top_stats if s.size_diff > 0)
|
|
assert total_added < 200 * 1024 * 1024, (
|
|
f"Unexpected memory growth: {total_added / 1024:.1f} KB"
|
|
)
|