perf: enqueue Cloud Tasks batches concurrently

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

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))