72 lines
3 KiB
Python
72 lines
3 KiB
Python
"""Overload protection in the executor: queue wait limit and cancellation (real worker processes)."""
|
|
|
|
import asyncio
|
|
import time
|
|
import unittest
|
|
|
|
from mcp_calculator.config import load_settings
|
|
from mcp_calculator.executor import ComputationError, Executor
|
|
from tests.support import slow
|
|
|
|
|
|
async def _warm(executor: Executor) -> None:
|
|
deadline = time.time() + 30
|
|
while not executor.healthy and time.time() < deadline:
|
|
await asyncio.sleep(0.05)
|
|
|
|
|
|
@slow
|
|
class QueueWaitTest(unittest.IsolatedAsyncioTestCase):
|
|
async def test_waiting_longer_than_the_limit_is_rejected(self):
|
|
env = {"MCP_CALC_WORKERS": "1", "MCP_CALC_TIMEOUT_SECONDS": "5", "MCP_CALC_MAX_QUEUE_WAIT_SECONDS": "0.5"}
|
|
executor = Executor(load_settings([], env=env))
|
|
executor.start()
|
|
self.addCleanup(executor.stop)
|
|
await _warm(executor)
|
|
|
|
blocker = asyncio.create_task(executor.run("algebra.evaluate", expression="9^9^9^9"))
|
|
await asyncio.sleep(0.2)
|
|
started = time.perf_counter()
|
|
with self.assertRaisesRegex(ComputationError, "no worker became free") as ctx:
|
|
await executor.run("algebra.evaluate", expression="1+1")
|
|
waited = time.perf_counter() - started
|
|
self.assertEqual(ctx.exception.kind, "busy")
|
|
self.assertGreaterEqual(waited, 0.45)
|
|
self.assertLess(waited, 2)
|
|
blocker.cancel()
|
|
with self.assertRaises(asyncio.CancelledError):
|
|
await blocker
|
|
|
|
|
|
@slow
|
|
class CancellationTest(unittest.IsolatedAsyncioTestCase):
|
|
async def asyncSetUp(self):
|
|
self.executor = Executor(load_settings([], env={"MCP_CALC_WORKERS": "1", "MCP_CALC_TIMEOUT_SECONDS": "30"}))
|
|
self.executor.start()
|
|
self.addCleanup(self.executor.stop)
|
|
await _warm(self.executor)
|
|
|
|
async def test_cancelling_a_running_computation_frees_the_worker(self):
|
|
running = asyncio.create_task(self.executor.run("algebra.evaluate", expression="9^9^9^9"))
|
|
await asyncio.sleep(0.5)
|
|
running.cancel()
|
|
with self.assertRaises(asyncio.CancelledError):
|
|
await running
|
|
self.assertEqual(self.executor.inflight, 0)
|
|
started = time.perf_counter()
|
|
self.assertEqual((await self.executor.run("algebra.evaluate", expression="1+1"))["exact"], "2")
|
|
self.assertLess(time.perf_counter() - started, 10, "the worker must be terminated, not left computing for 30 s")
|
|
|
|
async def test_cancelling_a_queued_computation_never_starts_it(self):
|
|
running = asyncio.create_task(self.executor.run("algebra.evaluate", expression="9^9^9^9"))
|
|
await asyncio.sleep(0.2)
|
|
queued = asyncio.create_task(self.executor.run("algebra.evaluate", expression="2+2"))
|
|
await asyncio.sleep(0.2)
|
|
self.assertEqual(self.executor.inflight, 2)
|
|
queued.cancel()
|
|
with self.assertRaises(asyncio.CancelledError):
|
|
await queued
|
|
self.assertEqual(self.executor.inflight, 1)
|
|
running.cancel()
|
|
with self.assertRaises(asyncio.CancelledError):
|
|
await running
|