RandomWalk: ensure we read updated termination flag in algorithm

Author: hindog

algo/src/main/java/org/neo4j/gds/traversal/RandomWalk.java

@@ -96,7 +96,7 @@ public Stream<long[]> compute() {
             ? new NextNodeSupplier.GraphNodeSupplier(graph.nodeCount())
             : NextNodeSupplier.ListNodeSupplier.of(config, graph);
 
-        var terminationFlag = new ExternalTerminationFlag(this.terminationFlag);
+        var terminationFlag = new ExternalTerminationFlag(this);
 
         BlockingQueue<long[]> walks = new ArrayBlockingQueue<>(config.walkBufferSize());
         long[] TOMB = new long[0];
@@ -198,15 +198,15 @@ private Stream<long[]> walksQueueConsumer(
 
     private static final class ExternalTerminationFlag implements TerminationFlag {
         private volatile boolean running = true;
-        private final TerminationFlag inner;
+        private final Algorithm<?> algo;
 
-        ExternalTerminationFlag(TerminationFlag inner) {
-            this.inner = inner;
+        ExternalTerminationFlag(Algorithm<?> algo) {
+            this.algo = algo;
         }
 
         @Override
         public boolean running() {
-            return this.running && this.inner.running();
+            return this.running && this.algo.getTerminationFlag().running();
         }
 
         void stop() {

algo/src/test/java/org/neo4j/gds/traversal/RandomWalkTest.java

@@ -36,6 +36,7 @@
 import org.neo4j.gds.compat.Neo4jProxy;
 import org.neo4j.gds.compat.TestLog;
 import org.neo4j.gds.core.concurrency.Pools;
+import org.neo4j.gds.core.utils.TerminationFlag;
 import org.neo4j.gds.core.utils.progress.GlobalTaskStore;
 import org.neo4j.gds.core.utils.progress.TaskRegistryFactory;
 import org.neo4j.gds.core.utils.progress.tasks.ProgressTracker;
@@ -411,6 +412,38 @@ void testWithConfiguredOffsetStartNodes() {
             .anyMatch(walk -> walk[0] == bInternalId);
     }
 
+    /**
+     * Ensure that when termination flag is set externally, we terminate the walk
+     * @throws InterruptedException
+     */
+    @Test
+    void testPartialReadMultipleRuns() {
+        for (int i = 0; i < 3; i++) {
+            Node2VecStreamConfig config = ImmutableNode2VecStreamConfig.builder()
+                    .walkBufferSize(1)
+                    .build();
+
+            var randomWalk = RandomWalk.create(
+                    graph,
+                    config,
+                    ProgressTracker.NULL_TRACKER,
+                    Pools.DEFAULT
+            );
+
+            var stream = randomWalk.compute();
+            long count = stream.limit(10).count();
+
+            randomWalk.setTerminationFlag(new TerminationFlag() {
+                @Override
+                public boolean running() {
+                    return false;
+                }
+            });
+
+            assertEquals(10, count);
+        }
+    }
+
     @Nested
     class ProgressTracking {