"fix" broken pagerank/ pregel encapsulation wrt terminationflag, and …

…thereby fix article rank termination behaviour
neo4j · May 31, 2024 · 621d2b1 · 621d2b1
1 parent 12b1a6a
commit 621d2b1
Show file tree

Hide file tree

Showing 3 changed files with 48 additions and 5 deletions.
diff --git a/algo/src/main/java/org/neo4j/gds/pagerank/PageRankAlgorithm.java b/algo/src/main/java/org/neo4j/gds/pagerank/PageRankAlgorithm.java
@@ -44,20 +44,45 @@ public class PageRankAlgorithm extends Algorithm<PageRankResult> {
     private final PageRankConfig config;
     private final ExecutorService executorService;
 
+    /**
+     * @deprecated Use the variant  that does proper injection of termination flag instead
+     */
+    @Deprecated
     public PageRankAlgorithm(
         Graph graph,
         PageRankConfig config,
         PregelComputation<PageRankConfig> pregelComputation,
         PageRankAlgorithmFactory.Mode mode,
         ExecutorService executorService,
         ProgressTracker progressTracker
+    ) {
+        this(
+            graph,
+            config,
+            pregelComputation,
+            mode,
+            executorService,
+            progressTracker,
+            TerminationFlag.RUNNING_TRUE
+        );
+    }
+
+    public PageRankAlgorithm(
+        Graph graph,
+        PageRankConfig config,
+        PregelComputation<PageRankConfig> pregelComputation,
+        PageRankAlgorithmFactory.Mode mode,
+        ExecutorService executorService,
+        ProgressTracker progressTracker,
+        TerminationFlag terminationFlag
     ) {
         super(progressTracker);
-        this.pregelJob = Pregel.create(graph, config, pregelComputation, executorService, progressTracker);
+        this.pregelJob = Pregel.create(graph, config, pregelComputation, executorService, progressTracker, terminationFlag);
         this.mode = mode;
         this.executorService = executorService;
         this.config = config;
         this.graph = graph;
+        this.terminationFlag = terminationFlag;
     }
 
     @Override

diff --git a/.../src/main/java/org/neo4j/gds/applications/algorithms/centrality/CentralityAlgorithms.java b/.../src/main/java/org/neo4j/gds/applications/algorithms/centrality/CentralityAlgorithms.java
@@ -92,7 +92,8 @@ PageRankResult articleRank(Graph graph, PageRankConfig configuration) {
             computation,
             mode,
             DefaultPool.INSTANCE,
-            progressTracker
+            progressTracker,
+            terminationFlag
         );
 
         return algorithmMachinery.runAlgorithmsAndManageProgressTracker(algorithm, progressTracker, true);

diff --git a/pregel/src/main/java/org/neo4j/gds/beta/pregel/Pregel.java b/pregel/src/main/java/org/neo4j/gds/beta/pregel/Pregel.java
@@ -58,12 +58,27 @@ public final class Pregel<CONFIG extends PregelConfig> {
 
     private final ExecutorService executor;
 
+    /**
+     * @deprecated Use the variant that does proper injection of termination flag instead
+     */
+    @Deprecated
     public static <CONFIG extends PregelConfig> Pregel<CONFIG> create(
         Graph graph,
         CONFIG config,
         BasePregelComputation<CONFIG> computation,
         ExecutorService executor,
         ProgressTracker progressTracker
+    ) {
+        return create(graph, config, computation, executor, progressTracker, TerminationFlag.RUNNING_TRUE);
+    }
+
+    public static <CONFIG extends PregelConfig> Pregel<CONFIG> create(
+        Graph graph,
+        CONFIG config,
+        BasePregelComputation<CONFIG> computation,
+        ExecutorService executor,
+        ProgressTracker progressTracker,
+        TerminationFlag terminationFlag
     ) {
         // This prevents users from disabling concurrency
         // validation in custom PregelConfig implementations.
@@ -86,7 +101,8 @@ public static <CONFIG extends PregelConfig> Pregel<CONFIG> create(
             computation,
             NodeValue.of(computation.schema(config), graph.nodeCount(), config.concurrency()),
             executor,
-            progressTracker
+            progressTracker,
+            terminationFlag
         );
     }
 
@@ -139,15 +155,16 @@ private Pregel(
         final BasePregelComputation<CONFIG> computation,
         final NodeValue initialNodeValue,
         final ExecutorService executor,
-        final ProgressTracker progressTracker
+        final ProgressTracker progressTracker,
+        TerminationFlag terminationFlag
     ) {
         this.graph = graph;
         this.config = config;
         this.computation = computation;
         this.nodeValues = initialNodeValue;
         this.executor = executor;
         this.progressTracker = progressTracker;
-        this.terminationFlag = TerminationFlag.RUNNING_TRUE;
+        this.terminationFlag = terminationFlag;
 
         var reducer = computation.reducer();