[CI] Parallelize unit-test phase with fork-safe port/temp-dir isolation

.github/workflows/scripts/pr-tests/.pinot_tests_unit.sh

@@ -26,30 +26,73 @@ ifconfig
 netstat -i
 
 # Unit Tests
-#   - TEST_SET#1 runs install and test together so the module list must ensure no additional modules were tested
-#     due to the -am flag (include dependency modules)
-#   - tests for pinot-plugins should not be ran multi-threaded
+#   - Both test sets run plain `mvn test` (no install, no -am): the modules were already built
+#     and installed by .pinot_tests_build.sh, so only the modules listed here are tested.
+#
+# Parallelism / memory:
+#   - UNIT_TEST_FORK_COUNT (default 3) sets surefire forkCount so test *classes* run in
+#     separate parallel JVMs (reuseForks=false keeps one class per JVM). This is
+#     process-level isolation, not TestNG intra-JVM threading, so tests that were unsafe
+#     to run multi-threaded within a single JVM (e.g. pinot-plugins) are unaffected.
+#     Cross-fork resource collisions (ZK/controller ports, temp dirs) are avoided by
+#     offsetting per surefire.forkNumber; embedded Kafka clusters use ephemeral ports.
+#     This is the main lever for shortening the unit-test phase. 3 forks on the 4-vCPU
+#     runner keeps a core free for the Maven reactor / GC while test JVMs spend much of
+#     their time blocked on ZK/Helix/socket startup, so the extra fork still pays off.
+#   - UNIT_TEST_FORK_HEAP (default 2500m) caps per-fork heap so N forks fit in the
+#     runner's memory (N * heap + the mvn JVM must stay under the runner's RAM).
+#   - UNIT_TEST_RERUN_COUNT (default 0) retries a failing test before failing the build. Left at 0
+#     because the load-sensitive flaky tests parallel forks exposed are fixed at the root cause
+#     (SegmentPreProcessorTest mtime granularity, LuceneMutableTextIndexTest NRT-refresh wait). It
+#     remains overridable as an escape hatch if a new flake appears, but is intentionally not a
+#     standing default so real failures are never masked.
+UNIT_TEST_FORK_COUNT="${UNIT_TEST_FORK_COUNT:-3}"
+# 2500m/fork: 3 forks * 2500m + the 2g Maven JVM (~9.5g) stays well under the runner's 16g.
+UNIT_TEST_FORK_HEAP="${UNIT_TEST_FORK_HEAP:-2500m}"
+UNIT_TEST_RERUN_COUNT="${UNIT_TEST_RERUN_COUNT:-0}"
+# Coverage adds ~30% to the test phase (JaCoCo agent per fork + aggregate report). Keep it on by
+# default to preserve Codecov behavior; set RUN_CODECOVERAGE=false (e.g. on PRs) to trade coverage
+# for a faster run.
+RUN_CODECOVERAGE="${RUN_CODECOVERAGE:-true}"
+# Fork-scope the JaCoCo exec file (jacoco-<forkNumber>.exec) so parallel forks don't append to
+# one shared jacoco.exec and corrupt coverage. Only the unit lane sets this; other lanes keep
+# the default empty suffix (target/jacoco.exec).
+FORK_OPTS="-Dunit.test.fork.count=${UNIT_TEST_FORK_COUNT} -Dunit.test.fork.heap=${UNIT_TEST_FORK_HEAP} -Dunit.test.rerun.count=${UNIT_TEST_RERUN_COUNT} -Djacoco.exec.suffix=-\${surefire.forkNumber}"
+if [ "$RUN_CODECOVERAGE" == "true" ]; then
+  COVERAGE_PROFILE=",codecoverage"
+else
+  COVERAGE_PROFILE=""
+fi
 if [ "$RUN_TEST_SET" == "1" ]; then
-  mvn test \
+  # pinot-segment-local's tests run in set #2 to balance pinot-core's longer test time in this
+  # shard against set #2's longer build. It remains built in set #1 as a pinot-core dependency.
+  # No -am on this command, so only the listed modules test.
+  mvn test ${FORK_OPTS} \
       -pl 'pinot-spi' \
       -pl 'pinot-segment-spi' \
       -pl 'pinot-common' \
       -pl ':pinot-yammer' \
       -pl 'pinot-core' \
       -pl 'pinot-query-planner' \
       -pl 'pinot-query-runtime' \
