mirrors/kafka
0
0 Fork 0
mirror of https://github.com/apache/kafka.git synced 2026-10-28 08:24:58 +08:00
Code Activity

KAFKA-21344: Handle partition rotation for share fetches (#23951)

Browse Source 
A share fetch hands its record budget to partitions in the order they
are read, so the partition read first  takes the largest share. Two
things decided that order earlier and together they could starve a
partition, hence changed the behaviour.

Earlier:
- `SharePartitionManager` rotated the fetch list by `sessionEpoch %
size`. The client bumps the epoch on both
  ShareFetch and ShareAcknowledge, so in a poll-and-commit loop fetch
epochs step by two and the rotation is
  constant. One unpaired request flips the parity, and with two
partitions the rotation then undoes the
  session reorder on every fetch, so the same partition is read first
until the other starves.
- `ShareSessionContext` moved served partitions to the end of the
session from inside the response iterator,
  in the iteration order of the response map. That map had passed
through `HashMap` copies, so when every
  partition returned records the session ended up in hash order and the
same partition led every fetch.

PR:

This change removes the epoch rotation from the fetch path and makes the
session reorder deterministic from  the fetch order. The fetch order is
preserved through `ShareFetchUtils`, `ShareFetch` and `KafkaApis` with
`LinkedHashMap`s. `ShareSessionContext` reorders the session in one pass
before pruning the response:  partitions without records keep their
position at the front, and partitions with records are appended in fetch
order rotated by one, so the partition read first is read last next
time. `PartitionRotateStrategy`  is reduced to that rotation and the
unused `sessionEpoch` parameter of `fetchMessages` is dropped.

`ShareConsumerFairnessTest` adds cases for the epoch parity change, for
the first position rotating when all partitions return records with
explicit and implicit acknowledgement, and for an idle partition being
read first once it has records.

`SharePartitionManagerTest` verifies the session order across subsequent
fetches.

Reviewers: Abhinav Dixit <adixit@confluent.io>, Andrew Schofield
 <aschofield@confluent.io>
This commit is contained in:
Apoorv Mittal authored and GitHub committed 2026-10-26 11:38:34 +00:00
1 parent e7a246f5a8
commit d193373f24
12 files changed
+478 -211
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

clients/clients-integration-tests/src/test/java/org/apache/kafka/clients/consumer/ShareConsumerFairnessTest.java
+200 -12
View File
@@ -17,7 +17,10 @@
package org.apache.kafka.clients.consumer;
import org.apache.kafka.clients.producer.Producer;
import org.apache.kafka.clients.producer.ProducerConfig;
import org.apache.kafka.clients.producer.ProducerRecord;
import org.apache.kafka.common.record.internal.DefaultRecordBatch;
import org.apache.kafka.common.record.internal.SimpleRecord;
import org.apache.kafka.common.test.ClusterInstance;
import org.apache.kafka.common.test.api.ClusterConfigProperty;
import org.apache.kafka.common.test.api.ClusterTest;
@@ -27,6 +30,8 @@ import org.apache.kafka.common.test.api.Type;
import org.junit.jupiter.api.Timeout;
import java.time.Duration;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.TreeMap;
@@ -59,24 +64,21 @@ public class ShareConsumerFairnessTest extends ShareConsumerTestBase {
private static final int PARTITIONS = 2;
private static final int RECORDS_PER_PARTITION = 8000;
private static final int MAX_POLL_RECORDS = 100;
private static final int POLLS_PER_PHASE = 20;
private static final int RECORD_VALUE_SIZE = 1024;
// Size on the log of a single record batch holding one record with a value of RECORD_VALUE_SIZE bytes.
private static final int RECORD_SIZE = DefaultRecordBatch.sizeInBytes(List.of(new SimpleRecord(new byte[RECORD_VALUE_SIZE])));
public ShareConsumerFairnessTest(ClusterInstance cluster) {
super(cluster);
}
@ClusterTest
public void testPartitionsWithDeepBacklogAreServedFairly() throws Exception {
public void testPartitionsServedFairlyInImplicitMode() {
String topic = "fairness";
String groupId = "fairness-group";
createTopic(topic, PARTITIONS, 1);
try (Producer<byte[], byte[]> producer = createProducer()) {
for (int i = 0; i < RECORDS_PER_PARTITION; i++) {
for (int p = 0; p < PARTITIONS; p++) {
producer.send(new ProducerRecord<>(topic, p, null, null, "value".getBytes()));
}
}
producer.flush();
}
produceRecords(topic, PARTITIONS);
alterShareAutoOffsetReset(groupId, "earliest");
Map<Integer, Integer> recordsByPartition = new TreeMap<>();
@@ -102,12 +104,198 @@ public class ShareConsumerFairnessTest extends ShareConsumerTestBase {
assertEquals(Map.of(0, RECORDS_PER_PARTITION, 1, RECORDS_PER_PARTITION), recordsByPartition);
for (int p = 0; p < PARTITIONS; p++) {
// Both partitions have a deep backlog, hence neither should have been starved while the other was
// being served, i.e. each has a meaningful share of the records when the first partition completes.
// Both partitions have a deep backlog and the fetches alternate between them, hence the other partition
// is only a few polls behind when the first partition completes. Allow a generous buffer, a starved
// partition is far below it.
int recordsOfPartition = recordsWhenFirstPartitionCompleted.getOrDefault(p, 0);
assertTrue(recordsOfPartition * 5 >= RECORDS_PER_PARTITION,
assertTrue(recordsOfPartition * 4 >= RECORDS_PER_PARTITION * 3,
"Partition " + p + " was starved, records by partition when the first partition completed: "
+ recordsWhenFirstPartitionCompleted);
}
}
@ClusterTest
public void testPartitionsServedFairlyWithExplicitAcknowledgements() {
verifyPartitionsServedFairlyWithExplicitAcknowledgements("fairness-epoch-parity", PARTITIONS);
verifyPartitionsServedFairlyWithExplicitAcknowledgements("fairness-epoch-parity-odd", 3);
}
@ClusterTest
public void testPartitionsRotateWhenAllPartitionsReturnRecordsWithExplicitAcknowledgement() {
verifyFirstPositionRotatesWhenAllPartitionsReturnRecords("fairness-rotation", PARTITIONS, "explicit");
verifyFirstPositionRotatesWhenAllPartitionsReturnRecords("fairness-rotation-odd", 3, "explicit");
}
@ClusterTest
public void testPartitionsRotateWhenAllPartitionsReturnRecordsWithImplicitAcknowledgement() {
verifyFirstPositionRotatesWhenAllPartitionsReturnRecords("fairness-rotation-implicit", 3, "implicit");
}
/**
* A partition which does not return records keeps its position at the front of the share session, hence
* it is read first as soon as it has records, ahead of the partition which was served all along.
*/
@ClusterTest
public void testPartitionWithoutRecordsIsReadFirstWhenRecordsArrive() {
String topic = "fairness-idle-partition";
String groupId = "fairness-idle-partition-group";
int lateRecords = MAX_POLL_RECORDS / 2;
createTopic(topic, PARTITIONS, 1);
// Only partition 0 has records initially.
try (Producer<byte[], byte[]> producer = createProducer()) {
for (int i = 0; i < RECORDS_PER_PARTITION; i++) {
producer.send(new ProducerRecord<>(topic, 0, null, null, "value".getBytes()));
}
producer.flush();
}
alterShareAutoOffsetReset(groupId, "earliest");
try (ShareConsumer<byte[], byte[]> consumer = createShareConsumer(groupId, Map.of(
ConsumerConfig.MAX_POLL_RECORDS_CONFIG, String.valueOf(MAX_POLL_RECORDS),
ConsumerConfig.SHARE_ACQUIRE_MODE_CONFIG, "record_limit",
ConsumerConfig.SHARE_ACKNOWLEDGEMENT_MODE_CONFIG, "explicit"))) {
consumer.subscribe(Set.of(topic));
Map<Integer, Integer> recordsByPartition = pollAndCommit(consumer);
assertEquals(Map.of(0, MAX_POLL_RECORDS * POLLS_PER_PHASE), recordsByPartition);
// Partition 1 gets records, fewer than a single poll can return.
try (Producer<byte[], byte[]> producer = createProducer()) {
for (int i = 0; i < lateRecords; i++) {
producer.send(new ProducerRecord<>(topic, 1, null, null, "value".getBytes()));
}
producer.flush();
}
// The next poll returns all the records of partition 1 and fills up from partition 0.
Map<Integer, Integer> nextPoll = new TreeMap<>();
pollRecords(consumer).forEach(record -> nextPoll.merge(record.partition(), 1, Integer::sum));
assertEquals(Map.of(0, MAX_POLL_RECORDS - lateRecords, 1, lateRecords), nextPoll);
}
}
private void verifyPartitionsServedFairlyWithExplicitAcknowledgements(String topic, int partitions) {
String groupId = topic + "-group";
createTopic(topic, partitions, 1);
produceRecords(topic, partitions);
alterShareAutoOffsetReset(groupId, "earliest");
Map<Integer, Integer> afterSingleAcknowledgement;
Map<Integer, Integer> afterMultipleAcknowledgements;
try (ShareConsumer<byte[], byte[]> consumer = createShareConsumer(groupId, Map.of(
ConsumerConfig.MAX_POLL_RECORDS_CONFIG, String.valueOf(MAX_POLL_RECORDS),
ConsumerConfig.SHARE_ACQUIRE_MODE_CONFIG, "record_limit",
ConsumerConfig.SHARE_ACKNOWLEDGEMENT_MODE_CONFIG, "explicit"))) {
consumer.subscribe(Set.of(topic));
afterSingleAcknowledgement = pollAndCommit(consumer);
// Commit the acknowledgements of one poll in two requests, which flips the parity of the epoch.
List<ConsumerRecord<byte[], byte[]>> records = pollRecords(consumer);
int half = records.size() / 2;
records.subList(0, half).forEach(consumer::acknowledge);
consumer.commitSync();
records.subList(half, records.size()).forEach(consumer::acknowledge);
consumer.commitSync();
afterMultipleAcknowledgements = pollAndCommit(consumer);
}
// Each phase consumes far fewer records than the backlog of either partition, hence both partitions have
// records throughout and each should get a meaningful share of the records of every phase.
String recordsByPhase = "records by partition before the epoch parity changed: " + afterSingleAcknowledgement
+ ", after: " + afterMultipleAcknowledgements;
for (Map<Integer, Integer> phase : List.of(afterSingleAcknowledgement, afterMultipleAcknowledgements)) {
int phaseRecords = phase.values().stream().mapToInt(Integer::intValue).sum();
for (int p = 0; p < partitions; p++) {
assertTrue(phase.getOrDefault(p, 0) * partitions * 2 >= phaseRecords,
"Partition " + p + " was starved, " + recordsByPhase);
}
}
}
private void verifyFirstPositionRotatesWhenAllPartitionsReturnRecords(String topic, int partitions, String acknowledgementMode) {
String groupId = topic + "-group";
// Records per partition a single fetch can return, capped by the fetch bytes of the partition.
int recordsPerPartitionPerFetch = 100;
createTopic(topic, partitions, 1);
// Enough records per partition to outlast the polls even if a partition were read first on every poll.
produceRecords(topic, partitions, POLLS_PER_PHASE * recordsPerPartitionPerFetch, RECORD_VALUE_SIZE);
alterShareAutoOffsetReset(groupId, "earliest");
Map<Integer, Integer> recordsByPartition;
try (ShareConsumer<byte[], byte[]> consumer = createShareConsumer(groupId, Map.of(
ConsumerConfig.FETCH_MAX_BYTES_CONFIG, String.valueOf(partitions * recordsPerPartitionPerFetch * RECORD_SIZE),
// Every partition but the last returns the records its fetch bytes allow, the last one gets half of it.
ConsumerConfig.MAX_POLL_RECORDS_CONFIG, String.valueOf(recordsPerPartitionPerFetch * partitions - recordsPerPartitionPerFetch / 2),
ConsumerConfig.SHARE_ACQUIRE_MODE_CONFIG, "record_limit",
ConsumerConfig.SHARE_ACKNOWLEDGEMENT_MODE_CONFIG, acknowledgementMode))) {
consumer.subscribe(Set.of(topic));
recordsByPartition = acknowledgementMode.equals("explicit") ? pollAndCommit(consumer) : pollImplicit(consumer);
}
// With the first position rotating, every partition takes every position in turn and the partitions
// consume the same number of records. Without rotation the last partition is stuck with half the share.
int totalRecords = recordsByPartition.values().stream().mapToInt(Integer::intValue).sum();
for (int p = 0; p < partitions; p++) {
assertTrue(recordsByPartition.getOrDefault(p, 0) * partitions * 4 >= totalRecords * 3,
"Partition " + p + " did not take the first position, records by partition: " + recordsByPartition);
}
}
private void produceRecords(String topic, int partitions) {
produceRecords(topic, partitions, RECORDS_PER_PARTITION, "value".getBytes().length);
}
/**
* Produces the given number of records of the given value size to every partition, each record in its own
* batch so the size of the records on the log is known.
*/
private void produceRecords(String topic, int partitions, int recordsPerPartition, int valueSize) {
try (Producer<byte[], byte[]> producer = createProducer(Map.of(ProducerConfig.BATCH_SIZE_CONFIG, 0))) {
for (int i = 0; i < recordsPerPartition; i++) {
for (int p = 0; p < partitions; p++) {
producer.send(new ProducerRecord<>(topic, p, null, null, new byte[valueSize]));
}
}
producer.flush();
}
}
/**
* Polls until {@link #POLLS_PER_PHASE} polls have returned records, and commits the acknowledgements of each poll
* in one request.
*/
private Map<Integer, Integer> pollAndCommit(ShareConsumer<byte[], byte[]> consumer) {
Map<Integer, Integer> recordsByPartition = new TreeMap<>();
for (int poll = 0; poll < POLLS_PER_PHASE; poll++) {
for (ConsumerRecord<byte[], byte[]> record : pollRecords(consumer)) {
recordsByPartition.merge(record.partition(), 1, Integer::sum);
consumer.acknowledge(record);
}
consumer.commitSync();
}
return recordsByPartition;
}
/**
* Polls until {@link #POLLS_PER_PHASE} polls have returned records, with implicit acknowledgement the records
* of a poll are acknowledged on the next poll.
*/
private Map<Integer, Integer> pollImplicit(ShareConsumer<byte[], byte[]> consumer) {
Map<Integer, Integer> recordsByPartition = new TreeMap<>();
for (int poll = 0; poll < POLLS_PER_PHASE; poll++) {
pollRecords(consumer).forEach(record -> recordsByPartition.merge(record.partition(), 1, Integer::sum));
}
return recordsByPartition;
}
/**
* Polls until records are returned, the test timeout fails the test if no records are delivered.
*/
private List<ConsumerRecord<byte[], byte[]>> pollRecords(ShareConsumer<byte[], byte[]> consumer) {
List<ConsumerRecord<byte[], byte[]>> records = new ArrayList<>();
while (records.isEmpty()) {
consumer.poll(Duration.ofMillis(200)).forEach(records::add);
}
return records;
}
}
core/src/main/java/kafka/server/share/ShareFetchUtils.java
+3 -2
View File
@@ -33,7 +33,6 @@ import org.apache.kafka.server.storage.log.FetchPartitionData;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.util.HashMap;
import java.util.Iterator;
import java.util.LinkedHashMap;
import java.util.List;
@@ -57,7 +56,9 @@ public class ShareFetchUtils {
PartitionMetadataProvider metadataProvider,
BiConsumer<SharePartitionKey, Throwable> exceptionHandler
) {
Map<TopicIdPartition, ShareFetchResponseData.PartitionData> response = new HashMap<>();
// Preserve the fetch order of the partitions in the response, the share session uses it to re-order
// the partitions for the subsequent fetch.
Map<TopicIdPartition, ShareFetchResponseData.PartitionData> response = new LinkedHashMap<>();
// Acquired records count for the share fetch request.
int acquiredRecordsCount = 0;
core/src/main/java/kafka/server/share/SharePartitionManager.java
+1 -9
View File
@@ -47,8 +47,6 @@ import org.apache.kafka.server.share.dlq.ShareGroupDLQManager;
import org.apache.kafka.server.share.fetch.DelayedShareFetchGroupKey;
import org.apache.kafka.server.share.fetch.DelayedShareFetchKey;
import org.apache.kafka.server.share.fetch.DelayedShareFetchPartitionKey;
import org.apache.kafka.server.share.fetch.PartitionRotateStrategy;
import org.apache.kafka.server.share.fetch.PartitionRotateStrategy.PartitionRotateMetadata;
import org.apache.kafka.server.share.fetch.ShareFetch;
import org.apache.kafka.server.share.metrics.ShareGroupMetrics;
import org.apache.kafka.server.share.persister.Persister;
@@ -309,7 +307,6 @@ public class SharePartitionManager implements AutoCloseable {
* @param memberId The member id, generated by the group-coordinator, this is used to identify the client.
* @param shareAcquireMode The share acquire mode from the share fetch request.
* @param fetchParams The fetch parameters from the share fetch request.
* @param sessionEpoch The session epoch for the member.
* @param maxFetchRecords The maximum number of records to fetch.
* @param batchSize The number of records per acquired records batch.
* @param topicIdPartitions The topic partitions to fetch for.
@@ -321,7 +318,6 @@ public class SharePartitionManager implements AutoCloseable {
String memberId,
FetchParams fetchParams,
byte shareAcquireMode,
int sessionEpoch,
int maxFetchRecords,
int batchSize,
List<TopicIdPartition> topicIdPartitions
@@ -329,12 +325,8 @@ public class SharePartitionManager implements AutoCloseable {
log.trace("Fetch request for topicIdPartitions: {} with groupId: {} fetch params: {}",
topicIdPartitions, groupId, fetchParams);
List<TopicIdPartition> rotatedTopicIdPartitions = PartitionRotateStrategy
.type(PartitionRotateStrategy.StrategyType.ROUND_ROBIN)
.rotate(topicIdPartitions, new PartitionRotateMetadata(sessionEpoch));
CompletableFuture<Map<TopicIdPartition, PartitionData>> future = new CompletableFuture<>();
processShareFetch(new ShareFetch(fetchParams, groupId, memberId, future, rotatedTopicIdPartitions, shareAcquireMode, batchSize, maxFetchRecords, brokerTopicStats));
processShareFetch(new ShareFetch(fetchParams, groupId, memberId, future, topicIdPartitions, shareAcquireMode, batchSize, maxFetchRecords, brokerTopicStats));
return future;
}
core/src/main/scala/kafka/server/KafkaApis.scala
+5 -5
View File
@@ -3420,7 +3420,7 @@ class KafkaApis(val requestChannel: RequestChannel,
if (shareFetchRequest.version() >= 2 && shareFetchRequest.data.isRenewAck)
CompletableFuture.completedFuture(mutable.Map.empty[TopicIdPartition, ShareFetchResponseData.PartitionData])
else
handleFetchFromShareFetchRequest(request, shareSessionEpoch, erroneousAndValidPartitionData, sharePartitionManager, authorizedTopics)
handleFetchFromShareFetchRequest(request, erroneousAndValidPartitionData, sharePartitionManager, authorizedTopics)
def combineShareFetchAndShareAcknowledgeResponses(fetchResult: CompletableFuture[Map[TopicIdPartition, ShareFetchResponseData.PartitionData]],
acknowledgeResult: CompletableFuture[Map[TopicIdPartition, ShareAcknowledgeResponseData.PartitionData]],
@@ -3516,7 +3516,6 @@ class KafkaApis(val requestChannel: RequestChannel,
// Visible for Testing
def handleFetchFromShareFetchRequest(request: Request,
shareSessionEpoch: Int,
erroneousAndValidPartitionData: ErroneousAndValidPartitionData,
sharePartitionManagerInstance: SharePartitionManager,
authorizedTopics: Set[String]
@@ -3577,19 +3576,20 @@ class KafkaApis(val requestChannel: RequestChannel,
shareFetchRequest.data.memberId,
params,
shareFetchRequest.data.shareAcquireMode,
shareSessionEpoch,
shareFetchRequest.data.maxRecords,
shareFetchRequest.data.batchSize,
interestedTopicPartitions
).thenApply{ result =>
val combinedResult = mutable.Map.empty[TopicIdPartition, ShareFetchResponseData.PartitionData]
// Preserve the fetch order of the partitions, the share session uses it to re-order the partitions
// for the subsequent fetch.
val combinedResult = mutable.LinkedHashMap.empty[TopicIdPartition, ShareFetchResponseData.PartitionData]
result.asScala.foreach { case (tp, data) =>
combinedResult += (tp -> data)
}
erroneous.foreach { case (tp, data) =>
combinedResult += (tp -> data)
}
combinedResult.toMap
combinedResult
}
}
}
core/src/test/java/kafka/server/share/ShareFetchUtilsTest.java
+41
View File
@@ -206,6 +206,47 @@ public class ShareFetchUtilsTest {
assertEquals(List.of(), resultData.get(tp1).acquiredRecords());
}
@Test
public void testProcessFetchResponsePreservesFetchOrder() {
String groupId = "grp";
String memberId = Uuid.randomUuid().toString();
TopicIdPartition tp0 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo", 0));
TopicIdPartition tp1 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo", 1));
TopicIdPartition tp2 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo", 2));
SharePartition sp0 = mock(SharePartition.class);
SharePartition sp1 = mock(SharePartition.class);
SharePartition sp2 = mock(SharePartition.class);
when(sp0.acquire(anyString(), any(ShareAcquireMode.class), anyInt(), anyInt(), anyLong(), any(FetchPartitionData.class), any())).thenReturn(ShareAcquiredRecords.empty());
when(sp1.acquire(anyString(), any(ShareAcquireMode.class), anyInt(), anyInt(), anyLong(), any(FetchPartitionData.class), any())).thenReturn(ShareAcquiredRecords.empty());
when(sp2.acquire(anyString(), any(ShareAcquireMode.class), anyInt(), anyInt(), anyLong(), any(FetchPartitionData.class), any())).thenReturn(ShareAcquiredRecords.empty());
LinkedHashMap<TopicIdPartition, SharePartition> sharePartitions = new LinkedHashMap<>();
sharePartitions.put(tp0, sp0);
sharePartitions.put(tp1, sp1);
sharePartitions.put(tp2, sp2);
ShareFetch shareFetch = new ShareFetch(FETCH_PARAMS, groupId, memberId,
new CompletableFuture<>(), List.of(tp0, tp1, tp2), BATCH_OPTIMIZED, BATCH_SIZE, 100, BROKER_TOPIC_STATS);
FetchPartitionData fetchPartitionData = new FetchPartitionData(Errors.NONE, 0L, 0L,
MemoryRecords.EMPTY, Optional.empty(), OptionalLong.empty(), Optional.empty(),
OptionalInt.empty(), false);
// The partitions are read in an order which differs from the request order.
List<ShareFetchPartitionData> responseData = List.of(
new ShareFetchPartitionData(tp2, 0, fetchPartitionData),
new ShareFetchPartitionData(tp0, 0, fetchPartitionData),
new ShareFetchPartitionData(tp1, 0, fetchPartitionData)
);
Map<TopicIdPartition, ShareFetchResponseData.PartitionData> resultData =
ShareFetchUtils.processFetchResponse(shareFetch, responseData, sharePartitions, mock(PartitionMetadataProvider.class), EXCEPTION_HANDLER);
// The share session re-orders its partitions from the response, hence the response must keep the
// order in which the partitions were read.
assertEquals(List.of(tp2, tp0, tp1), new ArrayList<>(resultData.keySet()));
}
@Test
public void testProcessFetchResponseWithLsoMovementForTopicPartition() {
String groupId = "grp";
core/src/test/java/kafka/server/share/SharePartitionManagerTest.java
+96 -67
View File
@@ -127,7 +127,6 @@ import scala.collection.Seq;
import scala.jdk.javaapi.CollectionConverters;
import static kafka.server.share.DelayedShareFetchTest.mockTopicIdPartitionToReturnDataEqualToMinBytes;
import static org.apache.kafka.server.share.fetch.ShareFetchTestUtils.validateRotatedListEquals;
import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
@@ -680,6 +679,51 @@ public class SharePartitionManagerTest {
assertEquals(1, resp2.responseData(topicNames).size());
}
@Test
public void testSubsequentShareSessionPartitionOrder() {
sharePartitionManager = SharePartitionManagerBuilder.builder().build();
Uuid fooId = Uuid.randomUuid();
TopicIdPartition tp0 = new TopicIdPartition(fooId, new TopicPartition("foo", 0));
TopicIdPartition tp1 = new TopicIdPartition(fooId, new TopicPartition("foo", 1));
TopicIdPartition tp2 = new TopicIdPartition(fooId, new TopicPartition("foo", 2));
String groupId = "grp";
String memberId = Uuid.randomUuid().toString();
// Create a new share session with foo-0, foo-1 and foo-2, the initial fetch does not re-order the partitions.
ShareFetchContext context = sharePartitionManager.newContext(groupId, List.of(tp0, tp1, tp2), EMPTY_PART_LIST,
memberId, ShareRequestMetadata.INITIAL_EPOCH, false, CONNECTION_ID);
assertInstanceOf(ShareSessionContext.class, context);
context.updateAndGenerateResponseData(groupId, memberId, responseData(List.of(tp0, tp1, tp2), Set.of(tp0)));
assertEquals(List.of(tp0, tp1, tp2), sharePartitionManager.cachedTopicIdPartitionsInShareSession(groupId, memberId));
// Only foo-0 returns records, hence it moves to the end and the other partitions keep their order.
context = sharePartitionManager.newContext(groupId, EMPTY_PART_LIST, EMPTY_PART_LIST, memberId, 1, true, CONNECTION_ID);
context.updateAndGenerateResponseData(groupId, memberId, responseData(List.of(tp0, tp1, tp2), Set.of(tp0)));
assertEquals(List.of(tp1, tp2, tp0), sharePartitionManager.cachedTopicIdPartitionsInShareSession(groupId, memberId));
// Fetched in the session order foo-1, foo-2, foo-0 where foo-1 and foo-2 return records. The partition
// without records is first, then the partitions with records where the one read first is the last.
context = sharePartitionManager.newContext(groupId, EMPTY_PART_LIST, EMPTY_PART_LIST, memberId, 2, true, CONNECTION_ID);
context.updateAndGenerateResponseData(groupId, memberId, responseData(List.of(tp1, tp2, tp0), Set.of(tp1, tp2)));
assertEquals(List.of(tp0, tp2, tp1), sharePartitionManager.cachedTopicIdPartitionsInShareSession(groupId, memberId));
// All the partitions return records, hence the partition read first moves to the end and the first
// position rotates.
context = sharePartitionManager.newContext(groupId, EMPTY_PART_LIST, EMPTY_PART_LIST, memberId, 3, true, CONNECTION_ID);
context.updateAndGenerateResponseData(groupId, memberId, responseData(List.of(tp0, tp2, tp1), Set.of(tp0, tp1, tp2)));
assertEquals(List.of(tp2, tp1, tp0), sharePartitionManager.cachedTopicIdPartitionsInShareSession(groupId, memberId));
context = sharePartitionManager.newContext(groupId, EMPTY_PART_LIST, EMPTY_PART_LIST, memberId, 4, true, CONNECTION_ID);
context.updateAndGenerateResponseData(groupId, memberId, responseData(List.of(tp2, tp1, tp0), Set.of(tp0, tp1, tp2)));
assertEquals(List.of(tp1, tp0, tp2), sharePartitionManager.cachedTopicIdPartitionsInShareSession(groupId, memberId));
// No partition returns records, hence the order is unchanged.
context = sharePartitionManager.newContext(groupId, EMPTY_PART_LIST, EMPTY_PART_LIST, memberId, 5, true, CONNECTION_ID);
context.updateAndGenerateResponseData(groupId, memberId, responseData(List.of(tp1, tp0, tp2), Set.of()));
assertEquals(List.of(tp1, tp0, tp2), sharePartitionManager.cachedTopicIdPartitionsInShareSession(groupId, memberId));
}
@Test
public void testZeroSizeShareSession() {
ShareSessionCache cache = new ShareSessionCache(10);
@@ -1208,18 +1252,18 @@ public class SharePartitionManagerTest {
doAnswer(invocation -> buildLogReadResult(topicIdPartitions)).when(mockReplicaManager).readFromLog(any(), any(), any(ReplicaQuota.class), anyBoolean());
CompletableFuture<Map<TopicIdPartition, PartitionData>> future = sharePartitionManager.fetchMessages(
groupId, memberId1.toString(), FETCH_PARAMS, BATCH_OPTIMIZED, 1, MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
groupId, memberId1.toString(), FETCH_PARAMS, BATCH_OPTIMIZED, MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
assertTrue(future.isDone());
Mockito.verify(mockReplicaManager, times(1)).readFromLog(
any(), any(), any(ReplicaQuota.class), anyBoolean());
future = sharePartitionManager.fetchMessages(groupId, memberId1.toString(), FETCH_PARAMS, BATCH_OPTIMIZED, 3,
future = sharePartitionManager.fetchMessages(groupId, memberId1.toString(), FETCH_PARAMS, BATCH_OPTIMIZED,
MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
assertTrue(future.isDone());
Mockito.verify(mockReplicaManager, times(2)).readFromLog(
any(), any(), any(ReplicaQuota.class), anyBoolean());
future = sharePartitionManager.fetchMessages(groupId, memberId1.toString(), FETCH_PARAMS, BATCH_OPTIMIZED, 10,
future = sharePartitionManager.fetchMessages(groupId, memberId1.toString(), FETCH_PARAMS, BATCH_OPTIMIZED,
MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
assertTrue(future.isDone());
Mockito.verify(mockReplicaManager, times(3)).readFromLog(
@@ -1263,7 +1307,7 @@ public class SharePartitionManagerTest {
.build();
CompletableFuture<Map<TopicIdPartition, ShareFetchResponseData.PartitionData>> future =
sharePartitionManager.fetchMessages(groupId, memberId, FETCH_PARAMS, BATCH_OPTIMIZED, 0,
sharePartitionManager.fetchMessages(groupId, memberId, FETCH_PARAMS, BATCH_OPTIMIZED,
MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
Mockito.verify(mockReplicaManager, times(0)).readFromLog(
any(), any(), any(ReplicaQuota.class), anyBoolean());
@@ -1302,7 +1346,7 @@ public class SharePartitionManagerTest {
doAnswer(invocation -> buildLogReadResult(topicIdPartitions)).when(mockReplicaManager).readFromLog(any(), any(), any(ReplicaQuota.class), anyBoolean());
sharePartitionManager.fetchMessages(groupId, memberId, FETCH_PARAMS, BATCH_OPTIMIZED, 0,
sharePartitionManager.fetchMessages(groupId, memberId, FETCH_PARAMS, BATCH_OPTIMIZED,
MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
// Since the nextFetchOffset does not point to endOffset + 1, i.e. some of the records in the cachedState are AVAILABLE,
// even though the maxInFlightMessages limit is exceeded, replicaManager.readFromLog should be called
@@ -2221,7 +2265,7 @@ public class SharePartitionManagerTest {
.build();
CompletableFuture<Map<TopicIdPartition, ShareFetchResponseData.PartitionData>> future =
sharePartitionManager.fetchMessages(groupId, memberId, FETCH_PARAMS, BATCH_OPTIMIZED, 0,
sharePartitionManager.fetchMessages(groupId, memberId, FETCH_PARAMS, BATCH_OPTIMIZED,
MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
// Verify that the fetch request is completed.
TestUtils.waitForCondition(
@@ -2291,7 +2335,7 @@ public class SharePartitionManagerTest {
.build();
CompletableFuture<Map<TopicIdPartition, ShareFetchResponseData.PartitionData>> future =
sharePartitionManager.fetchMessages(groupId, Uuid.randomUuid().toString(), FETCH_PARAMS, BATCH_OPTIMIZED, 0,
sharePartitionManager.fetchMessages(groupId, Uuid.randomUuid().toString(), FETCH_PARAMS, BATCH_OPTIMIZED,
MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
// Verify that the fetch request is completed.
TestUtils.waitForCondition(
@@ -2353,15 +2397,15 @@ public class SharePartitionManagerTest {
// Send 3 requests for share fetch for same share partition.
CompletableFuture<Map<TopicIdPartition, ShareFetchResponseData.PartitionData>> future1 =
sharePartitionManager.fetchMessages(groupId, memberId, FETCH_PARAMS, BATCH_OPTIMIZED, 0,
sharePartitionManager.fetchMessages(groupId, memberId, FETCH_PARAMS, BATCH_OPTIMIZED,
MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
CompletableFuture<Map<TopicIdPartition, ShareFetchResponseData.PartitionData>> future2 =
sharePartitionManager.fetchMessages(groupId, memberId, FETCH_PARAMS, BATCH_OPTIMIZED, 0,
sharePartitionManager.fetchMessages(groupId, memberId, FETCH_PARAMS, BATCH_OPTIMIZED,
MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
CompletableFuture<Map<TopicIdPartition, ShareFetchResponseData.PartitionData>> future3 =
sharePartitionManager.fetchMessages(groupId, memberId, FETCH_PARAMS, BATCH_OPTIMIZED, 0,
sharePartitionManager.fetchMessages(groupId, memberId, FETCH_PARAMS, BATCH_OPTIMIZED,
MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
Mockito.verify(sp0, times(3)).maybeInitialize();
@@ -2420,7 +2464,7 @@ public class SharePartitionManagerTest {
// Return LeaderNotAvailableException to simulate initialization failure.
when(sp0.maybeInitialize()).thenReturn(CompletableFuture.failedFuture(new LeaderNotAvailableException("Leader not available")));
CompletableFuture<Map<TopicIdPartition, ShareFetchResponseData.PartitionData>> future =
sharePartitionManager.fetchMessages(groupId, memberId, FETCH_PARAMS, BATCH_OPTIMIZED, 0,
sharePartitionManager.fetchMessages(groupId, memberId, FETCH_PARAMS, BATCH_OPTIMIZED,
MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
TestUtils.waitForCondition(
future::isDone,
@@ -2436,7 +2480,7 @@ public class SharePartitionManagerTest {
// Return IllegalStateException to simulate initialization failure.
when(sp0.maybeInitialize()).thenReturn(CompletableFuture.failedFuture(new IllegalStateException("Illegal state")));
future = sharePartitionManager.fetchMessages(groupId, memberId, FETCH_PARAMS, BATCH_OPTIMIZED, 0,
future = sharePartitionManager.fetchMessages(groupId, memberId, FETCH_PARAMS, BATCH_OPTIMIZED,
MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
TestUtils.waitForCondition(
future::isDone,
@@ -2450,7 +2494,7 @@ public class SharePartitionManagerTest {
partitionCache.put(new SharePartitionKey(groupId, tp0), sp0);
// Return CoordinatorNotAvailableException to simulate initialization failure.
when(sp0.maybeInitialize()).thenReturn(CompletableFuture.failedFuture(new CoordinatorNotAvailableException("Coordinator not available")));
future = sharePartitionManager.fetchMessages(groupId, memberId, FETCH_PARAMS, BATCH_OPTIMIZED, 0,
future = sharePartitionManager.fetchMessages(groupId, memberId, FETCH_PARAMS, BATCH_OPTIMIZED,
MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
TestUtils.waitForCondition(
future::isDone,
@@ -2464,7 +2508,7 @@ public class SharePartitionManagerTest {
partitionCache.put(new SharePartitionKey(groupId, tp0), sp0);
// Return InvalidRequestException to simulate initialization failure.
when(sp0.maybeInitialize()).thenReturn(CompletableFuture.failedFuture(new InvalidRequestException("Invalid request")));
future = sharePartitionManager.fetchMessages(groupId, memberId, FETCH_PARAMS, BATCH_OPTIMIZED, 0,
future = sharePartitionManager.fetchMessages(groupId, memberId, FETCH_PARAMS, BATCH_OPTIMIZED,
MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
TestUtils.waitForCondition(
future::isDone,
@@ -2478,7 +2522,7 @@ public class SharePartitionManagerTest {
partitionCache.put(new SharePartitionKey(groupId, tp0), sp0);
// Return FencedStateEpochException to simulate initialization failure.
when(sp0.maybeInitialize()).thenReturn(CompletableFuture.failedFuture(new FencedStateEpochException("Fenced state epoch")));
future = sharePartitionManager.fetchMessages(groupId, memberId, FETCH_PARAMS, BATCH_OPTIMIZED, 0,
future = sharePartitionManager.fetchMessages(groupId, memberId, FETCH_PARAMS, BATCH_OPTIMIZED,
MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
TestUtils.waitForCondition(
future::isDone,
@@ -2492,7 +2536,7 @@ public class SharePartitionManagerTest {
partitionCache.put(new SharePartitionKey(groupId, tp0), sp0);
// Return NotLeaderOrFollowerException to simulate initialization failure.
when(sp0.maybeInitialize()).thenReturn(CompletableFuture.failedFuture(new NotLeaderOrFollowerException("Not leader or follower")));
future = sharePartitionManager.fetchMessages(groupId, memberId, FETCH_PARAMS, BATCH_OPTIMIZED, 0,
future = sharePartitionManager.fetchMessages(groupId, memberId, FETCH_PARAMS, BATCH_OPTIMIZED,
MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
TestUtils.waitForCondition(
future::isDone,
@@ -2506,7 +2550,7 @@ public class SharePartitionManagerTest {
partitionCache.put(new SharePartitionKey(groupId, tp0), sp0);
// Return RuntimeException to simulate initialization failure.
when(sp0.maybeInitialize()).thenReturn(CompletableFuture.failedFuture(new RuntimeException("Runtime exception")));
future = sharePartitionManager.fetchMessages(groupId, memberId, FETCH_PARAMS, BATCH_OPTIMIZED, 0,
future = sharePartitionManager.fetchMessages(groupId, memberId, FETCH_PARAMS, BATCH_OPTIMIZED,
MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
TestUtils.waitForCondition(
future::isDone,
@@ -2541,7 +2585,7 @@ public class SharePartitionManagerTest {
.build();
CompletableFuture<Map<TopicIdPartition, ShareFetchResponseData.PartitionData>> future =
sharePartitionManager.fetchMessages(groupId, Uuid.randomUuid().toString(), FETCH_PARAMS, BATCH_OPTIMIZED, 0,
sharePartitionManager.fetchMessages(groupId, Uuid.randomUuid().toString(), FETCH_PARAMS, BATCH_OPTIMIZED,
MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
TestUtils.waitForCondition(
future::isDone,
@@ -2582,7 +2626,7 @@ public class SharePartitionManagerTest {
// Validate when exception is thrown.
CompletableFuture<Map<TopicIdPartition, ShareFetchResponseData.PartitionData>> future =
sharePartitionManager.fetchMessages(groupId, Uuid.randomUuid().toString(), FETCH_PARAMS, BATCH_OPTIMIZED, 0,
sharePartitionManager.fetchMessages(groupId, Uuid.randomUuid().toString(), FETCH_PARAMS, BATCH_OPTIMIZED,
MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
TestUtils.waitForCondition(
future::isDone,
@@ -2592,7 +2636,7 @@ public class SharePartitionManagerTest {
assertTrue(partitionCache.isEmpty());
// Validate when partition is not leader.
future = sharePartitionManager.fetchMessages(groupId, Uuid.randomUuid().toString(), FETCH_PARAMS, BATCH_OPTIMIZED, 0,
future = sharePartitionManager.fetchMessages(groupId, Uuid.randomUuid().toString(), FETCH_PARAMS, BATCH_OPTIMIZED,
MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
TestUtils.waitForCondition(
future::isDone,
@@ -2660,7 +2704,7 @@ public class SharePartitionManagerTest {
// Validate when exception is thrown.
CompletableFuture<Map<TopicIdPartition, ShareFetchResponseData.PartitionData>> future =
sharePartitionManager.fetchMessages(groupId, Uuid.randomUuid().toString(), FETCH_PARAMS, BATCH_OPTIMIZED, 0,
sharePartitionManager.fetchMessages(groupId, Uuid.randomUuid().toString(), FETCH_PARAMS, BATCH_OPTIMIZED,
MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
assertTrue(future.isDone());
assertFalse(future.isCompletedExceptionally());
@@ -2718,7 +2762,7 @@ public class SharePartitionManagerTest {
.build();
CompletableFuture<Map<TopicIdPartition, ShareFetchResponseData.PartitionData>> future =
sharePartitionManager.fetchMessages(groupId, memberId, FETCH_PARAMS, BATCH_OPTIMIZED, 0,
sharePartitionManager.fetchMessages(groupId, memberId, FETCH_PARAMS, BATCH_OPTIMIZED,
MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
validateShareFetchFutureException(future, tp0, Errors.UNKNOWN_SERVER_ERROR, "Exception");
// Verify that the share partition is still in the cache on exception.
@@ -2727,7 +2771,7 @@ public class SharePartitionManagerTest {
// Throw NotLeaderOrFollowerException from replica manager fetch which should evict instance from the cache.
doThrow(new NotLeaderOrFollowerException("Leader exception")).when(mockReplicaManager).readFromLog(any(), any(), any(ReplicaQuota.class), anyBoolean());
future = sharePartitionManager.fetchMessages(groupId, memberId, FETCH_PARAMS, BATCH_OPTIMIZED, 0,
future = sharePartitionManager.fetchMessages(groupId, memberId, FETCH_PARAMS, BATCH_OPTIMIZED,
MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
validateShareFetchFutureException(future, tp0, Errors.NOT_LEADER_OR_FOLLOWER, "Leader exception");
assertTrue(partitionCache.isEmpty());
@@ -2780,7 +2824,7 @@ public class SharePartitionManagerTest {
.build();
CompletableFuture<Map<TopicIdPartition, ShareFetchResponseData.PartitionData>> future =
sharePartitionManager.fetchMessages(groupId, memberId, FETCH_PARAMS, BATCH_OPTIMIZED, 0,
sharePartitionManager.fetchMessages(groupId, memberId, FETCH_PARAMS, BATCH_OPTIMIZED,
MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
validateShareFetchFutureException(future, tp0, Errors.FENCED_STATE_EPOCH, "Fenced exception");
// Verify that tp1 is still in the cache on exception.
@@ -2795,7 +2839,7 @@ public class SharePartitionManagerTest {
// Throw FencedStateEpochException from replica manager fetch which should evict instance from the cache.
doThrow(new FencedStateEpochException("Fenced exception again")).when(mockReplicaManager).readFromLog(any(), any(), any(ReplicaQuota.class), anyBoolean());
future = sharePartitionManager.fetchMessages(groupId, memberId, FETCH_PARAMS, BATCH_OPTIMIZED, 0,
future = sharePartitionManager.fetchMessages(groupId, memberId, FETCH_PARAMS, BATCH_OPTIMIZED,
MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
validateShareFetchFutureException(future, List.of(tp0, tp1), Errors.FENCED_STATE_EPOCH, "Fenced exception again");
assertTrue(partitionCache.isEmpty());
@@ -2828,7 +2872,7 @@ public class SharePartitionManagerTest {
.build();
CompletableFuture<Map<TopicIdPartition, PartitionData>> future = sharePartitionManager.fetchMessages(
groupId, memberId, FETCH_PARAMS, BATCH_OPTIMIZED, 0, MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
groupId, memberId, FETCH_PARAMS, BATCH_OPTIMIZED, MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
assertTrue(future.isDone());
// Validate that the listener is registered.
verify(mockReplicaManager, times(2)).maybeAddListener(any(), any());
@@ -2881,59 +2925,24 @@ public class SharePartitionManagerTest {
}
@Test
public void testFetchMessagesRotatePartitions() {
public void testFetchMessagesPreservesPartitionOrder() {
String groupId = "grp";
Uuid memberId1 = Uuid.randomUuid();
TopicIdPartition tp0 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo", 0));
TopicIdPartition tp1 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo", 1));
TopicIdPartition tp2 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("bar", 0));
TopicIdPartition tp3 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("bar", 1));
TopicIdPartition tp4 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo", 2));
TopicIdPartition tp5 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("bar", 2));
TopicIdPartition tp6 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo", 3));
List<TopicIdPartition> topicIdPartitions = List.of(tp0, tp1, tp2, tp3, tp4, tp5, tp6);
List<TopicIdPartition> topicIdPartitions = List.of(tp0, tp1, tp2, tp3);
sharePartitionManager = Mockito.spy(SharePartitionManagerBuilder.builder().withBrokerTopicStats(brokerTopicStats).build());
// Capture the arguments passed to processShareFetch.
ArgumentCaptor<ShareFetch> captor = ArgumentCaptor.forClass(ShareFetch.class);
sharePartitionManager.fetchMessages(groupId, memberId1.toString(), FETCH_PARAMS, BATCH_OPTIMIZED, 0,
// The partitions are fetched in the order of the share session.
sharePartitionManager.fetchMessages(groupId, memberId1.toString(), FETCH_PARAMS, BATCH_OPTIMIZED,
MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
verify(sharePartitionManager, times(1)).processShareFetch(captor.capture());
// Verify the partitions rotation, no rotation.
ShareFetch resultShareFetch = captor.getValue();
validateRotatedListEquals(resultShareFetch.topicIdPartitions(), topicIdPartitions, 0);
// Single rotation.
sharePartitionManager.fetchMessages(groupId, memberId1.toString(), FETCH_PARAMS, BATCH_OPTIMIZED, 1,
MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
verify(sharePartitionManager, times(2)).processShareFetch(captor.capture());
// Verify the partitions rotation, rotate by 1.
resultShareFetch = captor.getValue();
validateRotatedListEquals(topicIdPartitions, resultShareFetch.topicIdPartitions(), 1);
// Rotation by 3, less that the number of partitions.
sharePartitionManager.fetchMessages(groupId, memberId1.toString(), FETCH_PARAMS, BATCH_OPTIMIZED, 3,
MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
verify(sharePartitionManager, times(3)).processShareFetch(captor.capture());
// Verify the partitions rotation, rotate by 3.
resultShareFetch = captor.getValue();
validateRotatedListEquals(topicIdPartitions, resultShareFetch.topicIdPartitions(), 3);
// Rotation by 12, more than the number of partitions.
sharePartitionManager.fetchMessages(groupId, memberId1.toString(), FETCH_PARAMS, BATCH_OPTIMIZED, 12,
MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
verify(sharePartitionManager, times(4)).processShareFetch(captor.capture());
// Verify the partitions rotation, rotate by 5 (12 % 7).
resultShareFetch = captor.getValue();
validateRotatedListEquals(topicIdPartitions, resultShareFetch.topicIdPartitions(), 5);
// Rotation by Integer.MAX_VALUE, boundary test.
sharePartitionManager.fetchMessages(groupId, memberId1.toString(), FETCH_PARAMS, BATCH_OPTIMIZED, Integer.MAX_VALUE,
MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
verify(sharePartitionManager, times(5)).processShareFetch(captor.capture());
// Verify the partitions rotation, rotate by 1 (2147483647 % 7).
resultShareFetch = captor.getValue();
validateRotatedListEquals(topicIdPartitions, resultShareFetch.topicIdPartitions(), 1);
assertEquals(topicIdPartitions, captor.getValue().topicIdPartitions());
}
@Test
@@ -3328,6 +3337,26 @@ public class SharePartitionManagerTest {
long failedShareAcknowledgementRequestCount
) { }
/**
* Builds the response data in the given fetch order, with records for the given partitions only.
*/
private LinkedHashMap<TopicIdPartition, ShareFetchResponseData.PartitionData> responseData(
List<TopicIdPartition> fetchOrder,
Set<TopicIdPartition> partitionsWithRecords
) {
LinkedHashMap<TopicIdPartition, ShareFetchResponseData.PartitionData> responseData = new LinkedHashMap<>();
fetchOrder.forEach(tp -> {
ShareFetchResponseData.PartitionData partitionData = new ShareFetchResponseData.PartitionData()
.setPartitionIndex(tp.partition());
if (partitionsWithRecords.contains(tp)) {
partitionData.setRecords(MemoryRecords.withRecords(Compression.NONE,
new SimpleRecord("test-key".getBytes(), "test-value".getBytes())));
}
responseData.put(tp, partitionData);
});
return responseData;
}
static class SharePartitionManagerBuilder {
private final Persister persister = new NoOpStatePersister();
private ReplicaManager replicaManager = mock(ReplicaManager.class);
core/src/test/scala/unit/kafka/server/KafkaApisTest.scala
+51 -50
View File
@@ -2545,7 +2545,7 @@ class KafkaApisTest extends Logging {
// Mock share partition manager responses
val records = memoryRecords(10, 0)
when(sharePartitionManager.fetchMessages(any(), any(), any(), any(), anyInt(), anyInt(), anyInt(), any())).thenReturn(
when(sharePartitionManager.fetchMessages(any(), any(), any(), any(), anyInt(), anyInt(), any())).thenReturn(
CompletableFuture.completedFuture(util.Map.of[TopicIdPartition, ShareFetchResponseData.PartitionData](
new TopicIdPartition(topicId, new TopicPartition(topicName, partitionIndex)),
new ShareFetchResponseData.PartitionData()
@@ -2745,7 +2745,7 @@ class KafkaApisTest extends Logging {
// Mock share partition manager responses for both fetch and acknowledge
val records = memoryRecords(10, 0)
when(sharePartitionManager.fetchMessages(any(), any(), any(), any(), anyInt(), anyInt(), anyInt(), any())).thenReturn(
when(sharePartitionManager.fetchMessages(any(), any(), any(), any(), anyInt(), anyInt(), any())).thenReturn(
CompletableFuture.completedFuture(util.Map.of[TopicIdPartition, ShareFetchResponseData.PartitionData](
new TopicIdPartition(topicId, new TopicPartition(topicName, partitionIndex)),
new ShareFetchResponseData.PartitionData()
@@ -4763,7 +4763,7 @@ class KafkaApisTest extends Logging {
val records = memoryRecords(10, 0)
when(sharePartitionManager.fetchMessages(any(), any(), any(), any(), anyInt(), anyInt(), anyInt(), any())).thenReturn(
when(sharePartitionManager.fetchMessages(any(), any(), any(), any(), anyInt(), anyInt(), any())).thenReturn(
CompletableFuture.completedFuture(util.Map.of[TopicIdPartition, ShareFetchResponseData.PartitionData](
new TopicIdPartition(topicId, new TopicPartition(topicName, partitionIndex)),
new ShareFetchResponseData.PartitionData()
@@ -4829,7 +4829,7 @@ class KafkaApisTest extends Logging {
val records = memoryRecords(5, 0)
val ackErrorMessage = "custom acknowledgement error"
when(sharePartitionManager.fetchMessages(any(), any(), any(), any(), anyInt(), anyInt(), anyInt(), any())).thenReturn(
when(sharePartitionManager.fetchMessages(any(), any(), any(), any(), anyInt(), anyInt(), any())).thenReturn(
CompletableFuture.completedFuture(util.Map.of[TopicIdPartition, ShareFetchResponseData.PartitionData](
new TopicIdPartition(topicId, new TopicPartition(topicName, partitionIndex)),
new ShareFetchResponseData.PartitionData()
@@ -4911,7 +4911,7 @@ class KafkaApisTest extends Logging {
val records = memoryRecords(10, 0)
when(sharePartitionManager.fetchMessages(any(), any(), any(), any(), anyInt(), anyInt(), anyInt(), any())).thenReturn(
when(sharePartitionManager.fetchMessages(any(), any(), any(), any(), anyInt(), anyInt(), any())).thenReturn(
CompletableFuture.completedFuture(util.Map.of[TopicIdPartition, ShareFetchResponseData.PartitionData](
new TopicIdPartition(topicId, new TopicPartition(topicName, partitionIndex)),
new ShareFetchResponseData.PartitionData()
@@ -5012,7 +5012,7 @@ class KafkaApisTest extends Logging {
val records = memoryRecords(10, 0)
when(sharePartitionManager.fetchMessages(any(), any(), any(), any(), anyInt(), anyInt(), anyInt(), any())).thenReturn(
when(sharePartitionManager.fetchMessages(any(), any(), any(), any(), anyInt(), anyInt(), any())).thenReturn(
CompletableFuture.completedFuture(util.Map.of[TopicIdPartition, ShareFetchResponseData.PartitionData](
new TopicIdPartition(topicId, new TopicPartition(topicName, partitionIndex)),
new ShareFetchResponseData.PartitionData()
@@ -5110,7 +5110,7 @@ class KafkaApisTest extends Logging {
addTopicToMetadataCache(topicName, 1, topicId = topicId)
val memberId: String = Uuid.randomUuid().toString
when(sharePartitionManager.fetchMessages(any(), any(), any(), any(), anyInt(), anyInt(), anyInt(), any())).thenReturn(
when(sharePartitionManager.fetchMessages(any(), any(), any(), any(), anyInt(), anyInt(), any())).thenReturn(
CompletableFuture.failedFuture[util.Map[TopicIdPartition, ShareFetchResponseData.PartitionData]](Errors.UNKNOWN_SERVER_ERROR.exception())
)
@@ -5156,7 +5156,7 @@ class KafkaApisTest extends Logging {
val records = memoryRecords(10, 0)
when(sharePartitionManager.fetchMessages(any(), any(), any(), any(), anyInt(), anyInt(), anyInt(), any())).thenReturn(
when(sharePartitionManager.fetchMessages(any(), any(), any(), any(), anyInt(), anyInt(), any())).thenReturn(
CompletableFuture.completedFuture(util.Map.of[TopicIdPartition, ShareFetchResponseData.PartitionData](
new TopicIdPartition(topicId, new TopicPartition(topicName, partitionIndex)),
new ShareFetchResponseData.PartitionData()
@@ -5227,7 +5227,7 @@ class KafkaApisTest extends Logging {
val groupId = "group"
when(sharePartitionManager.fetchMessages(any(), any(), any(), any(), anyInt(), anyInt(), anyInt(), any())).thenReturn(
when(sharePartitionManager.fetchMessages(any(), any(), any(), any(), anyInt(), anyInt(), any())).thenReturn(
CompletableFuture.failedFuture[util.Map[TopicIdPartition, ShareFetchResponseData.PartitionData]](Errors.UNKNOWN_SERVER_ERROR.exception())
)
@@ -5334,7 +5334,7 @@ class KafkaApisTest extends Logging {
val records = MemoryRecords.EMPTY
when(sharePartitionManager.fetchMessages(any(), any(), any(), any(), anyInt(), anyInt(), anyInt(), any())).thenReturn(
when(sharePartitionManager.fetchMessages(any(), any(), any(), any(), anyInt(), anyInt(), any())).thenReturn(
CompletableFuture.completedFuture(util.Map.of[TopicIdPartition, ShareFetchResponseData.PartitionData](
new TopicIdPartition(topicId, new TopicPartition(topicName, partitionIndex)),
new ShareFetchResponseData.PartitionData()
@@ -5394,7 +5394,7 @@ class KafkaApisTest extends Logging {
val groupId = "group"
val records = memoryRecords(10, 0)
when(sharePartitionManager.fetchMessages(any(), any(), any(), any(), anyInt(), anyInt(), anyInt(), any())).thenReturn(
when(sharePartitionManager.fetchMessages(any(), any(), any(), any(), anyInt(), anyInt(), any())).thenReturn(
CompletableFuture.completedFuture(util.Map.of[TopicIdPartition, ShareFetchResponseData.PartitionData](
new TopicIdPartition(topicId, partitionIndex, topicName),
new ShareFetchResponseData.PartitionData()
@@ -5481,7 +5481,7 @@ class KafkaApisTest extends Logging {
val groupId = "group"
val records = memoryRecords(10, 0)
when(sharePartitionManager.fetchMessages(any(), any(), any(), any(), anyInt(), anyInt(), anyInt(), any())).thenReturn(
when(sharePartitionManager.fetchMessages(any(), any(), any(), any(), anyInt(), anyInt(), any())).thenReturn(
CompletableFuture.completedFuture(util.Map.of[TopicIdPartition, ShareFetchResponseData.PartitionData](
new TopicIdPartition(topicId, new TopicPartition(topicName, partitionIndex)),
new ShareFetchResponseData.PartitionData()
@@ -5600,7 +5600,7 @@ class KafkaApisTest extends Logging {
val records2 = memoryRecords(10, 10)
val records3 = memoryRecords(10, 20)
when(sharePartitionManager.fetchMessages(any(), any(), any(), any(), anyInt(), anyInt(), anyInt(), any())).thenReturn(
when(sharePartitionManager.fetchMessages(any(), any(), any(), any(), anyInt(), anyInt(), any())).thenReturn(
CompletableFuture.completedFuture(util.Map.of[TopicIdPartition, ShareFetchResponseData.PartitionData](
new TopicIdPartition(topicId, new TopicPartition(topicName, partitionIndex)),
new ShareFetchResponseData.PartitionData()
@@ -5823,7 +5823,7 @@ class KafkaApisTest extends Logging {
val groupId = "group"
when(sharePartitionManager.fetchMessages(any(), any(), any(), any(), anyInt(), anyInt(), anyInt(), any())).thenReturn(
when(sharePartitionManager.fetchMessages(any(), any(), any(), any(), anyInt(), anyInt(), any())).thenReturn(
CompletableFuture.completedFuture(util.Map.of[TopicIdPartition, ShareFetchResponseData.PartitionData](
new TopicIdPartition(topicId1, new TopicPartition(topicName1, 0)),
new ShareFetchResponseData.PartitionData()
@@ -6280,7 +6280,7 @@ class KafkaApisTest extends Logging {
val tp2 = new TopicIdPartition(topicId2, new TopicPartition(topicName2, 0))
val tp3 = new TopicIdPartition(topicId2, new TopicPartition(topicName2, 1))
when(sharePartitionManager.fetchMessages(any(), any(), any(), any(), anyInt(), anyInt(), anyInt(), any())).thenReturn(
when(sharePartitionManager.fetchMessages(any(), any(), any(), any(), anyInt(), anyInt(), any())).thenReturn(
CompletableFuture.completedFuture(util.Map.of[TopicIdPartition, ShareFetchResponseData.PartitionData](
tp1,
new ShareFetchResponseData.PartitionData()
@@ -6357,7 +6357,6 @@ class KafkaApisTest extends Logging {
val fetchResult: Map[TopicIdPartition, ShareFetchResponseData.PartitionData] =
kafkaApis.handleFetchFromShareFetchRequest(
request,
0,
erroneousAndValidPartitionData,
sharePartitionManager,
authorizedTopics
@@ -6426,7 +6425,7 @@ class KafkaApisTest extends Logging {
val tp2 = new TopicIdPartition(topicId1, new TopicPartition(topicName1, 1))
val tp3 = new TopicIdPartition(topicId2, new TopicPartition(topicName2, 0))
when(sharePartitionManager.fetchMessages(any(), any(), any(), any(), anyInt(), anyInt(), anyInt(), any())).thenReturn(
when(sharePartitionManager.fetchMessages(any(), any(), any(), any(), anyInt(), anyInt(), any())).thenReturn(
CompletableFuture.completedFuture(util.Map.of[TopicIdPartition, ShareFetchResponseData.PartitionData](
tp1,
new ShareFetchResponseData.PartitionData()
@@ -6490,7 +6489,6 @@ class KafkaApisTest extends Logging {
val fetchResult: Map[TopicIdPartition, ShareFetchResponseData.PartitionData] =
kafkaApis.handleFetchFromShareFetchRequest(
request,
0,
erroneousAndValidPartitionData,
sharePartitionManager,
authorizedTopics
@@ -6553,7 +6551,7 @@ class KafkaApisTest extends Logging {
val tp2 = new TopicIdPartition(topicId2, new TopicPartition(topicName2, 0))
val tp3 = new TopicIdPartition(topicId2, new TopicPartition(topicName2, 1))
when(sharePartitionManager.fetchMessages(any(), any(), any(), any(), anyInt(), anyInt(), anyInt(), any())).thenReturn(
when(sharePartitionManager.fetchMessages(any(), any(), any(), any(), anyInt(), anyInt(), any())).thenReturn(
CompletableFuture.completedFuture(util.Map.of[TopicIdPartition, ShareFetchResponseData.PartitionData](
tp1,
new ShareFetchResponseData.PartitionData()
@@ -6615,7 +6613,6 @@ class KafkaApisTest extends Logging {
val fetchResult: Map[TopicIdPartition, ShareFetchResponseData.PartitionData] =
kafkaApis.handleFetchFromShareFetchRequest(
request,
0,
erroneousAndValidPartitionData,
sharePartitionManager,
authorizedTopics
@@ -6689,31 +6686,32 @@ class KafkaApisTest extends Logging {
val tp3 = new TopicIdPartition(topicId2, new TopicPartition(topicName2, 1))
val tp4 = new TopicIdPartition(topicId3, new TopicPartition(topicName3, 0))
when(sharePartitionManager.fetchMessages(any(), any(), any(), any(), anyInt(), anyInt(), anyInt(), any())).thenReturn(
CompletableFuture.completedFuture(util.Map.of[TopicIdPartition, ShareFetchResponseData.PartitionData](
tp2,
new ShareFetchResponseData.PartitionData()
.setPartitionIndex(0)
.setErrorCode(Errors.NONE.code)
.setRecords(records1)
.setAcquiredRecords(new util.ArrayList(util.List.of(
new ShareFetchResponseData.AcquiredRecords()
.setFirstOffset(0)
.setLastOffset(9)
.setDeliveryCount(1)
))),
tp3,
new ShareFetchResponseData.PartitionData()
.setPartitionIndex(1)
.setErrorCode(Errors.NONE.code)
.setRecords(records2)
.setAcquiredRecords(new util.ArrayList(util.List.of(
new ShareFetchResponseData.AcquiredRecords()
.setFirstOffset(0)
.setLastOffset(19)
.setDeliveryCount(1)
)))
))
// The partitions are fetched in the order tp3, tp2 which differs from the request order.
val fetchedPartitions = new util.LinkedHashMap[TopicIdPartition, ShareFetchResponseData.PartitionData]
fetchedPartitions.put(tp3,
new ShareFetchResponseData.PartitionData()
.setPartitionIndex(1)
.setErrorCode(Errors.NONE.code)
.setRecords(records2)
.setAcquiredRecords(new util.ArrayList(util.List.of(
new ShareFetchResponseData.AcquiredRecords()
.setFirstOffset(0)
.setLastOffset(19)
.setDeliveryCount(1)
))))
fetchedPartitions.put(tp2,
new ShareFetchResponseData.PartitionData()
.setPartitionIndex(0)
.setErrorCode(Errors.NONE.code)
.setRecords(records1)
.setAcquiredRecords(new util.ArrayList(util.List.of(
new ShareFetchResponseData.AcquiredRecords()
.setFirstOffset(0)
.setLastOffset(9)
.setDeliveryCount(1)
))))
when(sharePartitionManager.fetchMessages(any(), any(), any(), any(), anyInt(), anyInt(), any())).thenReturn(
CompletableFuture.completedFuture(fetchedPartitions)
)
val erroneousPartitions: util.Map[TopicIdPartition, ShareFetchResponseData.PartitionData] = new util.HashMap()
@@ -6761,13 +6759,16 @@ class KafkaApisTest extends Logging {
val fetchResult: Map[TopicIdPartition, ShareFetchResponseData.PartitionData] =
kafkaApis.handleFetchFromShareFetchRequest(
request,
0,
erroneousAndValidPartitionData,
sharePartitionManager,
authorizedTopics
).get()
assertEquals(4, fetchResult.size)
// The share session re-orders its partitions from the result, hence the result must keep the order in
// which the partitions were fetched, with the erroneous partitions appended.
assertEquals(List(tp3, tp2), fetchResult.keys.toList.take(2))
assertEquals(Set(tp1, tp4), fetchResult.keys.toList.drop(2).toSet)
assertTrue(fetchResult.contains(tp1))
val partitionData1: PartitionData = fetchResult.getOrElse(tp1, null)
@@ -6858,7 +6859,7 @@ class KafkaApisTest extends Logging {
val groupId = "group"
when(sharePartitionManager.fetchMessages(any(), any(), any(), any(), anyInt(), anyInt(), anyInt(), any())).thenReturn(
when(sharePartitionManager.fetchMessages(any(), any(), any(), any(), anyInt(), anyInt(), any())).thenReturn(
CompletableFuture.completedFuture(util.Map.of[TopicIdPartition, ShareFetchResponseData.PartitionData](
new TopicIdPartition(topicId, new TopicPartition(topicName, partitionIndex)),
new ShareFetchResponseData.PartitionData()
@@ -6992,7 +6993,7 @@ class KafkaApisTest extends Logging {
val groupId = "group"
when(sharePartitionManager.fetchMessages(any(), any(), any(), any(), anyInt(), anyInt(), anyInt(), any())).thenReturn(
when(sharePartitionManager.fetchMessages(any(), any(), any(), any(), anyInt(), anyInt(), any())).thenReturn(
CompletableFuture.completedFuture(util.Map.of[TopicIdPartition, ShareFetchResponseData.PartitionData](
new TopicIdPartition(topicId, new TopicPartition(topicName, partitionIndex)),
new ShareFetchResponseData.PartitionData()
@@ -7103,7 +7104,7 @@ class KafkaApisTest extends Logging {
assertEquals(expectedRecords, topicResponse.partitions.get(0).records)
assertEquals(0, topicResponse.partitions.get(0).acquiredRecords.size())
// fetchMessages only called once for 1st ShareFetch.
verify(sharePartitionManager, times(1)).fetchMessages(any(), any(), any(), any(), anyInt(), anyInt(), anyInt(), any())
verify(sharePartitionManager, times(1)).fetchMessages(any(), any(), any(), any(), anyInt(), anyInt(), any())
}
@ParameterizedTest
@@ -14624,7 +14625,7 @@ class KafkaApisTest extends Logging {
val groupId = "group"
when(sharePartitionManager.fetchMessages(any(), any(), any(), any(), anyInt(), anyInt(), anyInt(), any())).thenReturn(
when(sharePartitionManager.fetchMessages(any(), any(), any(), any(), anyInt(), anyInt(), any())).thenReturn(
CompletableFuture.completedFuture(util.Map.of[TopicIdPartition, ShareFetchResponseData.PartitionData](
new TopicIdPartition(topicId, new TopicPartition(topicName, partitionIndex)),
new ShareFetchResponseData.PartitionData()
server/src/main/java/org/apache/kafka/server/share/context/ShareSessionContext.java
+39 -6
View File
@@ -26,6 +26,8 @@ import org.apache.kafka.common.requests.ShareFetchResponse;
import org.apache.kafka.common.requests.ShareRequestMetadata;
import org.apache.kafka.server.share.CachedSharePartition;
import org.apache.kafka.server.share.ErroneousAndValidPartitionData;
import org.apache.kafka.server.share.fetch.PartitionRotateStrategy;
import org.apache.kafka.server.share.fetch.PartitionRotateStrategy.StrategyType;
import org.apache.kafka.server.share.session.ShareSession;
import org.slf4j.Logger;
@@ -143,12 +145,6 @@ public class ShareSessionContext extends ShareFetchContext {
boolean mustRespond = cachedPart.maybeUpdateResponseData(respData, updateShareContextAndRemoveUnselected);
if (mustRespond) {
nextElement = element;
if (updateShareContextAndRemoveUnselected && ShareFetchResponse.recordsSize(respData) > 0) {
// Session.partitionMap is of type ImplicitLinkedHashCollection<> which tracks the order of insertion of elements.
// Since, we are updating an element in this case, we need to perform a remove and then a mustAdd to maintain the correct order
session.partitionMap().remove(cachedPart);
session.partitionMap().mustAdd(cachedPart);
}
} else {
if (updateShareContextAndRemoveUnselected) {
iterator.remove();
@@ -203,6 +199,7 @@ public class ShareSessionContext extends ShareFetchContext {
session.key(), expectedEpoch, sessionEpoch);
return ShareFetchResponse.of(Errors.INVALID_SHARE_SESSION_EPOCH, 0, new LinkedHashMap<>(), List.of(), 0);
}
updatePartitionOrder(updates);
// Iterate over the update list using PartitionIterator. This will prune updates which don't need to be sent
Iterator<Map.Entry<TopicIdPartition, ShareFetchResponseData.PartitionData>> partitionIterator = new PartitionIterator(
updates.entrySet().iterator(), true);
@@ -215,6 +212,42 @@ public class ShareSessionContext extends ShareFetchContext {
}
}
/**
* Re-orders the partitions of the session as per the records returned in the response. The subsequent
* fetch reads the partitions in the order of the session, hence the partitions which returned records
* are moved to the end so the partitions which did not return records are read first. Among the partitions
* which returned records, the partition read first takes the largest share of the records, hence they are
* appended in the order decided by the {@link PartitionRotateStrategy} so the first position rotates.
* <p>
* The updates are expected in the fetch order of the partitions.
*
* @param updates The response data for the partitions, in fetch order.
*/
private void updatePartitionOrder(LinkedHashMap<TopicIdPartition, ShareFetchResponseData.PartitionData> updates) {
List<TopicIdPartition> servedPartitions = new ArrayList<>();
updates.forEach((topicIdPartition, partitionData) -> {
if (ShareFetchResponse.recordsSize(partitionData) > 0) {
servedPartitions.add(topicIdPartition);
}
});
if (servedPartitions.isEmpty()) {
return;
}
List<TopicIdPartition> rotatedPartitions = PartitionRotateStrategy.type(StrategyType.ROUND_ROBIN)
.rotate(servedPartitions);
synchronized (session) {
rotatedPartitions.forEach(topicIdPartition -> {
CachedSharePartition cachedPart = session.partitionMap().find(new CachedSharePartition(topicIdPartition));
if (cachedPart != null) {
// Session.partitionMap is of type ImplicitLinkedHashCollection which tracks the order of
// insertion of elements, hence a remove and a mustAdd moves the element to the end.
session.partitionMap().remove(cachedPart);
session.partitionMap().mustAdd(cachedPart);
}
});
}
}
@Override
public ErroneousAndValidPartitionData getErroneousAndValidTopicIdPartitions() {
if (!isSubsequent) {
server/src/main/java/org/apache/kafka/server/share/fetch/PartitionRotateStrategy.java
+9 -28
View File
@@ -45,11 +45,10 @@ public interface PartitionRotateStrategy {
* Rotate the partitions based on the strategy.
*
* @param topicIdPartitions the topicIdPartitions to rotate
* @param metadata the metadata to rotate
*
* @return the rotated topicIdPartitions
*/
List<TopicIdPartition> rotate(List<TopicIdPartition> topicIdPartitions, PartitionRotateMetadata metadata);
List<TopicIdPartition> rotate(List<TopicIdPartition> topicIdPartitions);
static PartitionRotateStrategy type(StrategyType type) {
return switch (type) {
@@ -58,41 +57,23 @@ public interface PartitionRotateStrategy {
}
/**
* Rotate the partitions based on the round-robin strategy.
* Rotate the partitions based on the round-robin strategy: the partitions move left by one, hence the first
* partition becomes the last and every other partition moves one position towards the front.
*
* @param topicIdPartitions the topicIdPartitions to rotate
* @param metadata the metadata to rotate
*
* @return the rotated topicIdPartitions
*/
static List<TopicIdPartition> rotateRoundRobin(
List<TopicIdPartition> topicIdPartitions,
PartitionRotateMetadata metadata
) {
if (topicIdPartitions.isEmpty() || topicIdPartitions.size() == 1 || metadata.sessionEpoch < 1) {
// No need to rotate the partitions if there are no partitions, only one partition or the
// session epoch is initial or final.
static List<TopicIdPartition> rotateRoundRobin(List<TopicIdPartition> topicIdPartitions) {
if (topicIdPartitions.size() <= 1) {
// No need to rotate the partitions if there are no partitions or only one partition.
return topicIdPartitions;
}
int rotateAt = metadata.sessionEpoch % topicIdPartitions.size();
if (rotateAt == 0) {
// No need to rotate the partitions if the rotateAt is 0.
return topicIdPartitions;
}
// Avoid modifying the original list, create copy.
List<TopicIdPartition> rotatedPartitions = new ArrayList<>(topicIdPartitions);
// Elements from the list should move left by the distance provided i.e. if the original list is [1,2,3],
// and rotation is by 1, then output should be [2,3,1] and not [3,1,2]. Hence, negate the distance here.
Collections.rotate(rotatedPartitions, -1 * rotateAt);
// Elements from the list should move left by one i.e. if the original list is [1,2,3], then the
// output should be [2,3,1] and not [3,1,2]. Hence, negate the distance here.
Collections.rotate(rotatedPartitions, -1);
return rotatedPartitions;
}
/**
* The partition rotate metadata which can be used to store the metadata for the partition rotation.
*
* @param sessionEpoch the share session epoch.
*/
record PartitionRotateMetadata(int sessionEpoch) { }
}
server/src/main/java/org/apache/kafka/server/share/fetch/ShareFetch.java
+2 -1
View File
@@ -28,6 +28,7 @@ import org.apache.kafka.storage.log.metrics.BrokerTopicStats;
import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
@@ -177,7 +178,7 @@ public class ShareFetch {
return;
}
Map<TopicIdPartition, PartitionData> response = new HashMap<>(partitionData);
Map<TopicIdPartition, PartitionData> response = new LinkedHashMap<>(partitionData);
// Add any erroneous partitions to the response.
addErroneousToResponse(response);
future.complete(response);
server/src/test/java/org/apache/kafka/server/share/fetch/PartitionRotateStrategyTest.java
+8 -31
View File
@@ -18,8 +18,6 @@ package org.apache.kafka.server.share.fetch;
import org.apache.kafka.common.TopicIdPartition;
import org.apache.kafka.common.Uuid;
import org.apache.kafka.common.requests.ShareRequestMetadata;
import org.apache.kafka.server.share.fetch.PartitionRotateStrategy.PartitionRotateMetadata;
import org.apache.kafka.server.share.fetch.PartitionRotateStrategy.StrategyType;
import org.junit.jupiter.api.Test;
@@ -38,49 +36,28 @@ public class PartitionRotateStrategyTest {
PartitionRotateStrategy strategy = PartitionRotateStrategy.type(StrategyType.ROUND_ROBIN);
List<TopicIdPartition> partitions = createPartitions(3);
List<TopicIdPartition> result = strategy.rotate(partitions, new PartitionRotateMetadata(1));
// The partition read first moves to the end, the others keep their order.
List<TopicIdPartition> result = strategy.rotate(partitions);
assertEquals(3, result.size());
validateRotatedListEquals(partitions, result, 1);
// Session epoch is greater than the number of partitions.
result = strategy.rotate(partitions, new PartitionRotateMetadata(5));
assertEquals(3, result.size());
// Rotating the result again moves the next partition to the end.
result = strategy.rotate(result);
validateRotatedListEquals(partitions, result, 2);
// Session epoch is at Integer.MAX_VALUE.
result = strategy.rotate(partitions, new PartitionRotateMetadata(Integer.MAX_VALUE));
assertEquals(3, result.size());
validateRotatedListEquals(partitions, result, 1);
// No rotation at same size as epoch.
result = strategy.rotate(partitions, new PartitionRotateMetadata(3));
assertEquals(3, result.size());
validateRotatedListEquals(partitions, result, 0);
}
@Test
public void testRoundRobinStrategyWithSpecialSessionEpochs() {
public void testRoundRobinStrategyWithSinglePartition() {
PartitionRotateStrategy strategy = PartitionRotateStrategy.type(StrategyType.ROUND_ROBIN);
List<TopicIdPartition> partitions = createPartitions(3);
List<TopicIdPartition> result = strategy.rotate(
partitions,
new PartitionRotateMetadata(ShareRequestMetadata.INITIAL_EPOCH));
assertEquals(3, result.size());
validateRotatedListEquals(partitions, result, 0);
result = strategy.rotate(
partitions,
new PartitionRotateMetadata(ShareRequestMetadata.FINAL_EPOCH));
assertEquals(3, result.size());
validateRotatedListEquals(partitions, result, 0);
List<TopicIdPartition> partitions = createPartitions(1);
assertEquals(partitions, strategy.rotate(partitions));
}
@Test
public void testRoundRobinStrategyWithEmptyPartitions() {
PartitionRotateStrategy strategy = PartitionRotateStrategy.type(StrategyType.ROUND_ROBIN);
// Empty partitions.
List<TopicIdPartition> result = strategy.rotate(new ArrayList<>(), new PartitionRotateMetadata(5));
List<TopicIdPartition> result = strategy.rotate(new ArrayList<>());
// The result should be empty.
assertTrue(result.isEmpty());
}
server/src/test/java/org/apache/kafka/server/share/fetch/ShareFetchTest.java
+23
View File
@@ -36,6 +36,7 @@ import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
@@ -172,6 +173,28 @@ public class ShareFetchTest {
assertEquals(1, brokerTopicStats.topicStats("foo").failedShareFetchRequestRate().count());
}
@Test
public void testMaybeCompletePreservesPartitionOrder() {
TopicIdPartition topicIdPartition0 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo", 0));
TopicIdPartition topicIdPartition1 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo", 1));
TopicIdPartition topicIdPartition2 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo", 2));
CompletableFuture<Map<TopicIdPartition, PartitionData>> future = new CompletableFuture<>();
ShareFetch shareFetch = new ShareFetch(mock(FetchParams.class), GROUP_ID, MEMBER_ID, future,
List.of(topicIdPartition0, topicIdPartition1, topicIdPartition2), BATCH_OPTIMIZED, BATCH_SIZE, 100, brokerTopicStats);
// The partitions were read in an order which differs from the request order.
LinkedHashMap<TopicIdPartition, PartitionData> partitionData = new LinkedHashMap<>();
partitionData.put(topicIdPartition1, new PartitionData().setPartitionIndex(1));
partitionData.put(topicIdPartition0, new PartitionData().setPartitionIndex(0));
shareFetch.addErroneous(topicIdPartition2, new RuntimeException());
shareFetch.maybeComplete(partitionData);
// The share session re-orders its partitions from the response, hence the response must keep the
// order in which the partitions were read, with the erroneous partitions appended.
assertEquals(List.of(topicIdPartition1, topicIdPartition0, topicIdPartition2), new ArrayList<>(future.join().keySet()));
}
@Test
public void testMaybeCompleteWithException() {
TopicIdPartition topicIdPartition0 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo", 0));