8190974: Parallel stream execution within a custom ForkJoinPool should obey the parallelism

Reviewed-by: martin, tvaleev

3 files changed, 25 insertions, 15 deletions

diff --git a/src/share/classes/java/util/stream/AbstractTask.java b/src/share/classes/java/util/stream/AbstractTask.java
index 33de7d5c52..8c5be49aff 100644
--- a/src/share/classes/java/util/stream/AbstractTask.java
+++ b/src/share/classes/java/util/stream/AbstractTask.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2012, 2013, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2012, 2017, Oracle and/or its affiliates. All rights reserved.
  * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
  *
  * This code is free software; you can redistribute it and/or modify it
@@ -27,6 +27,7 @@ package java.util.stream;
 import java.util.Spliterator;
 import java.util.concurrent.CountedCompleter;
 import java.util.concurrent.ForkJoinPool;
+import java.util.concurrent.ForkJoinWorkerThread;
 
 /**
  * Abstract base class for most fork-join tasks used to implement stream ops.
@@ -88,13 +89,7 @@ abstract class AbstractTask<P_IN, P_OUT, R,
                             K extends AbstractTask<P_IN, P_OUT, R, K>>
         extends CountedCompleter<R> {
 
-    /**
-     * Default target factor of leaf tasks for parallel decomposition.
-     * To allow load balancing, we over-partition, currently to approximately
-     * four tasks per processor, which enables others to help out
-     * if leaf tasks are uneven or some processors are otherwise busy.
-     */
-    static final int LEAF_TARGET = ForkJoinPool.getCommonPoolParallelism() << 2;
+    private static final int LEAF_TARGET = ForkJoinPool.getCommonPoolParallelism() << 2;
 
     /** The pipeline helper, common to all tasks in a computation */
     protected final PipelineHelper<P_OUT> helper;
@@ -157,6 +152,22 @@ abstract class AbstractTask<P_IN, P_OUT, R,
     }
 
     /**
+     * Default target of leaf tasks for parallel decomposition.
+     * To allow load balancing, we over-partition, currently to approximately
+     * four tasks per processor, which enables others to help out
+     * if leaf tasks are uneven or some processors are otherwise busy.
+     */
+    public static int getLeafTarget() {
+        Thread t = Thread.currentThread();
+        if (t instanceof ForkJoinWorkerThread) {
+            return ((ForkJoinWorkerThread) t).getPool().getParallelism() << 2;
+        }
+        else {
+            return LEAF_TARGET;
+        }
+    }
+
+    /**
      * Constructs a new node of type T whose parent is the receiver; must call
      * the AbstractTask(T, Spliterator) constructor with the receiver and the
      * provided Spliterator.
@@ -181,7 +192,7 @@ abstract class AbstractTask<P_IN, P_OUT, R,
      * @return suggested target leaf size
      */
     public static long suggestTargetSize(long sizeEstimate) {
-        long est = sizeEstimate / LEAF_TARGET;
+        long est = sizeEstimate / getLeafTarget();
         return est > 0L ? est : 1L;
     }
 
diff --git a/src/share/classes/java/util/stream/ForEachOps.java b/src/share/classes/java/util/stream/ForEachOps.java
index b527f054f4..38d63e929d 100644
--- a/src/share/classes/java/util/stream/ForEachOps.java
+++ b/src/share/classes/java/util/stream/ForEachOps.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2012, 2013, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2012, 2017, Oracle and/or its affiliates. All rights reserved.
  * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
  *
  * This code is free software; you can redistribute it and/or modify it
@@ -28,7 +28,6 @@ import java.util.Objects;
 import java.util.Spliterator;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.CountedCompleter;
-import java.util.concurrent.ForkJoinTask;
 import java.util.function.Consumer;
 import java.util.function.DoubleConsumer;
 import java.util.function.IntConsumer;
@@ -378,7 +377,7 @@ final class ForEachOps {
             this.spliterator = spliterator;
             this.targetSize = AbstractTask.suggestTargetSize(spliterator.estimateSize());
             // Size map to avoid concurrent re-sizes
-            this.completionMap = new ConcurrentHashMap<>(Math.max(16, AbstractTask.LEAF_TARGET << 1));
+            this.completionMap = new ConcurrentHashMap<>(Math.max(16, AbstractTask.getLeafTarget() << 1));
             this.action = action;
             this.leftPredecessor = null;
         }
diff --git a/src/share/classes/java/util/stream/StreamSpliterators.java b/src/share/classes/java/util/stream/StreamSpliterators.java
index 4f2d2a2c56..42ad027f70 100644
--- a/src/share/classes/java/util/stream/StreamSpliterators.java
+++ b/src/share/classes/java/util/stream/StreamSpliterators.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2012, 2013, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2012, 2017, Oracle and/or its affiliates. All rights reserved.
  * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
  *
  * This code is free software; you can redistribute it and/or modify it
@@ -898,7 +898,7 @@ class StreamSpliterators {
      * Note: The source spliterator may report {@code ORDERED} since that
      * spliterator be the result of a previous pipeline stage that was
      * collected to a {@code Node}. It is the order of the pipeline stage
-     * that governs whether the this slice spliterator is to be used or not.
+     * that governs whether this slice spliterator is to be used or not.
      */
     static abstract class UnorderedSliceSpliterator<T, T_SPLITR extends Spliterator<T>> {
         static final int CHUNK_SIZE = 1 << 7;
@@ -915,7 +915,7 @@ class StreamSpliterators {
             this.unlimited = limit < 0;
             this.skipThreshold = limit >= 0 ? limit : 0;
             this.chunkSize = limit >= 0 ? (int)Math.min(CHUNK_SIZE,
-                ((skip + limit) / AbstractTask.LEAF_TARGET) + 1) : CHUNK_SIZE;
+                                                        ((skip + limit) / AbstractTask.getLeafTarget()) + 1) : CHUNK_SIZE;
             this.permits = new AtomicLong(limit >= 0 ? skip + limit : skip);
         }