KAFKA-17367: Share coordinator impl. New merge batches algorithm. [3/N]

checkstyle/import-control-share-coordinator.xml

@@ -62,6 +62,8 @@
             <allow pkg="org.apache.kafka.server.util.timer" />
             <allow pkg="org.apache.kafka.timeline" />
             <allow pkg="org.junit.jupiter.api" />
+            <allow pkg="org.junit.jupiter.params" />
+            <allow pkg="org.junit.jupiter.provider" />
             <allow pkg="org.mockito" />
             <allow pkg="org.slf4j" />
             <subpackage name="generated">

checkstyle/suppressions.xml

@@ -340,6 +340,12 @@
     <suppress checks="NPathComplexity"
               files="CoordinatorRuntime.java"/>
 
+    <!-- share coordinator -->
+    <suppress checks="NPathComplexity"
+              files="ShareCoordinatorShard.java"/>
+    <suppress checks="CyclomaticComplexity"
+              files="ShareCoordinatorShard.java"/>
+
     <!-- storage -->
     <suppress checks="CyclomaticComplexity"
               files="(LogLoader|LogValidator|RemoteLogManagerConfig|RemoteLogManager).java"/>

share-coordinator/src/main/java/org/apache/kafka/coordinator/share/ShareCoordinatorShard.java

@@ -48,19 +48,17 @@
 import org.apache.kafka.server.common.ApiMessageAndVersion;
 import org.apache.kafka.server.config.ShareCoordinatorConfig;
 import org.apache.kafka.server.share.PartitionFactory;
+import org.apache.kafka.server.share.PersisterStateBatch;
 import org.apache.kafka.server.share.SharePartitionKey;
 import org.apache.kafka.timeline.SnapshotRegistry;
 import org.apache.kafka.timeline.TimelineHashMap;
 
 import org.slf4j.Logger;
 
-import java.util.Collection;
 import java.util.Collections;
