fix(signals): bound poison batch retries in grouping v2

products/signals/backend/temporal/grouping.py

@@ -1062,6 +1062,7 @@ def _augment_candidates_with_batch(
 async def _process_signal_batch(
     batch: list[EmitSignalInputs],
     cached_type_examples: Optional[FetchSignalTypeExamplesOutput] = None,
+    emit_prep_drops: bool = True,
 ) -> tuple[int, FetchSignalTypeExamplesOutput]:
     """
     Process a batch of signals with parallel preparation (steps 1-4) and sequential
@@ -1071,6 +1072,10 @@ async def _process_signal_batch(
     Earlier signals in the batch are injected into later signals' candidate sets via
     local cosine distance comparison, eliminating the need for per-signal CH waits
     within a batch.
+
+    A prep-phase (steps 1-4) failure re-raises so the caller can decide whether to retry.
+    Callers that retry the batch pass emit_prep_drops=False and own the terminal
+    signal_dropped telemetry, so a deterministic failure is not counted once per attempt.
     """
     team_id = batch[0].team_id
     # Purely defensive
@@ -1195,7 +1200,8 @@ async def _process_signal_batch(
             team_id=team_id,
             batch_size=len(batch),
         )
-        await asyncio.gather(*(capture_signal_dropped(signal, e, stage="grouping_prep") for signal in batch))
+        if emit_prep_drops:
+            await asyncio.gather(*(capture_signal_dropped(signal, e, stage="grouping_prep") for signal in batch))
         raise
 
     # === SEQUENTIAL PHASE (steps 5-7) ===

products/signals/backend/temporal/grouping_v2.py

@@ -17,6 +17,7 @@
 from posthog.temporal.common.client import async_connect
 from posthog.temporal.common.scoped import scoped_temporal
 
+from products.signals.backend.temporal.drop_telemetry import capture_signal_dropped
 from products.signals.backend.temporal.grouping import (
     TYPE_EXAMPLES_CACHE_TTL,
     FetchSignalTypeExamplesOutput,
@@ -36,12 +37,18 @@
 BATCH_COLLECT_MAX_SIGNALS = 20
 BATCH_COLLECT_TIMEOUT = timedelta(seconds=30)
 RETRY_BACKOFF = timedelta(seconds=10)
+# A batch key that fails processing this many times is dead-lettered: its signals are dropped
+# and the key leaves the buffer, so a deterministic failure cannot jam the team's pipeline.
+MAX_BATCH_ATTEMPTS = 3
 
 # Patch ID for the multi-batch collection change. Once all in-flight workflows
 # that recorded history under the old single-batch code have drained (they
 # continue_as_new every iteration, so this is fast), replace this
 # workflow.patched() call with workflow.deprecate_patch() and eventually remove.
 _PATCH_COLLECT_BATCH = "collect-batch-signals-v1"
+# Patch ID for bounding retries and dead-lettering poison batches. In-flight runs replay the
+# old unbounded stash-and-sleep path; fresh runs take the bounded path.
+_PATCH_BOUNDED_RETRY = "bounded-batch-retry-v1"
 
 
 @dataclass
@@ -50,6 +57,8 @@ class CollectedBatch:
 
     signals: list[EmitSignalInputs] = field(default_factory=list)
     object_keys: list[str] = field(default_factory=list)
+    # Signals grouped by their source S3 key, so a dead-lettered key drops only its own signals.
+    signals_by_key: dict[str, list[EmitSignalInputs]] = field(default_factory=dict)
 
 
 @activity.defn
@@ -78,6 +87,7 @@ class TeamSignalGroupingV2Workflow:
 
     def __init__(self) -> None:
         self._batch_key_buffer: list[str] = []
+        self._batch_key_attempts: dict[str, int] = {}
         self._cached_type_examples: Optional[FetchSignalTypeExamplesOutput] = None
         self._type_examples_fetched_at: Optional[datetime] = None
         self._paused_until: Optional[datetime] = None
@@ -120,6 +130,7 @@ def _continue_as_new(self, input: TeamSignalGroupingV2Input) -> None:
                 team_id=input.team_id,
                 pending_batch_keys=list(self._batch_key_buffer),
                 paused_until=self._paused_until,
+                batch_key_attempts=dict(self._batch_key_attempts),
             )
         )
 
@@ -155,6 +166,7 @@ async def _collect_next_batch(self) -> CollectedBatch:
             )
 
             collected.signals.extend(read_result.signals)
+            collected.signals_by_key[object_key] = read_result.signals
 
         return collected
 
@@ -221,6 +233,84 @@ async def _run_new_collect_batch_path(self, input: TeamSignalGroupingV2Input) ->
                 self._batch_buffer_size_gauge.set(len(self._batch_key_buffer))
             return
 
