chore(ci): update gradle config, test infrastructure, and testng suit… (#15207)

Author: david-leifker

gradle.properties

@@ -5,9 +5,23 @@ org.gradle.caching=true
 # Cycle daemons after 30m
 org.gradle.daemon.idletimeout=1800000
 
-# Increase gradle JVM memory to 5GB to allow tests to run locally
-org.gradle.jvmargs=-Xmx5120m -XX:MaxMetaspaceSize=512m
-org.gradle.workers.max=4
+# Gradle JVM memory configuration
+# Default: 2GB heap (optimized for 8GB systems)
+# Override options:
+#   1. Memory Profile (recommended): Set GRADLE_MEMORY_PROFILE env var or -Dgradle.memory.profile
+#      - "low": 2GB heap, 384MB metaspace (8GB systems)
+#      - "default": 5GB heap, 512MB metaspace (16GB+ systems)
+#      - "high": 8GB heap, 768MB metaspace (32GB+ systems)
+#   2. Direct override for org.gradle.jvmargs (requires stopping daemon first):
+#      ./gradlew --stop
+#      ./gradlew -Dorg.gradle.jvmargs="-Xmx5g -XX:MaxMetaspaceSize=512m" build
+#   3. Via GRADLE_OPTS environment variable:
+#      export GRADLE_OPTS="-Xmx5g -XX:MaxMetaspaceSize=512m"
+#   4. Via ~/.gradle/gradle.properties (user-specific, takes precedence)
+# Note: org.gradle.jvmargs affects the Gradle daemon JVM, not test processes
+# Test process memory is controlled by build.gradle based on memory profile
+org.gradle.jvmargs=-Xmx2g -XX:MaxMetaspaceSize=384m
+org.gradle.workers.max=2
 
 # Increase retries to 5 (from default of 3) and increase interval from 125ms to 1s.
 # Based on this thread https://github.com/gradle/gradle/issues/4629, it's unclear

metadata-io/build.gradle

@@ -7,6 +7,98 @@ plugins {
 
 apply from: '../gradle/coverage/java-coverage.gradle'
 
+// Memory profile configuration - supports environment variable GRADLE_MEMORY_PROFILE or system property
+// Values: "low" (8GB systems), "default" (16GB+ systems), "high" (32GB+ systems)
+// Auto-detects system memory if not explicitly set
+def detectMemoryProfile() {
+  // Check if explicitly set via environment variable or system property
+  def explicitProfile = System.getenv("GRADLE_MEMORY_PROFILE") ?: System.getProperty("gradle.memory.profile")
+  if (explicitProfile) {
+    return explicitProfile
+  }
+
+  // Auto-detect based on actual system memory (total RAM) using OperatingSystemMXBean
+  try {
+    def osBean = java.lang.management.ManagementFactory.getOperatingSystemMXBean()
+    def sunOsBean = Class.forName("com.sun.management.OperatingSystemMXBean")
+    if (!sunOsBean.isInstance(osBean)) {
+      throw new Exception("OperatingSystemMXBean does not support getTotalPhysicalMemorySize()")
+    }
+
+    def method = sunOsBean.getMethod("getTotalPhysicalMemorySize")
+    def totalMemoryBytes = method.invoke(osBean) as Long
+    def totalMemoryGB = totalMemoryBytes / (1024.0 * 1024.0 * 1024.0)
+
+    def detectedProfile
+    if (totalMemoryGB < 12) {
+      detectedProfile = "low"  // < 12GB: use low profile (8GB systems)
+    } else if (totalMemoryGB < 24) {
+      detectedProfile = "default"  // 12-24GB: use default profile (16GB+ systems)
+    } else {
+      detectedProfile = "high"  // >= 24GB: use high profile (32GB+ systems)
+    }
+
+    println "Auto-detected memory profile: ${detectedProfile} (System RAM: ${String.format("%.1f", totalMemoryGB)}GB)"
+    return detectedProfile
+  } catch (Exception e) {
+    // Fallback to default if detection fails
+    println "Warning: Could not detect system memory (${e.message}), using 'default' profile"
+    return "default"
+  }
+}
+
+def memoryProfile = detectMemoryProfile()
+
+// Define memory configurations per profile
+// since it's read before build.gradle runs
+def memoryConfigs = [
+  "low": [
+    testHeap: "1g",
+    testMinHeap: "256m",
+    testMetaspace: "384m",
+    // Docker container memory settings (for ES/OS/Cassandra/Neo4j) - applied via system properties
+    elasticsearchHeap: "256m",
+    opensearchHeap: "512m",
+    cassandraHeap: "128m",
+    neo4jHeap: "256m"
+  ],
+  "default": [
+    testHeap: "2g",
+    testMinHeap: "512m",
+    testMetaspace: "512m",
+    elasticsearchHeap: "512m",
+    opensearchHeap: "1024m",
+    cassandraHeap: "128m",
+    neo4jHeap: "512m"
+  ],
+  "high": [
+    testHeap: "3g",
+    testMinHeap: "1g",
+    testMetaspace: "768m",
+    elasticsearchHeap: "768m",
+    opensearchHeap: "1536m",
+    cassandraHeap: "256m",
+    neo4jHeap: "768m"
+  ]
+]
+
+def memoryConfig = memoryConfigs[memoryProfile] ?: memoryConfigs["low"]
+println "Using memory profile: ${memoryProfile}"
+println "  Test processes: ${memoryConfig.testHeap} heap (min ${memoryConfig.testMinHeap}), ${memoryConfig.testMetaspace} metaspace"
+println "  Docker containers:"
+println "    - Elasticsearch: ${memoryConfig.elasticsearchHeap}"
+println "    - OpenSearch: ${memoryConfig.opensearchHeap}"
+println "    - Cassandra: ${memoryConfig.cassandraHeap}"
+println "    - Neo4j: ${memoryConfig.neo4jHeap}"
+println "  Note: Gradle daemon memory is set in gradle.properties (see file for details)"
+println "  Override: Set GRADLE_MEMORY_PROFILE environment variable or -Dgradle.memory.profile system property"
+
+// Set system properties for container memory configuration (containers read these at runtime)
+System.setProperty("testcontainers.elasticsearch.heap", memoryConfig.elasticsearchHeap)
+System.setProperty("testcontainers.opensearch.heap", memoryConfig.opensearchHeap)
+System.setProperty("testcontainers.cassandra.heap", memoryConfig.cassandraHeap)
+System.setProperty("testcontainers.neo4j.heap", memoryConfig.neo4jHeap)
+
 configurations {
   enhance
 }
@@ -151,12 +243,84 @@ test {
   environment 'STRICT_URN_VALIDATION_ENABLED', 'true'
 }
 
+task testEs7CassandraNeo4j(type: Test) {
+  doFirst {
+    maxParallelForks = 1
+  }
+  useTestNG() {
+    suites 'src/test/resources/testng-es7-cassandra-neo4j.xml'
+  }
+  testLogging {
+    events "failed", "skipped", "passed"
+    exceptionFormat "full"
+    showStandardStreams = true
+  }
+
+  environment 'STRICT_URN_VALIDATION_ENABLED', 'true'
+  
+  group = 'verification'
+  description = 'Runs only the testng-es7-cassandra-neo4j.xml test suite'
+}
+
+task testOpenSearch(type: Test) {
+  doFirst {
+    maxParallelForks = 1
+  }
+  useTestNG() {
+    suites 'src/test/resources/testng-opensearch.xml'
+  }
+  testLogging {
+    events "failed", "skipped", "passed"
+    exceptionFormat "full"
+    showStandardStreams = true
+  }
+
+  environment 'STRICT_URN_VALIDATION_ENABLED', 'true'
+  
+  group = 'verification'
+  description = 'Runs only the testng-opensearch.xml test suite'
+}
+
+task testOther(type: Test) {
+  doFirst {
+    maxParallelForks = 1
+  }
+  useTestNG() {
+    suites 'src/test/resources/testng-other.xml'
+  }
+  testLogging {
+    events "failed", "skipped", "passed", "started"
+    exceptionFormat "full"
+    showStandardStreams = true
+  }
+
+  environment 'STRICT_URN_VALIDATION_ENABLED', 'true'
+  
+  group = 'verification'
+  description = 'Runs only other test suite to debug the failing tests'
+}
+
 ebean {
   debugLevel = 1 // 0 - 9
 }
 
 tasks.withType(Test) {
   enableAssertions = false
+  // Set heap size for test JVM processes based on memory profile
+  // This is separate from gradle.properties which only sets memory for the Gradle daemon
+  maxHeapSize = memoryConfig.testHeap
+  minHeapSize = memoryConfig.testMinHeap
+
+  // Add JVM args for memory efficiency and to prevent OOM kills
+  // -XX:MaxMetaspaceSize: Limit metaspace to prevent unbounded growth
+  // -XX:+UseG1GC: Use G1 garbage collector for better memory efficiency
+  // -XX:MaxGCPauseMillis: Target pause time for GC (higher for low-memory systems)
+  def gcPauseMillis = memoryProfile == "low" ? "1000" : "500"
+  jvmArgs = [
+          "-XX:MaxMetaspaceSize=${memoryConfig.testMetaspace}",
+          "-XX:+UseG1GC",
+          "-XX:MaxGCPauseMillis=${gcPauseMillis}"
+  ]
 }
 
 clean {

metadata-io/src/test/java/com/linkedin/metadata/CassandraTestUtils.java

@@ -17,6 +17,7 @@
 import java.util.Map;
 import java.util.stream.Collectors;
 import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
 import javax.net.ssl.SSLContext;
 import lombok.extern.slf4j.Slf4j;
 import org.testcontainers.cassandra.CassandraContainer;
@@ -33,9 +34,18 @@ private CassandraTestUtils() {}
   public static CassandraContainer setupContainer() {
     final DockerImageName imageName = DockerImageName.parse(IMAGE_NAME);
 
+    // Read heap size from system property, default to 128m for backwards compatibility
+    // System property is set by build.gradle based on GRADLE_MEMORY_PROFILE
+    String cassandraHeap = System.getProperty("testcontainers.cassandra.heap", "128m");
     CassandraContainer container =
         new CassandraContainer(imageName)
-            .withEnv("JVM_OPTS", "-Xms128M -Xmx128M -Djava.net.preferIPv4Stack=true")
+            .withEnv(
+                "JVM_OPTS",
+                "-Xms"
+                    + cassandraHeap
+                    + " -Xmx"
+                    + cassandraHeap
+                    + " -Djava.net.preferIPv4Stack=true")
             .withStartupTimeout(Duration.ofMinutes(2))
             .withReuse(true);
 
@@ -185,4 +195,21 @@ public static void purgeData(CassandraContainer container) {
       assertEquals(rs.size(), 0);
     }
   }
+
+  /**
+   * Close a CqlSession to prevent Netty thread leaks (e.g., "s15-admin-0" threads). Each CqlSession
+   * creates Netty event loop threads that must be closed.
+   *
+   * @param session The CqlSession instance to close, or null if none
+   */
+  public static void closeSession(@Nullable CqlSession session) {
+    if (session != null && !session.isClosed()) {
+      try {
+        session.close();
+      } catch (Exception e) {
+        // Log but don't fail - this is called from test cleanup
+        System.err.println("Failed to close CqlSession: " + e.getMessage());
+      }
+    }
+  }
 }

metadata-io/src/test/java/com/linkedin/metadata/aspect/utils/DefaultAspectsUtilTest.java

@@ -23,16 +23,19 @@
 import java.util.List;
 import java.util.stream.Collectors;
 import org.testng.Assert;
+import org.testng.annotations.AfterMethod;
 import org.testng.annotations.Test;
 
 public class DefaultAspectsUtilTest {
 
+  private Database server;
+
   public DefaultAspectsUtilTest() {}
 
   @Test
   public void testAdditionalChanges() {
     OperationContext opContext = TestOperationContexts.systemContextNoSearchAuthorization();
-    Database server = EbeanTestUtils.createTestServer(DefaultAspectsUtilTest.class.getSimpleName());
+    server = EbeanTestUtils.createTestServer(DefaultAspectsUtilTest.class.getSimpleName());
     EbeanAspectDao aspectDao = new EbeanAspectDao(server, EbeanConfiguration.testDefault, null);
     aspectDao.setConnectionValidated(true);
     EventProducer mockProducer = mock(EventProducer.class);
@@ -75,4 +78,9 @@ public void testAdditionalChanges() {
             .collect(Collectors.toList()),
         List.of(ChangeType.CREATE, ChangeType.CREATE, ChangeType.CREATE, ChangeType.CREATE));
   }
+
+  @AfterMethod
+  public void cleanup() {
+    EbeanTestUtils.shutdownDatabase(server);
+  }
 }

metadata-io/src/test/java/com/linkedin/metadata/elasticsearch/update/ESBulkProcessorTest.java

@@ -27,6 +27,7 @@
 import org.opensearch.index.reindex.UpdateByQueryRequest;
 import org.opensearch.script.Script;
 import org.opensearch.script.ScriptType;
+import org.testng.annotations.AfterMethod;
 import org.testng.annotations.BeforeMethod;
 import org.testng.annotations.Test;
 
@@ -42,8 +43,27 @@ public class ESBulkProcessorTest {
   private String TEST_TASK_STRING = "nodeid123:1234";
 
   @BeforeMethod
-  public void setup() {
+  public void setup() throws IOException {
     mocks = MockitoAnnotations.openMocks(this);
+    // Stub generateBulkProcessor and generateAsyncBulkProcessor to do nothing
+    // This prevents any real BulkProcessor instances from being created
+    doNothing()
+        .when(mockSearchClient)
+        .generateBulkProcessor(any(), any(), anyInt(), anyLong(), anyLong(), anyInt(), anyInt());
+    doNothing()
+        .when(mockSearchClient)
+        .generateAsyncBulkProcessor(
+            any(), any(), anyInt(), anyLong(), anyLong(), anyInt(), anyInt());
+    // Ensure closeBulkProcessor can be called without issues
+    doNothing().when(mockSearchClient).closeBulkProcessor();
+    doNothing().when(mockSearchClient).flushBulkProcessor();
+  }
+
+  @AfterMethod
+  public void tearDown() throws Exception {
+    if (mocks != null) {
+      mocks.close();
+    }
   }
 
   @Test

metadata-io/src/test/java/com/linkedin/metadata/entity/EbeanAspectMigrationsDaoTest.java

@@ -21,17 +21,19 @@
 import java.util.Set;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
+import org.testng.annotations.AfterMethod;
 import org.testng.annotations.BeforeMethod;
 import org.testng.annotations.Test;
 
 public class EbeanAspectMigrationsDaoTest extends AspectMigrationsDaoTest<EbeanAspectDao> {
 
+  private Database server;
+
   public EbeanAspectMigrationsDaoTest() {}
 
   @BeforeMethod
   public void setupTest() {
-    Database server =
-        EbeanTestUtils.createTestServer(EbeanAspectMigrationsDaoTest.class.getSimpleName());
+    server = EbeanTestUtils.createTestServer(EbeanAspectMigrationsDaoTest.class.getSimpleName());
     _mockProducer = mock(EventProducer.class);
     EbeanAspectDao dao = new EbeanAspectDao(server, EbeanConfiguration.testDefault, null);
     dao.setConnectionValidated(true);
@@ -67,4 +69,11 @@ public void testStreamAspects() throws AssertionError {
       assertTrue(urnsFetched.contains(urn));
     }
   }
+
+  @AfterMethod
+  public void cleanup() {
+    // Shutdown Database instance to prevent thread pool and connection leaks
+    // This includes the "gma.heartBeat" thread and connection pools
+    EbeanTestUtils.shutdownDatabase(server);
+  }
 }

metadata-io/src/test/java/com/linkedin/metadata/entity/EbeanEntityServiceOptimizationTest.java

@@ -28,6 +28,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.stream.Collectors;
+import org.testng.annotations.AfterMethod;
 import org.testng.annotations.BeforeMethod;
 import org.testng.annotations.Test;
 
@@ -61,11 +62,13 @@ public class EbeanEntityServiceOptimizationTest {
       TestOperationContexts.systemContextNoSearchAuthorization();
 
   private EntityServiceImpl entityService;
+  private Database server;
 
   @BeforeMethod
   public void setupTest() {
-    Database server =
+    server =
         EbeanTestUtils.createTestServer(EbeanEntityServiceOptimizationTest.class.getSimpleName());
+
     AspectDao aspectDao = new EbeanAspectDao(server, EbeanConfiguration.testDefault, null);
     PreProcessHooks preProcessHooks = new PreProcessHooks();
     preProcessHooks.setUiEnabled(true);
@@ -316,4 +319,11 @@ private static String formatUnknownStatements(List<String> statements) {
     }
     return builder.toString();
   }
+
+  @AfterMethod
+  public void cleanup() {
+    // Shutdown Database instance to prevent thread pool and connection leaks
+    // This includes the "gma.heartBeat" thread and connection pools
+    EbeanTestUtils.shutdownDatabase(server);
+  }
 }