From d799f378a3fae5c83104350954dfb12bbf3c0d37 Mon Sep 17 00:00:00 2001
From: Lasse Westh-Nielsen <lassewesth@gmail.com>
Date: Wed, 22 May 2024 15:46:43 +0200
Subject: [PATCH] fix injection of termination flag into centrality algorithms

---
 .../gds/betweenness/BetweennessCentrality.java     |  6 ++++--
 .../betweenness/BetweennessCentralityFactory.java  |  4 +++-
 .../gds/betweenness/BetweennessCentralityTest.java |  7 +++++--
 .../WeightedBetweennessCentralityTest.java         | 10 +++++++---
 applications/algorithms/centrality/build.gradle    |  1 +
 .../centrality/CentralityAlgorithms.java           | 13 ++++++++++---
 .../gds/applications/CentralityApplications.java   |  2 +-
 ...eClassificationPredictPipelineExecutorTest.java | 14 +++++++++++++-
 8 files changed, 44 insertions(+), 13 deletions(-)

diff --git a/algo/src/main/java/org/neo4j/gds/betweenness/BetweennessCentrality.java b/algo/src/main/java/org/neo4j/gds/betweenness/BetweennessCentrality.java
index d7cd319550..d032c05bf3 100644
--- a/algo/src/main/java/org/neo4j/gds/betweenness/BetweennessCentrality.java
+++ b/algo/src/main/java/org/neo4j/gds/betweenness/BetweennessCentrality.java
@@ -33,6 +33,7 @@
 import org.neo4j.gds.core.utils.paged.HugeLongArrayStack;
 import org.neo4j.gds.core.utils.paged.ParallelDoublePageCreator;
 import org.neo4j.gds.core.utils.progress.tasks.ProgressTracker;
+import org.neo4j.gds.termination.TerminationFlag;
 
 import java.util.concurrent.ExecutorService;
 import java.util.function.Consumer;
@@ -57,7 +58,8 @@ public BetweennessCentrality(
         ForwardTraverser.Factory traverserFactory,
         ExecutorService executorService,
         Concurrency concurrency,
-        ProgressTracker progressTracker
+        ProgressTracker progressTracker,
+        TerminationFlag terminationFlag
     ) {
         super(progressTracker);
         this.graph = graph;
@@ -69,7 +71,7 @@ public BetweennessCentrality(
         this.selectionStrategy.init(graph, executorService, concurrency);
         this.divisor = graph.schema().isUndirected() ? 2.0 : 1.0;
         this.traverserFactory = traverserFactory;
-
+        this.terminationFlag = terminationFlag;
     }
 
     @Override
diff --git a/algo/src/main/java/org/neo4j/gds/betweenness/BetweennessCentralityFactory.java b/algo/src/main/java/org/neo4j/gds/betweenness/BetweennessCentralityFactory.java
index 688d9da354..b115e84463 100644
--- a/algo/src/main/java/org/neo4j/gds/betweenness/BetweennessCentralityFactory.java
+++ b/algo/src/main/java/org/neo4j/gds/betweenness/BetweennessCentralityFactory.java
@@ -26,6 +26,7 @@
 import org.neo4j.gds.core.utils.progress.tasks.ProgressTracker;
 import org.neo4j.gds.core.utils.progress.tasks.Task;
 import org.neo4j.gds.core.utils.progress.tasks.Tasks;
+import org.neo4j.gds.termination.TerminationFlag;
 
 import java.util.Optional;
 
@@ -58,7 +59,8 @@ public BetweennessCentrality build(
             traverserFactory,
             DefaultPool.INSTANCE,
             parameters.concurrency(),
-            progressTracker
+            progressTracker,
+            TerminationFlag.RUNNING_TRUE
         );
     }
 
diff --git a/algo/src/test/java/org/neo4j/gds/betweenness/BetweennessCentralityTest.java b/algo/src/test/java/org/neo4j/gds/betweenness/BetweennessCentralityTest.java
index 59c653ec4b..2caf89e665 100644
--- a/algo/src/test/java/org/neo4j/gds/betweenness/BetweennessCentralityTest.java
+++ b/algo/src/test/java/org/neo4j/gds/betweenness/BetweennessCentralityTest.java
@@ -33,6 +33,7 @@
 import org.neo4j.gds.core.utils.progress.EmptyTaskRegistryFactory;
 import org.neo4j.gds.core.utils.progress.tasks.ProgressTracker;
 import org.neo4j.gds.extension.TestGraph;
+import org.neo4j.gds.termination.TerminationFlag;
 
 import java.util.Map;
 import java.util.Optional;
@@ -139,7 +140,8 @@ void sampling(int concurrency, TestGraph graph, int samplingSize, Map<String, Do
             ForwardTraverser.Factory.unweighted(),
             DefaultPool.INSTANCE,
             new Concurrency(concurrency),
-            ProgressTracker.NULL_TRACKER
+            ProgressTracker.NULL_TRACKER,
+            TerminationFlag.RUNNING_TRUE
         ).compute().centralities();
 
         assertEquals(expectedResult.size(), actualResult.size());
