[router] Fix in-flight request check in router shutdown

internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/integration/utils/VeniceRouterWrapper.java

@@ -70,7 +70,7 @@ public class VeniceRouterWrapper extends ProcessWrapper implements MetricsAware
   private static final String ROUTER_SERVICE_METRIC_PREFIX = "router";
   private final VeniceProperties properties;
   private final String zkAddress;
-  private RouterServer service;
+  private RouterServer routerServer;
   private final String d2ClusterName;
   private final String clusterDiscoveryD2ClusterName;
   private final String regionName;
@@ -81,13 +81,13 @@ public class VeniceRouterWrapper extends ProcessWrapper implements MetricsAware
       String regionName,
       String serviceName,
       File dataDirectory,
-      RouterServer service,
+      RouterServer routerServer,
       VeniceProperties properties,
       String zkAddress,
       String d2ClusterName,
       String clusterDiscoveryD2ClusterName) {
     super(serviceName, dataDirectory);
-    this.service = service;
+    this.routerServer = routerServer;
     this.properties = properties;
     this.zkAddress = zkAddress;
     this.d2ClusterName = d2ClusterName;
@@ -224,22 +224,22 @@ public int getSslPort() {
   }
 
   public String getD2ServiceNameForCluster(String clusterName) {
-    return service.getConfig().getClusterToD2Map().get(clusterName);
+    return routerServer.getConfig().getClusterToD2Map().get(clusterName);
   }
 
   @Override
   protected void internalStart() throws Exception {
-    service.start();
+    routerServer.start();
     TestUtils.waitForNonDeterministicCompletion(
         IntegrationTestUtils.MAX_ASYNC_START_WAIT_TIME_MS,
         TimeUnit.MILLISECONDS,
-        () -> service.isRunning());
+        () -> routerServer.isRunning());
     LOGGER.info("Started VeniceRouterWrapper: {}", this);
   }
 
   @Override
   protected void internalStop() throws Exception {
-    service.stop();
+    routerServer.stop();
   }
 
   @Override
