From f45d04033d7ed1949fb39111c9b7a4c6f70fcf98 Mon Sep 17 00:00:00 2001 From: Tony Le Date: Thu, 22 Oct 2026 17:00:56 -0400 Subject: [PATCH] feat(consumer): Add cooperative rebalance functionality --- src/config/mod.rs | 21 +++++++ src/kafka/consumer.rs | 107 +++++++++++++++++++++++++++++++--- src/kafka/consumer_backend.rs | 52 +++++++++++------ 3 files changed, 154 insertions(+), 26 deletions(-) diff --git a/src/config/mod.rs b/src/config/mod.rs index 3c5289f3..d1d0d9b4 100644 --- a/src/config/mod.rs +++ b/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" diff --git a/src/kafka/consumer.rs b/src/kafka/consumer.rs index 629a3502..7e02d082 100644 --- a/src/kafka/consumer.rs +++ b/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, + events: UnboundedReceiver<(Event, SyncSender<()>)>, + activation_store: Arc, + mut spawn_actors: impl FnMut(Arc, 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; } diff --git a/src/kafka/consumer_backend.rs b/src/kafka/consumer_backend.rs index fb1b664b..5f42a117 100644 --- a/src/kafka/consumer_backend.rs +++ b/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) { 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) { - 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, } +impl Assignment { + pub fn split(self) -> impl Iterator { + 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()