Fix thread-queue-parallel consumer to handle Kafka rebalances correctly (#95753)

The thread-queue-parallel consumer was encountering errors during Kafka
rebalances. This was caused by the commit thread continuing to run after
partition revocation and attempting to commit offsets for partitions no
longer owned by the consumer.

I also noticed a few other problems with the structure of things and
improve those as well.

- Move commit thread ownership from individual strategies to
FixedQueuePool to survive rebalances
- Add partition assignment validation to reject messages from unassigned
partitions
- Implement OffsetTracker.clear() to reset state during partition
revocation
- Add update_assignments() to atomically update partitions and commit
function

<!-- Describe your PR here. -->

Author: wedamija

src/sentry/remote_subscriptions/consumers/queue_consumer.py

@@ -12,6 +12,7 @@
 import sentry_sdk
 from arroyo.backends.kafka.consumer import KafkaPayload
 from arroyo.processing.strategies import ProcessingStrategy
+from arroyo.processing.strategies.abstract import MessageRejected
 from arroyo.types import BrokerValue, FilteredPayload, Message, Partition
 
 from sentry.utils import metrics
@@ -21,6 +22,12 @@
 T = TypeVar("T")
 
 
+class UnassignedPartitionError(Exception):
+    """Raised when trying to track offsets for an unassigned partition."""
+
+    pass
+
+
 @dataclass
 class WorkItem(Generic[T]):
     """Work item that includes the original message for offset tracking."""
@@ -47,20 +54,25 @@ def __init__(self) -> None:
         self.partition_locks: dict[Partition, threading.Lock] = {}
 
     def _get_partition_lock(self, partition: Partition) -> threading.Lock:
-        """Get or create a lock for a partition."""
-        lock = self.partition_locks.get(partition)
-        if lock:
-            return lock
-        return self.partition_locks.setdefault(partition, threading.Lock())
+        """Get the lock for a partition."""
+        return self.partition_locks[partition]
 
     def add_offset(self, partition: Partition, offset: int) -> None:
         """Record that we've started processing an offset."""
+        if partition not in self.partition_locks:
+            raise UnassignedPartitionError(
+                f"Partition {partition} is not assigned to this consumer"
+            )
+
         with self._get_partition_lock(partition):
             self.all_offsets[partition].add(offset)
             self.outstanding[partition].add(offset)
 
     def complete_offset(self, partition: Partition, offset: int) -> None:
         """Mark an offset as completed."""
+        if partition not in self.partition_locks:
+            return
+
         with self._get_partition_lock(partition):
             self.outstanding[partition].discard(offset)
 
@@ -104,6 +116,18 @@ def mark_committed(self, partition: Partition, offset: int) -> None:
             # Remove all offsets <= committed offset
             self.all_offsets[partition] = {o for o in self.all_offsets[partition] if o > offset}
 
+    def clear(self) -> None:
+        """Clear all offset tracking state."""
+        self.all_offsets.clear()
+        self.outstanding.clear()
+        self.last_committed.clear()
+        self.partition_locks.clear()
+
+    def update_assignments(self, partitions: set[Partition]) -> None:
+        """Update partition assignments and reset all tracking state."""
+        self.clear()
+        self.partition_locks = {partition: threading.Lock() for partition in partitions}
+
 
 class OrderedQueueWorker(threading.Thread, Generic[T]):
     """Worker thread that processes items from a queue in order."""
@@ -138,9 +162,6 @@ def run(self) -> None:
                     name=f"monitors.{self.identifier}.worker_{self.worker_id}",
                 ):
                     self.result_processor(self.identifier, work_item.result)
