opensquilla--opensquilla
wehub-resource-sync
144 行
4.6 KiB
Python
144 行
4.6 KiB
Python
"""Graceful shutdown drains in-flight tasks without cancelling them.
|
|
|
|
shutdown(graceful=True) waits for in-flight tasks to finish; a long task
|
|
running when shutdown is called must complete before shutdown returns
|
|
without being cancelled. When ``graceful_timeout`` elapses the remaining
|
|
tasks are cancelled and shutdown still returns cleanly without raising.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
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
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _make_envelope(
|
|
agent_id: str = "agent-drain",
|
|
session_key: str = "agent-drain::sess-1",
|
|
) -> 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
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# graceful_shutdown_drains_inflight
|
|
# ---------------------------------------------------------------------------
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_graceful_shutdown_drains_inflight() -> None:
|
|
"""Enqueue one long task, call shutdown(graceful=True, timeout=10s).
|
|
|
|
The task must complete (not be cancelled) and shutdown must return after
|
|
the task finishes.
|
|
"""
|
|
task_duration = 0.3 # 300 ms — long enough to still be running when shutdown is called
|
|
completed: list[str] = []
|
|
was_cancelled: list[bool] = []
|
|
|
|
task_started = asyncio.Event()
|
|
|
|
async def turn_handler(run: Any) -> None:
|
|
task_started.set()
|
|
try:
|
|
await asyncio.sleep(task_duration)
|
|
completed.append(run.task_id)
|
|
except asyncio.CancelledError:
|
|
was_cancelled.append(True)
|
|
raise
|
|
|
|
runtime = TaskRuntime(
|
|
storage=_make_storage(),
|
|
turn_handler=turn_handler,
|
|
max_concurrency=4,
|
|
)
|
|
|
|
h = await runtime.enqueue(_make_envelope(), "long task")
|
|
|
|
# Wait until the task is actually running before we call shutdown.
|
|
await asyncio.wait_for(task_started.wait(), timeout=5.0)
|
|
|
|
# Graceful shutdown: must wait for the 300 ms task to finish.
|
|
await runtime.shutdown(graceful=True, graceful_timeout=10.0)
|
|
|
|
assert completed == [h.task_id], (
|
|
f"Task did not complete before shutdown returned: completed={completed}"
|
|
)
|
|
assert was_cancelled == [], (
|
|
"Task was cancelled during graceful shutdown — expected drain, not cancel"
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# graceful_timeout fallback
|
|
# ---------------------------------------------------------------------------
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_graceful_shutdown_timeout_fallback_is_clean() -> None:
|
|
"""When graceful_timeout expires, remaining tasks are cancelled cleanly.
|
|
|
|
shutdown() must return without raising an exception even when tasks are
|
|
still running at the timeout boundary.
|
|
"""
|
|
task_started = asyncio.Event()
|
|
|
|
async def slow_handler(run: Any) -> None:
|
|
task_started.set()
|
|
# This task will not finish within the tiny graceful_timeout.
|
|
await asyncio.sleep(60)
|
|
|
|
runtime = TaskRuntime(
|
|
storage=_make_storage(),
|
|
turn_handler=slow_handler,
|
|
max_concurrency=4,
|
|
)
|
|
|
|
await runtime.enqueue(_make_envelope(), "slow task")
|
|
await asyncio.wait_for(task_started.wait(), timeout=5.0)
|
|
|
|
# Tiny graceful_timeout so we hit the fallback-to-cancel path.
|
|
# Must not raise.
|
|
await runtime.shutdown(graceful=True, graceful_timeout=0.05, timeout=2.0)
|