Max-flow supports termination flag

IoannisPanagiotas · AlfredClemedtson · IoannisPanagiotas · commit 9b292e3ddee2 · 2025-10-07T10:41:56.000+02:00
Co-authored-by: Alfred Clemedtson &lt;alfred.clemedtson@neo4j.com&gt;
diff --git a/algo/src/main/java/org/neo4j/gds/maxflow/Discharging.java b/algo/src/main/java/org/neo4j/gds/maxflow/Discharging.java
@@ -26,6 +26,7 @@
 import org.neo4j.gds.core.concurrency.RunWithConcurrency;
 import org.neo4j.gds.core.utils.paged.HugeAtomicBitSet;
 import org.neo4j.gds.core.utils.paged.HugeLongArrayQueue;
+import org.neo4j.gds.termination.TerminationFlag;
 
 import java.util.ArrayList;
 import java.util.Collection;
@@ -36,6 +37,7 @@ final class Discharging {
     private final AtomicWorkingSet workingSet;
     private final Concurrency concurrency;
     private final Collection<Runnable> dischargeTasks;
+    private final  TerminationFlag terminationFlag;
 
 
     static Discharging createDischarging(
@@ -48,7 +50,8 @@ static Discharging createDischarging(
         long beta,
         AtomicLong workSinceLastGR,
         Concurrency concurrency,
-        HugeLongArrayQueue[] threadQueues
+        HugeLongArrayQueue[] threadQueues,
+        TerminationFlag terminationFlag
     )
     {
 
@@ -72,25 +75,38 @@ static Discharging createDischarging(
             ));
         }
 
-        return new Discharging(workingSet, dischargeTasks, concurrency);
+        return new Discharging(workingSet, dischargeTasks, concurrency, terminationFlag);
     }
 
-    private Discharging(AtomicWorkingSet workingSet, Collection<Runnable> dischargeTasks, Concurrency concurrency) {
+    private Discharging(AtomicWorkingSet workingSet, Collection<Runnable> dischargeTasks, Concurrency concurrency,
+        TerminationFlag terminationFlag
+    ) {
         this.workingSet = workingSet;
         this.dischargeTasks = dischargeTasks;
         this.concurrency = concurrency;
+        this.terminationFlag = terminationFlag;
     }
 
     void processWorkingSet() {
         //Discharge working set
-        RunWithConcurrency.builder().concurrency(concurrency).tasks(dischargeTasks).build().run();
+        runTasks();
         workingSet.resetIdx();
 
         //Sync working set
-        RunWithConcurrency.builder().concurrency(concurrency).tasks(dischargeTasks).build().run();
+        runTasks();
         workingSet.reset();
 
         //Update and sync new working set
-        RunWithConcurrency.builder().concurrency(concurrency).tasks(dischargeTasks).build().run();
+        runTasks();
+
+    }
+
+    private void runTasks(){
+        RunWithConcurrency.builder()
+            .concurrency(concurrency)
+            .terminationFlag(terminationFlag)
+            .tasks(dischargeTasks)
+            .build()
+            .run();
     }
 }
diff --git a/algo/src/main/java/org/neo4j/gds/maxflow/FlowGraph.java b/algo/src/main/java/org/neo4j/gds/maxflow/FlowGraph.java
@@ -26,6 +26,7 @@
 import org.neo4j.gds.collections.ha.HugeDoubleArray;
 import org.neo4j.gds.collections.ha.HugeLongArray;
 import org.neo4j.gds.collections.ha.HugeObjectArray;