-
-            except queue.ShutDown:
-                break
             except Exception:
                 logger.exception(
                     "Unexpected error in queue worker", extra={"worker_id": self.worker_id}
@@ -173,13 +194,20 @@ def __init__(
         result_processor: Callable[[str, T], None],
         identifier: str,
         num_queues: int = 20,
+        commit_interval: float = 1.0,
     ) -> None:
         self.result_processor = result_processor
         self.identifier = identifier
         self.num_queues = num_queues
+        self.commit_interval = commit_interval
         self.offset_tracker = OffsetTracker()
         self.queues: list[queue.Queue[WorkItem[T]]] = []
         self.workers: list[OrderedQueueWorker[T]] = []
+        self.commit_function: Callable[[dict[Partition, int]], None] | None = None
+        self.commit_shutdown_event = threading.Event()
+
+        self.commit_thread = threading.Thread(target=self._commit_loop, daemon=True)
+        self.commit_thread.start()
 
         for i in range(num_queues):
             work_queue: queue.Queue[WorkItem[T]] = queue.Queue()
@@ -195,6 +223,29 @@ def __init__(
             worker.start()
             self.workers.append(worker)
 
+    def _commit_loop(self) -> None:
+        """Background thread that periodically commits offsets."""
+        while not self.commit_shutdown_event.is_set():
+            try:
+                self.commit_shutdown_event.wait(self.commit_interval)
+                if self.commit_shutdown_event.is_set():
+                    break
+
+                committable = self.offset_tracker.get_committable_offsets()
+
+                if committable and self.commit_function:
+                    metrics.incr(
+                        "remote_subscriptions.queue_pool.offsets_committed",
+                        len(committable),
+                        tags={"identifier": self.identifier},
+                    )
+
+                    self.commit_function(committable)
+                    for partition, offset in committable.items():
+                        self.offset_tracker.mark_committed(partition, offset)
+            except Exception:
+                logger.exception("Error in commit loop")
+
     def get_queue_for_group(self, group_key: str) -> int:
         """
         Get queue index for a group using consistent hashing.
@@ -205,10 +256,25 @@ def submit(self, group_key: str, work_item: WorkItem[T]) -> None:
         """
         Submit a work item to the appropriate queue.
         """
+        try:
+            self.offset_tracker.add_offset(work_item.partition, work_item.offset)
+        except UnassignedPartitionError:
+            logger.exception(
+                "Received message for unassigned partition, skipping",
+                extra={
+                    "partition": work_item.partition,
+                    "offset": work_item.offset,
+                    "identifier": self.identifier,
+                },
+            )
+            metrics.incr(
+                "remote_subscriptions.queue_pool.submit.unassigned_partition",
+                tags={"identifier": self.identifier},
+            )
+            return
+
         queue_index = self.get_queue_for_group(group_key)
         work_queue = self.queues[queue_index]
-
-        self.offset_tracker.add_offset(work_item.partition, work_item.offset)
         work_queue.put(work_item)
 
     def get_stats(self) -> dict[str, Any]:
@@ -219,7 +285,7 @@ def get_stats(self) -> dict[str, Any]:
             "total_items": sum(queue_depths),
         }
 
-    def wait_until_empty(self, timeout: float = 5.0) -> bool:
+    def wait_until_empty(self, timeout: float) -> bool:
         """Wait until all queues are empty. Returns True if successful, False if timeout."""
         start_time = time.time()
         while time.time() - start_time < timeout:
@@ -228,8 +294,61 @@ def wait_until_empty(self, timeout: float = 5.0) -> bool:
             time.sleep(0.01)
         return False
 
+    def flush(self, timeout: float | None = None) -> bool:
+        """
+        Wait for all queues to be empty. Returns True if successful, False if timeout.
+        If timeout is None, immediately flush without waiting.
+        If timeout is reached, flushes all remaining work.
+        """
+        if timeout is None:
+            success = False
+        else:
+            success = self.wait_until_empty(timeout)
+        if not success:
+            metrics.incr(
+                "remote_subscriptions.queue_pool.flush.timeout",
+                tags={"identifier": self.identifier},
+            )
+            cleared_count = 0
+            for q in self.queues:
+                while not q.empty():
+                    try:
+                        q.get_nowait()
+                        cleared_count += 1
+                    except queue.Empty:
+                        break
+                    except Exception:
+                        logger.exception("Error clearing queue")
+            if cleared_count > 0:
+                metrics.incr(
+                    "remote_subscriptions.queue_pool.timeout_queue_size",
+                    cleared_count,
+                    tags={"identifier": self.identifier},
+                )
+
+        self.offset_tracker.clear()
+        return success
+
+    def update_assignments(
+        self,
+        partitions: set[Partition],
+        commit_function: Callable[[dict[Partition, int]], None],
+    ) -> None:
+        """
+        Update partition assignments and commit function atomically.
+        """
+        self.offset_tracker.update_assignments(partitions)
+        self.commit_function = commit_function
+
+        logger.info(
+            "Updated partition assignments",
+            extra={
+                "identifier": self.identifier,
+                "partitions": len(partitions),
+            },
+        )
+
     def shutdown(self) -> None:
-        """Gracefully shutdown all workers."""
         for worker in self.workers:
             worker.shutdown = True
 
@@ -240,7 +359,10 @@ def shutdown(self) -> None:
                 logger.exception("Error shutting down queue")
 
         for worker in self.workers:
-            worker.join(timeout=5.0)
+            worker.join(timeout=1.0)
+
+        self.commit_shutdown_event.set()
+        self.commit_thread.join(timeout=1.0)
 
 
 class SimpleQueueProcessingStrategy(ProcessingStrategy[KafkaPayload], Generic[T]):
@@ -260,37 +382,18 @@ def __init__(
         decoder: Callable[[KafkaPayload | FilteredPayload], T | None],
         grouping_fn: Callable[[T], str],
         commit_function: Callable[[dict[Partition, int]], None],
+        partitions: set[Partition],
     ) -> None:
         self.queue_pool = queue_pool
         self.decoder = decoder
         self.grouping_fn = grouping_fn
-        self.commit_function = commit_function
         self.shutdown_event = threading.Event()
-
-        self.commit_thread = threading.Thread(target=self._commit_loop, daemon=True)
-        self.commit_thread.start()
-
-    def _commit_loop(self) -> None:
-        while not self.shutdown_event.is_set():
-            try:
-                self.shutdown_event.wait(1.0)
-
-                committable = self.queue_pool.offset_tracker.get_committable_offsets()
-
-                if committable:
-                    metrics.incr(
-                        "remote_subscriptions.queue_pool.offsets_committed",
-                        len(committable),
-                        tags={"identifier": self.queue_pool.identifier},
-                    )
-
-                    self.commit_function(committable)
-                    for partition, offset in committable.items():
-                        self.queue_pool.offset_tracker.mark_committed(partition, offset)
-            except Exception:
-                logger.exception("Error in commit loop")
+        self.queue_pool.update_assignments(partitions, commit_function)
 
     def submit(self, message: Message[KafkaPayload | FilteredPayload]) -> None:
+        if self.shutdown_event.is_set():
+            raise MessageRejected("Strategy is shutting down")
+
         try:
             result = self.decoder(message.payload)
 
@@ -299,8 +402,11 @@ def submit(self, message: Message[KafkaPayload | FilteredPayload]) -> None:
             offset = message.value.offset
 
             if result is None:
-                self.queue_pool.offset_tracker.add_offset(partition, offset)
-                self.queue_pool.offset_tracker.complete_offset(partition, offset)
+                try:
+                    self.queue_pool.offset_tracker.add_offset(partition, offset)
+                    self.queue_pool.offset_tracker.complete_offset(partition, offset)
+                except UnassignedPartitionError:
+                    pass
                 return
 
             group_key = self.grouping_fn(result)
@@ -334,12 +440,10 @@ def poll(self) -> None:
 
     def close(self) -> None:
         self.shutdown_event.set()
-        self.commit_thread.join(timeout=5.0)
-        self.queue_pool.shutdown()
 
     def terminate(self) -> None:
         self.shutdown_event.set()
-        self.queue_pool.shutdown()
+        self.queue_pool.flush(timeout=0)
 
     def join(self, timeout: float | None = None) -> None:
-        self.close()
+        self.queue_pool.flush(timeout=timeout or 0)

src/sentry/remote_subscriptions/consumers/result_consumer.py

@@ -112,6 +112,7 @@ def __init__(
         num_processes: int | None = None,
         input_block_size: int | None = None,
         output_block_size: int | None = None,
+        commit_interval: float | None = None,
     ) -> None:
         self.mode = mode
         metric_tags = {"identifier": self.identifier, "mode": self.mode}
@@ -133,7 +134,8 @@ def __init__(
             self.queue_pool = FixedQueuePool(
                 result_processor=self.result_processor,
                 identifier=self.identifier,
-                num_queues=max_workers or 20,  # Number of parallel queues
+                num_queues=max_workers or 20,
+                commit_interval=commit_interval or 1.0,
             )
 
         metrics.incr(
@@ -207,7 +209,7 @@ def create_with_partitions(
         if self.parallel:
             return self.create_multiprocess_worker(commit)
         if self.thread_queue_parallel:
-            return self.create_thread_queue_parallel_worker(commit)
+            return self.create_thread_queue_parallel_worker(commit, partitions)
         else:
             return self.create_serial_worker(commit)
 
@@ -242,7 +244,7 @@ def create_thread_parallel_worker(self, commit: Commit) -> ProcessingStrategy[Ka
         )
 
     def create_thread_queue_parallel_worker(
-        self, commit: Commit
+        self, commit: Commit, partitions: Mapping[Partition, int]
     ) -> ProcessingStrategy[KafkaPayload]:
         assert self.queue_pool is not None
 
@@ -256,6 +258,7 @@ def commit_offsets(offsets: dict[Partition, int]):
             decoder=partial(self.decode_payload, self.topic_for_codec),
             grouping_fn=self.build_payload_grouping_key,
             commit_function=commit_offsets,
+            partitions=set(partitions.keys()),
         )
 
     def partition_message_batch(self, message: Message[ValuesBatch[KafkaPayload]]) -> list[list[T]]: