dajac commented on code in PR #13639: URL: https://github.com/apache/kafka/pull/13639#discussion_r1200149897
########## group-coordinator/src/test/java/org/apache/kafka/coordinator/group/GroupMetadataManagerTest.java: ########## @@ -0,0 +1,2080 @@ +/* + * 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.coordinator.group; + +import org.apache.kafka.common.Uuid; +import org.apache.kafka.common.errors.FencedMemberEpochException; +import org.apache.kafka.common.errors.GroupIdNotFoundException; +import org.apache.kafka.common.errors.GroupMaxSizeReachedException; +import org.apache.kafka.common.errors.InvalidRequestException; +import org.apache.kafka.common.errors.NotCoordinatorException; +import org.apache.kafka.common.errors.UnknownMemberIdException; +import org.apache.kafka.common.errors.UnknownServerException; +import org.apache.kafka.common.errors.UnsupportedAssignorException; +import org.apache.kafka.common.message.ConsumerGroupHeartbeatRequestData; +import org.apache.kafka.common.message.ConsumerGroupHeartbeatResponseData; +import org.apache.kafka.common.metadata.PartitionRecord; +import org.apache.kafka.common.metadata.TopicRecord; +import org.apache.kafka.common.network.ClientInformation; +import org.apache.kafka.common.network.ListenerName; +import org.apache.kafka.common.protocol.ApiKeys; +import org.apache.kafka.common.protocol.ApiMessage; +import org.apache.kafka.common.requests.RequestContext; +import org.apache.kafka.common.requests.RequestHeader; +import org.apache.kafka.common.security.auth.KafkaPrincipal; +import org.apache.kafka.common.security.auth.SecurityProtocol; +import org.apache.kafka.common.utils.LogContext; +import org.apache.kafka.coordinator.group.assignor.AssignmentSpec; +import org.apache.kafka.coordinator.group.assignor.GroupAssignment; +import org.apache.kafka.coordinator.group.assignor.PartitionAssignor; +import org.apache.kafka.coordinator.group.assignor.PartitionAssignorException; +import org.apache.kafka.coordinator.group.consumer.Assignment; +import org.apache.kafka.coordinator.group.consumer.ConsumerGroup; +import org.apache.kafka.coordinator.group.consumer.ConsumerGroupMember; +import org.apache.kafka.coordinator.group.consumer.TopicMetadata; +import org.apache.kafka.coordinator.group.generated.ConsumerGroupCurrentMemberAssignmentKey; +import org.apache.kafka.coordinator.group.generated.ConsumerGroupCurrentMemberAssignmentValue; +import org.apache.kafka.coordinator.group.generated.ConsumerGroupMemberMetadataKey; +import org.apache.kafka.coordinator.group.generated.ConsumerGroupMemberMetadataValue; +import org.apache.kafka.coordinator.group.generated.ConsumerGroupMetadataKey; +import org.apache.kafka.coordinator.group.generated.ConsumerGroupMetadataValue; +import org.apache.kafka.coordinator.group.generated.ConsumerGroupPartitionMetadataKey; +import org.apache.kafka.coordinator.group.generated.ConsumerGroupPartitionMetadataValue; +import org.apache.kafka.coordinator.group.generated.ConsumerGroupTargetAssignmentMemberKey; +import org.apache.kafka.coordinator.group.generated.ConsumerGroupTargetAssignmentMemberValue; +import org.apache.kafka.coordinator.group.generated.ConsumerGroupTargetAssignmentMetadataKey; +import org.apache.kafka.coordinator.group.generated.ConsumerGroupTargetAssignmentMetadataValue; +import org.apache.kafka.image.TopicImage; +import org.apache.kafka.image.TopicsDelta; +import org.apache.kafka.image.TopicsImage; +import org.apache.kafka.server.common.ApiMessageAndVersion; +import org.apache.kafka.timeline.SnapshotRegistry; +import org.junit.jupiter.api.Test; + +import java.net.InetAddress; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Set; + +import static org.apache.kafka.coordinator.group.consumer.AssignmentTestUtil.mkAssignment; +import static org.apache.kafka.coordinator.group.consumer.AssignmentTestUtil.mkTopicAssignment; +import static org.junit.jupiter.api.AssertionFailureBuilder.assertionFailure; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +public class GroupMetadataManagerTest { + static class MockPartitionAssignor implements PartitionAssignor { + private final String name; + private AssignmentSpec lastSpecReceived = null; + private GroupAssignment prepareGroupAssignment = null; + + MockPartitionAssignor(String name) { + this.name = name; + } + + public AssignmentSpec lastSpecReceived() { + return lastSpecReceived; + } + + public void prepareGroupAssignment(GroupAssignment prepareGroupAssignment) { + this.prepareGroupAssignment = prepareGroupAssignment; + } + + @Override + public String name() { + return name; + } + + @Override + public GroupAssignment assign(AssignmentSpec assignmentSpec) throws PartitionAssignorException { + lastSpecReceived = assignmentSpec; + return prepareGroupAssignment; + } + } + + public static class TopicsImageBuilder { + private TopicsDelta delta = new TopicsDelta(TopicsImage.EMPTY); + + public TopicsImageBuilder addTopic( + Uuid topicId, + String topicName, + int numPartitions + ) { + delta.replay(new TopicRecord().setTopicId(topicId).setName(topicName)); + for (int i = 0; i < numPartitions; i++) { + delta.replay(new PartitionRecord() + .setTopicId(topicId) + .setPartitionId(i)); + } + return this; + } + + public TopicsImage build() { + return delta.apply(); + } + } + + static class ConsumerGroupBuilder { + private final String groupId; + private final int groupEpoch; + private int assignmentEpoch; + private final Map<String, ConsumerGroupMember> members = new HashMap<>(); + private final Map<String, Assignment> assignments = new HashMap<>(); + + public ConsumerGroupBuilder(String groupId, int groupEpoch) { + this.groupId = groupId; + this.groupEpoch = groupEpoch; + this.assignmentEpoch = 0; + } + + public ConsumerGroupBuilder withMember(ConsumerGroupMember member) { + this.members.put(member.memberId(), member); + return this; + } + + public ConsumerGroupBuilder withAssignment(String memberId, Map<Uuid, Set<Integer>> assignment) { + this.assignments.put(memberId, new Assignment(assignment)); + return this; + } + + public ConsumerGroupBuilder withAssignment(String memberId, Assignment assignment) { + this.assignments.put(memberId, assignment); + return this; + } + + public ConsumerGroupBuilder withAssignmentEpoch(int assignmentEpoch) { + this.assignmentEpoch = assignmentEpoch; + return this; + } + + public List<Record> build(TopicsImage topicsImage) { + List<Record> records = new ArrayList<>(); + + // Add subscription records for members. + members.forEach((memberId, member) -> { + records.add(RecordHelpers.newMemberSubscriptionRecord(groupId, member)); + }); + + // Add subscription metadata. + Map<String, TopicMetadata> subscriptionMetadata = new HashMap<>(); + members.forEach((memberId, member) -> { + member.subscribedTopicNames().forEach(topicName -> + subscriptionMetadata.computeIfAbsent(topicName, __ -> { + TopicImage topicImage = topicsImage.getTopic(topicName); + if (topicImage == null) { + return null; + } else { + return new TopicMetadata( + topicImage.id(), + topicImage.name(), + topicImage.partitions().size() + ); + } + }) + ); + }); + if (!subscriptionMetadata.isEmpty()) { + records.add(RecordHelpers.newGroupSubscriptionMetadataRecord(groupId, subscriptionMetadata)); + } + + // Add group epoch record. + records.add(RecordHelpers.newGroupEpochRecord(groupId, groupEpoch)); + + // Add target assignment records. + assignments.forEach((memberId, assignment) -> { + records.add(RecordHelpers.newTargetAssignmentRecord(groupId, memberId, assignment.partitions())); + }); + + // Add target assignment epoch. + records.add(RecordHelpers.newTargetAssignmentEpochRecord(groupId, assignmentEpoch)); + + // Add current assignment records for members. + members.forEach((memberId, member) -> { + records.add(RecordHelpers.newCurrentAssignmentRecord(groupId, member)); + }); + + return records; + } + } + + static class GroupMetadataManagerTestContext { + static class Builder { + private LogContext logContext; + private SnapshotRegistry snapshotRegistry; + private TopicsImage topicsImage; + private List<PartitionAssignor> assignors; + private List<ConsumerGroupBuilder> consumerGroupBuilders = new ArrayList<>(); + private int consumerGroupMaxSize = Integer.MAX_VALUE; + + public Builder withLogContext(LogContext logContext) { + this.logContext = logContext; + return this; + } + + public Builder withSnapshotRegistry(SnapshotRegistry snapshotRegistry) { + this.snapshotRegistry = snapshotRegistry; + return this; + } + + public Builder withTopicsImage(TopicsImage topicsImage) { + this.topicsImage = topicsImage; + return this; + } + + public Builder withAssignors(List<PartitionAssignor> assignors) { + this.assignors = assignors; + return this; + } + + public Builder withConsumerGroup(ConsumerGroupBuilder builder) { + this.consumerGroupBuilders.add(builder); + return this; + } + + public Builder withConsumerGroupMaxSize(int consumerGroupMaxSize) { + this.consumerGroupMaxSize = consumerGroupMaxSize; + return this; + } + + public GroupMetadataManagerTestContext build() { + if (logContext == null) logContext = new LogContext(); + if (snapshotRegistry == null) snapshotRegistry = new SnapshotRegistry(logContext); + if (topicsImage == null) topicsImage = TopicsImage.EMPTY; + if (assignors == null) assignors = Collections.emptyList(); + + GroupMetadataManagerTestContext context = new GroupMetadataManagerTestContext( + snapshotRegistry, + new GroupMetadataManager.Builder() + .withSnapshotRegistry(snapshotRegistry) + .withLogContext(logContext) + .withTopicsImage(topicsImage) + .withConsumerGroupHeartbeatInterval(5000) + .withConsumerGroupMaxSize(consumerGroupMaxSize) + .withAssignors(assignors) + .build() + ); + + consumerGroupBuilders.forEach(builder -> { + builder.build(topicsImage).forEach(context::replay); + }); + + context.commit(); + + return context; + } + } + + final SnapshotRegistry snapshotRegistry; + final GroupMetadataManager groupMetadataManager; + + long lastCommittedOffset = 0L; + long lastWrittenOffset = 0L; + + public GroupMetadataManagerTestContext( + SnapshotRegistry snapshotRegistry, + GroupMetadataManager groupMetadataManager + ) { + this.snapshotRegistry = snapshotRegistry; + this.groupMetadataManager = groupMetadataManager; + } + + public void commit() { + long lastCommittedOffset = this.lastCommittedOffset; + this.lastCommittedOffset = lastWrittenOffset; + snapshotRegistry.deleteSnapshotsUpTo(lastCommittedOffset); + } + + public void rollback() { + lastWrittenOffset = lastCommittedOffset; + snapshotRegistry.revertToSnapshot(lastCommittedOffset); + } + + public ConsumerGroup.ConsumerGroupState consumerGroupState( + String groupId + ) { + return groupMetadataManager + .getOrMaybeCreateConsumerGroup(groupId, false) + .state(); + } + + public ConsumerGroupMember.MemberState consumerGroupMemberState( + String groupId, + String memberId + ) { + return groupMetadataManager + .getOrMaybeCreateConsumerGroup(groupId, false) + .getOrMaybeCreateMember(memberId, false) + .state(); + } + + public Result<ConsumerGroupHeartbeatResponseData> consumerGroupHeartbeat( + ConsumerGroupHeartbeatRequestData request + ) { + snapshotRegistry.getOrCreateSnapshot(lastCommittedOffset); + + RequestContext context = new RequestContext( + new RequestHeader( + ApiKeys.CONSUMER_GROUP_HEARTBEAT, + ApiKeys.CONSUMER_GROUP_HEARTBEAT.latestVersion(), + "client", + 0 + ), + "1", + InetAddress.getLoopbackAddress(), + KafkaPrincipal.ANONYMOUS, + ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT), + SecurityProtocol.PLAINTEXT, + ClientInformation.EMPTY, + false + ); + + Result<ConsumerGroupHeartbeatResponseData> result = groupMetadataManager.consumerGroupHeartbeat( + context, + request + ); + + result.records().forEach(this::replay); + return result; + } + + private ApiMessage messageOrNull(ApiMessageAndVersion apiMessageAndVersion) { + if (apiMessageAndVersion == null) { + return null; + } else { + return apiMessageAndVersion.message(); + } + } + + private void replay( + Record record + ) { + ApiMessageAndVersion key = record.key(); + ApiMessageAndVersion value = record.value(); + + if (key == null) { + throw new IllegalStateException("Received a null key in " + record); + } + + switch (key.version()) { + case ConsumerGroupMemberMetadataKey.HIGHEST_SUPPORTED_VERSION: + groupMetadataManager.replay( + (ConsumerGroupMemberMetadataKey) key.message(), + (ConsumerGroupMemberMetadataValue) messageOrNull(value) + ); + break; + + case ConsumerGroupMetadataKey.HIGHEST_SUPPORTED_VERSION: + groupMetadataManager.replay( + (ConsumerGroupMetadataKey) key.message(), + (ConsumerGroupMetadataValue) messageOrNull(value) + ); + break; + + case ConsumerGroupPartitionMetadataKey.HIGHEST_SUPPORTED_VERSION: + groupMetadataManager.replay( + (ConsumerGroupPartitionMetadataKey) key.message(), + (ConsumerGroupPartitionMetadataValue) messageOrNull(value) + ); + break; + + case ConsumerGroupTargetAssignmentMemberKey.HIGHEST_SUPPORTED_VERSION: + groupMetadataManager.replay( + (ConsumerGroupTargetAssignmentMemberKey) key.message(), + (ConsumerGroupTargetAssignmentMemberValue) messageOrNull(value) + ); + break; + + case ConsumerGroupTargetAssignmentMetadataKey.HIGHEST_SUPPORTED_VERSION: + groupMetadataManager.replay( + (ConsumerGroupTargetAssignmentMetadataKey) key.message(), + (ConsumerGroupTargetAssignmentMetadataValue) messageOrNull(value) + ); + break; + + case ConsumerGroupCurrentMemberAssignmentKey.HIGHEST_SUPPORTED_VERSION: + groupMetadataManager.replay( + (ConsumerGroupCurrentMemberAssignmentKey) key.message(), + (ConsumerGroupCurrentMemberAssignmentValue) messageOrNull(value) + ); + break; + + default: + throw new IllegalStateException("Received an unknown record type " + key.version() + + " in " + record); + } + + lastWrittenOffset++; + } + } + + @Test + public void testConsumerHeartbeatRequestValidation() { + MockPartitionAssignor assignor = new MockPartitionAssignor("range"); + GroupMetadataManagerTestContext context = new GroupMetadataManagerTestContext.Builder() + .withAssignors(Collections.singletonList(assignor)) + .build(); + Exception ex; + + ex = assertThrows(InvalidRequestException.class, () -> context.consumerGroupHeartbeat( + new ConsumerGroupHeartbeatRequestData())); + assertEquals("GroupId can't be empty.", ex.getMessage()); + + ex = assertThrows(InvalidRequestException.class, () -> context.consumerGroupHeartbeat( + new ConsumerGroupHeartbeatRequestData() + .setGroupId("foo") + .setMemberEpoch(0))); + assertEquals("RebalanceTimeoutMs must be provided in first request.", ex.getMessage()); + + ex = assertThrows(InvalidRequestException.class, () -> context.consumerGroupHeartbeat( + new ConsumerGroupHeartbeatRequestData() + .setGroupId("foo") + .setMemberEpoch(0) + .setRebalanceTimeoutMs(5000))); + assertEquals("TopicPartitions must be empty when (re-)joining.", ex.getMessage()); Review Comment: The topic partitions should be an empty list when joining or re-joining. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: jira-unsubscr...@kafka.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org
