Resiliency: mark a service as used even if the getClient call during warmUp fails (#1040)

* Mark service as _usedService regardless of whether getClient call succeeds

* Update changelog

* Remove empty test (added mistakenly)

Author: ZoabKapoor

CHANGELOG.md

@@ -14,6 +14,9 @@ and what APIs have changed, if applicable.
 
 ## [Unreleased]
 
+## [29.64.1] - 2025-02-15
+- Fix warmUp -- record service as used regardless of whether getClient succeeds
+
 ## [29.64.0] - 2025-01-31
 - Allow subscribing to a single D2URI
 
@@ -5770,7 +5773,8 @@ patch operations can re-use these classes for generating patch messages.
 
 ## [0.14.1]
 
-[Unreleased]: https://github.com/linkedin/rest.li/compare/v29.64.0...master
+[Unreleased]: https://github.com/linkedin/rest.li/compare/v29.64.1...master
+[29.64.1]: https://github.com/linkedin/rest.li/compare/v29.64.0...v29.64.1
 [29.64.0]: https://github.com/linkedin/rest.li/compare/v29.63.2...v29.64.0
 [29.63.2]: https://github.com/linkedin/rest.li/compare/v29.63.1...v29.63.2
 [29.63.1]: https://github.com/linkedin/rest.li/compare/v29.63.0...v29.63.1

d2-test-api/src/main/java/com/linkedin/d2/balancer/util/TestLoadBalancer.java

@@ -36,6 +36,7 @@
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicInteger;
 
 
@@ -49,12 +50,17 @@ public class TestLoadBalancer implements LoadBalancerWithFacilities, WarmUpServi
   private final AtomicInteger _completedRequestCount = new AtomicInteger();
   private int _warmUpDelayMs = 0;
   private int _serviceDataDelayMs = 0;
+  private boolean _shouldThrowOnGetClient = false;
 
   private final int DELAY_STANDARD_DEVIATION = 5; //ms
   private final ScheduledExecutorService _executorService = Executors.newSingleThreadScheduledExecutor();
 
   public TestLoadBalancer() {}
 
+  public TestLoadBalancer(boolean shouldThrowOnGetClient) {
+    _shouldThrowOnGetClient = shouldThrowOnGetClient;
+  }
+
   public TestLoadBalancer(int warmUpDelayMs)
   {
     this(warmUpDelayMs, 0);
@@ -69,7 +75,12 @@ public TestLoadBalancer(int warmUpDelayMs, int serviceDataDelayMs)
   @Override
   public void getClient(Request request, RequestContext requestContext, Callback<TransportClient> clientCallback)
   {
-    clientCallback.onSuccess(new TestClient());
+    if (_shouldThrowOnGetClient)
+    {
+      clientCallback.onError(new TimeoutException());
+    } else {
+      clientCallback.onSuccess(new TestClient());
+    }
   }
 
   @Override

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

@@ -441,10 +441,10 @@ private Set<String> getUsedClusters()
   @Override
   public TransportClient getClient(Request request, RequestContext requestContext) throws ServiceUnavailableException
   {
-    TransportClient client = _loadBalancer.getClient(request, requestContext);
-
+    // Add serviceName to _usedServices *before* making the call to _loadBalancer.getClient. Even if
+    // the call fails, we still *intend* to use serviceName, so it should be in _usedServices.
     String serviceName = LoadBalancerUtil.getServiceNameFromUri(request.getURI());
     _usedServices.add(serviceName);
-    return client;
+    return _loadBalancer.getClient(request, requestContext);
   }
 }

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

@@ -165,16 +165,24 @@ public void testDeletingFilesAfterShutdown() throws InterruptedException, Execut
     }
   }
 
+  @DataProvider(name = "shouldThrowOnGetClientDataProvider")
+  public Object[][] shouldThrowOnGetClientDataProvider()
+  {
+    return new Object[][]
+        {{true}, {false}};
+  }
+
   /**
    * Since the list might from the fetcher might not be complete (update service, old data, etc.., and the user might
    * require additional services at runtime, we have to check that those services are not cleared from the cache
-   * otherwise it would incur in a penalty at the next deployment
+   * otherwise it would incur in a penalty at the next deployment.
+   * Note that regardless of whether getClient returns successfully or not (if it times out, for example),
+   * we should still record the service/s we tried to warm up.
    */
-  @Test(timeOut = 10000, retryAnalyzer = ThreeRetries.class)
-  public void testNotDeletingFilesGetClient() throws InterruptedException, ExecutionException, TimeoutException, ServiceUnavailableException
-  {
+  @Test(dataProvider = "shouldThrowOnGetClientDataProvider", timeOut = 10000, retryAnalyzer = ThreeRetries.class)
+  public void testNotDeletingFilesGetClient(boolean shouldThrowOnGetClient) throws InterruptedException, ExecutionException, TimeoutException {
     createDefaultServicesIniFiles();
-    TestLoadBalancer balancer = new TestLoadBalancer();
+    TestLoadBalancer balancer = new TestLoadBalancer(shouldThrowOnGetClient);
 
     List<String> allServicesBeforeShutdown = getAllDownstreamServices();
     DownstreamServicesFetcher returnNoDownstreams = callback -> callback.onSuccess(Collections.emptyList());
@@ -190,14 +198,19 @@ public void testNotDeletingFilesGetClient() throws InterruptedException, Executi
     warmUpLoadBalancer.start(callback);
     callback.get(5000, TimeUnit.MILLISECONDS);
 
-    warmUpLoadBalancer.getClient(new URIRequest("d2://" + pickOneService), new RequestContext());
+    try {
+      warmUpLoadBalancer.getClient(new URIRequest("d2://" + pickOneService), new RequestContext());
+    } catch (Exception e) {
+      Assert.assertTrue(shouldThrowOnGetClient);
+    }
 
     FutureCallback<None> shutdownCallback = new FutureCallback<>();
     warmUpLoadBalancer.shutdown(() -> shutdownCallback.onSuccess(None.none()));
     shutdownCallback.get(5000, TimeUnit.MILLISECONDS);
 
     List<String> allServicesAfterShutdown = getAllDownstreamServices();
 
+    // regardless of whether getClient returned successfully or threw an Exception, we should still record the service we tried to warm up
     Assert.assertEquals(1, allServicesAfterShutdown.size(), "After shutdown there should be just one service, the one that we 'get the client' on");
   }
 

gradle.properties

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