try to fix d2 warmup unit test issue

CHANGELOG.md

@@ -14,6 +14,9 @@ and what APIs have changed, if applicable.
 
 ## [Unreleased]
 
+## [29.58.4] - 2024-08-12
+- Fix warmup unit test flaky unit test
+
 ## [29.58.3] - 2024-08-12
 - Disable the warmUp flaky unit test
 
@@ -5719,7 +5722,8 @@ patch operations can re-use these classes for generating patch messages.
 
 ## [0.14.1]
 
-[Unreleased]: https://github.com/linkedin/rest.li/compare/v29.58.3...master
+[Unreleased]: https://github.com/linkedin/rest.li/compare/v29.58.4...master
+[29.58.4]: https://github.com/linkedin/rest.li/compare/v29.58.3...v29.58.4
 [29.58.3]: https://github.com/linkedin/rest.li/compare/v29.58.2...v29.58.3
 [29.58.2]: https://github.com/linkedin/rest.li/compare/v29.58.1...v29.58.2
 [29.58.1]: https://github.com/linkedin/rest.li/compare/v29.58.0...v29.58.1

d2/src/main/java/com/linkedin/d2/balancer/util/WarmUpLoadBalancer.java

@@ -397,18 +397,47 @@ private static boolean isModeToWarmUp(DualReadModeProvider.DualReadMode mode, bo
   @Override
   public void shutdown(PropertyEventThread.PropertyEventShutdownCallback shutdown)
   {
-    // avoid cleaning when you risk to have partial results since some of the services have not loaded yet
+    // Indicate that shutdown has started
+    _shuttingDown = true;
+
+    // Cancel all outstanding requests
+    _outstandingRequests.forEach(future -> future.cancel(true));
+    _outstandingRequests.clear();
+
+    // Shut down the executor service and wait for it to terminate
+    _executorService.shutdown();
+    try
+    {
+      // Wait for termination with a timeout
+      if (!_executorService.awaitTermination(_warmUpTimeoutMillis, TimeUnit.MILLISECONDS))
+      {
+        // Force shutdown if termination takes too long
+        _executorService.shutdownNow();
+        // Wait again to ensure termination
+        if (!_executorService.awaitTermination(_warmUpTimeoutMillis, TimeUnit.MILLISECONDS))
+        {
+          // Log a warning if the executor service did not terminate
+          LOG.warn("Executor service did not terminate in a timely manner.");
+        }
+      }
+    }
+    catch (InterruptedException e)
+    {
+      // If interrupted, force shutdown and restore interrupt status
+      _executorService.shutdownNow();
+      Thread.currentThread().interrupt();
+      LOG.error("Interrupted while waiting for executor service to terminate.", e);
+    }
+
+    // Proceed with cleanup only if all outstanding requests are completed
     if (completedOutStandingRequests())
     {
-      // cleanup from unused services
       FileSystemDirectory fsDirectory = new FileSystemDirectory(_d2FsDirPath, _d2ServicePath);
       fsDirectory.removeAllServicesWithExcluded(_usedServices);
       fsDirectory.removeAllClustersWithExcluded(getUsedClusters());
     }
 
-    _shuttingDown = true;
-    _outstandingRequests.forEach(future -> future.cancel(true));
-    _outstandingRequests.clear();
+    // Finalize shutdown of the load balancer
     _loadBalancer.shutdown(shutdown);
   }
 

d2/src/test/java/com/linkedin/d2/balancer/util/WarmUpLoadBalancerTest.java

@@ -128,7 +128,7 @@ public void testMakingWarmUpRequests() throws URISyntaxException, InterruptedExc
     Assert.assertEquals(VALID_FILES.size(), requestCount.get());
   }
 
-  @Ignore("ingore this flaky test")
+  @Test(timeOut = 10000)
   public void testDeletingFilesAfterShutdown() throws InterruptedException, ExecutionException, TimeoutException
   {
     createDefaultServicesIniFiles();

gradle.properties

@@ -1,4 +1,4 @@
-version=29.58.3
+version=29.58.4
 group=com.linkedin.pegasus
 org.gradle.configureondemand=true
 org.gradle.parallel=true