+import org.neo4j.gds.termination.TerminationFlag;
 
 public final class FlowGraph {
     private final Graph graph;
@@ -60,7 +61,7 @@ private FlowGraph(
         this.demand = demand;
     }
 
-    public static FlowGraph create(Graph graph, NodeWithValue[] supply, NodeWithValue[] demand) {
+    public static FlowGraph create(Graph graph, NodeWithValue[] supply, NodeWithValue[] demand, TerminationFlag terminationFlag) {
         var superSource = graph.nodeCount();
         var superTarget = graph.nodeCount() + 1;
         var newNodeCount = graph.nodeCount() + 2;
@@ -69,6 +70,7 @@ public static FlowGraph create(Graph graph, NodeWithValue[] supply, NodeWithValu
         reverseDegree.setAll(x -> 0L);
 
         for (long nodeId = 0; nodeId < graph.nodeCount(); nodeId++) {
+            terminationFlag.assertRunning();
             graph.forEachRelationship(
                 nodeId, 0D, (s, t, capacity) -> {
                     if(capacity < 0D){
@@ -120,9 +122,11 @@ public static FlowGraph create(Graph graph, NodeWithValue[] supply, NodeWithValu
             return true;
         };
         for (long nodeId = 0; nodeId < graph.nodeCount(); nodeId++) {
+            terminationFlag.assertRunning();
             graph.forEachRelationship(nodeId, 0D, consumer);
         }
         for (var source : supply) {
+            terminationFlag.assertRunning();
             consumer.accept(superSource, source.node(), source.value());
         }
         for (var target : demand) {
diff --git a/algo/src/main/java/org/neo4j/gds/maxflow/GlobalRelabeling.java b/algo/src/main/java/org/neo4j/gds/maxflow/GlobalRelabeling.java
@@ -24,6 +24,7 @@
 import org.neo4j.gds.core.concurrency.RunWithConcurrency;
 import org.neo4j.gds.core.utils.paged.HugeAtomicBitSet;
 import org.neo4j.gds.core.utils.paged.HugeLongArrayQueue;
+import org.neo4j.gds.termination.TerminationFlag;
 
 import java.util.ArrayList;
 import java.util.Collection;
@@ -38,14 +39,16 @@ public final class GlobalRelabeling {
     private final long target;
     private final Concurrency concurrency;
     private final Collection<Runnable> globalRelabelingTasks;
+    private final TerminationFlag terminationFlag;
 
     static GlobalRelabeling createRelabeling(
         FlowGraph flowGraph,
         HugeLongArray label,
         long source,
         long target,
         Concurrency concurrency,
-        HugeLongArrayQueue[] threadQueues
+        HugeLongArrayQueue[] threadQueues,
+        TerminationFlag terminationFlag
     ) {
         var vertexIsDiscovered = HugeAtomicBitSet.create(flowGraph.nodeCount());
         var frontier = new AtomicWorkingSet(flowGraph.nodeCount());
@@ -55,10 +58,22 @@ static GlobalRelabeling createRelabeling(
             globalRelabelingTasks.add(new GlobalRelabellingBFSTask(flowGraph.concurrentCopy(), frontier, vertexIsDiscovered, label, threadQueues[i]));
         }
 
-        return new GlobalRelabeling(flowGraph.nodeCount(), label, frontier, vertexIsDiscovered, source, target, concurrency, globalRelabelingTasks);
+        return new GlobalRelabeling(
+            flowGraph.nodeCount(),
+            label,
+            frontier,
+            vertexIsDiscovered,
+            source,
+            target,
+            concurrency,
+            globalRelabelingTasks,
+            terminationFlag
+        );
     }
 
-    private GlobalRelabeling(long nodeCount, HugeLongArray label, AtomicWorkingSet frontier, HugeAtomicBitSet vertexIsDiscovered, long source, long target, Concurrency concurrency, Collection<Runnable> globalRelabelingTasks) {
+    private GlobalRelabeling(long nodeCount, HugeLongArray label, AtomicWorkingSet frontier, HugeAtomicBitSet vertexIsDiscovered, long source, long target, Concurrency concurrency, Collection<Runnable> globalRelabelingTasks,
+        TerminationFlag terminationFlag
+    ) {
         this.nodeCount = nodeCount;
         this.label = label;
         this.frontier = frontier;
@@ -67,6 +82,7 @@ private GlobalRelabeling(long nodeCount, HugeLongArray label, AtomicWorkingSet f
         this.target = target;
         this.concurrency = concurrency;
         this.globalRelabelingTasks = globalRelabelingTasks;
+        this.terminationFlag = terminationFlag;
     }
 
     public void globalRelabeling() {
@@ -78,11 +94,22 @@ public void globalRelabeling() {
         vertexIsDiscovered.set(source);
         vertexIsDiscovered.set(target);
         while (!frontier.isEmpty()) {
-            RunWithConcurrency.builder().concurrency(concurrency).tasks(globalRelabelingTasks).build().run();
+            //relax nodes in the frontier
+            runTasks();
             frontier.reset();
-            RunWithConcurrency.builder().concurrency(concurrency).tasks(globalRelabelingTasks).build().run();
+            //update the frontier
+            runTasks();
         }
         label.set(source, nodeCount);
 
     }
+
+   private void runTasks(){
+       RunWithConcurrency.builder()
+           .concurrency(concurrency)
+           .terminationFlag(terminationFlag)
+           .tasks(globalRelabelingTasks)
+           .build()
+           .run();
+    }
 }
diff --git a/algo/src/main/java/org/neo4j/gds/maxflow/MaxFlow.java b/algo/src/main/java/org/neo4j/gds/maxflow/MaxFlow.java
@@ -61,7 +61,7 @@ public FlowResult compute() {
 
     private Preflow initPreflow() {
         var supplyAndDemand = SupplyAndDemandFactory.create(graph, parameters.sourceNodes(), parameters.targetNodes());
-        var flowGraph = FlowGraph.create(graph, supplyAndDemand.getLeft(), supplyAndDemand.getRight());
+        var flowGraph = FlowGraph.create(graph, supplyAndDemand.getLeft(), supplyAndDemand.getRight(), terminationFlag);
         var excess = HugeDoubleArray.newArray(flowGraph.nodeCount());
         excess.setAll(x -> 0D);
         flowGraph.forEachRelationship(
@@ -107,7 +107,8 @@ private void maximizeFlow(Preflow preflow, long sourceNode, long targetNode) { /
             sourceNode,
             targetNode,
             parameters.concurrency(),
-            threadQueues
+            threadQueues,
+            terminationFlag
         );
 
         var discharging = Discharging.createDischarging(
@@ -120,7 +121,8 @@ private void maximizeFlow(Preflow preflow, long sourceNode, long targetNode) { /
             parameters.beta(),
             workSinceLastGR,
             parameters.concurrency(),
-            threadQueues
+            threadQueues,
+            terminationFlag
             );
 
         while (!workingSet.isEmpty()) {
diff --git a/algo/src/test/java/org/neo4j/gds/maxflow/FlowGraphTest.java b/algo/src/test/java/org/neo4j/gds/maxflow/FlowGraphTest.java
@@ -26,6 +26,7 @@
 import org.neo4j.gds.extension.GdlGraph;
 import org.neo4j.gds.extension.Inject;
 import org.neo4j.gds.extension.TestGraph;
+import org.neo4j.gds.termination.TerminationFlag;
 
 import java.util.HashMap;
 import java.util.HashSet;
@@ -61,7 +62,7 @@ static FlowGraph createFlowGraph(Graph graph, long source, long target) {
             .reduce(0D, Double::sum);
         NodeWithValue[] supply = {new NodeWithValue(source, outgoingCapacityFromSource)};
         NodeWithValue[] demand = {new NodeWithValue(target, outgoingCapacityFromSource)}; //more is useless since this is max in network
-        return FlowGraph.create(graph, supply, demand);
+        return FlowGraph.create(graph, supply, demand, TerminationFlag.RUNNING_TRUE);
     }
 
     @Test
diff --git a/algo/src/test/java/org/neo4j/gds/maxflow/GlobalRelabelingTest.java b/algo/src/test/java/org/neo4j/gds/maxflow/GlobalRelabelingTest.java
@@ -27,6 +27,7 @@
 import org.neo4j.gds.extension.GdlGraph;
 import org.neo4j.gds.extension.Inject;
 import org.neo4j.gds.extension.TestGraph;
+import org.neo4j.gds.termination.TerminationFlag;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.neo4j.gds.maxflow.FlowGraphTest.createFlowGraph;
@@ -76,7 +77,8 @@ void test() {
             graph.toMappedNodeId("c"),
             graph.toMappedNodeId("d"),
             new Concurrency(1),
-            threadQueues
+            threadQueues,
+            TerminationFlag.RUNNING_TRUE
         );
 
         globalRelabeling.globalRelabeling();
@@ -106,7 +108,8 @@ void test2() {
             graph.toMappedNodeId("a"),
             graph.toMappedNodeId("e"),
             new Concurrency(1),
-            threadQueues
+            threadQueues,
+            TerminationFlag.RUNNING_TRUE
         );
 
         globalRelabeling.globalRelabeling();
@@ -136,7 +139,8 @@ void test3() {
             flowGraph.superSource(),
             flowGraph.superTarget(),
             new Concurrency(1),
-            threadQueues
+            threadQueues,
+            TerminationFlag.RUNNING_TRUE
         );
 
         globalRelabeling.globalRelabeling();
