fix: close flat mapper iterator after usage

apache · Aug 14, 2023 · 3d343d0 · 3d343d0
1 parent b02c2bd
commit 3d343d0
Showing 1 changed file with 21 additions and 11 deletions.
diff --git a/hugegraph-core/src/main/java/org/apache/hugegraph/traversal/algorithm/CountTraverser.java b/hugegraph-core/src/main/java/org/apache/hugegraph/traversal/algorithm/CountTraverser.java
@@ -22,6 +22,7 @@
 import java.util.Set;
 
 import org.apache.commons.lang.mutable.MutableLong;
+import org.apache.hugegraph.HugeException;
 import org.apache.hugegraph.HugeGraph;
 import org.apache.hugegraph.backend.id.Id;
 import org.apache.hugegraph.backend.query.QueryResults;
@@ -78,19 +79,28 @@ public long count(Id source, List<EdgeStep> steps,
             });
         }
 
-        // The last step, just query count
-        EdgeStep lastStep = steps.get(stepNum - 1);
-        while (edges.hasNext()) {
-            Id target = ((HugeEdge) edges.next()).id().otherVertexId();
-            if (this.dedup(target)) {
-                continue;
+        try {
+            // The last step, just query count
+            EdgeStep lastStep = steps.get(stepNum - 1);
+            while (edges.hasNext()) {
+                Id target = ((HugeEdge) edges.next()).id().otherVertexId();
+                if (this.dedup(target)) {
+                    continue;
+                }
+                // Count last layer vertices(without dedup size)
+                long edgesCount = this.edgesCount(target, lastStep);
+                this.count.add(edgesCount);
             }
-            // Count last layer vertices(without dedup size)
-            long edgesCount = this.edgesCount(target, lastStep);
-            this.count.add(edgesCount);
-        }
 
-        return this.count.longValue();
+            return this.count.longValue();
+        } finally {
+            if (edges instanceof FlatMapperIterator) {
+                try {
+                    ((FlatMapperIterator<?, ?>)edges).close();
+                } catch (Exception ignored) {
+                }
+            }
+        }
     }
 
     private Iterator<Edge> edgesOfVertexWithCount(Id source, EdgeStep step) {