@@ -251,7 +251,7 @@ protected void newProcess() {
 
     d2Servers.addAll(D2TestUtils.getD2Servers(zkAddress, clusterDiscoveryD2ClusterName, httpURI, httpsURI));
 
-    service = new RouterServer(
+    routerServer = new RouterServer(
         properties,
         d2Servers,
         Optional.empty(),
@@ -272,28 +272,32 @@ public String getComponentTagForLogging() {
   }
 
   public HelixBaseRoutingRepository getRoutingDataRepository() {
-    return service.getRoutingDataRepository();
+    return routerServer.getRoutingDataRepository();
   }
 
   public ReadOnlyStoreRepository getMetaDataRepository() {
-    return service.getMetadataRepository();
+    return routerServer.getMetadataRepository();
   }
 
   public ReadOnlySchemaRepository getSchemaRepository() {
-    return service.getSchemaRepository();
+    return routerServer.getSchemaRepository();
   }
 
   public ZkRoutersClusterManager getRoutersClusterManager() {
-    return service.getRoutersClusterManager();
+    return routerServer.getRoutersClusterManager();
   }
 
   @Override
   public MetricsRepository getMetricsRepository() {
-    return service.getMetricsRepository();
+    return routerServer.getMetricsRepository();
+  }
+
+  public RouterServer getRouter() {
+    return routerServer;
   }
 
   public void refresh() {
-    service.refresh();
+    routerServer.refresh();
   }
 
   @Override

internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/router/TestRead.java

@@ -389,6 +389,7 @@ public void testRead() throws Exception {
               "There should be some idle connections since test queries are finished");
         }
 
+        Assert.assertTrue(getRouterMetricValue("total_inflight_request_count") > 0.0);
         Assert.assertTrue(getAggregateRouterMetricValue(".localhost--response_waiting_time.50thPercentile") > 0);
         Assert.assertTrue(
             getAggregateRouterMetricValue(".localhost--multiget_streaming_response_waiting_time.50thPercentile") > 0);
@@ -553,6 +554,9 @@ public void testRead() throws Exception {
         try {
           queriesSent++;
           storeClient.batchGet(keySet).get();
+          for (VeniceRouterWrapper routerWrapper: veniceCluster.getVeniceRouters()) {
+            Assert.assertTrue(routerWrapper.getRouter().hasInFlightRequest());
+          }
         } catch (ExecutionException e) {
           Throwable cause = e.getCause();
           Assert.assertTrue(
@@ -592,6 +596,12 @@ public void testRead() throws Exception {
           "The throttled_request metric is inconsistent with the number of quota exceptions received by the client!");
 
       getAggregateRouterMetricValue(".total--multiget_throttled_request_latency.Max");
+      for (VeniceRouterWrapper routerWrapper: veniceCluster.getVeniceRouters()) {
+        routerWrapper.getRouter().stop();
+      }
+      for (VeniceRouterWrapper routerWrapper: veniceCluster.getVeniceRouters()) {
+        Assert.assertFalse(routerWrapper.getRouter().hasInFlightRequest());
+      }
       /** TODO Re-enable this assertion once we stop throwing batch get quota exceptions from {@link com.linkedin.venice.router.api.VeniceDelegateMode} */
       // Assert.assertTrue(throttledRequestLatencyForBatchGetAfterQueries > 0.0, "There should be batch get throttled
       // request latency now!");
@@ -602,6 +612,10 @@ private double getMaxServerMetricValue(String metricName) {
     return MetricsUtils.getMax(metricName, veniceCluster.getVeniceServers());
   }
 
+  private double getRouterMetricValue(String metricName) {
+    return MetricsUtils.getMax(metricName, veniceCluster.getVeniceRouters());
+  }
+
   private double getMaxRouterMetricValue(String metricName) {
     return MetricsUtils.getMax(metricName, veniceCluster.getVeniceRouters());
   }

services/venice-router/src/main/java/com/linkedin/venice/router/RouterServer.java

@@ -852,10 +852,11 @@ public void stopInner() throws Exception {
      * correctly.
      */
 
+    LOGGER.info("Waiting to make sure all in-flight requests are drained");
     // Graceful shutdown: Wait till all the requests are drained
     try {
       RetryUtils.executeWithMaxAttempt(() -> {
-        if (RouterHttpRequestStats.hasInFlightRequests()) {
+        if (hasInFlightRequest()) {
           throw new VeniceException("There are still in-flight requests in router");
         }
       },
@@ -866,6 +867,7 @@ public void stopInner() throws Exception {
       LOGGER.error(
           "There are still in-flight request during router shutdown, still continuing shutdown, it might cause unhealthy request in client");
     }
+    LOGGER.info("Drained all in-flight requests, starting to shutdown the router.");
     storageNodeClient.close();
     workerEventLoopGroup.shutdownGracefully();
     serverEventLoopGroup.shutdownGracefully();
@@ -927,6 +929,10 @@ public ReadOnlySchemaRepository getSchemaRepository() {
     return schemaRepository;
   }
 
+  public boolean hasInFlightRequest() {
+    return RouterHttpRequestStats.hasInFlightRequests();
+  }
+
   private void handleExceptionInStartServices(VeniceException e, boolean async) throws VeniceException {
     if (async) {
       Utils.exit("Failed to start router services due to " + e);

services/venice-router/src/main/java/com/linkedin/venice/router/stats/RouterHttpRequestStats.java

@@ -66,6 +66,7 @@
 
 public class RouterHttpRequestStats extends AbstractVeniceHttpStats {
   private static final MetricConfig METRIC_CONFIG = new MetricConfig().timeWindow(10, TimeUnit.SECONDS);
+  private static final String TOTAL_INFLIGHT_REQUEST_COUNT = "total_inflight_request_count";
   private static final VeniceMetricsRepository localMetricRepo = new VeniceMetricsRepository(
       new VeniceMetricsConfig.Builder().setServiceName(ROUTER_SERVICE_NAME)
           .setMetricPrefix(ROUTER_SERVICE_METRIC_PREFIX)
@@ -74,7 +75,7 @@ public class RouterHttpRequestStats extends AbstractVeniceHttpStats {
 
   private final static Sensor totalInflightRequestSensor = localMetricRepo.sensor("total_inflight_request");
   static {
-    totalInflightRequestSensor.add("total_inflight_request_count", new Rate());
+    totalInflightRequestSensor.add(TOTAL_INFLIGHT_REQUEST_COUNT, new Gauge());
   }
 
   /** metrics to track incoming requests */
@@ -658,9 +659,9 @@ private Sensor registerSensorFinal(String sensorName, MeasurableStat... stats) {
   }
 
   static public boolean hasInFlightRequests() {
-    Metric metric = localMetricRepo.getMetric("total_inflight_request_count");
+    Metric metric = localMetricRepo.getMetric(TOTAL_INFLIGHT_REQUEST_COUNT);
     // max return -infinity when there are no samples. validate only against finite value
-    return Double.isFinite(metric.value()) ? metric.value() > 0.0 : false;
+    return Double.isFinite(metric.value()) && metric.value() > 0.0;
   }
 
   /** used only for testing */

services/venice-router/src/test/java/com/linkedin/venice/router/RouteHttpRequestStatsTest.java

@@ -53,5 +53,7 @@ public void routerMetricsTest() {
     Assert.assertEquals(stats.getPendingRequestCount("my_host2"), 0);
     routerHttpRequestStats.recordIncomingRequest();
     Assert.assertTrue(RouterHttpRequestStats.hasInFlightRequests());
+    routerHttpRequestStats.recordResponse();
+    Assert.assertFalse(RouterHttpRequestStats.hasInFlightRequests());
   }
 }