Shared Threadpool for normal/datadriven tests.

Closes #2019

Author: krmahadevan

CHANGES.txt

@@ -1,4 +1,5 @@
 Current
+Fixed: GITHUB-2019: Total thread count in testng parallel tests with dataproviders (Krishnan Mahadevan)
 Fixed: GITHUB-3006: ITestResult injected at @AfterMethod incorrect when a configuration method failed (Krishnan Mahadevan)
 Fixed: GITHUB-2980: Data Provider Threads configuration in the suite don't match the documentation (Krishnan Mahadevan)
 Fixed: GITHUB-3003: BeforeClass|AfterClass with inheritedGroups triggers cyclic dependencies (Krishnan Mahadevan)

testng-core-api/src/main/java/org/testng/xml/XmlSuite.java

@@ -145,6 +145,8 @@ public String toString() {
 
   private boolean shareThreadPoolForDataProviders = false;
 
+  private boolean useGlobalThreadPool = false;
+
   /** The thread count. */
   public static final Integer DEFAULT_THREAD_COUNT = 5;
 
@@ -253,6 +255,14 @@ public void setShareThreadPoolForDataProviders(boolean shareThreadPoolForDataPro
     this.shareThreadPoolForDataProviders = shareThreadPoolForDataProviders;
   }
 
+  public boolean useGlobalThreadPool() {
+    return this.useGlobalThreadPool;
+  }
+
+  public void shouldUseGlobalThreadPool(boolean flag) {
+    this.useGlobalThreadPool = flag;
+  }
+
   public boolean isShareThreadPoolForDataProviders() {
     return shareThreadPoolForDataProviders;
   }

testng-core/src/main/java/org/testng/CommandLineArgs.java

@@ -283,4 +283,12 @@ public class CommandLineArgs {
       description =
           "Should TestNG use a global Shared ThreadPool (At suite level) for running data providers.")
   public Boolean shareThreadPoolForDataProviders = false;
+
+  public static final String USE_GLOBAL_THREAD_POOL = "-useGlobalThreadPool";
+
+  @Parameter(
+      names = USE_GLOBAL_THREAD_POOL,
+      description =
+          "Should TestNG use a global Shared ThreadPool (At suite level) for running regular and data driven tests.")
+  public Boolean useGlobalThreadPool = false;
 }

testng-core/src/main/java/org/testng/TestNG.java

@@ -616,6 +616,14 @@ public boolean isShareThreadPoolForDataProviders() {
     return this.m_configuration.isShareThreadPoolForDataProviders();
   }
 
