KAFKA-12154: Snapshot Loading API

checkstyle/import-control.xml

@@ -286,6 +286,7 @@
     <allow pkg="org.apache.kafka.queue"/>
     <allow pkg="org.apache.kafka.raft"/>
     <allow pkg="org.apache.kafka.shell"/>
+    <allow pkg="org.apache.kafka.snapshot"/>
     <allow pkg="org.jline"/>
     <allow pkg="scala.compat"/>
   </subpackage>

core/src/main/scala/kafka/tools/TestRaftServer.scala

@@ -37,6 +37,7 @@ import org.apache.kafka.common.utils.{Time, Utils}
 import org.apache.kafka.common.{TopicPartition, Uuid, protocol}
 import org.apache.kafka.raft.BatchReader.Batch
 import org.apache.kafka.raft.{BatchReader, RaftClient, RaftConfig, RecordSerde}
+import org.apache.kafka.snapshot.SnapshotReader
 
 import scala.jdk.CollectionConverters._
 
@@ -150,6 +151,7 @@ class TestRaftServer(
     case class HandleClaim(epoch: Int) extends RaftEvent
     case object HandleResign extends RaftEvent
     case class HandleCommit(reader: BatchReader[Array[Byte]]) extends RaftEvent
+    case class HandleSnapshot(reader: SnapshotReader[Array[Byte]]) extends RaftEvent
     case object Shutdown extends RaftEvent
 
     private val eventQueue = new LinkedBlockingDeque[RaftEvent]()
@@ -175,6 +177,10 @@ class TestRaftServer(
       eventQueue.offer(HandleCommit(reader))
     }
 
+    override def handleSnapshot(reader: SnapshotReader[Array[Byte]]): Unit = {
+      eventQueue.offer(HandleSnapshot(reader))
+    }
+
     override def initiateShutdown(): Boolean = {
       val initiated = super.initiateShutdown()
       eventQueue.offer(Shutdown)
@@ -226,7 +232,11 @@ class TestRaftServer(
             reader.close()
           }
 
-        case _ =>
+        case HandleSnapshot(reader) =>
+          // Ignore snapshots; only interested on records appended by this leader
+          reader.close()
+
+        case Shutdown => // Ignore shutdown command
       }
     }
 

raft/src/main/java/org/apache/kafka/raft/BatchReader.java

@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.raft;
 
+import java.util.Collections;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Objects;
@@ -60,19 +61,21 @@ public interface BatchReader<T> extends Iterator<BatchReader.Batch<T>>, AutoClos
     @Override
     void close();
 
