Remove contention on hasSendMessage atomic boolean

Author: DarthMax

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

@@ -43,8 +43,9 @@ public final class ForkJoinComputeStep<
 
     private final InitFunction<CONFIG, INIT_CONTEXT> initFunction;
     private final ComputeFunction<CONFIG, COMPUTE_CONTEXT> computeFunction;
-    private final Supplier<INIT_CONTEXT> initContext;
-    private final Supplier<COMPUTE_CONTEXT> computeContext;
+    private final Supplier<INIT_CONTEXT> initContextSupplier;
+    private final Supplier<COMPUTE_CONTEXT> computeContextSupplier;
+    private final COMPUTE_CONTEXT computeContext;
     private final NodeValue nodeValue;
     private final HugeAtomicBitSet voteBits;
     private final Messenger<ITERATOR> messenger;
@@ -56,8 +57,8 @@ public final class ForkJoinComputeStep<
     ForkJoinComputeStep(
         InitFunction<CONFIG, INIT_CONTEXT> initFunction,
         ComputeFunction<CONFIG, COMPUTE_CONTEXT> computeFunction,
-        Supplier<INIT_CONTEXT> initContext,
-        Supplier<COMPUTE_CONTEXT> computeContext,
+        Supplier<INIT_CONTEXT> initContextSupplier,
+        Supplier<COMPUTE_CONTEXT> computeContextSupplier,
         MutableInt iteration,
         Partition nodeBatch,
         NodeValue nodeValue,
@@ -70,15 +71,16 @@ public final class ForkJoinComputeStep<
         super(parent);
         this.initFunction = initFunction;
         this.computeFunction = computeFunction;
-        this.initContext = initContext;
-        this.computeContext = computeContext;
+        this.initContextSupplier = initContextSupplier;
+        this.computeContextSupplier = computeContextSupplier;
         this.iteration = iteration;
         this.voteBits = voteBits;
         this.nodeBatch = nodeBatch;
         this.nodeValue = nodeValue;
         this.messenger = messenger;
         this.hasSentMessage = sentMessage;
         this.progressTracker = progressTracker;
+        this.computeContext = computeContextSupplier.get();
     }
 
     @Override
@@ -99,8 +101,8 @@ public void compute() {
             var leftTask = new ForkJoinComputeStep<>(
                 initFunction,
                 computeFunction,
-                initContext,
-                computeContext,
+                initContextSupplier,
+                computeContextSupplier,
                 iteration,
                 leftBatch,
                 nodeValue,
@@ -119,6 +121,7 @@ public void compute() {
             this.compute();
         } else {
             computeBatch();
+            hasSentMessage.set(computeContext.hasSentMessage());
             tryComplete();
         }
     }
@@ -155,12 +158,12 @@ public Partition nodeBatch() {
 
     @Override
     public INIT_CONTEXT initContext() {
-        return initContext.get();
+        return initContextSupplier.get();
     }
 
     @Override
     public COMPUTE_CONTEXT computeContext() {
-        return computeContext.get();
+        return computeContext;
     }
 
     @Override

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

@@ -108,7 +108,6 @@ void release() {
             messenger,
             voteBits,
             iteration,
-            hasSentMessages,
             progressTracker
         );
 
@@ -149,7 +148,6 @@ void release() {
             messenger,
             voteBits,
             iteration,
-            hasSentMessages,
             progressTracker
         );
 

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

@@ -78,6 +78,7 @@ public final class PartitionedComputeStep<
     @Override
     public void run() {
         computeBatch();
+        hasSentMessage.set(computeContext().hasSentMessage());
     }
 
     @Override

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

@@ -154,7 +154,6 @@ void release() {
             messenger,
             voteBits,
             iteration,
-            hasSentMessages,
             progressTracker
         );
 
@@ -197,7 +196,6 @@ void release() {
             messenger,
             voteBits,
             iteration,
-            hasSentMessages,
             progressTracker
         );
 

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

@@ -28,8 +28,6 @@
 import org.neo4j.gds.core.utils.paged.HugeAtomicBitSet;
 import org.neo4j.gds.core.utils.progress.tasks.ProgressTracker;
 
-import java.util.concurrent.atomic.AtomicBoolean;
-
 /**
  * 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
@@ -41,7 +39,8 @@ public class ComputeContext<CONFIG extends PregelConfig> extends NodeCentricCont
 
     private final Messenger<?> messenger;
     private final MutableInt iteration;
-    private final AtomicBoolean hasSendMessage;
+    private boolean hasSendMessage;
+
     protected BasePregelComputation<CONFIG> computation;
 
     public ComputeContext(Graph graph,
@@ -51,7 +50,6 @@ public ComputeContext(Graph graph,
                           Messenger<?> messenger,
                           HugeAtomicBitSet voteBits,
                           MutableInt iteration,
-                          AtomicBoolean hasSendMessage,
                           ProgressTracker progressTracker) {
         super(graph, config, nodeValue, progressTracker);
         this.computation = computation;
@@ -61,7 +59,7 @@ public ComputeContext(Graph graph,
         this.messenger = messenger;
         this.voteBits = voteBits;
         this.iteration = iteration;
-        this.hasSendMessage = hasSendMessage;
+        this.hasSendMessage = false;
     }
 
     private final SendMessagesFunction sendMessagesFunction;
@@ -152,7 +150,7 @@ public void sendToNeighbors(double message) {
      */
     public void sendTo(long targetNodeId, double message) {
         messenger.sendTo(targetNodeId, message);
-        this.hasSendMessage.set(true);
+        this.hasSendMessage = true;
     }
 
     private void sendToNeighbors(long sourceNodeId, double message) {
@@ -169,6 +167,10 @@ private void sendToNeighborsWeighted(long sourceNodeId, double message) {
         });
     }
 
+    public boolean hasSentMessage() {
+        return hasSendMessage;
+    }
+
     @FunctionalInterface
     interface SendMessagesFunction {
         void sendToNeighbors(long sourceNodeId, double message);
@@ -186,7 +188,6 @@ public BidirectionalComputeContext(
             Messenger<?> messenger,
             HugeAtomicBitSet voteBits,
             MutableInt iteration,
-            AtomicBoolean hasSendMessage,
             ProgressTracker progressTracker
         ) {
             super(
@@ -197,7 +198,6 @@ public BidirectionalComputeContext(
                 messenger,
                 voteBits,
                 iteration,
-                hasSendMessage,
                 progressTracker
             );