izzyharker commented on code in PR #23220: URL: https://github.com/apache/kafka/pull/23220#discussion_r4028076590
########## group-coordinator/src/test/java/org/apache/kafka/coordinator/group/CompactionReplayTestContext.java: ########## @@ -0,0 +1,516 @@ +/* + * 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.clients.consumer.ConsumerPartitionAssignor; +import org.apache.kafka.clients.consumer.internals.ConsumerProtocol; +import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.common.Uuid; +import org.apache.kafka.common.message.ConsumerGroupHeartbeatRequestData; +import org.apache.kafka.common.message.ConsumerGroupHeartbeatResponseData; +import org.apache.kafka.common.message.JoinGroupRequestData; +import org.apache.kafka.common.message.JoinGroupResponseData; +import org.apache.kafka.common.message.LeaveGroupRequestData; +import org.apache.kafka.common.message.StreamsGroupHeartbeatRequestData; +import org.apache.kafka.common.message.StreamsGroupHeartbeatRequestData.Subtopology; +import org.apache.kafka.common.message.StreamsGroupHeartbeatRequestData.Topology; +import org.apache.kafka.common.message.StreamsGroupHeartbeatResponseData; +import org.apache.kafka.common.message.SyncGroupRequestData; +import org.apache.kafka.common.requests.ConsumerGroupHeartbeatRequest; +import org.apache.kafka.common.utils.Utils; +import org.apache.kafka.coordinator.common.runtime.CoordinatorMetadataImage; +import org.apache.kafka.coordinator.common.runtime.CoordinatorRecord; +import org.apache.kafka.coordinator.common.runtime.CoordinatorResult; +import org.apache.kafka.coordinator.group.api.assignor.GroupAssignment; +import org.apache.kafka.coordinator.group.modern.MemberAssignmentImpl; +import org.apache.kafka.coordinator.group.streams.MockTaskAssignor; +import org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil; +import org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.TaskRole; +import org.apache.kafka.coordinator.group.streams.TasksTuple; + +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.OptionalInt; +import java.util.OptionalLong; +import java.util.Set; +import java.util.function.Function; +import java.util.stream.Collectors; + +import static java.lang.Math.max; +import static org.apache.kafka.common.requests.JoinGroupRequest.UNKNOWN_MEMBER_ID; +import static org.apache.kafka.common.requests.StreamsGroupHeartbeatRequest.LEAVE_GROUP_MEMBER_EPOCH; +import static org.apache.kafka.coordinator.group.StreamsGroupTestUtil.staticHeartbeat; +import static org.apache.kafka.coordinator.group.StreamsGroupTestUtil.staticJoinHeartbeat; + +/** + * Drives group coordinator scenarios and records the resulting log for {@link + * GroupCoordinatorShardCompactionReplayTest}. Each request helper runs an operation against a live + * {@link GroupMetadataManagerTestContext} and appends whatever records it produced, so scenarios + * read as a sequence of coordinator operations rather than of log bookkeeping. + * + * Records are grouped into batches. The records from an atomic write are grouped into a single + * batch, and a non-atomic write gets one batch per record. + */ +final class CompactionReplayTestContext { + + static final String FOO_TOPIC_NAME = "foo"; + static final String BAR_TOPIC_NAME = "bar"; + static final String SUBTOPOLOGY_ID = "subtopology-1"; + + private static final int MAX_RECONCILIATION_ROUNDS = 10; + private static final int LONG_TIMEOUT_MS = 60000; + + private final GroupMetadataManagerTestContext context; + private final MockPartitionAssignor consumerAssignor; + private final MockTaskAssignor streamsAssignor; + private final CoordinatorMetadataImage metadataImage; + + private final List<List<CoordinatorRecord>> batches = new ArrayList<>(); + + CompactionReplayTestContext( + GroupMetadataManagerTestContext context, + MockPartitionAssignor consumerAssignor, + MockTaskAssignor streamsAssignor, + CoordinatorMetadataImage metadataImage + ) { + this.context = context; + this.consumerAssignor = consumerAssignor; + this.streamsAssignor = streamsAssignor; + this.metadataImage = metadataImage; + } + + /** + * Appends a CoordinatorResult which may be atomic or non-atomic. A + * non-atomic result puts each record in a unique batch, while an + * atomic result puts all records in a single batch + */ + private void append(CoordinatorResult<?, CoordinatorRecord> result) { + if (result.isAtomic()) { + append(result.records()); + } else { + result.records().forEach(this::append); + } + } + + private void append(List<CoordinatorRecord> batch) { + if (!batch.isEmpty()) { + batches.add(List.copyOf(batch)); + } + } + + private void append(CoordinatorRecord record) { + append(List.of(record)); + } + + /** + * The records every scenario step wrote, in order. + */ + List<CoordinatorRecord> records() { + return batches.stream().flatMap(List::stream).toList(); + } + + /** + * The positions in {@link #records()} at which a batch starts, plus the length of the log: + * the boundaries a cleaning window may fall on. + */ + List<Integer> batchBoundaries() { + List<Integer> boundaries = new ArrayList<>(); + int position = 0; + for (List<CoordinatorRecord> batch : batches) { + boundaries.add(position); + position += batch.size(); + } + boundaries.add(position); + return boundaries; + } + + /** + * The current type of {@code groupId}, used to assert that a scenario upgraded or downgraded the + * group as intended. + */ + Group.GroupType groupType(String groupId) { + return context.groupMetadataManager.group(groupId).type(); + } + + // Request helpers. + + /** + * A classic join request using the consumer embedded protocol. + */ + private JoinGroupRequestData classicJoinRequest(String groupId, String memberId) { + return classicJoinRequest(groupId, memberId, null); + } + + /** + * A classic join request for a static member (non-null {@code instanceId}) using the consumer + * embedded protocol. + */ + private JoinGroupRequestData classicJoinRequest(String groupId, String memberId, String instanceId) { + return new GroupMetadataManagerTestContext.JoinGroupRequestBuilder() + .withGroupId(groupId) + .withMemberId(memberId) + .withGroupInstanceId(instanceId) + .withProtocolType("consumer") + .withProtocols(GroupMetadataManagerTestContext.toConsumerProtocol( + List.of(FOO_TOPIC_NAME, BAR_TOPIC_NAME), List.of())) + .withRebalanceTimeoutMs(LONG_TIMEOUT_MS) + .withSessionTimeoutMs(LONG_TIMEOUT_MS) + .build(); + } + + /** + * Creates a classic group when the first member joins. + */ + JoinGroupResponseData joinFirstClassicMember(String groupId) throws Exception { + var firstJoin = context.sendClassicGroupJoin(classicJoinRequest(groupId, UNKNOWN_MEMBER_ID), true); + append(firstJoin.records); + firstJoin.appendFuture.complete(null); + String memberId = firstJoin.joinFuture.get().memberId(); + + var secondJoin = context.sendClassicGroupJoin(classicJoinRequest(groupId, memberId), true); + append(secondJoin.records); + secondJoin.appendFuture.complete(null); + // The first generation only forms once the initial rebalance delay has elapsed. + sleepCapturing(context.classicGroupInitialRebalanceDelayMs); + return secondJoin.joinFuture.get(); + } + + /** + * Joins a new member to an existing classic group, triggering a rebalance. The member first joins + * with an unknown id to be assigned one, then rejoins with it. Returns the assigned member id. + */ + String joinClassicMember(String groupId) throws Exception { + var firstJoin = context.sendClassicGroupJoin(classicJoinRequest(groupId, UNKNOWN_MEMBER_ID), true); + append(firstJoin.records); Review Comment: Sorry, I missed this one. It's because the group is created in `joinFirstClassicMember` so the coordinator result has a pending write. `joinClassicMember` doesn't need it because the group already exists. -- 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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
