diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/api/resources/PinotBrokerHealthCheck.java b/pinot-broker/src/main/java/org/apache/pinot/broker/api/resources/PinotBrokerHealthCheck.java index 8743f961407b..d4eac37b32b5 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/api/resources/PinotBrokerHealthCheck.java +++ b/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"; } diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java index e0ccad18e369..2def89340612 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java @@ -1258,6 +1258,19 @@ public Set 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 diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/api/resources/PinotBrokerHealthCheckTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/api/resources/PinotBrokerHealthCheckTest.java new file mode 100644 index 000000000000..2201f47fd3c7 --- /dev/null +++ b/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); + } +} diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/broker/HelixBrokerStarterTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/broker/HelixBrokerStarterTest.java index 180d72a13d22..0c82826d8da0 100644 --- a/pinot-broker/src/test/java/org/apache/pinot/broker/broker/HelixBrokerStarterTest.java +++ b/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 { diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/routing/manager/BrokerRoutingManagerTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/routing/manager/BrokerRoutingManagerTest.java index bc4134ebb615..89bbff76b7be 100644 --- a/pinot-broker/src/test/java/org/apache/pinot/broker/routing/manager/BrokerRoutingManagerTest.java +++ b/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 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 diff --git a/pinot-server/src/main/java/org/apache/pinot/server/api/AdminApiApplication.java b/pinot-server/src/main/java/org/apache/pinot/server/api/AdminApiApplication.java index e80b2da117e4..ac2ef6d3ecb3 100644 --- a/pinot-server/src/main/java/org/apache/pinot/server/api/AdminApiApplication.java +++ b/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); diff --git a/pinot-server/src/main/java/org/apache/pinot/server/api/resources/HealthCheckResource.java b/pinot-server/src/main/java/org/apache/pinot/server/api/resources/HealthCheckResource.java index 9291070dd977..65c611148255 100644 --- a/pinot-server/src/main/java/org/apache/pinot/server/api/resources/HealthCheckResource.java +++ b/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); diff --git a/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/BaseServerStarter.java b/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/BaseServerStarter.java index c156d3a20777..aba5eea7008f 100644 --- a/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/BaseServerStarter.java +++ b/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. diff --git a/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/BrokerRoutingReadyChecker.java b/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/BrokerRoutingReadyChecker.java new file mode 100644 index 000000000000..f6bfa046937a --- /dev/null +++ b/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/BrokerRoutingReadyChecker.java @@ -0,0 +1,191 @@ +/** + * 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.server.starter.helix; + +import com.google.common.annotations.VisibleForTesting; +import com.google.common.util.concurrent.ThreadFactoryBuilder; +import java.net.URI; +import java.net.URLEncoder; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.List; +import java.util.Set; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; +import java.util.function.LongSupplier; +import java.util.function.Predicate; +import java.util.function.Supplier; +import org.apache.helix.HelixAdmin; +import org.apache.helix.HelixManager; +import org.apache.helix.model.ExternalView; +import org.apache.helix.model.InstanceConfig; +import org.apache.pinot.common.utils.SimpleHttpResponse; +import org.apache.pinot.common.utils.config.InstanceUtils; +import org.apache.pinot.common.utils.helix.HelixHelper; +import org.apache.pinot.common.utils.http.HttpClient; +import org.apache.pinot.spi.utils.CommonConstants; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + + +/// Checks broker routing state in the background while a server starts. The checker remains false until every online +/// broker reports the server as routable, or until the configured timeout when fail-open behavior is enabled. Health +/// endpoints only read the cached result. Brokers must return the routing-specific response, so an older broker's +/// normal health response cannot be mistaken for an acknowledgement. All mutable state is confined to the scheduler +/// thread except for the volatile ready flag read by health-check threads. +public class BrokerRoutingReadyChecker implements AutoCloseable { + private static final Logger LOGGER = LoggerFactory.getLogger(BrokerRoutingReadyChecker.class); + private static final long CHECK_INTERVAL_MS = 1_000L; + private static final long CHECK_TIMEOUT_MS = 5_000L; + + private final String _serverInstanceId; + private final Supplier> _onlineBrokersSupplier; + private final Predicate> _allBrokersReady; + private final ScheduledExecutorService _scheduler; + private final ExecutorService _requestExecutor; + private final LongSupplier _currentTimeMs; + private final long _deadlineMs; + private final boolean _failOpen; + private volatile boolean _ready; + private boolean _timeoutLogged; + + public BrokerRoutingReadyChecker(HelixManager helixManager, String serverInstanceId, long timeoutMs, + boolean failOpen) { + _serverInstanceId = serverInstanceId; + HelixAdmin helixAdmin = helixManager.getClusterManagmentTool(); + String clusterName = helixManager.getClusterName(); + _onlineBrokersSupplier = () -> { + ExternalView brokerResource = helixAdmin.getResourceExternalView(clusterName, + CommonConstants.Helix.BROKER_RESOURCE_INSTANCE); + return Set.copyOf(HelixHelper.getOnlineInstanceFromExternalView(brokerResource)); + }; + _requestExecutor = Executors.newCachedThreadPool( + new ThreadFactoryBuilder().setNameFormat("broker-routing-ready-request-%d").setDaemon(true).build()); + _allBrokersReady = brokers -> checkAllBrokers(helixAdmin, clusterName, brokers); + _scheduler = Executors.newSingleThreadScheduledExecutor( + new ThreadFactoryBuilder().setNameFormat("broker-routing-ready-checker").setDaemon(true).build()); + _currentTimeMs = System::currentTimeMillis; + _deadlineMs = _currentTimeMs.getAsLong() + timeoutMs; + _failOpen = failOpen; + } + + @VisibleForTesting + BrokerRoutingReadyChecker(String serverInstanceId, Supplier> onlineBrokersSupplier, + Predicate> allBrokersReady) { + this(serverInstanceId, onlineBrokersSupplier, allBrokersReady, Long.MAX_VALUE, false, () -> 0L); + } + + @VisibleForTesting + BrokerRoutingReadyChecker(String serverInstanceId, Supplier> onlineBrokersSupplier, + Predicate> allBrokersReady, long timeoutMs, boolean failOpen, LongSupplier currentTimeMs) { + _serverInstanceId = serverInstanceId; + _onlineBrokersSupplier = onlineBrokersSupplier; + _allBrokersReady = allBrokersReady; + _requestExecutor = null; + _scheduler = null; + _currentTimeMs = currentTimeMs; + _deadlineMs = currentTimeMs.getAsLong() + timeoutMs; + _failOpen = failOpen; + } + + public void start() { + _scheduler.scheduleWithFixedDelay(this::check, 0L, CHECK_INTERVAL_MS, TimeUnit.MILLISECONDS); + } + + public boolean isReady() { + return _ready; + } + + @VisibleForTesting + void check() { + if (_ready) { + return; + } + try { + Set onlineBrokers = _onlineBrokersSupplier.get(); + if (!onlineBrokers.isEmpty() && _allBrokersReady.test(onlineBrokers) + // Do not mark the server ready if broker membership changed while acknowledgements were collected. + && onlineBrokers.equals(_onlineBrokersSupplier.get())) { + _ready = true; + LOGGER.info("All online brokers report server {} as routable: {}", _serverInstanceId, onlineBrokers); + return; + } + } catch (Exception e) { + LOGGER.debug("Failed to check broker routing readiness for server {}", _serverInstanceId, e); + } + + if (_currentTimeMs.getAsLong() >= _deadlineMs) { + if (!_timeoutLogged) { + LOGGER.warn("Timed out waiting for all online brokers to report server {} as routable; failOpen={}", + _serverInstanceId, _failOpen); + _timeoutLogged = true; + } + if (_failOpen) { + _ready = true; + } + } + } + + private boolean checkAllBrokers(HelixAdmin helixAdmin, String clusterName, Set brokers) { + List> futures = new ArrayList<>(brokers.size()); + for (String broker : brokers) { + futures.add(CompletableFuture.supplyAsync(() -> checkBroker(helixAdmin, clusterName, broker), _requestExecutor)); + } + try { + CompletableFuture.allOf(futures.toArray(new CompletableFuture[0])) + .get(CHECK_TIMEOUT_MS, TimeUnit.MILLISECONDS); + return futures.stream().allMatch(CompletableFuture::join); + } catch (Exception e) { + futures.forEach(future -> future.cancel(true)); + LOGGER.debug("Failed to collect routing readiness from all online brokers for server {}", _serverInstanceId, e); + return false; + } + } + + private boolean checkBroker(HelixAdmin helixAdmin, String clusterName, String broker) { + try { + InstanceConfig instanceConfig = helixAdmin.getInstanceConfig(clusterName, broker); + if (instanceConfig == null) { + return false; + } + String serverInstance = URLEncoder.encode(_serverInstanceId, StandardCharsets.UTF_8); + URI uri = URI.create(InstanceUtils.getInstanceBaseUri(instanceConfig) + "/health?serverInstance=" + + serverInstance); + SimpleHttpResponse response = HttpClient.getInstance().sendGetRequest(uri); + return response.getStatusCode() == 200 + && CommonConstants.Broker.SERVER_ROUTING_READY_RESPONSE.equals(response.getResponse()); + } catch (Exception e) { + LOGGER.debug("Broker {} has not confirmed routing readiness for server {}", broker, _serverInstanceId, e); + return false; + } + } + + @Override + public void close() { + if (_scheduler != null) { + _scheduler.shutdownNow(); + } + if (_requestExecutor != null) { + _requestExecutor.shutdownNow(); + } + } +} diff --git a/pinot-server/src/test/java/org/apache/pinot/server/api/AccessControlTest.java b/pinot-server/src/test/java/org/apache/pinot/server/api/AccessControlTest.java index 4c4231bab111..b3e09b7d7c30 100644 --- a/pinot-server/src/test/java/org/apache/pinot/server/api/AccessControlTest.java +++ b/pinot-server/src/test/java/org/apache/pinot/server/api/AccessControlTest.java @@ -87,7 +87,7 @@ public void setUp() CommonConstants.Helix.PREFIX_OF_SERVER_INSTANCE + hostname + "_" + port); _adminApiApplication = new AdminApiApplication(serverInstance, new DenyAllAccessFactory(), mock(ServerReloadJobStatusCache.class), - serverConf); + serverConf, () -> true); int adminApiApplicationPort = getAvailablePort(); _adminApiApplication.start(List.of( diff --git a/pinot-server/src/test/java/org/apache/pinot/server/api/BaseResourceTest.java b/pinot-server/src/test/java/org/apache/pinot/server/api/BaseResourceTest.java index 200fcce5cb0c..094f95f149fe 100644 --- a/pinot-server/src/test/java/org/apache/pinot/server/api/BaseResourceTest.java +++ b/pinot-server/src/test/java/org/apache/pinot/server/api/BaseResourceTest.java @@ -27,6 +27,7 @@ import java.util.List; import java.util.Map; import java.util.concurrent.Executors; +import java.util.concurrent.atomic.AtomicBoolean; import javax.ws.rs.client.ClientBuilder; import javax.ws.rs.client.WebTarget; import org.apache.commons.io.FileUtils; @@ -89,6 +90,7 @@ public abstract class BaseResourceTest { protected final List _realtimeIndexSegments = new ArrayList<>(); protected final List _offlineIndexSegments = new ArrayList<>(); protected File _tempDir; + protected final AtomicBoolean _isServerReadyToServeQueries = new AtomicBoolean(true); protected File _avroFile; protected AdminApiApplication _adminApiApplication; protected WebTarget _webTarget; @@ -162,7 +164,7 @@ public void setUp() configureServerConf(serverConf); _adminApiApplication = new AdminApiApplication(_serverInstance, new AllowAllAccessFactory(), mock(ServerReloadJobStatusCache.class), - serverConf); + serverConf, _isServerReadyToServeQueries::get); _adminApiApplication.start(List.of( new ListenerConfig(CommonConstants.HTTP_PROTOCOL, "0.0.0.0", 0, CommonConstants.HTTP_PROTOCOL, new TlsConfig(), HttpServerThreadPoolConfig.defaultInstance()))); diff --git a/pinot-server/src/test/java/org/apache/pinot/server/api/HealthCheckResourceTest.java b/pinot-server/src/test/java/org/apache/pinot/server/api/HealthCheckResourceTest.java index 1d13b8761fee..07bfba9ef50c 100644 --- a/pinot-server/src/test/java/org/apache/pinot/server/api/HealthCheckResourceTest.java +++ b/pinot-server/src/test/java/org/apache/pinot/server/api/HealthCheckResourceTest.java @@ -59,6 +59,15 @@ public void checkHealthProbes() { _webTarget.path(healthPath).queryParam("checkType", "readiness").request().get(Response.class).getStatus(), 200); + _isServerReadyToServeQueries.set(false); + assertEquals(_webTarget.path(livenessPath).request().get(Response.class).getStatus(), 200); + assertEquals(_webTarget.path(healthPath).request().get(Response.class).getStatus(), 503); + assertEquals(_webTarget.path(readinessPath).request().get(Response.class).getStatus(), 503); + assertEquals( + _webTarget.path(healthPath).queryParam("checkType", "readiness").request().get(Response.class).getStatus(), + 503); + _isServerReadyToServeQueries.set(true); + ServiceStatus.setServiceStatusCallback(_instanceId, mockFailureCallback); assertEquals(_webTarget.path(livenessPath).request().get(Response.class).getStatus(), 200); assertEquals( diff --git a/pinot-server/src/test/java/org/apache/pinot/server/starter/helix/BrokerRoutingReadyCheckerTest.java b/pinot-server/src/test/java/org/apache/pinot/server/starter/helix/BrokerRoutingReadyCheckerTest.java new file mode 100644 index 000000000000..1183b20e9b72 --- /dev/null +++ b/pinot-server/src/test/java/org/apache/pinot/server/starter/helix/BrokerRoutingReadyCheckerTest.java @@ -0,0 +1,103 @@ +/** + * 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.server.starter.helix; + +import java.util.Set; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.atomic.AtomicReference; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertTrue; + + +public class BrokerRoutingReadyCheckerTest { + private static final String SERVER_INSTANCE = "Server_localhost_8098"; + + @Test + public void testWaitsForOnlineBrokers() { + AtomicReference> onlineBrokers = new AtomicReference<>(Set.of()); + BrokerRoutingReadyChecker checker = + new BrokerRoutingReadyChecker(SERVER_INSTANCE, onlineBrokers::get, brokers -> true); + + checker.check(); + assertFalse(checker.isReady()); + + onlineBrokers.set(Set.of("Broker_localhost_8099")); + checker.check(); + assertTrue(checker.isReady()); + } + + @Test + public void testRetriesUntilAllBrokersConfirm() { + AtomicInteger attempts = new AtomicInteger(); + BrokerRoutingReadyChecker checker = new BrokerRoutingReadyChecker(SERVER_INSTANCE, + () -> Set.of("Broker_localhost_8099", "Broker_localhost_8100"), + brokers -> attempts.incrementAndGet() > 1); + + checker.check(); + assertFalse(checker.isReady()); + + checker.check(); + assertTrue(checker.isReady()); + } + + @Test + public void testBrokerMembershipMustRemainStable() { + AtomicInteger reads = new AtomicInteger(); + BrokerRoutingReadyChecker checker = new BrokerRoutingReadyChecker(SERVER_INSTANCE, + () -> reads.incrementAndGet() == 1 ? Set.of("Broker_localhost_8099") + : Set.of("Broker_localhost_8099", "Broker_localhost_8100"), + brokers -> true); + + checker.check(); + assertFalse(checker.isReady()); + } + + @Test + public void testTimeoutFailsOpen() { + AtomicLong currentTimeMs = new AtomicLong(1_000L); + BrokerRoutingReadyChecker checker = new BrokerRoutingReadyChecker(SERVER_INSTANCE, + () -> Set.of("Broker_localhost_8099"), brokers -> false, 5_000L, true, currentTimeMs::get); + + checker.check(); + assertFalse(checker.isReady()); + + currentTimeMs.set(6_000L); + checker.check(); + assertTrue(checker.isReady()); + } + + @Test + public void testTimeoutFailsClosedButRecovers() { + AtomicLong currentTimeMs = new AtomicLong(1_000L); + AtomicReference brokerReady = new AtomicReference<>(false); + BrokerRoutingReadyChecker checker = new BrokerRoutingReadyChecker(SERVER_INSTANCE, + () -> Set.of("Broker_localhost_8099"), brokers -> brokerReady.get(), 5_000L, false, currentTimeMs::get); + + currentTimeMs.set(6_000L); + checker.check(); + assertFalse(checker.isReady()); + + brokerReady.set(true); + checker.check(); + assertTrue(checker.isReady()); + } +} diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java index c117713b8c81..dc4ddc13a7f2 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java @@ -416,6 +416,7 @@ public static class Broker { // Comma separated list of packages that contains javax service resources. public static final String BROKER_RESOURCE_PACKAGES = "broker.restlet.api.resource.packages"; public static final String DEFAULT_BROKER_RESOURCE_PACKAGES = "org.apache.pinot.broker.api.resources"; + public static final String SERVER_ROUTING_READY_RESPONSE = "ROUTABLE"; // Configuration to consider the broker ServiceStatus as being STARTED if the percent of resources (tables) that // are ONLINE for this broker has crossed the threshold percentage of the total number of tables @@ -1589,6 +1590,19 @@ public static class Server { public static final String CONFIG_OF_STARTUP_SERVICE_STATUS_CHECK_INTERVAL_MS = "pinot.server.startup.serviceStatusCheckIntervalMs"; public static final long DEFAULT_STARTUP_SERVICE_STATUS_CHECK_INTERVAL_MS = 10_000L; + // Startup: wait for brokers to add the server to routing before reporting the server as ready. Disabled by default + // for the initial rollout so mixed-version clusters retain the existing readiness behavior. + public static final String CONFIG_OF_STARTUP_ENABLE_BROKER_ROUTING_CHECK = + "pinot.server.startup.enableBrokerRoutingCheck"; + public static final boolean DEFAULT_STARTUP_ENABLE_BROKER_ROUTING_CHECK = false; + public static final String CONFIG_OF_STARTUP_BROKER_ROUTING_CHECK_TIMEOUT_MS = + "pinot.server.startup.brokerRoutingCheckTimeoutMs"; + public static final long DEFAULT_STARTUP_BROKER_ROUTING_CHECK_TIMEOUT_MS = 60_000L; + // When true, report ready after the timeout even if one or more brokers have not confirmed routing. When false, + // keep reporting unready and continue checking until the brokers recover. + public static final String CONFIG_OF_STARTUP_BROKER_ROUTING_CHECK_FAIL_OPEN = + "pinot.server.startup.brokerRoutingCheckFailOpen"; + public static final boolean DEFAULT_STARTUP_BROKER_ROUTING_CHECK_FAIL_OPEN = true; // Shutdown: timeout for the shutdown checks public static final String CONFIG_OF_SHUTDOWN_TIMEOUT_MS = "pinot.server.shutdown.timeoutMs"; public static final long DEFAULT_SHUTDOWN_TIMEOUT_MS = 600_000L;