From 621d2b180e8989735189c8327002b31bd431307a Mon Sep 17 00:00:00 2001
From: Lasse Westh-Nielsen <lassewesth@gmail.com>
Date: Fri, 31 May 2024 15:57:47 +0200
Subject: [PATCH] "fix" broken pagerank/ pregel encapsulation wrt
 terminationflag, and thereby fix article rank termination behaviour

---
 .../neo4j/gds/pagerank/PageRankAlgorithm.java | 27 ++++++++++++++++++-
 .../centrality/CentralityAlgorithms.java      |  3 ++-
 .../org/neo4j/gds/beta/pregel/Pregel.java     | 23 +++++++++++++---
 3 files changed, 48 insertions(+), 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
index 71ea929cb7..a9b40b2438 100644
--- a/algo/src/main/java/org/neo4j/gds/pagerank/PageRankAlgorithm.java
+++ b/algo/src/main/java/org/neo4j/gds/pagerank/PageRankAlgorithm.java
@@ -44,6 +44,10 @@ 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,
@@ -51,13 +55,34 @@ public PageRankAlgorithm(
         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/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 125ad37046..35bc1f83b6 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
@@ -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
index 9fbf61715e..320704fb7e 100644
--- 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,7 +155,8 @@ 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;
@@ -147,7 +164,7 @@ private Pregel(
         this.nodeValues = initialNodeValue;
         this.executor = executor;
         this.progressTracker = progressTracker;
-        this.terminationFlag = TerminationFlag.RUNNING_TRUE;
+        this.terminationFlag = terminationFlag;
 
         var reducer = computation.reducer();