+        if workflow.patched(_PATCH_BOUNDED_RETRY):
+            await self._process_collected(input, collected)
+        else:
+            await self._process_collected_unbounded(input, collected)
+
+        # continue_as_new after each processing round to keep history bounded.
+        # Carry over any pending keys that arrived while we were processing.
+        self._continue_as_new(input)
+
+    async def _process_collected(self, input: TeamSignalGroupingV2Input, collected: CollectedBatch) -> None:
+        """Process a collected batch, bounding retries so a poison batch cannot jam the buffer."""
+        try:
+            # emit_prep_drops=False: a prep failure is retried, so this path owns the terminal
+            # signal_dropped telemetry (see _dead_letter) rather than emitting once per attempt.
+            dropped, _type_examples = await _process_signal_batch(collected.signals, emit_prep_drops=False)
+        except Exception as error:
+            logger.exception(
+                "Failed to process signal batch",
+                team_id=input.team_id,
+                batch_size=len(collected.signals),
+                batch_keys=collected.object_keys,
+            )
+            await self._handle_batch_failure(input, collected, error)
+            return
+
+        for key in collected.object_keys:
+            self._batch_key_attempts.pop(key, None)
+        if self._signals_processed_counter is not None:
+            self._signals_processed_counter.add(len(collected.signals))
+        if self._signals_dropped_counter is not None and dropped > 0:
+            self._signals_dropped_counter.add(dropped)
+
+    async def _handle_batch_failure(
+        self, input: TeamSignalGroupingV2Input, collected: CollectedBatch, error: BaseException
+    ) -> None:
+        """Count the attempt per key, dead-letter keys at the cap, and re-stash the rest for retry."""
+        retry_keys: list[str] = []
+        dead_keys: list[str] = []
+        for key in collected.object_keys:
+            attempts = self._batch_key_attempts.get(key, 0) + 1
+            self._batch_key_attempts[key] = attempts
+            (dead_keys if attempts >= MAX_BATCH_ATTEMPTS else retry_keys).append(key)
+
+        if dead_keys:
+            await self._dead_letter(input, collected, dead_keys, error)
+            for key in dead_keys:
+                self._batch_key_attempts.pop(key, None)
+
+        # Stash keys still under the attempt cap back at the head so they retry after continue_as_new.
+        self._batch_key_buffer = retry_keys + self._batch_key_buffer
+        if self._batch_buffer_size_gauge is not None:
+            self._batch_buffer_size_gauge.set(len(self._batch_key_buffer))
+        if retry_keys:
+            # Sleep before retrying to avoid hot-looping on a deterministic failure.
+            await workflow.sleep(RETRY_BACKOFF)
+
+    async def _dead_letter(
+        self,
+        input: TeamSignalGroupingV2Input,
+        collected: CollectedBatch,
+        dead_keys: list[str],
+        error: BaseException,
+    ) -> None:
+        """Drop signals whose batch key hit the attempt cap, and emit their drop telemetry once."""
+        dead_signals = [signal for key in dead_keys for signal in collected.signals_by_key.get(key, [])]
+        logger.error(
+            "Dropping poison signal batch after repeated failures",
+            team_id=input.team_id,
+            dead_keys=dead_keys,
+            signal_count=len(dead_signals),
+        )
+        if self._signals_dropped_counter is not None and dead_signals:
+            self._signals_dropped_counter.add(len(dead_signals))
+        for signal in dead_signals:
+            await capture_signal_dropped(signal, error, stage="grouping_prep")
+
+    async def _process_collected_unbounded(self, input: TeamSignalGroupingV2Input, collected: CollectedBatch) -> None:
+        """Legacy unbounded stash-and-sleep path, kept for in-flight runs replaying old history."""
         try:
             dropped, _type_examples = await _process_signal_batch(collected.signals)
             if self._signals_processed_counter is not None:
@@ -241,10 +331,6 @@ async def _run_new_collect_batch_path(self, input: TeamSignalGroupingV2Input) ->
                 self._batch_buffer_size_gauge.set(len(self._batch_key_buffer))
             await workflow.sleep(RETRY_BACKOFF)
 
-        # continue_as_new after each processing round to keep history bounded.
-        # Carry over any pending keys that arrived while we were processing.
-        self._continue_as_new(input)
-
     @temporalio.workflow.run
     async def run(self, input: TeamSignalGroupingV2Input) -> None:
         with posthoganalytics.new_context(capture_exceptions=False):
@@ -255,6 +341,7 @@ async def run(self, input: TeamSignalGroupingV2Input) -> None:
     async def _run_impl(self, input: TeamSignalGroupingV2Input) -> None:
         # Restore state carried over from continue_as_new
         self._batch_key_buffer.extend(input.pending_batch_keys)
+        self._batch_key_attempts = dict(input.batch_key_attempts)
         self._paused_until = input.paused_until
         start_time = workflow.now()
 

