feat(zerozone2): async start RouterChannel and ConfirmWAL (#2960)

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

Author: superhx

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

@@ -21,8 +21,10 @@
 
 import com.automq.stream.s3.wal.WriteAheadLog;
 
+import java.util.concurrent.CompletableFuture;
+
 public interface ConfirmWALProvider {
 
-    WriteAheadLog readOnly(String walConfig, int nodeId);
+    CompletableFuture<WriteAheadLog> readOnly(String walConfig, int nodeId);
 
 }

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

@@ -31,9 +31,12 @@
 import com.automq.stream.utils.Time;
 
 import java.util.Map;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
 
 public class DefaultConfirmWALProvider implements ConfirmWALProvider {
+    private static final ExecutorService START_EXECUTOR = java.util.concurrent.Executors.newCachedThreadPool();
     private final Map<Short, ObjectStorage> objectStorages = new ConcurrentHashMap<>();
     private final String clusterId;
     private final Time time = Time.SYSTEM;
@@ -43,7 +46,7 @@ public DefaultConfirmWALProvider(String clusterId) {
     }
 
     @Override
-    public WriteAheadLog readOnly(String walConfig, int nodeId) {
+    public CompletableFuture<WriteAheadLog> readOnly(String walConfig, int nodeId) {
         BucketURI bucketURI = BucketURI.parse(walConfig);
         ObjectStorage objectStorage = objectStorages.computeIfAbsent(bucketURI.bucketId(), id -> {
                 try {
@@ -64,6 +67,13 @@ public WriteAheadLog readOnly(String walConfig, int nodeId) {
             .withNodeId(nodeId)
             .withOpenMode(OpenMode.READ_ONLY)
             .build();
-        return new ObjectWALService(time, objectStorage, objectWALConfig);
+        ObjectWALService wal = new ObjectWALService(time, objectStorage, objectWALConfig);
+        return CompletableFuture.runAsync(() -> {
+            try {
+                wal.start();
+            } catch (Throwable e) {
+                throw new RuntimeException(e);
+            }
+        }, START_EXECUTOR).thenApply(v -> wal);
     }
 }

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

@@ -37,7 +37,6 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.util.List;
 import java.util.Map;
@@ -84,11 +83,6 @@ public RouterChannel channel() {
                     .withType(WAL_TYPE)
                     .build();
                 ObjectWALService wal = new ObjectWALService(Time.SYSTEM, objectStorage(), config);
-                try {
-                    wal.start();
-                } catch (IOException e) {
-                    throw new RuntimeException(e);
-                }
                 RouterChannel routerChannel = new ObjectRouterChannel(this.nodeId, channelId, wal);
                 routerChannel.nextEpoch(epoch.getCurrent());
                 routerChannel.trim(epoch.getCommitted());
@@ -106,11 +100,6 @@ public RouterChannel readOnlyChannel(int node) {
         return routerChannels.computeIfAbsent(node, nodeId -> {
             ObjectWALConfig config = ObjectWALConfig.builder().withClusterId(clusterId).withNodeId(node).withOpenMode(OpenMode.READ_ONLY).withType(WAL_TYPE).build();
             ObjectWALService wal = new ObjectWALService(Time.SYSTEM, objectStorage(), config);
-            try {
-                wal.start();
-            } catch (IOException e) {
-                throw new RuntimeException(e);
-            }
             return new ObjectRouterChannel(nodeId, channelId, wal);
         });
     }
@@ -127,8 +116,8 @@ public void addEpochListener(EpochListener listener) {
 
     @Override
     public void close() {
-        FutureUtil.suppress(() -> routerChannel.close(), LOGGER);
-        routerChannels.forEach((nodeId, channel) -> FutureUtil.suppress(channel::close, LOGGER));
+        FutureUtil.suppress(() -> routerChannel.close().get(), LOGGER);
+        routerChannels.forEach((nodeId, channel) -> FutureUtil.suppress(() -> channel.close().get(), LOGGER));
     }
 
     @Override

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

@@ -25,6 +25,7 @@
 import com.automq.stream.s3.wal.exception.OverCapacityException;
 import com.automq.stream.s3.wal.impl.DefaultRecordOffset;
 import com.automq.stream.s3.wal.impl.object.ObjectWALService;
+import com.automq.stream.utils.FutureUtil;
 import com.automq.stream.utils.LogContext;
 
 import org.slf4j.Logger;
@@ -34,12 +35,15 @@
 import java.util.Map;
 import java.util.Queue;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.concurrent.locks.ReentrantReadWriteLock;
 
 import io.netty.buffer.ByteBuf;
 
 public class ObjectRouterChannel implements RouterChannel {
+    private static final ExecutorService ASYNC_EXECUTOR = Executors.newCachedThreadPool();
     private final Logger logger;
     private final AtomicLong mockOffset = new AtomicLong(0);
     private final ReentrantReadWriteLock lock = new ReentrantReadWriteLock();
@@ -54,15 +58,29 @@ public class ObjectRouterChannel implements RouterChannel {
     private final Queue<Long> channelEpochQueue = new LinkedList<>();
     private final Map<Long, RecordOffset> channelEpoch2LastRecordOffset = new HashMap<>();
 
+    private final CompletableFuture<Void> startCf;
+
     public ObjectRouterChannel(int nodeId, short channelId, ObjectWALService wal) {
         this.logger = new LogContext(String.format("[OBJECT_ROUTER_CHANNEL-%s-%s] ", channelId, nodeId)).logger(ObjectRouterChannel.class);
         this.nodeId = nodeId;
         this.channelId = channelId;
         this.wal = wal;
+        this.startCf = CompletableFuture.runAsync(() -> {
+            try {
+                wal.start();
+            } catch (Throwable e) {
+                logger.error("start object router channel failed.", e);
+                throw new RuntimeException(e);
+            }
+        }, ASYNC_EXECUTOR);
     }
 
     @Override
     public CompletableFuture<AppendResult> append(int targetNodeId, short orderHint, ByteBuf data) {
+        return startCf.thenCompose(nil -> append0(targetNodeId, orderHint, data));
+    }
+
+    CompletableFuture<AppendResult> append0(int targetNodeId, short orderHint, ByteBuf data) {
         StreamRecordBatch record = new StreamRecordBatch(targetNodeId, 0, mockOffset.incrementAndGet(), 1, data);
         try {
             return wal.append(TraceContext.DEFAULT, record).thenApply(walRst -> {
@@ -83,6 +101,10 @@ public CompletableFuture<AppendResult> append(int targetNodeId, short orderHint,
 
     @Override
     public CompletableFuture<ByteBuf> get(ByteBuf channelOffset) {
+        return startCf.thenCompose(nil -> get0(channelOffset));
+    }
+
+    CompletableFuture<ByteBuf> get0(ByteBuf channelOffset) {
         return wal.get(DefaultRecordOffset.of(ChannelOffset.of(channelOffset).walRecordOffset())).thenApply(streamRecordBatch -> {
             ByteBuf payload = streamRecordBatch.getPayload().retainedSlice();
             streamRecordBatch.release();
@@ -108,7 +130,7 @@ public void trim(long epoch) {
         writeLock.lock();
         try {
             RecordOffset recordOffset = null;
-            for (;;) {
+            for (; ; ) {
                 Long channelEpoch = channelEpochQueue.peek();
                 if (channelEpoch == null || channelEpoch > epoch) {
                     break;
@@ -129,7 +151,7 @@ public void trim(long epoch) {
     }
 
     @Override
-    public void close() {
-        wal.shutdownGracefully();
+    public CompletableFuture<Void> close() {
+        return startCf.thenAcceptAsync(nil -> FutureUtil.suppress(wal::shutdownGracefully, logger), ASYNC_EXECUTOR);
     }
 }

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

@@ -33,7 +33,7 @@ public interface RouterChannel {
 
     void trim(long epoch);
 
-    void close();
+    CompletableFuture<Void> close();
 
     class AppendResult {
         private final long epoch;

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

@@ -40,6 +40,7 @@
 
 import com.automq.stream.s3.cache.SnapshotReadCache;
 import com.automq.stream.s3.wal.RecordOffset;
+import com.automq.stream.utils.FutureUtil;
 import com.automq.stream.utils.Threads;
 import com.automq.stream.utils.threads.EventLoop;
 
@@ -262,10 +263,12 @@ class Subscriber {
         public Subscriber(Node node, AutoMQVersion version) {
             this.node = node;
             this.version = version;
-            this.requester = new SubscriberRequester(this, node, version, asyncSender, SnapshotReadPartitionsManager.this::getTopicName, eventLoop, time);
             this.replayer = new SubscriberReplayer(confirmWALProvider, SnapshotReadPartitionsManager.this.replayer, node, metadataCache);
-            LOGGER.info("[SNAPSHOT_READ_SUBSCRIBE],node={}", node);
+            this.requester = new SubscriberRequester(this, node, version, asyncSender, SnapshotReadPartitionsManager.this::getTopicName, eventLoop, time);
+            // start the tasks after initialized.
+            this.requester.start();
             run();
+            LOGGER.info("[SNAPSHOT_READ_SUBSCRIBE],node={}", node);
         }
 
         // only for test
@@ -314,9 +317,9 @@ public CompletableFuture<Void> close() {
                     partitions.forEach(SnapshotReadPartitionsManager.this::removePartition);
                     partitions.clear();
                     snapshotWithOperations.clear();
-                    replayer.close();
+                    CompletableFuture<Void> replayerCloseCf = replayer.close();
                     requester.nextSnapshotCf().complete(null);
-                    cf.complete(null);
+                    FutureUtil.propagate(replayerCloseCf, cf);
                 } catch (Throwable e) {
                     cf.completeExceptionally(e);
                 }

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

@@ -38,22 +38,26 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
 import java.util.function.LongConsumer;
 import java.util.stream.Collectors;
 
 class SubscriberReplayer {
     private static final Logger LOGGER = LoggerFactory.getLogger(SubscriberReplayer.class);
+    private static final ExecutorService CLOSE_EXECUTOR = Executors.newCachedThreadPool();
     private long loadedObjectOrderId = -1L;
     private CompletableFuture<Void> lastDataLoadCf = CompletableFuture.completedFuture(null);
-    private WriteAheadLog wal;
+    private CompletableFuture<WriteAheadLog> wal;
     private RecordOffset loadedEndOffset = null;
 
     private final Replayer replayer;
     private final Node node;
     private final MetadataCache metadataCache;
     private final ConfirmWALProvider confirmWALProvider;
 
-    public SubscriberReplayer(ConfirmWALProvider confirmWALProvider, Replayer replayer, Node node, MetadataCache metadataCache) {
+    public SubscriberReplayer(ConfirmWALProvider confirmWALProvider, Replayer replayer, Node node,
+        MetadataCache metadataCache) {
         this.confirmWALProvider = confirmWALProvider;
         this.replayer = replayer;
         this.node = node;
@@ -73,11 +77,11 @@ public void onNewWalEndOffset(String walConfig, RecordOffset endOffset) {
             return;
         }
         // The replayer will ensure the order of replay
-        this.lastDataLoadCf = replayer.replay(wal, startOffset, endOffset).thenAccept(nil -> {
+        this.lastDataLoadCf = wal.thenCompose(w -> replayer.replay(w, startOffset, endOffset).thenAccept(nil -> {
             if (LOGGER.isTraceEnabled()) {
                 LOGGER.trace("replay {} confirm wal [{}, {})", node, startOffset, endOffset);
             }
-        });
+        }));
     }
 
     public CompletableFuture<Void> relayObject() {
@@ -104,11 +108,12 @@ public CompletableFuture<Void> replayWal() {
         return lastDataLoadCf;
     }
 
-    public void close() {
-        WriteAheadLog wal = this.wal;
+    public CompletableFuture<Void> close() {
+        CompletableFuture<WriteAheadLog> wal = this.wal;
         if (wal != null) {
-            FutureUtil.suppress(wal::shutdownGracefully, LOGGER);
+            return CompletableFuture.runAsync(() -> FutureUtil.suppress(() -> wal.get().shutdownGracefully(), LOGGER), CLOSE_EXECUTOR);
         }
+        return CompletableFuture.completedFuture(null);
     }
 
     public void reset() {

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

@@ -81,6 +81,9 @@ public SubscriberRequester(SnapshotReadPartitionsManager.Subscriber subscriber,
         this.topicNameGetter = topicNameGetter;
         this.eventLoop = eventLoop;
         this.time = time;
+    }
+
+    public void start() {
         request();
     }
 
@@ -126,7 +129,11 @@ private void request0() {
         AutomqGetPartitionSnapshotRequest.Builder builder = new AutomqGetPartitionSnapshotRequest.Builder(data);
         asyncSender.sendRequest(node, builder)
             .thenAcceptAsync(rst -> {
-                handleResponse(rst, snapshotCf);
+                try {
+                    handleResponse(rst, snapshotCf);
+                } catch (Exception e) {
+                    subscriber.reset("Exception when handle snapshot response: " + e.getMessage());
+                }
                 subscriber.unsafeRun();
             }, eventLoop)
             .exceptionally(ex -> {