@Override public Subscription subscription(Set<String> topics) { return new Subscription(new ArrayList<>(topics)); }
private void prepopulateCurrentAssignments(Map<String, Subscription> subscriptions, Map<String, List<TopicPartition>> currentAssignment) { for (Map.Entry<String, Subscription> subscriptionEntry : subscriptions.entrySet()) { ByteBuffer userData = subscriptionEntry.getValue().userData(); if (userData != null && userData.hasRemaining()) currentAssignment.put(subscriptionEntry.getKey(), deserializeTopicPartitionAssignment(userData)); } }
String consumer = entry.getKey(); consumer2AllPotentialPartitions.put(consumer, new ArrayList<TopicPartition>()); for (String topic: entry.getValue().topics()) { for (int i = 0; i < partitionsPerTopic.get(topic); ++i) { TopicPartition topicPartition = new TopicPartition(topic, i); } else if (!subscriptions.get(entry.getKey()).topics().contains(partition.topic())) {
final Subscription subscription = entry.getValue(); final SubscriptionInfo info = SubscriptionInfo.decode(subscription.userData()); final int usedVersion = info.version(); supportedVersions.add(info.latestSupportedVersion());
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; }
@Override public Subscription subscription(final Set<String> topics) { // Adds the following information to subscription // 1. Client UUID (a unique id assigned to an instance of KafkaStreams) // 2. Task ids of previously running tasks // 3. Task ids of valid local states on the client's state directory. final Set<TaskId> previousActiveTasks = taskManager.prevActiveTaskIds(); final Set<TaskId> standbyTasks = taskManager.cachedTasksIds(); standbyTasks.removeAll(previousActiveTasks); final SubscriptionInfo data = new SubscriptionInfo( usedSubscriptionMetadataVersion, taskManager.processId(), previousActiveTasks, standbyTasks, this.userEndPoint); taskManager.updateSubscriptionsFromMetadata(topics); return new Subscription(new ArrayList<>(topics), data.encode()); }
Subscription subscription = ConsumerProtocol.deserializeSubscription(subscriptionEntry.getValue()); subscriptions.put(subscriptionEntry.getKey(), subscription); allSubscribedTopics.addAll(subscription.topics());
assertTrue("Error: Partition " + partition + "is assigned to c" + i + ", but it is not subscribed to Topic t" + partition.topic() + "\nSubscriptions: " + subscriptions.toString() + "\nAssignments: " + assignments.toString(), subscriptions.get(consumer).topics().contains(partition.topic()));
partitionsPerTopic.put(topic, 3); Map<String, Subscription> subscriptions = new HashMap<>(); subscriptions.put(consumer1, new Subscription(topics(topic))); subscriptions.put(consumer2, new Subscription(topics(topic))); partitionsPerTopic.put(topic2, 3); subscriptions.put(consumer1, new Subscription(topics(topic, topic2), StickyAssignor.serializeTopicPartitionAssignment(assignment.get(consumer1)))); subscriptions.put(consumer2, new Subscription(topics(topic, topic2), StickyAssignor.serializeTopicPartitionAssignment(assignment.get(consumer2)))); assignment = assignor.assign(partitionsPerTopic, subscriptions); new Subscription(topics(topic2), StickyAssignor.serializeTopicPartitionAssignment(assignment.get(consumer1)))); subscriptions.put(consumer2, new Subscription(topics(topic2), StickyAssignor.serializeTopicPartitionAssignment(assignment.get(consumer2)))); assignment = assignor.assign(partitionsPerTopic, subscriptions);
final ByteBuffer byteBuffer = ConsumerProtocol.serializeSubscription(new PartitionAssignor.Subscription(singletonList(topic)));
partitionsPerTopic.put("topic01", 3); Map<String, Subscription> subscriptions = new HashMap<>(); subscriptions.put("consumer01", new Subscription(topics("topic01"))); subscriptions.put("consumer02", new Subscription(topics("topic01"))); subscriptions.put("consumer03", new Subscription(topics("topic01"))); subscriptions.put("consumer04", new Subscription(topics("topic01"))); new Subscription(topics("topic01"), StickyAssignor.serializeTopicPartitionAssignment(assignment.get("consumer02")))); subscriptions.put("consumer03", new Subscription(topics("topic01"), StickyAssignor.serializeTopicPartitionAssignment(assignment.get("consumer03")))); subscriptions.put("consumer04", new Subscription(topics("topic01"), StickyAssignor.serializeTopicPartitionAssignment(assignment.get("consumer04"))));
for (int i = 0; i < numConsumers; ++i) { List<String> sub = Utils.sorted(getRandomSublist(topics)); subscriptions.put(getConsumerName(i, maxNumConsumers), new Subscription(sub)); String consumer = getConsumerName(i, maxNumConsumers); subscriptions.put(consumer, new Subscription(sub, StickyAssignor.serializeTopicPartitionAssignment(assignment.get(consumer))));
partitionsPerTopic.put(topic, 3); Map<String, Subscription> subscriptions = new HashMap<>(); subscriptions.put(consumer1, new Subscription(topics(topic))); new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment(assignment.get(consumer1)))); subscriptions.put(consumer2, new Subscription(topics(topic))); assignment = assignor.assign(partitionsPerTopic, subscriptions); assertEquals(partitions(tp(topic, 1), tp(topic, 2)), assignment.get(consumer1)); new Subscription(topics(topic), StickyAssignor.serializeTopicPartitionAssignment(assignment.get(consumer2)))); assignment = assignor.assign(partitionsPerTopic, subscriptions); assertTrue(assignment.get(consumer2).contains(tp(topic, 0)));
@Test public void testLargeAssignmentWithMultipleConsumersLeaving() { Random rand = new Random(); int topicCount = 40; int consumerCount = 200; Map<String, Integer> partitionsPerTopic = new HashMap<>(); for (int i = 0; i < topicCount; i++) partitionsPerTopic.put(getTopicName(i, topicCount), rand.nextInt(10) + 1); Map<String, Subscription> subscriptions = new HashMap<>(); for (int i = 0; i < consumerCount; i++) { List<String> topics = new ArrayList<String>(); for (int j = 0; j < rand.nextInt(20); j++) topics.add(getTopicName(rand.nextInt(topicCount), topicCount)); subscriptions.put(getConsumerName(i, consumerCount), new Subscription(topics)); } Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, subscriptions); verifyValidityAndBalance(subscriptions, assignment); for (int i = 1; i < consumerCount; i++) { String consumer = getConsumerName(i, consumerCount); subscriptions.put(consumer, new Subscription(subscriptions.get(consumer).topics(), StickyAssignor.serializeTopicPartitionAssignment(assignment.get(consumer)))); } for (int i = 0; i < 50; ++i) { String c = getConsumerName(rand.nextInt(consumerCount), consumerCount); subscriptions.remove(c); } assignment = assignor.assign(partitionsPerTopic, subscriptions); verifyValidityAndBalance(subscriptions, assignment); assertTrue(assignor.isSticky()); }
subscriptions.put("consumer1", new Subscription(topics("topic1", "topic2", "topic3", "topic4", "topic5"))); subscriptions.put("consumer2", new Subscription(topics("topic1", "topic3", "topic5"))); subscriptions.put("consumer3", new Subscription(topics("topic1", "topic3", "topic5"))); subscriptions.put("consumer4", new Subscription(topics("topic1", "topic2", "topic3", "topic4", "topic5")));
@Test public void testSameSubscriptions() { Map<String, Integer> partitionsPerTopic = new HashMap<>(); for (int i = 1; i < 15; i++) partitionsPerTopic.put(getTopicName(i, 15), i); Map<String, Subscription> subscriptions = new HashMap<>(); for (int i = 1; i < 9; i++) { List<String> topics = new ArrayList<String>(); for (int j = 1; j <= partitionsPerTopic.size(); j++) topics.add(getTopicName(j, 15)); subscriptions.put(getConsumerName(i, 9), new Subscription(topics)); } Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, subscriptions); verifyValidityAndBalance(subscriptions, assignment); for (int i = 1; i < 9; i++) { String consumer = getConsumerName(i, 9); subscriptions.put(consumer, new Subscription(subscriptions.get(consumer).topics(), StickyAssignor.serializeTopicPartitionAssignment(assignment.get(consumer)))); } subscriptions.remove(getConsumerName(5, 9)); assignment = assignor.assign(partitionsPerTopic, subscriptions); verifyValidityAndBalance(subscriptions, assignment); assertTrue(assignor.isSticky()); }
@Test public void testReassignmentAfterOneConsumerLeaves() { Map<String, Integer> partitionsPerTopic = new HashMap<>(); for (int i = 1; i < 20; i++) partitionsPerTopic.put(getTopicName(i, 20), i); Map<String, Subscription> subscriptions = new HashMap<>(); for (int i = 1; i < 20; i++) { List<String> topics = new ArrayList<String>(); for (int j = 1; j <= i; j++) topics.add(getTopicName(j, 20)); subscriptions.put(getConsumerName(i, 20), new Subscription(topics)); } Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, subscriptions); verifyValidityAndBalance(subscriptions, assignment); for (int i = 1; i < 20; i++) { String consumer = getConsumerName(i, 20); subscriptions.put(consumer, new Subscription(subscriptions.get(consumer).topics(), StickyAssignor.serializeTopicPartitionAssignment(assignment.get(consumer)))); } subscriptions.remove("consumer10"); assignment = assignor.assign(partitionsPerTopic, subscriptions); verifyValidityAndBalance(subscriptions, assignment); assertTrue(assignor.isSticky()); }
@Test public void testMultipleConsumersMixedTopicSubscriptions() { String topic1 = "topic1"; String topic2 = "topic2"; String consumer1 = "consumer1"; String consumer2 = "consumer2"; String consumer3 = "consumer3"; Map<String, Integer> partitionsPerTopic = new HashMap<>(); partitionsPerTopic.put(topic1, 3); partitionsPerTopic.put(topic2, 2); Map<String, Subscription> subscriptions = new HashMap<>(); subscriptions.put(consumer1, new Subscription(topics(topic1))); subscriptions.put(consumer2, new Subscription(topics(topic1, topic2))); subscriptions.put(consumer3, new Subscription(topics(topic1))); Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, subscriptions); assertEquals(partitions(tp(topic1, 0), tp(topic1, 2)), assignment.get(consumer1)); assertEquals(partitions(tp(topic2, 0), tp(topic2, 1)), assignment.get(consumer2)); assertEquals(partitions(tp(topic1, 1)), assignment.get(consumer3)); verifyValidityAndBalance(subscriptions, assignment); assertTrue(isFullyBalanced(assignment)); }
@Test public void deserializeNewSubscriptionVersion() { // verify that a new version which adds a field is still parseable short version = 100; Schema subscriptionSchemaV100 = new Schema( new Field(ConsumerProtocol.TOPICS_KEY_NAME, new ArrayOf(Type.STRING)), new Field(ConsumerProtocol.USER_DATA_KEY_NAME, Type.BYTES), new Field("foo", Type.STRING)); Struct subscriptionV100 = new Struct(subscriptionSchemaV100); subscriptionV100.set(ConsumerProtocol.TOPICS_KEY_NAME, new Object[]{"topic"}); subscriptionV100.set(ConsumerProtocol.USER_DATA_KEY_NAME, ByteBuffer.wrap(new byte[0])); subscriptionV100.set("foo", "bar"); Struct headerV100 = new Struct(ConsumerProtocol.CONSUMER_PROTOCOL_HEADER_SCHEMA); headerV100.set(ConsumerProtocol.VERSION_KEY_NAME, version); ByteBuffer buffer = ByteBuffer.allocate(subscriptionV100.sizeOf() + headerV100.sizeOf()); headerV100.writeTo(buffer); subscriptionV100.writeTo(buffer); buffer.flip(); Subscription subscription = ConsumerProtocol.deserializeSubscription(buffer); assertEquals(Arrays.asList("topic"), subscription.topics()); }
@Test public void testMultipleConsumersMixedTopics() { String topic1 = "topic1"; String topic2 = "topic2"; String consumer1 = "consumer1"; String consumer2 = "consumer2"; String consumer3 = "consumer3"; Map<String, Integer> partitionsPerTopic = new HashMap<>(); partitionsPerTopic.put(topic1, 3); partitionsPerTopic.put(topic2, 2); Map<String, Subscription> consumers = new HashMap<>(); consumers.put(consumer1, new Subscription(topics(topic1))); consumers.put(consumer2, new Subscription(topics(topic1, topic2))); consumers.put(consumer3, new Subscription(topics(topic1))); Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, consumers); assertAssignment(partitions(tp(topic1, 0)), assignment.get(consumer1)); assertAssignment(partitions(tp(topic1, 1), tp(topic2, 0), tp(topic2, 1)), assignment.get(consumer2)); assertAssignment(partitions(tp(topic1, 2)), assignment.get(consumer3)); }