feat(consumer): Add cooperative rebalance functionality

src/config/mod.rs

@@ -142,6 +142,11 @@ pub struct Config {
     /// `auto_offset_reset`.
     pub kafka_auto_offset_reset: String,
 
+    /// Use the `cooperative-sticky` assignor instead of eager rebalancing. Members on the new
+    /// assignor wait until every member on the old one has left the group, so flip this with a
+    /// rolling deploy, in either direction.
+    pub kafka_cooperative_rebalance: bool,
+
     /// The number of ms for timeouts when publishing messages to kafka.
     pub kafka_send_timeout_ms: u64,
 
@@ -294,6 +299,7 @@ impl Default for Config {
             kafka_session_timeout_ms: 6000,
             kafka_auto_commit_interval_ms: 5000,
             kafka_auto_offset_reset: "latest".to_owned(),
+            kafka_cooperative_rebalance: false,
             kafka_send_timeout_ms: 500,
             store: StoreConfig::default(),
             runtime_config_path: None,
@@ -955,6 +961,9 @@ impl Config {
             )
             .set("auto.offset.reset", auto_offset_reset)
             .set("enable.auto.offset.store", "false");
+        if self.kafka_cooperative_rebalance {
+            config.set("partition.assignment.strategy", "cooperative-sticky");
+        }
 
         config
     }
@@ -1438,6 +1447,11 @@ mod tests {
         );
         assert_eq!(consumer_config.get("group.id").unwrap(), "taskworker");
         assert!(consumer_config.get("session.timeout.ms").is_some());
+        assert!(
+            consumer_config
+                .get("partition.assignment.strategy")
+                .is_none()
+        );
     }
 
     #[test]
@@ -2180,6 +2194,7 @@ kafka_deadletter_topic: profiles-dlq
 kafka_session_timeout_ms: 7000
 kafka_auto_commit_interval_ms: 2000
 kafka_auto_offset_reset: earliest
+kafka_cooperative_rebalance: true
 
 kafka_topics:
   profiles:
@@ -2205,6 +2220,12 @@ kafka_clusters:
 
             // No per-topic overrides, so the globals are used.
             assert_eq!(consumer_config.get("session.timeout.ms").unwrap(), "7000");
