feat(controller): add has opening streams method (#1707)

Signed-off-by: Robin Han <[email protected]>
AutoMQ · Aug 2, 2024 · be31452 · be31452
1 parent 756fe41
commit be31452
Show file tree

Hide file tree

Showing 3 changed files with 50 additions and 1 deletion.
diff --git a/metadata/src/main/java/org/apache/kafka/controller/stream/DefaultNodeRuntimeInfoGetter.java b/metadata/src/main/java/org/apache/kafka/controller/stream/DefaultNodeRuntimeInfoGetter.java
@@ -40,6 +40,6 @@ public NodeState state(int nodeId) {
 
     @Override
     public boolean hasOpeningStreams(int nodeId) {
-        return !streamControlManager.getOpeningStreams(nodeId).isEmpty();
+        return streamControlManager.hasOpeningStreams(nodeId);
     }
 }
diff --git a/metadata/src/main/java/org/apache/kafka/controller/stream/StreamControlManager.java b/metadata/src/main/java/org/apache/kafka/controller/stream/StreamControlManager.java
@@ -37,6 +37,8 @@
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.function.Function;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 import org.apache.kafka.common.Uuid;
@@ -1077,6 +1079,29 @@ public List<StreamRuntimeMetadata> getOpeningStreams(int nodeId) {
         return streams;
     }
 
+    public boolean hasOpeningStreams(int nodeId) {
+        DeltaList<Long> streamIdList = node2streams.get(nodeId);
+        if (streamIdList == null) {
+            return false;
+        }
+        AtomicBoolean hasOpeningStreams = new AtomicBoolean(false);
+        streamIdList.reverseForEachWithBreak(new Function<Long, Boolean>() {
+            @Override
+            public Boolean apply(Long streamId) {
+                StreamRuntimeMetadata streamRuntimeMetadata = streamsMetadata.get(streamId);
+                if (streamRuntimeMetadata == null) {
+                    return false;
+                }
+                if (streamRuntimeMetadata.currentState() == StreamState.OPENED) {
+                    hasOpeningStreams.set(true);
+                    return true;
+                }
+                return false;
+            }
+        });
+        return hasOpeningStreams.get();
+    }
+
     /**
      * Check whether this node is the owner of this stream.
      */

diff --git a/metadata/src/main/java/org/apache/kafka/image/DeltaList.java b/metadata/src/main/java/org/apache/kafka/image/DeltaList.java
@@ -16,6 +16,7 @@
 import java.util.List;
 import java.util.Objects;
 import java.util.function.Consumer;
+import java.util.function.Function;
 import java.util.stream.Collectors;
 
 public class DeltaList<T> {
@@ -115,6 +116,29 @@ public void reverseForEach(Consumer<T> consumer) {
         }
     }
 
+    /**
+     * Reverse iterate the list
+     *
+     * @param func accept a list element, return break loop mark.
+     */
+    public void reverseForEachWithBreak(Function<T, Boolean> func) {
+        synchronized (operations) {
+            List<Operation<T>> removedList = new ArrayList<>();
+            for (int i = snapshotIndex - 1; i >= 0; i--) {
+                Operation<T> operation = operations.get(i);
+                if (operation.tombstoneMatcher != null) {
+                    removedList.add(operation);
+                } else {
+                    if (!isRemoved(operation, removedList)) {
+                        if (func.apply(operation.t)) {
+                            return;
+                        }
+                    }
+                }
+            }
+        }
+    }
+
     public void forEach(Consumer<T> consumer) {
         synchronized (operations) {
             List<Operation<T>> removedList = new ArrayList<>();