-      -P github-actions,codecoverage,no-integration-tests || exit 1
+      -P github-actions,no-integration-tests${COVERAGE_PROFILE} || exit 1
 fi
 if [ "$RUN_TEST_SET" == "2" ]; then
-  mvn test \
+  mvn test ${FORK_OPTS} \
     -pl '!pinot-spi' \
     -pl '!pinot-segment-spi' \
     -pl '!pinot-common' \
     -pl '!pinot-core' \
     -pl '!pinot-query-planner' \
     -pl '!pinot-query-runtime' \
     -pl '!:pinot-yammer' \
-    -P github-actions,codecoverage,no-integration-tests || exit 1
+    -P github-actions,no-integration-tests${COVERAGE_PROFILE} || exit 1
 fi
 
-mvn jacoco:report-aggregate@report -P codecoverage || exit 1
+# Aggregate coverage across all per-fork exec files (jacoco-*.exec) written under forkCount>1,
+# while still matching the single-fork jacoco.exec produced by non-parallel runs. Skipped when
+# coverage is disabled.
+if [ "$RUN_CODECOVERAGE" == "true" ]; then
+  mvn jacoco:report-aggregate@report -P codecoverage \
+    -Djacoco.dataFileIncludes='**/target/jacoco-*.exec,**/target/jacoco.exec' || exit 1
+fi

pinot-broker/src/test/java/org/apache/pinot/broker/queryquota/HelixExternalViewBasedQueryQuotaManagerTest.java

@@ -327,10 +327,10 @@ public void testWhenOnlyDefaultAppQuotaIsSetItAffectsAllApplications()
     Assert.assertEquals(_queryQuotaManager.getDatabaseRateLimiterMap().size(), 1);
     Assert.assertEquals(_queryQuotaManager.getApplicationRateLimiterMap().size(), 1);
 
-    runQueries(100, true, APP_NAME);
-    runQueries(100, true, "otherApp");
+    assertApplicationRateLimitedInBurst(APP_NAME, 100);
+    assertApplicationRateLimitedInBurst("otherApp", 100);
     runQueries(100, false, "someApp");
-    runQueries(201, true, "someApp");
+    assertApplicationRateLimitedInBurst("someApp", 201);
 
     Assert.assertEquals(_queryQuotaManager.getApplicationRateLimiterMap().size(), 3);
     _queryQuotaManager.dropTableQueryQuota(OFFLINE_TABLE_NAME);
@@ -697,4 +697,13 @@ private void runQueries(double qps, boolean shouldFail, String appName)
       Assert.assertTrue(failCount == 0, "Expected no failure with qps: " + qps + " and app :" + appName);
     }
   }
+
+  private void assertApplicationRateLimitedInBurst(String appName, int numQueries) {
+    for (int i = 0; i < numQueries; i++) {
+      if (!_queryQuotaManager.acquireApplication(appName)) {
+        return;
+      }
+    }
+    Assert.fail("Expected application rate limiting for " + numQueries + " queries and app: " + appName);
+  }
 }

pinot-common/pom.xml

@@ -59,7 +59,8 @@
         <groupId>org.apache.maven.plugins</groupId>
         <artifactId>maven-surefire-plugin</artifactId>
         <configuration>
-          <forkCount>1</forkCount>
+          <!-- Keep this module's historical fork reuse, but honor the root parallel-fork knob. -->
+          <forkCount>${unit.test.fork.count}</forkCount>
           <reuseForks>true</reuseForks>
           <properties>
             <property>

pinot-common/src/main/java/org/apache/pinot/common/utils/ZkStarter.java

@@ -21,6 +21,7 @@
 import java.io.File;
 import java.io.IOException;
 import java.net.InetSocketAddress;
