for (TopicPartition tp : assigned.partitions()) assignedTopics.add(tp.topic());
@Override public Map<String, Assignment> assign(Cluster metadata, Map<String, Subscription> subscriptions) { Set<String> allSubscribedTopics = new HashSet<>(); for (Map.Entry<String, Subscription> subscriptionEntry : subscriptions.entrySet()) allSubscribedTopics.addAll(subscriptionEntry.getValue().topics()); Map<String, Integer> partitionsPerTopic = new HashMap<>(); for (String topic : allSubscribedTopics) { Integer numPartitions = metadata.partitionCountForTopic(topic); if (numPartitions != null && numPartitions > 0) partitionsPerTopic.put(topic, numPartitions); else log.debug("Skipping assignment for topic {} since no metadata is available", topic); } Map<String, List<TopicPartition>> rawAssignments = assign(partitionsPerTopic, subscriptions); // this class maintains no user data, so just wrap the results Map<String, Assignment> assignments = new HashMap<>(); for (Map.Entry<String, List<TopicPartition>> assignmentEntry : rawAssignments.entrySet()) assignments.put(assignmentEntry.getKey(), new Assignment(assignmentEntry.getValue())); return assignments; }
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.groupPartitionsByTopic(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; }
final List<TopicPartition> partitions = new ArrayList<>(assignment.partitions()); Collections.sort(partitions, PARTITION_COMPARATOR); final AssignmentInfo info = AssignmentInfo.decode(assignment.userData()); if (info.errCode() != Error.NONE.code) {
assignment.put(consumer, new Assignment( activePartitions, new AssignmentInfo(minUserMetadataVersion, active, standby, partitionsByHostState, 0).encode()));
assignment.put(consumerId, new Assignment( assignedPartitions, new AssignmentInfo( assignment.put(consumerId, new Assignment( Collections.emptyList(), new AssignmentInfo().encode()
protected static List<List<LogPartition>> assignments(PartitionAssignor assignor, int threads, Map<String, Integer> streams) { final List<PartitionInfo> parts = new ArrayList<>(); streams.forEach((streamName, size) -> parts.addAll(getPartsFor(streamName, size))); Map<String, PartitionAssignor.Subscription> subscriptions = new HashMap<>(); List<String> streamNames = streams.keySet().stream().sorted().collect(Collectors.toList()); for (int i = 0; i < threads; i++) { subscriptions.put(String.valueOf(i), new PartitionAssignor.Subscription(streamNames)); } Cluster cluster = new Cluster("kafka-cluster", Collections.emptyList(), parts, Collections.emptySet(), Collections.emptySet()); Map<String, PartitionAssignor.Assignment> assignments = assignor.assign(cluster, subscriptions); List<List<LogPartition>> ret = new ArrayList<>(threads); for (int i = 0; i < threads; i++) { ret.add(assignments.get(String.valueOf(i)) .partitions() .stream() .map(part -> new LogPartition(part.topic(), part.partition())) .collect(Collectors.toList())); } return ret; }
Map<String, Assignment> errorAssignment(final Map<UUID, ClientMetadata> clientsMetadata, final String topic, final int errorCode) { log.error("{} is unknown yet during rebalance," + " please make sure they have been pre-created before starting the Streams application.", topic); final Map<String, Assignment> assignment = new HashMap<>(); for (final ClientMetadata clientMetadata : clientsMetadata.values()) { for (final String consumerId : clientMetadata.consumers) { assignment.put(consumerId, new Assignment( Collections.emptyList(), new AssignmentInfo(AssignmentInfo.LATEST_SUPPORTED_VERSION, Collections.emptyList(), Collections.emptyMap(), Collections.emptyMap(), errorCode).encode() )); } } return assignment; } /*
topicPartitions.add(2, myTopicPartition2); final ByteBuffer memberAssignment = ConsumerProtocol.serializeAssignment(new PartitionAssignor.Assignment(topicPartitions));
subscriptions.assignFromSubscribed(assignment.partitions());
@Test public void deserializeNewAssignmentVersion() { // verify that a new version which adds a field is still parseable short version = 100; Schema assignmentSchemaV100 = new Schema( new Field(ConsumerProtocol.TOPIC_PARTITIONS_KEY_NAME, new ArrayOf(ConsumerProtocol.TOPIC_ASSIGNMENT_V0)), new Field(ConsumerProtocol.USER_DATA_KEY_NAME, Type.BYTES), new Field("foo", Type.STRING)); Struct assignmentV100 = new Struct(assignmentSchemaV100); assignmentV100.set(ConsumerProtocol.TOPIC_PARTITIONS_KEY_NAME, new Object[]{new Struct(ConsumerProtocol.TOPIC_ASSIGNMENT_V0) .set(ConsumerProtocol.TOPIC_KEY_NAME, "foo") .set(ConsumerProtocol.PARTITIONS_KEY_NAME, new Object[]{1})}); assignmentV100.set(ConsumerProtocol.USER_DATA_KEY_NAME, ByteBuffer.wrap(new byte[0])); assignmentV100.set("foo", "bar"); Struct headerV100 = new Struct(ConsumerProtocol.CONSUMER_PROTOCOL_HEADER_SCHEMA); headerV100.set(ConsumerProtocol.VERSION_KEY_NAME, version); ByteBuffer buffer = ByteBuffer.allocate(assignmentV100.sizeOf() + headerV100.sizeOf()); headerV100.writeTo(buffer); assignmentV100.writeTo(buffer); buffer.flip(); PartitionAssignor.Assignment assignment = ConsumerProtocol.deserializeAssignment(buffer); assertEquals(toSet(Arrays.asList(new TopicPartition("foo", 1))), toSet(assignment.partitions())); }
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); }
private SyncGroupResponse syncGroupResponse(List<TopicPartition> partitions, Errors error) { ByteBuffer buf = ConsumerProtocol.serializeAssignment(new PartitionAssignor.Assignment(partitions)); return new SyncGroupResponse(error, buf); }
private SyncGroupResponse syncGroupResponse(List<TopicPartition> partitions, Errors error) { ByteBuffer buf = ConsumerProtocol.serializeAssignment(new PartitionAssignor.Assignment(partitions)); return new SyncGroupResponse(error, buf); }
@Override public void onAssignment(Assignment assignment) { memberAssignment = assignment.partitions(); }
@Test public void serializeDeserializeAssignment() { List<TopicPartition> partitions = Arrays.asList(new TopicPartition("foo", 0), new TopicPartition("bar", 2)); ByteBuffer buffer = ConsumerProtocol.serializeAssignment(new PartitionAssignor.Assignment(partitions)); PartitionAssignor.Assignment parsedAssignment = ConsumerProtocol.deserializeAssignment(buffer); assertEquals(toSet(partitions), toSet(parsedAssignment.partitions())); }
@Test public void deserializeNullAssignmentUserData() { List<TopicPartition> partitions = Arrays.asList(new TopicPartition("foo", 0), new TopicPartition("bar", 2)); ByteBuffer buffer = ConsumerProtocol.serializeAssignment(new PartitionAssignor.Assignment(partitions, null)); PartitionAssignor.Assignment parsedAssignment = ConsumerProtocol.deserializeAssignment(buffer); assertEquals(toSet(partitions), toSet(parsedAssignment.partitions())); assertNull(parsedAssignment.userData()); }