Fix check for sendMessages for PartitionedComputer

Author: DarthMax

pregel/src/main/java/org/neo4j/gds/beta/pregel/ForkJoinComputeStep.java

@@ -121,7 +121,7 @@ public void compute() {
             this.compute();
         } else {
             computeBatch();
-            hasSentMessage.set(computeContext.hasSentMessage());
+            hasSentMessage.compareAndSet(false, computeContext.hasSentMessage());
             tryComplete();
         }
     }

pregel/src/main/java/org/neo4j/gds/beta/pregel/ForkJoinComputer.java

@@ -30,6 +30,7 @@
 import org.neo4j.gds.core.utils.partition.Partition;
 import org.neo4j.gds.core.utils.progress.tasks.ProgressTracker;
 
+import java.util.Optional;
 import java.util.concurrent.ForkJoinPool;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Supplier;
@@ -108,6 +109,7 @@ void release() {
             messenger,
             voteBits,
             iteration,
+            Optional.empty(),
             progressTracker
         );
 
@@ -148,6 +150,7 @@ void release() {
             messenger,
             voteBits,
             iteration,
+            Optional.empty(),
             progressTracker
         );
 

pregel/src/main/java/org/neo4j/gds/beta/pregel/PartitionedComputeStep.java

@@ -19,15 +19,14 @@
  */
 package org.neo4j.gds.beta.pregel;
 
+import org.apache.commons.lang3.mutable.MutableBoolean;
 import org.apache.commons.lang3.mutable.MutableInt;
 import org.neo4j.gds.beta.pregel.context.ComputeContext;
 import org.neo4j.gds.beta.pregel.context.InitContext;
 import org.neo4j.gds.core.utils.paged.HugeAtomicBitSet;
 import org.neo4j.gds.core.utils.partition.Partition;
 import org.neo4j.gds.core.utils.progress.tasks.ProgressTracker;
 
-import java.util.concurrent.atomic.AtomicBoolean;
-
 public final class PartitionedComputeStep<
     CONFIG extends PregelConfig,
     ITERATOR extends Messages.MessageIterator,
@@ -45,7 +44,7 @@ public final class PartitionedComputeStep<
     private final Messenger<ITERATOR> messenger;
 
     private final MutableInt iteration;
