Fix getFJPoolWithConcurrency issue

Co-authored-by: Martin Junghanns <[email protected]>
Co-Authored-By: Paul Horn <[email protected]>
Co-Authored-By: Jonatan Jäderberg <[email protected]>

Author: 4 people

core/src/main/java/org/neo4j/gds/core/concurrency/ParallelUtil.java

@@ -76,7 +76,7 @@ private ParallelUtil() {}
      * The concurrency value is assumed to already be validated towards the edition limitation.
      */
     public static <T extends BaseStream<?, T>, R> R parallelStream(T data, int concurrency, Function<T, R> fn) {
-        ForkJoinPool pool = getFJPoolWithConcurrency(concurrency);
+        ForkJoinPool pool = Pools.createForkJoinPool(concurrency);
         try {
             return pool.submit(() -> fn.apply(data.parallel())).get();
         } catch (Exception e) {
@@ -1487,20 +1487,4 @@ void pushBack(final T element) {
             pushedElement = element;
         }
     }
-
-    public static ForkJoinPool getFJPoolWithConcurrency(int concurrency) {
-        return new ForkJoinPool(concurrency, forkJoinPoolWorkerThreadFactory, null, false);
-    }
-
-    private static final ForkJoinPool.ForkJoinWorkerThreadFactory forkJoinPoolWorkerThreadFactory;
-
-    private static final String FORK_JOIN_INFIX = "-forkjoin-";
-
-    static {
-        forkJoinPoolWorkerThreadFactory = pool -> {
-            var worker = ForkJoinPool.defaultForkJoinWorkerThreadFactory.newThread(pool);
-            worker.setName(Pools.THREAD_NAME_PREFIX + FORK_JOIN_INFIX + worker.getPoolIndex());
-            return worker;
-        };
-    }
 }

core/src/main/java/org/neo4j/gds/core/concurrency/Pools.java

@@ -28,6 +28,7 @@
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
+import java.util.concurrent.ForkJoinPool;
 import java.util.concurrent.FutureTask;
 import java.util.concurrent.RejectedExecutionHandler;
 import java.util.concurrent.ThreadPoolExecutor;
@@ -94,4 +95,14 @@ public void rejectedExecution(Runnable r, ThreadPoolExecutor executor) {
             }
         }
     }
+
+    public static final ForkJoinPool.ForkJoinWorkerThreadFactory FJ_WORKER_THREAD_FACTORY = pool -> {
+        var worker = ForkJoinPool.defaultForkJoinWorkerThreadFactory.newThread(pool);
+        worker.setName(Pools.THREAD_NAME_PREFIX + "-forkjoin-" + worker.getPoolIndex());
+        return worker;
+    };
+
+    public static ForkJoinPool createForkJoinPool(int concurrency) {
+        return new ForkJoinPool(concurrency, FJ_WORKER_THREAD_FACTORY, null, false);
+    }
 }

core/src/main/java/org/neo4j/gds/core/utils/paged/HugeMergeSort.java

@@ -20,7 +20,7 @@
 package org.neo4j.gds.core.utils.paged;
 
 import org.jetbrains.annotations.Nullable;
-import org.neo4j.gds.core.concurrency.ParallelUtil;
+import org.neo4j.gds.core.concurrency.Pools;
 import org.neo4j.gds.core.utils.mem.AllocationTracker;
 
 import java.util.concurrent.CountedCompleter;
@@ -31,8 +31,12 @@ public final class HugeMergeSort {
 
     public static void sort(HugeLongArray array, int concurrency, AllocationTracker allocationTracker) {
         var temp = HugeLongArray.newArray(array.size(), allocationTracker);
-        var forkJoinPool = ParallelUtil.getFJPoolWithConcurrency(concurrency);
-        forkJoinPool.invoke(new MergeSortTask(null, array, temp, 0, array.size() - 1));
+        var forkJoinPool = Pools.createForkJoinPool(concurrency);
+        try {
+            forkJoinPool.invoke(new MergeSortTask(null, array, temp, 0, array.size() - 1));
+        } finally {
+            forkJoinPool.shutdown();
+        }
     }
 
     static class MergeSortTask extends CountedCompleter<Void> {

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

@@ -80,4 +80,9 @@ public void runIteration() {
     public boolean hasConverged() {
         return !sentMessage.get() && voteBits.allSet();
     }
+
+    @Override
+    void release() {
+        forkJoinPool.shutdown();
+    }
 }

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

@@ -20,10 +20,10 @@
 package org.neo4j.gds.beta.pregel;
 
 import org.jetbrains.annotations.NotNull;
-import org.neo4j.gds.core.utils.partition.Partition;
 import org.neo4j.gds.api.Graph;
 import org.neo4j.gds.core.concurrency.ParallelUtil;
 import org.neo4j.gds.core.utils.paged.HugeAtomicBitSet;
+import org.neo4j.gds.core.utils.partition.Partition;
 import org.neo4j.gds.core.utils.partition.PartitionUtils;
 import org.neo4j.gds.core.utils.progress.tasks.ProgressTracker;
 
@@ -83,6 +83,12 @@ public boolean hasConverged() {
 
     }
 
+    @Override
+    void release() {
+        // Unlike in the sibling ForkJoinComputer, we will not shut down the
+        // executor service (thread pool), since we use the shared global thread pool.
+    }
+
     @NotNull
     private List<PartitionedComputeStep<CONFIG, ?>> createComputeSteps(HugeAtomicBitSet voteBits) {
         Function<Partition, PartitionedComputeStep<CONFIG, ?>> partitionFunction = partition -> new PartitionedComputeStep<>(

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

@@ -22,7 +22,7 @@
 import org.immutables.value.Value;
 import org.neo4j.gds.api.Graph;
 import org.neo4j.gds.beta.pregel.context.MasterComputeContext;
-import org.neo4j.gds.core.concurrency.ParallelUtil;
+import org.neo4j.gds.core.concurrency.Pools;
 import org.neo4j.gds.core.utils.mem.AllocationTracker;
 import org.neo4j.gds.core.utils.mem.MemoryEstimation;
 import org.neo4j.gds.core.utils.mem.MemoryEstimations;
@@ -149,7 +149,7 @@ private Pregel(
             .messenger(messenger)
             .voteBits(HugeAtomicBitSet.create(graph.nodeCount(), allocationTracker))
             .executorService(config.useForkJoin()
-                ? ParallelUtil.getFJPoolWithConcurrency(config.concurrency())
+                ? Pools.createForkJoinPool(config.concurrency())
                 : executor)
             .progressTracker(progressTracker)
             .build();
@@ -194,6 +194,7 @@ public PregelResult run() {
                 .build();
         } finally {
             progressTracker.endSubTask();
+            computer.release();
         }
     }
 

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

@@ -64,6 +64,8 @@ abstract class PregelComputer<CONFIG extends PregelConfig> {
 
     abstract boolean hasConverged();
 
+    abstract void release();
+
     static <CONFIG extends PregelConfig> ComputerBuilder<CONFIG> builder() {
         return new ComputerBuilder<>();
     }