jineshparakh commented on code in PR #19500: URL: https://github.com/apache/pinot/pull/19500#discussion_r3965519139
########## pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BrokerWarmupConfig.java: ########## @@ -0,0 +1,110 @@ +/** + * 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.requesthandler; + +import java.util.Arrays; +import java.util.List; +import java.util.stream.Collectors; +import org.apache.pinot.spi.env.PinotConfiguration; +import org.apache.pinot.spi.utils.CommonConstants.Broker; + + +/// Settings for the broker startup warmup performed before readiness is granted. +/// +/// The exit condition is a **depth floor OR a budget ceiling**, never a latency guess: +/// - [#budgetMs] is a hard ceiling. Warmup releases the readiness gate when it expires whatever the +/// probe progress, so a slow or unreachable server cannot stall a rolling restart. +/// - [#minIterations] is the floor for the default probe: it declares the broker warm once this many +/// probe queries have completed successfully -- enough invocations to drive the query path's JIT to +/// its top tier. (A latency target was deliberately avoided: a trivial probe reaches low latency +/// after a handful of iterations while the code is still only partially compiled, so latency is a +/// false early-exit signal.) +/// +/// Probe selection has two optional overrides, each falling back to auto-selection when unset: +/// - [#queries]: one or more complete probe queries (each names its own table), separated by `;` in +/// the config value when more than one. When set they are used instead of the default +/// `SELECT * FROM "<t>" LIMIT 1`, so representative production query shapes are warmed -- and several +/// shapes together warm several serve-path branches (scan, aggregation, group-by, join), which a +/// single shape does not. Every configured query is run each round, exiting on the same +/// [#minIterations] probe-count depth floor as the default probe (a cost-independent warm signal), +/// or the budget. [#tables] is ignored +/// when this is set. Note: the config layer treats a bare comma as its own list delimiter, so a comma Review Comment: No longer applicable — `warmup.queries` has been removed, so there's no operator-supplied SQL for the config layer to mangle. The probe is now the fixed `SELECT * FROM "<t>" LIMIT 1`. ########## pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BrokerRequestHandler.java: ########## @@ -47,6 +47,20 @@ public interface BrokerRequestHandler { void shutDown(); + /// Warms this handler's data plane so the first real query does not pay for it, and reports whether the + /// handler reached its warmth floor. Called during startup after Helix convergence, before readiness is + /// granted. Implementations must be **best effort and bounded**: they must return by `deadlineMs` and + /// must never throw, because the caller gates readiness on this and a stuck warmup would stall a rolling + /// restart. `deadlineMs` is an absolute [System#currentTimeMillis] value; time already spent waiting for + /// the cluster view to converge counts against the same budget, so the rolling-restart cost stays + /// bounded by one number. + /// + /// @return `true` if the handler reached its warmth floor, `false` if the deadline passed first. Either + /// way the caller proceeds; the result is for logging and metrics. + default boolean warmUp(BrokerWarmupConfig config, long deadlineMs) { Review Comment: Done — rewrote the interface javadoc to state all three contracts the caller depends on: must return by `deadlineMs` (with every wait bounded by the remaining budget), must never throw, and must be safe to run against a handler already serving traffic. Also documented how the gate composes with `preConnectServers` (readiness = all gates satisfied). ########## pinot-common/src/main/java/org/apache/pinot/common/metrics/BrokerGauge.java: ########## @@ -28,6 +28,15 @@ public enum BrokerGauge implements AbstractMetrics.Gauge { MAX_BURST_QPS("tables", false), QUERY_RATE_LIMIT_DISABLED("queryQuota", true), NETTY_CONNECTION_CONNECT_TIME_MS("nettyConnection", true), + // 1 once startup warmup has completed and readiness may be granted, 0 while still warming. Always 1 Review Comment: Fixed — renamed to `STARTUP_WARMUP_COMPLETE` and `STARTUP_WARMUP_UNCOVERED_SERVERS`, matching the `startup.warmup.*` config namespace and the `STARTUP_WARMUP_DURATION_MS` timer. ########## pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/BrokerStartupWarmupIntegrationTest.java: ########## @@ -0,0 +1,162 @@ +/** + * 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.integration.tests; + +import java.io.File; +import java.util.List; +import org.apache.commons.io.FileUtils; +import org.apache.pinot.broker.requesthandler.BrokerWarmupConfig; +import org.apache.pinot.common.utils.ServiceStatus; +import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.data.Schema; +import org.apache.pinot.spi.env.PinotConfiguration; +import org.apache.pinot.spi.utils.CommonConstants; +import org.apache.pinot.spi.utils.builder.TableNameBuilder; +import org.apache.pinot.util.TestUtils; +import org.testng.Assert; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Test; + + +/// Integration test for the broker startup data-plane warmup feature +/// (`pinot.broker.startup.warmup.*`). Brings up a real ZK + controller + server + broker with an offline +/// table, then verifies two things end-to-end: +/// +/// 1. With warmup enabled the readiness gate opens: the broker's [ServiceStatus] reaches `GOOD`. +/// Readiness is held at `STARTING` until warmup completes, and the broker's query path serves +/// regardless of [ServiceStatus], so a gate that never released would not fail `setUp` -- asserting +/// `GOOD` is what actually proves the gate released after warmup. +/// 2. The production warmup path (`RoutingManager` set-cover -> [QueryRouter] probe -> reduce) reaches +/// its depth floor and returns `true` within the budget against a live server. +public class BrokerStartupWarmupIntegrationTest extends BaseClusterIntegrationTest { + private static final long WARMUP_BUDGET_MS = 30_000L; + // Small floor so the integration test warms and opens readiness quickly (the depth floor is only large + // in production, to drive JIT to its top tier). + private static final int WARMUP_MIN_ITERATIONS = 10; + + @Override + protected void overrideBrokerConf(PinotConfiguration brokerConf) { + brokerConf.setProperty(CommonConstants.Broker.CONFIG_OF_BROKER_STARTUP_WARMUP_ENABLED, true); + brokerConf.setProperty(CommonConstants.Broker.CONFIG_OF_BROKER_STARTUP_WARMUP_BUDGET_MS, WARMUP_BUDGET_MS); + brokerConf.setProperty(CommonConstants.Broker.CONFIG_OF_BROKER_STARTUP_WARMUP_MIN_ITERATIONS, + WARMUP_MIN_ITERATIONS); + } + + @BeforeClass + public void setUp() + throws Exception { + TestUtils.ensureDirectoriesExistAndEmpty(_tempDir, _segmentDir, _tarDir); + startZk(); + startController(); + startBroker(); + startServer(); + + Schema schema = createSchema(); + addSchema(schema); + TableConfig tableConfig = createOfflineTableConfig(); + addTableConfig(tableConfig); + + // Build and upload segments so the broker has a live server to route probe queries to. + List<File> avroFiles = unpackAvroData(_tempDir); + ClusterIntegrationTestUtils.buildSegmentsFromAvro(avroFiles, tableConfig, schema, 0, _segmentDir, _tarDir); + uploadSegments(getTableName(), _tarDir); + + waitForAllDocsLoaded(600_000L); + } + + @AfterClass + public void tearDown() + throws Exception { + dropOfflineTable(getTableName()); + stopBroker(); + stopServer(); + stopController(); + stopZk(); + FileUtils.deleteDirectory(_tempDir); + } + + @Test + public void warmupEnabledBrokerReachesGoodServiceStatus() { Review Comment: Fixed. Added `BrokerWarmupGateTest` — a deterministic unit test of the gate's STARTING → GOOD transition (STARTING + "Warming up broker data plane" while cold, GOOD + no description once warm), since a healthy cluster warms faster than an IT can observe mid-flight. Also strengthened this IT assertion to confirm the warming description is gone once GOOD. ########## pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/BrokerStartupWarmupIntegrationTest.java: ########## @@ -0,0 +1,162 @@ +/** + * 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.integration.tests; + +import java.io.File; +import java.util.List; +import org.apache.commons.io.FileUtils; +import org.apache.pinot.broker.requesthandler.BrokerWarmupConfig; +import org.apache.pinot.common.utils.ServiceStatus; +import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.data.Schema; +import org.apache.pinot.spi.env.PinotConfiguration; +import org.apache.pinot.spi.utils.CommonConstants; +import org.apache.pinot.spi.utils.builder.TableNameBuilder; +import org.apache.pinot.util.TestUtils; +import org.testng.Assert; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Test; + + +/// Integration test for the broker startup data-plane warmup feature +/// (`pinot.broker.startup.warmup.*`). Brings up a real ZK + controller + server + broker with an offline +/// table, then verifies two things end-to-end: +/// +/// 1. With warmup enabled the readiness gate opens: the broker's [ServiceStatus] reaches `GOOD`. +/// Readiness is held at `STARTING` until warmup completes, and the broker's query path serves +/// regardless of [ServiceStatus], so a gate that never released would not fail `setUp` -- asserting +/// `GOOD` is what actually proves the gate released after warmup. +/// 2. The production warmup path (`RoutingManager` set-cover -> [QueryRouter] probe -> reduce) reaches +/// its depth floor and returns `true` within the budget against a live server. +public class BrokerStartupWarmupIntegrationTest extends BaseClusterIntegrationTest { + private static final long WARMUP_BUDGET_MS = 30_000L; + // Small floor so the integration test warms and opens readiness quickly (the depth floor is only large + // in production, to drive JIT to its top tier). + private static final int WARMUP_MIN_ITERATIONS = 10; + + @Override + protected void overrideBrokerConf(PinotConfiguration brokerConf) { + brokerConf.setProperty(CommonConstants.Broker.CONFIG_OF_BROKER_STARTUP_WARMUP_ENABLED, true); + brokerConf.setProperty(CommonConstants.Broker.CONFIG_OF_BROKER_STARTUP_WARMUP_BUDGET_MS, WARMUP_BUDGET_MS); + brokerConf.setProperty(CommonConstants.Broker.CONFIG_OF_BROKER_STARTUP_WARMUP_MIN_ITERATIONS, + WARMUP_MIN_ITERATIONS); + } + + @BeforeClass + public void setUp() + throws Exception { + TestUtils.ensureDirectoriesExistAndEmpty(_tempDir, _segmentDir, _tarDir); + startZk(); + startController(); + startBroker(); + startServer(); + + Schema schema = createSchema(); + addSchema(schema); + TableConfig tableConfig = createOfflineTableConfig(); + addTableConfig(tableConfig); + + // Build and upload segments so the broker has a live server to route probe queries to. + List<File> avroFiles = unpackAvroData(_tempDir); + ClusterIntegrationTestUtils.buildSegmentsFromAvro(avroFiles, tableConfig, schema, 0, _segmentDir, _tarDir); + uploadSegments(getTableName(), _tarDir); + + waitForAllDocsLoaded(600_000L); + } + + @AfterClass + public void tearDown() + throws Exception { + dropOfflineTable(getTableName()); + stopBroker(); + stopServer(); + stopController(); + stopZk(); + FileUtils.deleteDirectory(_tempDir); + } + + @Test + public void warmupEnabledBrokerReachesGoodServiceStatus() { + String instanceId = _brokerStarters.get(0).getInstanceId(); + TestUtils.waitForCondition(aVoid -> ServiceStatus.getServiceStatus(instanceId) == ServiceStatus.Status.GOOD, + WARMUP_BUDGET_MS, "Broker with warmup enabled never reported GOOD service status"); + } + + @Test + public void warmUpReachesFloorAgainstLiveServer() { + BrokerWarmupConfig config = new BrokerWarmupConfig(true, WARMUP_BUDGET_MS, WARMUP_MIN_ITERATIONS, 5, + List.of(), List.of(), 1); + boolean reachedFloor = _brokerStarters.get(0).getBrokerRequestHandler() + .warmUp(config, System.currentTimeMillis() + WARMUP_BUDGET_MS); + Assert.assertTrue(reachedFloor, + "Warmup should reach its probe-count floor against a live server within the budget"); + } + + @Test + public void warmUpReachesFloorWithMultipleCustomQueries() { + // Two distinct query shapes drive the custom-query path (each run every round, then reduced) against a + // live server; warmup should reach its probe-count floor before the budget and return true. + BrokerWarmupConfig config = new BrokerWarmupConfig(true, WARMUP_BUDGET_MS, WARMUP_MIN_ITERATIONS, 5, + List.of("SELECT COUNT(*) FROM " + getTableName(), "SELECT * FROM " + getTableName() + " LIMIT 1"), + List.of(), 1); + boolean warmed = _brokerStarters.get(0).getBrokerRequestHandler() + .warmUp(config, System.currentTimeMillis() + WARMUP_BUDGET_MS); + Assert.assertTrue(warmed, + "Custom multi-query warmup should reach its floor against a live server within the budget"); + } + + @Test + public void warmUpReachesFloorWithConcurrentProbes() { + // Concurrency 3: probes fire in parallel on a 3-thread pool, exercising the concurrent scatter/gather + // path the serial arms do not. + BrokerWarmupConfig config = new BrokerWarmupConfig(true, WARMUP_BUDGET_MS, WARMUP_MIN_ITERATIONS, 5, + List.of(), List.of(), 3); + boolean reachedFloor = _brokerStarters.get(0).getBrokerRequestHandler() + .warmUp(config, System.currentTimeMillis() + WARMUP_BUDGET_MS); + Assert.assertTrue(reachedFloor, "Concurrent warmup should reach its probe-count floor within the budget"); + } + + @Test + public void warmUpWithExplicitTableListReachesFloor() { + // Custom tables path: filterRoutableTables keeps the configured (routable) table, then the default probe + // runs against it. Exercises the hasCustomTables branch that auto-select does not. + String offlineTable = TableNameBuilder.OFFLINE.tableNameWithType(getTableName()); + BrokerWarmupConfig config = new BrokerWarmupConfig(true, WARMUP_BUDGET_MS, WARMUP_MIN_ITERATIONS, 5, + List.of(), List.of(offlineTable), 1); + boolean reachedFloor = _brokerStarters.get(0).getBrokerRequestHandler() + .warmUp(config, System.currentTimeMillis() + WARMUP_BUDGET_MS); + Assert.assertTrue(reachedFloor, "Warmup over an explicit table list should reach its floor within the budget"); + } + + @Test + public void warmUpReturnsFalseWhenBudgetExpiresBeforeFloor() { + // An unreachable floor with a short budget must exit on the budget (returning false) and must not hang. + long shortBudgetMs = 2_000L; + BrokerWarmupConfig config = new BrokerWarmupConfig(true, shortBudgetMs, 100_000_000, 5, + List.of(), List.of(), 1); + long start = System.currentTimeMillis(); + boolean reachedFloor = _brokerStarters.get(0).getBrokerRequestHandler() + .warmUp(config, start + shortBudgetMs); + long elapsed = System.currentTimeMillis() - start; + Assert.assertFalse(reachedFloor, "Warmup must return false when the budget expires before the floor"); + Assert.assertTrue(elapsed < shortBudgetMs + 10_000L, Review Comment: Fixed — tightened to `budget + 2s`. With the budget now a hard ceiling (each get bounded by the remaining budget, stragglers cancelled), the return is near-immediate after the budget, so the loose slack was hiding exactly the overrun flagged above. ########## pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/BrokerWarmupTest.java: ########## @@ -0,0 +1,355 @@ +/** + * 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.requesthandler; + +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; +import org.apache.pinot.common.metrics.BrokerGauge; +import org.apache.pinot.common.metrics.BrokerMetrics; +import org.apache.pinot.core.routing.RoutingManager; +import org.apache.pinot.core.routing.TableRouteInfo; +import org.apache.pinot.core.transport.ServerInstance; +import org.apache.pinot.spi.config.table.TableType; +import org.apache.pinot.spi.env.PinotConfiguration; +import org.apache.pinot.spi.utils.CommonConstants.Broker; +import org.testng.annotations.Test; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertTrue; + + +/// Unit coverage for broker startup warmup: config parsing (including the concurrency and minIterations +/// knobs and their floors), the no-op interface default, the greedy table selection that keeps warmup +/// bounded on a large tenant, and the server-coverage check that flags when `maxTables` leaves servers +/// unprobed. +public class BrokerWarmupTest { Review Comment: Fixed. Extracted `runConcurrentRound` (package-private, static) and added unit tests with injected probe callables covering: collects only successful latencies; returns by the budget when probes queue past it and cancels the stragglers; an already-expired deadline collects nothing; and an interrupt ends the round promptly while preserving the interrupt flag. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
