Return server readiness only after query serving starts

pinot-broker/src/main/java/org/apache/pinot/broker/api/resources/PinotBrokerHealthCheck.java

@@ -34,17 +34,20 @@
 import javax.ws.rs.GET;
 import javax.ws.rs.Path;
 import javax.ws.rs.Produces;
+import javax.ws.rs.QueryParam;
 import javax.ws.rs.WebApplicationException;
 import javax.ws.rs.core.HttpHeaders;
 import javax.ws.rs.core.MediaType;
 import javax.ws.rs.core.Response;
 import org.apache.pinot.broker.broker.BrokerAdminApiApplication;
+import org.apache.pinot.broker.routing.manager.BrokerRoutingManager;
 import org.apache.pinot.common.metrics.BrokerMeter;
 import org.apache.pinot.common.metrics.BrokerMetrics;
 import org.apache.pinot.common.utils.ServiceStatus;
 import org.apache.pinot.core.auth.Actions;
 import org.apache.pinot.core.auth.Authorize;
 import org.apache.pinot.core.auth.TargetType;
+import org.apache.pinot.spi.utils.CommonConstants;
 
 import static org.apache.pinot.spi.utils.CommonConstants.SWAGGER_AUTHORIZATION_KEY;
 
@@ -62,6 +65,9 @@ public class PinotBrokerHealthCheck {
   @Inject
   private BrokerMetrics _brokerMetrics;
 
+  @Inject
+  private BrokerRoutingManager _routingManager;
+
   @Inject
   @Named(BrokerAdminApiApplication.START_TIME)
   private Instant _startTime;
@@ -75,9 +81,18 @@ public class PinotBrokerHealthCheck {
       @ApiResponse(code = 200, message = "Broker is healthy"),
       @ApiResponse(code = 503, message = "Broker is not healthy")
   })
