wehub-resource-sync
131 lines
4.4 KiB
Python
131 lines
4.4 KiB
Python
"""No head-session blocking when idle slots are available.
|
|
|
|
Different session_keys for the same agent run concurrently when global
|
|
slots are free — session A holding a slot does NOT block B/C/D from
|
|
taking the remaining idle slots. With ``max_concurrency=4`` and one
|
|
agent owning four sessions ABCD enqueued simultaneously, all four must
|
|
start within 4 seconds rather than serialising behind the first session.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import time
|
|
from typing import Any
|
|
from unittest.mock import MagicMock
|
|
|
|
import pytest
|
|
|
|
from opensquilla.gateway.routing import RouteEnvelope, SourceKind
|
|
from opensquilla.gateway.task_runtime import TaskRuntime
|
|
from opensquilla.session.models import AgentTaskRecord
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers (same pattern as test_fair_queuing.py)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _make_envelope(agent_id: str, session_key: str) -> RouteEnvelope:
|
|
return RouteEnvelope(
|
|
source_kind=SourceKind.WEB,
|
|
source_name="test",
|
|
agent_id=agent_id,
|
|
session_key=session_key,
|
|
input_provenance={"kind": "test"},
|
|
)
|
|
|
|
|
|
def _make_storage() -> Any:
|
|
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(**kwargs: 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
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# no_head_blocking_with_idle_slots
|
|
# ---------------------------------------------------------------------------
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_no_head_blocking_with_idle_slots() -> None:
|
|
"""max_concurrency=4 + 1 agent + 4 sessions ABCD → all 4 start concurrently.
|
|
|
|
Each task sleeps for 1 s. If head-blocking were present the tasks would
|
|
serialise (total ~4 s); with the fix they all start immediately (total ~1 s).
|
|
We assert all 4 tasks start within 2 s of enqueue to give CI headroom.
|
|
"""
|
|
start_deadline = 2.0 # all 4 must have *started* within this many seconds
|
|
|
|
agent_id = "agent-no-block"
|
|
started_at: dict[str, float] = {}
|
|
enqueue_time: float = 0.0
|
|
|
|
gate = asyncio.Event() # keeps tasks alive until we release them
|
|
|
|
async def turn_handler(run: Any) -> None:
|
|
started_at[run.session_key] = time.monotonic()
|
|
await gate.wait()
|
|
|
|
runtime = TaskRuntime(
|
|
storage=_make_storage(),
|
|
turn_handler=turn_handler,
|
|
max_concurrency=4,
|
|
max_pending_per_session=None,
|
|
)
|
|
|
|
sessions = [f"{agent_id}::sess-{label}" for label in ("A", "B", "C", "D")]
|
|
envs = [_make_envelope(agent_id, sk) for sk in sessions]
|
|
|
|
enqueue_time = time.monotonic()
|
|
handles = []
|
|
for env in envs:
|
|
h = await runtime.enqueue(env, "hello")
|
|
handles.append(h)
|
|
|
|
# Give the event loop enough ticks for all 4 tasks to reach their handler.
|
|
deadline = asyncio.get_event_loop().time() + start_deadline
|
|
while len(started_at) < 4 and asyncio.get_event_loop().time() < deadline:
|
|
await asyncio.sleep(0.05)
|
|
|
|
# Release all tasks so they can finish cleanly.
|
|
gate.set()
|
|
|
|
# Wait for completion (with generous timeout).
|
|
for h in handles:
|
|
try:
|
|
await asyncio.wait_for(runtime.wait(h.task_id), timeout=10.0)
|
|
except (TimeoutError, KeyError):
|
|
pass
|
|
|
|
assert len(started_at) == 4, (
|
|
f"Only {len(started_at)}/4 sessions started within {start_deadline}s: "
|
|
f"{list(started_at.keys())}"
|
|
)
|
|
|
|
# All 4 must have started within START_DEADLINE seconds of the first enqueue.
|
|
last_start = max(started_at.values()) - enqueue_time
|
|
assert last_start <= start_deadline, (
|
|
f"Last session started {last_start:.2f}s after enqueue — head-blocking "
|
|
f"suspected (expected all 4 within {start_deadline}s): {started_at}"
|
|
)
|