products/signals/backend/temporal/types.py

@@ -125,6 +125,10 @@ class TeamSignalGroupingV2Input:
     team_id: int
     pending_batch_keys: list[str] = field(default_factory=list)
     paused_until: Optional[datetime] = None
+    # Failed-processing attempt count per S3 batch key, carried across continue_as_new so a
+    # poison batch is bounded rather than retried forever. Defaults empty so histories written
+    # before this field replay unchanged.
+    batch_key_attempts: dict[str, int] = field(default_factory=dict)
 
 
 @dataclass

products/signals/backend/test/test_grouping_v2_workflow.py

@@ -0,0 +1,89 @@
+import uuid
+import asyncio
+from datetime import timedelta
+
+import pytest
+from unittest.mock import patch
+
+from temporalio import activity
+from temporalio.exceptions import ApplicationError
+from temporalio.testing import WorkflowEnvironment
+from temporalio.worker import UnsandboxedWorkflowRunner, Worker
+
+from products.signals.backend.temporal.drop_telemetry import CaptureSignalDroppedInput
+from products.signals.backend.temporal.grouping_v2 import MAX_BATCH_ATTEMPTS, TeamSignalGroupingV2Workflow
+from products.signals.backend.temporal.signal_queries import FetchSignalTypeExamplesInput
+from products.signals.backend.temporal.types import (
+    EmitSignalInputs,
+    ReadSignalsFromS3Input,
+    ReadSignalsFromS3Output,
+    TeamSignalGroupingV2Input,
+)
+
+TASK_QUEUE = "test-grouping-v2-queue"
+GROUPING_V2_MODULE = "products.signals.backend.temporal.grouping_v2"
+
+
+class _Recorder:
+    def __init__(self) -> None:
+        self.reads = 0
+        self.drops = 0
+        self.dropped = asyncio.Event()
+
+
+@pytest.mark.asyncio
+async def test_poison_batch_is_dead_lettered_after_attempt_cap():
+    # A batch whose preparation always fails must not retry forever: it is retried up to the cap,
+    # then dead-lettered so the buffer drains, and its one signal drops exactly once — not once per
+    # attempt, which is what jammed a team's pipeline and inflated the signal_dropped metric.
+    recorder = _Recorder()
+    signal = EmitSignalInputs(
+        team_id=1,
+        source_product="error_tracking",
+        source_type="issue",
+        source_id=str(uuid.uuid4()),
+        description="poison signal",
+    )
+
+    @activity.defn(name="read_signals_from_s3_activity")
+    async def fake_read(_input: ReadSignalsFromS3Input) -> ReadSignalsFromS3Output:
+        recorder.reads += 1
+        return ReadSignalsFromS3Output(signals=[signal])
+
+    @activity.defn(name="fetch_signal_type_examples_activity")
+    async def fake_prep_fails(_input: FetchSignalTypeExamplesInput) -> None:
+        raise ApplicationError("embedding API 500", non_retryable=True)
+
+    @activity.defn(name="capture_signal_dropped_activity")
+    async def fake_drop(_input: CaptureSignalDroppedInput) -> None:
+        recorder.drops += 1
+        recorder.dropped.set()
+
+    # Shrink the collection window and retry backoff so the retry loop runs in milliseconds instead
+    # of the ~44s-per-round wall clock it takes in production.
+    with (
+        patch(f"{GROUPING_V2_MODULE}.BATCH_COLLECT_TIMEOUT", timedelta(seconds=1)),
+        patch(f"{GROUPING_V2_MODULE}.RETRY_BACKOFF", timedelta(seconds=0)),
+    ):
+        async with await WorkflowEnvironment.start_time_skipping() as env:
+            async with Worker(
+                env.client,
+                task_queue=TASK_QUEUE,
+                workflows=[TeamSignalGroupingV2Workflow],
+                activities=[fake_read, fake_prep_fails, fake_drop],
+                workflow_runner=UnsandboxedWorkflowRunner(),
+            ):
+                handle = await env.client.start_workflow(
+                    TeamSignalGroupingV2Workflow.run,
+                    TeamSignalGroupingV2Input(team_id=1),
+                    id=f"grouping-v2-{uuid.uuid4()}",
+                    task_queue=TASK_QUEUE,
+                )
+                await handle.signal(TeamSignalGroupingV2Workflow.submit_batch, "signals/batch/poison")
+                await asyncio.wait_for(recorder.dropped.wait(), timeout=30)
+                await env.client.get_workflow_handle(handle.id).terminate()
+
+    # Read once per attempt: the key is re-collected each round until the cap dead-letters it.
+    assert recorder.reads == MAX_BATCH_ATTEMPTS
+    # Emitted once for the signal, not once per attempt.
+    assert recorder.drops == 1