fix(zerozone2): fix the upgrade from v1 to v2 (#2985)

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

Author: superhx

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

@@ -171,7 +171,7 @@ public synchronized CompletableFuture<AutomqGetPartitionSnapshotResponse> snapsh
             long finalSessionEpoch = sessionEpoch;
             CompletableFuture<Void> collectPartitionSnapshotsCf;
             if (!requestCommit && inflightCommitCfSet.isEmpty()) {
-                collectPartitionSnapshotsCf = collectPartitionSnapshots(resp);
+                collectPartitionSnapshotsCf = collectPartitionSnapshots(requestVersion, resp);
             } else {
                 collectPartitionSnapshotsCf = CompletableFuture.completedFuture(null);
             }
@@ -203,15 +203,15 @@ public synchronized boolean expired() {
             return time.milliseconds() - lastGetSnapshotsTimestamp > 60000;
         }
 
-        private CompletableFuture<Void> collectPartitionSnapshots(AutomqGetPartitionSnapshotResponseData resp) {
+        private CompletableFuture<Void> collectPartitionSnapshots(short requestVersion, AutomqGetPartitionSnapshotResponseData resp) {
             Map<Uuid, List<PartitionSnapshot>> topic2partitions = new HashMap<>();
             List<CompletableFuture<Void>> completeCfList = COMPLETE_CF_LIST_LOCAL.get();
             completeCfList.clear();
             removed.forEach(partition -> {
                 PartitionSnapshotVersion version = synced.remove(partition);
                 if (version != null) {
                     List<PartitionSnapshot> partitionSnapshots = topic2partitions.computeIfAbsent(partition.topicId().get(), topic -> new ArrayList<>());
-                    partitionSnapshots.add(snapshot(partition, version, null, completeCfList));
+                    partitionSnapshots.add(snapshot(requestVersion, partition, version, null, completeCfList));
                 }
             });
             removed.clear();
@@ -221,7 +221,7 @@ private CompletableFuture<Void> collectPartitionSnapshots(AutomqGetPartitionSnap
                 if (!Objects.equals(p.version, oldVersion)) {
                     List<PartitionSnapshot> partitionSnapshots = topic2partitions.computeIfAbsent(p.partition.topicId().get(), topic -> new ArrayList<>());
                     PartitionSnapshotVersion newVersion = p.version.copy();
-                    PartitionSnapshot partitionSnapshot = snapshot(p.partition, oldVersion, newVersion, completeCfList);
+                    PartitionSnapshot partitionSnapshot = snapshot(requestVersion, p.partition, oldVersion, newVersion, completeCfList);
                     partitionSnapshots.add(partitionSnapshot);
                     synced.put(p.partition, newVersion);
                 }
@@ -239,7 +239,7 @@ private CompletableFuture<Void> collectPartitionSnapshots(AutomqGetPartitionSnap
             return retCf;
         }
 
-        private PartitionSnapshot snapshot(Partition partition, PartitionSnapshotVersion oldVersion,
+        private PartitionSnapshot snapshot(short requestVersion, Partition partition, PartitionSnapshotVersion oldVersion,
             PartitionSnapshotVersion newVersion, List<CompletableFuture<Void>> completeCfList) {
             if (newVersion == null) {
                 // partition is closed
@@ -268,7 +268,9 @@ private PartitionSnapshot snapshot(Partition partition, PartitionSnapshotVersion
                 if (includeSegments) {
                     snapshot.setLogMetadata(logMetadata(src.logMeta()));
                 }
-                snapshot.setLastTimestampOffset(timestampOffset(src.lastTimestampOffset()));
+                if (requestVersion > ZERO_ZONE_V0_REQUEST_VERSION) {
+                    snapshot.setLastTimestampOffset(timestampOffset(src.lastTimestampOffset()));
+                }
                 return snapshot;
             });
         }

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

@@ -118,7 +118,7 @@ public BrokersAsyncSender(
                 brokerConfig.connectionSetupTimeoutMs(),
                 brokerConfig.connectionSetupTimeoutMaxMs(),
                 time,
-                false,
+                true,
                 new ApiVersions(),
                 logContext,
                 MetadataRecoveryStrategy.REBOOTSTRAP

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

@@ -112,11 +112,7 @@ public CompletableFuture<AutomqZoneRouterResponse> handleZoneRouterRequest0(Rout
             .thenCompose(rst -> prevLastRouterCf.thenApply(nil -> rst))
             .thenComposeAsync(produces -> {
                 List<CompletableFuture<AutomqZoneRouterResponseData.Response>> cfList = new ArrayList<>();
-                produces.stream().map(request -> {
-                    try (request) {
-                        return append(request);
-                    }
-                }).forEach(cfList::add);
+                produces.stream().map(this::append).forEach(cfList::add);
                 return CompletableFuture.allOf(cfList.toArray(new CompletableFuture[0])).thenApply(nil -> {
                     AutomqZoneRouterResponseData response = new AutomqZoneRouterResponseData();
                     cfList.forEach(cf -> response.responses().add(cf.join()));

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

@@ -115,7 +115,7 @@ private void request0() {
 
         tryReset0();
         lastRequestTime = time.milliseconds();
-        AutomqGetPartitionSnapshotRequestData data = new AutomqGetPartitionSnapshotRequestData().setSessionId(sessionId).setSessionEpoch(sessionEpoch).setVersion((short) 1);
+        AutomqGetPartitionSnapshotRequestData data = new AutomqGetPartitionSnapshotRequestData().setSessionId(sessionId).setSessionEpoch(sessionEpoch);
         if (version.isZeroZoneV2Supported()) {
             data.setVersion((short) 1);
         }
@@ -198,7 +198,10 @@ private void handleResponse(ClientResponse clientResponse, CompletableFuture<Voi
             int c2 = o2.operation.code() == SnapshotOperation.REMOVE.code() ? 0 : 1;
             return c1 - c2;
         });
-        subscriber.onNewWalEndOffset(resp.confirmWalConfig(), DefaultRecordOffset.of(Unpooled.wrappedBuffer(resp.confirmWalEndOffset())));
+        if (resp.confirmWalEndOffset() != null && resp.confirmWalEndOffset().length > 0) {
+            // zerozone v2
+            subscriber.onNewWalEndOffset(resp.confirmWalConfig(), DefaultRecordOffset.of(Unpooled.wrappedBuffer(resp.confirmWalEndOffset())));
+        }
         batch.operations.add(SnapshotWithOperation.snapshotMark(snapshotCf));
         subscriber.onNewOperationBatch(batch);
     }