/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You under the Apache License, Version 2.0 * (the "License"); you may not use this file except in compliance with * the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ package org.apache.kafka.clients.consumer.internals; import org.apache.kafka.common.TopicPartition; import org.apache.kafka.common.protocol.types.ArrayOf; import org.apache.kafka.common.protocol.types.Field; import org.apache.kafka.common.protocol.types.Schema; import org.apache.kafka.common.protocol.types.SchemaException; import org.apache.kafka.common.protocol.types.Struct; import org.apache.kafka.common.protocol.types.Type; import org.apache.kafka.common.utils.CollectionUtils; import java.nio.ByteBuffer; import java.util.ArrayList; import java.util.List; import java.util.Map; /** * ConsumerProtocol contains the schemas for consumer subscriptions and assignments for use with * Kafka's generalized group management protocol. Below is the version 0 format: * * <pre> * Subscription => Version Topics * Version => Int16 * Topics => [String] * UserData => Bytes * * Assignment => Version TopicPartitions * Version => int16 * TopicPartitions => [Topic Partitions] * Topic => String * Partitions => [int32] * </pre> * * The current implementation assumes that future versions will not break compatibility. When * it encounters a newer version, it parses it using the current format. This basically means * that new versions cannot remove or reorder any of the existing fields. */ public class ConsumerProtocol { public static final String PROTOCOL_TYPE = "consumer"; public static final String VERSION_KEY_NAME = "version"; public static final String TOPICS_KEY_NAME = "topics"; public static final String TOPIC_KEY_NAME = "topic"; public static final String PARTITIONS_KEY_NAME = "partitions"; public static final String TOPIC_PARTITIONS_KEY_NAME = "topic_partitions"; public static final String USER_DATA_KEY_NAME = "user_data"; public static final short CONSUMER_PROTOCOL_V0 = 0; public static final Schema CONSUMER_PROTOCOL_HEADER_SCHEMA = new Schema( new Field(VERSION_KEY_NAME, Type.INT16)); private static final Struct CONSUMER_PROTOCOL_HEADER_V0 = new Struct(CONSUMER_PROTOCOL_HEADER_SCHEMA) .set(VERSION_KEY_NAME, CONSUMER_PROTOCOL_V0); public static final Schema SUBSCRIPTION_V0 = new Schema( new Field(TOPICS_KEY_NAME, new ArrayOf(Type.STRING)), new Field(USER_DATA_KEY_NAME, Type.NULLABLE_BYTES)); public static final Schema TOPIC_ASSIGNMENT_V0 = new Schema( new Field(TOPIC_KEY_NAME, Type.STRING), new Field(PARTITIONS_KEY_NAME, new ArrayOf(Type.INT32))); public static final Schema ASSIGNMENT_V0 = new Schema( new Field(TOPIC_PARTITIONS_KEY_NAME, new ArrayOf(TOPIC_ASSIGNMENT_V0)), new Field(USER_DATA_KEY_NAME, Type.NULLABLE_BYTES)); public static ByteBuffer serializeSubscription(PartitionAssignor.Subscription subscription) { Struct struct = new Struct(SUBSCRIPTION_V0); struct.set(USER_DATA_KEY_NAME, subscription.userData()); struct.set(TOPICS_KEY_NAME, subscription.topics().toArray()); ByteBuffer buffer = ByteBuffer.allocate(CONSUMER_PROTOCOL_HEADER_V0.sizeOf() + SUBSCRIPTION_V0.sizeOf(struct)); CONSUMER_PROTOCOL_HEADER_V0.writeTo(buffer); SUBSCRIPTION_V0.write(buffer, struct); buffer.flip(); return buffer; } public static PartitionAssignor.Subscription deserializeSubscription(ByteBuffer buffer) { Struct header = CONSUMER_PROTOCOL_HEADER_SCHEMA.read(buffer); Short version = header.getShort(VERSION_KEY_NAME); checkVersionCompatibility(version); Struct struct = SUBSCRIPTION_V0.read(buffer); ByteBuffer userData = struct.getBytes(USER_DATA_KEY_NAME); List<String> topics = new ArrayList<>(); for (Object topicObj : struct.getArray(TOPICS_KEY_NAME)) topics.add((String) topicObj); return new PartitionAssignor.Subscription(topics, userData); } public static PartitionAssignor.Assignment deserializeAssignment(ByteBuffer buffer) { Struct header = CONSUMER_PROTOCOL_HEADER_SCHEMA.read(buffer); Short version = header.getShort(VERSION_KEY_NAME); checkVersionCompatibility(version); Struct struct = ASSIGNMENT_V0.read(buffer); ByteBuffer userData = struct.getBytes(USER_DATA_KEY_NAME); List<TopicPartition> partitions = new ArrayList<>(); for (Object structObj : struct.getArray(TOPIC_PARTITIONS_KEY_NAME)) { Struct assignment = (Struct) structObj; String topic = assignment.getString(TOPIC_KEY_NAME); for (Object partitionObj : assignment.getArray(PARTITIONS_KEY_NAME)) { Integer partition = (Integer) partitionObj; partitions.add(new TopicPartition(topic, partition)); } } return new PartitionAssignor.Assignment(partitions, userData); } public static ByteBuffer serializeAssignment(PartitionAssignor.Assignment assignment) { Struct struct = new Struct(ASSIGNMENT_V0); struct.set(USER_DATA_KEY_NAME, assignment.userData()); List<Struct> topicAssignments = new ArrayList<>(); Map<String, List<Integer>> partitionsByTopic = CollectionUtils.groupDataByTopic(assignment.partitions()); for (Map.Entry<String, List<Integer>> topicEntry : partitionsByTopic.entrySet()) { Struct topicAssignment = new Struct(TOPIC_ASSIGNMENT_V0); topicAssignment.set(TOPIC_KEY_NAME, topicEntry.getKey()); topicAssignment.set(PARTITIONS_KEY_NAME, topicEntry.getValue().toArray()); topicAssignments.add(topicAssignment); } struct.set(TOPIC_PARTITIONS_KEY_NAME, topicAssignments.toArray()); ByteBuffer buffer = ByteBuffer.allocate(CONSUMER_PROTOCOL_HEADER_V0.sizeOf() + ASSIGNMENT_V0.sizeOf(struct)); CONSUMER_PROTOCOL_HEADER_V0.writeTo(buffer); ASSIGNMENT_V0.write(buffer, struct); buffer.flip(); return buffer; } private static void checkVersionCompatibility(short version) { // check for invalid versions if (version < CONSUMER_PROTOCOL_V0) throw new SchemaException("Unsupported subscription version: " + version); // otherwise, assume versions can be parsed as V0 } }