kevin-wu24 commented on code in PR #23136: URL: https://github.com/apache/kafka/pull/23136#discussion_r3763339061
########## raft/src/testFixtures/java/org/apache/kafka/raft/RaftClientContextBuilder.java: ########## @@ -0,0 +1,439 @@ +/* + * 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.raft; + +import org.apache.kafka.common.Node; +import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.common.Uuid; +import org.apache.kafka.common.config.AbstractConfig; +import org.apache.kafka.common.memory.MemoryPool; +import org.apache.kafka.common.metrics.Metrics; +import org.apache.kafka.common.record.internal.MemoryRecords; +import org.apache.kafka.common.utils.MockTime; +import org.apache.kafka.common.utils.internals.LogContext; +import org.apache.kafka.raft.internals.StringSerde; +import org.apache.kafka.server.common.Feature; +import org.apache.kafka.server.common.KRaftVersion; +import org.apache.kafka.server.common.OffsetAndEpoch; +import org.apache.kafka.server.common.serialization.RecordSerde; +import org.apache.kafka.snapshot.RecordsSnapshotWriter; +import org.apache.kafka.snapshot.Snapshots; + +import org.mockito.Mockito; + +import java.io.IOException; +import java.net.InetSocketAddress; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.OptionalInt; +import java.util.Set; +import java.util.function.Consumer; +import java.util.function.Function; +import java.util.stream.Collectors; +import java.util.stream.IntStream; + +import static org.apache.kafka.raft.SharedRaftClientContext.RaftProtocol.KIP_853_PROTOCOL; + +public final class RaftClientContextBuilder { + static final int DEFAULT_ELECTION_TIMEOUT_MS = 10000; + + static final RecordSerde<String> SERDE = new StringSerde(); + static final TopicPartition METADATA_PARTITION = new TopicPartition("metadata", 0); + static final int ELECTION_BACKOFF_MAX_MS = 100; + static final int FETCH_MAX_WAIT_MS = 0; + // fetch timeout is usually larger than election timeout + static final int FETCH_TIMEOUT_MS = 50000; + private static final int DEFAULT_REQUEST_TIMEOUT_MS = 5000; + static final int RETRY_BACKOFF_MS = 50; + private static final int DEFAULT_APPEND_LINGER_MS = 0; + + private final MockMessageQueue messageQueue = new MockMessageQueue(); + private final MockTime time = new MockTime(); + private final MockQuorumStateStore quorumStateStore = new MockQuorumStateStore(); + private final MockableRandom random = new MockableRandom(1L); + private final LogContext logContext = new LogContext(); + private final MockLog log = new MockLog(METADATA_PARTITION, Uuid.METADATA_TOPIC_ID, logContext); + private final String clusterId = Uuid.randomUuid().toString(); + private final OptionalInt localId; + private KRaftVersion kraftVersion = KRaftVersion.KRAFT_VERSION_0; + private final Uuid localDirectoryId; + + private int requestTimeoutMs = DEFAULT_REQUEST_TIMEOUT_MS; + private int electionTimeoutMs = DEFAULT_ELECTION_TIMEOUT_MS; + private int appendLingerMs = DEFAULT_APPEND_LINGER_MS; + private MemoryPool memoryPool = MemoryPool.NONE; + private Optional<List<InetSocketAddress>> bootstrapServers = Optional.empty(); + private SharedRaftClientContext.RaftProtocol raftProtocol = SharedRaftClientContext.RaftProtocol.KIP_595_PROTOCOL; + private boolean canBecomeVoter = false; + private VoterSet startingVoters = VoterSet.empty(); + private Endpoints localListeners = Endpoints.empty(); + private boolean isStartingVotersStatic = false; + private boolean autoJoin = false; + private int fetchSnapshotMaxBytes = QuorumConfig.DEFAULT_QUORUM_FETCH_SNAPSHOT_MAX_BYTES; + private int fetchMaxBytes = QuorumConfig.DEFAULT_QUORUM_FETCH_MAX_BYTES; + + public RaftClientContextBuilder(int localId, Set<Integer> staticVoters) { + this(OptionalInt.of(localId), staticVoters); + } + + public RaftClientContextBuilder(OptionalInt localId, Set<Integer> staticVoters) { + this(localId, Uuid.randomUuid()); + + withStaticVoters(staticVoters); + } + + public RaftClientContextBuilder(int localId, Uuid localDirectoryId) { + this(OptionalInt.of(localId), localDirectoryId); + } + + public RaftClientContextBuilder(OptionalInt localId, Uuid localDirectoryId) { + this.localId = localId; + this.localDirectoryId = localDirectoryId; + } + + RaftClientContextBuilder withElectedLeader(int epoch, int leaderId) { + quorumStateStore.writeElectionState( + ElectionState.withElectedLeader(epoch, leaderId, Optional.empty(), startingVoters.voterIds()), + kraftVersion + ); + return this; + } + + RaftClientContextBuilder withUnknownLeader(int epoch) { + quorumStateStore.writeElectionState( + ElectionState.withUnknownLeader(epoch, startingVoters.voterIds()), + kraftVersion + ); + return this; + } + + RaftClientContextBuilder withVotedCandidate(int epoch, ReplicaKey votedKey) { + quorumStateStore.writeElectionState( + ElectionState.withVotedCandidate(epoch, votedKey, startingVoters.voterIds()), + kraftVersion + ); + return this; + } + + RaftClientContextBuilder updateRandom(Consumer<MockableRandom> consumer) { + consumer.accept(random); + return this; + } + + RaftClientContextBuilder withMemoryPool(MemoryPool pool) { + this.memoryPool = pool; + return this; + } + + RaftClientContextBuilder withAppendLingerMs(int appendLingerMs) { + this.appendLingerMs = appendLingerMs; + return this; + } + + public RaftClientContextBuilder appendToLog(int epoch, List<String> records) { + MemoryRecords batch = RaftClientTestContext.buildBatch( + time.milliseconds(), + log.endOffset().offset(), + epoch, + records + ); + log.appendAsLeader(batch, epoch); + // Need to flush the log to update the last flushed offset. This is always correct + // because append operation was done in the Builder which represent the state of the + // log before the replica starts. + log.flush(false); + + // Reset the value of this method since "flush" before the replica start should not + // count when checking for flushes by the KRaft client. + log.flushedSinceLastChecked(); + return this; + } + + RaftClientContextBuilder withEmptySnapshot(OffsetAndEpoch snapshotId) { + try (RecordsSnapshotWriter<?> snapshot = new RecordsSnapshotWriter.Builder() + .setTime(time) + .setKraftVersion(KRaftVersion.KRAFT_VERSION_0) + .setRawSnapshotWriter(log.createNewSnapshotUnchecked(snapshotId).get()) + .build(SERDE) + ) { + snapshot.freeze(); + } + + return this; + } + + RaftClientContextBuilder deleteBeforeSnapshot(OffsetAndEpoch snapshotId) { + if (snapshotId.offset() > log.highWatermark().offset()) { + log.updateHighWatermark(new LogOffsetMetadata(snapshotId.offset())); + } + log.deleteBeforeSnapshot(snapshotId); + + return this; + } + + RaftClientContextBuilder withElectionTimeoutMs(int electionTimeoutMs) { + this.electionTimeoutMs = electionTimeoutMs; + return this; + } + + RaftClientContextBuilder withRequestTimeoutMs(int requestTimeoutMs) { + this.requestTimeoutMs = requestTimeoutMs; + return this; + } + + RaftClientContextBuilder withBootstrapServers(Optional<List<InetSocketAddress>> bootstrapServers) { + this.bootstrapServers = bootstrapServers; + return this; + } + + // deprecated, use withRpc instead + RaftClientContextBuilder withKip853Rpc(boolean withKip853Rpc) { + if (withKip853Rpc) { + this.raftProtocol = KIP_853_PROTOCOL; + } + return this; + } + + RaftClientContextBuilder withRaftProtocol(SharedRaftClientContext.RaftProtocol raftProtocol) { + this.raftProtocol = raftProtocol; + return this; + } + + RaftClientContextBuilder withCanBecomeVoter(boolean canBecomeVoter) { + this.canBecomeVoter = canBecomeVoter; + return this; + } + + RaftClientContextBuilder withStartingVoters(VoterSet voters, KRaftVersion kraftVersion) { + if (kraftVersion.isReconfigSupported()) { + return withBootstrapSnapshot(Optional.of(voters)); + } else { + return withStaticVoters(voters.voterIds()); + } + } + + RaftClientContextBuilder withStaticVoters(Set<Integer> staticVoters) { + Map<Integer, InetSocketAddress> staticVoterAddressMap = staticVoters + .stream() + .collect( + Collectors.toMap(Function.identity(), RaftClientTestContext::mockAddress) + ); + + return withStaticVoters( + VoterSet.fromInetSocketAddresses( + MockNetworkChannel.LISTENER_NAME, + staticVoterAddressMap + ) + ); + } + + RaftClientContextBuilder withStaticVoters(VoterSet staticVoters) { + startingVoters = staticVoters; + isStartingVotersStatic = true; + kraftVersion = KRaftVersion.KRAFT_VERSION_0; + + return this; + } + + RaftClientContextBuilder withBootstrapSnapshot(Optional<VoterSet> voters) { + return withBootstrapSnapshotRecords(voters, List.of()); + } + + RaftClientContextBuilder withBootstrapSnapshotRecords(Optional<VoterSet> voters, List<String> records) { + startingVoters = voters.orElse(VoterSet.empty()); + isStartingVotersStatic = false; + + if (voters.isPresent()) { + kraftVersion = KRaftVersion.LATEST_PRODUCTION; + + RecordsSnapshotWriter.Builder builder = new RecordsSnapshotWriter.Builder() + .setRawSnapshotWriter( + log.createNewSnapshotUnchecked(Snapshots.BOOTSTRAP_SNAPSHOT_ID).get() + ) + .setKraftVersion(kraftVersion) + .setVoterSet(voters); + + try (RecordsSnapshotWriter<String> writer = builder.build(SERDE)) { + if (!records.isEmpty()) { + writer.append(records); + } + writer.freeze(); + } + } else { + // Create an empty bootstrap snapshot if there is no voter set + kraftVersion = KRaftVersion.KRAFT_VERSION_0; + withEmptySnapshot(Snapshots.BOOTSTRAP_SNAPSHOT_ID); + } + + return this; + } + + RaftClientContextBuilder withLocalListeners(Endpoints localListeners) { + this.localListeners = localListeners; + return this; + } + + RaftClientContextBuilder withAutoJoin(boolean autoJoin) { + this.autoJoin = autoJoin; + return this; + } + + RaftClientContextBuilder withFetchSnapshotMaxBytes(int fetchSnapshotMaxSizeBytes) { + this.fetchSnapshotMaxBytes = fetchSnapshotMaxSizeBytes; + return this; + } + + RaftClientContextBuilder withFetchMaxBytes(int fetchMaxBytes) { + this.fetchMaxBytes = fetchMaxBytes; + return this; + } + + public RaftClientTestContext build() throws IOException { Review Comment: Can we rename this to `buildTestContext()`? ########## raft/src/testFixtures/java/org/apache/kafka/raft/SharedRaftClientContext.java: ########## @@ -0,0 +1,704 @@ +/* + * 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.raft; + +import org.apache.kafka.common.Node; +import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.common.Uuid; +import org.apache.kafka.common.message.AddRaftVoterRequestData; +import org.apache.kafka.common.message.AddRaftVoterResponseData; +import org.apache.kafka.common.message.ApiVersionsResponseData; +import org.apache.kafka.common.message.BeginQuorumEpochRequestData; +import org.apache.kafka.common.message.BeginQuorumEpochResponseData; +import org.apache.kafka.common.message.DescribeQuorumRequestData; +import org.apache.kafka.common.message.DescribeQuorumResponseData; +import org.apache.kafka.common.message.EndQuorumEpochRequestData; +import org.apache.kafka.common.message.EndQuorumEpochResponseData; +import org.apache.kafka.common.message.FetchRequestData; +import org.apache.kafka.common.message.FetchResponseData; +import org.apache.kafka.common.message.FetchSnapshotRequestData; +import org.apache.kafka.common.message.FetchSnapshotResponseData; +import org.apache.kafka.common.message.RemoveRaftVoterRequestData; +import org.apache.kafka.common.message.RemoveRaftVoterResponseData; +import org.apache.kafka.common.message.UpdateRaftVoterRequestData; +import org.apache.kafka.common.message.UpdateRaftVoterResponseData; +import org.apache.kafka.common.message.VoteRequestData; +import org.apache.kafka.common.message.VoteResponseData; +import org.apache.kafka.common.protocol.ApiKeys; +import org.apache.kafka.common.protocol.ApiMessage; +import org.apache.kafka.common.protocol.ByteBufferAccessor; +import org.apache.kafka.common.protocol.DataOutputStreamWritable; +import org.apache.kafka.common.protocol.Errors; +import org.apache.kafka.common.protocol.ObjectSerializationCache; +import org.apache.kafka.common.utils.MockTime; +import org.apache.kafka.server.common.KRaftVersion; +import org.apache.kafka.test.TestCondition; + +import java.io.ByteArrayOutputStream; +import java.io.DataOutputStream; +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.Iterator; +import java.util.List; +import java.util.Optional; +import java.util.OptionalInt; +import java.util.OptionalLong; + +/** + * The shared machinery for driving a mock {@link KafkaRaftClient} through the raft protocol. + * {@link RaftClientTestContext} extends this and overrides the relevant helpers to add protocol + * assertions, calling {@code super} for the shared work. + */ +public abstract class SharedRaftClientContext { + final TopicPartition metadataPartition = RaftClientContextBuilder.METADATA_PARTITION; + final Uuid metadataTopicId = Uuid.METADATA_TOPIC_ID; + final int fetchMaxBytes; + + int electionTimeoutMs; + + final MockQuorumStateStore quorumStateStore; + final String clusterId; + final OptionalInt localId; + public final Uuid localDirectoryId; + public final KRaftVersion kraftVersion; + public final KafkaRaftClient<String> client; + public final MockLog log; + final MockNetworkChannel channel; + final MockTime time; + final VoterSet startingVoters; + // Used to determine which RPC request and response to construct + final RaftProtocol raftProtocol; + + private final List<RaftResponse.Outbound> sentResponses = new ArrayList<>(); + private final List<Throwable> uncaughtExceptions = new ArrayList<>(); + + private static final int MAX_POLLS = 50; + + @SuppressWarnings("ParameterNumber") + SharedRaftClientContext( + String clusterId, + OptionalInt localId, + Uuid localDirectoryId, + KRaftVersion kraftVersion, + KafkaRaftClient<String> client, + MockLog log, + MockNetworkChannel channel, + MockTime time, + MockQuorumStateStore quorumStateStore, + VoterSet startingVoters, + RaftProtocol raftProtocol, + int fetchMaxBytes + ) { + this.clusterId = clusterId; + this.localId = localId; + this.localDirectoryId = localDirectoryId; + this.kraftVersion = kraftVersion; + this.client = client; + this.log = log; + this.channel = channel; + this.time = time; + this.quorumStateStore = quorumStateStore; + this.startingVoters = startingVoters; + this.raftProtocol = raftProtocol; + this.fetchMaxBytes = fetchMaxBytes; + } + + public void unattachedToCandidate() throws Exception { + time.sleep(electionTimeoutMs * 2L); + expectAndGrantPreVotes(currentEpoch()); + } + + public void unattachedToLeader() throws Exception { + int currentEpoch = currentEpoch(); + unattachedToCandidate(); + expectAndGrantVotes(currentEpoch + 1); + expectBeginEpoch(currentEpoch + 1); + } + + public OptionalInt currentLeader() { + return currentLeaderAndEpoch().leaderId(); + } + + public int currentEpoch() { + return currentLeaderAndEpoch().epoch(); + } + + LeaderAndEpoch currentLeaderAndEpoch() { + ElectionState election = quorumStateStore.readElectionState().get(); + return new LeaderAndEpoch(election.optionalLeaderId(), election.epoch()); + } + + void expectAndGrantVotes(int epoch) throws Exception { + pollUntilRequest(); + + List<RaftRequest.Outbound> voteRequests = collectVoteRequests(epoch, + log.lastFetchedEpoch(), log.endOffset().offset()); + + for (RaftRequest.Outbound request : voteRequests) { + VoteResponseData voteResponse = voteResponse(true, OptionalInt.empty(), epoch); + deliverResponse(request.correlationId(), request.destination(), voteResponse); + } + + pollUntil(() -> client.quorum().isLeader()); + } + + void expectAndGrantPreVotes(int epoch) throws Exception { Review Comment: I do not see an Override for this method in `RaftClientTestContext` that asserts ``` assertTrue(client.quorum().isCandidate()); ``` at the end of the method. I see an override for `expectAndGrantVotes` that asserts an elected leader after calling the super class' method. ########## raft/src/testFixtures/java/org/apache/kafka/raft/SharedRaftClientContext.java: ########## @@ -0,0 +1,704 @@ +/* + * 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.raft; + +import org.apache.kafka.common.Node; +import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.common.Uuid; +import org.apache.kafka.common.message.AddRaftVoterRequestData; +import org.apache.kafka.common.message.AddRaftVoterResponseData; +import org.apache.kafka.common.message.ApiVersionsResponseData; +import org.apache.kafka.common.message.BeginQuorumEpochRequestData; +import org.apache.kafka.common.message.BeginQuorumEpochResponseData; +import org.apache.kafka.common.message.DescribeQuorumRequestData; +import org.apache.kafka.common.message.DescribeQuorumResponseData; +import org.apache.kafka.common.message.EndQuorumEpochRequestData; +import org.apache.kafka.common.message.EndQuorumEpochResponseData; +import org.apache.kafka.common.message.FetchRequestData; +import org.apache.kafka.common.message.FetchResponseData; +import org.apache.kafka.common.message.FetchSnapshotRequestData; +import org.apache.kafka.common.message.FetchSnapshotResponseData; +import org.apache.kafka.common.message.RemoveRaftVoterRequestData; +import org.apache.kafka.common.message.RemoveRaftVoterResponseData; +import org.apache.kafka.common.message.UpdateRaftVoterRequestData; +import org.apache.kafka.common.message.UpdateRaftVoterResponseData; +import org.apache.kafka.common.message.VoteRequestData; +import org.apache.kafka.common.message.VoteResponseData; +import org.apache.kafka.common.protocol.ApiKeys; +import org.apache.kafka.common.protocol.ApiMessage; +import org.apache.kafka.common.protocol.ByteBufferAccessor; +import org.apache.kafka.common.protocol.DataOutputStreamWritable; +import org.apache.kafka.common.protocol.Errors; +import org.apache.kafka.common.protocol.ObjectSerializationCache; +import org.apache.kafka.common.utils.MockTime; +import org.apache.kafka.server.common.KRaftVersion; +import org.apache.kafka.test.TestCondition; + +import java.io.ByteArrayOutputStream; +import java.io.DataOutputStream; +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.Iterator; +import java.util.List; +import java.util.Optional; +import java.util.OptionalInt; +import java.util.OptionalLong; + +/** + * The shared machinery for driving a mock {@link KafkaRaftClient} through the raft protocol. + * {@link RaftClientTestContext} extends this and overrides the relevant helpers to add protocol + * assertions, calling {@code super} for the shared work. + */ +public abstract class SharedRaftClientContext { + final TopicPartition metadataPartition = RaftClientContextBuilder.METADATA_PARTITION; + final Uuid metadataTopicId = Uuid.METADATA_TOPIC_ID; + final int fetchMaxBytes; + + int electionTimeoutMs; + + final MockQuorumStateStore quorumStateStore; + final String clusterId; + final OptionalInt localId; + public final Uuid localDirectoryId; + public final KRaftVersion kraftVersion; + public final KafkaRaftClient<String> client; + public final MockLog log; + final MockNetworkChannel channel; + final MockTime time; + final VoterSet startingVoters; + // Used to determine which RPC request and response to construct + final RaftProtocol raftProtocol; + + private final List<RaftResponse.Outbound> sentResponses = new ArrayList<>(); + private final List<Throwable> uncaughtExceptions = new ArrayList<>(); + + private static final int MAX_POLLS = 50; + + @SuppressWarnings("ParameterNumber") + SharedRaftClientContext( + String clusterId, + OptionalInt localId, + Uuid localDirectoryId, + KRaftVersion kraftVersion, + KafkaRaftClient<String> client, + MockLog log, + MockNetworkChannel channel, + MockTime time, + MockQuorumStateStore quorumStateStore, + VoterSet startingVoters, + RaftProtocol raftProtocol, + int fetchMaxBytes + ) { + this.clusterId = clusterId; + this.localId = localId; + this.localDirectoryId = localDirectoryId; + this.kraftVersion = kraftVersion; + this.client = client; + this.log = log; + this.channel = channel; + this.time = time; + this.quorumStateStore = quorumStateStore; + this.startingVoters = startingVoters; + this.raftProtocol = raftProtocol; + this.fetchMaxBytes = fetchMaxBytes; + } + + public void unattachedToCandidate() throws Exception { + time.sleep(electionTimeoutMs * 2L); + expectAndGrantPreVotes(currentEpoch()); + } + + public void unattachedToLeader() throws Exception { + int currentEpoch = currentEpoch(); + unattachedToCandidate(); + expectAndGrantVotes(currentEpoch + 1); + expectBeginEpoch(currentEpoch + 1); + } + + public OptionalInt currentLeader() { + return currentLeaderAndEpoch().leaderId(); + } + + public int currentEpoch() { + return currentLeaderAndEpoch().epoch(); + } + + LeaderAndEpoch currentLeaderAndEpoch() { + ElectionState election = quorumStateStore.readElectionState().get(); + return new LeaderAndEpoch(election.optionalLeaderId(), election.epoch()); + } + + void expectAndGrantVotes(int epoch) throws Exception { + pollUntilRequest(); + + List<RaftRequest.Outbound> voteRequests = collectVoteRequests(epoch, + log.lastFetchedEpoch(), log.endOffset().offset()); + + for (RaftRequest.Outbound request : voteRequests) { + VoteResponseData voteResponse = voteResponse(true, OptionalInt.empty(), epoch); + deliverResponse(request.correlationId(), request.destination(), voteResponse); + } + + pollUntil(() -> client.quorum().isLeader()); + } + + void expectAndGrantPreVotes(int epoch) throws Exception { + pollUntilRequest(); + + List<RaftRequest.Outbound> voteRequests = collectPreVoteRequests( + epoch, + log.lastFetchedEpoch(), + log.endOffset().offset() + ); + + for (RaftRequest.Outbound request : voteRequests) { + if (!raftProtocol.isPreVoteSupported()) { + deliverResponse( + request.correlationId(), + request.destination(), + RaftUtil.errorResponse(ApiKeys.VOTE, Errors.UNSUPPORTED_VERSION) + ); + } else { + VoteResponseData voteResponse = voteResponse(true, OptionalInt.empty(), epoch); + deliverResponse(request.correlationId(), request.destination(), voteResponse); + } + } + + pollUntil(() -> client.quorum().isCandidate()); + } + + int localIdOrThrow() { + return localId.orElseThrow(() -> new AssertionError("Required local id is not defined")); + } + + private void expectBeginEpoch(int epoch) throws Exception { + pollUntilRequest(); + for (RaftRequest.Outbound request : collectBeginEpochRequests(epoch)) { + BeginQuorumEpochResponseData beginEpochResponse = beginEpochResponse(epoch, localIdOrThrow()); + deliverResponse(request.correlationId(), request.destination(), beginEpochResponse); + poll(); + } + } + + /** + * Asserts that no uncaught exceptions occurred in async callbacks (e.g., CompletionStage.whenComplete). + * This method is automatically called by the poll() wrapper method, but can also be called directly + * by tests to check for async exceptions at any point. + * + * @throws AssertionError if any uncaught exceptions were captured + */ + public void assertNoAsyncExceptions() { + if (!uncaughtExceptions.isEmpty()) { + Throwable first = uncaughtExceptions.get(0); + uncaughtExceptions.clear(); + throw new AssertionError("Uncaught exception in async callback", first); + } + } + + /** + * Poll for new events and check for any uncaught exceptions in async callbacks. + * This is a wrapper around client.poll() that also calls assertNoAsyncExceptions(). + */ + public void poll() { + client.poll(); + assertNoAsyncExceptions(); + } + + public void pollUntil(TestCondition condition) throws InterruptedException { + try { + for (int remaining = MAX_POLLS; remaining > 0; remaining--) { + poll(); + if (condition.conditionMet()) { + return; + } + } + } catch (Exception e) { + throw new RuntimeException(e); + } + throw new IllegalStateException( + String.format("Condition not met within %d polls", MAX_POLLS) + ); + } + + public void pollUntilResponse() throws InterruptedException { + pollUntil(() -> !sentResponses.isEmpty()); + } + + void pollUntilRequest() throws InterruptedException { + pollUntil(channel::hasSentRequests); + } + + // RaftClientTestContext overrides this to also assert the collected requests are pre-votes. + List<RaftRequest.Outbound> collectPreVoteRequests( + int epoch, + int lastEpoch, + long lastEpochOffset + ) { + return collectVoteRequestMessages(); + } + + // RaftClientTestContext overrides this to also assert the collected requests are (standard) votes. + List<RaftRequest.Outbound> collectVoteRequests( + int epoch, + int lastEpoch, + long lastEpochOffset + ) { + return collectVoteRequestMessages(); + } + + private List<RaftRequest.Outbound> collectVoteRequestMessages() { + List<RaftRequest.Outbound> voteRequests = new ArrayList<>(); + for (RaftRequest.Outbound raftMessage : channel.drainSendQueue()) { + if (raftMessage.data() instanceof VoteRequestData) { + voteRequests.add(raftMessage); + } + } + return voteRequests; + } + + // Round-trips a message through serialization to mimic the network, exercising the client's + // request/response encoding on every delivery. The benchmark context overrides this to skip the + // round trip so it is not measured as the client's own work. + ApiMessage maybeRoundTrip(ApiMessage message, short version) { + return roundTripApiMessage(message, version); + } + + private ApiMessage roundTripApiMessage(ApiMessage message, short version) { + ObjectSerializationCache cache = new ObjectSerializationCache(); + ByteArrayOutputStream buffer = new ByteArrayOutputStream(message.size(cache, version)); + + // Encode the message to a byte array with the given version + DataOutputStreamWritable writer = new DataOutputStreamWritable(new DataOutputStream(buffer)); + message.write(writer, cache, version); + + // Decode the message from the byte array + ByteBufferAccessor reader = new ByteBufferAccessor(ByteBuffer.wrap(buffer.toByteArray())); + message.read(reader, version); + + return message; + } + + public void deliverRequest(ApiMessage request) { + short version = raftRequestVersion(request); + deliverRequest(request, version); + } + + void deliverRequest(ApiMessage request, short version) { + deliverRequest(inboundRequest(request, version)); + } + + public RaftRequest.Inbound inboundRequest(ApiMessage request) { + return inboundRequest(request, raftRequestVersion(request)); + } + + private RaftRequest.Inbound inboundRequest(ApiMessage request, short version) { + ApiMessage versionedRequest = maybeRoundTrip(request, version); + return new RaftRequest.Inbound( + channel.listenerName(), + channel.newCorrelationId(), + version, + versionedRequest, + time.milliseconds() + ); + } + + private void deliverRequest(RaftRequest.Inbound inboundRequest) { + client.handle(inboundRequest).whenComplete((response, exception) -> { + if (exception != null) { + uncaughtExceptions.add(exception); + } else { + sentResponses.add(response); + } + }); + } + + void deliverResponse(int correlationId, Node source, ApiMessage response) { + short version = raftResponseVersion(response); + ApiMessage versionedResponse = maybeRoundTrip(response, version); + channel.mockReceive(new RaftResponse.Inbound(correlationId, versionedResponse, source)); + } + + List<RaftResponse.Outbound> drainSentResponses( + ApiKeys apiKey + ) { + List<RaftResponse.Outbound> res = new ArrayList<>(); + Iterator<RaftResponse.Outbound> iterator = sentResponses.iterator(); + while (iterator.hasNext()) { + RaftResponse.Outbound response = iterator.next(); + if (response.data().apiKey() == apiKey.id) { + res.add(response); + iterator.remove(); + } + } + return res; + } + + // RaftClientTestContext overrides this to also assert each collected request. + List<RaftRequest.Outbound> collectBeginEpochRequests(int epoch) { + return new ArrayList<>(channel.drainSentRequests(Optional.of(ApiKeys.BEGIN_QUORUM_EPOCH))); + } + + BeginQuorumEpochResponseData beginEpochResponse(int epoch, int leaderId) { + return RaftUtil.singletonBeginQuorumEpochResponse( + channel.listenerName(), + raftProtocol.beginQuorumEpochRpcVersion(), + Errors.NONE, + metadataPartition, + Errors.NONE, + epoch, + leaderId, + startingVoters.listeners(leaderId) + ); + } + + VoteRequestData voteRequest( + int epoch, + ReplicaKey candidateKey, + int lastEpoch, + long lastEpochOffset + ) { + return voteRequest( + clusterId, + epoch, + candidateKey, + lastEpoch, + lastEpochOffset, + false + ); + } + + VoteRequestData voteRequest( + String clusterId, + int epoch, + ReplicaKey candidateKey, + int lastEpoch, + long lastEpochOffset, + boolean preVote + ) { + ReplicaKey localReplicaKey = raftProtocol.isReconfigSupported() ? + ReplicaKey.of(localIdOrThrow(), localDirectoryId) : + ReplicaKey.of(-1, ReplicaKey.NO_DIRECTORY_ID); + + return voteRequest( + clusterId, + epoch, + candidateKey, + localReplicaKey, + lastEpoch, + lastEpochOffset, + preVote + ); + } + + VoteRequestData voteRequest( + String clusterId, + int epoch, + ReplicaKey candidateKey, + ReplicaKey voterKey, + int lastEpoch, + long lastEpochOffset, + boolean preVote + ) { + return RaftUtil.singletonVoteRequest( + metadataPartition, + clusterId, + epoch, + candidateKey, + voterKey, + lastEpoch, + lastEpochOffset, + preVote + ); + } + + VoteResponseData voteResponse(boolean voteGranted, OptionalInt leaderId, int epoch) { + return voteResponse(Errors.NONE, voteGranted, leaderId, epoch, raftProtocol.voteRpcVersion()); + } + + VoteResponseData voteResponse(Errors error, OptionalInt leaderId, int epoch) { + return voteResponse(error, false, leaderId, epoch, raftProtocol.voteRpcVersion()); + } + + VoteResponseData voteResponse(Errors error, boolean voteGranted, OptionalInt leaderId, int epoch, short version) { + return RaftUtil.singletonVoteResponse( + channel.listenerName(), + version, + Errors.NONE, + metadataPartition, + error, + epoch, + leaderId.orElse(-1), + voteGranted, + leaderId.isPresent() ? startingVoters.listeners(leaderId.getAsInt()) : Endpoints.empty() + ); + } + + public FetchRequestData fetchRequest( + int epoch, + ReplicaKey replicaKey, + long fetchOffset, + int lastFetchedEpoch, + int maxWaitTimeMs + ) { + return fetchRequest( + epoch, + replicaKey, + fetchOffset, + lastFetchedEpoch, + OptionalLong.of(Long.MAX_VALUE), + maxWaitTimeMs + ); + } + + FetchRequestData fetchRequest( + int epoch, + ReplicaKey replicaKey, + long fetchOffset, + int lastFetchedEpoch, + OptionalLong highWatermark, + int maxWaitTimeMs + ) { + return fetchRequest( + epoch, + clusterId, + replicaKey, + fetchOffset, + lastFetchedEpoch, + highWatermark, + maxWaitTimeMs + ); + } + + FetchRequestData fetchRequest( + int epoch, + String clusterId, + ReplicaKey replicaKey, + long fetchOffset, + int lastFetchedEpoch, + OptionalLong highWatermark, + int maxWaitTimeMs + ) { + FetchRequestData request = RaftUtil.singletonFetchRequest( + metadataPartition, + metadataTopicId, + fetchPartition -> { + fetchPartition + .setCurrentLeaderEpoch(epoch) + .setLastFetchedEpoch(lastFetchedEpoch) + .setFetchOffset(fetchOffset) + .setHighWatermark(highWatermark.orElse(-1)); + if (raftProtocol.isReconfigSupported()) { + fetchPartition + .setReplicaDirectoryId(replicaKey.directoryId().orElse(ReplicaKey.NO_DIRECTORY_ID)); + } + } + ); + return request + .setMaxWaitMs(maxWaitTimeMs) + .setClusterId(clusterId) + .setMaxBytes(fetchMaxBytes) + .setReplicaState( + new FetchRequestData.ReplicaState().setReplicaId(replicaKey.id()) + ); + } + + public DescribeQuorumRequestData describeQuorumRequest() { + return RaftUtil.singletonDescribeQuorumRequest(metadataPartition); + } + + private short raftRequestVersion(ApiMessage request) { + if (request instanceof FetchRequestData) { + return raftProtocol.fetchRpcVersion(); + } else if (request instanceof FetchSnapshotRequestData) { + return raftProtocol.fetchSnapshotRpcVersion(); + } else if (request instanceof VoteRequestData) { + return raftProtocol.voteRpcVersion(); + } else if (request instanceof BeginQuorumEpochRequestData) { + return raftProtocol.beginQuorumEpochRpcVersion(); + } else if (request instanceof EndQuorumEpochRequestData) { + return raftProtocol.endQuorumEpochRpcVersion(); + } else if (request instanceof DescribeQuorumRequestData) { + return raftProtocol.describeQuorumRpcVersion(); + } else if (request instanceof AddRaftVoterRequestData) { + return raftProtocol.addVoterRpcVersion(); + } else if (request instanceof RemoveRaftVoterRequestData) { + return raftProtocol.removeVoterRpcVersion(); + } else if (request instanceof UpdateRaftVoterRequestData) { + return raftProtocol.updateVoterRpcVersion(); + } else { + throw new IllegalArgumentException(String.format("Request %s is not a raft request", request)); + } + } + + private short raftResponseVersion(ApiMessage response) { + if (response instanceof FetchResponseData) { + return raftProtocol.fetchRpcVersion(); + } else if (response instanceof FetchSnapshotResponseData) { + return raftProtocol.fetchSnapshotRpcVersion(); + } else if (response instanceof VoteResponseData) { + return raftProtocol.voteRpcVersion(); + } else if (response instanceof BeginQuorumEpochResponseData) { + return raftProtocol.beginQuorumEpochRpcVersion(); + } else if (response instanceof EndQuorumEpochResponseData) { + return raftProtocol.endQuorumEpochRpcVersion(); + } else if (response instanceof DescribeQuorumResponseData) { + return raftProtocol.describeQuorumRpcVersion(); + } else if (response instanceof AddRaftVoterResponseData) { + return raftProtocol.addVoterRpcVersion(); + } else if (response instanceof RemoveRaftVoterResponseData) { + return raftProtocol.removeVoterRpcVersion(); + } else if (response instanceof UpdateRaftVoterResponseData) { + return raftProtocol.updateVoterRpcVersion(); + } else if (response instanceof ApiVersionsResponseData) { + return 4; + } else { + throw new IllegalArgumentException(String.format("Request %s is not a raft response", response)); + } + } + + // RaftClientTestContext overrides this to also assert the local node is the leader. + public void advanceLocalLeaderHighWatermarkToLogEndOffset() throws InterruptedException { + long localLogEndOffset = log.endOffset().offset(); + + Iterable<ReplicaKey> followers = () -> startingVoters + .voterKeys() + .stream() + .filter(voterKey -> voterKey.id() != localId.getAsInt()) + .iterator(); + + // Send a request from every voter + for (ReplicaKey follower : followers) { + deliverRequest( + fetchRequest(currentEpoch(), follower, localLogEndOffset, currentEpoch(), 0) + ); + + pollUntilResponse(); + // Drain the leader's fetch response so the next follower's poll sees a fresh one. + // RaftClientTestContext's override additionally asserts the response. Review Comment: We can remove these in-line comments. ########## raft/src/testFixtures/java/org/apache/kafka/raft/SharedRaftClientContext.java: ########## @@ -0,0 +1,704 @@ +/* + * 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.raft; + +import org.apache.kafka.common.Node; +import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.common.Uuid; +import org.apache.kafka.common.message.AddRaftVoterRequestData; +import org.apache.kafka.common.message.AddRaftVoterResponseData; +import org.apache.kafka.common.message.ApiVersionsResponseData; +import org.apache.kafka.common.message.BeginQuorumEpochRequestData; +import org.apache.kafka.common.message.BeginQuorumEpochResponseData; +import org.apache.kafka.common.message.DescribeQuorumRequestData; +import org.apache.kafka.common.message.DescribeQuorumResponseData; +import org.apache.kafka.common.message.EndQuorumEpochRequestData; +import org.apache.kafka.common.message.EndQuorumEpochResponseData; +import org.apache.kafka.common.message.FetchRequestData; +import org.apache.kafka.common.message.FetchResponseData; +import org.apache.kafka.common.message.FetchSnapshotRequestData; +import org.apache.kafka.common.message.FetchSnapshotResponseData; +import org.apache.kafka.common.message.RemoveRaftVoterRequestData; +import org.apache.kafka.common.message.RemoveRaftVoterResponseData; +import org.apache.kafka.common.message.UpdateRaftVoterRequestData; +import org.apache.kafka.common.message.UpdateRaftVoterResponseData; +import org.apache.kafka.common.message.VoteRequestData; +import org.apache.kafka.common.message.VoteResponseData; +import org.apache.kafka.common.protocol.ApiKeys; +import org.apache.kafka.common.protocol.ApiMessage; +import org.apache.kafka.common.protocol.ByteBufferAccessor; +import org.apache.kafka.common.protocol.DataOutputStreamWritable; +import org.apache.kafka.common.protocol.Errors; +import org.apache.kafka.common.protocol.ObjectSerializationCache; +import org.apache.kafka.common.utils.MockTime; +import org.apache.kafka.server.common.KRaftVersion; +import org.apache.kafka.test.TestCondition; + +import java.io.ByteArrayOutputStream; +import java.io.DataOutputStream; +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.Iterator; +import java.util.List; +import java.util.Optional; +import java.util.OptionalInt; +import java.util.OptionalLong; + +/** + * The shared machinery for driving a mock {@link KafkaRaftClient} through the raft protocol. + * {@link RaftClientTestContext} extends this and overrides the relevant helpers to add protocol + * assertions, calling {@code super} for the shared work. + */ +public abstract class SharedRaftClientContext { + final TopicPartition metadataPartition = RaftClientContextBuilder.METADATA_PARTITION; + final Uuid metadataTopicId = Uuid.METADATA_TOPIC_ID; + final int fetchMaxBytes; + + int electionTimeoutMs; + + final MockQuorumStateStore quorumStateStore; + final String clusterId; + final OptionalInt localId; + public final Uuid localDirectoryId; + public final KRaftVersion kraftVersion; + public final KafkaRaftClient<String> client; + public final MockLog log; + final MockNetworkChannel channel; + final MockTime time; + final VoterSet startingVoters; + // Used to determine which RPC request and response to construct + final RaftProtocol raftProtocol; + + private final List<RaftResponse.Outbound> sentResponses = new ArrayList<>(); + private final List<Throwable> uncaughtExceptions = new ArrayList<>(); + + private static final int MAX_POLLS = 50; + + @SuppressWarnings("ParameterNumber") + SharedRaftClientContext( + String clusterId, + OptionalInt localId, + Uuid localDirectoryId, + KRaftVersion kraftVersion, + KafkaRaftClient<String> client, + MockLog log, + MockNetworkChannel channel, + MockTime time, + MockQuorumStateStore quorumStateStore, + VoterSet startingVoters, + RaftProtocol raftProtocol, + int fetchMaxBytes + ) { + this.clusterId = clusterId; + this.localId = localId; + this.localDirectoryId = localDirectoryId; + this.kraftVersion = kraftVersion; + this.client = client; + this.log = log; + this.channel = channel; + this.time = time; + this.quorumStateStore = quorumStateStore; + this.startingVoters = startingVoters; + this.raftProtocol = raftProtocol; + this.fetchMaxBytes = fetchMaxBytes; + } + + public void unattachedToCandidate() throws Exception { + time.sleep(electionTimeoutMs * 2L); + expectAndGrantPreVotes(currentEpoch()); + } + + public void unattachedToLeader() throws Exception { + int currentEpoch = currentEpoch(); + unattachedToCandidate(); + expectAndGrantVotes(currentEpoch + 1); + expectBeginEpoch(currentEpoch + 1); + } + + public OptionalInt currentLeader() { + return currentLeaderAndEpoch().leaderId(); + } + + public int currentEpoch() { + return currentLeaderAndEpoch().epoch(); + } + + LeaderAndEpoch currentLeaderAndEpoch() { + ElectionState election = quorumStateStore.readElectionState().get(); + return new LeaderAndEpoch(election.optionalLeaderId(), election.epoch()); + } + + void expectAndGrantVotes(int epoch) throws Exception { + pollUntilRequest(); + + List<RaftRequest.Outbound> voteRequests = collectVoteRequests(epoch, + log.lastFetchedEpoch(), log.endOffset().offset()); + + for (RaftRequest.Outbound request : voteRequests) { + VoteResponseData voteResponse = voteResponse(true, OptionalInt.empty(), epoch); + deliverResponse(request.correlationId(), request.destination(), voteResponse); + } + + pollUntil(() -> client.quorum().isLeader()); + } + + void expectAndGrantPreVotes(int epoch) throws Exception { + pollUntilRequest(); + + List<RaftRequest.Outbound> voteRequests = collectPreVoteRequests( + epoch, + log.lastFetchedEpoch(), + log.endOffset().offset() + ); + + for (RaftRequest.Outbound request : voteRequests) { + if (!raftProtocol.isPreVoteSupported()) { + deliverResponse( + request.correlationId(), + request.destination(), + RaftUtil.errorResponse(ApiKeys.VOTE, Errors.UNSUPPORTED_VERSION) + ); + } else { + VoteResponseData voteResponse = voteResponse(true, OptionalInt.empty(), epoch); + deliverResponse(request.correlationId(), request.destination(), voteResponse); + } + } + + pollUntil(() -> client.quorum().isCandidate()); + } + + int localIdOrThrow() { + return localId.orElseThrow(() -> new AssertionError("Required local id is not defined")); + } + + private void expectBeginEpoch(int epoch) throws Exception { + pollUntilRequest(); + for (RaftRequest.Outbound request : collectBeginEpochRequests(epoch)) { + BeginQuorumEpochResponseData beginEpochResponse = beginEpochResponse(epoch, localIdOrThrow()); + deliverResponse(request.correlationId(), request.destination(), beginEpochResponse); + poll(); + } + } + + /** + * Asserts that no uncaught exceptions occurred in async callbacks (e.g., CompletionStage.whenComplete). + * This method is automatically called by the poll() wrapper method, but can also be called directly + * by tests to check for async exceptions at any point. + * + * @throws AssertionError if any uncaught exceptions were captured + */ + public void assertNoAsyncExceptions() { + if (!uncaughtExceptions.isEmpty()) { + Throwable first = uncaughtExceptions.get(0); + uncaughtExceptions.clear(); + throw new AssertionError("Uncaught exception in async callback", first); + } + } + + /** + * Poll for new events and check for any uncaught exceptions in async callbacks. + * This is a wrapper around client.poll() that also calls assertNoAsyncExceptions(). + */ + public void poll() { + client.poll(); + assertNoAsyncExceptions(); + } + + public void pollUntil(TestCondition condition) throws InterruptedException { + try { + for (int remaining = MAX_POLLS; remaining > 0; remaining--) { + poll(); + if (condition.conditionMet()) { + return; + } + } + } catch (Exception e) { + throw new RuntimeException(e); + } + throw new IllegalStateException( + String.format("Condition not met within %d polls", MAX_POLLS) + ); + } + + public void pollUntilResponse() throws InterruptedException { + pollUntil(() -> !sentResponses.isEmpty()); + } + + void pollUntilRequest() throws InterruptedException { + pollUntil(channel::hasSentRequests); + } + + // RaftClientTestContext overrides this to also assert the collected requests are pre-votes. + List<RaftRequest.Outbound> collectPreVoteRequests( + int epoch, + int lastEpoch, + long lastEpochOffset + ) { + return collectVoteRequestMessages(); + } + + // RaftClientTestContext overrides this to also assert the collected requests are (standard) votes. + List<RaftRequest.Outbound> collectVoteRequests( + int epoch, + int lastEpoch, + long lastEpochOffset + ) { + return collectVoteRequestMessages(); + } + + private List<RaftRequest.Outbound> collectVoteRequestMessages() { + List<RaftRequest.Outbound> voteRequests = new ArrayList<>(); + for (RaftRequest.Outbound raftMessage : channel.drainSendQueue()) { + if (raftMessage.data() instanceof VoteRequestData) { + voteRequests.add(raftMessage); + } + } + return voteRequests; + } + + // Round-trips a message through serialization to mimic the network, exercising the client's + // request/response encoding on every delivery. The benchmark context overrides this to skip the + // round trip so it is not measured as the client's own work. + ApiMessage maybeRoundTrip(ApiMessage message, short version) { + return roundTripApiMessage(message, version); + } + + private ApiMessage roundTripApiMessage(ApiMessage message, short version) { + ObjectSerializationCache cache = new ObjectSerializationCache(); + ByteArrayOutputStream buffer = new ByteArrayOutputStream(message.size(cache, version)); + + // Encode the message to a byte array with the given version + DataOutputStreamWritable writer = new DataOutputStreamWritable(new DataOutputStream(buffer)); + message.write(writer, cache, version); + + // Decode the message from the byte array + ByteBufferAccessor reader = new ByteBufferAccessor(ByteBuffer.wrap(buffer.toByteArray())); + message.read(reader, version); + + return message; + } + + public void deliverRequest(ApiMessage request) { + short version = raftRequestVersion(request); + deliverRequest(request, version); + } + + void deliverRequest(ApiMessage request, short version) { + deliverRequest(inboundRequest(request, version)); + } + + public RaftRequest.Inbound inboundRequest(ApiMessage request) { + return inboundRequest(request, raftRequestVersion(request)); + } + + private RaftRequest.Inbound inboundRequest(ApiMessage request, short version) { + ApiMessage versionedRequest = maybeRoundTrip(request, version); + return new RaftRequest.Inbound( + channel.listenerName(), + channel.newCorrelationId(), + version, + versionedRequest, + time.milliseconds() + ); + } + + private void deliverRequest(RaftRequest.Inbound inboundRequest) { + client.handle(inboundRequest).whenComplete((response, exception) -> { + if (exception != null) { + uncaughtExceptions.add(exception); + } else { + sentResponses.add(response); + } + }); + } + + void deliverResponse(int correlationId, Node source, ApiMessage response) { + short version = raftResponseVersion(response); + ApiMessage versionedResponse = maybeRoundTrip(response, version); + channel.mockReceive(new RaftResponse.Inbound(correlationId, versionedResponse, source)); + } + + List<RaftResponse.Outbound> drainSentResponses( + ApiKeys apiKey + ) { + List<RaftResponse.Outbound> res = new ArrayList<>(); + Iterator<RaftResponse.Outbound> iterator = sentResponses.iterator(); + while (iterator.hasNext()) { + RaftResponse.Outbound response = iterator.next(); + if (response.data().apiKey() == apiKey.id) { + res.add(response); + iterator.remove(); + } + } + return res; + } + + // RaftClientTestContext overrides this to also assert each collected request. Review Comment: I don' think we need the comments that say "oh this subclass will override this." This applies to multiple places in this file. ########## raft/src/testFixtures/java/org/apache/kafka/raft/RaftClientContextBuilder.java: ########## @@ -0,0 +1,439 @@ +/* + * 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.raft; + +import org.apache.kafka.common.Node; +import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.common.Uuid; +import org.apache.kafka.common.config.AbstractConfig; +import org.apache.kafka.common.memory.MemoryPool; +import org.apache.kafka.common.metrics.Metrics; +import org.apache.kafka.common.record.internal.MemoryRecords; +import org.apache.kafka.common.utils.MockTime; +import org.apache.kafka.common.utils.internals.LogContext; +import org.apache.kafka.raft.internals.StringSerde; +import org.apache.kafka.server.common.Feature; +import org.apache.kafka.server.common.KRaftVersion; +import org.apache.kafka.server.common.OffsetAndEpoch; +import org.apache.kafka.server.common.serialization.RecordSerde; +import org.apache.kafka.snapshot.RecordsSnapshotWriter; +import org.apache.kafka.snapshot.Snapshots; + +import org.mockito.Mockito; + +import java.io.IOException; +import java.net.InetSocketAddress; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.OptionalInt; +import java.util.Set; +import java.util.function.Consumer; +import java.util.function.Function; +import java.util.stream.Collectors; +import java.util.stream.IntStream; + +import static org.apache.kafka.raft.SharedRaftClientContext.RaftProtocol.KIP_853_PROTOCOL; + +public final class RaftClientContextBuilder { + static final int DEFAULT_ELECTION_TIMEOUT_MS = 10000; + + static final RecordSerde<String> SERDE = new StringSerde(); + static final TopicPartition METADATA_PARTITION = new TopicPartition("metadata", 0); + static final int ELECTION_BACKOFF_MAX_MS = 100; + static final int FETCH_MAX_WAIT_MS = 0; + // fetch timeout is usually larger than election timeout + static final int FETCH_TIMEOUT_MS = 50000; + private static final int DEFAULT_REQUEST_TIMEOUT_MS = 5000; + static final int RETRY_BACKOFF_MS = 50; + private static final int DEFAULT_APPEND_LINGER_MS = 0; + + private final MockMessageQueue messageQueue = new MockMessageQueue(); + private final MockTime time = new MockTime(); + private final MockQuorumStateStore quorumStateStore = new MockQuorumStateStore(); + private final MockableRandom random = new MockableRandom(1L); + private final LogContext logContext = new LogContext(); + private final MockLog log = new MockLog(METADATA_PARTITION, Uuid.METADATA_TOPIC_ID, logContext); + private final String clusterId = Uuid.randomUuid().toString(); + private final OptionalInt localId; + private KRaftVersion kraftVersion = KRaftVersion.KRAFT_VERSION_0; + private final Uuid localDirectoryId; + + private int requestTimeoutMs = DEFAULT_REQUEST_TIMEOUT_MS; + private int electionTimeoutMs = DEFAULT_ELECTION_TIMEOUT_MS; + private int appendLingerMs = DEFAULT_APPEND_LINGER_MS; + private MemoryPool memoryPool = MemoryPool.NONE; + private Optional<List<InetSocketAddress>> bootstrapServers = Optional.empty(); + private SharedRaftClientContext.RaftProtocol raftProtocol = SharedRaftClientContext.RaftProtocol.KIP_595_PROTOCOL; + private boolean canBecomeVoter = false; + private VoterSet startingVoters = VoterSet.empty(); + private Endpoints localListeners = Endpoints.empty(); + private boolean isStartingVotersStatic = false; + private boolean autoJoin = false; + private int fetchSnapshotMaxBytes = QuorumConfig.DEFAULT_QUORUM_FETCH_SNAPSHOT_MAX_BYTES; + private int fetchMaxBytes = QuorumConfig.DEFAULT_QUORUM_FETCH_MAX_BYTES; + + public RaftClientContextBuilder(int localId, Set<Integer> staticVoters) { + this(OptionalInt.of(localId), staticVoters); + } + + public RaftClientContextBuilder(OptionalInt localId, Set<Integer> staticVoters) { + this(localId, Uuid.randomUuid()); + + withStaticVoters(staticVoters); + } + + public RaftClientContextBuilder(int localId, Uuid localDirectoryId) { + this(OptionalInt.of(localId), localDirectoryId); + } + + public RaftClientContextBuilder(OptionalInt localId, Uuid localDirectoryId) { + this.localId = localId; + this.localDirectoryId = localDirectoryId; + } + + RaftClientContextBuilder withElectedLeader(int epoch, int leaderId) { + quorumStateStore.writeElectionState( + ElectionState.withElectedLeader(epoch, leaderId, Optional.empty(), startingVoters.voterIds()), + kraftVersion + ); + return this; + } + + RaftClientContextBuilder withUnknownLeader(int epoch) { + quorumStateStore.writeElectionState( + ElectionState.withUnknownLeader(epoch, startingVoters.voterIds()), + kraftVersion + ); + return this; + } + + RaftClientContextBuilder withVotedCandidate(int epoch, ReplicaKey votedKey) { + quorumStateStore.writeElectionState( + ElectionState.withVotedCandidate(epoch, votedKey, startingVoters.voterIds()), + kraftVersion + ); + return this; + } + + RaftClientContextBuilder updateRandom(Consumer<MockableRandom> consumer) { + consumer.accept(random); + return this; + } + + RaftClientContextBuilder withMemoryPool(MemoryPool pool) { + this.memoryPool = pool; + return this; + } + + RaftClientContextBuilder withAppendLingerMs(int appendLingerMs) { + this.appendLingerMs = appendLingerMs; + return this; + } + + public RaftClientContextBuilder appendToLog(int epoch, List<String> records) { + MemoryRecords batch = RaftClientTestContext.buildBatch( + time.milliseconds(), + log.endOffset().offset(), + epoch, + records + ); + log.appendAsLeader(batch, epoch); + // Need to flush the log to update the last flushed offset. This is always correct + // because append operation was done in the Builder which represent the state of the + // log before the replica starts. + log.flush(false); + + // Reset the value of this method since "flush" before the replica start should not + // count when checking for flushes by the KRaft client. + log.flushedSinceLastChecked(); + return this; + } + + RaftClientContextBuilder withEmptySnapshot(OffsetAndEpoch snapshotId) { + try (RecordsSnapshotWriter<?> snapshot = new RecordsSnapshotWriter.Builder() + .setTime(time) + .setKraftVersion(KRaftVersion.KRAFT_VERSION_0) + .setRawSnapshotWriter(log.createNewSnapshotUnchecked(snapshotId).get()) + .build(SERDE) + ) { + snapshot.freeze(); + } + + return this; + } + + RaftClientContextBuilder deleteBeforeSnapshot(OffsetAndEpoch snapshotId) { + if (snapshotId.offset() > log.highWatermark().offset()) { + log.updateHighWatermark(new LogOffsetMetadata(snapshotId.offset())); + } + log.deleteBeforeSnapshot(snapshotId); + + return this; + } + + RaftClientContextBuilder withElectionTimeoutMs(int electionTimeoutMs) { + this.electionTimeoutMs = electionTimeoutMs; + return this; + } + + RaftClientContextBuilder withRequestTimeoutMs(int requestTimeoutMs) { + this.requestTimeoutMs = requestTimeoutMs; + return this; + } + + RaftClientContextBuilder withBootstrapServers(Optional<List<InetSocketAddress>> bootstrapServers) { + this.bootstrapServers = bootstrapServers; + return this; + } + + // deprecated, use withRpc instead + RaftClientContextBuilder withKip853Rpc(boolean withKip853Rpc) { + if (withKip853Rpc) { + this.raftProtocol = KIP_853_PROTOCOL; + } + return this; + } + + RaftClientContextBuilder withRaftProtocol(SharedRaftClientContext.RaftProtocol raftProtocol) { + this.raftProtocol = raftProtocol; + return this; + } + + RaftClientContextBuilder withCanBecomeVoter(boolean canBecomeVoter) { + this.canBecomeVoter = canBecomeVoter; + return this; + } + + RaftClientContextBuilder withStartingVoters(VoterSet voters, KRaftVersion kraftVersion) { + if (kraftVersion.isReconfigSupported()) { + return withBootstrapSnapshot(Optional.of(voters)); + } else { + return withStaticVoters(voters.voterIds()); + } + } + + RaftClientContextBuilder withStaticVoters(Set<Integer> staticVoters) { + Map<Integer, InetSocketAddress> staticVoterAddressMap = staticVoters + .stream() + .collect( + Collectors.toMap(Function.identity(), RaftClientTestContext::mockAddress) + ); + + return withStaticVoters( + VoterSet.fromInetSocketAddresses( + MockNetworkChannel.LISTENER_NAME, + staticVoterAddressMap + ) + ); + } + + RaftClientContextBuilder withStaticVoters(VoterSet staticVoters) { + startingVoters = staticVoters; + isStartingVotersStatic = true; + kraftVersion = KRaftVersion.KRAFT_VERSION_0; + + return this; + } + + RaftClientContextBuilder withBootstrapSnapshot(Optional<VoterSet> voters) { + return withBootstrapSnapshotRecords(voters, List.of()); + } + + RaftClientContextBuilder withBootstrapSnapshotRecords(Optional<VoterSet> voters, List<String> records) { + startingVoters = voters.orElse(VoterSet.empty()); + isStartingVotersStatic = false; + + if (voters.isPresent()) { + kraftVersion = KRaftVersion.LATEST_PRODUCTION; + + RecordsSnapshotWriter.Builder builder = new RecordsSnapshotWriter.Builder() + .setRawSnapshotWriter( + log.createNewSnapshotUnchecked(Snapshots.BOOTSTRAP_SNAPSHOT_ID).get() + ) + .setKraftVersion(kraftVersion) + .setVoterSet(voters); + + try (RecordsSnapshotWriter<String> writer = builder.build(SERDE)) { + if (!records.isEmpty()) { + writer.append(records); + } + writer.freeze(); + } + } else { + // Create an empty bootstrap snapshot if there is no voter set + kraftVersion = KRaftVersion.KRAFT_VERSION_0; + withEmptySnapshot(Snapshots.BOOTSTRAP_SNAPSHOT_ID); + } + + return this; + } + + RaftClientContextBuilder withLocalListeners(Endpoints localListeners) { + this.localListeners = localListeners; + return this; + } + + RaftClientContextBuilder withAutoJoin(boolean autoJoin) { + this.autoJoin = autoJoin; + return this; + } + + RaftClientContextBuilder withFetchSnapshotMaxBytes(int fetchSnapshotMaxSizeBytes) { + this.fetchSnapshotMaxBytes = fetchSnapshotMaxSizeBytes; + return this; + } + + RaftClientContextBuilder withFetchMaxBytes(int fetchMaxBytes) { + this.fetchMaxBytes = fetchMaxBytes; + return this; + } + + public RaftClientTestContext build() throws IOException { + Metrics metrics = new Metrics(time); + MockNetworkChannel channel = new MockNetworkChannel(); + ExternalKRaftMetrics externalKRaftMetrics = Mockito.mock(ExternalKRaftMetrics.class); + RaftClientTestContext.MockListener listener = new RaftClientTestContext.MockListener(localId); + KafkaRaftClient<String> client = buildClient(listener, channel, metrics, externalKRaftMetrics); + + RaftClientTestContext context = new RaftClientTestContext( + clusterId, + localId, + localDirectoryId, + kraftVersion, + client, + log, + channel, + messageQueue, + time, + quorumStateStore, + startingVoters, + bootstrapIds(), + raftProtocol, + canBecomeVoter, + metrics, + externalKRaftMetrics, + listener, + fetchMaxBytes + ); + + applyOverrides(context); + context.requestTimeoutMs = requestTimeoutMs; + context.appendLingerMs = appendLingerMs; + return context; + } + + private KafkaRaftClient<String> buildClient( + RaftClient.Listener<String> registeredListener, + MockNetworkChannel channel, + Metrics metrics, + ExternalKRaftMetrics externalKRaftMetrics + ) { + Map<Integer, InetSocketAddress> staticVoterAddressMap = Map.of(); + if (isStartingVotersStatic) { + staticVoterAddressMap = startingVoters + .voterNodes(startingVoters.voterIds().stream(), channel.listenerName()) + .stream() + .collect( + Collectors.toMap( + Node::id, + node -> InetSocketAddress.createUnresolved(node.host(), node.port()) + ) + ); + } + + /* + * Compute the local listeners if the test didn't override it. + * Only potential voters/leader need to provide the local listeners. + * If the local id is not set (must be observer), the local listener can be empty. + */ + Endpoints localListeners = this.localListeners.isEmpty() ? + localId.isPresent() ? + startingVoters.listeners(localId.getAsInt()) : + Endpoints.empty() : + this.localListeners; + + Map<String, Object> configMap = new HashMap<>(); + configMap.put(QuorumConfig.QUORUM_REQUEST_TIMEOUT_MS_CONFIG, requestTimeoutMs); + configMap.put(QuorumConfig.QUORUM_RETRY_BACKOFF_MS_CONFIG, RETRY_BACKOFF_MS); + configMap.put(QuorumConfig.QUORUM_ELECTION_TIMEOUT_MS_CONFIG, electionTimeoutMs); + configMap.put(QuorumConfig.QUORUM_ELECTION_BACKOFF_MAX_MS_CONFIG, ELECTION_BACKOFF_MAX_MS); + configMap.put(QuorumConfig.QUORUM_FETCH_TIMEOUT_MS_CONFIG, FETCH_TIMEOUT_MS); + configMap.put(QuorumConfig.QUORUM_LINGER_MS_CONFIG, appendLingerMs); + configMap.put(QuorumConfig.QUORUM_AUTO_JOIN_ENABLE_CONFIG, autoJoin); + configMap.put(QuorumConfig.QUORUM_FETCH_SNAPSHOT_MAX_BYTES_CONFIG, fetchSnapshotMaxBytes); + configMap.put(QuorumConfig.QUORUM_FETCH_MAX_BYTES_CONFIG, fetchMaxBytes); + QuorumConfig quorumConfig = new QuorumConfig(new AbstractConfig(QuorumConfig.CONFIG_DEF, configMap)); + + List<InetSocketAddress> computedBootstrapServers = bootstrapServers.orElseGet(() -> { + if (isStartingVotersStatic) { + return List.of(); + } else { + return startingVoters + .voterNodes(startingVoters.voterIds().stream(), channel.listenerName()) + .stream() + .map(node -> InetSocketAddress.createUnresolved(node.host(), node.port())) + .collect(Collectors.toList()); + } + }); + + KafkaRaftClient<String> client = new KafkaRaftClient<>( + localId, + localDirectoryId, + SERDE, + channel, + messageQueue, + log, + memoryPool, + time, + new MockExpirationService(time), + FETCH_MAX_WAIT_MS, + canBecomeVoter, + clusterId, + computedBootstrapServers, + localListeners, + Feature.KRAFT_VERSION.supportedVersionRange(), + logContext, + random, + quorumConfig + ); + + client.register(registeredListener); + client.initialize( + staticVoterAddressMap, + quorumStateStore, + metrics, + externalKRaftMetrics + ); + + return client; + } + + private Set<Integer> bootstrapIds() { + return IntStream + .iterate(-2, id -> id - 1) + .limit(bootstrapServers.map(List::size).orElse(0)) + .boxed() + .collect(Collectors.toSet()); + } + + private void applyOverrides(SharedRaftClientContext context) { + context.electionTimeoutMs = electionTimeoutMs; + } Review Comment: Can we remove this method? ########## raft/src/testFixtures/java/org/apache/kafka/raft/SharedRaftClientContext.java: ########## @@ -0,0 +1,704 @@ +/* + * 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.raft; + +import org.apache.kafka.common.Node; +import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.common.Uuid; +import org.apache.kafka.common.message.AddRaftVoterRequestData; +import org.apache.kafka.common.message.AddRaftVoterResponseData; +import org.apache.kafka.common.message.ApiVersionsResponseData; +import org.apache.kafka.common.message.BeginQuorumEpochRequestData; +import org.apache.kafka.common.message.BeginQuorumEpochResponseData; +import org.apache.kafka.common.message.DescribeQuorumRequestData; +import org.apache.kafka.common.message.DescribeQuorumResponseData; +import org.apache.kafka.common.message.EndQuorumEpochRequestData; +import org.apache.kafka.common.message.EndQuorumEpochResponseData; +import org.apache.kafka.common.message.FetchRequestData; +import org.apache.kafka.common.message.FetchResponseData; +import org.apache.kafka.common.message.FetchSnapshotRequestData; +import org.apache.kafka.common.message.FetchSnapshotResponseData; +import org.apache.kafka.common.message.RemoveRaftVoterRequestData; +import org.apache.kafka.common.message.RemoveRaftVoterResponseData; +import org.apache.kafka.common.message.UpdateRaftVoterRequestData; +import org.apache.kafka.common.message.UpdateRaftVoterResponseData; +import org.apache.kafka.common.message.VoteRequestData; +import org.apache.kafka.common.message.VoteResponseData; +import org.apache.kafka.common.protocol.ApiKeys; +import org.apache.kafka.common.protocol.ApiMessage; +import org.apache.kafka.common.protocol.ByteBufferAccessor; +import org.apache.kafka.common.protocol.DataOutputStreamWritable; +import org.apache.kafka.common.protocol.Errors; +import org.apache.kafka.common.protocol.ObjectSerializationCache; +import org.apache.kafka.common.utils.MockTime; +import org.apache.kafka.server.common.KRaftVersion; +import org.apache.kafka.test.TestCondition; + +import java.io.ByteArrayOutputStream; +import java.io.DataOutputStream; +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.Iterator; +import java.util.List; +import java.util.Optional; +import java.util.OptionalInt; +import java.util.OptionalLong; + +/** + * The shared machinery for driving a mock {@link KafkaRaftClient} through the raft protocol. + * {@link RaftClientTestContext} extends this and overrides the relevant helpers to add protocol + * assertions, calling {@code super} for the shared work. + */ +public abstract class SharedRaftClientContext { + final TopicPartition metadataPartition = RaftClientContextBuilder.METADATA_PARTITION; + final Uuid metadataTopicId = Uuid.METADATA_TOPIC_ID; + final int fetchMaxBytes; + + int electionTimeoutMs; + + final MockQuorumStateStore quorumStateStore; + final String clusterId; + final OptionalInt localId; + public final Uuid localDirectoryId; + public final KRaftVersion kraftVersion; + public final KafkaRaftClient<String> client; + public final MockLog log; + final MockNetworkChannel channel; + final MockTime time; + final VoterSet startingVoters; + // Used to determine which RPC request and response to construct + final RaftProtocol raftProtocol; + + private final List<RaftResponse.Outbound> sentResponses = new ArrayList<>(); + private final List<Throwable> uncaughtExceptions = new ArrayList<>(); + + private static final int MAX_POLLS = 50; + + @SuppressWarnings("ParameterNumber") + SharedRaftClientContext( + String clusterId, + OptionalInt localId, + Uuid localDirectoryId, + KRaftVersion kraftVersion, + KafkaRaftClient<String> client, + MockLog log, + MockNetworkChannel channel, + MockTime time, + MockQuorumStateStore quorumStateStore, + VoterSet startingVoters, + RaftProtocol raftProtocol, + int fetchMaxBytes + ) { + this.clusterId = clusterId; + this.localId = localId; + this.localDirectoryId = localDirectoryId; + this.kraftVersion = kraftVersion; + this.client = client; + this.log = log; + this.channel = channel; + this.time = time; + this.quorumStateStore = quorumStateStore; + this.startingVoters = startingVoters; + this.raftProtocol = raftProtocol; + this.fetchMaxBytes = fetchMaxBytes; + } + + public void unattachedToCandidate() throws Exception { + time.sleep(electionTimeoutMs * 2L); + expectAndGrantPreVotes(currentEpoch()); + } + + public void unattachedToLeader() throws Exception { + int currentEpoch = currentEpoch(); + unattachedToCandidate(); + expectAndGrantVotes(currentEpoch + 1); + expectBeginEpoch(currentEpoch + 1); + } + + public OptionalInt currentLeader() { + return currentLeaderAndEpoch().leaderId(); + } + + public int currentEpoch() { + return currentLeaderAndEpoch().epoch(); + } + + LeaderAndEpoch currentLeaderAndEpoch() { + ElectionState election = quorumStateStore.readElectionState().get(); + return new LeaderAndEpoch(election.optionalLeaderId(), election.epoch()); + } + + void expectAndGrantVotes(int epoch) throws Exception { + pollUntilRequest(); + + List<RaftRequest.Outbound> voteRequests = collectVoteRequests(epoch, + log.lastFetchedEpoch(), log.endOffset().offset()); + + for (RaftRequest.Outbound request : voteRequests) { + VoteResponseData voteResponse = voteResponse(true, OptionalInt.empty(), epoch); + deliverResponse(request.correlationId(), request.destination(), voteResponse); + } + + pollUntil(() -> client.quorum().isLeader()); + } + + void expectAndGrantPreVotes(int epoch) throws Exception { + pollUntilRequest(); + + List<RaftRequest.Outbound> voteRequests = collectPreVoteRequests( + epoch, + log.lastFetchedEpoch(), + log.endOffset().offset() + ); + + for (RaftRequest.Outbound request : voteRequests) { + if (!raftProtocol.isPreVoteSupported()) { + deliverResponse( + request.correlationId(), + request.destination(), + RaftUtil.errorResponse(ApiKeys.VOTE, Errors.UNSUPPORTED_VERSION) + ); + } else { + VoteResponseData voteResponse = voteResponse(true, OptionalInt.empty(), epoch); + deliverResponse(request.correlationId(), request.destination(), voteResponse); + } + } + + pollUntil(() -> client.quorum().isCandidate()); + } + + int localIdOrThrow() { + return localId.orElseThrow(() -> new AssertionError("Required local id is not defined")); + } + + private void expectBeginEpoch(int epoch) throws Exception { + pollUntilRequest(); + for (RaftRequest.Outbound request : collectBeginEpochRequests(epoch)) { + BeginQuorumEpochResponseData beginEpochResponse = beginEpochResponse(epoch, localIdOrThrow()); + deliverResponse(request.correlationId(), request.destination(), beginEpochResponse); + poll(); + } + } + + /** + * Asserts that no uncaught exceptions occurred in async callbacks (e.g., CompletionStage.whenComplete). + * This method is automatically called by the poll() wrapper method, but can also be called directly + * by tests to check for async exceptions at any point. + * + * @throws AssertionError if any uncaught exceptions were captured + */ + public void assertNoAsyncExceptions() { + if (!uncaughtExceptions.isEmpty()) { + Throwable first = uncaughtExceptions.get(0); + uncaughtExceptions.clear(); + throw new AssertionError("Uncaught exception in async callback", first); + } + } + + /** + * Poll for new events and check for any uncaught exceptions in async callbacks. + * This is a wrapper around client.poll() that also calls assertNoAsyncExceptions(). + */ + public void poll() { + client.poll(); + assertNoAsyncExceptions(); + } + + public void pollUntil(TestCondition condition) throws InterruptedException { + try { + for (int remaining = MAX_POLLS; remaining > 0; remaining--) { + poll(); + if (condition.conditionMet()) { + return; + } + } + } catch (Exception e) { + throw new RuntimeException(e); + } + throw new IllegalStateException( + String.format("Condition not met within %d polls", MAX_POLLS) + ); + } + + public void pollUntilResponse() throws InterruptedException { + pollUntil(() -> !sentResponses.isEmpty()); + } + + void pollUntilRequest() throws InterruptedException { + pollUntil(channel::hasSentRequests); + } + + // RaftClientTestContext overrides this to also assert the collected requests are pre-votes. + List<RaftRequest.Outbound> collectPreVoteRequests( + int epoch, + int lastEpoch, + long lastEpochOffset + ) { + return collectVoteRequestMessages(); + } + + // RaftClientTestContext overrides this to also assert the collected requests are (standard) votes. + List<RaftRequest.Outbound> collectVoteRequests( + int epoch, + int lastEpoch, + long lastEpochOffset + ) { + return collectVoteRequestMessages(); + } + + private List<RaftRequest.Outbound> collectVoteRequestMessages() { + List<RaftRequest.Outbound> voteRequests = new ArrayList<>(); + for (RaftRequest.Outbound raftMessage : channel.drainSendQueue()) { + if (raftMessage.data() instanceof VoteRequestData) { + voteRequests.add(raftMessage); + } + } + return voteRequests; + } + + // Round-trips a message through serialization to mimic the network, exercising the client's + // request/response encoding on every delivery. The benchmark context overrides this to skip the + // round trip so it is not measured as the client's own work. + ApiMessage maybeRoundTrip(ApiMessage message, short version) { + return roundTripApiMessage(message, version); + } + + private ApiMessage roundTripApiMessage(ApiMessage message, short version) { + ObjectSerializationCache cache = new ObjectSerializationCache(); + ByteArrayOutputStream buffer = new ByteArrayOutputStream(message.size(cache, version)); + + // Encode the message to a byte array with the given version + DataOutputStreamWritable writer = new DataOutputStreamWritable(new DataOutputStream(buffer)); + message.write(writer, cache, version); + + // Decode the message from the byte array + ByteBufferAccessor reader = new ByteBufferAccessor(ByteBuffer.wrap(buffer.toByteArray())); + message.read(reader, version); + + return message; + } + + public void deliverRequest(ApiMessage request) { + short version = raftRequestVersion(request); + deliverRequest(request, version); + } + + void deliverRequest(ApiMessage request, short version) { + deliverRequest(inboundRequest(request, version)); + } + + public RaftRequest.Inbound inboundRequest(ApiMessage request) { + return inboundRequest(request, raftRequestVersion(request)); + } + + private RaftRequest.Inbound inboundRequest(ApiMessage request, short version) { + ApiMessage versionedRequest = maybeRoundTrip(request, version); + return new RaftRequest.Inbound( + channel.listenerName(), + channel.newCorrelationId(), + version, + versionedRequest, + time.milliseconds() + ); + } Review Comment: ```suggestion RaftRequest.Inbound inboundRequest(ApiMessage request, short version) { return new RaftRequest.Inbound( channel.listenerName(), channel.newCorrelationId(), version, versionedRequest, time.milliseconds() ); } ``` Then in `RaftClientTestContext` you can override this and say: ``` @Override RaftRequest.Inbound inboundRequest(ApiMessage request, short version) { ApiMessage msg = roundTripApiMessage(request, version); return super.inboundRequest(msg, version); } ``` That would get rid of `maybeRoundTrip` and `roundTripApiMessage`. ########## raft/src/testFixtures/java/org/apache/kafka/raft/SharedRaftClientContext.java: ########## @@ -0,0 +1,704 @@ +/* + * 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.raft; + +import org.apache.kafka.common.Node; +import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.common.Uuid; +import org.apache.kafka.common.message.AddRaftVoterRequestData; +import org.apache.kafka.common.message.AddRaftVoterResponseData; +import org.apache.kafka.common.message.ApiVersionsResponseData; +import org.apache.kafka.common.message.BeginQuorumEpochRequestData; +import org.apache.kafka.common.message.BeginQuorumEpochResponseData; +import org.apache.kafka.common.message.DescribeQuorumRequestData; +import org.apache.kafka.common.message.DescribeQuorumResponseData; +import org.apache.kafka.common.message.EndQuorumEpochRequestData; +import org.apache.kafka.common.message.EndQuorumEpochResponseData; +import org.apache.kafka.common.message.FetchRequestData; +import org.apache.kafka.common.message.FetchResponseData; +import org.apache.kafka.common.message.FetchSnapshotRequestData; +import org.apache.kafka.common.message.FetchSnapshotResponseData; +import org.apache.kafka.common.message.RemoveRaftVoterRequestData; +import org.apache.kafka.common.message.RemoveRaftVoterResponseData; +import org.apache.kafka.common.message.UpdateRaftVoterRequestData; +import org.apache.kafka.common.message.UpdateRaftVoterResponseData; +import org.apache.kafka.common.message.VoteRequestData; +import org.apache.kafka.common.message.VoteResponseData; +import org.apache.kafka.common.protocol.ApiKeys; +import org.apache.kafka.common.protocol.ApiMessage; +import org.apache.kafka.common.protocol.ByteBufferAccessor; +import org.apache.kafka.common.protocol.DataOutputStreamWritable; +import org.apache.kafka.common.protocol.Errors; +import org.apache.kafka.common.protocol.ObjectSerializationCache; +import org.apache.kafka.common.utils.MockTime; +import org.apache.kafka.server.common.KRaftVersion; +import org.apache.kafka.test.TestCondition; + +import java.io.ByteArrayOutputStream; +import java.io.DataOutputStream; +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.Iterator; +import java.util.List; +import java.util.Optional; +import java.util.OptionalInt; +import java.util.OptionalLong; + +/** + * The shared machinery for driving a mock {@link KafkaRaftClient} through the raft protocol. + * {@link RaftClientTestContext} extends this and overrides the relevant helpers to add protocol + * assertions, calling {@code super} for the shared work. + */ +public abstract class SharedRaftClientContext { + final TopicPartition metadataPartition = RaftClientContextBuilder.METADATA_PARTITION; + final Uuid metadataTopicId = Uuid.METADATA_TOPIC_ID; + final int fetchMaxBytes; + + int electionTimeoutMs; + + final MockQuorumStateStore quorumStateStore; + final String clusterId; + final OptionalInt localId; + public final Uuid localDirectoryId; + public final KRaftVersion kraftVersion; + public final KafkaRaftClient<String> client; + public final MockLog log; + final MockNetworkChannel channel; + final MockTime time; + final VoterSet startingVoters; + // Used to determine which RPC request and response to construct + final RaftProtocol raftProtocol; + + private final List<RaftResponse.Outbound> sentResponses = new ArrayList<>(); + private final List<Throwable> uncaughtExceptions = new ArrayList<>(); + + private static final int MAX_POLLS = 50; + + @SuppressWarnings("ParameterNumber") + SharedRaftClientContext( + String clusterId, + OptionalInt localId, + Uuid localDirectoryId, + KRaftVersion kraftVersion, + KafkaRaftClient<String> client, + MockLog log, + MockNetworkChannel channel, + MockTime time, + MockQuorumStateStore quorumStateStore, + VoterSet startingVoters, + RaftProtocol raftProtocol, + int fetchMaxBytes + ) { + this.clusterId = clusterId; + this.localId = localId; + this.localDirectoryId = localDirectoryId; + this.kraftVersion = kraftVersion; + this.client = client; + this.log = log; + this.channel = channel; + this.time = time; + this.quorumStateStore = quorumStateStore; + this.startingVoters = startingVoters; + this.raftProtocol = raftProtocol; + this.fetchMaxBytes = fetchMaxBytes; + } + + public void unattachedToCandidate() throws Exception { + time.sleep(electionTimeoutMs * 2L); + expectAndGrantPreVotes(currentEpoch()); + } + + public void unattachedToLeader() throws Exception { + int currentEpoch = currentEpoch(); + unattachedToCandidate(); + expectAndGrantVotes(currentEpoch + 1); + expectBeginEpoch(currentEpoch + 1); + } + + public OptionalInt currentLeader() { + return currentLeaderAndEpoch().leaderId(); + } + + public int currentEpoch() { + return currentLeaderAndEpoch().epoch(); + } + + LeaderAndEpoch currentLeaderAndEpoch() { + ElectionState election = quorumStateStore.readElectionState().get(); + return new LeaderAndEpoch(election.optionalLeaderId(), election.epoch()); + } + + void expectAndGrantVotes(int epoch) throws Exception { + pollUntilRequest(); + + List<RaftRequest.Outbound> voteRequests = collectVoteRequests(epoch, + log.lastFetchedEpoch(), log.endOffset().offset()); + + for (RaftRequest.Outbound request : voteRequests) { + VoteResponseData voteResponse = voteResponse(true, OptionalInt.empty(), epoch); + deliverResponse(request.correlationId(), request.destination(), voteResponse); + } + + pollUntil(() -> client.quorum().isLeader()); + } + + void expectAndGrantPreVotes(int epoch) throws Exception { + pollUntilRequest(); + + List<RaftRequest.Outbound> voteRequests = collectPreVoteRequests( + epoch, + log.lastFetchedEpoch(), + log.endOffset().offset() + ); + + for (RaftRequest.Outbound request : voteRequests) { + if (!raftProtocol.isPreVoteSupported()) { + deliverResponse( + request.correlationId(), + request.destination(), + RaftUtil.errorResponse(ApiKeys.VOTE, Errors.UNSUPPORTED_VERSION) + ); + } else { + VoteResponseData voteResponse = voteResponse(true, OptionalInt.empty(), epoch); + deliverResponse(request.correlationId(), request.destination(), voteResponse); + } + } + + pollUntil(() -> client.quorum().isCandidate()); + } + + int localIdOrThrow() { + return localId.orElseThrow(() -> new AssertionError("Required local id is not defined")); + } + + private void expectBeginEpoch(int epoch) throws Exception { + pollUntilRequest(); + for (RaftRequest.Outbound request : collectBeginEpochRequests(epoch)) { + BeginQuorumEpochResponseData beginEpochResponse = beginEpochResponse(epoch, localIdOrThrow()); + deliverResponse(request.correlationId(), request.destination(), beginEpochResponse); + poll(); + } + } + + /** + * Asserts that no uncaught exceptions occurred in async callbacks (e.g., CompletionStage.whenComplete). + * This method is automatically called by the poll() wrapper method, but can also be called directly + * by tests to check for async exceptions at any point. + * + * @throws AssertionError if any uncaught exceptions were captured + */ + public void assertNoAsyncExceptions() { + if (!uncaughtExceptions.isEmpty()) { + Throwable first = uncaughtExceptions.get(0); + uncaughtExceptions.clear(); + throw new AssertionError("Uncaught exception in async callback", first); + } + } + + /** + * Poll for new events and check for any uncaught exceptions in async callbacks. + * This is a wrapper around client.poll() that also calls assertNoAsyncExceptions(). + */ + public void poll() { + client.poll(); + assertNoAsyncExceptions(); + } + + public void pollUntil(TestCondition condition) throws InterruptedException { + try { + for (int remaining = MAX_POLLS; remaining > 0; remaining--) { + poll(); + if (condition.conditionMet()) { + return; + } + } + } catch (Exception e) { + throw new RuntimeException(e); + } + throw new IllegalStateException( + String.format("Condition not met within %d polls", MAX_POLLS) + ); + } + + public void pollUntilResponse() throws InterruptedException { + pollUntil(() -> !sentResponses.isEmpty()); + } + + void pollUntilRequest() throws InterruptedException { + pollUntil(channel::hasSentRequests); + } + + // RaftClientTestContext overrides this to also assert the collected requests are pre-votes. + List<RaftRequest.Outbound> collectPreVoteRequests( + int epoch, + int lastEpoch, + long lastEpochOffset + ) { + return collectVoteRequestMessages(); + } + + // RaftClientTestContext overrides this to also assert the collected requests are (standard) votes. + List<RaftRequest.Outbound> collectVoteRequests( + int epoch, + int lastEpoch, + long lastEpochOffset + ) { + return collectVoteRequestMessages(); + } + + private List<RaftRequest.Outbound> collectVoteRequestMessages() { + List<RaftRequest.Outbound> voteRequests = new ArrayList<>(); + for (RaftRequest.Outbound raftMessage : channel.drainSendQueue()) { + if (raftMessage.data() instanceof VoteRequestData) { + voteRequests.add(raftMessage); + } + } + return voteRequests; + } + + // Round-trips a message through serialization to mimic the network, exercising the client's + // request/response encoding on every delivery. The benchmark context overrides this to skip the + // round trip so it is not measured as the client's own work. + ApiMessage maybeRoundTrip(ApiMessage message, short version) { + return roundTripApiMessage(message, version); + } + + private ApiMessage roundTripApiMessage(ApiMessage message, short version) { Review Comment: This method should not be in this class because it would not be shared by `RaftClientBenchmarkContext`. ########## raft/src/testFixtures/java/org/apache/kafka/raft/RaftClientContextBuilder.java: ########## @@ -0,0 +1,439 @@ +/* + * 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.raft; + +import org.apache.kafka.common.Node; +import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.common.Uuid; +import org.apache.kafka.common.config.AbstractConfig; +import org.apache.kafka.common.memory.MemoryPool; +import org.apache.kafka.common.metrics.Metrics; +import org.apache.kafka.common.record.internal.MemoryRecords; +import org.apache.kafka.common.utils.MockTime; +import org.apache.kafka.common.utils.internals.LogContext; +import org.apache.kafka.raft.internals.StringSerde; +import org.apache.kafka.server.common.Feature; +import org.apache.kafka.server.common.KRaftVersion; +import org.apache.kafka.server.common.OffsetAndEpoch; +import org.apache.kafka.server.common.serialization.RecordSerde; +import org.apache.kafka.snapshot.RecordsSnapshotWriter; +import org.apache.kafka.snapshot.Snapshots; + +import org.mockito.Mockito; + +import java.io.IOException; +import java.net.InetSocketAddress; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.OptionalInt; +import java.util.Set; +import java.util.function.Consumer; +import java.util.function.Function; +import java.util.stream.Collectors; +import java.util.stream.IntStream; + +import static org.apache.kafka.raft.SharedRaftClientContext.RaftProtocol.KIP_853_PROTOCOL; + +public final class RaftClientContextBuilder { + static final int DEFAULT_ELECTION_TIMEOUT_MS = 10000; + + static final RecordSerde<String> SERDE = new StringSerde(); + static final TopicPartition METADATA_PARTITION = new TopicPartition("metadata", 0); + static final int ELECTION_BACKOFF_MAX_MS = 100; + static final int FETCH_MAX_WAIT_MS = 0; + // fetch timeout is usually larger than election timeout + static final int FETCH_TIMEOUT_MS = 50000; + private static final int DEFAULT_REQUEST_TIMEOUT_MS = 5000; + static final int RETRY_BACKOFF_MS = 50; + private static final int DEFAULT_APPEND_LINGER_MS = 0; + + private final MockMessageQueue messageQueue = new MockMessageQueue(); + private final MockTime time = new MockTime(); + private final MockQuorumStateStore quorumStateStore = new MockQuorumStateStore(); + private final MockableRandom random = new MockableRandom(1L); + private final LogContext logContext = new LogContext(); + private final MockLog log = new MockLog(METADATA_PARTITION, Uuid.METADATA_TOPIC_ID, logContext); + private final String clusterId = Uuid.randomUuid().toString(); + private final OptionalInt localId; + private KRaftVersion kraftVersion = KRaftVersion.KRAFT_VERSION_0; + private final Uuid localDirectoryId; + + private int requestTimeoutMs = DEFAULT_REQUEST_TIMEOUT_MS; + private int electionTimeoutMs = DEFAULT_ELECTION_TIMEOUT_MS; + private int appendLingerMs = DEFAULT_APPEND_LINGER_MS; + private MemoryPool memoryPool = MemoryPool.NONE; + private Optional<List<InetSocketAddress>> bootstrapServers = Optional.empty(); + private SharedRaftClientContext.RaftProtocol raftProtocol = SharedRaftClientContext.RaftProtocol.KIP_595_PROTOCOL; + private boolean canBecomeVoter = false; + private VoterSet startingVoters = VoterSet.empty(); + private Endpoints localListeners = Endpoints.empty(); + private boolean isStartingVotersStatic = false; + private boolean autoJoin = false; + private int fetchSnapshotMaxBytes = QuorumConfig.DEFAULT_QUORUM_FETCH_SNAPSHOT_MAX_BYTES; + private int fetchMaxBytes = QuorumConfig.DEFAULT_QUORUM_FETCH_MAX_BYTES; + + public RaftClientContextBuilder(int localId, Set<Integer> staticVoters) { + this(OptionalInt.of(localId), staticVoters); + } + + public RaftClientContextBuilder(OptionalInt localId, Set<Integer> staticVoters) { + this(localId, Uuid.randomUuid()); + + withStaticVoters(staticVoters); + } + + public RaftClientContextBuilder(int localId, Uuid localDirectoryId) { + this(OptionalInt.of(localId), localDirectoryId); + } + + public RaftClientContextBuilder(OptionalInt localId, Uuid localDirectoryId) { + this.localId = localId; + this.localDirectoryId = localDirectoryId; + } + + RaftClientContextBuilder withElectedLeader(int epoch, int leaderId) { + quorumStateStore.writeElectionState( + ElectionState.withElectedLeader(epoch, leaderId, Optional.empty(), startingVoters.voterIds()), + kraftVersion + ); + return this; + } + + RaftClientContextBuilder withUnknownLeader(int epoch) { + quorumStateStore.writeElectionState( + ElectionState.withUnknownLeader(epoch, startingVoters.voterIds()), + kraftVersion + ); + return this; + } + + RaftClientContextBuilder withVotedCandidate(int epoch, ReplicaKey votedKey) { + quorumStateStore.writeElectionState( + ElectionState.withVotedCandidate(epoch, votedKey, startingVoters.voterIds()), + kraftVersion + ); + return this; + } + + RaftClientContextBuilder updateRandom(Consumer<MockableRandom> consumer) { + consumer.accept(random); + return this; + } + + RaftClientContextBuilder withMemoryPool(MemoryPool pool) { + this.memoryPool = pool; + return this; + } + + RaftClientContextBuilder withAppendLingerMs(int appendLingerMs) { + this.appendLingerMs = appendLingerMs; + return this; + } + + public RaftClientContextBuilder appendToLog(int epoch, List<String> records) { + MemoryRecords batch = RaftClientTestContext.buildBatch( + time.milliseconds(), + log.endOffset().offset(), + epoch, + records + ); + log.appendAsLeader(batch, epoch); + // Need to flush the log to update the last flushed offset. This is always correct + // because append operation was done in the Builder which represent the state of the + // log before the replica starts. + log.flush(false); + + // Reset the value of this method since "flush" before the replica start should not + // count when checking for flushes by the KRaft client. + log.flushedSinceLastChecked(); + return this; + } + + RaftClientContextBuilder withEmptySnapshot(OffsetAndEpoch snapshotId) { + try (RecordsSnapshotWriter<?> snapshot = new RecordsSnapshotWriter.Builder() + .setTime(time) + .setKraftVersion(KRaftVersion.KRAFT_VERSION_0) + .setRawSnapshotWriter(log.createNewSnapshotUnchecked(snapshotId).get()) + .build(SERDE) + ) { + snapshot.freeze(); + } + + return this; + } + + RaftClientContextBuilder deleteBeforeSnapshot(OffsetAndEpoch snapshotId) { + if (snapshotId.offset() > log.highWatermark().offset()) { + log.updateHighWatermark(new LogOffsetMetadata(snapshotId.offset())); + } + log.deleteBeforeSnapshot(snapshotId); + + return this; + } + + RaftClientContextBuilder withElectionTimeoutMs(int electionTimeoutMs) { + this.electionTimeoutMs = electionTimeoutMs; + return this; + } + + RaftClientContextBuilder withRequestTimeoutMs(int requestTimeoutMs) { + this.requestTimeoutMs = requestTimeoutMs; + return this; + } + + RaftClientContextBuilder withBootstrapServers(Optional<List<InetSocketAddress>> bootstrapServers) { + this.bootstrapServers = bootstrapServers; + return this; + } + + // deprecated, use withRpc instead + RaftClientContextBuilder withKip853Rpc(boolean withKip853Rpc) { + if (withKip853Rpc) { + this.raftProtocol = KIP_853_PROTOCOL; + } + return this; + } + + RaftClientContextBuilder withRaftProtocol(SharedRaftClientContext.RaftProtocol raftProtocol) { + this.raftProtocol = raftProtocol; + return this; + } + + RaftClientContextBuilder withCanBecomeVoter(boolean canBecomeVoter) { + this.canBecomeVoter = canBecomeVoter; + return this; + } + + RaftClientContextBuilder withStartingVoters(VoterSet voters, KRaftVersion kraftVersion) { + if (kraftVersion.isReconfigSupported()) { + return withBootstrapSnapshot(Optional.of(voters)); + } else { + return withStaticVoters(voters.voterIds()); + } + } + + RaftClientContextBuilder withStaticVoters(Set<Integer> staticVoters) { + Map<Integer, InetSocketAddress> staticVoterAddressMap = staticVoters + .stream() + .collect( + Collectors.toMap(Function.identity(), RaftClientTestContext::mockAddress) + ); + + return withStaticVoters( + VoterSet.fromInetSocketAddresses( + MockNetworkChannel.LISTENER_NAME, + staticVoterAddressMap + ) + ); + } + + RaftClientContextBuilder withStaticVoters(VoterSet staticVoters) { + startingVoters = staticVoters; + isStartingVotersStatic = true; + kraftVersion = KRaftVersion.KRAFT_VERSION_0; + + return this; + } + + RaftClientContextBuilder withBootstrapSnapshot(Optional<VoterSet> voters) { + return withBootstrapSnapshotRecords(voters, List.of()); + } + + RaftClientContextBuilder withBootstrapSnapshotRecords(Optional<VoterSet> voters, List<String> records) { + startingVoters = voters.orElse(VoterSet.empty()); + isStartingVotersStatic = false; + + if (voters.isPresent()) { + kraftVersion = KRaftVersion.LATEST_PRODUCTION; + + RecordsSnapshotWriter.Builder builder = new RecordsSnapshotWriter.Builder() + .setRawSnapshotWriter( + log.createNewSnapshotUnchecked(Snapshots.BOOTSTRAP_SNAPSHOT_ID).get() + ) + .setKraftVersion(kraftVersion) + .setVoterSet(voters); + + try (RecordsSnapshotWriter<String> writer = builder.build(SERDE)) { + if (!records.isEmpty()) { + writer.append(records); + } + writer.freeze(); + } + } else { + // Create an empty bootstrap snapshot if there is no voter set + kraftVersion = KRaftVersion.KRAFT_VERSION_0; + withEmptySnapshot(Snapshots.BOOTSTRAP_SNAPSHOT_ID); + } + + return this; + } + + RaftClientContextBuilder withLocalListeners(Endpoints localListeners) { + this.localListeners = localListeners; + return this; + } + + RaftClientContextBuilder withAutoJoin(boolean autoJoin) { + this.autoJoin = autoJoin; + return this; + } + + RaftClientContextBuilder withFetchSnapshotMaxBytes(int fetchSnapshotMaxSizeBytes) { + this.fetchSnapshotMaxBytes = fetchSnapshotMaxSizeBytes; + return this; + } + + RaftClientContextBuilder withFetchMaxBytes(int fetchMaxBytes) { + this.fetchMaxBytes = fetchMaxBytes; + return this; + } + + public RaftClientTestContext build() throws IOException { + Metrics metrics = new Metrics(time); + MockNetworkChannel channel = new MockNetworkChannel(); + ExternalKRaftMetrics externalKRaftMetrics = Mockito.mock(ExternalKRaftMetrics.class); + RaftClientTestContext.MockListener listener = new RaftClientTestContext.MockListener(localId); + KafkaRaftClient<String> client = buildClient(listener, channel, metrics, externalKRaftMetrics); + + RaftClientTestContext context = new RaftClientTestContext( + clusterId, + localId, + localDirectoryId, + kraftVersion, + client, + log, + channel, + messageQueue, + time, + quorumStateStore, + startingVoters, + bootstrapIds(), + raftProtocol, + canBecomeVoter, + metrics, + externalKRaftMetrics, + listener, + fetchMaxBytes + ); + + applyOverrides(context); Review Comment: Why does this method only apply to the electionTimeoutMs? It seems like it should not. -- 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]
