feat(cloud): harden multi-tenant runtime resources

tests/unit_tests/utils/test_bounded_executor.py

@@ -0,0 +1,278 @@
+from __future__ import annotations
+
+import asyncio
+import threading
+
+import pytest
+
+from langbot.pkg.utils.bounded_executor import (
+    BlockingWorkCapacityError,
+    BoundedThreadPoolExecutor,
+    blocking_work_scope,
+    configure_bounded_default_executor,
+    run_blocking_atomic,
+    run_blocking_cleanup,
+)
+
+
+def test_bounded_executor_rejects_instead_of_queueing_without_limit():
+    executor = BoundedThreadPoolExecutor(
+        max_workers=1,
+        max_pending=1,
+        max_inflight_per_scope=1,
+    )
+    started = threading.Event()
+    release = threading.Event()
+
+    def block() -> str:
+        started.set()
+        release.wait(timeout=5)
+        return 'done'
+
+    first = executor.submit(block)
+    assert started.wait(timeout=1)
+    second = executor.submit(lambda: 'queued')
+
+    with pytest.raises(
+        BlockingWorkCapacityError,
+        match='capacity reached',
+    ):
+        executor.submit(lambda: 'rejected')
+
+    assert executor.snapshot() == {
+        'max_workers': 1,
+        'max_pending': 1,
+        'max_inflight_per_scope': 1,
+        'inflight': 2,
+        'running': 1,
+        'pending': 1,
+        'active_scopes': 0,
+        'submitted_total': 2,
+        'completed_total': 0,
+        'rejected_total': 1,
+        'global_rejected_total': 1,
+        'scope_rejected_total': 0,
+    }
+
+    release.set()
+    assert first.result(timeout=1) == 'done'
+    assert second.result(timeout=1) == 'queued'
+    assert executor.snapshot()['inflight'] == 0
+    executor.shutdown()
+
+
+def test_workspace_scope_cannot_monopolize_global_workers():
+    executor = BoundedThreadPoolExecutor(
+        max_workers=2,
+        max_pending=2,
+        max_inflight_per_scope=1,
+    )
+    release = threading.Event()
+    workspace_a_started = threading.Event()
+    workspace_b_started = threading.Event()
+
+    def block(started: threading.Event) -> str:
+        started.set()
+        release.wait(timeout=5)
+        return 'done'
+
+    try:
+        with blocking_work_scope('workspace-a'):
+            workspace_a = executor.submit(block, workspace_a_started)
+            assert workspace_a_started.wait(timeout=1)
+            with pytest.raises(
+                BlockingWorkCapacityError,
+                match='Workspace blocking executor capacity reached',
+            ):
+                executor.submit(lambda: 'rejected')
+
+        with blocking_work_scope('workspace-b'):
+            workspace_b = executor.submit(block, workspace_b_started)
+            assert workspace_b_started.wait(timeout=1)
+
+        snapshot = executor.snapshot()
+        assert snapshot['inflight'] == 2
+        assert snapshot['active_scopes'] == 2
+        assert snapshot['scope_rejected_total'] == 1
+        assert snapshot['global_rejected_total'] == 0
+    finally:
+        release.set()
+        assert workspace_a.result(timeout=1) == 'done'
+        assert workspace_b.result(timeout=1) == 'done'
+        executor.shutdown()
+
+
+def test_default_executor_bounds_asyncio_to_thread():
+    loop = asyncio.new_event_loop()
+    executor = configure_bounded_default_executor(
+        loop,
+        max_workers=2,
+        max_pending=3,
+    )
+    try:
+        assert loop.run_until_complete(asyncio.to_thread(lambda: 'bounded')) == 'bounded'
+        assert executor.snapshot()['completed_total'] == 1
+        assert (
+            configure_bounded_default_executor(
+                loop,
+                max_workers=2,
+                max_pending=3,
+            )
+            is executor
+        )
+    finally:
+        executor.shutdown()
+        loop.close()
+
+
+def test_workspace_scope_is_enforced_for_asyncio_to_thread():
+    loop = asyncio.new_event_loop()
+    executor = configure_bounded_default_executor(
+        loop,
+        max_workers=2,
+        max_pending=2,
+        max_inflight_per_scope=1,
+    )
+    started = threading.Event()
+    release = threading.Event()
+
+    def block() -> str:
+        started.set()
+        release.wait(timeout=5)
+        return 'workspace-a'
+
+    async def exercise() -> None:
+        with blocking_work_scope('workspace-a'):
+            workspace_a = asyncio.create_task(asyncio.to_thread(block))
+        try:
+            while not started.is_set():
+                await asyncio.sleep(0)
+
+            with blocking_work_scope('workspace-a'):
+                with pytest.raises(
+                    BlockingWorkCapacityError,
+                    match='Workspace blocking executor capacity reached',
+                ):
+                    await asyncio.to_thread(lambda: 'rejected')
+
+            with blocking_work_scope('workspace-b'):
+                assert await asyncio.to_thread(lambda: 'workspace-b') == 'workspace-b'
+        finally:
+            release.set()
+            assert await workspace_a == 'workspace-a'
+
+    try:
+        loop.run_until_complete(exercise())
+    finally:
+        executor.shutdown()
+        loop.close()
+
+
+def test_blocking_cleanup_waits_for_capacity_instead_of_leaking_work():
+    loop = asyncio.new_event_loop()
+    executor = configure_bounded_default_executor(
+        loop,
+        max_workers=1,
+        max_pending=0,
+        max_inflight_per_scope=1,
+    )
+    started = threading.Event()
+    release = threading.Event()
+    cleaned = threading.Event()
+
+    def block() -> None:
+        started.set()
+        release.wait(timeout=5)
+
+    async def exercise() -> None:
+        blocker = asyncio.create_task(asyncio.to_thread(block))
+        while not started.is_set():
+            await asyncio.sleep(0)
+        cleanup = asyncio.create_task(run_blocking_cleanup(cleaned.set))
+        await asyncio.sleep(0.03)
+        assert not cleanup.done()
+        release.set()
+        await blocker
+        await cleanup
+
+    try:
+        loop.run_until_complete(exercise())
+        assert cleaned.is_set()
+        assert executor.snapshot()['global_rejected_total'] >= 1
+    finally:
+        release.set()
+        executor.shutdown()
+        loop.close()
+
+
+def test_blocking_atomic_waits_for_thread_before_propagating_cancellation():
+    loop = asyncio.new_event_loop()
+    executor = configure_bounded_default_executor(
+        loop,
+        max_workers=1,
+        max_pending=1,
+        max_inflight_per_scope=1,
+    )
+    started = threading.Event()
+    release = threading.Event()
+    completed = threading.Event()
+
+    def block() -> None:
+        started.set()
+        release.wait(timeout=5)
+        completed.set()
+
+    async def exercise() -> None:
+        operation = asyncio.create_task(run_blocking_atomic(block))
+        while not started.is_set():
+            await asyncio.sleep(0)
+        operation.cancel()
+        await asyncio.sleep(0)
+        assert not operation.done()
+        release.set()
+        with pytest.raises(asyncio.CancelledError):
+            await operation
+
+    try:
+        loop.run_until_complete(exercise())
+        assert completed.is_set()
+    finally:
+        release.set()
+        executor.shutdown()
+        loop.close()
+
+
+@pytest.mark.parametrize(
+    ('max_workers', 'max_pending'),
+    [
+        (0, 1),
+        (65, 1),
+        (1, -1),
+        (1, 4097),
+        (True, 1),
+    ],
+)
+def test_bounded_executor_rejects_unsafe_limits(
+    max_workers,
+    max_pending,
+):
+    with pytest.raises(ValueError):
+        BoundedThreadPoolExecutor(
+            max_workers=max_workers,
+            max_pending=max_pending,
+        )
+
+
+@pytest.mark.parametrize(
+    ('max_workers', 'max_inflight_per_scope'),
+    [(8, 0), (8, 4097), (8, True), (8, 5), (2, 2)],
+)
+def test_bounded_executor_rejects_unsafe_scope_limits(
+    max_workers,
+    max_inflight_per_scope,
+):
+    with pytest.raises(ValueError):
+        BoundedThreadPoolExecutor(
+            max_workers=max_workers,
+            max_inflight_per_scope=max_inflight_per_scope,
+        )