-  public String getBrokerHealth() {
+  public String getBrokerHealth(@QueryParam("serverInstance") String serverInstance) {
     ServiceStatus.Status status = ServiceStatus.getServiceStatus(_instanceId);
     if (status == ServiceStatus.Status.GOOD) {
+      if (serverInstance != null) {
+        if (!_routingManager.isServerRoutable(serverInstance)) {
+          String errMessage = String.format("Server %s is not available for routing", serverInstance);
+          throw new WebApplicationException(errMessage,
+              Response.status(Response.Status.SERVICE_UNAVAILABLE).entity(errMessage).build());
+        }
+        _brokerMetrics.addMeteredGlobalValue(BrokerMeter.HEALTHCHECK_OK_CALLS, 1);
+        return CommonConstants.Broker.SERVER_ROUTING_READY_RESPONSE;
+      }
       _brokerMetrics.addMeteredGlobalValue(BrokerMeter.HEALTHCHECK_OK_CALLS, 1);
       return "OK";
     }

pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java

@@ -1258,6 +1258,19 @@ public Set<String> getServingInstances(String tableNameWithType) {
     return routingEntry._instanceSelector.getServingInstances();
   }
 
+  /// Returns whether the server is currently available to the broker routing entries.
+  ///
+  /// The read lock ensures this method cannot observe `_routableServerInstanceMap` while an instance-config callback
+  /// is still applying the corresponding change to individual routing entries.
+  public boolean isServerRoutable(String instanceId) {
+    _globalLock.readLock().lock();
+    try {
+      return _routableServerInstanceMap.containsKey(instanceId);
+    } finally {
+      _globalLock.readLock().unlock();
+    }
+  }
+
   /// Returns the table-level query timeout in milliseconds for the given table, or `null` if the timeout is not
   /// configured in the table config.
   @Nullable

pinot-broker/src/test/java/org/apache/pinot/broker/api/resources/PinotBrokerHealthCheckTest.java

@@ -0,0 +1,82 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.broker.api.resources;
+
+import java.lang.reflect.Field;
+import javax.ws.rs.WebApplicationException;
+import org.apache.pinot.broker.routing.manager.BrokerRoutingManager;
+import org.apache.pinot.common.metrics.BrokerMetrics;
+import org.apache.pinot.common.utils.ServiceStatus;
+import org.apache.pinot.spi.utils.CommonConstants;
+import org.testng.annotations.AfterMethod;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.expectThrows;
+
+
+public class PinotBrokerHealthCheckTest {
+  private static final String BROKER_INSTANCE = "Broker_localhost_8099";
+  private static final String SERVER_INSTANCE = "Server_localhost_8098";
+
+  private BrokerRoutingManager _routingManager;
+  private PinotBrokerHealthCheck _healthCheck;
+
+  @BeforeMethod
+  public void setUp()
+      throws Exception {
+    _routingManager = mock(BrokerRoutingManager.class);
+    _healthCheck = new PinotBrokerHealthCheck();
+    setField("_instanceId", BROKER_INSTANCE);
+    setField("_brokerMetrics", mock(BrokerMetrics.class));
+    setField("_routingManager", _routingManager);
+
+    ServiceStatus.ServiceStatusCallback callback = mock(ServiceStatus.ServiceStatusCallback.class);
+    when(callback.getServiceStatus()).thenReturn(ServiceStatus.Status.GOOD);
+    ServiceStatus.setServiceStatusCallback(BROKER_INSTANCE, callback);
+  }
+
+  @AfterMethod
+  public void tearDown() {
+    ServiceStatus.removeServiceStatusCallback(BROKER_INSTANCE);
+  }
+
+  @Test
+  public void testRoutingReadiness() {
+    assertEquals(_healthCheck.getBrokerHealth(null), "OK");
+
+    when(_routingManager.isServerRoutable(SERVER_INSTANCE)).thenReturn(false);
+    WebApplicationException exception =
+        expectThrows(WebApplicationException.class, () -> _healthCheck.getBrokerHealth(SERVER_INSTANCE));
+    assertEquals(exception.getResponse().getStatus(), 503);
+
+    when(_routingManager.isServerRoutable(SERVER_INSTANCE)).thenReturn(true);
+    assertEquals(_healthCheck.getBrokerHealth(SERVER_INSTANCE), CommonConstants.Broker.SERVER_ROUTING_READY_RESPONSE);
+  }
+
+  private void setField(String name, Object value)
+      throws Exception {
+    Field field = PinotBrokerHealthCheck.class.getDeclaredField(name);
+    field.setAccessible(true);
+    field.set(_healthCheck, value);
+  }
+}

pinot-broker/src/test/java/org/apache/pinot/broker/broker/HelixBrokerStarterTest.java

@@ -21,6 +21,10 @@
 import java.util.HashMap;
 import java.util.Map;
 import java.util.concurrent.TimeUnit;
+import javax.ws.rs.client.Client;
+import javax.ws.rs.client.ClientBuilder;
+import javax.ws.rs.client.WebTarget;
+import javax.ws.rs.core.Response;
 import org.apache.helix.model.ExternalView;
 import org.apache.helix.model.HelixConfigScope;
 import org.apache.helix.model.IdealState;
@@ -148,6 +152,29 @@ public void testClusterConfigOverride() {
     assertEquals(config.getProperty(Broker.CONFIG_OF_BROKER_DEFAULT_QUERY_LIMIT, 1), 1000);
   }
 
+  @Test
+  public void testServerRoutingHealthCheck() {
+    BrokerRoutingManager routingManager = _brokerStarter.getRoutingManager();
+    TestUtils.waitForCondition(aVoid -> !routingManager.getEnabledServerInstanceMap().isEmpty(), 30_000L,
+        "Failed to find an enabled server");
+    String serverInstance = routingManager.getEnabledServerInstanceMap().keySet().iterator().next();
+    assertTrue(routingManager.isServerRoutable(serverInstance));
+
+    Client client = ClientBuilder.newClient();
+    try {
+      WebTarget healthTarget = client.target("http://localhost:18099/health");
+      try (Response response = healthTarget.queryParam("serverInstance", serverInstance).request().get()) {
+        assertEquals(response.getStatus(), 200);
+        assertEquals(response.readEntity(String.class), CommonConstants.Broker.SERVER_ROUTING_READY_RESPONSE);
+      }
+      try (Response response = healthTarget.queryParam("serverInstance", "Server_unknown_8098").request().get()) {
+        assertEquals(response.getStatus(), 503);
+      }
+    } finally {
+      client.close();
+    }
+  }
+
   @Test
   public void testResourceAndTagAssignment()
       throws Exception {

pinot-broker/src/test/java/org/apache/pinot/broker/routing/manager/BrokerRoutingManagerTest.java

@@ -165,6 +165,23 @@ public void testNoErrorWhenCallbackNotSet() {
     assertTrue(_routingManager.getEnabledServerInstanceMap().containsKey(SERVER_INSTANCE_ID));
   }
 
+  @Test
+  public void testServerRoutableState() {
+    assertFalse(_routingManager.isServerRoutable(SERVER_INSTANCE_ID));
+
+    List<ZNRecord> instanceConfigs = List.of(createEnabledServerZNRecord(SERVER_INSTANCE_ID));
+    when(_zkDataAccessor.getChildren(eq(INSTANCE_CONFIGS_PATH), any(), eq(AccessOption.PERSISTENT),
+        anyInt(), anyInt())).thenReturn(instanceConfigs);
+    _routingManager.processClusterChange(ChangeType.INSTANCE_CONFIG);
+    assertTrue(_routingManager.isServerRoutable(SERVER_INSTANCE_ID));
+
+    _routingManager.excludeServerFromRouting(SERVER_INSTANCE_ID);
+    assertFalse(_routingManager.isServerRoutable(SERVER_INSTANCE_ID));
+
+    _routingManager.includeServerToRouting(SERVER_INSTANCE_ID);
+    assertTrue(_routingManager.isServerRoutable(SERVER_INSTANCE_ID));
+  }
+
   @Test
   public void testServerReenableCallbackInvokedWhenExcludedServerReenabled() {
     // Set up callback

pinot-server/src/main/java/org/apache/pinot/server/api/AdminApiApplication.java

@@ -23,6 +23,7 @@
 import java.time.Instant;
 import java.util.List;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.function.BooleanSupplier;
 import javax.ws.rs.container.ContainerRequestContext;
 import javax.ws.rs.container.ContainerResponseContext;
 import javax.ws.rs.container.ContainerResponseFilter;
@@ -53,6 +54,7 @@ public class AdminApiApplication extends ResourceConfig {
   private static final Logger LOGGER = LoggerFactory.getLogger(AdminApiApplication.class);
   public static final String PINOT_CONFIGURATION = "pinotConfiguration";
   public static final String SERVER_INSTANCE_ID = "serverInstanceId";
+  public static final String SERVER_READY_TO_SERVE_QUERIES = "serverReadyToServeQueries";
 
   public static final String START_TIME = "serverStartTime";
 
@@ -63,8 +65,8 @@ public class AdminApiApplication extends ResourceConfig {
 
 
   public AdminApiApplication(ServerInstance instance, AccessControlFactory accessControlFactory,
-      ServerReloadJobStatusCache reloadJobStatusCache,
-      PinotConfiguration serverConf) {
+      ServerReloadJobStatusCache reloadJobStatusCache, PinotConfiguration serverConf,
+      BooleanSupplier isServerReadyToServeQueries) {
     _serverInstance = instance;
 
     _adminApiResourcePackages = serverConf.getProperty(CommonConstants.Server.CONFIG_OF_SERVER_RESOURCE_PACKAGES,
@@ -77,6 +79,7 @@ public AdminApiApplication(ServerInstance instance, AccessControlFactory accessC
       @Override
       protected void configure() {
         bind(_shutDownInProgress).to(AtomicBoolean.class);
+        bind(isServerReadyToServeQueries).to(BooleanSupplier.class).named(SERVER_READY_TO_SERVE_QUERIES);
         bind(_serverInstance).to(ServerInstance.class);
         bind(_serverInstance.getHelixManager()).to(HelixManager.class);
         bind(_serverInstance.getServerMetrics()).to(ServerMetrics.class);

pinot-server/src/main/java/org/apache/pinot/server/api/resources/HealthCheckResource.java

@@ -27,6 +27,7 @@
 import java.time.Instant;
 import java.time.format.DateTimeFormatter;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.function.BooleanSupplier;
 import javax.annotation.Nullable;
 import javax.inject.Inject;
 import javax.inject.Named;
@@ -55,6 +56,10 @@ public class HealthCheckResource {
   @Inject
   private AtomicBoolean _shutDownInProgress;
 
+  @Inject
+  @Named(AdminApiApplication.SERVER_READY_TO_SERVE_QUERIES)
+  private BooleanSupplier _isServerReadyToServeQueries;
+
   @Inject
   @Named(AdminApiApplication.SERVER_INSTANCE_ID)
   private String _instanceId;
@@ -111,6 +116,11 @@ public String checkReadiness() {
       throw new WebApplicationException(errMessage,
           Response.status(Response.Status.SERVICE_UNAVAILABLE).entity(errMessage).build());
     }
+    if (!_isServerReadyToServeQueries.getAsBoolean()) {
+      String errMessage = "Server is not ready to serve queries";
+      throw new WebApplicationException(errMessage,
+          Response.status(Response.Status.SERVICE_UNAVAILABLE).entity(errMessage).build());
+    }
     Status status = ServiceStatus.getServiceStatus(_instanceId);
     if (status == Status.GOOD) {
       _serverMetrics.addMeteredGlobalValue(ServerMeter.READINESS_CHECK_OK_CALLS, 1);

pinot-server/src/main/java/org/apache/pinot/server/starter/helix/BaseServerStarter.java

@@ -182,6 +182,7 @@ public abstract class BaseServerStarter implements ServiceStartable {
   protected QueryKillingManager _queryKillingManager;
   protected DefaultClusterConfigChangeHandler _clusterConfigChangeHandler;
   protected volatile boolean _isServerReadyToServeQueries = false;
+  protected volatile BrokerRoutingReadyChecker _brokerRoutingReadyChecker;
   protected ScheduledExecutorService _helixMessageCountScheduler;
   protected ServerReloadJobStatusCache _reloadJobStatusCache;
   // Override this to provide custom thread pool for Helix state transitions. Null means using Helix's default
@@ -911,6 +912,13 @@ public void start()
     _serverInstance.startQueryServer();
     _helixAdmin.setConfig(_instanceConfigScope,
         Map.of(Helix.IS_SHUTDOWN_IN_PROGRESS, Boolean.toString(false)));
+    if (_serverConf.getProperty(Server.CONFIG_OF_STARTUP_ENABLE_BROKER_ROUTING_CHECK,
+        Server.DEFAULT_STARTUP_ENABLE_BROKER_ROUTING_CHECK)) {
+      _brokerRoutingReadyChecker = createBrokerRoutingReadyChecker();
+      _brokerRoutingReadyChecker.start();
+    }
+    // Publish query readiness only after the optional routing checker is installed. Otherwise a concurrent health
+    // request could observe a transient ready state while the checker is still null.
     _isServerReadyToServeQueries = true;
     // Throttling for realtime consumption is disabled up to this point to allow maximum consumption during startup time
     RealtimeConsumptionRateManager.getInstance().enablePartitionRateLimiter();
@@ -968,6 +976,19 @@ protected boolean isServerReadyToServeQueries() {
     return _isServerReadyToServeQueries;
   }
 
+  protected boolean isServerReadyForHealthCheck() {
+    return isServerReadyToServeQueries()
+        && (_brokerRoutingReadyChecker == null || _brokerRoutingReadyChecker.isReady());
+  }
+
+  protected BrokerRoutingReadyChecker createBrokerRoutingReadyChecker() {
+    long timeoutMs = _serverConf.getProperty(Server.CONFIG_OF_STARTUP_BROKER_ROUTING_CHECK_TIMEOUT_MS,
+        Server.DEFAULT_STARTUP_BROKER_ROUTING_CHECK_TIMEOUT_MS);
+    boolean failOpen = _serverConf.getProperty(Server.CONFIG_OF_STARTUP_BROKER_ROUTING_CHECK_FAIL_OPEN,
+        Server.DEFAULT_STARTUP_BROKER_ROUTING_CHECK_FAIL_OPEN);
+    return new BrokerRoutingReadyChecker(_helixManager, _instanceId, timeoutMs, failOpen);
+  }
+
   protected SegmentOperationsThrottler createMultiColumnIndexPreprocessThrottler() {
     int maxConcurrency = Integer.parseInt(
         _serverConf.getProperty(Helix.CONFIG_OF_MAX_SEGMENT_MULTICOL_TEXT_INDEX_PREPROCESS_PARALLELISM,
@@ -1035,6 +1056,9 @@ public void stop() {
     _adminApiApplication.startShuttingDown();
     _helixAdmin.setConfig(_instanceConfigScope,
         Map.of(Helix.IS_SHUTDOWN_IN_PROGRESS, Boolean.toString(true)));
+    if (_brokerRoutingReadyChecker != null) {
+      _brokerRoutingReadyChecker.close();
+    }
     if (_transitionThreadPoolManager != null) {
       _transitionThreadPoolManager.shutdown();
     }
@@ -1259,7 +1283,8 @@ private void initSegmentFetcher(PinotConfiguration config)
   }
 
   protected AdminApiApplication createServerAdminApp() {
-    return new AdminApiApplication(_serverInstance, _accessControlFactory, _reloadJobStatusCache, _serverConf);
+    return new AdminApiApplication(_serverInstance, _accessControlFactory, _reloadJobStatusCache, _serverConf,
+        this::isServerReadyForHealthCheck);
   }
 
   /// Creates the [SegmentMessageHandlerFactory] used to handle user-defined Helix messages for segments.