Support suite level thread pools for data provider

Closes testng-team#2980

We can now configure TestNG such that it uses a
Suite level thread pool when running data driven
Tests in parallel.

This can be enabled via the configuration 

“-shareThreadPoolForDataProviders” 
with a value of “true”

Alternatively one can also use the suite level 
attribute “share-thread-pool-for-data-providers”

Author: krmahadevan

CHANGES.txt

@@ -1,5 +1,6 @@
 Current
 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)
 New:   Added @Inherited to the Listeners annotation, allowing it to be used in forming meta-annotations. (Pavlo Glushchenko)
 Fixed: GITHUB-2991: Suite attributes map should be thread safe (Krishnan Mahadevan)

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

@@ -6,6 +6,7 @@
 import java.util.Locale;
 import java.util.Map;
 import java.util.Set;
+import java.util.UUID;
 import org.testng.ITestObjectFactory;
 import org.testng.collections.Lists;
 import org.testng.collections.Maps;
@@ -123,6 +124,9 @@ public String toString() {
   private String m_parentModule = "";
   private String m_guiceStage = "";
 
+  /** Represents a unique id for this suite. Can be used for uniquely identifying the xml suite. */
+  public final UUID SUITE_ID = UUID.randomUUID();
+
   /** Whether to SKIP or CONTINUE to re-attempt failed configuration methods. */
   public static final FailurePolicy DEFAULT_CONFIG_FAILURE_POLICY = FailurePolicy.SKIP;
 
@@ -139,6 +143,8 @@ public String toString() {
   public static final Boolean DEFAULT_SKIP_FAILED_INVOCATION_COUNTS = Boolean.FALSE;
   private Boolean m_skipFailedInvocationCounts = DEFAULT_SKIP_FAILED_INVOCATION_COUNTS;
 
+  private boolean shareThreadPoolForDataProviders = false;
+
   /** The thread count. */
   public static final Integer DEFAULT_THREAD_COUNT = 5;
 
@@ -243,6 +249,14 @@ public Class<? extends ITestObjectFactory> getObjectFactoryClass() {
     return m_objectFactoryClass;
   }
 
+  public void setShareThreadPoolForDataProviders(boolean shareThreadPoolForDataProviders) {
+    this.shareThreadPoolForDataProviders = shareThreadPoolForDataProviders;
+  }
+
+  public boolean isShareThreadPoolForDataProviders() {
+    return shareThreadPoolForDataProviders;
+  }
+
   @Deprecated
   public void setObjectFactory(ITestObjectFactory objectFactory) {
     setObjectFactoryClass(objectFactory.getClass());

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

@@ -274,4 +274,13 @@ public class CommandLineArgs {
       names = GENERATE_RESULTS_PER_SUITE,
       description = "Should TestNG consider failures in Data Providers  as test failures.")
   public Boolean generateResultsPerSuite = false;
+
+  public static final String SHARE_THREAD_POOL_FOR_DATA_PROVIDERS =
+      "-shareThreadPoolForDataProviders";
+
+  @Parameter(
+      names = SHARE_THREAD_POOL_FOR_DATA_PROVIDERS,
+      description =
+          "Should TestNG use a global Shared ThreadPool (At suite level) for running data providers.")
+  public Boolean shareThreadPoolForDataProviders = false;
 }

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

@@ -31,6 +31,7 @@
 import org.testng.internal.ExitCode;
 import org.testng.internal.IConfiguration;
 import org.testng.internal.ListenerOrderDeterminer;
+import org.testng.internal.ObjectBag;
 import org.testng.internal.OverrideProcessor;
 import org.testng.internal.ReporterConfig;
 import org.testng.internal.RuntimeBehavior;
@@ -609,6 +610,14 @@ public boolean isPropagateDataProviderFailureAsTestFailure() {
     return this.m_configuration.isPropagateDataProviderFailureAsTestFailure();
   }
 
+  public void shareThreadPoolForDataProviders(boolean flag) {
+    this.m_configuration.shareThreadPoolForDataProviders(flag);
+  }
+
+  public boolean isShareThreadPoolForDataProviders() {
+    return this.m_configuration.isShareThreadPoolForDataProviders();
+  }
+
   /**
    * 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.
@@ -1082,6 +1091,7 @@ public void run() {
     m_end = System.currentTimeMillis();
 
     if (null != suiteRunners) {
+      suiteRunners.forEach(ObjectBag::cleanup);
       generateReports(suiteRunners);
     }
 
@@ -1186,6 +1196,9 @@ public List<ISuite> runSuitesLocally() {
     // First initialize the suite runners to ensure there are no configuration issues.
     // Create a map with XmlSuite as key and corresponding SuiteRunner as value
     for (XmlSuite xmlSuite : m_suites) {
+      if (m_configuration.isShareThreadPoolForDataProviders()) {
+        xmlSuite.setShareThreadPoolForDataProviders(true);
+      }
       createSuiteRunners(suiteRunnerMap, xmlSuite);
     }
 
@@ -1454,6 +1467,8 @@ public static TestNG privateMain(String[] argv, ITestListener listener) {
    * @param cla The command line parameters
    */
   protected void configure(CommandLineArgs cla) {
+    Optional.ofNullable(cla.shareThreadPoolForDataProviders)
+        .ifPresent(this::shareThreadPoolForDataProviders);
     Optional.ofNullable(cla.propagateDataProviderFailureAsTestFailure)
         .ifPresent(value -> propagateDataProviderFailureAsTestFailure());
     setReportAllDataDrivenTestsAsSkipped(cla.includeAllDataDrivenTestsWhenSkipping);

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

@@ -23,6 +23,8 @@ public class Configuration implements IConfiguration {
   private ITestObjectFactory m_objectFactory;
   private IHookable m_hookable;
   private IConfigurable m_configurable;
+
+  private boolean shareThreadPoolForDataProviders = false;
   private final Map<Class<? extends IExecutionListener>, IExecutionListener> m_executionListeners =
       Maps.newLinkedHashMap();
   private final Map<Class<? extends IConfigurationListener>, IConfigurationListener>
@@ -168,4 +170,14 @@ public void propagateDataProviderFailureAsTestFailure() {
   public boolean isPropagateDataProviderFailureAsTestFailure() {
     return propagateDataProviderFailureAsTestFailure;
   }
+
+  @Override
+  public boolean isShareThreadPoolForDataProviders() {
+    return this.shareThreadPoolForDataProviders;
+  }
+
+  @Override
+  public void shareThreadPoolForDataProviders(boolean flag) {
+    this.shareThreadPoolForDataProviders = flag;
+  }
 }

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

@@ -59,4 +59,8 @@ default boolean getReportAllDataDrivenTestsAsSkipped() {
   void propagateDataProviderFailureAsTestFailure();
 
   boolean isPropagateDataProviderFailureAsTestFailure();
+
+  boolean isShareThreadPoolForDataProviders();
+
+  void shareThreadPoolForDataProviders(boolean flag);
 }

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

@@ -0,0 +1,57 @@
+package org.testng.internal;
+
+import java.io.Closeable;
+import java.io.IOException;
+import java.util.Map;
+import java.util.Optional;
+import java.util.UUID;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.function.Supplier;
+import org.testng.ISuite;
+import org.testng.log4testng.Logger;
+
+/**
+ * A simple bean bag that is intended to help share objects during the lifetime of TestNG without
+ * needing it to be a singleton.
+ */
+public final class ObjectBag {
+
+  private static final Logger logger = Logger.getLogger(ObjectBag.class);
+  private final Map<Class<?>, Object> bag = new ConcurrentHashMap<>();
+
+  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());
+  }
+
+  public static void cleanup(ISuite suite) {
+    UUID uid = suite.getXmlSuite().SUITE_ID;
+    Optional.ofNullable(instances.get(uid)).ifPresent(ObjectBag::cleanup);
+    instances.remove(uid);
+  }
+
+  /**
+   * @param type - The type of the object to be created
+   * @param supplier - A {@link Supplier} that should be used to produce a new instance
+   * @return - Either the newly produced instance or the existing instance.
+   */
+  public Object createIfRequired(Class<?> type, Supplier<Object> supplier) {
+    return bag.computeIfAbsent(type, t -> supplier.get());
+  }
+
+  public void cleanup() {
+    bag.values().stream()
+        .filter(it -> it instanceof Closeable)
+        .map(it -> (Closeable) it)
+        .forEach(
+            it -> {
+              try {
+                it.close();
+              } catch (IOException e) {
+                logger.debug("Could not clean-up " + it, e);
+              }
+            });
+    bag.clear();
+  }
+}

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

@@ -1,5 +1,7 @@
 package org.testng.internal;
 
+import java.io.Closeable;
+import java.io.IOException;
 import java.util.List;
 import java.util.concurrent.Callable;
 import java.util.concurrent.ExecutionException;
@@ -15,12 +17,18 @@
 import org.testng.internal.thread.ThreadUtil;
 
 /** Simple wrapper for an ExecutorCompletionService. */
-public class PoolService<FutureType> {
+public class PoolService<FutureType> implements Closeable {
 
   private final ExecutorCompletionService<FutureType> m_completionService;
   private final ExecutorService m_executor;
 
+  private final boolean shutdownAfterExecution;
+
   public PoolService(int threadPoolSize) {
+    this(threadPoolSize, true);
+  }
+
+  public PoolService(int threadPoolSize, boolean shutdownAfterExecution) {
 
     ThreadFactory threadFactory =
         new ThreadFactory() {
@@ -35,6 +43,7 @@ public Thread newThread(@Nonnull Runnable r) {
         };
     m_executor = Executors.newFixedThreadPool(threadPoolSize, threadFactory);
     m_completionService = new ExecutorCompletionService<>(m_executor);
+    this.shutdownAfterExecution = shutdownAfterExecution;
   }
 
   public List<FutureType> submitTasksAndWait(List<? extends Callable<FutureType>> tasks) {
@@ -53,7 +62,16 @@ public List<FutureType> submitTasksAndWait(List<? extends Callable<FutureType>>
       }
     }
 
-    m_executor.shutdown();
+    if (shutdownAfterExecution) {
+      m_executor.shutdown();
+    }
     return result;
   }
+
+  @Override
+  public void close() throws IOException {
+    if (!shutdownAfterExecution) {
+      m_executor.shutdown();
+    }
+  }
 }

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

@@ -7,6 +7,7 @@
 import org.testng.ITestResult;
 import org.testng.collections.CollectionUtils;
 import org.testng.collections.Lists;
+import org.testng.internal.ObjectBag;
 import org.testng.internal.Parameters;
 import org.testng.internal.PoolService;
 import org.testng.internal.invokers.ITestInvoker.FailureContext;
@@ -129,7 +130,18 @@ public List<ITestResult> runInParallel(
       // testng387: increment the param index in the bag.
       parametersIndex += 1;
     }
-    PoolService<List<ITestResult>> ps = new PoolService<>(suite.getDataProviderThreadCount());
+
+    ObjectBag objectBag = ObjectBag.getInstance(context.getSuite());
+    boolean sharedThreadPool = context.getSuite().getXmlSuite().isShareThreadPoolForDataProviders();
+
+    @SuppressWarnings("unchecked")
+    PoolService<List<ITestResult>> ps =
+        sharedThreadPool
+            ? (PoolService<List<ITestResult>>)
+                objectBag.createIfRequired(
+                    PoolService.class,
+                    () -> new PoolService<>(suite.getDataProviderThreadCount(), false))
+            : new PoolService<>(suite.getDataProviderThreadCount());
     List<List<ITestResult>> r = ps.submitTasksAndWait(workers);
     for (List<ITestResult> l2 : r) {
       result.addAll(l2);

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

@@ -18,6 +18,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Optional;
 import java.util.Stack;
 import org.testng.ITestObjectFactory;
 import org.testng.TestNGException;
@@ -267,6 +268,15 @@ private void xmlSuite(boolean start, Attributes attributes) {
       if (null != dataProviderThreadCount) {
         m_currentSuite.setDataProviderThreadCount(Integer.parseInt(dataProviderThreadCount));
       }
+
+      String shareThreadPoolForDataProviders =
+          attributes.getValue("share-thread-pool-for-data-providers");
+      Optional.ofNullable(shareThreadPoolForDataProviders)
+          .ifPresent(
+              it ->
+                  m_currentSuite.setShareThreadPoolForDataProviders(
+                      Boolean.parseBoolean(shareThreadPoolForDataProviders)));
+
       String timeOut = attributes.getValue("time-out");
       if (null != timeOut) {
         m_currentSuite.setTimeOut(timeOut);

testng-core/src/main/resources/testng-1.0.dtd

@@ -56,11 +56,13 @@ Cedric Beust & Alexandru Popescu
 @attr  skipfailedinvocationcounts Whether to skip failed invocations.
 @attr  data-provider-thread-count An integer giving the size of the thread pool to use
        for parallel data providers.
+@attr share-thread-pool-for-data-providers - Whether TestNG should use a common thread pool
+      for running parallel data providers. (Works only with TestNG versions 7.9.0 or higher)
 @attr  object-factory A class that implements IObjectFactory that will be used to
        instantiate the test objects.
 @attr allow-return-values If true, tests that return a value will be run as well
 -->
-<!ATTLIST suite 
+<!ATTLIST suite
     name CDATA #REQUIRED
     junit (true | false) "false"
     verbose CDATA #IMPLIED
@@ -73,6 +75,7 @@ Cedric Beust & Alexandru Popescu
     time-out CDATA #IMPLIED
     skipfailedinvocationcounts (true | false) "false"
     data-provider-thread-count CDATA "10"
+    share-thread-pool-for-data-providers (true | false) "false"
     object-factory CDATA #IMPLIED
     group-by-instances (true | false) "false"
     preserve-order (true | false) "true"