KafkaSinkCluster: route ConsumerGroupDescribe request

shotover-proxy/tests/kafka_int_tests/mod.rs

@@ -629,6 +629,16 @@ async fn cluster_2_racks_multi_shotover_kafka_3_9(#[case] driver: KafkaDriver) {
     let connection_builder = KafkaConnectionBuilder::new(driver, "127.0.0.1:9192");
     test_cases::cluster_test_suite(&connection_builder).await;
 
+    #[allow(irrefutable_let_patterns)]
+    if let KafkaDriver::Java = driver {
+        // new consumer group protocol is only on java driver
+        test_cases::produce_consume_partitions_new_consumer_group_protocol(
+            &connection_builder,
+            "partitions3_new_consumer_group_protocol",
+        )
+        .await;
+    }
+
     for shotover in shotovers {
         tokio::time::timeout(
             Duration::from_secs(10),

shotover-proxy/tests/kafka_int_tests/test_cases.rs

@@ -3,11 +3,12 @@ use std::{collections::HashMap, time::Duration};
 use test_helpers::{
     connection::kafka::{
         Acl, AclOperation, AclPermissionType, AlterConfig, ConfigEntry, ConsumerConfig,
-        ConsumerGroupDescription, DescribeReplicaLogDirInfo, ExpectedResponse, IsolationLevel,
-        KafkaAdmin, KafkaConnectionBuilder, KafkaConsumer, KafkaDriver, KafkaProducer,
-        ListOffsetsResultInfo, NewPartition, NewPartitionReassignment, NewTopic, OffsetAndMetadata,
-        OffsetSpec, Record, RecordsToDelete, ResourcePatternType, ResourceSpecifier, ResourceType,
-        TopicPartition, TopicPartitionReplica, TransactionDescription,
+        ConsumerGroupDescription, ConsumerProtocol, DescribeReplicaLogDirInfo, ExpectedResponse,
+        IsolationLevel, KafkaAdmin, KafkaConnectionBuilder, KafkaConsumer, KafkaDriver,
+        KafkaProducer, ListOffsetsResultInfo, NewPartition, NewPartitionReassignment, NewTopic,
+        OffsetAndMetadata, OffsetSpec, Record, RecordsToDelete, ResourcePatternType,
+        ResourceSpecifier, ResourceType, TopicPartition, TopicPartitionReplica,
+        TransactionDescription,
     },
     docker_compose::DockerCompose,
 };
@@ -47,6 +48,11 @@ async fn admin_setup(connection_builder: &KafkaConnectionBuilder) {
                 num_partitions: 3,
                 replication_factor: 1,
             },
+            NewTopic {
+                name: "partitions3_new_consumer_group_protocol",
+                num_partitions: 3,
+                replication_factor: 1,
+            },
             NewTopic {
                 name: "acks0",
                 num_partitions: 1,
@@ -1091,6 +1097,64 @@ pub async fn produce_consume_partitions3(
     }
 }
 
+/// The new consumer protocol must be specifically enabled in the broker config.
+/// We only do this for the kafka 3.9 docker-compose.yaml so this test case is
+/// manually called for that test and not included in the standard test suite.
+pub async fn produce_consume_partitions_new_consumer_group_protocol(
+    connection_builder: &KafkaConnectionBuilder,
+    topic_name: &str,
+) {
+    let producer = connection_builder.connect_producer("1", 0).await;
+    let mut consumer = connection_builder
+        .connect_consumer(
+            ConsumerConfig::consume_from_topics(vec![topic_name.to_owned()])
+                .with_group("some_group")
+                .with_protocol(ConsumerProtocol::Consumer),
+        )
+        .await;
+
+    for _ in 0..5 {
+        producer
+            .assert_produce(
+                Record {
+                    payload: "Message1",
+                    topic_name,
+                    key: Some("Key".into()),
+                },
+                // We cant predict the offsets since that will depend on which partition the keyless record ends up in
+                None,
+            )
+            .await;
+        producer
+            .assert_produce(
+                Record {
+                    payload: "Message2",
+                    topic_name,
+                    key: None,
+                },
+                None,
+            )
+            .await;
+
+        consumer
+            .assert_consume_in_any_order(vec![
+                ExpectedResponse {
+                    message: "Message1".to_owned(),
+                    key: Some("Key".to_owned()),
+                    topic_name: topic_name.to_owned(),
+                    offset: None,
+                },
+                ExpectedResponse {
+                    message: "Message2".to_owned(),
+                    key: None,
+                    topic_name: topic_name.to_owned(),
+                    offset: None,
+                },
+            ])
+            .await;
+    }
+}
+
 pub async fn produce_consume_multi_topic_consumer(connection_builder: &KafkaConnectionBuilder) {
     let producer = connection_builder.connect_producer("1", 0).await;
     let mut consumer = connection_builder

shotover-proxy/tests/test-configs/kafka/cluster-2-racks/docker-compose-kafka-3.9.yaml

@@ -36,6 +36,8 @@ services:
       #
       # However for an integration test workload we are constantly spinning up single consumer groups, so the default value makes the tests take twice as long to run.
       KAFKA_CFG_GROUP_INITIAL_REBALANCE_DELAY_MS: "0"
+
+      KAFKA_CFG_GROUP_COORDINATOR_REBALANCE_PROTOCOLS: "consumer, classic"
     volumes: &volumes
       - type: tmpfs
         target: /bitnami/kafka

shotover/src/transforms/kafka/sink_cluster/mod.rs

@@ -32,7 +32,8 @@ use kafka_protocol::messages::produce_response::{
 };
 use kafka_protocol::messages::{
     AddOffsetsToTxnRequest, AddPartitionsToTxnRequest, AddPartitionsToTxnResponse, ApiKey,
-    BrokerId, DeleteGroupsRequest, DeleteGroupsResponse, DeleteRecordsRequest,
+    BrokerId, ConsumerGroupDescribeRequest, ConsumerGroupDescribeResponse,
+    ConsumerGroupHeartbeatRequest, DeleteGroupsRequest, DeleteGroupsResponse, DeleteRecordsRequest,
     DeleteRecordsResponse, DescribeClusterResponse, DescribeGroupsRequest, DescribeGroupsResponse,
     DescribeLogDirsResponse, DescribeProducersRequest, DescribeProducersResponse,
     DescribeTransactionsRequest, DescribeTransactionsResponse, EndTxnRequest, FetchRequest,
@@ -57,8 +58,8 @@ use scram_over_mtls::{
 use serde::{Deserialize, Serialize};
 use shotover_node::{ShotoverNode, ShotoverNodeConfig};
 use split::{
-    AddPartitionsToTxnRequestSplitAndRouter, DeleteGroupsSplitAndRouter,
-    DeleteRecordsRequestSplitAndRouter, DescribeGroupsSplitAndRouter,
+    AddPartitionsToTxnRequestSplitAndRouter, ConsumerGroupDescribeSplitAndRouter,
+    DeleteGroupsSplitAndRouter, DeleteRecordsRequestSplitAndRouter, DescribeGroupsSplitAndRouter,
     DescribeLogDirsSplitAndRouter, DescribeProducersRequestSplitAndRouter,
     DescribeTransactionsSplitAndRouter, ListGroupsSplitAndRouter, ListOffsetsRequestSplitAndRouter,
     ListTransactionsSplitAndRouter, OffsetFetchSplitAndRouter,
@@ -736,6 +737,10 @@ impl KafkaSinkCluster {
                 Some(Frame::Kafka(KafkaFrame::Request {
                     body:
                         RequestBody::Heartbeat(HeartbeatRequest { group_id, .. })
+                        | RequestBody::ConsumerGroupHeartbeat(ConsumerGroupHeartbeatRequest {
+                            group_id,
+                            ..
+                        })
                         | RequestBody::SyncGroup(SyncGroupRequest { group_id, .. })
                         | RequestBody::JoinGroup(JoinGroupRequest { group_id, .. })
                         | RequestBody::LeaveGroup(LeaveGroupRequest { group_id, .. })
@@ -744,6 +749,14 @@ impl KafkaSinkCluster {
                 })) => {
                     self.store_group(&mut groups, group_id.clone());
                 }
+                Some(Frame::Kafka(KafkaFrame::Request {
+                    body: RequestBody::ConsumerGroupDescribe(describe),
+                    ..
+                })) => {
+                    for group_id in &describe.group_ids {
+                        self.store_group(&mut groups, group_id.clone());
+                    }
+                }
                 Some(Frame::Kafka(KafkaFrame::Request {
                     body: RequestBody::DeleteGroups(delete_groups),
                     ..
@@ -955,6 +968,13 @@ impl KafkaSinkCluster {
                     let group_id = heartbeat.group_id.clone();
                     self.route_to_group_coordinator(request, group_id);
                 }
+                Some(Frame::Kafka(KafkaFrame::Request {
+                    body: RequestBody::ConsumerGroupHeartbeat(heartbeat),
+                    ..
+                })) => {
+                    let group_id = heartbeat.group_id.clone();
+                    self.route_to_group_coordinator(request, group_id);
+                }
                 Some(Frame::Kafka(KafkaFrame::Request {
                     body: RequestBody::SyncGroup(sync_group),
                     ..
@@ -1006,6 +1026,12 @@ impl KafkaSinkCluster {
                 })) => {
                     self.split_and_route_request::<DescribeGroupsSplitAndRouter>(request)?;
                 }
+                Some(Frame::Kafka(KafkaFrame::Request {
+                    body: RequestBody::ConsumerGroupDescribe(_),
+                    ..
+                })) => {
+                    self.split_and_route_request::<ConsumerGroupDescribeSplitAndRouter>(request)?;
+                }
                 Some(Frame::Kafka(KafkaFrame::Request {
                     body: RequestBody::OffsetDelete(offset_delete),
                     ..
@@ -1728,6 +1754,29 @@ The connection to the client has been closed."
         result
     }
 
+    /// This method removes all groups from the ConsumerGroupDescribe request and returns them split up by their destination.
+    /// If any groups are unroutable they will have their BrokerId set to -1
+    fn split_consumer_group_describe_request_by_destination(
+        &mut self,
+        body: &mut ConsumerGroupDescribeRequest,
+    ) -> HashMap<BrokerId, Vec<GroupId>> {
+        let mut result: HashMap<BrokerId, Vec<GroupId>> = Default::default();
+
+        for group in body.group_ids.drain(..) {
+            if let Some(destination) = self.group_to_coordinator_broker.get(&group) {
+                let dest_groups = result.entry(*destination).or_default();
+                dest_groups.push(group);
+            } else {
+                tracing::warn!("no known coordinator for group {group:?}, routing request to a random broker so that a NOT_COORDINATOR or similar error is returned to the client");
+                let destination = BrokerId(-1);
+                let dest_groups = result.entry(destination).or_default();
+                dest_groups.push(group);
+            }
+        }
+
+        result
+    }
+
     /// This method removes all topics from the list offsets request and returns them split up by their destination
     /// If any topics are unroutable they will have their BrokerId set to -1
     fn split_offset_for_leader_epoch_request_by_destination(
@@ -2359,6 +2408,10 @@ The connection to the client has been closed."
                 body: ResponseBody::DescribeGroups(base),
                 ..
             })) => Self::combine_describe_groups(base, drain)?,
+            Some(Frame::Kafka(KafkaFrame::Response {
+                body: ResponseBody::ConsumerGroupDescribe(base),
+                ..
+            })) => Self::combine_consumer_group_describe(base, drain)?,
             Some(Frame::Kafka(KafkaFrame::Response {
                 body: ResponseBody::DescribeLogDirs(base_body),
                 ..
@@ -2799,6 +2852,23 @@ The connection to the client has been closed."
         Ok(())
     }
 
+    fn combine_consumer_group_describe(
+        base: &mut ConsumerGroupDescribeResponse,
+        drain: impl Iterator<Item = ResponseToBeCombined>,
+    ) -> Result<()> {
+        for mut next in drain {
+            if let Some(Frame::Kafka(KafkaFrame::Response {
+                body: ResponseBody::ConsumerGroupDescribe(next),
+                ..
+            })) = next.response.frame()
+            {
+                base.groups.extend(std::mem::take(&mut next.groups));
+            }
+        }
+
+        Ok(())
+    }
+
     fn combine_add_partitions_to_txn(
         base_add_partitions_to_txn: &mut AddPartitionsToTxnResponse,
         drain: impl Iterator<Item = ResponseToBeCombined>,
@@ -2999,6 +3069,10 @@ The connection to the client has been closed."
                 body: ResponseBody::Heartbeat(heartbeat),
                 ..
             })) => self.handle_group_coordinator_routing_error(&request_ty, heartbeat.error_code),
+            Some(Frame::Kafka(KafkaFrame::Response {
+                body: ResponseBody::ConsumerGroupHeartbeat(heartbeat),
+                ..
+            })) => self.handle_group_coordinator_routing_error(&request_ty, heartbeat.error_code),
             Some(Frame::Kafka(KafkaFrame::Response {
                 body: ResponseBody::SyncGroup(sync_group),
                 ..
@@ -3173,9 +3247,6 @@ The connection to the client has been closed."
                     // We can consider supporting it when kafka itself starts to support it but we will need to be very
                     // careful to correctly implement the pagination/cursor logic.
                     ApiKey::DescribeTopicPartitionsKey as i16,
-                    // This message type is part of the new consumer group API, we should implement support for it in the future.
-                    // I've disabled it for now to keep the scope down for kafka 3.9 support.
-                    ApiKey::ConsumerGroupDescribeKey as i16,
                 ];
                 api_versions
                     .api_keys

shotover/src/transforms/kafka/sink_cluster/split.rs

@@ -11,8 +11,8 @@ use kafka_protocol::messages::{
     delete_records_request::DeleteRecordsTopic, describe_producers_request::TopicRequest,
     list_offsets_request::ListOffsetsTopic, offset_fetch_request::OffsetFetchRequestGroup,
     offset_for_leader_epoch_request::OffsetForLeaderTopic, produce_request::TopicProduceData,
-    AddPartitionsToTxnRequest, BrokerId, DeleteGroupsRequest, DeleteRecordsRequest,
-    DescribeGroupsRequest, DescribeLogDirsRequest, DescribeProducersRequest,
+    AddPartitionsToTxnRequest, BrokerId, ConsumerGroupDescribeRequest, DeleteGroupsRequest,
+    DeleteRecordsRequest, DescribeGroupsRequest, DescribeLogDirsRequest, DescribeProducersRequest,
     DescribeTransactionsRequest, GroupId, ListGroupsRequest, ListOffsetsRequest,
     ListTransactionsRequest, OffsetFetchRequest, OffsetForLeaderEpochRequest, ProduceRequest,
     TopicName, TransactionalId,
@@ -393,3 +393,31 @@ impl RequestSplitAndRouter for DescribeGroupsSplitAndRouter {
         request.groups = item;
     }
 }
+
+pub struct ConsumerGroupDescribeSplitAndRouter;
+
+impl RequestSplitAndRouter for ConsumerGroupDescribeSplitAndRouter {
+    type Request = ConsumerGroupDescribeRequest;
+    type SubRequests = Vec<GroupId>;
+
+    fn split_by_destination(
+        transform: &mut KafkaSinkCluster,
+        request: &mut Self::Request,
+    ) -> HashMap<BrokerId, Self::SubRequests> {
+        transform.split_consumer_group_describe_request_by_destination(request)
+    }
+
+    fn get_request_frame(request: &mut Message) -> &mut Self::Request {
+        match request.frame() {
+            Some(Frame::Kafka(KafkaFrame::Request {
+                body: RequestBody::ConsumerGroupDescribe(request),
+                ..
+            })) => request,
+            _ => unreachable!(),
+        }
+    }
+
+    fn reassemble(request: &mut Self::Request, item: Self::SubRequests) {
+        request.group_ids = item;
+    }
+}

test-helpers/src/connection/kafka/cpp.rs

@@ -3,8 +3,8 @@ use std::collections::{HashMap, HashSet};
 pub use rdkafka;
 
 use super::{
-    ConsumerConfig, ExpectedResponse, NewPartition, OffsetAndMetadata, ProduceResult, Record,
-    TopicPartition,
+    ConsumerConfig, ConsumerProtocol, ExpectedResponse, NewPartition, OffsetAndMetadata,
+    ProduceResult, Record, TopicPartition,
 };
 use anyhow::Result;
 use pretty_assertions::assert_eq;
@@ -96,6 +96,10 @@ impl KafkaConnectionBuilderCpp {
             .create()
             .unwrap();
 
+        if let ConsumerProtocol::Consumer = config.protocol {
+            panic!("New consumer protocol not support by rdkafka driver");
+        }
+
         let topic_names: Vec<&str> = config.topic_names.iter().map(|x| x.as_str()).collect();
         consumer.subscribe(&topic_names).unwrap();
         KafkaConsumerCpp { consumer }

test-helpers/src/connection/kafka/java.rs

@@ -173,6 +173,14 @@ impl KafkaConnectionBuilderJava {
             "org.apache.kafka.common.serialization.StringDeserializer".to_owned(),
         );
 
+        config.insert(
+            "group.protocol".to_owned(),
+            match consumer_config.protocol {
+                super::ConsumerProtocol::Classic => "CLASSIC".to_owned(),
+                super::ConsumerProtocol::Consumer => "CONSUMER".to_owned(),
+            },
+        );
+
         let consumer = self.jvm.construct(
             "org.apache.kafka.clients.consumer.KafkaConsumer",
             vec![properties(&self.jvm, &config)],