@@ -158,7 +160,8 @@ void noSampling(int concurrency) {
             ForwardTraverser.Factory.unweighted(),
             DefaultPool.INSTANCE,
             new Concurrency(concurrency),
-            ProgressTracker.NULL_TRACKER
+            ProgressTracker.NULL_TRACKER,
+            TerminationFlag.RUNNING_TRUE
         ).compute().centralities();
 
         assertEquals(5, actualResult.size(), "Expected 5 centrality values");
diff --git a/algo/src/test/java/org/neo4j/gds/betweenness/WeightedBetweennessCentralityTest.java b/algo/src/test/java/org/neo4j/gds/betweenness/WeightedBetweennessCentralityTest.java
index 84a8e13656..efd6593d08 100644
--- a/algo/src/test/java/org/neo4j/gds/betweenness/WeightedBetweennessCentralityTest.java
+++ b/algo/src/test/java/org/neo4j/gds/betweenness/WeightedBetweennessCentralityTest.java
@@ -30,6 +30,7 @@
 import org.neo4j.gds.extension.IdFunction;
 import org.neo4j.gds.extension.Inject;
 import org.neo4j.gds.extension.TestGraph;
+import org.neo4j.gds.termination.TerminationFlag;
 
 import java.util.Optional;
 
@@ -80,7 +81,8 @@ void shouldEqualWithUnweightedWhenWeightsAreEqual() {
             ForwardTraverser.Factory.weighted(),
             DefaultPool.INSTANCE,
             new Concurrency(8),
-            ProgressTracker.NULL_TRACKER
+            ProgressTracker.NULL_TRACKER,
+            TerminationFlag.RUNNING_TRUE
         );
         var algoUnweighted = new BetweennessCentrality(
             equallyWeightedGraph,
@@ -88,7 +90,8 @@ void shouldEqualWithUnweightedWhenWeightsAreEqual() {
             ForwardTraverser.Factory.unweighted(),
             DefaultPool.INSTANCE,
             new Concurrency(8),
-            ProgressTracker.NULL_TRACKER
+            ProgressTracker.NULL_TRACKER,
+            TerminationFlag.RUNNING_TRUE
         );
         var resultWeighted = algoWeighted.compute().centralities();
         var resultUnweighted = algoUnweighted.compute().centralities();
@@ -114,7 +117,8 @@ void shouldComputeWithWeights() {
              ForwardTraverser.Factory.weighted(),
              DefaultPool.INSTANCE,
              new Concurrency(8),
-             ProgressTracker.NULL_TRACKER
+             ProgressTracker.NULL_TRACKER,
+             TerminationFlag.RUNNING_TRUE
          );
         var result = bc.compute().centralities();
         var softAssertions = new SoftAssertions();
diff --git a/applications/algorithms/centrality/build.gradle b/applications/algorithms/centrality/build.gradle
index 3787f2486d..6dd009f181 100644
--- a/applications/algorithms/centrality/build.gradle
+++ b/applications/algorithms/centrality/build.gradle
@@ -13,4 +13,5 @@ dependencies {
     implementation project(":logging")
     implementation project(":memory-usage")
     implementation project(":progress-tracking")
+    implementation project(":termination")
 }
diff --git a/applications/algorithms/centrality/src/main/java/org/neo4j/gds/applications/algorithms/centrality/CentralityAlgorithms.java b/applications/algorithms/centrality/src/main/java/org/neo4j/gds/applications/algorithms/centrality/CentralityAlgorithms.java
index 54914f17ad..845398ea0e 100644
--- a/applications/algorithms/centrality/src/main/java/org/neo4j/gds/applications/algorithms/centrality/CentralityAlgorithms.java
+++ b/applications/algorithms/centrality/src/main/java/org/neo4j/gds/applications/algorithms/centrality/CentralityAlgorithms.java
@@ -38,12 +38,15 @@
 import org.neo4j.gds.degree.DegreeCentrality;
 import org.neo4j.gds.degree.DegreeCentralityConfig;
 import org.neo4j.gds.degree.DegreeCentralityResult;
+import org.neo4j.gds.termination.TerminationFlag;
 
 public class CentralityAlgorithms {
     private final ProgressTrackerCreator progressTrackerCreator;
+    private final TerminationFlag terminationFlag;
 
-    public CentralityAlgorithms(ProgressTrackerCreator progressTrackerCreator) {
+    public CentralityAlgorithms(ProgressTrackerCreator progressTrackerCreator, TerminationFlag terminationFlag) {
         this.progressTrackerCreator = progressTrackerCreator;
+        this.terminationFlag = terminationFlag;
     }
 
     BetwennessCentralityResult betweennessCentrality(Graph graph, BetweennessCentralityBaseConfig configuration) {
@@ -60,7 +63,10 @@ BetwennessCentralityResult betweennessCentrality(Graph graph, BetweennessCentral
             ? ForwardTraverser.Factory.weighted()
             : ForwardTraverser.Factory.unweighted();
 
-        var task = Tasks.leaf(LabelForProgressTracking.BetweennessCentrality.value, samplingSize.orElse(graph.nodeCount()));
+        var task = Tasks.leaf(
+            LabelForProgressTracking.BetweennessCentrality.value,
+            samplingSize.orElse(graph.nodeCount())
+        );
         var progressTracker = progressTrackerCreator.createProgressTracker(configuration, task);
 
         var algorithm = new BetweennessCentrality(
@@ -69,7 +75,8 @@ BetwennessCentralityResult betweennessCentrality(Graph graph, BetweennessCentral
             traverserFactory,
             DefaultPool.INSTANCE,
             parameters.concurrency(),
-            progressTracker
+            progressTracker,
+            terminationFlag
         );
 
         return algorithm.compute();
diff --git a/applications/facade/src/main/java/org/neo4j/gds/applications/CentralityApplications.java b/applications/facade/src/main/java/org/neo4j/gds/applications/CentralityApplications.java
index bc55d0932d..6d8a07fc51 100644
--- a/applications/facade/src/main/java/org/neo4j/gds/applications/CentralityApplications.java
+++ b/applications/facade/src/main/java/org/neo4j/gds/applications/CentralityApplications.java
@@ -62,7 +62,7 @@ static CentralityApplications create(
         ProgressTrackerCreator progressTrackerCreator
     ) {
         var estimation = new CentralityAlgorithmsEstimationModeBusinessFacade(estimationTemplate);
-        var algorithms = new CentralityAlgorithms(progressTrackerCreator);
+        var algorithms = new CentralityAlgorithms(progressTrackerCreator, requestScopedDependencies.getTerminationFlag());
         var mutateNodePropertyService = new MutateNodePropertyService(log);
         var mutateNodeProperty = new MutateNodeProperty(mutateNodePropertyService);
         var mutation = new CentralityAlgorithmsMutateModeBusinessFacade(
diff --git a/proc/machine-learning/src/test/java/org/neo4j/gds/ml/pipeline/node/classification/predict/NodeClassificationPredictPipelineExecutorTest.java b/proc/machine-learning/src/test/java/org/neo4j/gds/ml/pipeline/node/classification/predict/NodeClassificationPredictPipelineExecutorTest.java
index a91d21334a..7a79f5e435 100644
--- a/proc/machine-learning/src/test/java/org/neo4j/gds/ml/pipeline/node/classification/predict/NodeClassificationPredictPipelineExecutorTest.java
+++ b/proc/machine-learning/src/test/java/org/neo4j/gds/ml/pipeline/node/classification/predict/NodeClassificationPredictPipelineExecutorTest.java
@@ -34,6 +34,7 @@
 import org.neo4j.gds.api.GraphStore;
 import org.neo4j.gds.api.schema.GraphSchema;
 import org.neo4j.gds.applications.ApplicationsFacade;
+import org.neo4j.gds.applications.algorithms.machinery.RequestScopedDependencies;
 import org.neo4j.gds.catalog.GraphProjectProc;
 import org.neo4j.gds.catalog.GraphStreamNodePropertiesProc;
 import org.neo4j.gds.core.CypherMapWrapper;
@@ -62,6 +63,7 @@
 import org.neo4j.gds.procedures.algorithms.centrality.CentralityProcedureFacade;
 import org.neo4j.gds.procedures.algorithms.configuration.ConfigurationParser;
 import org.neo4j.gds.procedures.algorithms.stubs.GenericStub;
+import org.neo4j.gds.termination.TerminationFlag;
 import org.neo4j.gds.test.TestProc;
 
 import java.util.ArrayList;
@@ -423,7 +425,17 @@ void shouldEstimateMemory() {
      * But let's be honest, there is enough work in front of us that such a fix is lower priority right now.
      */
     private static AlgorithmsProcedureFacade createAlgorithmsProcedureFacade() {
-        var applicationsFacade = ApplicationsFacade.create(null, Optional.empty(), null, null, null, null, null, null);
+        var applicationsFacade = ApplicationsFacade.create(
+            null,
+            Optional.empty(),
+            null,
+            null,
+            null,
+            null,
+            RequestScopedDependencies.builder().with(
+                TerminationFlag.RUNNING_TRUE).build(),
+            null
+        );
         var configurationParser = new ConfigurationParser(null, null);
         var genericStub = new GenericStub(null, null, null, configurationParser, null, null);
         var centralityProcedureFacade = CentralityProcedureFacade.create(