chore(zerozone2): isolate the record encode allocator (#2944)

Signed-off-by: Robin Han <hanxvdovehx@gmail.com>

Author: superhx

core/src/main/java/kafka/automq/zerozone/DefaultLinkRecordDecoder.java

@@ -22,6 +22,7 @@
 import org.apache.kafka.common.record.MemoryRecords;
 import org.apache.kafka.common.record.MutableRecordBatch;
 
+import com.automq.stream.s3.cache.SnapshotReadCache;
 import com.automq.stream.s3.model.StreamRecordBatch;
 
 import org.slf4j.Logger;
@@ -62,8 +63,8 @@ public CompletableFuture<StreamRecordBatch> decode(StreamRecordBatch src) {
                     recordBatch.setPartitionLeaderEpoch(linkRecord.partitionLeaderEpoch());
                     StreamRecordBatch streamRecordBatch = new StreamRecordBatch(src.getStreamId(), src.getEpoch(), src.getBaseOffset(),
                         -src.getCount(), Unpooled.wrappedBuffer(records.buffer()));
-                    // duplicated copy the payload
-                    streamRecordBatch.encoded();
+                    // The buf will be release after the finally block, so we need copy the data by #encoded.
+                    streamRecordBatch.encoded(SnapshotReadCache.ENCODE_ALLOC);
                     return streamRecordBatch;
                 } finally {
                     src.release();

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

@@ -60,6 +60,7 @@ public class ByteBufAlloc {
     public static final int BLOCK_CACHE = 11;
     public static final int S3_WAL = 12;
     public static final int POOLED_MEMORY_RECORDS = 13;
+    public static final int SNAPSHOT_READ_CACHE = 14;
 
     // the MAX_TYPE_NUMBER may change when new type added.
     public static final int MAX_TYPE_NUMBER = 20;
@@ -104,6 +105,7 @@ public class ByteBufAlloc {
         registerAllocType(BLOCK_CACHE, "block_cache");
         registerAllocType(S3_WAL, "s3_wal");
         registerAllocType(POOLED_MEMORY_RECORDS, "pooled_memory_records");
+        registerAllocType(SNAPSHOT_READ_CACHE, "snapshot_read_cache");
 
     }
 

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

@@ -24,8 +24,6 @@
 
 import io.netty.buffer.ByteBuf;
 
-import static com.automq.stream.s3.ByteBufAlloc.ENCODE_RECORD;
-
 public class StreamRecordBatchCodec {
     public static final byte MAGIC_V0 = 0x22;
     public static final int HEADER_SIZE =
@@ -35,12 +33,11 @@ public class StreamRecordBatchCodec {
             + 8 // baseOffset
             + 4 // lastOffsetDelta
             + 4; // payload length
-    private static final ByteBufSeqAlloc ENCODE_ALLOC = new ByteBufSeqAlloc(ENCODE_RECORD, 8);
 
-    public static ByteBuf encode(StreamRecordBatch streamRecord) {
+    public static ByteBuf encode(StreamRecordBatch streamRecord, ByteBufSeqAlloc alloc) {
         int totalLength = HEADER_SIZE + streamRecord.size(); // payload
         // use sequential allocator to avoid memory fragmentation
-        ByteBuf buf = ENCODE_ALLOC.byteBuffer(totalLength);
+        ByteBuf buf = alloc.byteBuffer(totalLength);
         buf.writeByte(MAGIC_V0);
         buf.writeLong(streamRecord.getStreamId());
         buf.writeLong(streamRecord.getEpoch());

s3stream/src/main/java/com/automq/stream/s3/cache/LogCache.java

@@ -345,7 +345,7 @@ public void clearStreamRecords(long streamId) {
         readLock.lock();
         try {
             for (LogCacheBlock block : blocks) {
-                size.addAndGet(-block.free(streamId));
+                block.free(streamId);
             }
         } finally {
             readLock.unlock();
@@ -478,16 +478,13 @@ public void free() {
             }, LOGGER);
         }
 
-        public long free(long streamId) {
-            AtomicLong size = new AtomicLong();
+        public void free(long streamId) {
             suppress(() -> {
                 StreamCache streamCache = map.remove(streamId);
                 if (streamCache != null) {
-                    size.addAndGet(streamCache.free());
+                    streamCache.free();
                 }
             }, LOGGER);
-            this.size.addAndGet(-size.get());
-            return size.get();
         }
 
         public void addFreeListener(FreeListener freeListener) {
@@ -625,14 +622,9 @@ synchronized StreamRange range() {
             return new StreamRange(startOffset, endOffset);
         }
 
-        synchronized long free() {
-            AtomicLong size = new AtomicLong();
-            records.forEach(record -> {
-                size.addAndGet(record.occupiedSize());
-                record.release();
-            });
+        synchronized void free() {
+            records.forEach(StreamRecordBatch::release);
             records.clear();
-            return size.get();
         }
 
         synchronized long startOffset() {

s3stream/src/main/java/com/automq/stream/s3/cache/SnapshotReadCache.java

@@ -1,5 +1,6 @@
 package com.automq.stream.s3.cache;
 
+import com.automq.stream.ByteBufSeqAlloc;
 import com.automq.stream.api.LinkRecordDecoder;
 import com.automq.stream.s3.DataBlockIndex;
 import com.automq.stream.s3.ObjectReader;
@@ -47,7 +48,10 @@
 import io.opentelemetry.api.common.AttributeKey;
 import io.opentelemetry.api.common.Attributes;
 
+import static com.automq.stream.s3.ByteBufAlloc.SNAPSHOT_READ_CACHE;
+
 public class SnapshotReadCache {
+    public static final ByteBufSeqAlloc ENCODE_ALLOC = new ByteBufSeqAlloc(SNAPSHOT_READ_CACHE, 8);
     private static final Logger LOGGER = LoggerFactory.getLogger(SnapshotReadCache.class);
     private static final long MAX_INFLIGHT_LOAD_BYTES = 100L * 1024 * 1024;
 
@@ -232,8 +236,7 @@ public void run() {
                         return;
                     }
                     records.addAll(cfList.stream().map(CompletableFuture::join).toList());
-                    // move to mem pool
-                    records.forEach(StreamRecordBatch::encoded);
+                    records.forEach(r -> r.encoded(ENCODE_ALLOC));
                     loadCf.complete(null);
                 });
             }).whenComplete((rst, ex) -> {

s3stream/src/main/java/com/automq/stream/s3/model/StreamRecordBatch.java

@@ -19,13 +19,17 @@
 
 package com.automq.stream.s3.model;
 
+import com.automq.stream.ByteBufSeqAlloc;
 import com.automq.stream.s3.StreamRecordBatchCodec;
 import com.automq.stream.utils.biniarysearch.ComparableItem;
 
 import io.netty.buffer.ByteBuf;
 
+import static com.automq.stream.s3.ByteBufAlloc.ENCODE_RECORD;
+
 public class StreamRecordBatch implements Comparable<StreamRecordBatch>, ComparableItem<Long> {
     private static final int OBJECT_OVERHEAD = 48 /* fields */ + 48 /* ByteBuf payload */ + 48 /* ByteBuf encoded */;
+    private static final ByteBufSeqAlloc ENCODE_ALLOC = new ByteBufSeqAlloc(ENCODE_RECORD, 8);
     private final long streamId;
     private final long epoch;
     private final long baseOffset;
@@ -42,9 +46,13 @@ public StreamRecordBatch(long streamId, long epoch, long baseOffset, int count,
     }
 
     public ByteBuf encoded() {
+        return encoded(ENCODE_ALLOC);
+    }
+
+    public ByteBuf encoded(ByteBufSeqAlloc alloc) {
         // TODO: keep the ref count
         if (encoded == null) {
-            encoded = StreamRecordBatchCodec.encode(this);
+            encoded = StreamRecordBatchCodec.encode(this, alloc);
             ByteBuf oldPayload = payload;
             payload = encoded.slice(encoded.readerIndex() + encoded.readableBytes() - payload.readableBytes(), payload.readableBytes());
             oldPayload.release();

s3stream/src/test/java/com/automq/stream/s3/cache/LogCacheTest.java

@@ -96,15 +96,17 @@ public void testClearStreamRecords() {
         logCache.archiveCurrentBlock();
         logCache.put(new StreamRecordBatch(233L, 0L, 13L, 2, TestUtils.random(20)));
 
-        long size0 = logCache.size();
         logCache.put(new StreamRecordBatch(234L, 0L, 13L, 2, TestUtils.random(20)));
-        long size1 = logCache.size();
 
+        assertTrue(logCache.blocks.get(0).containsStream(233L));
+        assertTrue(logCache.blocks.get(1).containsStream(234L));
         logCache.clearStreamRecords(233L);
-        assertEquals(size1 - size0, logCache.size());
+        assertFalse(logCache.blocks.get(0).containsStream(233L));
+        assertTrue(logCache.blocks.get(1).containsStream(234L));
 
         logCache.clearStreamRecords(234L);
-        assertEquals(0, logCache.size());
+        assertFalse(logCache.blocks.get(0).containsStream(233L));
+        assertFalse(logCache.blocks.get(1).containsStream(234L));
     }
 
     @Test