KAFKA-10394: generate snapshot

checkstyle/import-control.xml

@@ -318,6 +318,7 @@
 
   <subpackage name="raft">
     <allow pkg="org.apache.kafka.raft" />
+    <allow pkg="org.apache.kafka.snapshot" />
     <allow pkg="org.apache.kafka.clients" />
     <allow pkg="org.apache.kafka.common.config" />
     <allow pkg="org.apache.kafka.common.message" />
@@ -329,6 +330,12 @@
     <allow pkg="com.fasterxml.jackson" />
   </subpackage>
 
+  <subpackage name="snapshot">
+    <allow pkg="org.apache.kafka.common.record" />
+    <allow pkg="org.apache.kafka.raft" />
+    <allow pkg="org.apache.kafka.test"/>
+  </subpackage>
+
   <subpackage name="connect">
     <allow pkg="org.apache.kafka.common" />
     <allow pkg="org.apache.kafka.connect.data" />

core/src/main/scala/kafka/raft/KafkaMetadataLog.scala

@@ -16,6 +16,7 @@
  */
 package kafka.raft
 
+import java.nio.file.NoSuchFileException
 import java.util.Optional
 
 import kafka.log.{AppendOrigin, Log}
@@ -24,6 +25,10 @@ import org.apache.kafka.common.record.{MemoryRecords, Records}
 import org.apache.kafka.common.{KafkaException, TopicPartition}
 import org.apache.kafka.raft
 import org.apache.kafka.raft.{LogAppendInfo, LogFetchInfo, LogOffsetMetadata, Isolation, ReplicatedLog}
+import org.apache.kafka.snapshot.FileRawSnapshotReader
+import org.apache.kafka.snapshot.FileRawSnapshotWriter
+import org.apache.kafka.snapshot.RawSnapshotReader
+import org.apache.kafka.snapshot.RawSnapshotWriter
 
 import scala.compat.java8.OptionConverters._
 
@@ -141,6 +146,18 @@ class KafkaMetadataLog(
     topicPartition
   }
 
+  override def createSnapshot(snapshotId: raft.OffsetAndEpoch): RawSnapshotWriter = {
+    FileRawSnapshotWriter.create(log.dir.toPath, snapshotId)
+  }
+
+  override def readSnapshot(snapshotId: raft.OffsetAndEpoch): Optional[RawSnapshotReader] = {
+    try {
+      Optional.of(FileRawSnapshotReader.open(log.dir.toPath, snapshotId))
+    } catch {
+      case e: NoSuchFileException => Optional.empty()
+    }
+  }
+
   override def close(): Unit = {
     log.close()
   }

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

@@ -23,14 +23,14 @@
 import org.apache.kafka.common.message.BeginQuorumEpochRequestData;
 import org.apache.kafka.common.message.BeginQuorumEpochResponseData;
 import org.apache.kafka.common.message.DescribeQuorumRequestData;
-import org.apache.kafka.common.message.DescribeQuorumResponseData;
 import org.apache.kafka.common.message.DescribeQuorumResponseData.ReplicaState;
+import org.apache.kafka.common.message.DescribeQuorumResponseData;
 import org.apache.kafka.common.message.EndQuorumEpochRequestData;
 import org.apache.kafka.common.message.EndQuorumEpochResponseData;
 import org.apache.kafka.common.message.FetchRequestData;
 import org.apache.kafka.common.message.FetchResponseData;
-import org.apache.kafka.common.message.LeaderChangeMessage;
 import org.apache.kafka.common.message.LeaderChangeMessage.Voter;
+import org.apache.kafka.common.message.LeaderChangeMessage;
 import org.apache.kafka.common.message.VoteRequestData;
 import org.apache.kafka.common.message.VoteResponseData;
 import org.apache.kafka.common.metrics.Metrics;
@@ -61,6 +61,7 @@
 import org.apache.kafka.raft.internals.MemoryBatchReader;
 import org.apache.kafka.raft.internals.RecordsBatchReader;
 import org.apache.kafka.raft.internals.ThresholdPurgatory;
+import org.apache.kafka.snapshot.SnapshotWriter;
 import org.slf4j.Logger;
 
 import java.io.IOException;
@@ -1815,6 +1816,18 @@ public CompletableFuture<Void> shutdown(int timeoutMs) {
         return shutdownComplete;
     }
 
+    @Override
+    public SnapshotWriter<T> createSnapshot(OffsetAndEpoch snapshotId) throws IOException {
+        return new SnapshotWriter<>(
+            log.createSnapshot(snapshotId),
+            MAX_BATCH_SIZE,
+            memoryPool,
+            time,
+            CompressionType.NONE,
+            serde
+        );
+    }
+
     private void close() {
         kafkaRaftMetrics.close();
     }

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

@@ -35,26 +35,26 @@
  * only valid state transitions. Below we define the possible state transitions and
  * how they are triggered:
  *
- * Unattached|Resigned =>
+ * Unattached|Resigned transitions to:
  *    Unattached: After learning of a new election with a higher epoch
  *    Voted: After granting a vote to a candidate
  *    Candidate: After expiration of the election timeout
  *    Follower: After discovering a leader with an equal or larger epoch
  *
- * Voted =>
+ * Voted transitions to:
  *    Unattached: After learning of a new election with a higher epoch
  *    Candidate: After expiration of the election timeout
  *
- * Candidate =>
+ * Candidate transitions to:
  *    Unattached: After learning of a new election with a higher epoch
  *    Candidate: After expiration of the election timeout
  *    Leader: After receiving a majority of votes
  *
- * Leader =>
+ * Leader transitions to:
  *    Unattached: After learning of a new election with a higher epoch
  *    Resigned: When shutting down gracefully
  *
