fix(image): guard streams image access with lock to prevent data loss (#2653)

fix(image): guard streams image access with lock to prevent compaction skip data

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

Author: superhx

core/src/main/scala/kafka/log/stream/s3/metadata/StreamMetadataManager.java

@@ -221,19 +221,21 @@ public List<StreamMetadata> getStreamMetadataList(List<Long> streamIds) {
         try (Image image = getImage()) {
             final S3StreamsMetadataImage streamsImage = image.streamsMetadata();
 
-            List<StreamMetadata> streamMetadataList = new ArrayList<>();
-            for (Long streamId : streamIds) {
-                S3StreamMetadataImage streamImage = streamsImage.timelineStreamMetadata().get(streamId);
-                if (streamImage == null) {
-                    LOGGER.warn("[GetStreamMetadataList]: stream: {} not exists", streamId);
-                    continue;
+            List<StreamMetadata> streamMetadataList = new ArrayList<>(streamIds.size());
+            streamsImage.inLockRun(() -> {
+                for (Long streamId : streamIds) {
+                    S3StreamMetadataImage streamImage = streamsImage.timelineStreamMetadata().get(streamId);
+                    if (streamImage == null) {
+                        LOGGER.warn("[GetStreamMetadataList]: stream: {} not exists", streamId);
+                        continue;
+                    }
+                    // If there is a streamImage, it means the stream exists.
+                    @SuppressWarnings("OptionalGetWithoutIsPresent") long endOffset = streamsImage.streamEndOffset(streamId).getAsLong();
+                    StreamMetadata streamMetadata = new StreamMetadata(streamId, streamImage.getEpoch(),
+                        streamImage.getStartOffset(), endOffset, streamImage.state());
+                    streamMetadataList.add(streamMetadata);
                 }
-                // If there is a streamImage, it means the stream exists.
-                @SuppressWarnings("OptionalGetWithoutIsPresent") long endOffset = streamsImage.streamEndOffset(streamId).getAsLong();
-                StreamMetadata streamMetadata = new StreamMetadata(streamId, streamImage.getEpoch(),
-                    streamImage.getStartOffset(), endOffset, streamImage.state());
-                streamMetadataList.add(streamMetadata);
-            }
+            });
             return streamMetadataList;
         }
     }

metadata/src/main/java/org/apache/kafka/image/S3StreamsMetadataImage.java

@@ -595,6 +595,7 @@ public int hashCode() {
         return Objects.hash(nextAssignedStreamId, streamMetadataList(), nodeMetadataList(), streamEndOffsets());
     }
 
+    // caller use this value should be protected by registryRef lock
     public TimelineHashMap<Integer, NodeS3StreamSetObjectMetadataImage> timelineNodeMetadata() {
         return nodeMetadataMap;
     }
@@ -610,10 +611,15 @@ List<NodeS3StreamSetObjectMetadataImage> nodeMetadataList() {
         });
     }
 
+    // caller use this value should be protected by registryRef lock
     public TimelineHashMap<Long, S3StreamMetadataImage> timelineStreamMetadata() {
         return streamMetadataMap;
     }
 
+    public void inLockRun(Runnable runnable) {
+        registryRef.inLock(runnable);
+    }
+
     List<S3StreamMetadataImage> streamMetadataList() {
         if (registryRef == RegistryRef.NOOP) {
             return Collections.emptyList();
@@ -630,30 +636,33 @@ public long nextAssignedStreamId() {
     }
 
     public OptionalLong streamEndOffset(long streamId) {
-        Long endOffset = streamEndOffsets.get(streamId);
-        if (endOffset != null) {
-            return OptionalLong.of(endOffset);
-        }
-        // There is no record in a new stream
-        if (streamMetadataMap.containsKey(streamId)) {
-            return OptionalLong.of(0L);
-        } else {
+        if (registryRef == RegistryRef.NOOP) {
             return OptionalLong.empty();
         }
+        return registryRef.inLock(() -> {
+            Long endOffset = streamEndOffsets.get(streamId);
+            if (endOffset != null) {
+                return OptionalLong.of(endOffset);
+            }
+            // There is no record in a new stream
+            if (streamMetadataMap.containsKey(streamId)) {
+                return OptionalLong.of(0L);
+            } else {
+                return OptionalLong.empty();
+            }
+        });
     }
 
+    // caller use this value should be protected by registryRef lock
     TimelineHashMap<TopicIdPartition, Set<Long>> partition2streams() {
         return partition2streams;
     }
 
+    // caller use this value should be protected by registryRef lock
     TimelineHashMap<Long, TopicIdPartition> stream2partition() {
         return stream2partition;
     }
 
-    RegistryRef registryRef() {
-        return registryRef;
-    }
-
     // caller use this value should be protected by registryRef lock
     TimelineHashMap<Long, Long> timelineStreamEndOffsets() {
         return streamEndOffsets;
@@ -671,6 +680,10 @@ Map<Long, Long> streamEndOffsets() {
         });
     }
 
+    RegistryRef registryRef() {
+        return registryRef;
+    }
+
     @Override
     public String toString() {
         return "S3StreamsMetadataImage{nextAssignedStreamId=" + nextAssignedStreamId + '}';