Simplify working-set update

Co-authored-by: Ioannis Panagiotas <ioannis.panagiotas@neo4j.com>

Author: AlfredClemedtson

algo/src/main/java/org/neo4j/gds/maxflow/AtomicWorkingSet.java

@@ -66,6 +66,15 @@ void batchPush(HugeLongArrayQueue queue) {
         }
     }
 
+    void batchPushAndConsume(HugeLongArrayQueue queue, Consumer<Long> consumer) {
+        long idx = size.getAndAdd(queue.size());
+        while (!queue.isEmpty()) {
+            var node = queue.remove();
+            workingSet.set(idx++, node);
+            consumer.accept(node);
+        }
+    }
+
     long getAndAdd(long batchSize) {
         return index.getAndAdd(batchSize);
     }

algo/src/main/java/org/neo4j/gds/maxflow/DischargeTask.java

@@ -84,8 +84,7 @@ public void run() {
         switch (phase) {
             case PHASE.DISCHARGE -> dischargeWorkingSet();
             case PHASE.SYNC_WORKING_SET -> syncWorkingSet();
-            case PHASE.UPDATE_WORKING_SET -> updateWorkingSet();
-            case PHASE.SYNC_NEW_WORKING_SET -> syncNewWorkingSet();
+            case PHASE.UPDATE_WORKING_SET -> updateAndSyncNewWorkingSet();
         }
     }
 
@@ -178,24 +177,18 @@ void syncWorkingSet() {
         phase = PHASE.UPDATE_WORKING_SET;
     }
 
-    void updateWorkingSet() {
-        workingSet.batchPush(localDiscoveredVertices);
-        phase = PHASE.SYNC_NEW_WORKING_SET;
-    }
-
-    void syncNewWorkingSet() {
-        batchConsumeWorkingSet((v) -> {
+    void updateAndSyncNewWorkingSet() {
+        workingSet.batchPushAndConsume(localDiscoveredVertices, (v) -> {
             excess.addTo(v, addedExcess.get(v));
             addedExcess.set(v, 0);
             isDiscovered.clear(v);
         });
-        phase = PHASE.UPDATE_WORKING_SET;
+        phase = PHASE.DISCHARGE;
     }
 
     enum PHASE {
         DISCHARGE,
         SYNC_WORKING_SET,
         UPDATE_WORKING_SET,
-        SYNC_NEW_WORKING_SET
     }
 }

algo/src/main/java/org/neo4j/gds/maxflow/Discharging.java

@@ -67,11 +67,7 @@ public static void processWorkingSet(
         RunWithConcurrency.builder().concurrency(concurrency).tasks(dischargeTasks).build().run();
         workingSet.reset();
 
-        //Update working set
+        //Update and sync new working set
         RunWithConcurrency.builder().concurrency(concurrency).tasks(dischargeTasks).build().run();
-
-        //Sync new working set
-        RunWithConcurrency.builder().concurrency(concurrency).tasks(dischargeTasks).build().run();
-        workingSet.resetIdx();
     }
 }

algo/src/test/java/org/neo4j/gds/maxflow/DischargeTaskTest.java

@@ -110,9 +110,8 @@ void discharge() {
         assertThat(excess.get(graph.toMappedNodeId("a"))).isEqualTo(0D);
         assertThat(excess.get(graph.toMappedNodeId("c"))).isEqualTo(8D);
 
-        task.updateWorkingSet();
+        task.updateAndSyncNewWorkingSet();
 
-        task.syncNewWorkingSet();
         workingSet.resetIdx();
 
         assertThat(label.get(graph.toMappedNodeId("a"))).isEqualTo(2L);