feat(snapshot_read): snapshot-read cache (#2453) (#2454)

feat(snapshot_read): snapshot read cache

Signed-off-by: Robin Han <[email protected]>

Author: superhx

bin/kafka-server-start.sh

@@ -47,7 +47,7 @@ if [ "x$KAFKA_HEAP_OPTS" = "x" ]; then
 fi
 
 if [ "x$KAFKA_OPTS" = "x" ]; then
-    export KAFKA_OPTS="-Dio.netty.allocator.maxOrder=11"
+    export KAFKA_OPTS="-XX:+ExitOnOutOfMemoryError -XX:+HeapDumpOnOutOfMemoryError -Dio.netty.allocator.maxOrder=11"
 fi
 
 EXTRA_ARGS=${EXTRA_ARGS-'-name kafkaServer -loggc'}

core/src/main/scala/kafka/server/MetadataCache.scala

@@ -22,6 +22,7 @@ import org.apache.kafka.admin.BrokerMetadata
 import org.apache.kafka.common.message.{MetadataResponseData, UpdateMetadataRequestData}
 import org.apache.kafka.common.network.ListenerName
 import org.apache.kafka.common.{Cluster, Node, TopicPartition, Uuid}
+import org.apache.kafka.image.MetadataImage
 import org.apache.kafka.metadata.BrokerRegistration
 import org.apache.kafka.server.common.automq.AutoMQVersion
 import org.apache.kafka.server.common.{FinalizedFeatures, KRaftVersion, MetadataVersion}
@@ -126,6 +127,8 @@ trait MetadataCache {
   def getValue(key: String): ByteBuffer
 
   def getStreamEndOffset(streamId: Long): OptionalLong
+
+  def safeRun[T](func: Function[MetadataImage, T]): T
   // AutoMQ inject end
 }
 

core/src/main/scala/kafka/server/metadata/KRaftMetadataCache.scala

@@ -623,7 +623,7 @@ class KRaftMetadataCache(
     _currentImage.features().autoMQVersion()
   }
 
-  def safeRun[T](func: Function[MetadataImage, T]): T = {
+  override def safeRun[T](func: Function[MetadataImage, T]): T = {
     val image = retainedImage()
     try {
       func.apply(image)

core/src/main/scala/kafka/server/metadata/ZkMetadataCache.scala

@@ -40,6 +40,7 @@ import org.apache.kafka.common.network.ListenerName
 import org.apache.kafka.common.protocol.Errors
 import org.apache.kafka.common.requests.{AbstractControlRequest, ApiVersionsResponse, MetadataResponse, UpdateMetadataRequest}
 import org.apache.kafka.common.security.auth.SecurityProtocol
+import org.apache.kafka.image.MetadataImage
 import org.apache.kafka.metadata.BrokerRegistration
 import org.apache.kafka.server.common.automq.AutoMQVersion
 import org.apache.kafka.server.common.{FinalizedFeatures, MetadataVersion}
@@ -736,6 +737,10 @@ class ZkMetadataCache(
   override def getStreamEndOffset(streamId: Long): OptionalLong = {
     throw new UnsupportedOperationException()
   }
+
+  override def safeRun[T](func: Function[MetadataImage, T]): T = {
+    throw new UnsupportedOperationException()
+  }
   // AutoMQ inject end
 
 }

s3stream/src/main/java/com/automq/stream/api/ReadOptions.java

@@ -27,6 +27,7 @@ public class ReadOptions {
     private boolean fastRead;
     private boolean pooledBuf;
     private boolean prioritizedRead;
+    private boolean snapshotRead;
 
     public static Builder builder() {
         return new Builder();
@@ -44,6 +45,15 @@ public boolean prioritizedRead() {
         return prioritizedRead;
     }
 
+    public boolean snapshotRead() {
+        return snapshotRead;
+    }
+
+    public ReadOptions snapshotRead(boolean snapshotRead) {
+        this.snapshotRead = snapshotRead;
+        return this;
+    }
+
     public static class Builder {
         private final ReadOptions options = new ReadOptions();
 
@@ -68,6 +78,11 @@ public Builder prioritizedRead(boolean prioritizedRead) {
             return this;
         }
 
+        public Builder snapshotRead(boolean snapshotRead) {
+            options.snapshotRead = snapshotRead;
+            return this;
+        }
+
         public ReadOptions build() {
             return options;
         }

s3stream/src/main/java/com/automq/stream/s3/Config.java

@@ -59,6 +59,7 @@ public class Config {
     private int refillPeriodMs = 10;
     private long objectRetentionTimeInSecond = 10 * 60; // 10min
     private boolean failoverEnable = false;
+    private boolean snapshotReadEnable = false;
     private Supplier<Version> version = () -> {
         throw new UnsupportedOperationException();
     };
@@ -324,6 +325,15 @@ public boolean failoverEnable() {
         return failoverEnable;
     }
 
+    public Config snapshotReadEnable(boolean snapshotReadEnable) {
+        this.snapshotReadEnable = snapshotReadEnable;
+        return this;
+    }
+
+    public boolean snapshotReadEnable() {
+        return snapshotReadEnable;
+    }
+
     public Config version(Supplier<Version> version) {
         this.version = version;
         return this;

s3stream/src/main/java/com/automq/stream/s3/ObjectReader.java

@@ -159,7 +159,7 @@ public CompletableFuture<DataBlockGroup> read(ReadOptions readOptions, DataBlock
         }
 
         void asyncGetBasicObjectInfo() {
-            int guessIndexBlockSize = 1024 + (int) (metadata.objectSize() / (1024 * 1024 /* 1MB */) * 36 /* index unit size*/);
+            int guessIndexBlockSize = 8192 + (int) (metadata.objectSize() / (1024 * 1024 /* 1MB */) * 36 /* index unit size*/);
             asyncGetBasicObjectInfo0(Math.max(0, metadata.objectSize() - guessIndexBlockSize), true);
         }
 
@@ -578,6 +578,10 @@ private static int check(ByteBuf buf) {
         }
 
         public CloseableIterator<StreamRecordBatch> iterator() {
+            return iterator(true);
+        }
+
+        public CloseableIterator<StreamRecordBatch> iterator(boolean copy) {
             ByteBuf buf = this.buf.duplicate();
             AtomicInteger currentBlockRecordCount = new AtomicInteger(0);
             AtomicInteger remainingRecordCount = new AtomicInteger(recordCount);
@@ -599,7 +603,7 @@ public StreamRecordBatch next() {
                         buf.skipBytes(4);
                     }
                     currentBlockRecordCount.decrementAndGet();
-                    return StreamRecordBatchCodec.duplicateDecode(buf);
+                    return copy ? StreamRecordBatchCodec.duplicateDecode(buf) : StreamRecordBatchCodec.sliceRetainDecode(buf);
                 }
 
                 @Override

s3stream/src/main/java/com/automq/stream/s3/S3Storage.java

@@ -24,6 +24,7 @@
 import com.automq.stream.s3.cache.LogCache;
 import com.automq.stream.s3.cache.ReadDataBlock;
 import com.automq.stream.s3.cache.S3BlockCache;
+import com.automq.stream.s3.cache.SnapshotReadCache;
 import com.automq.stream.s3.context.AppendContext;
 import com.automq.stream.s3.context.FetchContext;
 import com.automq.stream.s3.failover.Failover;
@@ -99,6 +100,7 @@ public class S3Storage implements Storage {
      * WAL log cache
      */
     private final LogCache deltaWALCache;
+    private final LogCache snapshotReadCache;
     /**
      * WAL out of order callback sequencer. {@link #streamCallbackLocks} will ensure the memory safety.
      */
@@ -150,7 +152,16 @@ public S3Storage(Config config, WriteAheadLog deltaWAL, StreamManager streamMana
         this.maxDeltaWALCacheSize = config.walCacheSize();
         this.deltaWAL = deltaWAL;
         this.blockCache = blockCache;
-        this.deltaWALCache = new LogCache(config.walCacheSize(), config.walUploadThreshold(), config.maxStreamNumPerStreamSetObject());
+        long deltaWALCacheSize = config.walCacheSize();
+        long snapshotReadCacheSize = 0;
+        if (config.snapshotReadEnable()) {
+            deltaWALCacheSize = Math.max(config.walCacheSize() / 3, 10L * 1024 * 1024);
+            snapshotReadCacheSize = Math.max(config.walCacheSize() / 3 * 2, 10L * 1024 * 1024);
+        }
+        this.deltaWALCache = new LogCache(deltaWALCacheSize, config.walUploadThreshold(), config.maxStreamNumPerStreamSetObject());
+        this.snapshotReadCache = new LogCache(snapshotReadCacheSize, Math.max(snapshotReadCacheSize / 6, 1));
+        S3StreamMetricsManager.registerDeltaWalCacheSizeSupplier(() -> deltaWALCache.size() + snapshotReadCache.size());
+        SnapshotReadCache.instance().setup(this.snapshotReadCache, objectStorage);
         this.streamManager = streamManager;
         this.objectManager = objectManager;
         this.objectStorage = objectStorage;
@@ -547,13 +558,15 @@ public CompletableFuture<ReadDataBlock> read(FetchContext context,
         return cf;
     }
 
+    @SuppressWarnings({"checkstyle:npathcomplexity"})
     @WithSpan
     private CompletableFuture<ReadDataBlock> read0(FetchContext context,
         @SpanAttribute long streamId,
         @SpanAttribute long startOffset,
         @SpanAttribute long endOffset,
         @SpanAttribute int maxBytes) {
-        List<StreamRecordBatch> logCacheRecords = deltaWALCache.get(context, streamId, startOffset, endOffset, maxBytes);
+        LogCache firstCache = context.readOptions().snapshotRead() ? snapshotReadCache : deltaWALCache;
+        List<StreamRecordBatch> logCacheRecords = firstCache.get(context, streamId, startOffset, endOffset, maxBytes);
         if (!logCacheRecords.isEmpty() && logCacheRecords.get(0).getBaseOffset() <= startOffset) {
             return CompletableFuture.completedFuture(new ReadDataBlock(logCacheRecords, CacheAccessType.DELTA_WAL_CACHE_HIT));
         }

s3stream/src/main/java/com/automq/stream/s3/S3Stream.java

@@ -224,12 +224,16 @@ private CompletableFuture<AppendResult> append0(AppendContext context, RecordBat
         });
     }
 
+    @SuppressWarnings({"checkstyle:npathcomplexity"})
     @Override
     @WithSpan
     public CompletableFuture<FetchResult> fetch(FetchContext context,
         @SpanAttribute long startOffset,
         @SpanAttribute long endOffset,
         @SpanAttribute int maxBytes) {
+        if (snapshotRead()) {
+            context.readOptions().snapshotRead(true);
+        }
         TimerUtil timerUtil = new TimerUtil();
         readLock.lock();
         try {

s3stream/src/main/java/com/automq/stream/s3/StreamRecordBatchCodec.java

@@ -74,15 +74,23 @@ public static StreamRecordBatch duplicateDecode(ByteBuf buf) {
      * Decode a stream record batch from a byte buffer and move the reader index.
      * The returned stream record batch shares the payload buffer with the input buffer.
      */
-    public static StreamRecordBatch decode(ByteBuf buf) {
+    public static StreamRecordBatch decode(ByteBuf buf, boolean retain) {
         buf.readByte(); // magic
         long streamId = buf.readLong();
         long epoch = buf.readLong();
         long baseOffset = buf.readLong();
         int lastOffsetDelta = buf.readInt();
         int payloadLength = buf.readInt();
-        ByteBuf payload = buf.slice(buf.readerIndex(), payloadLength);
+        ByteBuf payload = retain ? buf.retainedSlice(buf.readerIndex(), payloadLength) : buf.slice(buf.readerIndex(), payloadLength);
         buf.skipBytes(payloadLength);
         return new StreamRecordBatch(streamId, epoch, baseOffset, lastOffsetDelta, payload);
     }
+
+    public static StreamRecordBatch decode(ByteBuf buf) {
+        return decode(buf, false);
+    }
+
+    public static StreamRecordBatch sliceRetainDecode(ByteBuf buf) {
+        return decode(buf, true);
+    }
 }