Merge branch 'AutoMQ:main' into main

Author: brittytino

core/src/main/java/kafka/automq/partition/snapshot/PartitionSnapshotsManager.java

@@ -101,6 +101,7 @@ public void onPartitionClose(Partition partition) {
 
     public CompletableFuture<AutomqGetPartitionSnapshotResponse> handle(AutomqGetPartitionSnapshotRequest request) {
         Session session;
+        boolean newSession = false;
         synchronized (this) {
             AutomqGetPartitionSnapshotRequestData requestData = request.data();
             int sessionId = requestData.sessionId();
@@ -115,10 +116,11 @@ public CompletableFuture<AutomqGetPartitionSnapshotResponse> handle(AutomqGetPar
                 sessionId = nextSessionId();
                 session = new Session(sessionId);
                 sessions.put(sessionId, session);
+                newSession = true;
             }
         }
         CompletableFuture<AutomqGetPartitionSnapshotResponse> resp = session.snapshotsDelta(request.data().version());
-        CompletableFuture<Void> commitCf = request.data().requestCommit() ? confirmWAL.commit(0, false) : CompletableFuture.completedFuture(null);
+        CompletableFuture<Void> commitCf = (request.data().requestCommit() || newSession) ? confirmWAL.commit(0, false) : CompletableFuture.completedFuture(null);
         return commitCf.exceptionally(nil -> null).thenCompose(nil -> resp);
     }
 

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

@@ -360,6 +360,7 @@ void reset(String reason) {
             snapshotWithOperations.clear();
             waitingDataLoadedQueue.clear();
             requester.reset();
+            replayer.reset();
         }
 
         void onNewWalEndOffset(String walConfig, RecordOffset endOffset) {

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

@@ -111,6 +111,11 @@ public void close() {
         }
     }
 
+    public void reset() {
+        this.loadedObjectOrderId = -1L;
+        this.loadedEndOffset = null;
+    }
+
     private List<S3ObjectMetadata> nextObjects() {
         return nextObjects0(metadataCache, node.id(), loadedObjectOrderId, value -> loadedObjectOrderId = value);
     }

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

@@ -170,6 +170,10 @@ private void handleResponse(ClientResponse clientResponse, CompletableFuture<Voi
         if (sessionId != 0 && resp.sessionId() != sessionId) {
             // switch to a new session
             subscriber.reset(String.format("switch sessionId from %s to %s", sessionId, resp.sessionId()));
+            // reset immediately to the new session.
+            tryReset0();
+            sessionId = resp.sessionId();
+            sessionEpoch = resp.sessionEpoch();
         }
         SnapshotReadPartitionsManager.OperationBatch batch = new SnapshotReadPartitionsManager.OperationBatch();
         resp.topics().forEach(topic -> topic.partitions().forEach(partition -> {
@@ -190,8 +194,6 @@ private void handleResponse(ClientResponse clientResponse, CompletableFuture<Voi
         subscriber.onNewWalEndOffset(resp.confirmWalConfig(), DefaultRecordOffset.of(Unpooled.wrappedBuffer(resp.confirmWalEndOffset())));
         batch.operations.add(SnapshotWithOperation.snapshotMark(snapshotCf));
         subscriber.onNewOperationBatch(batch);
-        sessionId = resp.sessionId();
-        sessionEpoch = resp.sessionEpoch();
     }
 
     private boolean tryReset0() {

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

@@ -33,7 +33,9 @@
 import org.apache.kafka.common.message.AutomqZoneRouterRequestData;
 import org.apache.kafka.common.message.MetadataResponseData;
 import org.apache.kafka.common.message.ProduceRequestData;
+import org.apache.kafka.common.protocol.Errors;
 import org.apache.kafka.common.record.MemoryRecords;
+import org.apache.kafka.common.requests.ProduceResponse;
 import org.apache.kafka.common.requests.s3.AutomqZoneRouterResponse;
 import org.apache.kafka.common.utils.LogContext;
 import org.apache.kafka.common.utils.Time;
@@ -154,6 +156,13 @@ public void close() {
 
     @Override
     public void handleProduceRequest(ProduceRequestArgs args) {
+        if (closed.get()) {
+            Map<TopicPartition, ProduceResponse.PartitionResponse> responseMap = new HashMap<>(args.entriesPerPartition().size());
+            args.entriesPerPartition().forEach((tp, records) ->
+                responseMap.put(tp, new ProduceResponse.PartitionResponse(Errors.NOT_LEADER_OR_FOLLOWER)));
+            args.responseCallback().accept(responseMap);
+            return;
+        }
         ClientIdMetadata clientId = args.clientId();
         fillRackIfMissing(clientId);
         if (version.isZeroZoneV2Supported()) {