+            assert_eq!(
+                consumer_config
+                    .get("partition.assignment.strategy")
+                    .unwrap(),
+                "cooperative-sticky"
+            );
             assert_eq!(
                 consumer_config.get("auto.commit.interval.ms").unwrap(),
                 "2000"

src/kafka/consumer.rs

@@ -1,4 +1,4 @@
-use std::collections::{BTreeSet, HashMap};
+use std::collections::{BTreeMap, BTreeSet, HashMap};
 use std::fmt::Debug;
 use std::future::Future;
 use std::mem::take;
@@ -103,14 +103,27 @@ pub async fn start_consumer(
         "application" => "taskbroker",
     )
     .set(0);
-    handle_events(
-        consumer.clone(),
-        event_receiver,
-        activation_store,
-        spawn_actors,
-        topics_tag,
-    )
-    .await?;
+    let cooperative =
+        kafka_client_config.get("partition.assignment.strategy") == Some("cooperative-sticky");
+    if cooperative {
+        handle_events_cooperative(
+            consumer.clone(),
+            event_receiver,
+            activation_store,
+            spawn_actors,
+            topics_tag,
+        )
+        .await?;
+    } else {
+        handle_events(
+            consumer.clone(),
+            event_receiver,
+            activation_store,
+            spawn_actors,
+            topics_tag,
+        )
+        .await?;
+    }
     // Dropping the consumer blocks while it closes, and Arroyo also makes its final commit then,
     // so keep it off the async workers.
     task::spawn_blocking(move || drop(consumer)).await?;
@@ -363,6 +376,82 @@ pub async fn handle_events(
     Ok(())
 }
 
+#[instrument(skip_all)]
+pub async fn handle_events_cooperative(
+    consumer: Arc<AsyncKafkaConsumer>,
+    events: UnboundedReceiver<(Event, SyncSender<()>)>,
+    activation_store: Arc<dyn ActivationStore>,
+    mut spawn_actors: impl FnMut(Arc<AsyncKafkaConsumer>, Assignment) -> ActorHandles,
+    topics_tag: String,
+) -> Result<(), anyhow::Error> {
+    const CALLBACK_DURATION: Duration = Duration::from_secs(4);
+
+    let _guard = elegant_departure::get_shutdown_guard().shutdown_on_drop();
+    let mut events_stream = UnboundedReceiverStream::new(events);
+
+    let mut pipelines: BTreeMap<(String, i32), ActorHandles> = BTreeMap::new();
+
+    loop {
+        select! {
+            res = async {
+                if pipelines.is_empty() {
+                    future::pending().await
+                } else {
+                    future::select_all(pipelines.values_mut().map(|handles| Box::pin(handles.join_next()))).await.0
+                }
+            } => {
+                error!("Actor exited unexpectedly with {:?}, shutting down...", res);
+                drop(elegant_departure::shutdown());
+            }
+
+            event = events_stream.next() => {
+                let Some((event, _rendezvous_guard)) = event else {
+                    unreachable!("Unexpected end to event stream");
+                };
+                info!("Received event: {:?}", event);
+                let shutdown = matches!(event, Event::Shutdown);
+                let revoked = match event {
+                    Event::Assign(assignment) => {
+                        activation_store.assign_partitions(&mut assignment.partitions.iter().map(TopicPartition::from));
+                        for (partition, assignment) in assignment.split() {
+                            let previous = pipelines.insert(partition, spawn_actors(consumer.clone(), assignment));
+                            assert!(previous.is_none(), "Got assignment for a partition we're already consuming from");
+                        }
+                        vec![]
+                    }
+                    Event::Revoke(revoked) => {
+                        activation_store.revoke_partitions(&mut revoked.iter().map(TopicPartition::from));
+                        revoked
+                            .iter()
+                            .map(|partition| {
+                                pipelines
+                                    .remove(partition)
+                                    .expect("Got revocation for a partition we're not consuming from")
+                            })
+                            .collect()
+                    }
+                    Event::Shutdown => {
+                        activation_store.revoke_partitions(&mut pipelines.keys().map(TopicPartition::from));
+                        take(&mut pipelines).into_values().collect()
+                    }
+                };
+                future::join_all(revoked.into_iter().map(|handles| handles.shutdown(CALLBACK_DURATION))).await;
+                metrics::gauge!(
+                    "arroyo.consumer.current_partitions",
+                    "topic" => topics_tag.clone(),
+                    "application" => "taskbroker",
+                )
+                .set(pipelines.len() as f64);
+                if shutdown {
+                    break;
+                }
+            }
+        }
+    }
+    debug!("Shutdown complete");
+    Ok(())
+}
+
 pub trait KafkaMessage {
     fn into_message(self) -> Result<MessageBackend, Error>;
 }

src/kafka/consumer_backend.rs

@@ -5,6 +5,7 @@ use std::sync::mpsc::{SyncSender, sync_channel};
 use anyhow::{Context, Error};
 use futures::Stream;
 use rdkafka::error::KafkaError;
+use rdkafka::types::RDKafkaErrorCode;
 use rdkafka::{ClientConfig, TopicPartitionList};
 use sentry_arroyo::backends::kafka::async_consumer::{
     AsyncAssignmentCallbacks, AsyncKafkaConsumer, AsyncPartitionQueue,
@@ -122,36 +123,40 @@ impl KafkaContext {
     }
 }
 
+fn partition_key(partition: Partition) -> (String, i32) {
+    (
+        partition.topic.as_str().to_owned(),
+        i32::from(partition.index),
+    )
+}
+
 impl AsyncAssignmentCallbacks for KafkaContext {
     fn on_assign(&self, queues: Vec<AsyncPartitionQueue>) {
         let partitions = queues
             .iter()
-            .map(|queue| {
-                let partition = queue.partition();
-                (
-                    partition.topic.as_str().to_owned(),
-                    i32::from(partition.index),
-                )
-            })
+            .map(|queue| partition_key(queue.partition()))
             .collect();
         self.handle_rebalance(Event::Assign(Assignment { partitions, queues }));
     }
 
     fn on_revoke(&self, partitions: Vec<Partition>) {
-        let partitions = partitions
-            .iter()
-            .map(|partition| {
-                (
-                    partition.topic.as_str().to_owned(),
-                    i32::from(partition.index),
-                )
-            })
-            .collect();
+        let partitions = partitions.into_iter().map(partition_key).collect();
         self.handle_rebalance(Event::Revoke(partitions));
     }
 
     /// Arroyo has already logged the error. Shut down on unexpected consumer errors.
-    fn on_error(&self, _error: KafkaError) {
+    fn on_error(&self, error: KafkaError) {
+        if error.rdkafka_error_code() == Some(RDKafkaErrorCode::InconsistentGroupProtocol) {
+            // for switching to cooperative-sticky
+            warn!("Waiting for group members on the other rebalance protocol to leave");
+            metrics::counter!(
+                "arroyo.consumer.inconsistent_group_protocol",
+                "topic" => self.topics_tag.clone(),
+                "application" => "taskbroker",
+            )
+            .increment(1);
+            return;
+        }
         drop(elegant_departure::shutdown());
     }
 }
@@ -162,6 +167,19 @@ pub struct Assignment {
     pub queues: Vec<AsyncPartitionQueue>,
 }
 
+impl Assignment {
+    pub fn split(self) -> impl Iterator<Item = ((String, i32), Assignment)> {
+        self.queues.into_iter().map(|queue| {
+            let partition = partition_key(queue.partition());
+            let assignment = Assignment {
+                partitions: BTreeSet::from([partition.clone()]),
+                queues: vec![queue],
+            };
+            (partition, assignment)
+        })
+    }
+}
+
 impl Debug for Assignment {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         f.debug_tuple("Assignment").field(&self.partitions).finish()