-    class Batch<T> {
+    final class Batch<T> implements Iterable<T> {
         private final long baseOffset;
         private final int epoch;
+        private final long lastOffset;
         private final List<T> records;
 
-        public Batch(long baseOffset, int epoch, List<T> records) {
+        private Batch(long baseOffset, int epoch, long lastOffset, List<T> records) {
             this.baseOffset = baseOffset;
             this.epoch = epoch;
+            this.lastOffset = lastOffset;
             this.records = records;
         }
 
         public long lastOffset() {
-            return baseOffset + records.size() - 1;
+            return lastOffset;
         }
 
         public long baseOffset() {
@@ -87,11 +90,17 @@ public int epoch() {
             return epoch;
         }
 
+        @Override
+        public Iterator<T> iterator() {
+            return records.iterator();
+        }
+
         @Override
         public String toString() {
             return "Batch(" +
                 "baseOffset=" + baseOffset +
                 ", epoch=" + epoch +
+                ", lastOffset=" + lastOffset +
                 ", records=" + records +
                 ')';
         }
@@ -110,6 +119,13 @@ public boolean equals(Object o) {
         public int hashCode() {
             return Objects.hash(baseOffset, epoch, records);
         }
-    }
 
+        public static <T> Batch<T> empty(long baseOffset, int epoch, long lastOffset) {
+            return new Batch<>(baseOffset, epoch, lastOffset, Collections.emptyList());
+        }
+
+        public static <T> Batch<T> of(long baseOffset, int epoch, List<T> records) {
+            return new Batch<>(baseOffset, epoch, baseOffset + records.size() - 1, records);
+        }
+    }
 }

raft/src/main/java/org/apache/kafka/raft/FollowerState.java

@@ -159,6 +159,8 @@ public String toString() {
             ", epoch=" + epoch +
             ", leaderId=" + leaderId +
             ", voters=" + voters +
+            ", highWatermark=" + highWatermark +
+            ", fetchingSnapshot=" + fetchingSnapshot +
             ')';
     }
 

raft/src/main/java/org/apache/kafka/raft/KafkaRaftClient.java

@@ -72,6 +72,7 @@
 import org.apache.kafka.raft.internals.ThresholdPurgatory;
 import org.apache.kafka.snapshot.RawSnapshotReader;
 import org.apache.kafka.snapshot.RawSnapshotWriter;
+import org.apache.kafka.snapshot.SnapshotReader;
 import org.apache.kafka.snapshot.SnapshotWriter;
 import org.slf4j.Logger;
 
@@ -137,7 +138,6 @@
  *    than the leader's log start offset. This API is similar to the Fetch API since the snapshot is stored
  *    as FileRecords, but we use {@link UnalignedRecords} in FetchSnapshotResponse because the records
  *    are not necessarily offset-aligned.
- *
  */
 public class KafkaRaftClient<T> implements RaftClient<T> {
     private static final int RETRY_BACKOFF_BASE_MS = 100;
@@ -311,8 +311,18 @@ private void updateListenersProgress(long highWatermark) {
     private void updateListenersProgress(List<ListenerContext> listenerContexts, long highWatermark) {
         for (ListenerContext listenerContext : listenerContexts) {
             listenerContext.nextExpectedOffset().ifPresent(nextExpectedOffset -> {
-                if (nextExpectedOffset < log.startOffset()) {
-                    listenerContext.fireHandleSnapshot(log.startOffset());
+                if (nextExpectedOffset < log.startOffset() && nextExpectedOffset < highWatermark) {
+                    SnapshotReader<T> snapshot = latestSnapshot().orElseThrow(() -> {
+                        return new IllegalStateException(
+                            String.format(
+                                "Snapshot expected when next offset is %s, log start offset is %s and high-watermark is %s",
+                                nextExpectedOffset,
+                                log.startOffset(),
+                                highWatermark
+                            )
+                        );
+                    });
+                    listenerContext.fireHandleSnapshot(snapshot);
                 }
             });
 
@@ -326,6 +336,16 @@ private void updateListenersProgress(List<ListenerContext> listenerContexts, lon
         }
     }
 
+    private Optional<SnapshotReader<T>> latestSnapshot() {
+        return log.latestSnapshotId().flatMap(snapshotId ->
+            log
+            .readSnapshot(snapshotId)
+            .map(reader ->
+                SnapshotReader.of(reader, serde, BufferSupplier.create(), MAX_BATCH_SIZE_BYTES)
+            )
+        );
+    }
+
     private void maybeFireHandleCommit(long baseOffset, int epoch, List<T> records) {
         for (ListenerContext listenerContext : listenerContexts) {
             OptionalLong nextExpectedOffsetOpt = listenerContext.nextExpectedOffset();
@@ -355,24 +375,28 @@ private void fireHandleResign(int epoch) {
     }
 
     @Override
-    public void initialize() throws IOException {
-        quorum.initialize(new OffsetAndEpoch(log.endOffset().offset, log.lastFetchedEpoch()));
+    public void initialize() {
+        try {
+            quorum.initialize(new OffsetAndEpoch(log.endOffset().offset, log.lastFetchedEpoch()));
 
-        long currentTimeMs = time.milliseconds();
-        if (quorum.isLeader()) {
-            throw new IllegalStateException("Voter cannot initialize as a Leader");
-        } else if (quorum.isCandidate()) {
-            onBecomeCandidate(currentTimeMs);
-        } else if (quorum.isFollower()) {
-            onBecomeFollower(currentTimeMs);
-        }
+            long currentTimeMs = time.milliseconds();
+            if (quorum.isLeader()) {
+                throw new IllegalStateException("Voter cannot initialize as a Leader");
+            } else if (quorum.isCandidate()) {
+                onBecomeCandidate(currentTimeMs);
+            } else if (quorum.isFollower()) {
+                onBecomeFollower(currentTimeMs);
+            }
 
-        // When there is only a single voter, become candidate immediately
-        if (quorum.isVoter()
-            && quorum.remoteVoters().isEmpty()
-            && !quorum.isLeader()
-            && !quorum.isCandidate()) {
-            transitionToCandidate(currentTimeMs);
+            // When there is only a single voter, become candidate immediately
+            if (quorum.isVoter()
+                && quorum.remoteVoters().isEmpty()
+                && !quorum.isCandidate()) {
+
+                transitionToCandidate(currentTimeMs);
+            }
+        } catch (IOException e) {
+            throw new RuntimeException(e);
         }
     }
 
@@ -1131,6 +1155,7 @@ private boolean handleFetchResponse(
                 if (records.sizeInBytes() > 0) {
                     appendAsFollower(records);
                 }
+
                 OptionalLong highWatermark = partitionResponse.highWatermark() < 0 ?
                     OptionalLong.empty() : OptionalLong.of(partitionResponse.highWatermark());
                 updateFollowerHighWatermark(state, highWatermark);
@@ -1260,28 +1285,34 @@ private FetchSnapshotResponseData handleFetchSnapshotRequest(
         }
 
         try (RawSnapshotReader snapshot = snapshotOpt.get()) {
-            if (partitionSnapshot.position() < 0 || partitionSnapshot.position() >= snapshot.sizeInBytes()) {
+            long snapshotSize = snapshot.sizeInBytes();
+            if (partitionSnapshot.position() < 0 || partitionSnapshot.position() >= snapshotSize) {
                 return FetchSnapshotResponse.singleton(
                     log.topicPartition(),
                     responsePartitionSnapshot -> addQuorumLeader(responsePartitionSnapshot)
                         .setErrorCode(Errors.POSITION_OUT_OF_RANGE.code())
                 );
             }
 
+            if (partitionSnapshot.position() > Integer.MAX_VALUE) {
+                throw new IllegalStateException(
+                    String.format(
+                        "Trying to fetch a snapshot with size (%s) and a position (%s) larger than %s",
+                        snapshotSize,
+                        partitionSnapshot.position(),
+                        Integer.MAX_VALUE
+                    )
+                );
+            }
+
             int maxSnapshotSize;
             try {
-                maxSnapshotSize = Math.toIntExact(snapshot.sizeInBytes());
+                maxSnapshotSize = Math.toIntExact(snapshotSize);
             } catch (ArithmeticException e) {
                 maxSnapshotSize = Integer.MAX_VALUE;
             }
 
-            if (partitionSnapshot.position() > Integer.MAX_VALUE) {
-                throw new IllegalStateException(String.format("Trying to fetch a snapshot with position: %d lager than Int.MaxValue", partitionSnapshot.position()));
-            }
-
-            UnalignedRecords records = snapshot.read(partitionSnapshot.position(), Math.min(data.maxBytes(), maxSnapshotSize));
-
-            long snapshotSize = snapshot.sizeInBytes();
+            UnalignedRecords records = snapshot.slice(partitionSnapshot.position(), Math.min(data.maxBytes(), maxSnapshotSize));
 
             return FetchSnapshotResponse.singleton(
                 log.topicPartition(),
@@ -1386,10 +1417,15 @@ private boolean handleFetchSnapshotResponse(
             );
         }
 
-        if (!(partitionSnapshot.unalignedRecords() instanceof MemoryRecords)) {
+        final UnalignedMemoryRecords records;
+        if (partitionSnapshot.unalignedRecords() instanceof MemoryRecords) {
+            records = new UnalignedMemoryRecords(((MemoryRecords) partitionSnapshot.unalignedRecords()).buffer());
+        } else if (partitionSnapshot.unalignedRecords() instanceof UnalignedMemoryRecords) {
+            records = (UnalignedMemoryRecords) partitionSnapshot.unalignedRecords();
+        } else {
             throw new IllegalStateException(String.format("Received unexpected fetch snapshot response: %s", partitionSnapshot));
         }
-        snapshot.append(new UnalignedMemoryRecords(((MemoryRecords) partitionSnapshot.unalignedRecords()).buffer()));
+        snapshot.append(records);
 
         if (snapshot.sizeInBytes() == partitionSnapshot.size()) {
             // Finished fetching the snapshot.
@@ -2103,7 +2139,7 @@ private long pollUnattachedAsObserver(UnattachedState state, long currentTimeMs)
     }
 
     private long pollCurrentState(long currentTimeMs) throws IOException {
-        maybeUpdateOldestSnapshotId();
+        maybeDeleteBeforeSnapshot();
 
         if (quorum.isLeader()) {
             return pollLeader(currentTimeMs);
@@ -2167,8 +2203,14 @@ private boolean maybeCompleteShutdown(long currentTimeMs) {
         return false;
     }
 
-    private void maybeUpdateOldestSnapshotId() {
-        log.latestSnapshotId().ifPresent(log::deleteBeforeSnapshot);
+    private void maybeDeleteBeforeSnapshot() {
+        log.latestSnapshotId().ifPresent(snapshotId -> {
+            quorum.highWatermark().ifPresent(highWatermark -> {
+                if (highWatermark.offset >= snapshotId.offset) {
+                    log.deleteBeforeSnapshot(snapshotId);
+                }
+            });
+        });
     }
 
     private void wakeup() {
@@ -2257,7 +2299,7 @@ public CompletableFuture<Void> shutdown(int timeoutMs) {
     }
 
     @Override
-    public SnapshotWriter<T> createSnapshot(OffsetAndEpoch snapshotId) throws IOException {
+    public SnapshotWriter<T> createSnapshot(OffsetAndEpoch snapshotId) {
         return new SnapshotWriter<>(
             log.createSnapshot(snapshotId),
             MAX_BATCH_SIZE_BYTES,
@@ -2365,14 +2407,16 @@ public synchronized OptionalLong nextExpectedOffset() {
         }
 
         /**
-         * This API is used when the Listener needs to be notified of a new Snapshot. This happens
-         * when the context last acked end offset is less that then log start offset.
+         * This API is used when the Listener needs to be notified of a new snapshot. This happens
+         * when the context's next offset is less than the log start offset.
          */
-        public void fireHandleSnapshot(long logStartOffset) {
+        public void fireHandleSnapshot(SnapshotReader<T> reader) {
             synchronized (this) {
-                nextOffset = logStartOffset;
+                nextOffset = reader.snapshotId().offset;
                 lastSent = null;
             }
+
+            listener.handleSnapshot(reader);
         }
 
         /**
@@ -2382,10 +2426,16 @@ public void fireHandleSnapshot(long logStartOffset) {
          * data in memory, we let the state machine read the records from disk.
          */
         public void fireHandleCommit(long baseOffset, Records records) {
-            BufferSupplier bufferSupplier = BufferSupplier.create();
-            RecordsBatchReader<T> reader = new RecordsBatchReader<>(baseOffset, records,
-                serde, bufferSupplier, this);
-            fireHandleCommit(reader);
+            fireHandleCommit(
+                RecordsBatchReader.of(
+                    baseOffset,
+                    records,
+                    serde,
+                    BufferSupplier.create(),
+                    MAX_BATCH_SIZE_BYTES,
+                    this
+                )
+            );
         }
 
         /**
@@ -2396,7 +2446,7 @@ public void fireHandleCommit(long baseOffset, Records records) {
          * followers.
          */
         public void fireHandleCommit(long baseOffset, int epoch, List<T> records) {
-            BatchReader.Batch<T> batch = new BatchReader.Batch<>(baseOffset, epoch, records);
+            BatchReader.Batch<T> batch = BatchReader.Batch.of(baseOffset, epoch, records);
             MemoryBatchReader<T> reader = new MemoryBatchReader<>(Collections.singletonList(batch), this);
             fireHandleCommit(reader);
         }
@@ -2425,6 +2475,7 @@ void fireHandleResign(int epoch) {
 
         public synchronized void onClose(BatchReader<T> reader) {
             OptionalLong lastOffset = reader.lastOffset();
+
             if (lastOffset.isPresent()) {
                 nextOffset = lastOffset.getAsLong() + 1;
             }