-import java.util.LinkedHashSet;
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Optional;
-import java.util.Set;
 import java.util.stream.Collectors;
 
 public class ShareCoordinatorShard implements CoordinatorShard<CoordinatorRecord> {
@@ -293,11 +291,11 @@ public CoordinatorResult<WriteShareGroupStateResponseData, CoordinatorRecord> wr
 
             // Since the number of update records for this share part key exceeds snapshotUpdateRecordsPerSnapshot,
             // we should be creating a share snapshot record.
-            List<PersisterOffsetsStateBatch> batchesToAdd = combineStateBatches(
-                shareStateMap.get(key).stateBatchAsSet(),
+            List<PersisterStateBatch> batchesToAdd = StateBatchUtil.combineStateBatches(
+                shareStateMap.get(key).stateBatches(),
                 partitionData.stateBatches().stream()
-                    .map(PersisterOffsetsStateBatch::from)
-                    .collect(Collectors.toCollection(LinkedHashSet::new)),
+                    .map(PersisterStateBatch::from)
+                    .collect(Collectors.toList()),
                 newStartOffset);
 
             recordList = Collections.singletonList(ShareCoordinatorRecordHelpers.newShareSnapshotRecord(
@@ -527,7 +525,7 @@ CoordinatorMetricsShard getMetricsShard() {
     private static ShareGroupOffset merge(ShareGroupOffset soFar, ShareUpdateValue newData) {
         // snapshot epoch should be same as last share snapshot
         // state epoch is not present
-        Set<PersisterOffsetsStateBatch> currentBatches = soFar.stateBatchAsSet();
+        List<PersisterStateBatch> currentBatches = soFar.stateBatches();
         long newStartOffset = newData.startOffset() == -1 ? soFar.startOffset() : newData.startOffset();
         int newLeaderEpoch = newData.leaderEpoch() == -1 ? soFar.leaderEpoch() : newData.leaderEpoch();
 
@@ -536,41 +534,12 @@ private static ShareGroupOffset merge(ShareGroupOffset soFar, ShareUpdateValue n
             .setStateEpoch(soFar.stateEpoch())
             .setStartOffset(newStartOffset)
             .setLeaderEpoch(newLeaderEpoch)
-            .setStateBatches(combineStateBatches(currentBatches, newData.stateBatches().stream()
-                .map(PersisterOffsetsStateBatch::from)
-                .collect(Collectors.toCollection(LinkedHashSet::new)), newStartOffset))
+            .setStateBatches(StateBatchUtil.combineStateBatches(currentBatches, newData.stateBatches().stream()
+                .map(StateBatchUtil::toPersisterStateBatch)
+                .collect(Collectors.toList()), newStartOffset))
             .build();
     }
 
-    /**
-     * Util method which takes in 2 collections containing {@link PersisterOffsetsStateBatch}
-     * and the startOffset.
-     * It removes all batches from the 1st collection which have the same first and last offset
-     * as the batches in 2nd collection. It then creates a final list of batches which contains the
-     * former result and all the batches in the 2nd collection. In set notation (A - B) U B (we prefer batches in B
-     * which have same first and last offset in A).
-     * Finally, it removes any batches where the lastOffset < startOffset, if the startOffset > -1.
-     * @param currentBatch - collection containing current soft state of batches
-     * @param newBatch - collection containing batches in incoming request
-     * @param startOffset - startOffset to consider when removing old batches.
-     * @return List containing combined batches
-     */
-    private static List<PersisterOffsetsStateBatch> combineStateBatches(
-        Collection<PersisterOffsetsStateBatch> currentBatch,
-        Collection<PersisterOffsetsStateBatch> newBatch,
-        long startOffset
-    ) {
-        currentBatch.removeAll(newBatch);
-        List<PersisterOffsetsStateBatch> batchesToAdd = new LinkedList<>(currentBatch);
-        batchesToAdd.addAll(newBatch);
-        // Any batches where the last offset is < the current start offset
-        // are now expired. We should remove them from the persister.
-        if (startOffset != -1) {
-            batchesToAdd.removeIf(batch -> batch.lastOffset() < startOffset);
-        }
-        return batchesToAdd;
-    }
-
     private static ApiMessage messageOrNull(ApiMessageAndVersion apiMessageAndVersion) {
         if (apiMessageAndVersion == null) {
             return null;

share-coordinator/src/main/java/org/apache/kafka/coordinator/share/ShareGroupOffset.java

@@ -20,12 +20,12 @@
 import org.apache.kafka.common.message.WriteShareGroupStateRequestData;
 import org.apache.kafka.coordinator.share.generated.ShareSnapshotValue;
 import org.apache.kafka.coordinator.share.generated.ShareUpdateValue;
+import org.apache.kafka.server.share.PersisterStateBatch;
 
 import java.util.Collections;
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Objects;
-import java.util.Set;
 import java.util.stream.Collectors;
 
 /**
@@ -37,13 +37,13 @@ public class ShareGroupOffset {
     private final int stateEpoch;
     private final int leaderEpoch;
     private final long startOffset;
-    private final List<PersisterOffsetsStateBatch> stateBatches;
+    private final List<PersisterStateBatch> stateBatches;
 
     private ShareGroupOffset(int snapshotEpoch,
                             int stateEpoch,
                             int leaderEpoch,
                             long startOffset,
-                            List<PersisterOffsetsStateBatch> stateBatches) {
+                            List<PersisterStateBatch> stateBatches) {
         this.snapshotEpoch = snapshotEpoch;
         this.stateEpoch = stateEpoch;
         this.leaderEpoch = leaderEpoch;
@@ -67,16 +67,16 @@ public long startOffset() {
         return startOffset;
     }
 
-    public List<PersisterOffsetsStateBatch> stateBatches() {
+    public List<PersisterStateBatch> stateBatches() {
         return Collections.unmodifiableList(stateBatches);
     }
 
-    private static PersisterOffsetsStateBatch toPersisterOffsetsStateBatch(ShareSnapshotValue.StateBatch stateBatch) {
-        return new PersisterOffsetsStateBatch(stateBatch.firstOffset(), stateBatch.lastOffset(), stateBatch.deliveryState(), stateBatch.deliveryCount());
+    private static PersisterStateBatch toPersisterOffsetsStateBatch(ShareSnapshotValue.StateBatch stateBatch) {
+        return new PersisterStateBatch(stateBatch.firstOffset(), stateBatch.lastOffset(), stateBatch.deliveryState(), stateBatch.deliveryCount());
     }
 
-    private static PersisterOffsetsStateBatch toPersisterOffsetsStateBatch(ShareUpdateValue.StateBatch stateBatch) {
-        return new PersisterOffsetsStateBatch(stateBatch.firstOffset(), stateBatch.lastOffset(), stateBatch.deliveryState(), stateBatch.deliveryCount());
+    private static PersisterStateBatch toPersisterOffsetsStateBatch(ShareUpdateValue.StateBatch stateBatch) {
+        return new PersisterStateBatch(stateBatch.firstOffset(), stateBatch.lastOffset(), stateBatch.deliveryState(), stateBatch.deliveryCount());
     }
 
     public static ShareGroupOffset fromRecord(ShareSnapshotValue record) {
@@ -99,11 +99,11 @@ public static ShareGroupOffset fromRequest(WriteShareGroupStateRequestData.Parti
             data.leaderEpoch(),
             data.startOffset(),
             data.stateBatches().stream()
-                .map(PersisterOffsetsStateBatch::from)
+                .map(PersisterStateBatch::from)
                 .collect(Collectors.toList()));
     }
 
-    public Set<PersisterOffsetsStateBatch> stateBatchAsSet() {
+    public LinkedHashSet<PersisterStateBatch> stateBatchAsSet() {
         return new LinkedHashSet<>(stateBatches);
     }
 
@@ -112,7 +112,7 @@ public static class Builder {
         private int stateEpoch;
         private int leaderEpoch;
         private long startOffset;
-        private List<PersisterOffsetsStateBatch> stateBatches;
+        private List<PersisterStateBatch> stateBatches;
 
         public Builder setSnapshotEpoch(int snapshotEpoch) {
             this.snapshotEpoch = snapshotEpoch;
@@ -134,7 +134,7 @@ public Builder setStartOffset(long startOffset) {
             return this;
         }
 
-        public Builder setStateBatches(List<PersisterOffsetsStateBatch> stateBatches) {
+        public Builder setStateBatches(List<PersisterStateBatch> stateBatches) {
             this.stateBatches = stateBatches;
             return this;
         }