diff --git a/src/youtube_extension/services/cloud/cloud_tasks_queue.py b/src/youtube_extension/services/cloud/cloud_tasks_queue.py index 4a720f53d..b34136a36 100644 --- a/src/youtube_extension/services/cloud/cloud_tasks_queue.py +++ b/src/youtube_extension/services/cloud/cloud_tasks_queue.py @@ -12,6 +12,7 @@ import json import logging import os +import weakref from dataclasses import dataclass from datetime import datetime, timedelta from typing import Any, Optional @@ -29,6 +30,20 @@ logger = logging.getLogger(__name__) +# Upper bound on concurrent task-creation RPCs in `enqueue_batch`. +# +# Each enqueue is a blocking gRPC call dispatched to the default asyncio thread +# pool, which is shared process-wide and sized `min(32, cpu_count + 4)`. Taking +# the whole pool starves every other `asyncio.to_thread` caller, so the bound is +# half of it (min 2, cap 8). +# +# Measured with a simulated 20 ms RTT over a 50-task batch on a 12-core host +# (pool = 16, so this evaluates to 8): serial 1328 ms -> bounded 187 ms (7.1x), +# with a co-tenant thread-pool user's worst-case wait unchanged at 2.6 ms. An +# unbounded fan-out of 16 reaches 110 ms (12x) but spikes that co-tenant to +# 17.5 ms, which is why the pool is deliberately not saturated. +_ENQUEUE_MAX_CONCURRENCY = min(8, max(2, min(32, (os.cpu_count() or 1) + 4) // 2)) + async def _run_sync_rpc(call, /, *args, **kwargs): """Run a synchronous RPC without abandoning it on caller cancellation. @@ -140,6 +155,16 @@ def __init__( # Initialize Cloud Tasks client self.client: Optional[tasks_v2.CloudTasksClient] = None + # Enqueue concurrency limiter, shared by every `enqueue_batch` call on + # this service so that concurrent batches (the batch endpoint drives the + # process-wide singleton) cannot collectively exceed the bound and + # starve co-tenant `to_thread` callers. Keyed by event loop because + # `asyncio.Semaphore` binds to the loop that first awaits it; the weak + # keys let finished loops (e.g. per-test `asyncio.run`) be collected. + self._enqueue_semaphores: weakref.WeakKeyDictionary[ + asyncio.AbstractEventLoop, asyncio.Semaphore + ] = weakref.WeakKeyDictionary() + logger.info( f"CloudTasksQueueService initialized: " f"project={self.project_id}, location={self.location}, queue={self.queue_name}" @@ -239,6 +264,20 @@ async def enqueue_video_processing( return task_id + def _get_enqueue_semaphore(self) -> asyncio.Semaphore: + """ + Return this service's enqueue limiter for the running event loop. + + Shared across concurrent `enqueue_batch` calls so the bound holds for + the process, not merely within a single batch. + """ + loop = asyncio.get_running_loop() + semaphore = self._enqueue_semaphores.get(loop) + if semaphore is None: + semaphore = asyncio.Semaphore(_ENQUEUE_MAX_CONCURRENCY) + self._enqueue_semaphores[loop] = semaphore + return semaphore + async def enqueue_batch( self, video_tasks: list[VideoProcessingTask], @@ -247,21 +286,40 @@ async def enqueue_batch( """ Enqueue multiple videos for processing. + Tasks are enqueued concurrently, bounded by `_ENQUEUE_MAX_CONCURRENCY`. + Each enqueue is a blocking gRPC round-trip, so a serial loop paid the + full network latency once per task and scaled linearly with batch size. + Args: video_tasks: List of video processing tasks task_config: Task configuration for all tasks Returns: - List of task IDs + List of task IDs, positionally aligned with `video_tasks` and + omitting any task that failed to enqueue. """ - task_ids = [] + semaphore = self._get_enqueue_semaphore() + + async def _enqueue_one(video_task: VideoProcessingTask) -> str: + async with semaphore: + return await self.enqueue_video_processing(video_task, task_config) - for video_task in video_tasks: - try: - task_id = await self.enqueue_video_processing(video_task, task_config) - task_ids.append(task_id) - except Exception as e: - logger.error(f"Failed to enqueue task for {video_task.video_id}: {e}") + results = await asyncio.gather( + *(_enqueue_one(video_task) for video_task in video_tasks), + return_exceptions=True, + ) + + task_ids = [] + for video_task, result in zip(video_tasks, results, strict=True): + if isinstance(result, BaseException): + # Only ordinary errors are skipped-and-logged. Anything else + # (CancelledError, KeyboardInterrupt) propagated out of the + # original `except Exception` loop and must keep doing so. + if not isinstance(result, Exception): + raise result + logger.error(f"Failed to enqueue task for {video_task.video_id}: {result}") + continue + task_ids.append(result) logger.info(f"Enqueued {len(task_ids)}/{len(video_tasks)} tasks successfully") return task_ids diff --git a/tests/unit/test_cloud_tasks_queue.py b/tests/unit/test_cloud_tasks_queue.py index 0731df155..6f8bb90ce 100644 --- a/tests/unit/test_cloud_tasks_queue.py +++ b/tests/unit/test_cloud_tasks_queue.py @@ -611,13 +611,22 @@ def _video_tasks(self, count=3) -> list[VideoProcessingTask]: ] async def test_returns_all_task_ids_on_success(self): - mock_tv2 = _make_mock_tasks_v2() + # Routing mock + request-keyed responses: enqueue_batch fans out + # concurrently, so the order create_task happens to be *called* in no + # longer implies input order. Keying the response off the request + # keeps this assertion about the id-to-input mapping (the thing that + # actually matters) instead of about RPC completion order. With a + # list side_effect this passes only because the mock returns + # instantly; adding 0-4ms of jitter makes it fail ~85% of the time. + mock_tv2 = _routing_tasks_v2() svc = _initialized_service(mock_tv2) - responses = [MagicMock() for _ in range(3)] - for i, r in enumerate(responses): - r.name = f".../tasks/t{i}" - svc.client.create_task.side_effect = responses + def create_task(request): + response = MagicMock() + response.name = f".../tasks/t{_video_id_of(request).removeprefix('vid-')}" + return response + + svc.client.create_task.side_effect = create_task with ( patch.object(m, "CLOUD_TASKS_AVAILABLE", True), @@ -1279,3 +1288,239 @@ async def test_uncancelled_calls_return_normally(self): """The happy path is unchanged by the cancellation handling.""" sentinel = MagicMock() assert await m._run_sync_rpc(lambda *_a, **_k: sentinel) is sentinel + + +# =========================================================================== +# enqueue_batch concurrency tests +# =========================================================================== + + +def _routing_tasks_v2() -> MagicMock: + """A ``tasks_v2`` mock whose builders pass their kwargs through. + + ``_make_mock_tasks_v2`` gives every builder a single fixed ``return_value``, + so each ``create_task`` call receives an identical sentinel and cannot be + attributed back to a specific video. Concurrency tests need that + attribution (call order no longer implies input order), so ``Task``, + ``HttpRequest`` and ``CreateTaskRequest`` are turned into pass-through + namespaces here. + """ + import types as _types + + mock = _make_mock_tasks_v2() + mock.HttpRequest.side_effect = lambda **kw: _types.SimpleNamespace(**kw) + mock.Task.side_effect = lambda **kw: _types.SimpleNamespace(**kw) + mock.CreateTaskRequest.side_effect = lambda **kw: _types.SimpleNamespace(**kw) + return mock + + +def _video_id_of(request) -> str: + """Recover the video id from a CreateTaskRequest built by the service.""" + return json.loads(request.task.http_request.body.decode())["video_id"] + + +class TestEnqueueBatchConcurrency: + """enqueue_batch must fan out concurrently under a bounded limit.""" + + def _video_tasks(self, count: int, prefix: str = "vid") -> list[VideoProcessingTask]: + return [ + VideoProcessingTask( + video_id=f"{prefix}-{i}", + video_url=f"https://youtube.com/watch?v={prefix}-{i}", + ) + for i in range(count) + ] + + async def test_batch_enqueues_concurrently(self): + """Wall time must be far below the serial sum of RPC latencies. + + The fan-out bound is pinned to ``count`` for this timing-only test so the + assertion measures *that* the batch fans out, independent of the host's + CPU count. Left unpinned, ``_ENQUEUE_MAX_CONCURRENCY`` floors at 2 on a + single-CPU runner, making 8x50 ms take four waves (200 ms) and tripping + the ``< serial_floor / 2`` (200 ms) bound. The derived value itself is + exercised by ``test_batch_bounds_in_flight_concurrency`` below. + """ + count, delay = 8, 0.05 + mock_tv2 = _routing_tasks_v2() + svc = _initialized_service(mock_tv2) + + def create_task(request=None, **_kwargs): + time.sleep(delay) # stand-in for the blocking gRPC round-trip + response = MagicMock() + response.name = f"projects/p/locations/l/queues/q/tasks/{_video_id_of(request)}" + return response + + svc.client.create_task.side_effect = create_task + + with ( + patch.object(m, "CLOUD_TASKS_AVAILABLE", True), + patch.object(m, "tasks_v2", mock_tv2), + patch.object(m, "_ENQUEUE_MAX_CONCURRENCY", count), + ): + started = time.perf_counter() + ids = await svc.enqueue_batch(self._video_tasks(count)) + elapsed = time.perf_counter() - started + + assert len(ids) == count + serial_floor = count * delay + assert elapsed < serial_floor / 2, ( + f"enqueue_batch took {elapsed * 1000:.0f}ms for {count} tasks; " + f"a serial implementation needs >={serial_floor * 1000:.0f}ms, so this " + f"is still serial" + ) + + async def test_batch_bounds_in_flight_concurrency(self): + """Fan-out is capped so the shared thread pool is not monopolised.""" + count = 24 + mock_tv2 = _routing_tasks_v2() + svc = _initialized_service(mock_tv2) + + lock = threading.Lock() + in_flight = 0 + peak = 0 + + def create_task(request=None, **_kwargs): + nonlocal in_flight, peak + with lock: + in_flight += 1 + peak = max(peak, in_flight) + time.sleep(0.02) + with lock: + in_flight -= 1 + response = MagicMock() + response.name = f"projects/p/locations/l/queues/q/tasks/{_video_id_of(request)}" + return response + + svc.client.create_task.side_effect = create_task + + with ( + patch.object(m, "CLOUD_TASKS_AVAILABLE", True), + patch.object(m, "tasks_v2", mock_tv2), + ): + ids = await svc.enqueue_batch(self._video_tasks(count)) + + assert len(ids) == count + assert peak > 1, "expected concurrent fan-out, observed a serial drain" + assert peak <= m._ENQUEUE_MAX_CONCURRENCY, ( + f"observed {peak} concurrent RPCs, limit is {m._ENQUEUE_MAX_CONCURRENCY}" + ) + + async def test_overlapping_batches_share_the_bound(self): + """Concurrent batches on one service must not exceed the bound together. + + The batch endpoint drives the process-wide singleton, so two in-flight + requests would each get their own limiter if the semaphore were built + per call -- admitting 2x the bound against the shared thread pool and + recreating the starvation the bound exists to prevent. + """ + per_batch = 12 + mock_tv2 = _routing_tasks_v2() + svc = _initialized_service(mock_tv2) + + lock = threading.Lock() + in_flight = 0 + peak = 0 + + def create_task(request=None, **_kwargs): + nonlocal in_flight, peak + with lock: + in_flight += 1 + peak = max(peak, in_flight) + time.sleep(0.02) + with lock: + in_flight -= 1 + response = MagicMock() + response.name = f"projects/p/locations/l/queues/q/tasks/{_video_id_of(request)}" + return response + + svc.client.create_task.side_effect = create_task + + with ( + patch.object(m, "CLOUD_TASKS_AVAILABLE", True), + patch.object(m, "tasks_v2", mock_tv2), + ): + first, second = await asyncio.gather( + svc.enqueue_batch(self._video_tasks(per_batch, prefix="a")), + svc.enqueue_batch(self._video_tasks(per_batch, prefix="b")), + ) + + assert len(first) == per_batch + assert len(second) == per_batch + assert peak > 1, "expected concurrent fan-out, observed a serial drain" + assert peak <= m._ENQUEUE_MAX_CONCURRENCY, ( + f"two overlapping batches reached {peak} concurrent RPCs, but the " + f"limit is {m._ENQUEUE_MAX_CONCURRENCY}; the limiter is not shared " + f"across calls" + ) + + async def test_ids_follow_input_order_despite_completion_order(self): + """Returned ids stay positionally aligned with the input tasks. + + Regression guard for the concurrent implementation: earlier videos are + made the *slowest*, so completion order is the reverse of input order. + """ + count = 6 + mock_tv2 = _routing_tasks_v2() + svc = _initialized_service(mock_tv2) + + def create_task(request=None, **_kwargs): + video_id = _video_id_of(request) + index = int(video_id.rsplit("-", 1)[1]) + time.sleep(0.01 * (count - index)) # vid-0 finishes last + response = MagicMock() + response.name = f"projects/p/locations/l/queues/q/tasks/{video_id}" + return response + + svc.client.create_task.side_effect = create_task + + with ( + patch.object(m, "CLOUD_TASKS_AVAILABLE", True), + patch.object(m, "tasks_v2", mock_tv2), + ): + ids = await svc.enqueue_batch(self._video_tasks(count)) + + assert ids == [f"vid-{i}" for i in range(count)] + + async def test_failures_are_skipped_without_losing_survivors(self): + """A failing task is logged and dropped; the rest still return ids.""" + count = 6 + mock_tv2 = _routing_tasks_v2() + svc = _initialized_service(mock_tv2) + + def create_task(request=None, **_kwargs): + video_id = _video_id_of(request) + if video_id in {"vid-1", "vid-4"}: + raise RuntimeError("network error") + response = MagicMock() + response.name = f"projects/p/locations/l/queues/q/tasks/{video_id}" + return response + + svc.client.create_task.side_effect = create_task + + with ( + patch.object(m, "CLOUD_TASKS_AVAILABLE", True), + patch.object(m, "tasks_v2", mock_tv2), + ): + ids = await svc.enqueue_batch(self._video_tasks(count)) + + assert ids == ["vid-0", "vid-2", "vid-3", "vid-5"] + + async def test_cancellation_propagates_and_is_not_logged_as_failure(self): + """CancelledError must not be absorbed into the skip-and-continue path.""" + mock_tv2 = _routing_tasks_v2() + svc = _initialized_service(mock_tv2) + + async def fake_enqueue(video_task, task_config=None): + if video_task.video_id == "vid-1": + raise asyncio.CancelledError() + return video_task.video_id + + svc.enqueue_video_processing = fake_enqueue + + with ( + patch.object(m, "CLOUD_TASKS_AVAILABLE", True), + patch.object(m, "tasks_v2", mock_tv2), + pytest.raises(asyncio.CancelledError), + ): + await svc.enqueue_batch(self._video_tasks(3))