-    private final AtomicBoolean hasSentMessage;
+    private final MutableBoolean hasSentMessage;
     private final NodeValue nodeValue;
 
     PartitionedComputeStep(
@@ -58,7 +57,7 @@ public final class PartitionedComputeStep<
         Messenger<ITERATOR> messenger,
         HugeAtomicBitSet voteBits,
         MutableInt iteration,
-        AtomicBoolean hasSentMessage,
+        MutableBoolean hasSentMessage,
         ProgressTracker progressTracker
     ) {
         this.initFunction = initFunction;
@@ -78,7 +77,6 @@ public final class PartitionedComputeStep<
     @Override
     public void run() {
         computeBatch();
-        hasSentMessage.set(computeContext().hasSentMessage());
     }
 
     @Override
@@ -128,10 +126,10 @@ public ProgressTracker progressTracker() {
 
     void init(int iteration) {
         this.iteration.setValue(iteration);
-        this.hasSentMessage.set(false);
+        hasSentMessage.setValue(false);
     }
 
     boolean hasSentMessage() {
-        return hasSentMessage.get();
+        return hasSentMessage.getValue();
     }
 }

pregel/src/main/java/org/neo4j/gds/beta/pregel/PartitionedComputer.java

@@ -19,6 +19,7 @@
  */
 package org.neo4j.gds.beta.pregel;
 
+import org.apache.commons.lang3.mutable.MutableBoolean;
 import org.apache.commons.lang3.mutable.MutableInt;
 import org.jetbrains.annotations.NotNull;
 import org.neo4j.gds.api.Graph;
@@ -35,7 +36,6 @@
 import java.util.List;
 import java.util.Optional;
 import java.util.concurrent.ExecutorService;
-import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Function;
 
 import static org.neo4j.gds.utils.StringFormatting.formatWithLocale;
@@ -137,7 +137,7 @@ void release() {
         Partition partition
     ) {
         MutableInt iteration = new MutableInt(0);
-        var hasSentMessages = new AtomicBoolean(false);
+        var hasSentMessages = new MutableBoolean(false);
 
         var initContext = new InitContext<>(
             graph,
@@ -154,6 +154,7 @@ void release() {
             messenger,
             voteBits,
             iteration,
+            Optional.of(hasSentMessages),
             progressTracker
         );
 
@@ -179,7 +180,7 @@ void release() {
         Partition partition
     ) {
         MutableInt iteration = new MutableInt(0);
-        var hasSentMessages = new AtomicBoolean(false);
+        var hasSentMessages = new MutableBoolean(false);
 
         var initContext = new BidirectionalInitContext<>(
             graph,
@@ -196,6 +197,7 @@ void release() {
             messenger,
             voteBits,
             iteration,
+            Optional.of(hasSentMessages),
             progressTracker
         );
 

pregel/src/main/java/org/neo4j/gds/beta/pregel/context/ComputeContext.java

@@ -19,6 +19,7 @@
  */
 package org.neo4j.gds.beta.pregel.context;
 
+import org.apache.commons.lang3.mutable.MutableBoolean;
 import org.apache.commons.lang3.mutable.MutableInt;
 import org.neo4j.gds.api.Graph;
 import org.neo4j.gds.beta.pregel.BasePregelComputation;
@@ -28,6 +29,8 @@
 import org.neo4j.gds.core.utils.paged.HugeAtomicBitSet;
 import org.neo4j.gds.core.utils.progress.tasks.ProgressTracker;
 
+import java.util.Optional;
+
 /**
  * A context that is used during the computation. It allows an implementation
  * to send messages to other nodes and change the state of the currently
@@ -39,7 +42,7 @@ public class ComputeContext<CONFIG extends PregelConfig> extends NodeCentricCont
 
     private final Messenger<?> messenger;
     private final MutableInt iteration;
-    private boolean hasSendMessage;
+    private final MutableBoolean hasSendMessage;
 
     protected BasePregelComputation<CONFIG> computation;
 
@@ -50,6 +53,7 @@ public ComputeContext(Graph graph,
                           Messenger<?> messenger,
                           HugeAtomicBitSet voteBits,
                           MutableInt iteration,
+                          Optional<MutableBoolean> hasSendMessage,
                           ProgressTracker progressTracker) {
         super(graph, config, nodeValue, progressTracker);
         this.computation = computation;
@@ -59,7 +63,7 @@ public ComputeContext(Graph graph,
         this.messenger = messenger;
         this.voteBits = voteBits;
         this.iteration = iteration;
-        this.hasSendMessage = false;
+        this.hasSendMessage = hasSendMessage.orElse(new MutableBoolean(false));
     }
 
     private final SendMessagesFunction sendMessagesFunction;
@@ -150,7 +154,7 @@ public void sendToNeighbors(double message) {
      */
     public void sendTo(long targetNodeId, double message) {
         messenger.sendTo(targetNodeId, message);
-        this.hasSendMessage = true;
+        this.hasSendMessage.setValue(true);
     }
 
     private void sendToNeighbors(long sourceNodeId, double message) {
@@ -168,7 +172,7 @@ private void sendToNeighborsWeighted(long sourceNodeId, double message) {
     }
 
     public boolean hasSentMessage() {
-        return hasSendMessage;
+        return hasSendMessage.getValue();
     }
 
     @FunctionalInterface
@@ -188,6 +192,7 @@ public BidirectionalComputeContext(
             Messenger<?> messenger,
             HugeAtomicBitSet voteBits,
             MutableInt iteration,
+            Optional<MutableBoolean> hasSendMessage,
             ProgressTracker progressTracker
         ) {
             super(
@@ -198,6 +203,7 @@ public BidirectionalComputeContext(
                 messenger,
                 voteBits,
                 iteration,
+                hasSendMessage,
                 progressTracker
             );