Repository: AutoMQ/automq (language: java)

Issue:
[Enhancement] Optimize the performance of WALUtil#generateRecord
The performance of WALUtil#generateRecord, executed in the critical section, will significantly impact latency. We should avoid performing CPU-bound operations.

<img width="2201" alt="image" src="https://github.com/user-attachments/assets/f4bf8315-dacb-4757-b8c3-1859a3c4adb0">

The flame graph indicates that allocating the buffer and calculating the checksum is the most time-consuming task. We could optimize this function by preallocating the buffer and using the more effective checksum algorithm.


Discussion / hints:

Fix patch:
diff --git a/s3stream/src/main/java/com/automq/stream/FixedSizeByteBufPool.java b/s3stream/src/main/java/com/automq/stream/FixedSizeByteBufPool.java
index 5dbfc696cb..20f2f61ad1 100644
--- a/s3stream/src/main/java/com/automq/stream/FixedSizeByteBufPool.java
+++ b/s3stream/src/main/java/com/automq/stream/FixedSizeByteBufPool.java
@@ -66,6 +66,7 @@ private ByteBuf allocate() {
      */
     public void release(ByteBuf buffer) {
         assert buffer.capacity() == bufferSize;
+        assert buffer.refCnt() == 1;
 
         if (poolSize.get() >= maxPoolSize) {
             buffer.release();
diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/common/RecordHeader.java b/s3stream/src/main/java/com/automq/stream/s3/wal/common/RecordHeader.java
index ca8c3fb62e..17e41487d6 100644
--- a/s3stream/src/main/java/com/automq/stream/s3/wal/common/RecordHeader.java
+++ b/s3stream/src/main/java/com/automq/stream/s3/wal/common/RecordHeader.java
@@ -96,10 +96,16 @@ private ByteBuf marshalHeaderExceptCRC(ByteBuf buf) {
         return buf;
     }
 
-    public ByteBuf marshal(ByteBuf emptyBuf) {
+    public ByteBuf marshal(ByteBuf emptyBuf, boolean calculateCRC) {
         assert emptyBuf.writableBytes() == RECORD_HEADER_SIZE;
         ByteBuf buf = marshalHeaderExceptCRC(emptyBuf);
-        buf.writeInt(WALUtil.crc32(buf, RECORD_HEADER_WITHOUT_CRC_SIZE));
+
+        if (calculateCRC) {
+            buf.writeInt(WALUtil.crc32(buf, RECORD_HEADER_WITHOUT_CRC_SIZE));
+        } else {
+            buf.writeInt(-1);
+        }
+
         return buf;
     }
 }
diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/ObjectWALService.java b/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/ObjectWALService.java
index ac9307117a..16caa0881d 100644
--- a/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/ObjectWALService.java
+++ b/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/ObjectWALService.java
@@ -11,6 +11,7 @@
 
 package com.automq.stream.s3.wal.impl.object;
 
+import com.automq.stream.FixedSizeByteBufPool;
 import com.automq.stream.s3.ByteBufAlloc;
 import com.automq.stream.s3.network.ThrottleStrategy;
 import com.automq.stream.s3.operator.ObjectStorage;
@@ -19,6 +20,7 @@
 import com.automq.stream.s3.wal.RecoverResult;
 import com.automq.stream.s3.wal.WriteAheadLog;
 import com.automq.stream.s3.wal.common.AppendResultImpl;
+import com.automq.stream.s3.wal.common.Record;
 import com.automq.stream.s3.wal.common.RecordHeader;
 import com.automq.stream.s3.wal.common.RecoverResultImpl;
 import com.automq.stream.s3.wal.common.WALMetadata;
@@ -28,6 +30,7 @@
 import com.automq.stream.s3.wal.util.WALUtil;
 import com.automq.stream.utils.Time;
 import io.netty.buffer.ByteBuf;
+import io.netty.buffer.CompositeByteBuf;
 import io.netty.buffer.Unpooled;
 import java.io.IOException;
 import java.util.ArrayDeque;
@@ -36,17 +39,16 @@
 import java.util.Objects;
 import java.util.Queue;
 import java.util.concurrent.CompletableFuture;
-
 import org.apache.commons.lang3.exception.ExceptionUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import static com.automq.stream.s3.wal.common.RecordHeader.RECORD_HEADER_SIZE;
-import static com.automq.stream.s3.wal.common.RecordHeader.RECORD_HEADER_WITHOUT_CRC_SIZE;
 
 public class ObjectWALService implements WriteAheadLog {
     private static final Logger log = LoggerFactory.getLogger(ObjectWALService.class);
 
+    protected FixedSizeByteBufPool byteBufPool = new FixedSizeByteBufPool(RECORD_HEADER_SIZE, 1024 * Runtime.getRuntime().availableProcessors());
     protected ObjectStorage objectStorage;
     protected ObjectWALConfig config;
 
@@ -84,20 +86,45 @@ public WALMetadata metadata() {
 
     @Override
     public AppendResult append(TraceContext context, ByteBuf data, int crc) throws OverCapacityException {
+        ByteBuf header = byteBufPool.get();
+        assert header.refCnt() == 1;
+
+        final CompletableFuture<AppendResult.CallbackResult> appendResultFuture = new CompletableFuture<>();
+        appendResultFuture.whenComplete((result, cause) -> {
+            // Return the header buffer to the buffer pool.
+            assert header.refCnt() == 1;
+            byteBufPool.release(header);
+        });
+
         try {
             final long recordSize = RECORD_HEADER_SIZE + data.readableBytes();
-            final CompletableFuture<AppendResult.CallbackResult> appendResultFuture = new CompletableFuture<>();
-            long expectedWriteOffset = accumulator.append(recordSize, start -> WALUtil.generateRecord(data, crc, start), appendResultFuture);
+
+            long expectedWriteOffset = accumulator.append(recordSize, start -> {
+                CompositeByteBuf recordByteBuf = ByteBufAlloc.compositeByteBuffer();
+
+                // To prevent duplication with the s3 client, disable the checksum.
+                Record record = WALUtil.generateRecord(data, header.retainedDuplicate(), -1, start, false);
+                recordByteBuf.addComponents(true, record.header(), record.body());
+                return recordByteBuf;
+            }, appendResultFuture);
 
             return new AppendResultImpl(expectedWriteOffset, appendResultFuture);
         } catch (Exception e) {
-            // Make sure the data buffer is released.
+            // Make sure the header buffer and data buffer is released.
+            if (header.refCnt() > 1) {
+                header.release();
+            } else {
+                log.error("[Bug] The header buffer is already released.", e);
+            }
+
             if (data.refCnt() > 0) {
                 data.release();
             } else {
                 log.error("[Bug] The data buffer is already released.", e);
             }
 
+            // Complete the future with exception, ensure the whenComplete method is executed.
+            appendResultFuture.completeExceptionally(e);
             Throwable cause = ExceptionUtils.getRootCause(e);
             if (cause instanceof OverCapacityException) {
                 if (((OverCapacityException) cause).error()) {
@@ -109,7 +136,7 @@ public AppendResult append(TraceContext context, ByteBuf data, int crc) throws O
                 throw new OverCapacityException("Append record to S3 WAL failed, due to accumulator is full: " + cause.getMessage());
             } else {
                 log.error("Append record to S3 WAL failed, due to unrecoverable exception.", e);
-                return new AppendResultImpl(-1, CompletableFuture.failedFuture(e));
+                return new AppendResultImpl(-1, appendResultFuture);
             }
         }
     }
@@ -231,10 +258,6 @@ public RecoverResult next() {
             ByteBuf recordHeaderBuf = dataBuffer.readBytes(RECORD_HEADER_SIZE);
             RecordHeader header = RecordHeader.unmarshal(recordHeaderBuf);
 
-            if (header.getRecordHeaderCRC() != WALUtil.crc32(recordHeaderBuf, RECORD_HEADER_WITHOUT_CRC_SIZE)) {
-                recordHeaderBuf.release();
-                throw new IllegalStateException("Record header crc check failed.");
-            }
             recordHeaderBuf.release();
 
             if (header.getMagicCode() != RecordHeader.RECORD_HEADER_MAGIC_CODE) {
@@ -263,11 +286,6 @@ public RecoverResult next() {
                 dataBuffer.release();
             }
 
-            if (header.getRecordBodyCRC() != WALUtil.crc32(recordBuf)) {
-                recordBuf.release();
-                throw new IllegalStateException("Record body crc check failed.");
-            }
-
             return new RecoverResultImpl(recordBuf, header.getRecordBodyCRC());
         }
     }
diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/util/WALUtil.java b/s3stream/src/main/java/com/automq/stream/s3/wal/util/WALUtil.java
index a69e537e82..8f1b1a9be6 100644
--- a/s3stream/src/main/java/com/automq/stream/s3/wal/util/WALUtil.java
+++ b/s3stream/src/main/java/com/automq/stream/s3/wal/util/WALUtil.java
@@ -11,13 +11,11 @@
 
 package com.automq.stream.s3.wal.util;
 
-import com.automq.stream.s3.ByteBufAlloc;
 import com.automq.stream.s3.wal.common.Record;
 import com.automq.stream.s3.wal.common.RecordHeader;
 import com.automq.stream.utils.CommandResult;
 import com.automq.stream.utils.CommandUtils;
 import io.netty.buffer.ByteBuf;
-import io.netty.buffer.CompositeByteBuf;
 import java.io.File;
 import java.io.IOException;
 import java.io.RandomAccessFile;
@@ -37,28 +35,18 @@ public class WALUtil {
     ));
 
     public static Record generateRecord(ByteBuf body, ByteBuf emptyHeader, int crc, long start) {
-        crc = 0 == crc ? WALUtil.crc32(body) : crc;
-        ByteBuf header = new RecordHeader()
-            .setMagicCode(RECORD_HEADER_MAGIC_CODE)
-            .setRecordBodyLength(body.readableBytes())
-            .setRecordBodyOffset(start + RECORD_HEADER_SIZE)
-            .setRecordBodyCRC(crc)
-            .marshal(emptyHeader);
-        return new Record(header, body);
+        return generateRecord(body, emptyHeader, crc, start, true);
     }
 
-    public static ByteBuf generateRecord(ByteBuf body, int crc, long start) {
-        CompositeByteBuf record = ByteBufAlloc.compositeByteBuffer();
-
+    public static Record generateRecord(ByteBuf body, ByteBuf emptyHeader, int crc, long start, boolean calculateCRC) {
         crc = 0 == crc ? WALUtil.crc32(body) : crc;
         ByteBuf header = new RecordHeader()
             .setMagicCode(RECORD_HEADER_MAGIC_CODE)
             .setRecordBodyLength(body.readableBytes())
             .setRecordBodyOffset(start + RECORD_HEADER_SIZE)
             .setRecordBodyCRC(crc)
-            .marshal(ByteBufAlloc.byteBuffer(RECORD_HEADER_SIZE));
-        record.addComponents(true, header, body);
-        return record;
+            .marshal(emptyHeader, calculateCRC);
+        return new Record(header, body);
     }
 
     /**


Test patch:
diff --git a/s3stream/src/test/java/com/automq/stream/s3/wal/impl/block/BlockWALServiceTest.java b/s3stream/src/test/java/com/automq/stream/s3/wal/impl/block/BlockWALServiceTest.java
index 9561ad9804..afb064479f 100644
--- a/s3stream/src/test/java/com/automq/stream/s3/wal/impl/block/BlockWALServiceTest.java
+++ b/s3stream/src/test/java/com/automq/stream/s3/wal/impl/block/BlockWALServiceTest.java
@@ -903,7 +903,7 @@ private ByteBuf recordHeader(ByteBuf body, long offset) {
             .setRecordBodyLength(body.readableBytes())
             .setRecordBodyOffset(offset + RECORD_HEADER_SIZE)
             .setRecordBodyCRC(WALUtil.crc32(body))
-            .marshal(ByteBufAlloc.byteBuffer(RECORD_HEADER_SIZE));
+            .marshal(ByteBufAlloc.byteBuffer(RECORD_HEADER_SIZE), true);
     }
 
     private void write(WALChannel walChannel, long logicOffset, int recordSize) throws IOException {
diff --git a/s3stream/src/test/java/com/automq/stream/s3/wal/impl/object/ObjectWALServiceTest.java b/s3stream/src/test/java/com/automq/stream/s3/wal/impl/object/ObjectWALServiceTest.java
index 07f6fe5675..3bffe43680 100644
--- a/s3stream/src/test/java/com/automq/stream/s3/wal/impl/object/ObjectWALServiceTest.java
+++ b/s3stream/src/test/java/com/automq/stream/s3/wal/impl/object/ObjectWALServiceTest.java
@@ -45,6 +45,7 @@ public void setUp() throws IOException {
     @AfterEach
     public void tearDown() {
         wal.shutdownGracefully();
+        objectStorage.close();
     }
 
     private ByteBuf generateByteBuf(int size) {
@@ -97,6 +98,7 @@ public void test() throws OverCapacityException, IOException {
 
             ByteBuf recoveredByteBuf = iterator.next().record();
             assertEquals(byteBuf, recoveredByteBuf);
+            recoveredByteBuf.release();
         }
         assertFalse(iterator.hasNext());
 
@@ -108,7 +110,8 @@ public void test() throws OverCapacityException, IOException {
         iterator = wal.recover();
         long count = 0;
         while (iterator.hasNext()) {
-            iterator.next();
+            ByteBuf record = iterator.next().record();
+            record.release();
             count++;
         }
         assertEquals(98, count);
@@ -120,7 +123,8 @@ public void test() throws OverCapacityException, IOException {
         iterator = wal.recover();
         count = 0;
         while (iterator.hasNext()) {
-            iterator.next();
+            ByteBuf record = iterator.next().record();
+            record.release();
             count++;
         }
         assertEquals(97, count);