Add shutdown_strategy_before_consumer flag (#467)

* Add shutdown_strategy_before_consumer flag

Full stacktrace we are encountering:

      File "sentry/utils/kafka.py", line 52, in run_processor_with_signals
        processor.run()
      File "arroyo/processing/processor.py", line 335, in run
        self._run_once()
      File "arroyo/processing/processor.py", line 418, in _run_once
        self.__processing_strategy.poll()
      File "arroyo/processing/strategies/healthcheck.py", line 32, in poll
        self.__next_step.poll()
      File "arroyo/processing/strategies/guard.py", line 101, in poll
        self.__inner_strategy.poll()
      File "arroyo/processing/strategies/run_task.py", line 55, in poll
        self.__next_step.poll()
      File "arroyo/processing/strategies/guard.py", line 37, in poll
        self.__next_step.poll()
      File "arroyo/processing/strategies/batching.py", line 82, in poll
        self.__reduce_step.poll()
      File "arroyo/processing/strategies/reduce.py", line 110, in poll
        self.__buffer_step.poll()
      File "arroyo/processing/strategies/buffer.py", line 165, in poll
        self.__flush(force=False)
      File "arroyo/processing/strategies/buffer.py", line 125, in __flush
        self.__next_step.submit(buffer_msg)
      File "arroyo/processing/strategies/guard.py", line 82, in submit
        self.__inner_strategy.submit(message)
      File "arroyo/processing/strategies/run_task.py", line 52, in submit
        self.__next_step.submit(Message(value))
      File "arroyo/processing/strategies/guard.py", line 34, in submit
        self.__next_step.submit(message)
      File "sentry/spans/consumers/process/flusher.py", line 316, in submit
        self.next_step.submit(message)
      File "arroyo/processing/strategies/commit.py", line 34, in submit
        self.__commit(message.committable)
      File "arroyo/processing/processor.py", line 321, in __commit
        self.__consumer.commit_offsets()
      File "arroyo/backends/kafka/consumer.py", line 624, in commit_offsets
        return self.__commit_retry_policy.call(self.__commit)
      File "arroyo/utils/retries.py", line 88, in call
        return callable()
      File "arroyo/backends/kafka/consumer.py", line 582, in __commit
        result = self.__consumer.commit(

    KafkaException: KafkaError{code=UNKNOWN_MEMBER_ID,val=25,str="Commit failed: Broker: Unknown member"}

* dedupe code

Author: untitaker

arroyo/processing/processor.py

@@ -139,6 +139,7 @@ def __init__(
         commit_policy: CommitPolicy = ONCE_PER_SECOND,
         dlq_policy: Optional[DlqPolicy[TStrategyPayload]] = None,
         join_timeout: Optional[float] = None,
+        shutdown_strategy_before_consumer: bool = False,
     ) -> None:
         self.__consumer = consumer
         self.__processor_factory = processor_factory
@@ -158,6 +159,7 @@ def __init__(
         self.__commit_policy_state = commit_policy.get_state_machine()
         self.__join_timeout = join_timeout
         self.__shutdown_requested = False
+        self.__shutdown_strategy_before_consumer = shutdown_strategy_before_consumer
 
         # Buffers messages for DLQ. Messages are added when they are submitted for processing and
         # removed once the commit callback is fired as they are guaranteed to be valid at that point.
@@ -170,49 +172,7 @@ def __init__(
         )
 
         def _close_strategy() -> None:
-            start_close = time.time()
-
-            if self.__processing_strategy is None:
-                # Partitions are revoked when the consumer is shutting down, at
-                # which point we already have closed the consumer.
-                return
-
-            logger.info("Closing %r...", self.__processing_strategy)
-            self.__processing_strategy.close()
-
-            logger.info("Waiting for %r to exit...", self.__processing_strategy)
-
-            while True:
-                start_join = time.time()
-
-                try:
-                    self.__processing_strategy.join(self.__join_timeout)
-                    self.__metrics_buffer.incr_timing(
-                        "arroyo.consumer.join.time", time.time() - start_join
-                    )
-                    break
-                except InvalidMessage as e:
-                    self.__metrics_buffer.incr_timing(
-                        "arroyo.consumer.join.time", time.time() - start_join
-                    )
-                    self._handle_invalid_message(e)
-
-            logger.info(
-                "%r exited successfully, releasing assignment.",
-                self.__processing_strategy,
-            )
-            self.__processing_strategy = None
-            self.__message = None  # avoid leaking buffered messages across assignments
-            self.__is_paused = False
-            self._clear_backpressure()
-
-            value = time.time() - start_close
-
-            self.__metrics_buffer.metrics.timing(
-                "arroyo.consumer.run.close_strategy", value
-            )
-
-            self.__metrics_buffer.incr_timing("arroyo.consumer.shutdown.time", value)
+            self._close_processing_strategy()
 
         def _create_strategy(partitions: Mapping[Partition, int]) -> None:
             start_create = time.time()
@@ -302,6 +262,47 @@ def on_partitions_revoked(partitions: Sequence[Partition]) -> None:
             [topic], on_assign=on_partitions_assigned, on_revoke=on_partitions_revoked
         )
 
+    def _close_processing_strategy(self) -> None:
+        """Close the processing strategy and wait for it to exit."""
+        start_close = time.time()
+
+        if self.__processing_strategy is None:
+            # Partitions are revoked when the consumer is shutting down, at
+            # which point we already have closed the consumer.
+            return
+
+        logger.info("Closing %r...", self.__processing_strategy)
+        self.__processing_strategy.close()
+
+        logger.info("Waiting for %r to exit...", self.__processing_strategy)
+
+        while True:
+            start_join = time.time()
+
+            try:
+                self.__processing_strategy.join(self.__join_timeout)
+                self.__metrics_buffer.incr_timing(
+                    "arroyo.consumer.join.time", time.time() - start_join
+                )
+                break
+            except InvalidMessage as e:
+                self.__metrics_buffer.incr_timing(
+                    "arroyo.consumer.join.time", time.time() - start_join
+                )
+                self._handle_invalid_message(e)
+
+        logger.info("%r exited successfully", self.__processing_strategy)
+        self.__processing_strategy = None
+        self.__message = None
+        self.__is_paused = False
+        self._clear_backpressure()
+
+        value = time.time() - start_close
+        self.__metrics_buffer.metrics.timing(
+            "arroyo.consumer.run.close_strategy", value
+        )
+        self.__metrics_buffer.incr_timing("arroyo.consumer.shutdown.time", value)
+
     def __commit(self, offsets: Mapping[Partition, int], force: bool = False) -> None:
         """
         If force is passed, commit immediately and do not throttle. This should
@@ -518,6 +519,13 @@ def signal_shutdown(self) -> None:
         self.__shutdown_requested = True
 
     def _shutdown(self) -> None:
+        # If shutdown_strategy_before_consumer is set, work around an issue
+        # where rdkafka would revoke our partition, but then also immediately
+        # revoke our member ID as well, causing join() of the CommitStrategy
+        # (that is running in the partition revocation callback) to crash.
+        if self.__shutdown_strategy_before_consumer:
+            self._close_processing_strategy()
+
         # close the consumer
         logger.info("Stopping consumer")
         self.__metrics_buffer.flush()