- * Follower =>
+ * Follower transitions to:
  *    Unattached: After learning of a new election with a higher epoch
  *    Candidate: After expiration of the fetch timeout
  *    Follower: After discovering a leader with a larger epoch
@@ -63,11 +63,11 @@
  * states are not possible for observers, so the only transitions that are possible
  * are between Unattached and Follower.
  *
- * Unattached =>
+ * Unattached transitions to:
  *    Unattached: After learning of a new election with a higher epoch
  *    Follower: After discovering a leader with an equal or larger epoch
  *
- * Follower =>
+ * Follower transitions to:
  *    Unattached: After learning of a new election with a higher epoch
  *    Follower: After discovering a leader with a larger epoch
  *

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

@@ -16,6 +16,8 @@
  */
 package org.apache.kafka.raft;
 
+import org.apache.kafka.snapshot.SnapshotWriter;
+
 import java.io.IOException;
 import java.util.List;
 import java.util.concurrent.CompletableFuture;
@@ -100,4 +102,15 @@ default void handleResign() {}
      */
     CompletableFuture<Void> shutdown(int timeoutMs);
 
+    /**
+     * Create a writable snapshot file for a given offset and epoch.
+     *
+     * The RaftClient assumes that the snapshot return will contain the records up to but
+     * not including the end offset in the snapshot id. See {@link SnapshotWriter} for
+     * details on how to use this object.
+     *
+     * @param snapshotId the end offset and epoch that identifies the snapshot
+     * @return a writable snapshot
+     */
+    SnapshotWriter<T> createSnapshot(OffsetAndEpoch snapshotId) throws IOException;
 }

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

@@ -18,8 +18,11 @@
 
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.record.Records;
+import org.apache.kafka.snapshot.RawSnapshotReader;
+import org.apache.kafka.snapshot.RawSnapshotWriter;
 
 import java.io.Closeable;
+import java.io.IOException;
 import java.util.Optional;
 import java.util.OptionalLong;
 
@@ -149,6 +152,29 @@ default OptionalLong truncateToEndOffset(OffsetAndEpoch endOffset) {
         return OptionalLong.of(truncationOffset);
     }
 
+    /**
+     * Create a writable snapshot for the given snapshot id.
+     *
+     * See {@link RawSnapshotWriter} for details on how to use this object.
+     *
+     * @param snapshotId the end offset and epoch that identifies the snapshot
+     * @return a writable snapshot
+     */
+    RawSnapshotWriter createSnapshot(OffsetAndEpoch snapshotId) throws IOException;
+
+    /**
+     * Opens a readable snapshot for the given snapshot id.
+     *
+     * Returns an Optional with a readable snapshot, if the snapshot exists, otherwise
+     * returns an empty Optional. See {@link RawSnapshotReader} for details on how to
+     * use this object.
+     *
+     * @param snapshotId the end offset and epoch that identifies the snapshot
+     * @return an Optional with a readable snapshot, if the snapshot exists, otherwise
+     *         returns an empty Optional
+     */
+    Optional<RawSnapshotReader> readSnapshot(OffsetAndEpoch snapshotId) throws IOException;
+
     default void close() {}
 
 }

raft/src/main/java/org/apache/kafka/snapshot/FileRawSnapshotReader.java

@@ -0,0 +1,94 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.snapshot;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.nio.file.Path;
+import java.util.Iterator;
+import org.apache.kafka.common.record.FileLogInputStream.FileChannelRecordBatch;
+import org.apache.kafka.common.record.FileRecords;
+import org.apache.kafka.common.record.RecordBatch;
+import org.apache.kafka.common.utils.AbstractIterator;
+import org.apache.kafka.raft.OffsetAndEpoch;
+
+public final class FileRawSnapshotReader implements RawSnapshotReader {
+    private final FileRecords fileRecords;
+    private final OffsetAndEpoch snapshotId;
+
+    private FileRawSnapshotReader(FileRecords fileRecords, OffsetAndEpoch snapshotId) {
+        this.fileRecords = fileRecords;
+        this.snapshotId = snapshotId;
+    }
+
+    @Override
+    public OffsetAndEpoch snapshotId() {
+        return snapshotId;
+    }
+
+    @Override
+    public long sizeInBytes() {
+        return fileRecords.sizeInBytes();
+    }
+
+    @Override
+    public Iterator<RecordBatch> iterator() {
+        return new Iterator<RecordBatch>() {
+            private final AbstractIterator<FileChannelRecordBatch> iterator = fileRecords.batchIterator();
+
+            @Override
+            public boolean hasNext() {
+                return iterator.hasNext();
+            }
+
+            @Override
+            public RecordBatch next() {
+                return iterator.next();
+            }
+        };
+    }
+
+    @Override
+    public int read(ByteBuffer buffer, long position) throws IOException {
+        return fileRecords.channel().read(buffer, position);
+    }
+
+    @Override
+    public void close() throws IOException {
+        fileRecords.close();
+    }
+
+    /**
+     * Opens a snapshot for reading.
+     *
+     * @param logDir the directory for the topic partition
+     * @param snapshotId the end offset and epoch for the snapshotId
+     * @throws java.nio.file.NoSuchFileException if the snapshot doesn't exist
+     * @throws IOException for any IO error while opening the snapshot
+     */
+    public static FileRawSnapshotReader open(Path logDir, OffsetAndEpoch snapshotId) throws IOException {
+        FileRecords fileRecords = FileRecords.open(
+            Snapshots.snapshotPath(logDir, snapshotId).toFile(),
+            false, // mutable
+            true, // fileAlreadyExists
+            0, // initFileSize
+            false // preallocate
+        );
+
+        return new FileRawSnapshotReader(fileRecords, snapshotId);
+    }
+}