+import java.util.UUID;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.TimeUnit;
@@ -42,6 +43,33 @@ private ZkStarter() {
   public static final int DEFAULT_ZK_TEST_PORT = 2191;
   private static final int DEFAULT_ZK_CLIENT_RETRIES = 10;
 
+  /// Per-fork offset applied to the default test port so that concurrent surefire forks
+  /// (forkCount > 1, reuseForks=false) do not scan from the same base port and collide.
+  ///
+  /// This is purely a test-harness hook: the offset is derived from the `surefire.forkNumber`
+  /// system property, which surefire injects only inside a forked test JVM (1-based, so it is 1
+  /// even at forkCount=1, and 1..N under parallel forks). In any production process that property
+  /// is absent, so `forkNumber()` returns 0 and the no-arg {@link #startLocalZkServer()} scans
+  /// from the historical {@link #DEFAULT_ZK_TEST_PORT}. Under tests each fork scans from a distinct
+  /// base ({@code DEFAULT_ZK_TEST_PORT + forkNumber*1000}); the exact port is still chosen by
+  /// findOpenPort and read back via {@code getZkUrl()}, so no caller depends on the literal base.
+  /// The stride (1000) is large enough that a fork exhausting ports below the next boundary
+  /// (findOpenPort scans upward) does not spill into the neighboring fork's band.
+  ///
+  /// The offset is deliberately centralized on the no-arg entry point rather than pushed into each
+  /// test: several tests across different modules call {@link #startLocalZkServer()} directly, and
+  /// duplicating the fork math into each caller (or introducing a parallel test-only start helper)
+  /// is more surface and more error-prone than one guarded, production-inert read here.
+  private static final int FORK_PORT_OFFSET = forkNumber() * 1000;
+
+  private static int forkNumber() {
+    try {
+      return Integer.parseInt(System.getProperty("surefire.forkNumber", "0"));
+    } catch (NumberFormatException e) {
+      return 0;
+    }
+  }
+
   public static class ZookeeperInstance {
     private PublicZooKeeperServerMain _serverMain;
     private String _dataDirPath;
@@ -135,9 +163,10 @@ public void shutdown() {
     }
   }
 
-  /// Starts an empty local Zk instance on the default port
+  /// Starts an empty local Zk instance on the default port (offset per surefire fork so that
+  /// concurrent forks bind disjoint port ranges).
   public static ZookeeperInstance startLocalZkServer() {
-    return startLocalZkServer(NetUtils.findOpenPort(DEFAULT_ZK_TEST_PORT));
+    return startLocalZkServer(NetUtils.findOpenPort(DEFAULT_ZK_TEST_PORT + FORK_PORT_OFFSET));
   }
 
   public static String getDefaultZkStr() {
@@ -147,8 +176,10 @@ public static String getDefaultZkStr() {
   /// Starts a local Zk instance with a generated empty data directory
   /// @param port The port to listen on
   public static ZookeeperInstance startLocalZkServer(final int port) {
+    // Use a random UUID rather than a timestamp so that concurrent forks/threads never share a
+    // ZK data directory (System.currentTimeMillis() collides when two instances start in the same ms).
     return startLocalZkServer(port,
-        org.apache.commons.io.FileUtils.getTempDirectoryPath() + File.separator + "test-" + System.currentTimeMillis());
+        org.apache.commons.io.FileUtils.getTempDirectoryPath() + File.separator + "test-" + UUID.randomUUID());
   }
 
   /// Starts a local Zk instance

pinot-controller/pom.xml

@@ -172,12 +172,41 @@
         <groupId>org.apache.maven.plugins</groupId>
         <artifactId>maven-surefire-plugin</artifactId>
         <configuration>
+          <!-- The stateful and stateless TestNG suites run in separate forks but Surefire names
+               both reports TEST-TestSuite.xml. Keep their reports in fork-specific subdirectories
+               so one suite cannot overwrite the other. -->
+          <reportsDirectory>${project.build.directory}/surefire-reports/$${surefire.forkNumber}</reportsDirectory>
+          <!-- Surefire collapses repeated TestNG invocations onto one class/method key when it
+               computes the XML suite count. Use TestNG's per-class JUnit reporter for this module;
+               it reports every invocation and does not write the shared HTML assets. -->
+          <disableXmlReport>true</disableXmlReport>
+          <properties combine.children="append">
+            <property>
+              <name>reporter</name>
+              <value>org.testng.reporters.JUnitReportReporter</value>
+            </property>
+          </properties>
           <suiteXmlFiles>
             <suiteXmlFile>testng-statefull.xml</suiteXmlFile>
             <suiteXmlFile>testng-stateless.xml</suiteXmlFile>
           </suiteXmlFiles>
         </configuration>
       </plugin>
+      <plugin>
+        <groupId>org.apache.maven.plugins</groupId>
+        <artifactId>maven-surefire-report-plugin</artifactId>
+        <version>${surefire.version}</version>
+        <configuration>
+          <!-- CI uses three possible fork indices, while Controller has only two TestNG suite test
+               sets, so at most two directories are populated. List all three; the report plugin
+               ignores missing directories. -->
+          <reportsDirectories>
+            <reportsDirectory>${project.build.directory}/surefire-reports/1/junitreports</reportsDirectory>
+            <reportsDirectory>${project.build.directory}/surefire-reports/2/junitreports</reportsDirectory>
+            <reportsDirectory>${project.build.directory}/surefire-reports/3/junitreports</reportsDirectory>
+          </reportsDirectories>
+        </configuration>
+      </plugin>
     </plugins>
   </build>
 

pinot-controller/src/test/java/org/apache/pinot/controller/ControllerStarterDynamicEnvTest.java

@@ -43,6 +43,7 @@
 /// This class tests env variables when starting controller from configs
 public class ControllerStarterDynamicEnvTest extends ControllerTest {
   private final Map<String, Object> _configOverride = new HashMap<>();
+  private int _controllerPortOverride;
 
   @Override
   protected void overrideControllerConf(Map<String, Object> properties) {
@@ -53,10 +54,11 @@ protected void overrideControllerConf(Map<String, Object> properties) {
   @Test
   public void testNoVariable()
       throws Exception {
+    int controllerPort = findControllerPortInForkRange();
     _configOverride.clear();
     _configOverride.put(CONTROLLER_HOST, "myHost");
     _configOverride.put(CONFIG_OF_INSTANCE_ID, "Controller_myInstance");
-    _configOverride.put(CONTROLLER_PORT, 1234);
+    _configOverride.put(CONTROLLER_PORT, controllerPort);
 
     startZk();
     this.startController();
@@ -66,7 +68,7 @@ public void testNoVariable()
     InstanceConfig instanceConfig = HelixHelper.getInstanceConfig(_helixManager, instanceId);
     assertEquals(instanceConfig.getInstanceName(), instanceId);
     assertEquals(instanceConfig.getHostName(), "myHost");
-    assertEquals(instanceConfig.getPort(), "1234");
+    assertEquals(instanceConfig.getPort(), Integer.toString(controllerPort));
     assertEquals(instanceConfig.getTags(), Set.of(CONTROLLER_INSTANCE));
 
     stopController();
@@ -76,11 +78,12 @@ public void testNoVariable()
   @Test
   public void testOneVariable()
       throws Exception {
+    int controllerPort = findControllerPortInForkRange();
     _configOverride.clear();
     _configOverride.put("dynamic.env.config", "controller.host");
     _configOverride.put(CONTROLLER_HOST, "HOST");
     _configOverride.put(CONFIG_OF_INSTANCE_ID, "Controller_myInstance");
-    _configOverride.put(CONTROLLER_PORT, 1234);
+    _configOverride.put(CONTROLLER_PORT, controllerPort);
 
     startZk();
     this.startController();
@@ -90,7 +93,7 @@ public void testOneVariable()
     InstanceConfig instanceConfig = HelixHelper.getInstanceConfig(_helixManager, instanceId);
     assertEquals(instanceConfig.getInstanceName(), instanceId);
     assertEquals(instanceConfig.getHostName(), "myHost");
-    assertEquals(instanceConfig.getPort(), "1234");
+    assertEquals(instanceConfig.getPort(), Integer.toString(controllerPort));
     assertEquals(instanceConfig.getTags(), Set.of(CONTROLLER_INSTANCE));
 
     stopController();
@@ -100,6 +103,7 @@ public void testOneVariable()
   @Test
   public void testMultipleVariables()
       throws Exception {
+    int controllerPort = findControllerPortInForkRange();
     _configOverride.clear();
     _configOverride.put("dynamic.env.config", "controller.host,controller.port");
     _configOverride.put(CONTROLLER_HOST, "HOST");
@@ -114,7 +118,7 @@ public void testMultipleVariables()
     InstanceConfig instanceConfig = HelixHelper.getInstanceConfig(_helixManager, instanceId);
     assertEquals(instanceConfig.getInstanceName(), instanceId);
     assertEquals(instanceConfig.getHostName(), "myHost");
-    assertEquals(instanceConfig.getPort(), "1234");
+    assertEquals(instanceConfig.getPort(), Integer.toString(controllerPort));
     assertEquals(instanceConfig.getTags(), Set.of(CONTROLLER_INSTANCE));
 
     stopController();
@@ -150,7 +154,7 @@ public void startController(Map<String, Object> properties)
       throws Exception {
     Map<String, String> envVariables = new HashMap<>();
     envVariables.put("HOST", "myHost");
-    envVariables.put("PORT", "1234");
+    envVariables.put("PORT", Integer.toString(_controllerPortOverride));
     _controllerStarter = createControllerStarter();
     _controllerStarter.init(new PinotConfiguration(properties, envVariables));
     _controllerStarter.start();
@@ -184,4 +188,9 @@ public void startController(Map<String, Object> properties)
     }
     assertEquals(System.getProperty("user.timezone"), "UTC");
   }
+
+  private int findControllerPortInForkRange() {
+    _controllerPortOverride = NetUtils.findOpenPort(_nextControllerPort);
+    return _controllerPortOverride;
+  }
 }

pinot-controller/src/test/java/org/apache/pinot/controller/ControllerStarterStatelessTest.java

@@ -24,6 +24,7 @@
 import org.apache.helix.model.InstanceConfig;
 import org.apache.pinot.common.utils.helix.HelixHelper;
 import org.apache.pinot.controller.helix.ControllerTest;
+import org.apache.pinot.spi.utils.NetUtils;
 import org.testng.annotations.Test;
 
 import static org.apache.pinot.controller.ControllerConf.CONTROLLER_HOST;
@@ -46,10 +47,11 @@ protected void overrideControllerConf(Map<String, Object> properties) {
   @Test
   public void testHostnamePortOverride()
       throws Exception {
+    int controllerPort = NetUtils.findOpenPort(_nextControllerPort);
     _configOverride.clear();
     _configOverride.put(CONFIG_OF_INSTANCE_ID, "Controller_myInstance");
     _configOverride.put(CONTROLLER_HOST, "myHost");
-    _configOverride.put(CONTROLLER_PORT, 1234);
+    _configOverride.put(CONTROLLER_PORT, controllerPort);
 
     startZk();
     startController();
@@ -59,7 +61,7 @@ public void testHostnamePortOverride()
     InstanceConfig instanceConfig = HelixHelper.getInstanceConfig(_helixManager, instanceId);
     assertEquals(instanceConfig.getInstanceName(), instanceId);
     assertEquals(instanceConfig.getHostName(), "myHost");
-    assertEquals(instanceConfig.getPort(), "1234");
+    assertEquals(instanceConfig.getPort(), Integer.toString(controllerPort));
     assertEquals(instanceConfig.getTags(), Set.of(CONTROLLER_INSTANCE));
 
     stopController();
@@ -69,10 +71,11 @@ public void testHostnamePortOverride()
   @Test
   public void testInvalidInstanceId()
       throws Exception {
+    int controllerPort = NetUtils.findOpenPort(_nextControllerPort);
     _configOverride.clear();
     _configOverride.put(CONFIG_OF_INSTANCE_ID, "myInstance");
     _configOverride.put(CONTROLLER_HOST, "myHost");
-    _configOverride.put(CONTROLLER_PORT, 1234);
+    _configOverride.put(CONTROLLER_PORT, controllerPort);
 
     startZk();
     try {
@@ -88,19 +91,20 @@ public void testInvalidInstanceId()
   @Test
   public void testDefaultInstanceId()
       throws Exception {
+    int controllerPort = NetUtils.findOpenPort(_nextControllerPort);
     _configOverride.clear();
     _configOverride.put(CONTROLLER_HOST, "myHost");
-    _configOverride.put(CONTROLLER_PORT, 1234);
+    _configOverride.put(CONTROLLER_PORT, controllerPort);
 
     startZk();
     startController();
 
     String instanceId = _controllerStarter.getInstanceId();
-    assertEquals(instanceId, "Controller_myHost_1234");
+    assertEquals(instanceId, "Controller_myHost_" + controllerPort);
     InstanceConfig instanceConfig = HelixHelper.getInstanceConfig(_helixManager, instanceId);
     assertEquals(instanceConfig.getInstanceName(), instanceId);
     assertEquals(instanceConfig.getHostName(), "myHost");
-    assertEquals(instanceConfig.getPort(), "1234");
+    assertEquals(instanceConfig.getPort(), Integer.toString(controllerPort));
     assertEquals(instanceConfig.getTags(), Set.of(CONTROLLER_INSTANCE));
 
     stopController();