feat(zerozone2): S3WAL supports sequential append and distributed read (#2764)

- S3WAL supports sequential append & callback
- Simple the S3Storage logic by S3WAl sequential append
- S3WAL supports distributed read from another nodes
- S3Storage supports linked record

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

Author: superhx

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

@@ -51,7 +51,7 @@
 import org.apache.kafka.metadata.stream.S3StreamSetObject;
 import org.apache.kafka.storage.internals.log.LogOffsetMetadata;
 
-import com.automq.stream.s3.cache.SnapshotReadCache;
+import com.automq.stream.Context;
 import com.automq.stream.s3.metadata.S3ObjectMetadata;
 import com.automq.stream.utils.Threads;
 import com.automq.stream.utils.threads.EventLoop;
@@ -103,7 +103,8 @@ public SnapshotReadPartitionsManager(KafkaConfig config, Metrics metrics, Time t
         this.time = time;
         this.replicaManager = replicaManager;
         this.metadataCache = metadataCache;
-        this.dataLoader = objects -> SnapshotReadCache.instance().load(objects);
+        // FIXME: tmp code in merging
+        this.dataLoader = objects -> Context.instance().snapshotReadCache().replay(objects);
         this.asyncSender = new AsyncSender.BrokersAsyncSender(config, metrics, "snapshot_read", Time.SYSTEM, "AUTOMQ_SNAPSHOT_READ", new LogContext());
     }
 

core/src/main/scala/kafka/log/stream/s3/wal/BootstrapWalV1.java

@@ -25,8 +25,11 @@
 import org.apache.kafka.controller.stream.NodeMetadata;
 
 import com.automq.stream.s3.exceptions.AutoMQException;
+import com.automq.stream.s3.model.StreamRecordBatch;
 import com.automq.stream.s3.trace.context.TraceContext;
 import com.automq.stream.s3.wal.AppendResult;
+import com.automq.stream.s3.wal.OpenMode;
+import com.automq.stream.s3.wal.RecordOffset;
 import com.automq.stream.s3.wal.RecoverResult;
 import com.automq.stream.s3.wal.WalFactory;
 import com.automq.stream.s3.wal.WalFactory.BuildOptions;
@@ -44,11 +47,12 @@
 import java.io.IOException;
 import java.util.Collections;
 import java.util.Iterator;
+import java.util.List;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.TimeUnit;
 
-import io.netty.buffer.ByteBuf;
+
 
 public class BootstrapWalV1 implements WriteAheadLog {
     private static final Logger LOGGER = LoggerFactory.getLogger(BootstrapWalV1.class);
@@ -114,8 +118,23 @@ public WALMetadata metadata() {
     }
 
     @Override
-    public AppendResult append(TraceContext context, ByteBuf data, int crc) throws OverCapacityException {
-        return wal.append(context, data, crc);
+    public CompletableFuture<AppendResult> append(TraceContext context, StreamRecordBatch streamRecordBatch) throws OverCapacityException {
+        return wal.append(context, streamRecordBatch);
+    }
+
+    @Override
+    public CompletableFuture<StreamRecordBatch> get(RecordOffset recordOffset) {
+        return wal.get(recordOffset);
+    }
+
+    @Override
+    public CompletableFuture<List<StreamRecordBatch>> get(RecordOffset startOffset, RecordOffset endOffset) {
+        return wal.get(startOffset, endOffset);
+    }
+
+    @Override
+    public RecordOffset confirmOffset() {
+        return wal.confirmOffset();
     }
 
     @Override
@@ -157,15 +176,15 @@ public CompletableFuture<Void> reset() {
     }
 
     @Override
-    public CompletableFuture<Void> trim(long offset) {
+    public CompletableFuture<Void> trim(RecordOffset offset) {
         return wal.trim(offset);
     }
 
     private CompletableFuture<? extends WriteAheadLog> buildRecoverWal(String kraftWalConfigs, long oldNodeEpoch) {
         IdURI uri = IdURI.parse(kraftWalConfigs);
         CompletableFuture<Void> cf = walHandle
             .acquirePermission(nodeId, oldNodeEpoch, uri, new WalHandle.AcquirePermissionOptions().failoverMode(failoverMode));
-        return cf.thenApplyAsync(nil -> factory.build(uri, BuildOptions.builder().nodeEpoch(oldNodeEpoch).failoverMode(failoverMode).build()), executor);
+        return cf.thenApplyAsync(nil -> factory.build(uri, BuildOptions.builder().nodeEpoch(oldNodeEpoch).openMode(failoverMode ? OpenMode.FAILOVER : OpenMode.READ_WRITE).build()), executor);
     }
 
     private CompletableFuture<? extends WriteAheadLog> buildWal(String kraftWalConfigs) {
@@ -175,7 +194,7 @@ private CompletableFuture<? extends WriteAheadLog> buildWal(String kraftWalConfi
             .failoverMode(false);
         CompletableFuture<Void> cf = walHandle
             .acquirePermission(nodeId, nodeEpoch, uri, options);
-        return cf.thenApplyAsync(nil -> factory.build(uri, BuildOptions.builder().nodeEpoch(nodeEpoch).failoverMode(false).build()), executor);
+        return cf.thenApplyAsync(nil -> factory.build(uri, BuildOptions.builder().nodeEpoch(nodeEpoch).openMode(OpenMode.READ_WRITE).build()), executor);
     }
 
     private CompletableFuture<Void> releasePermission(String kraftWalConfigs) {

core/src/main/scala/kafka/log/stream/s3/wal/DefaultWalFactory.java

@@ -24,8 +24,10 @@
 import com.automq.stream.s3.operator.BucketURI;
 import com.automq.stream.s3.operator.ObjectStorage;
 import com.automq.stream.s3.operator.ObjectStorageFactory;
+import com.automq.stream.s3.wal.ReservationService;
 import com.automq.stream.s3.wal.WalFactory;
 import com.automq.stream.s3.wal.WriteAheadLog;
+import com.automq.stream.s3.wal.impl.object.ObjectReservationService;
 import com.automq.stream.s3.wal.impl.object.ObjectWALConfig;
 import com.automq.stream.s3.wal.impl.object.ObjectWALService;
 import com.automq.stream.utils.IdURI;
@@ -65,8 +67,9 @@ public WriteAheadLog build(IdURI uri, BuildOptions options) {
                     .withClusterId(AutoMQApplication.getClusterId())
                     .withNodeId(nodeId)
                     .withEpoch(options.nodeEpoch())
-                    .withFailover(options.failoverMode());
-
+                    .withOpenMode(options.openMode());
+                ReservationService reservationService = new ObjectReservationService(AutoMQApplication.getClusterId(), walObjectStorage, walObjectStorage.bucketId());
+                configBuilder.withReservationService(reservationService);
                 return new ObjectWALService(Time.SYSTEM, walObjectStorage, configBuilder.build());
             default:
                 throw new IllegalArgumentException("Unsupported WAL protocol: " + uri.protocol());

core/src/main/scala/kafka/log/streamaspect/ClientWrapper.java

@@ -231,6 +231,11 @@ public long confirmOffset() {
             return stream.confirmOffset();
         }
 
+        @Override
+        public void confirmOffset(long offset) {
+            stream.confirmOffset(offset);
+        }
+
         @Override
         public long nextOffset() {
             return stream.nextOffset();

core/src/main/scala/kafka/log/streamaspect/LazyStream.java

@@ -114,6 +114,11 @@ public long confirmOffset() {
         return inner.confirmOffset();
     }
 
+    @Override
+    public void confirmOffset(long offset) {
+        inner.confirmOffset(offset);
+    }
+
     @Override
     public long nextOffset() {
         return inner.nextOffset();
@@ -199,6 +204,10 @@ public long confirmOffset() {
             return 0;
         }
 
+        @Override
+        public void confirmOffset(long offset) {
+        }
+
         @Override
         public long nextOffset() {
             return 0;

core/src/main/scala/kafka/log/streamaspect/MemoryClient.java

@@ -109,6 +109,11 @@ public long confirmOffset() {
             return nextOffsetAlloc.get();
         }
 
+        @Override
+        public void confirmOffset(long offset) {
+            nextOffsetAlloc.set(offset);
+        }
+
         @Override
         public long nextOffset() {
             return nextOffsetAlloc.get();

core/src/main/scala/kafka/log/streamaspect/MetaStream.java

@@ -117,6 +117,11 @@ public long confirmOffset() {
         return innerStream.confirmOffset();
     }
 
+    @Override
+    public void confirmOffset(long offset) {
+        innerStream.confirmOffset(offset);
+    }
+
     @Override
     public long nextOffset() {
         return innerStream.nextOffset();

s3stream/src/main/java/com/automq/stream/Context.java

@@ -0,0 +1,61 @@
+/*
+ * Copyright 2025, AutoMQ HK Limited.
+ *
+ * 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 com.automq.stream;
+
+import com.automq.stream.api.KVClient;
+import com.automq.stream.s3.ConfirmWAL;
+import com.automq.stream.s3.cache.SnapshotReadCache;
+
+public class Context {
+    private SnapshotReadCache snapshotReadCache;
+    private ConfirmWAL confirmWAL;
+    private KVClient kvClient;
+
+    public static final Context INSTANCE = new Context();
+
+    public static Context instance() {
+        return INSTANCE;
+    }
+
+    public KVClient kvClient() {
+        return kvClient;
+    }
+
+    public void kvClient(KVClient kvClient) {
+        this.kvClient = kvClient;
+    }
+
+    public void snapshotReadCache(SnapshotReadCache snapshotReadCache) {
+        this.snapshotReadCache = snapshotReadCache;
+    }
+
+    public SnapshotReadCache snapshotReadCache() {
+        return snapshotReadCache;
+    }
+
+    public void confirmWAL(ConfirmWAL confirmWAL) {
+        this.confirmWAL = confirmWAL;
+    }
+
+    public ConfirmWAL confirmWAL() {
+        return confirmWAL;
+    }
+
+}

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

@@ -50,6 +50,11 @@ public interface Stream {
      */
     long confirmOffset();
 
+    /**
+     * Set confirm offset. Only support in snapshotRead mode
+     */
+    void confirmOffset(long offset);
+
     /**
      * Get stream next append record offset.
      */

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

@@ -0,0 +1,55 @@
+/*
+ * Copyright 2025, AutoMQ HK Limited.
+ *
+ * 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 com.automq.stream.s3;
+
+import com.automq.stream.s3.S3Storage.LazyCommit;
+import com.automq.stream.s3.wal.RecordOffset;
+import com.automq.stream.s3.wal.WriteAheadLog;
+
+import java.util.concurrent.CompletableFuture;
+import java.util.function.Function;
+
+public class ConfirmWAL {
+    private final WriteAheadLog log;
+    private final Function<LazyCommit, CompletableFuture<Void>> commitHandle;
+
+    public ConfirmWAL(WriteAheadLog log, Function<LazyCommit, CompletableFuture<Void>> commitHandle) {
+        this.log = log;
+        this.commitHandle = commitHandle;
+    }
+
+    public RecordOffset confirmOffset() {
+        return log.confirmOffset();
+    }
+
+    /**
+     * Commit with lazy timeout.
+     * If in [0, lazyLingerMs), there is no other commit happened, then trigger a new commit.
+     * @param lazyLingerMs lazy linger milliseconds.
+     */
+    public CompletableFuture<Void> commit(long lazyLingerMs, boolean awaitTrim) {
+        return commitHandle.apply(new LazyCommit(lazyLingerMs, awaitTrim));
+    }
+
+    public CompletableFuture<Void> commit(long lazyLingerMs) {
+        return commit(lazyLingerMs, true);
+    }
+
+}