+  public boolean useGlobalThreadPool() {
+    return this.m_configuration.useGlobalThreadPool();
+  }
+
+  public void shouldUseGlobalThreadPool(boolean flag) {
+    this.m_configuration.shouldUseGlobalThreadPool(flag);
+  }
+
   /**
    * Set the suites file names to be run by this TestNG object. This method tries to load and parse
    * the specified TestNG suite xml files. If a file is missing, it is ignored.
@@ -1203,6 +1211,9 @@ public List<ISuite> runSuitesLocally() {
       if (m_configuration.isShareThreadPoolForDataProviders()) {
         xmlSuite.setShareThreadPoolForDataProviders(true);
       }
+      if (m_configuration.useGlobalThreadPool()) {
+        xmlSuite.shouldUseGlobalThreadPool(true);
+      }
       createSuiteRunners(suiteRunnerMap, xmlSuite);
     }
 
@@ -1457,6 +1468,7 @@ public static TestNG privateMain(String[] argv, ITestListener listener) {
    * @param cla The command line parameters
    */
   protected void configure(CommandLineArgs cla) {
+    Optional.ofNullable(cla.useGlobalThreadPool).ifPresent(this::shouldUseGlobalThreadPool);
     Optional.ofNullable(cla.shareThreadPoolForDataProviders)
         .ifPresent(this::shareThreadPoolForDataProviders);
     Optional.ofNullable(cla.propagateDataProviderFailureAsTestFailure)

testng-core/src/main/java/org/testng/TestTaskExecutor.java

@@ -5,7 +5,9 @@
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.ThreadPoolExecutor;
 import java.util.concurrent.TimeUnit;
+import java.util.function.Supplier;
 import org.testng.internal.IConfiguration;
+import org.testng.internal.ObjectBag;
 import org.testng.internal.RuntimeBehavior;
 import org.testng.internal.Utils;
 import org.testng.internal.thread.TestNGThreadFactory;
@@ -47,8 +49,7 @@ public TestTaskExecutor(
 
   public void execute() {
     String name = "test-" + xmlTest.getName();
-    int threadCount = xmlTest.getThreadCount();
-    threadCount = Math.max(threadCount, 1);
+    int threadCount = Math.max(xmlTest.getThreadCount(), 1);
     if (RuntimeBehavior.favourCustomThreadPoolExecutor()) {
       IExecutorFactory execFactory = configuration.getExecutorFactory();
       ITestNGThreadPoolExecutor executor =
@@ -65,14 +66,22 @@ public void execute() {
       executor.run();
       service = executor;
     } else {
-      service =
-          new ThreadPoolExecutor(
-              threadCount,
-              threadCount,
-              0,
-              TimeUnit.MILLISECONDS,
-              queue,
-              new TestNGThreadFactory(name));
+      boolean reUse = xmlTest.getSuite().useGlobalThreadPool();
+      Supplier<Object> supplier =
+          () ->
+              new ThreadPoolExecutor(
+                  threadCount,
+                  threadCount,
+                  0,
+                  TimeUnit.MILLISECONDS,
+                  queue,
+                  new TestNGThreadFactory(name));
+      if (reUse) {
+        ObjectBag bag = ObjectBag.getInstance(xmlTest.getSuite());
+        service = (ExecutorService) bag.createIfRequired(ExecutorService.class, supplier);
+      } else {
+        service = (ExecutorService) supplier.get();
+      }
       GraphOrchestrator<ITestNGMethod> executor =
           new GraphOrchestrator<>(service, factory, graph, comparator);
       executor.run();

testng-core/src/main/java/org/testng/internal/Configuration.java

@@ -39,6 +39,8 @@ public class Configuration implements IConfiguration {
 
   private boolean propagateDataProviderFailureAsTestFailure;
 
+  private boolean useGlobalThreadPool = false;
+
   public Configuration() {
     init(new JDK15AnnotationFinder(new DefaultAnnotationTransformer()));
   }
@@ -180,4 +182,14 @@ public boolean isShareThreadPoolForDataProviders() {
   public void shareThreadPoolForDataProviders(boolean flag) {
     this.shareThreadPoolForDataProviders = flag;
   }
+
+  @Override
+  public boolean useGlobalThreadPool() {
+    return this.useGlobalThreadPool;
+  }
+
+  @Override
+  public void shouldUseGlobalThreadPool(boolean flag) {
+    this.useGlobalThreadPool = flag;
+  }
 }

testng-core/src/main/java/org/testng/internal/IConfiguration.java

@@ -63,4 +63,8 @@ default boolean getReportAllDataDrivenTestsAsSkipped() {
   boolean isShareThreadPoolForDataProviders();
 
   void shareThreadPoolForDataProviders(boolean flag);
+
+  boolean useGlobalThreadPool();
+
+  void shouldUseGlobalThreadPool(boolean flag);
 }

testng-core/src/main/java/org/testng/internal/ObjectBag.java

@@ -8,6 +8,7 @@
 import java.util.function.Supplier;
 import org.testng.ISuite;
 import org.testng.log4testng.Logger;
+import org.testng.xml.XmlSuite;
 
 /**
  * A simple bean bag that is intended to help share objects during the lifetime of TestNG without
@@ -21,7 +22,11 @@ public final class ObjectBag {
   private static final Map<UUID, ObjectBag> instances = new ConcurrentHashMap<>();
 
   public static ObjectBag getInstance(ISuite suite) {
-    return instances.computeIfAbsent(suite.getXmlSuite().SUITE_ID, k -> new ObjectBag());
+    return getInstance(suite.getXmlSuite());
+  }
+
+  public static ObjectBag getInstance(XmlSuite xmlSuite) {
+    return instances.computeIfAbsent(xmlSuite.SUITE_ID, k -> new ObjectBag());
   }
 
   public static void cleanup(ISuite suite) {

testng-core/src/main/java/org/testng/internal/invokers/MethodRunner.java

@@ -12,6 +12,8 @@
 import java.util.concurrent.Executors;
 import java.util.concurrent.ThreadFactory;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Supplier;
 import java.util.stream.Collectors;
 import org.testng.ITestContext;
 import org.testng.ITestResult;
@@ -21,7 +23,7 @@
 import org.testng.internal.Parameters;
 import org.testng.internal.invokers.ITestInvoker.FailureContext;
 import org.testng.internal.invokers.TestMethodArguments.Builder;
-import org.testng.internal.thread.ThreadUtil;
+import org.testng.internal.thread.TestNGThreadFactory;
 import org.testng.xml.XmlSuite;
 
 public class MethodRunner implements IMethodRunner {
@@ -107,9 +109,9 @@ public List<ITestResult> runInParallel(
     XmlSuite suite = context.getSuite().getXmlSuite();
     int parametersIndex = 0;
     ObjectBag objectBag = ObjectBag.getInstance(context.getSuite());
-    boolean reUse = suite.isShareThreadPoolForDataProviders();
+    boolean reUse = suite.isShareThreadPoolForDataProviders() || suite.useGlobalThreadPool();
 
-    ExecutorService service = getOrCreate(reUse, suite.getDataProviderThreadCount(), objectBag);
+    ExecutorService service = getOrCreate(reUse, suite, objectBag);
     List<CompletableFuture<List<ITestResult>>> all = new ArrayList<>();
     for (Object[] next : CollectionUtils.asIterable(allParamValues)) {
       if (next == null) {
@@ -163,18 +165,20 @@ public List<ITestResult> runInParallel(
     return result;
   }
 
-  private static ExecutorService getOrCreate(boolean reUse, int count, ObjectBag objectBag) {
+  private static ExecutorService getOrCreate(boolean reUse, XmlSuite suite, ObjectBag objectBag) {
+    AtomicReference<Integer> count = new AtomicReference<>();
+    count.set(suite.getDataProviderThreadCount());
+    Supplier<Object> supplier = () -> Executors.newFixedThreadPool(count.get(), threadFactory());
     if (reUse) {
-      return (ExecutorService)
-          objectBag.createIfRequired(
-              ExecutorService.class, () -> Executors.newFixedThreadPool(count, threadFactory()));
+      if (suite.useGlobalThreadPool()) {
+        count.set(suite.getThreadCount());
+      }
+      return (ExecutorService) objectBag.createIfRequired(ExecutorService.class, supplier);
     }
-    return Executors.newFixedThreadPool(count, threadFactory());
+    return (ExecutorService) supplier.get();
   }
 
   private static ThreadFactory threadFactory() {
-    AtomicInteger threadNumber = new AtomicInteger(0);
-    return r ->
-        new Thread(r, ThreadUtil.THREAD_NAME + "-PoolService-" + threadNumber.getAndIncrement());
+    return new TestNGThreadFactory("PoolService");
   }
 }

testng-core/src/main/java/org/testng/xml/TestNGContentHandler.java

@@ -277,6 +277,13 @@ private void xmlSuite(boolean start, Attributes attributes) {
                   m_currentSuite.setShareThreadPoolForDataProviders(
                       Boolean.parseBoolean(shareThreadPoolForDataProviders)));
 
+      String useGlobalThreadPool = attributes.getValue("use-global-thread-pool");
+      Optional.ofNullable(useGlobalThreadPool)
+          .ifPresent(
+              it ->
+                  m_currentSuite.shouldUseGlobalThreadPool(
+                      Boolean.parseBoolean(useGlobalThreadPool)));
+
       String timeOut = attributes.getValue("time-out");
       if (null != timeOut) {
         m_currentSuite.setTimeOut(timeOut);