gortiz commented on code in PR #19500:
URL: https://github.com/apache/pinot/pull/19500#discussion_r3955478615
##########
pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/SingleConnectionBrokerRequestHandler.java:
##########
@@ -117,6 +138,438 @@ public void shutDown() {
_brokerReduceService.shutDown();
}
+ /// Warms the scatter-gather path before readiness is granted: repeatedly
runs real probe queries until
+ /// they have run enough times (the `minIterations` depth floor) or the
budget expires. Broker-to-server
+ /// channels are opened lazily by the probes themselves.
+ ///
+ /// Probes go through [QueryRouter] rather than [#handleRequest],
deliberately bypassing access control,
+ /// query quota and the query log -- warmup then needs no synthetic identity
and pollutes no
+ /// customer-facing surface. Never throws, and always returns within the
configured budget.
+ @Override
+ public boolean warmUp(BrokerWarmupConfig config, long deadlineMs) {
+ // Stage 1: local, no-network warmup of the compile +
response-serialization paths, run before the
+ // network probe so they are warm even if no server is reachable.
+ warmUpLocal();
+ int concurrency = Math.max(1, config.concurrency());
+ ExecutorService probePool = Executors.newFixedThreadPool(concurrency,
+ new
ThreadFactoryBuilder().setNameFormat("broker-warmup-probe-%d").setDaemon(true).build());
+ try {
+ return config.hasCustomQuery() ? warmUpWithCustomQuery(config,
deadlineMs, probePool, concurrency)
+ : warmUpWithAutoSelect(config, deadlineMs, probePool, concurrency);
+ } catch (Exception e) {
+ LOGGER.warn("Broker warmup failed; proceeding without it", e);
+ return false;
+ } finally {
+ probePool.shutdownNow();
+ }
+ }
+
+ /// Stage 1: local, no-network warmup. The network probe goes through
[QueryRouter] and stops at the
+ /// gathered DataTables, so it never exercises the SQL compile path's cold
start nor the response build +
+ /// JSON serialization the first real query pays. This compiles a throwaway
query and serializes a small
+ /// synthetic [BrokerResponseNative], warming both. Never throws.
+ private void warmUpLocal() {
+ try {
+ CalciteSqlCompiler.compileToBrokerRequest("SELECT 1");
+ BrokerResponseNative response = new BrokerResponseNative();
+ response.setResultTable(new ResultTable(
+ new DataSchema(new String[]{"warmup"}, new
DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.LONG}),
+ Collections.singletonList(new Object[]{1L})));
+ response.setNumDocsScanned(1);
+ response.toJsonString();
+ } catch (Exception e) {
+ LOGGER.debug("Local (stage 1) warmup failed; continuing", e);
+ }
+ }
+
+ /// Fires the given probe tasks concurrently on the shared pool and returns
the latencies (ms) of the
+ /// probes that completed successfully. Empty means the whole round was
unproductive (nothing routable
+ /// yet / all failed); a failed probe simply does not contribute.
+ private List<Long> runConcurrentRound(List<Callable<Long>> tasks,
ExecutorService pool) {
+ List<Future<Long>> futures = new ArrayList<>(tasks.size());
+ for (Callable<Long> task : tasks) {
+ futures.add(pool.submit(task));
+ }
+ List<Long> latencies = new ArrayList<>(futures.size());
+ for (Future<Long> future : futures) {
+ try {
+ Long elapsedMs = future.get(WARMUP_PROBE_TIMEOUT_MS + 1000L,
TimeUnit.MILLISECONDS);
Review Comment:
**Correctness / major — the budget is not the hard ceiling the design relies
on.**
This `get()` uses the constant `WARMUP_PROBE_TIMEOUT_MS + 1000` (6s) rather
than the remaining budget, and the gets are sequential. Two ways that
overshoots `deadlineMs`:
1. A round can be entered with `remainingMs == 1`. `probeTimeoutMs` is then
1ms, but this `get()` still blocks up to 6s.
2. In `warmUpWithCustomQuery`, `tasks = queries.size() * concurrency` while
the pool has only `concurrency` threads, so tasks queue. Each queued future's
6s window starts when *its* `get()` is called, so the worst-case round is
`tasks x 6s`, not 6s. With 4 configured queries that is ~24s past a 15s budget,
and readiness stays `STARTING` for the whole overrun.
The rolling-restart safety argument in the description rests on "readiness
opens when it expires whatever the probe progress", so I think this one needs
fixing before merge. Suggestion: bound the wait by the deadline, e.g.
`future.get(Math.max(0, deadlineMs - System.currentTimeMillis()) + slack,
MILLISECONDS)`, and once the deadline has passed `cancel(true)` the remaining
futures instead of waiting on each.
##########
pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/SingleConnectionBrokerRequestHandler.java:
##########
@@ -117,6 +138,438 @@ public void shutDown() {
_brokerReduceService.shutDown();
}
+ /// Warms the scatter-gather path before readiness is granted: repeatedly
runs real probe queries until
+ /// they have run enough times (the `minIterations` depth floor) or the
budget expires. Broker-to-server
+ /// channels are opened lazily by the probes themselves.
+ ///
+ /// Probes go through [QueryRouter] rather than [#handleRequest],
deliberately bypassing access control,
+ /// query quota and the query log -- warmup then needs no synthetic identity
and pollutes no
+ /// customer-facing surface. Never throws, and always returns within the
configured budget.
+ @Override
+ public boolean warmUp(BrokerWarmupConfig config, long deadlineMs) {
+ // Stage 1: local, no-network warmup of the compile +
response-serialization paths, run before the
+ // network probe so they are warm even if no server is reachable.
+ warmUpLocal();
+ int concurrency = Math.max(1, config.concurrency());
+ ExecutorService probePool = Executors.newFixedThreadPool(concurrency,
+ new
ThreadFactoryBuilder().setNameFormat("broker-warmup-probe-%d").setDaemon(true).build());
+ try {
+ return config.hasCustomQuery() ? warmUpWithCustomQuery(config,
deadlineMs, probePool, concurrency)
+ : warmUpWithAutoSelect(config, deadlineMs, probePool, concurrency);
+ } catch (Exception e) {
+ LOGGER.warn("Broker warmup failed; proceeding without it", e);
+ return false;
+ } finally {
+ probePool.shutdownNow();
+ }
+ }
+
+ /// Stage 1: local, no-network warmup. The network probe goes through
[QueryRouter] and stops at the
+ /// gathered DataTables, so it never exercises the SQL compile path's cold
start nor the response build +
+ /// JSON serialization the first real query pays. This compiles a throwaway
query and serializes a small
+ /// synthetic [BrokerResponseNative], warming both. Never throws.
+ private void warmUpLocal() {
+ try {
+ CalciteSqlCompiler.compileToBrokerRequest("SELECT 1");
+ BrokerResponseNative response = new BrokerResponseNative();
+ response.setResultTable(new ResultTable(
+ new DataSchema(new String[]{"warmup"}, new
DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.LONG}),
+ Collections.singletonList(new Object[]{1L})));
+ response.setNumDocsScanned(1);
+ response.toJsonString();
+ } catch (Exception e) {
+ LOGGER.debug("Local (stage 1) warmup failed; continuing", e);
+ }
+ }
+
+ /// Fires the given probe tasks concurrently on the shared pool and returns
the latencies (ms) of the
+ /// probes that completed successfully. Empty means the whole round was
unproductive (nothing routable
+ /// yet / all failed); a failed probe simply does not contribute.
+ private List<Long> runConcurrentRound(List<Callable<Long>> tasks,
ExecutorService pool) {
+ List<Future<Long>> futures = new ArrayList<>(tasks.size());
+ for (Callable<Long> task : tasks) {
+ futures.add(pool.submit(task));
+ }
+ List<Long> latencies = new ArrayList<>(futures.size());
+ for (Future<Long> future : futures) {
+ try {
+ Long elapsedMs = future.get(WARMUP_PROBE_TIMEOUT_MS + 1000L,
TimeUnit.MILLISECONDS);
+ if (elapsedMs != null && elapsedMs >= 0) {
+ latencies.add(elapsedMs);
+ }
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ return latencies;
+ } catch (Exception e) {
+ // Individual probe failed/timed out; it just does not count toward
the round. Cancel it so a probe
+ // that outlived the round's get()-timeout does not keep running
behind the next round (no-op if it
+ // already completed).
+ future.cancel(true);
+ }
+ }
+ return latencies;
+ }
+
+ /// Sleeps [#WARMUP_FAILURE_BACKOFF_MS] after an unproductive probe round.
Returns `false` if interrupted
+ /// (shutdown) so the caller stops promptly and readiness opens.
+ private boolean warmupBackoff() {
+ try {
+ Thread.sleep(WARMUP_FAILURE_BACKOFF_MS);
+ return true;
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ return false;
+ }
+ }
+
+ /// Default probe path: probe `SELECT * FROM "<t>" LIMIT 1` over the
configured tables or a greedy
+ /// set-cover of tables covering every routable server. Exits once at least
`minIterations` probes have
+ /// completed successfully -- a depth floor that guarantees enough
invocations to drive the query path's
+ /// JIT to its top tier -- OR the budget expires. A latency target is
deliberately not used: 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.
+ private boolean warmUpWithAutoSelect(BrokerWarmupConfig config, long
deadlineMs, ExecutorService pool,
+ int concurrency) {
+ long successfulProbes = 0;
+ int rounds = 0;
+ // Select the probe tables once and reuse them across rounds. Only
re-select while the result is empty
+ // (routing not populated yet); recomputing the greedy set-cover every
round would repeat O(tables) work
+ // up to minIterations times on a large-table tenant, for a result that
does not change once non-empty.
+ List<String> tables = List.of();
+ boolean coverageReported = false;
+ // Monotonic across rounds so the round-robin actually advances through
every covered table even at
+ // concurrency 1; resetting per round would probe only the first
`concurrency` tables forever.
+ int probeSeq = 0;
+ while (System.currentTimeMillis() < deadlineMs &&
!Thread.currentThread().isInterrupted()) {
+ if (tables.isEmpty()) {
+ tables = config.hasCustomTables() ?
filterRoutableTables(_routingManager, config.tables())
+ : selectProbeTables(_routingManager, config.maxTables());
+ }
+ if (tables.isEmpty()) {
+ // Converged but nothing routable yet: back off and retry rather than
declaring the broker warm --
+ // an empty routing table here would otherwise make the gate a no-op
precisely on a cold broker.
+ if (!warmupBackoff()) {
+ return false;
+ }
+ continue;
+ }
+ if (!coverageReported) {
+ // Report once, now that the probe tables are settled, whether they
route to every server.
+ coverageReported = true;
+ reportServerCoverage(_routingManager, _brokerMetrics,
config.hasCustomTables(), tables);
+ }
+ long remainingMs = deadlineMs - System.currentTimeMillis();
+ if (remainingMs <= 0) {
+ break;
+ }
+ rounds++;
+ long probeTimeoutMs = Math.min(remainingMs, WARMUP_PROBE_TIMEOUT_MS);
+ // A batch of `concurrency` probes, round-robin over the covered tables
(advancing across rounds via
+ // probeSeq): exercises every server (coverage) AND real concurrency
(contention) at once.
+ List<Callable<Long>> tasks = new ArrayList<>(concurrency);
+ for (String tableNameWithType : roundRobinBatch(tables, probeSeq,
concurrency)) {
+ tasks.add(() -> probe(compileProbe(tableNameWithType),
probeTimeoutMs));
+ }
+ probeSeq += concurrency;
+ List<Long> latencies = runConcurrentRound(tasks, pool);
+ if (latencies.isEmpty()) {
+ if (!warmupBackoff()) {
+ return false;
+ }
+ continue;
+ }
+ successfulProbes += latencies.size();
+ if (successfulProbes >= config.minIterations()) {
+ LOGGER.info("Broker warmup completed after {} round(s) at concurrency
{}; {} probes (floor {})", rounds,
+ concurrency, successfulProbes, config.minIterations());
+ return true;
+ }
+ }
+ LOGGER.warn("Broker warmup budget expired after {} round(s) at concurrency
{}; {} probes (floor {}). Proceeding "
+ + "to serve traffic.", rounds, concurrency, successfulProbes,
config.minIterations());
+ return false;
+ }
+
+ /// Custom-query probe path: run the operator-supplied queries repeatedly,
all of them every round, and
+ /// exit on the same `minIterations` probe-count floor as the default probe,
or the budget. A probe-count
+ /// floor is a cost-independent warm signal; latency plateauing is not (a
trivial query reaches a stable
+ /// latency while still cold), so it is deliberately not used here. An
expensive query that cannot reach
+ /// the floor within the budget exits on the budget, having warmed as much
as the budget allowed.
+ private boolean warmUpWithCustomQuery(BrokerWarmupConfig config, long
deadlineMs, ExecutorService pool,
+ int concurrency) {
+ // Validate-compile each configured query once so a bad one fails fast
(and never busy-spins), keeping
+ // only those that compile. Each probe below compiles its own fresh
request so nothing mutable is shared
+ // across the concurrent tasks.
+ List<String> queries = new ArrayList<>(config.queries().size());
+ for (String query : config.queries()) {
+ try {
+ CalciteSqlCompiler.compileToBrokerRequest(query);
+ queries.add(query);
+ } catch (Exception e) {
+ LOGGER.warn("Broker warmup query failed to compile; skipping it: {}",
query, e);
+ }
+ }
+ if (queries.isEmpty()) {
+ LOGGER.warn("Broker warmup: no configured query compiled; skipping
warmup");
+ return false;
+ }
+ long successfulProbes = 0;
+ int rounds = 0;
+ while (System.currentTimeMillis() < deadlineMs &&
!Thread.currentThread().isInterrupted()) {
+ long remainingMs = deadlineMs - System.currentTimeMillis();
+ if (remainingMs <= 0) {
+ break;
+ }
+ rounds++;
+ long probeTimeoutMs = Math.min(remainingMs, WARMUP_PROBE_TIMEOUT_MS);
+ // Run every configured query each round (each `concurrency` times, to
warm the concurrency step too).
+ List<Callable<Long>> tasks = new ArrayList<>(queries.size() *
concurrency);
+ for (String query : queries) {
+ for (int i = 0; i < concurrency; i++) {
+ tasks.add(() ->
probe(CalciteSqlCompiler.compileToBrokerRequest(query), probeTimeoutMs));
Review Comment:
**Correctness — same root cause as the `future.get()` comment above.**
`probeTimeoutMs` is computed once at round start, but these tasks queue
behind a pool of only `concurrency` threads (`tasks.size() == queries.size() *
concurrency`). A task that starts 20s into the round still runs with the
`probeTimeoutMs` computed at t=0.
Either size the pool to `tasks.size()`, or compute the per-probe timeout
inside the lambda from `deadlineMs` and skip the probe when it is already `<=
0`.
##########
pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/SingleConnectionBrokerRequestHandler.java:
##########
@@ -117,6 +138,438 @@ public void shutDown() {
_brokerReduceService.shutDown();
}
+ /// Warms the scatter-gather path before readiness is granted: repeatedly
runs real probe queries until
+ /// they have run enough times (the `minIterations` depth floor) or the
budget expires. Broker-to-server
+ /// channels are opened lazily by the probes themselves.
+ ///
+ /// Probes go through [QueryRouter] rather than [#handleRequest],
deliberately bypassing access control,
+ /// query quota and the query log -- warmup then needs no synthetic identity
and pollutes no
+ /// customer-facing surface. Never throws, and always returns within the
configured budget.
+ @Override
+ public boolean warmUp(BrokerWarmupConfig config, long deadlineMs) {
+ // Stage 1: local, no-network warmup of the compile +
response-serialization paths, run before the
+ // network probe so they are warm even if no server is reachable.
+ warmUpLocal();
+ int concurrency = Math.max(1, config.concurrency());
+ ExecutorService probePool = Executors.newFixedThreadPool(concurrency,
+ new
ThreadFactoryBuilder().setNameFormat("broker-warmup-probe-%d").setDaemon(true).build());
+ try {
+ return config.hasCustomQuery() ? warmUpWithCustomQuery(config,
deadlineMs, probePool, concurrency)
+ : warmUpWithAutoSelect(config, deadlineMs, probePool, concurrency);
+ } catch (Exception e) {
+ LOGGER.warn("Broker warmup failed; proceeding without it", e);
+ return false;
+ } finally {
+ probePool.shutdownNow();
+ }
+ }
+
+ /// Stage 1: local, no-network warmup. The network probe goes through
[QueryRouter] and stops at the
+ /// gathered DataTables, so it never exercises the SQL compile path's cold
start nor the response build +
+ /// JSON serialization the first real query pays. This compiles a throwaway
query and serializes a small
+ /// synthetic [BrokerResponseNative], warming both. Never throws.
+ private void warmUpLocal() {
+ try {
+ CalciteSqlCompiler.compileToBrokerRequest("SELECT 1");
+ BrokerResponseNative response = new BrokerResponseNative();
+ response.setResultTable(new ResultTable(
+ new DataSchema(new String[]{"warmup"}, new
DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.LONG}),
+ Collections.singletonList(new Object[]{1L})));
+ response.setNumDocsScanned(1);
+ response.toJsonString();
+ } catch (Exception e) {
+ LOGGER.debug("Local (stage 1) warmup failed; continuing", e);
+ }
+ }
+
+ /// Fires the given probe tasks concurrently on the shared pool and returns
the latencies (ms) of the
+ /// probes that completed successfully. Empty means the whole round was
unproductive (nothing routable
+ /// yet / all failed); a failed probe simply does not contribute.
+ private List<Long> runConcurrentRound(List<Callable<Long>> tasks,
ExecutorService pool) {
+ List<Future<Long>> futures = new ArrayList<>(tasks.size());
+ for (Callable<Long> task : tasks) {
+ futures.add(pool.submit(task));
+ }
+ List<Long> latencies = new ArrayList<>(futures.size());
+ for (Future<Long> future : futures) {
+ try {
+ Long elapsedMs = future.get(WARMUP_PROBE_TIMEOUT_MS + 1000L,
TimeUnit.MILLISECONDS);
+ if (elapsedMs != null && elapsedMs >= 0) {
+ latencies.add(elapsedMs);
+ }
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ return latencies;
+ } catch (Exception e) {
+ // Individual probe failed/timed out; it just does not count toward
the round. Cancel it so a probe
+ // that outlived the round's get()-timeout does not keep running
behind the next round (no-op if it
+ // already completed).
+ future.cancel(true);
+ }
+ }
+ return latencies;
+ }
+
+ /// Sleeps [#WARMUP_FAILURE_BACKOFF_MS] after an unproductive probe round.
Returns `false` if interrupted
+ /// (shutdown) so the caller stops promptly and readiness opens.
+ private boolean warmupBackoff() {
+ try {
+ Thread.sleep(WARMUP_FAILURE_BACKOFF_MS);
+ return true;
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ return false;
+ }
+ }
+
+ /// Default probe path: probe `SELECT * FROM "<t>" LIMIT 1` over the
configured tables or a greedy
+ /// set-cover of tables covering every routable server. Exits once at least
`minIterations` probes have
+ /// completed successfully -- a depth floor that guarantees enough
invocations to drive the query path's
+ /// JIT to its top tier -- OR the budget expires. A latency target is
deliberately not used: 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.
+ private boolean warmUpWithAutoSelect(BrokerWarmupConfig config, long
deadlineMs, ExecutorService pool,
+ int concurrency) {
+ long successfulProbes = 0;
+ int rounds = 0;
+ // Select the probe tables once and reuse them across rounds. Only
re-select while the result is empty
+ // (routing not populated yet); recomputing the greedy set-cover every
round would repeat O(tables) work
+ // up to minIterations times on a large-table tenant, for a result that
does not change once non-empty.
+ List<String> tables = List.of();
+ boolean coverageReported = false;
+ // Monotonic across rounds so the round-robin actually advances through
every covered table even at
+ // concurrency 1; resetting per round would probe only the first
`concurrency` tables forever.
+ int probeSeq = 0;
+ while (System.currentTimeMillis() < deadlineMs &&
!Thread.currentThread().isInterrupted()) {
+ if (tables.isEmpty()) {
+ tables = config.hasCustomTables() ?
filterRoutableTables(_routingManager, config.tables())
+ : selectProbeTables(_routingManager, config.maxTables());
+ }
+ if (tables.isEmpty()) {
+ // Converged but nothing routable yet: back off and retry rather than
declaring the broker warm --
+ // an empty routing table here would otherwise make the gate a no-op
precisely on a cold broker.
+ if (!warmupBackoff()) {
+ return false;
+ }
+ continue;
+ }
+ if (!coverageReported) {
+ // Report once, now that the probe tables are settled, whether they
route to every server.
+ coverageReported = true;
+ reportServerCoverage(_routingManager, _brokerMetrics,
config.hasCustomTables(), tables);
+ }
+ long remainingMs = deadlineMs - System.currentTimeMillis();
+ if (remainingMs <= 0) {
+ break;
+ }
+ rounds++;
+ long probeTimeoutMs = Math.min(remainingMs, WARMUP_PROBE_TIMEOUT_MS);
+ // A batch of `concurrency` probes, round-robin over the covered tables
(advancing across rounds via
+ // probeSeq): exercises every server (coverage) AND real concurrency
(contention) at once.
+ List<Callable<Long>> tasks = new ArrayList<>(concurrency);
+ for (String tableNameWithType : roundRobinBatch(tables, probeSeq,
concurrency)) {
+ tasks.add(() -> probe(compileProbe(tableNameWithType),
probeTimeoutMs));
+ }
+ probeSeq += concurrency;
+ List<Long> latencies = runConcurrentRound(tasks, pool);
+ if (latencies.isEmpty()) {
+ if (!warmupBackoff()) {
+ return false;
+ }
+ continue;
+ }
+ successfulProbes += latencies.size();
+ if (successfulProbes >= config.minIterations()) {
+ LOGGER.info("Broker warmup completed after {} round(s) at concurrency
{}; {} probes (floor {})", rounds,
+ concurrency, successfulProbes, config.minIterations());
+ return true;
+ }
+ }
+ LOGGER.warn("Broker warmup budget expired after {} round(s) at concurrency
{}; {} probes (floor {}). Proceeding "
+ + "to serve traffic.", rounds, concurrency, successfulProbes,
config.minIterations());
+ return false;
+ }
+
+ /// Custom-query probe path: run the operator-supplied queries repeatedly,
all of them every round, and
+ /// exit on the same `minIterations` probe-count floor as the default probe,
or the budget. A probe-count
+ /// floor is a cost-independent warm signal; latency plateauing is not (a
trivial query reaches a stable
+ /// latency while still cold), so it is deliberately not used here. An
expensive query that cannot reach
+ /// the floor within the budget exits on the budget, having warmed as much
as the budget allowed.
+ private boolean warmUpWithCustomQuery(BrokerWarmupConfig config, long
deadlineMs, ExecutorService pool,
+ int concurrency) {
+ // Validate-compile each configured query once so a bad one fails fast
(and never busy-spins), keeping
+ // only those that compile. Each probe below compiles its own fresh
request so nothing mutable is shared
+ // across the concurrent tasks.
+ List<String> queries = new ArrayList<>(config.queries().size());
+ for (String query : config.queries()) {
+ try {
+ CalciteSqlCompiler.compileToBrokerRequest(query);
+ queries.add(query);
+ } catch (Exception e) {
+ LOGGER.warn("Broker warmup query failed to compile; skipping it: {}",
query, e);
+ }
+ }
+ if (queries.isEmpty()) {
+ LOGGER.warn("Broker warmup: no configured query compiled; skipping
warmup");
+ return false;
+ }
+ long successfulProbes = 0;
+ int rounds = 0;
+ while (System.currentTimeMillis() < deadlineMs &&
!Thread.currentThread().isInterrupted()) {
+ long remainingMs = deadlineMs - System.currentTimeMillis();
+ if (remainingMs <= 0) {
+ break;
+ }
+ rounds++;
+ long probeTimeoutMs = Math.min(remainingMs, WARMUP_PROBE_TIMEOUT_MS);
+ // Run every configured query each round (each `concurrency` times, to
warm the concurrency step too).
+ List<Callable<Long>> tasks = new ArrayList<>(queries.size() *
concurrency);
+ for (String query : queries) {
+ for (int i = 0; i < concurrency; i++) {
+ tasks.add(() ->
probe(CalciteSqlCompiler.compileToBrokerRequest(query), probeTimeoutMs));
+ }
+ }
+ List<Long> latencies = runConcurrentRound(tasks, pool);
+ if (latencies.isEmpty()) {
+ if (!warmupBackoff()) {
+ return false;
+ }
+ continue;
+ }
+ successfulProbes += latencies.size();
+ // Same depth floor as the default probe: a probe-count floor is a
cost-independent warm signal (it
+ // drives the query path's JIT to its top tier), whereas latency
plateauing is not -- a trivial query
+ // reaches a stable latency while still cold. An expensive query that
cannot reach the floor within the
+ // budget simply exits on the budget, having warmed as much as the
budget allowed.
+ if (successfulProbes >= config.minIterations()) {
+ LOGGER.info("Broker warmup (custom query) completed after {} round(s)
at concurrency {}; {} probes "
+ + "(floor {})", rounds, concurrency, successfulProbes,
config.minIterations());
+ return true;
+ }
+ }
+ LOGGER.warn("Broker warmup (custom query) budget expired after {} round(s)
at concurrency {}; {} probes "
+ + "(floor {}). Proceeding to serve traffic.", rounds, concurrency,
successfulProbes, config.minIterations());
+ return false;
+ }
+
+ /// Compiles the default probe query for a single physical table.
+ private BrokerRequest compileProbe(String tableNameWithType) {
+ return CalciteSqlCompiler.compileToBrokerRequest("SELECT * FROM \"" +
tableNameWithType + "\" LIMIT 1");
+ }
+
+ /// Issues one probe query through [QueryRouter] and returns its wall-clock
duration in ms, or -1 on
+ /// failure. Builds the route the same way the normal path does, minus
auth/quota/logging.
+ private long probe(BrokerRequest brokerRequest, long timeoutMs) {
+ String tableName = brokerRequest.getQuerySource().getTableName();
+ try {
+ String rawTableName = TableNameBuilder.extractRawTableName(tableName);
+ TableType tableType =
TableNameBuilder.getTableTypeFromTableName(tableName);
+ long requestId = _requestIdGenerator.get();
+ TableRouteInfo routeInfo =
+ _implicitHybridTableRouteProvider.getTableRouteInfo(rawTableName,
_tableCache, _routingManager);
+ if (!(routeInfo instanceof ImplicitHybridTableRouteInfo) ||
!routeInfo.isExists()) {
+ return -1;
+ }
+ // A type-suffixed table (the default probe) targets exactly one type; a
raw name (possible with a
+ // custom query) targets whichever types actually exist. Each leg's
request must carry the
+ // type-suffixed table name -- that is what routing resolves against --
so build a per-leg request
+ // from the base query (as the normal path does) rather than reusing the
raw-named request, which
+ // would route to no server and yield an empty request map.
+ BrokerRequest offlineBrokerRequest = shouldRouteOffline(tableType,
routeInfo)
+ ? typedRequest(brokerRequest, routeInfo.getOfflineTableName()) :
null;
+ BrokerRequest realtimeBrokerRequest = shouldRouteRealtime(tableType,
routeInfo)
+ ? typedRequest(brokerRequest, routeInfo.getRealtimeTableName()) :
null;
+ if (offlineBrokerRequest == null && realtimeBrokerRequest == null) {
+ // Neither leg is routable right now (e.g. a hybrid table probed via
its offline name before the
+ // time boundary exists). submitQuery asserts at least one leg is
non-null, so skip cleanly rather
+ // than build an empty route -- the table is retried on the next round
once its leg appears.
+ return -1;
+ }
+ ImplicitHybridTableRouteInfo hybridRouteInfo =
(ImplicitHybridTableRouteInfo) routeInfo;
+ hybridRouteInfo.setOfflineBrokerRequest(offlineBrokerRequest);
+ hybridRouteInfo.setRealtimeBrokerRequest(realtimeBrokerRequest);
+ _implicitHybridTableRouteProvider.calculateRoutes(routeInfo,
_routingManager, offlineBrokerRequest,
+ realtimeBrokerRequest, requestId);
+ if (!routeInfo.isRouteExists()) {
+ return -1;
+ }
+ long startMs = System.currentTimeMillis();
+ // Scatter/gather (+ DataTable deserialize on receive) is the dominant
cost and does not need a
+ // QueryThreadContext; run it directly so its warmth always counts
toward the probe latency.
+ AsyncQueryResponse response = _queryRouter.submitQuery(requestId,
rawTableName, routeInfo, timeoutMs);
+ Map<ServerRoutingInstance, ServerResponse> finalResponses =
response.getFinalResponses();
+ // Best-effort: also warm the broker reduce path (result discarded) so
the first real query of this
+ // shape does not pay it. Isolated in its own try -- reduceOnDataTable
requires a QueryThreadContext
+ // and could fail for edge cases; a reduce failure must never fail the
probe, since the
+ // scatter/gather/deserialize warmth has already happened.
+ try (QueryThreadContext ignore = QueryThreadContext.open(
+ new QueryExecutionContext(QueryExecutionContext.QueryType.SSE,
requestId, Long.toString(requestId),
+ "warmup", startMs, Long.MAX_VALUE, Long.MAX_VALUE, _brokerId,
_brokerId, ""), _threadAccountant)) {
+ Map<ServerRoutingInstance, DataTable> dataTableMap = new
HashMap<>(finalResponses.size());
+ for (Map.Entry<ServerRoutingInstance, ServerResponse> entry :
finalResponses.entrySet()) {
+ DataTable dataTable = entry.getValue().getDataTable();
+ if (dataTable != null) {
+ dataTableMap.put(entry.getKey(), dataTable);
+ }
+ }
+ if (!dataTableMap.isEmpty()) {
+ _brokerReduceService.reduceOnDataTable(brokerRequest, brokerRequest,
dataTableMap, timeoutMs,
+ _brokerMetrics);
+ }
+ } catch (Exception reduceEx) {
+ LOGGER.debug("Warmup reduce step failed (continuing); scatter/gather
already warmed", reduceEx);
+ }
+ return System.currentTimeMillis() - startMs;
+ } catch (Exception e) {
+ if (e instanceof InterruptedException) {
+ // Shutdown interrupted this probe (via shutdownNow); restore the flag
so the pool worker unwinds.
+ Thread.currentThread().interrupt();
+ }
+ LOGGER.debug("Warmup probe failed for query on table: {}", tableName, e);
+ return -1;
+ }
+ }
+
+ /// Builds a per-leg probe request carrying the type-suffixed table name
that routing resolves against,
+ /// from the base (possibly raw-named) request. Deep-copies so the base
request is not mutated.
+ private static BrokerRequest typedRequest(BrokerRequest base, String
tableNameWithType) {
+ PinotQuery pinotQuery = base.getPinotQuery().deepCopy();
+ pinotQuery.getDataSource().setTableName(tableNameWithType);
+ return CalciteSqlCompiler.convertToBrokerRequest(pinotQuery);
+ }
+
+ /// Whether a probe routes the offline leg: the query targets offline (a
type-suffixed offline name) or is
+ /// untyped (a raw name, `tableType` null), AND an offline table exists.
Static so the offline / realtime /
+ /// hybrid decision is unit-testable against a mocked route without a live
cluster.
+ @VisibleForTesting
+ static boolean shouldRouteOffline(@Nullable TableType tableType,
TableRouteInfo routeInfo) {
+ return tableType != TableType.REALTIME && routeInfo.hasOffline();
+ }
+
+ /// Whether a probe routes the realtime leg: the query targets realtime or
is untyped, AND a realtime
+ /// table exists.
+ @VisibleForTesting
+ static boolean shouldRouteRealtime(@Nullable TableType tableType,
TableRouteInfo routeInfo) {
+ return tableType != TableType.OFFLINE && routeInfo.hasRealtime();
+ }
+
+ /// Keeps only the configured tables that this broker actually routes.
Static for unit testing.
+ @VisibleForTesting
+ static List<String> filterRoutableTables(RoutingManager routingManager,
List<String> tables) {
Review Comment:
**Correctness — a raw (type-less) table name in `warmup.tables` is always
dropped, but the config comment documents the list as accepting "raw or with
type".**
`routingManager.getRoutableTables()` returns type-suffixed names only
(`myTable_OFFLINE` / `myTable_REALTIME`), so `routable.contains(table)` can
never match a raw name. An operator who follows the documented config gets
`configured table X is not routable on this broker; skipping` for every entry,
an empty table list, and then a full budget of 100ms backoff loops — i.e.
silently no warmup at all, while the feature reports as enabled.
Either expand a raw name to both types before the containment check, or
change the comment in `CommonConstants` to say type-suffixed names are
required. A unit test for the raw-name case would pin whichever you pick.
##########
pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/SingleConnectionBrokerRequestHandler.java:
##########
@@ -117,6 +138,438 @@ public void shutDown() {
_brokerReduceService.shutDown();
}
+ /// Warms the scatter-gather path before readiness is granted: repeatedly
runs real probe queries until
+ /// they have run enough times (the `minIterations` depth floor) or the
budget expires. Broker-to-server
+ /// channels are opened lazily by the probes themselves.
+ ///
+ /// Probes go through [QueryRouter] rather than [#handleRequest],
deliberately bypassing access control,
+ /// query quota and the query log -- warmup then needs no synthetic identity
and pollutes no
+ /// customer-facing surface. Never throws, and always returns within the
configured budget.
+ @Override
+ public boolean warmUp(BrokerWarmupConfig config, long deadlineMs) {
+ // Stage 1: local, no-network warmup of the compile +
response-serialization paths, run before the
+ // network probe so they are warm even if no server is reachable.
+ warmUpLocal();
+ int concurrency = Math.max(1, config.concurrency());
+ ExecutorService probePool = Executors.newFixedThreadPool(concurrency,
+ new
ThreadFactoryBuilder().setNameFormat("broker-warmup-probe-%d").setDaemon(true).build());
+ try {
+ return config.hasCustomQuery() ? warmUpWithCustomQuery(config,
deadlineMs, probePool, concurrency)
+ : warmUpWithAutoSelect(config, deadlineMs, probePool, concurrency);
+ } catch (Exception e) {
+ LOGGER.warn("Broker warmup failed; proceeding without it", e);
+ return false;
+ } finally {
+ probePool.shutdownNow();
+ }
+ }
+
+ /// Stage 1: local, no-network warmup. The network probe goes through
[QueryRouter] and stops at the
+ /// gathered DataTables, so it never exercises the SQL compile path's cold
start nor the response build +
+ /// JSON serialization the first real query pays. This compiles a throwaway
query and serializes a small
+ /// synthetic [BrokerResponseNative], warming both. Never throws.
+ private void warmUpLocal() {
+ try {
+ CalciteSqlCompiler.compileToBrokerRequest("SELECT 1");
+ BrokerResponseNative response = new BrokerResponseNative();
+ response.setResultTable(new ResultTable(
+ new DataSchema(new String[]{"warmup"}, new
DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.LONG}),
+ Collections.singletonList(new Object[]{1L})));
+ response.setNumDocsScanned(1);
+ response.toJsonString();
+ } catch (Exception e) {
+ LOGGER.debug("Local (stage 1) warmup failed; continuing", e);
+ }
+ }
+
+ /// Fires the given probe tasks concurrently on the shared pool and returns
the latencies (ms) of the
+ /// probes that completed successfully. Empty means the whole round was
unproductive (nothing routable
+ /// yet / all failed); a failed probe simply does not contribute.
+ private List<Long> runConcurrentRound(List<Callable<Long>> tasks,
ExecutorService pool) {
+ List<Future<Long>> futures = new ArrayList<>(tasks.size());
+ for (Callable<Long> task : tasks) {
+ futures.add(pool.submit(task));
+ }
+ List<Long> latencies = new ArrayList<>(futures.size());
+ for (Future<Long> future : futures) {
+ try {
+ Long elapsedMs = future.get(WARMUP_PROBE_TIMEOUT_MS + 1000L,
TimeUnit.MILLISECONDS);
+ if (elapsedMs != null && elapsedMs >= 0) {
+ latencies.add(elapsedMs);
+ }
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ return latencies;
+ } catch (Exception e) {
+ // Individual probe failed/timed out; it just does not count toward
the round. Cancel it so a probe
+ // that outlived the round's get()-timeout does not keep running
behind the next round (no-op if it
+ // already completed).
+ future.cancel(true);
+ }
+ }
+ return latencies;
+ }
+
+ /// Sleeps [#WARMUP_FAILURE_BACKOFF_MS] after an unproductive probe round.
Returns `false` if interrupted
+ /// (shutdown) so the caller stops promptly and readiness opens.
+ private boolean warmupBackoff() {
+ try {
+ Thread.sleep(WARMUP_FAILURE_BACKOFF_MS);
+ return true;
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ return false;
+ }
+ }
+
+ /// Default probe path: probe `SELECT * FROM "<t>" LIMIT 1` over the
configured tables or a greedy
+ /// set-cover of tables covering every routable server. Exits once at least
`minIterations` probes have
+ /// completed successfully -- a depth floor that guarantees enough
invocations to drive the query path's
+ /// JIT to its top tier -- OR the budget expires. A latency target is
deliberately not used: 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.
+ private boolean warmUpWithAutoSelect(BrokerWarmupConfig config, long
deadlineMs, ExecutorService pool,
+ int concurrency) {
+ long successfulProbes = 0;
+ int rounds = 0;
+ // Select the probe tables once and reuse them across rounds. Only
re-select while the result is empty
+ // (routing not populated yet); recomputing the greedy set-cover every
round would repeat O(tables) work
+ // up to minIterations times on a large-table tenant, for a result that
does not change once non-empty.
+ List<String> tables = List.of();
+ boolean coverageReported = false;
+ // Monotonic across rounds so the round-robin actually advances through
every covered table even at
+ // concurrency 1; resetting per round would probe only the first
`concurrency` tables forever.
+ int probeSeq = 0;
+ while (System.currentTimeMillis() < deadlineMs &&
!Thread.currentThread().isInterrupted()) {
+ if (tables.isEmpty()) {
+ tables = config.hasCustomTables() ?
filterRoutableTables(_routingManager, config.tables())
+ : selectProbeTables(_routingManager, config.maxTables());
+ }
+ if (tables.isEmpty()) {
+ // Converged but nothing routable yet: back off and retry rather than
declaring the broker warm --
+ // an empty routing table here would otherwise make the gate a no-op
precisely on a cold broker.
+ if (!warmupBackoff()) {
+ return false;
+ }
+ continue;
+ }
+ if (!coverageReported) {
+ // Report once, now that the probe tables are settled, whether they
route to every server.
+ coverageReported = true;
+ reportServerCoverage(_routingManager, _brokerMetrics,
config.hasCustomTables(), tables);
+ }
+ long remainingMs = deadlineMs - System.currentTimeMillis();
+ if (remainingMs <= 0) {
+ break;
+ }
+ rounds++;
+ long probeTimeoutMs = Math.min(remainingMs, WARMUP_PROBE_TIMEOUT_MS);
+ // A batch of `concurrency` probes, round-robin over the covered tables
(advancing across rounds via
+ // probeSeq): exercises every server (coverage) AND real concurrency
(contention) at once.
+ List<Callable<Long>> tasks = new ArrayList<>(concurrency);
+ for (String tableNameWithType : roundRobinBatch(tables, probeSeq,
concurrency)) {
+ tasks.add(() -> probe(compileProbe(tableNameWithType),
probeTimeoutMs));
+ }
+ probeSeq += concurrency;
+ List<Long> latencies = runConcurrentRound(tasks, pool);
+ if (latencies.isEmpty()) {
+ if (!warmupBackoff()) {
+ return false;
+ }
+ continue;
+ }
+ successfulProbes += latencies.size();
+ if (successfulProbes >= config.minIterations()) {
+ LOGGER.info("Broker warmup completed after {} round(s) at concurrency
{}; {} probes (floor {})", rounds,
+ concurrency, successfulProbes, config.minIterations());
+ return true;
+ }
+ }
+ LOGGER.warn("Broker warmup budget expired after {} round(s) at concurrency
{}; {} probes (floor {}). Proceeding "
+ + "to serve traffic.", rounds, concurrency, successfulProbes,
config.minIterations());
+ return false;
+ }
+
+ /// Custom-query probe path: run the operator-supplied queries repeatedly,
all of them every round, and
+ /// exit on the same `minIterations` probe-count floor as the default probe,
or the budget. A probe-count
+ /// floor is a cost-independent warm signal; latency plateauing is not (a
trivial query reaches a stable
+ /// latency while still cold), so it is deliberately not used here. An
expensive query that cannot reach
+ /// the floor within the budget exits on the budget, having warmed as much
as the budget allowed.
+ private boolean warmUpWithCustomQuery(BrokerWarmupConfig config, long
deadlineMs, ExecutorService pool,
+ int concurrency) {
+ // Validate-compile each configured query once so a bad one fails fast
(and never busy-spins), keeping
+ // only those that compile. Each probe below compiles its own fresh
request so nothing mutable is shared
+ // across the concurrent tasks.
+ List<String> queries = new ArrayList<>(config.queries().size());
+ for (String query : config.queries()) {
+ try {
+ CalciteSqlCompiler.compileToBrokerRequest(query);
+ queries.add(query);
+ } catch (Exception e) {
+ LOGGER.warn("Broker warmup query failed to compile; skipping it: {}",
query, e);
+ }
+ }
+ if (queries.isEmpty()) {
+ LOGGER.warn("Broker warmup: no configured query compiled; skipping
warmup");
+ return false;
+ }
+ long successfulProbes = 0;
+ int rounds = 0;
+ while (System.currentTimeMillis() < deadlineMs &&
!Thread.currentThread().isInterrupted()) {
+ long remainingMs = deadlineMs - System.currentTimeMillis();
+ if (remainingMs <= 0) {
+ break;
+ }
+ rounds++;
+ long probeTimeoutMs = Math.min(remainingMs, WARMUP_PROBE_TIMEOUT_MS);
+ // Run every configured query each round (each `concurrency` times, to
warm the concurrency step too).
+ List<Callable<Long>> tasks = new ArrayList<>(queries.size() *
concurrency);
+ for (String query : queries) {
+ for (int i = 0; i < concurrency; i++) {
+ tasks.add(() ->
probe(CalciteSqlCompiler.compileToBrokerRequest(query), probeTimeoutMs));
+ }
+ }
+ List<Long> latencies = runConcurrentRound(tasks, pool);
+ if (latencies.isEmpty()) {
+ if (!warmupBackoff()) {
+ return false;
+ }
+ continue;
+ }
+ successfulProbes += latencies.size();
+ // Same depth floor as the default probe: a probe-count floor is a
cost-independent warm signal (it
+ // drives the query path's JIT to its top tier), whereas latency
plateauing is not -- a trivial query
+ // reaches a stable latency while still cold. An expensive query that
cannot reach the floor within the
+ // budget simply exits on the budget, having warmed as much as the
budget allowed.
+ if (successfulProbes >= config.minIterations()) {
+ LOGGER.info("Broker warmup (custom query) completed after {} round(s)
at concurrency {}; {} probes "
+ + "(floor {})", rounds, concurrency, successfulProbes,
config.minIterations());
+ return true;
+ }
+ }
+ LOGGER.warn("Broker warmup (custom query) budget expired after {} round(s)
at concurrency {}; {} probes "
+ + "(floor {}). Proceeding to serve traffic.", rounds, concurrency,
successfulProbes, config.minIterations());
+ return false;
+ }
+
+ /// Compiles the default probe query for a single physical table.
+ private BrokerRequest compileProbe(String tableNameWithType) {
+ return CalciteSqlCompiler.compileToBrokerRequest("SELECT * FROM \"" +
tableNameWithType + "\" LIMIT 1");
+ }
+
+ /// Issues one probe query through [QueryRouter] and returns its wall-clock
duration in ms, or -1 on
+ /// failure. Builds the route the same way the normal path does, minus
auth/quota/logging.
+ private long probe(BrokerRequest brokerRequest, long timeoutMs) {
+ String tableName = brokerRequest.getQuerySource().getTableName();
+ try {
+ String rawTableName = TableNameBuilder.extractRawTableName(tableName);
+ TableType tableType =
TableNameBuilder.getTableTypeFromTableName(tableName);
+ long requestId = _requestIdGenerator.get();
+ TableRouteInfo routeInfo =
+ _implicitHybridTableRouteProvider.getTableRouteInfo(rawTableName,
_tableCache, _routingManager);
+ if (!(routeInfo instanceof ImplicitHybridTableRouteInfo) ||
!routeInfo.isExists()) {
+ return -1;
+ }
+ // A type-suffixed table (the default probe) targets exactly one type; a
raw name (possible with a
+ // custom query) targets whichever types actually exist. Each leg's
request must carry the
+ // type-suffixed table name -- that is what routing resolves against --
so build a per-leg request
+ // from the base query (as the normal path does) rather than reusing the
raw-named request, which
+ // would route to no server and yield an empty request map.
+ BrokerRequest offlineBrokerRequest = shouldRouteOffline(tableType,
routeInfo)
+ ? typedRequest(brokerRequest, routeInfo.getOfflineTableName()) :
null;
+ BrokerRequest realtimeBrokerRequest = shouldRouteRealtime(tableType,
routeInfo)
+ ? typedRequest(brokerRequest, routeInfo.getRealtimeTableName()) :
null;
+ if (offlineBrokerRequest == null && realtimeBrokerRequest == null) {
+ // Neither leg is routable right now (e.g. a hybrid table probed via
its offline name before the
+ // time boundary exists). submitQuery asserts at least one leg is
non-null, so skip cleanly rather
+ // than build an empty route -- the table is retried on the next round
once its leg appears.
+ return -1;
+ }
+ ImplicitHybridTableRouteInfo hybridRouteInfo =
(ImplicitHybridTableRouteInfo) routeInfo;
+ hybridRouteInfo.setOfflineBrokerRequest(offlineBrokerRequest);
+ hybridRouteInfo.setRealtimeBrokerRequest(realtimeBrokerRequest);
+ _implicitHybridTableRouteProvider.calculateRoutes(routeInfo,
_routingManager, offlineBrokerRequest,
+ realtimeBrokerRequest, requestId);
+ if (!routeInfo.isRouteExists()) {
+ return -1;
+ }
+ long startMs = System.currentTimeMillis();
+ // Scatter/gather (+ DataTable deserialize on receive) is the dominant
cost and does not need a
+ // QueryThreadContext; run it directly so its warmth always counts
toward the probe latency.
+ AsyncQueryResponse response = _queryRouter.submitQuery(requestId,
rawTableName, routeInfo, timeoutMs);
+ Map<ServerRoutingInstance, ServerResponse> finalResponses =
response.getFinalResponses();
Review Comment:
**Risk — `getFinalResponses()` also feeds `ServerRoutingStatsManager`, so
warmup seeds adaptive server selection.**
`AsyncQueryResponse.getFinalResponses()` unconditionally calls
`_serverRoutingStatsManager.recordStatsUponResponseArrival()` for every server,
and when a server did not respond or returned an exception it records `latency
= _timeoutMs` (up to 5s here).
Submission and arrival both live inside `AsyncQueryResponse`, so the
in-flight counters stay balanced — no leak. But the EMA is not free:
- 1000 probes of `SELECT * ... LIMIT 1` seed the EMA with latencies that are
not representative of real traffic.
- A server that is slow or unreachable *during warmup* collects up to 1000 x
5000ms samples, so with adaptive server selection enabled it starts real
traffic heavily penalised — the opposite of the intent.
Worth a line in the description at minimum. Have you checked against the
adaptive-selection EMA alpha whether the bias has decayed by the time traffic
arrives?
##########
pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/SingleConnectionBrokerRequestHandler.java:
##########
@@ -117,6 +138,438 @@ public void shutDown() {
_brokerReduceService.shutDown();
}
+ /// Warms the scatter-gather path before readiness is granted: repeatedly
runs real probe queries until
+ /// they have run enough times (the `minIterations` depth floor) or the
budget expires. Broker-to-server
+ /// channels are opened lazily by the probes themselves.
+ ///
+ /// Probes go through [QueryRouter] rather than [#handleRequest],
deliberately bypassing access control,
+ /// query quota and the query log -- warmup then needs no synthetic identity
and pollutes no
+ /// customer-facing surface. Never throws, and always returns within the
configured budget.
+ @Override
+ public boolean warmUp(BrokerWarmupConfig config, long deadlineMs) {
+ // Stage 1: local, no-network warmup of the compile +
response-serialization paths, run before the
+ // network probe so they are warm even if no server is reachable.
+ warmUpLocal();
+ int concurrency = Math.max(1, config.concurrency());
+ ExecutorService probePool = Executors.newFixedThreadPool(concurrency,
+ new
ThreadFactoryBuilder().setNameFormat("broker-warmup-probe-%d").setDaemon(true).build());
+ try {
+ return config.hasCustomQuery() ? warmUpWithCustomQuery(config,
deadlineMs, probePool, concurrency)
+ : warmUpWithAutoSelect(config, deadlineMs, probePool, concurrency);
+ } catch (Exception e) {
+ LOGGER.warn("Broker warmup failed; proceeding without it", e);
+ return false;
+ } finally {
+ probePool.shutdownNow();
+ }
+ }
+
+ /// Stage 1: local, no-network warmup. The network probe goes through
[QueryRouter] and stops at the
+ /// gathered DataTables, so it never exercises the SQL compile path's cold
start nor the response build +
+ /// JSON serialization the first real query pays. This compiles a throwaway
query and serializes a small
+ /// synthetic [BrokerResponseNative], warming both. Never throws.
+ private void warmUpLocal() {
+ try {
+ CalciteSqlCompiler.compileToBrokerRequest("SELECT 1");
+ BrokerResponseNative response = new BrokerResponseNative();
+ response.setResultTable(new ResultTable(
+ new DataSchema(new String[]{"warmup"}, new
DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.LONG}),
+ Collections.singletonList(new Object[]{1L})));
+ response.setNumDocsScanned(1);
+ response.toJsonString();
+ } catch (Exception e) {
+ LOGGER.debug("Local (stage 1) warmup failed; continuing", e);
+ }
+ }
+
+ /// Fires the given probe tasks concurrently on the shared pool and returns
the latencies (ms) of the
+ /// probes that completed successfully. Empty means the whole round was
unproductive (nothing routable
+ /// yet / all failed); a failed probe simply does not contribute.
+ private List<Long> runConcurrentRound(List<Callable<Long>> tasks,
ExecutorService pool) {
+ List<Future<Long>> futures = new ArrayList<>(tasks.size());
+ for (Callable<Long> task : tasks) {
+ futures.add(pool.submit(task));
+ }
+ List<Long> latencies = new ArrayList<>(futures.size());
+ for (Future<Long> future : futures) {
+ try {
+ Long elapsedMs = future.get(WARMUP_PROBE_TIMEOUT_MS + 1000L,
TimeUnit.MILLISECONDS);
+ if (elapsedMs != null && elapsedMs >= 0) {
+ latencies.add(elapsedMs);
+ }
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ return latencies;
+ } catch (Exception e) {
+ // Individual probe failed/timed out; it just does not count toward
the round. Cancel it so a probe
+ // that outlived the round's get()-timeout does not keep running
behind the next round (no-op if it
+ // already completed).
+ future.cancel(true);
+ }
+ }
+ return latencies;
+ }
+
+ /// Sleeps [#WARMUP_FAILURE_BACKOFF_MS] after an unproductive probe round.
Returns `false` if interrupted
+ /// (shutdown) so the caller stops promptly and readiness opens.
+ private boolean warmupBackoff() {
+ try {
+ Thread.sleep(WARMUP_FAILURE_BACKOFF_MS);
+ return true;
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ return false;
+ }
+ }
+
+ /// Default probe path: probe `SELECT * FROM "<t>" LIMIT 1` over the
configured tables or a greedy
+ /// set-cover of tables covering every routable server. Exits once at least
`minIterations` probes have
+ /// completed successfully -- a depth floor that guarantees enough
invocations to drive the query path's
+ /// JIT to its top tier -- OR the budget expires. A latency target is
deliberately not used: 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.
+ private boolean warmUpWithAutoSelect(BrokerWarmupConfig config, long
deadlineMs, ExecutorService pool,
+ int concurrency) {
+ long successfulProbes = 0;
+ int rounds = 0;
+ // Select the probe tables once and reuse them across rounds. Only
re-select while the result is empty
+ // (routing not populated yet); recomputing the greedy set-cover every
round would repeat O(tables) work
+ // up to minIterations times on a large-table tenant, for a result that
does not change once non-empty.
+ List<String> tables = List.of();
+ boolean coverageReported = false;
+ // Monotonic across rounds so the round-robin actually advances through
every covered table even at
+ // concurrency 1; resetting per round would probe only the first
`concurrency` tables forever.
+ int probeSeq = 0;
+ while (System.currentTimeMillis() < deadlineMs &&
!Thread.currentThread().isInterrupted()) {
+ if (tables.isEmpty()) {
+ tables = config.hasCustomTables() ?
filterRoutableTables(_routingManager, config.tables())
+ : selectProbeTables(_routingManager, config.maxTables());
+ }
+ if (tables.isEmpty()) {
+ // Converged but nothing routable yet: back off and retry rather than
declaring the broker warm --
+ // an empty routing table here would otherwise make the gate a no-op
precisely on a cold broker.
+ if (!warmupBackoff()) {
+ return false;
+ }
+ continue;
+ }
+ if (!coverageReported) {
+ // Report once, now that the probe tables are settled, whether they
route to every server.
+ coverageReported = true;
+ reportServerCoverage(_routingManager, _brokerMetrics,
config.hasCustomTables(), tables);
+ }
+ long remainingMs = deadlineMs - System.currentTimeMillis();
+ if (remainingMs <= 0) {
+ break;
+ }
+ rounds++;
+ long probeTimeoutMs = Math.min(remainingMs, WARMUP_PROBE_TIMEOUT_MS);
+ // A batch of `concurrency` probes, round-robin over the covered tables
(advancing across rounds via
+ // probeSeq): exercises every server (coverage) AND real concurrency
(contention) at once.
+ List<Callable<Long>> tasks = new ArrayList<>(concurrency);
+ for (String tableNameWithType : roundRobinBatch(tables, probeSeq,
concurrency)) {
+ tasks.add(() -> probe(compileProbe(tableNameWithType),
probeTimeoutMs));
+ }
+ probeSeq += concurrency;
+ List<Long> latencies = runConcurrentRound(tasks, pool);
+ if (latencies.isEmpty()) {
+ if (!warmupBackoff()) {
+ return false;
+ }
+ continue;
+ }
+ successfulProbes += latencies.size();
+ if (successfulProbes >= config.minIterations()) {
+ LOGGER.info("Broker warmup completed after {} round(s) at concurrency
{}; {} probes (floor {})", rounds,
+ concurrency, successfulProbes, config.minIterations());
+ return true;
+ }
+ }
+ LOGGER.warn("Broker warmup budget expired after {} round(s) at concurrency
{}; {} probes (floor {}). Proceeding "
+ + "to serve traffic.", rounds, concurrency, successfulProbes,
config.minIterations());
+ return false;
+ }
+
+ /// Custom-query probe path: run the operator-supplied queries repeatedly,
all of them every round, and
+ /// exit on the same `minIterations` probe-count floor as the default probe,
or the budget. A probe-count
+ /// floor is a cost-independent warm signal; latency plateauing is not (a
trivial query reaches a stable
+ /// latency while still cold), so it is deliberately not used here. An
expensive query that cannot reach
+ /// the floor within the budget exits on the budget, having warmed as much
as the budget allowed.
+ private boolean warmUpWithCustomQuery(BrokerWarmupConfig config, long
deadlineMs, ExecutorService pool,
+ int concurrency) {
+ // Validate-compile each configured query once so a bad one fails fast
(and never busy-spins), keeping
+ // only those that compile. Each probe below compiles its own fresh
request so nothing mutable is shared
+ // across the concurrent tasks.
+ List<String> queries = new ArrayList<>(config.queries().size());
+ for (String query : config.queries()) {
+ try {
+ CalciteSqlCompiler.compileToBrokerRequest(query);
+ queries.add(query);
+ } catch (Exception e) {
+ LOGGER.warn("Broker warmup query failed to compile; skipping it: {}",
query, e);
+ }
+ }
+ if (queries.isEmpty()) {
+ LOGGER.warn("Broker warmup: no configured query compiled; skipping
warmup");
+ return false;
+ }
+ long successfulProbes = 0;
+ int rounds = 0;
+ while (System.currentTimeMillis() < deadlineMs &&
!Thread.currentThread().isInterrupted()) {
+ long remainingMs = deadlineMs - System.currentTimeMillis();
+ if (remainingMs <= 0) {
+ break;
+ }
+ rounds++;
+ long probeTimeoutMs = Math.min(remainingMs, WARMUP_PROBE_TIMEOUT_MS);
+ // Run every configured query each round (each `concurrency` times, to
warm the concurrency step too).
+ List<Callable<Long>> tasks = new ArrayList<>(queries.size() *
concurrency);
+ for (String query : queries) {
+ for (int i = 0; i < concurrency; i++) {
+ tasks.add(() ->
probe(CalciteSqlCompiler.compileToBrokerRequest(query), probeTimeoutMs));
+ }
+ }
+ List<Long> latencies = runConcurrentRound(tasks, pool);
+ if (latencies.isEmpty()) {
+ if (!warmupBackoff()) {
+ return false;
+ }
+ continue;
+ }
+ successfulProbes += latencies.size();
+ // Same depth floor as the default probe: a probe-count floor is a
cost-independent warm signal (it
+ // drives the query path's JIT to its top tier), whereas latency
plateauing is not -- a trivial query
+ // reaches a stable latency while still cold. An expensive query that
cannot reach the floor within the
+ // budget simply exits on the budget, having warmed as much as the
budget allowed.
+ if (successfulProbes >= config.minIterations()) {
+ LOGGER.info("Broker warmup (custom query) completed after {} round(s)
at concurrency {}; {} probes "
+ + "(floor {})", rounds, concurrency, successfulProbes,
config.minIterations());
+ return true;
+ }
+ }
+ LOGGER.warn("Broker warmup (custom query) budget expired after {} round(s)
at concurrency {}; {} probes "
+ + "(floor {}). Proceeding to serve traffic.", rounds, concurrency,
successfulProbes, config.minIterations());
+ return false;
+ }
+
+ /// Compiles the default probe query for a single physical table.
+ private BrokerRequest compileProbe(String tableNameWithType) {
+ return CalciteSqlCompiler.compileToBrokerRequest("SELECT * FROM \"" +
tableNameWithType + "\" LIMIT 1");
+ }
+
+ /// Issues one probe query through [QueryRouter] and returns its wall-clock
duration in ms, or -1 on
+ /// failure. Builds the route the same way the normal path does, minus
auth/quota/logging.
+ private long probe(BrokerRequest brokerRequest, long timeoutMs) {
+ String tableName = brokerRequest.getQuerySource().getTableName();
+ try {
+ String rawTableName = TableNameBuilder.extractRawTableName(tableName);
+ TableType tableType =
TableNameBuilder.getTableTypeFromTableName(tableName);
+ long requestId = _requestIdGenerator.get();
+ TableRouteInfo routeInfo =
+ _implicitHybridTableRouteProvider.getTableRouteInfo(rawTableName,
_tableCache, _routingManager);
+ if (!(routeInfo instanceof ImplicitHybridTableRouteInfo) ||
!routeInfo.isExists()) {
+ return -1;
+ }
+ // A type-suffixed table (the default probe) targets exactly one type; a
raw name (possible with a
+ // custom query) targets whichever types actually exist. Each leg's
request must carry the
+ // type-suffixed table name -- that is what routing resolves against --
so build a per-leg request
+ // from the base query (as the normal path does) rather than reusing the
raw-named request, which
+ // would route to no server and yield an empty request map.
+ BrokerRequest offlineBrokerRequest = shouldRouteOffline(tableType,
routeInfo)
+ ? typedRequest(brokerRequest, routeInfo.getOfflineTableName()) :
null;
+ BrokerRequest realtimeBrokerRequest = shouldRouteRealtime(tableType,
routeInfo)
+ ? typedRequest(brokerRequest, routeInfo.getRealtimeTableName()) :
null;
+ if (offlineBrokerRequest == null && realtimeBrokerRequest == null) {
+ // Neither leg is routable right now (e.g. a hybrid table probed via
its offline name before the
+ // time boundary exists). submitQuery asserts at least one leg is
non-null, so skip cleanly rather
+ // than build an empty route -- the table is retried on the next round
once its leg appears.
+ return -1;
+ }
+ ImplicitHybridTableRouteInfo hybridRouteInfo =
(ImplicitHybridTableRouteInfo) routeInfo;
+ hybridRouteInfo.setOfflineBrokerRequest(offlineBrokerRequest);
+ hybridRouteInfo.setRealtimeBrokerRequest(realtimeBrokerRequest);
+ _implicitHybridTableRouteProvider.calculateRoutes(routeInfo,
_routingManager, offlineBrokerRequest,
+ realtimeBrokerRequest, requestId);
+ if (!routeInfo.isRouteExists()) {
+ return -1;
+ }
+ long startMs = System.currentTimeMillis();
+ // Scatter/gather (+ DataTable deserialize on receive) is the dominant
cost and does not need a
+ // QueryThreadContext; run it directly so its warmth always counts
toward the probe latency.
+ AsyncQueryResponse response = _queryRouter.submitQuery(requestId,
rawTableName, routeInfo, timeoutMs);
+ Map<ServerRoutingInstance, ServerResponse> finalResponses =
response.getFinalResponses();
+ // Best-effort: also warm the broker reduce path (result discarded) so
the first real query of this
+ // shape does not pay it. Isolated in its own try -- reduceOnDataTable
requires a QueryThreadContext
+ // and could fail for edge cases; a reduce failure must never fail the
probe, since the
+ // scatter/gather/deserialize warmth has already happened.
+ try (QueryThreadContext ignore = QueryThreadContext.open(
+ new QueryExecutionContext(QueryExecutionContext.QueryType.SSE,
requestId, Long.toString(requestId),
+ "warmup", startMs, Long.MAX_VALUE, Long.MAX_VALUE, _brokerId,
_brokerId, ""), _threadAccountant)) {
+ Map<ServerRoutingInstance, DataTable> dataTableMap = new
HashMap<>(finalResponses.size());
+ for (Map.Entry<ServerRoutingInstance, ServerResponse> entry :
finalResponses.entrySet()) {
+ DataTable dataTable = entry.getValue().getDataTable();
+ if (dataTable != null) {
+ dataTableMap.put(entry.getKey(), dataTable);
+ }
+ }
+ if (!dataTableMap.isEmpty()) {
+ _brokerReduceService.reduceOnDataTable(brokerRequest, brokerRequest,
dataTableMap, timeoutMs,
Review Comment:
**Risk / major — this does pollute a customer-facing surface.**
`BrokerReduceService.reduceOnDataTable` calls
`aggregator.setStats(rawTableName, response, brokerMetrics)`, which records
`BrokerMeter.DOCUMENTS_SCANNED`, `ENTRIES_SCANNED_IN_FILTER`,
`ENTRIES_SCANNED_POST_FILTER` and the `OFFLINE`/`REALTIME_*_CPU_TIME_NS` +
`FRESHNESS_LAG_MS` timers — all tagged with the real raw table name.
With the default `minIterations=1000` that is up to 1000 synthetic samples
injected into a user's per-table dashboards and alerts in the first seconds
after every restart. The description says warmup "pollutes no customer facing
surface"; that holds for access control / quota / query log, but not for query
metrics.
Options: pass a throwaway `BrokerMetrics` into the warmup reduce, or
document the behaviour explicitly in the config comment so operators can
exclude the window.
##########
pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/BaseBrokerStarter.java:
##########
@@ -894,10 +911,76 @@ private void registerServiceStatusHandler() {
new
ServiceStatus.IdealStateAndExternalViewMatchServiceStatusCallback(_participantHelixManager,
_clusterName, _instanceId, resourcesToMonitor,
minResourcePercentForStartup)));
+ List<ServiceStatus.ServiceStatusCallback> callbacks = new ArrayList<>(3);
+ callbacks.add(new
ServiceStatus.LifecycleServiceStatusCallback(this::isStarting,
this::isShuttingDown));
+ callbacks.add(_helixConvergenceCallback);
+ if (_warmupConfig.enabled()) {
+ // The warmup readiness gate. Reports STARTING (not a new status value)
until warmup completes:
+ // callers throughout the codebase test for GOOD, and a new enum
constant would be visible to older
+ // mixed-version peers. MultipleCallbackServiceStatusCallback surfaces
the first non-GOOD callback, so
+ // this composes without touching the Helix or lifecycle callbacks. It
gates the existing readiness
+ // endpoint (getBrokerHealth -> ServiceStatus): a warming broker stays
out of the Service until it
+ // reports GOOD.
+ callbacks.add(new ServiceStatus.ServiceStatusCallback() {
+ @Override
+ public ServiceStatus.Status getServiceStatus() {
+ return _isWarm ? ServiceStatus.Status.GOOD :
ServiceStatus.Status.STARTING;
+ }
+
+ @Override
+ public String getStatusDescription() {
+ return _isWarm ? ServiceStatus.STATUS_DESCRIPTION_NONE : "Warming up
broker data plane";
+ }
+ });
+ }
ServiceStatus.setServiceStatusCallback(_instanceId,
- new ServiceStatus.MultipleCallbackServiceStatusCallback(List.of(
- new ServiceStatus.LifecycleServiceStatusCallback(this::isStarting,
this::isShuttingDown),
- _helixConvergenceCallback)));
+ new ServiceStatus.MultipleCallbackServiceStatusCallback(callbacks));
+ }
+
+ /// Runs the data-plane warmup on a background thread and flips [#_isWarm]
when it finishes.
+ ///
+ /// Asynchronous so `start()` still returns promptly -- the gate is enforced
through `ServiceStatus`, not
+ /// by blocking startup. The flag is set in a `finally` so the gate opens
even if warmup throws:
+ /// readiness held open indefinitely would stall a rolling restart, a worse
failure than serving a cold
+ /// broker. When disabled this is a no-op and readiness behaves exactly as
before.
+ private void startWarmup() {
+ // The gauge is published in both branches so dashboards can rely on it
always existing; with warmup
+ // disabled it simply reads 1 from the start, matching pre-change
behaviour.
+ _brokerMetrics.setOrUpdateGlobalGauge(BrokerGauge.BROKER_WARM, () ->
_isWarm ? 1L : 0L);
+ if (!_warmupConfig.enabled()) {
+ _isWarm = true;
+ return;
+ }
+ _warmupThread = new Thread(() -> {
+ // Set once Helix converges. Both the budget and the duration metric are
measured from here, not from
+ // thread start, so the deliberately unbounded convergence wait is
charged against neither: readiness
+ // is already withheld until convergence by the Helix callbacks, so it
costs nothing.
+ long warmStartMs = 0L;
+ try {
+ long threadStartMs = System.currentTimeMillis();
+ awaitHelixConvergence();
+ warmStartMs = System.currentTimeMillis();
+ LOGGER.info("Helix converged after {} ms; starting data-plane warmup",
warmStartMs - threadStartMs);
+ boolean reachedFloor = _brokerRequestHandler.warmUp(_warmupConfig,
warmStartMs + _warmupConfig.budgetMs());
+ LOGGER.info("Broker warmup finished in {} ms (reachedFloor={})",
System.currentTimeMillis() - warmStartMs,
+ reachedFloor);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ LOGGER.info("Broker warmup interrupted before completion");
+ } catch (Throwable t) {
+ LOGGER.warn("Broker warmup threw; opening readiness anyway", t);
+ } finally {
+ _isWarm = true;
Review Comment:
**Minor — `_isWarm` in a `finally` is right. Separate point: `stop()` does
not wait for the probe pool.**
`stopWarmup()` interrupts the warmup thread, but the probe threads are only
interrupted by `probePool.shutdownNow()` inside `warmUp()`'s `finally`, which
runs on the warmup thread *after* it notices the interrupt. `stop()` meanwhile
proceeds to disconnect Helix and shut down `_brokerRequestHandler`
(`_queryRouter`, `_brokerReduceService`).
Nothing corrupts, but a broker stopped mid-warmup will log NPEs /
rejected-execution from probes racing the teardown — noise that on-call has to
learn to ignore. A bounded `warmupThread.join(1000)` after the interrupt would
keep shutdown logs clean without meaningfully slowing shutdown.
##########
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:
**Tests — this assertion is close to vacuous, and it is the only coverage of
the feature's central behaviour.**
By the time it runs, `setUp()` has built segments, uploaded them and waited
for all docs to load, so warmup finished long ago. The test proves "eventually
GOOD", which is also the behaviour with warmup *disabled*. It would pass if the
callback always returned GOOD, or if it were never registered at all.
What would actually pin it: assert the transition. Bring up a broker with a
large budget and an unreachable floor, assert
`ServiceStatus.getServiceStatus(instanceId) == STARTING` and that
`getStatusDescription()` contains "Warming up" *before* warmup can finish, then
assert GOOD afterwards. That closed-then-open transition is the feature, and
nothing currently covers it — `BrokerWarmupTest` only covers the static helpers.
##########
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:
**API contract — two things worth adding to this javadoc, because the
implementation depends on them.**
1. "Must return by `deadlineMs`" is the load-bearing clause of the whole
feature, and the shipped implementation does not honour it (see my comment on
`runConcurrentRound`). Either the contract or the implementation needs to move.
2. Nothing says whether the caller may invoke this concurrently with live
traffic. `BrokerStartupWarmupIntegrationTest` does exactly that — it calls
`warmUp()` on an already-started broker — and so would anyone wiring this to an
admin endpoint later. Worth stating that implementations must be safe against a
handler that is already serving.
Also: mirroring `preConnectServers()` right below is a clean precedent, but
the broker now has two independent startup gates with two independent budgets.
How they compose (readiness = both satisfied) is documented in the PR
description but not in the code; the interface seems like the right home for it.
##########
pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/SingleConnectionBrokerRequestHandler.java:
##########
@@ -117,6 +138,438 @@ public void shutDown() {
_brokerReduceService.shutDown();
}
+ /// Warms the scatter-gather path before readiness is granted: repeatedly
runs real probe queries until
+ /// they have run enough times (the `minIterations` depth floor) or the
budget expires. Broker-to-server
+ /// channels are opened lazily by the probes themselves.
+ ///
+ /// Probes go through [QueryRouter] rather than [#handleRequest],
deliberately bypassing access control,
+ /// query quota and the query log -- warmup then needs no synthetic identity
and pollutes no
+ /// customer-facing surface. Never throws, and always returns within the
configured budget.
+ @Override
+ public boolean warmUp(BrokerWarmupConfig config, long deadlineMs) {
+ // Stage 1: local, no-network warmup of the compile +
response-serialization paths, run before the
+ // network probe so they are warm even if no server is reachable.
+ warmUpLocal();
+ int concurrency = Math.max(1, config.concurrency());
+ ExecutorService probePool = Executors.newFixedThreadPool(concurrency,
+ new
ThreadFactoryBuilder().setNameFormat("broker-warmup-probe-%d").setDaemon(true).build());
+ try {
+ return config.hasCustomQuery() ? warmUpWithCustomQuery(config,
deadlineMs, probePool, concurrency)
+ : warmUpWithAutoSelect(config, deadlineMs, probePool, concurrency);
+ } catch (Exception e) {
+ LOGGER.warn("Broker warmup failed; proceeding without it", e);
+ return false;
+ } finally {
+ probePool.shutdownNow();
+ }
+ }
+
+ /// Stage 1: local, no-network warmup. The network probe goes through
[QueryRouter] and stops at the
+ /// gathered DataTables, so it never exercises the SQL compile path's cold
start nor the response build +
+ /// JSON serialization the first real query pays. This compiles a throwaway
query and serializes a small
+ /// synthetic [BrokerResponseNative], warming both. Never throws.
+ private void warmUpLocal() {
+ try {
+ CalciteSqlCompiler.compileToBrokerRequest("SELECT 1");
+ BrokerResponseNative response = new BrokerResponseNative();
+ response.setResultTable(new ResultTable(
+ new DataSchema(new String[]{"warmup"}, new
DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.LONG}),
+ Collections.singletonList(new Object[]{1L})));
+ response.setNumDocsScanned(1);
+ response.toJsonString();
+ } catch (Exception e) {
+ LOGGER.debug("Local (stage 1) warmup failed; continuing", e);
+ }
+ }
+
+ /// Fires the given probe tasks concurrently on the shared pool and returns
the latencies (ms) of the
+ /// probes that completed successfully. Empty means the whole round was
unproductive (nothing routable
+ /// yet / all failed); a failed probe simply does not contribute.
+ private List<Long> runConcurrentRound(List<Callable<Long>> tasks,
ExecutorService pool) {
+ List<Future<Long>> futures = new ArrayList<>(tasks.size());
+ for (Callable<Long> task : tasks) {
+ futures.add(pool.submit(task));
+ }
+ List<Long> latencies = new ArrayList<>(futures.size());
+ for (Future<Long> future : futures) {
+ try {
+ Long elapsedMs = future.get(WARMUP_PROBE_TIMEOUT_MS + 1000L,
TimeUnit.MILLISECONDS);
+ if (elapsedMs != null && elapsedMs >= 0) {
+ latencies.add(elapsedMs);
+ }
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ return latencies;
+ } catch (Exception e) {
+ // Individual probe failed/timed out; it just does not count toward
the round. Cancel it so a probe
+ // that outlived the round's get()-timeout does not keep running
behind the next round (no-op if it
+ // already completed).
+ future.cancel(true);
+ }
+ }
+ return latencies;
+ }
+
+ /// Sleeps [#WARMUP_FAILURE_BACKOFF_MS] after an unproductive probe round.
Returns `false` if interrupted
+ /// (shutdown) so the caller stops promptly and readiness opens.
+ private boolean warmupBackoff() {
+ try {
+ Thread.sleep(WARMUP_FAILURE_BACKOFF_MS);
+ return true;
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ return false;
+ }
+ }
+
+ /// Default probe path: probe `SELECT * FROM "<t>" LIMIT 1` over the
configured tables or a greedy
+ /// set-cover of tables covering every routable server. Exits once at least
`minIterations` probes have
+ /// completed successfully -- a depth floor that guarantees enough
invocations to drive the query path's
+ /// JIT to its top tier -- OR the budget expires. A latency target is
deliberately not used: 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.
+ private boolean warmUpWithAutoSelect(BrokerWarmupConfig config, long
deadlineMs, ExecutorService pool,
+ int concurrency) {
+ long successfulProbes = 0;
+ int rounds = 0;
+ // Select the probe tables once and reuse them across rounds. Only
re-select while the result is empty
+ // (routing not populated yet); recomputing the greedy set-cover every
round would repeat O(tables) work
+ // up to minIterations times on a large-table tenant, for a result that
does not change once non-empty.
+ List<String> tables = List.of();
+ boolean coverageReported = false;
+ // Monotonic across rounds so the round-robin actually advances through
every covered table even at
+ // concurrency 1; resetting per round would probe only the first
`concurrency` tables forever.
+ int probeSeq = 0;
+ while (System.currentTimeMillis() < deadlineMs &&
!Thread.currentThread().isInterrupted()) {
+ if (tables.isEmpty()) {
+ tables = config.hasCustomTables() ?
filterRoutableTables(_routingManager, config.tables())
+ : selectProbeTables(_routingManager, config.maxTables());
+ }
+ if (tables.isEmpty()) {
+ // Converged but nothing routable yet: back off and retry rather than
declaring the broker warm --
+ // an empty routing table here would otherwise make the gate a no-op
precisely on a cold broker.
+ if (!warmupBackoff()) {
+ return false;
+ }
+ continue;
+ }
+ if (!coverageReported) {
+ // Report once, now that the probe tables are settled, whether they
route to every server.
+ coverageReported = true;
+ reportServerCoverage(_routingManager, _brokerMetrics,
config.hasCustomTables(), tables);
+ }
+ long remainingMs = deadlineMs - System.currentTimeMillis();
+ if (remainingMs <= 0) {
+ break;
+ }
+ rounds++;
+ long probeTimeoutMs = Math.min(remainingMs, WARMUP_PROBE_TIMEOUT_MS);
+ // A batch of `concurrency` probes, round-robin over the covered tables
(advancing across rounds via
+ // probeSeq): exercises every server (coverage) AND real concurrency
(contention) at once.
+ List<Callable<Long>> tasks = new ArrayList<>(concurrency);
+ for (String tableNameWithType : roundRobinBatch(tables, probeSeq,
concurrency)) {
+ tasks.add(() -> probe(compileProbe(tableNameWithType),
probeTimeoutMs));
+ }
+ probeSeq += concurrency;
+ List<Long> latencies = runConcurrentRound(tasks, pool);
+ if (latencies.isEmpty()) {
+ if (!warmupBackoff()) {
+ return false;
+ }
+ continue;
+ }
+ successfulProbes += latencies.size();
+ if (successfulProbes >= config.minIterations()) {
+ LOGGER.info("Broker warmup completed after {} round(s) at concurrency
{}; {} probes (floor {})", rounds,
+ concurrency, successfulProbes, config.minIterations());
+ return true;
+ }
+ }
+ LOGGER.warn("Broker warmup budget expired after {} round(s) at concurrency
{}; {} probes (floor {}). Proceeding "
+ + "to serve traffic.", rounds, concurrency, successfulProbes,
config.minIterations());
+ return false;
+ }
+
+ /// Custom-query probe path: run the operator-supplied queries repeatedly,
all of them every round, and
+ /// exit on the same `minIterations` probe-count floor as the default probe,
or the budget. A probe-count
+ /// floor is a cost-independent warm signal; latency plateauing is not (a
trivial query reaches a stable
+ /// latency while still cold), so it is deliberately not used here. An
expensive query that cannot reach
+ /// the floor within the budget exits on the budget, having warmed as much
as the budget allowed.
+ private boolean warmUpWithCustomQuery(BrokerWarmupConfig config, long
deadlineMs, ExecutorService pool,
+ int concurrency) {
+ // Validate-compile each configured query once so a bad one fails fast
(and never busy-spins), keeping
+ // only those that compile. Each probe below compiles its own fresh
request so nothing mutable is shared
+ // across the concurrent tasks.
+ List<String> queries = new ArrayList<>(config.queries().size());
+ for (String query : config.queries()) {
+ try {
+ CalciteSqlCompiler.compileToBrokerRequest(query);
+ queries.add(query);
+ } catch (Exception e) {
+ LOGGER.warn("Broker warmup query failed to compile; skipping it: {}",
query, e);
+ }
+ }
+ if (queries.isEmpty()) {
+ LOGGER.warn("Broker warmup: no configured query compiled; skipping
warmup");
+ return false;
+ }
+ long successfulProbes = 0;
+ int rounds = 0;
+ while (System.currentTimeMillis() < deadlineMs &&
!Thread.currentThread().isInterrupted()) {
+ long remainingMs = deadlineMs - System.currentTimeMillis();
+ if (remainingMs <= 0) {
+ break;
+ }
+ rounds++;
+ long probeTimeoutMs = Math.min(remainingMs, WARMUP_PROBE_TIMEOUT_MS);
+ // Run every configured query each round (each `concurrency` times, to
warm the concurrency step too).
+ List<Callable<Long>> tasks = new ArrayList<>(queries.size() *
concurrency);
+ for (String query : queries) {
+ for (int i = 0; i < concurrency; i++) {
+ tasks.add(() ->
probe(CalciteSqlCompiler.compileToBrokerRequest(query), probeTimeoutMs));
+ }
+ }
+ List<Long> latencies = runConcurrentRound(tasks, pool);
+ if (latencies.isEmpty()) {
+ if (!warmupBackoff()) {
+ return false;
+ }
+ continue;
+ }
+ successfulProbes += latencies.size();
+ // Same depth floor as the default probe: a probe-count floor is a
cost-independent warm signal (it
+ // drives the query path's JIT to its top tier), whereas latency
plateauing is not -- a trivial query
+ // reaches a stable latency while still cold. An expensive query that
cannot reach the floor within the
+ // budget simply exits on the budget, having warmed as much as the
budget allowed.
+ if (successfulProbes >= config.minIterations()) {
+ LOGGER.info("Broker warmup (custom query) completed after {} round(s)
at concurrency {}; {} probes "
+ + "(floor {})", rounds, concurrency, successfulProbes,
config.minIterations());
+ return true;
+ }
+ }
+ LOGGER.warn("Broker warmup (custom query) budget expired after {} round(s)
at concurrency {}; {} probes "
+ + "(floor {}). Proceeding to serve traffic.", rounds, concurrency,
successfulProbes, config.minIterations());
+ return false;
+ }
+
+ /// Compiles the default probe query for a single physical table.
+ private BrokerRequest compileProbe(String tableNameWithType) {
+ return CalciteSqlCompiler.compileToBrokerRequest("SELECT * FROM \"" +
tableNameWithType + "\" LIMIT 1");
+ }
+
+ /// Issues one probe query through [QueryRouter] and returns its wall-clock
duration in ms, or -1 on
+ /// failure. Builds the route the same way the normal path does, minus
auth/quota/logging.
+ private long probe(BrokerRequest brokerRequest, long timeoutMs) {
+ String tableName = brokerRequest.getQuerySource().getTableName();
+ try {
+ String rawTableName = TableNameBuilder.extractRawTableName(tableName);
+ TableType tableType =
TableNameBuilder.getTableTypeFromTableName(tableName);
+ long requestId = _requestIdGenerator.get();
+ TableRouteInfo routeInfo =
+ _implicitHybridTableRouteProvider.getTableRouteInfo(rawTableName,
_tableCache, _routingManager);
+ if (!(routeInfo instanceof ImplicitHybridTableRouteInfo) ||
!routeInfo.isExists()) {
+ return -1;
+ }
+ // A type-suffixed table (the default probe) targets exactly one type; a
raw name (possible with a
+ // custom query) targets whichever types actually exist. Each leg's
request must carry the
+ // type-suffixed table name -- that is what routing resolves against --
so build a per-leg request
+ // from the base query (as the normal path does) rather than reusing the
raw-named request, which
+ // would route to no server and yield an empty request map.
+ BrokerRequest offlineBrokerRequest = shouldRouteOffline(tableType,
routeInfo)
+ ? typedRequest(brokerRequest, routeInfo.getOfflineTableName()) :
null;
+ BrokerRequest realtimeBrokerRequest = shouldRouteRealtime(tableType,
routeInfo)
+ ? typedRequest(brokerRequest, routeInfo.getRealtimeTableName()) :
null;
+ if (offlineBrokerRequest == null && realtimeBrokerRequest == null) {
+ // Neither leg is routable right now (e.g. a hybrid table probed via
its offline name before the
+ // time boundary exists). submitQuery asserts at least one leg is
non-null, so skip cleanly rather
+ // than build an empty route -- the table is retried on the next round
once its leg appears.
+ return -1;
+ }
+ ImplicitHybridTableRouteInfo hybridRouteInfo =
(ImplicitHybridTableRouteInfo) routeInfo;
+ hybridRouteInfo.setOfflineBrokerRequest(offlineBrokerRequest);
Review Comment:
**Nit — these two setters are dead writes.**
`ImplicitHybridTableRouteProvider.calculateRoutes()` takes
`offlineBrokerRequest`/`realtimeBrokerRequest` as parameters, may null them
locally when the routing table comes back empty, and ends with
`hybridTableRouteInfo.setOfflineBrokerRequest(...)` /
`setRealtimeBrokerRequest(...)`. So setting them here is redundant, and
slightly misleading since it reads as though `calculateRoutes` consumed the
route info's state. Passing them as parameters is enough.
##########
pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java:
##########
@@ -435,6 +435,42 @@ public static class Broker {
"pinot.broker.startup.minResourcePercent";
public static final double DEFAULT_BROKER_MIN_RESOURCE_PERCENT_FOR_START =
100.0;
+ // Startup data-plane warmup: before readiness is granted, run probe
queries so the JIT-compiled query
+ // path and per-query caches are warm before the first real traffic. Off
by default; opt-in per
+ // deployment.
+ public static final String CONFIG_OF_BROKER_STARTUP_WARMUP_ENABLED =
"pinot.broker.startup.warmup.enabled";
+ public static final boolean DEFAULT_BROKER_STARTUP_WARMUP_ENABLED = false;
+ // Hard ceiling on the whole warmup (measured from Helix convergence):
readiness opens when it expires
+ // whatever the probe progress, so a slow or unreachable server cannot
stall a rolling restart.
+ public static final String CONFIG_OF_BROKER_STARTUP_WARMUP_BUDGET_MS =
"pinot.broker.startup.warmup.budgetMs";
+ public static final long DEFAULT_BROKER_STARTUP_WARMUP_BUDGET_MS = 15_000L;
+ // Minimum number of successful probe queries before the default probe
declares the broker warm. This is
+ // a depth floor, not a latency guess: enough probe invocations to drive
the query path's JIT to its top
+ // tier. Warmup exits when this many probes have run OR the budget expires
-- whichever comes first.
+ public static final String CONFIG_OF_BROKER_STARTUP_WARMUP_MIN_ITERATIONS =
+ "pinot.broker.startup.warmup.minIterations";
+ public static final int DEFAULT_BROKER_STARTUP_WARMUP_MIN_ITERATIONS =
1000;
Review Comment:
**Defaults — `minIterations=1000` and `budgetMs=15000` look mutually
unreachable.**
With `concurrency=1` (the default) each round is one probe: compile + route
+ scatter + gather + deserialize + reduce against real segments. 1000 of those
in 15s means 15ms per full round trip, which leaves very little headroom on a
busy or TLS-enabled cluster.
The consequence is that most deployments exit on the budget, and that path
logs at WARN (`budget expired ... Proceeding to serve traffic`) — so enabling
the feature produces a warning on every restart, which reads as a failure.
Could you say which of the two the 2-vCPU benchmark arm actually hit? If it
exited on the budget rather than the floor, the floor default is decorative and
one of the two numbers should move. Alternatively, demote that log to INFO when
`successfulProbes` is already substantial.
##########
pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/BaseBrokerStarter.java:
##########
@@ -894,10 +911,76 @@ private void registerServiceStatusHandler() {
new
ServiceStatus.IdealStateAndExternalViewMatchServiceStatusCallback(_participantHelixManager,
_clusterName, _instanceId, resourcesToMonitor,
minResourcePercentForStartup)));
+ List<ServiceStatus.ServiceStatusCallback> callbacks = new ArrayList<>(3);
+ callbacks.add(new
ServiceStatus.LifecycleServiceStatusCallback(this::isStarting,
this::isShuttingDown));
+ callbacks.add(_helixConvergenceCallback);
+ if (_warmupConfig.enabled()) {
+ // The warmup readiness gate. Reports STARTING (not a new status value)
until warmup completes:
+ // callers throughout the codebase test for GOOD, and a new enum
constant would be visible to older
+ // mixed-version peers. MultipleCallbackServiceStatusCallback surfaces
the first non-GOOD callback, so
+ // this composes without touching the Helix or lifecycle callbacks. It
gates the existing readiness
+ // endpoint (getBrokerHealth -> ServiceStatus): a warming broker stays
out of the Service until it
+ // reports GOOD.
+ callbacks.add(new ServiceStatus.ServiceStatusCallback() {
Review Comment:
**Docs — worth stating the gate's scope: it gates `ServiceStatus` / the HTTP
readiness endpoint, not Helix-based broker discovery.**
The gate itself is nicely done — reusing `STARTING` instead of adding an
enum constant, and composing through `MultipleCallbackServiceStatusCallback`,
is the right call for mixed-version peers.
But a warming broker is already in the broker resource's ExternalView by
then — that is exactly the signal `awaitHelixConvergence()` waits for. Clients
that select brokers from Helix (the Java client's Helix-based broker selector,
and anything reading the broker resource directly) will route to this broker
while it reports `STARTING`. That is fine for the k8s Service `readinessProbe`
deployment this targets, but the description reads as if readiness were
universal.
One sentence in the config comment would prevent a surprised operator:
"gates the HTTP readiness endpoint; clients using Helix-based broker discovery
are unaffected".
##########
pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/SingleConnectionBrokerRequestHandler.java:
##########
@@ -117,6 +138,438 @@ public void shutDown() {
_brokerReduceService.shutDown();
}
+ /// Warms the scatter-gather path before readiness is granted: repeatedly
runs real probe queries until
+ /// they have run enough times (the `minIterations` depth floor) or the
budget expires. Broker-to-server
+ /// channels are opened lazily by the probes themselves.
+ ///
+ /// Probes go through [QueryRouter] rather than [#handleRequest],
deliberately bypassing access control,
+ /// query quota and the query log -- warmup then needs no synthetic identity
and pollutes no
+ /// customer-facing surface. Never throws, and always returns within the
configured budget.
+ @Override
+ public boolean warmUp(BrokerWarmupConfig config, long deadlineMs) {
+ // Stage 1: local, no-network warmup of the compile +
response-serialization paths, run before the
+ // network probe so they are warm even if no server is reachable.
+ warmUpLocal();
+ int concurrency = Math.max(1, config.concurrency());
+ ExecutorService probePool = Executors.newFixedThreadPool(concurrency,
+ new
ThreadFactoryBuilder().setNameFormat("broker-warmup-probe-%d").setDaemon(true).build());
+ try {
+ return config.hasCustomQuery() ? warmUpWithCustomQuery(config,
deadlineMs, probePool, concurrency)
+ : warmUpWithAutoSelect(config, deadlineMs, probePool, concurrency);
+ } catch (Exception e) {
+ LOGGER.warn("Broker warmup failed; proceeding without it", e);
+ return false;
+ } finally {
+ probePool.shutdownNow();
+ }
+ }
+
+ /// Stage 1: local, no-network warmup. The network probe goes through
[QueryRouter] and stops at the
+ /// gathered DataTables, so it never exercises the SQL compile path's cold
start nor the response build +
+ /// JSON serialization the first real query pays. This compiles a throwaway
query and serializes a small
+ /// synthetic [BrokerResponseNative], warming both. Never throws.
+ private void warmUpLocal() {
Review Comment:
**Design — stage 1 runs exactly once, which by this PR's own argument cannot
warm anything.**
The central claim is that ~1000 invocations are needed to reach top-tier
JIT, and that a handful leaves the path only partially compiled (the
29%-cancellation shallow arm). `warmUpLocal()` calls `compileToBrokerRequest`
and `toJsonString` once each — squarely in the still-interpreted regime.
Related: the probes go through `QueryRouter`, not `handleRequest`, so
everything between the HTTP layer and the router stays cold after 1000 probes —
`BaseSingleStageBrokerRequestHandler.doHandleRequest`, query rewriting,
time-boundary attachment, `BrokerResponse` construction and the JSON
serialization of the real response. "Each probe exercises the full serve path"
is a bit strong: it exercises compile + route + scatter/gather + reduce, which
is most of the CPU but not the request/response shell.
Either loop `warmUpLocal()` to the same floor (it is local and cheap, so
this is nearly free), or soften the claim that stage 1 warms those paths.
##########
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:
**Naming — `BROKER_WARM` and `BROKER_WARMUP_UNCOVERED_SERVERS` use two
different prefixes for one feature, and the surrounding enum uses neither.**
Neighbours are `QUERY_RATE_LIMIT_DISABLED`,
`NETTY_CONNECTION_CONNECT_TIME_MS`, `UNHEALTHY_SERVERS` — the `BROKER_` prefix
is redundant inside `BrokerGauge`. `BrokerTimer.STARTUP_WARMUP_DURATION_MS`
already picks the convention that matches the `pinot.broker.startup.warmup.*`
config prefix.
Suggest `STARTUP_WARM` (or `WARMUP_COMPLETE`) and
`STARTUP_WARMUP_UNCOVERED_SERVERS`. Gauge names are effectively permanent once
dashboards and alerts reference them, so this is much cheaper to settle now
than after release.
##########
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:
**Tests — coverage of the static helpers is genuinely good; the round loop
has none.**
Everything covered here (`selectProbeTables`, `uncoveredRoutableServers`,
`roundRobinBatch`, `shouldRouteOffline`/`Realtime`, `filterRoutableTables`,
config parsing) is a pure function, and the set-cover and tenant-scoping tests
in particular are well chosen.
Untested, and where I think the risk actually lives:
- `warmUpWithAutoSelect` / `warmUpWithCustomQuery` return `false` and return
*promptly* when the budget expires. The integration test covers this with a 10s
slack, which is not a unit-level guarantee.
- An interrupt mid-round unwinds and returns rather than continuing.
- The "nothing routable yet" path backs off instead of busy-spinning, and
re-selects once routing populates.
- `runConcurrentRound` cancels futures it timed out on.
All four are reachable with an injectable probe function — extracting the
loop from the network probe would make them unit-testable the same way the
helpers already are.
##########
pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/SingleConnectionBrokerRequestHandler.java:
##########
@@ -117,6 +138,438 @@ public void shutDown() {
_brokerReduceService.shutDown();
}
+ /// Warms the scatter-gather path before readiness is granted: repeatedly
runs real probe queries until
+ /// they have run enough times (the `minIterations` depth floor) or the
budget expires. Broker-to-server
+ /// channels are opened lazily by the probes themselves.
+ ///
+ /// Probes go through [QueryRouter] rather than [#handleRequest],
deliberately bypassing access control,
+ /// query quota and the query log -- warmup then needs no synthetic identity
and pollutes no
+ /// customer-facing surface. Never throws, and always returns within the
configured budget.
+ @Override
+ public boolean warmUp(BrokerWarmupConfig config, long deadlineMs) {
+ // Stage 1: local, no-network warmup of the compile +
response-serialization paths, run before the
+ // network probe so they are warm even if no server is reachable.
+ warmUpLocal();
+ int concurrency = Math.max(1, config.concurrency());
+ ExecutorService probePool = Executors.newFixedThreadPool(concurrency,
+ new
ThreadFactoryBuilder().setNameFormat("broker-warmup-probe-%d").setDaemon(true).build());
+ try {
+ return config.hasCustomQuery() ? warmUpWithCustomQuery(config,
deadlineMs, probePool, concurrency)
+ : warmUpWithAutoSelect(config, deadlineMs, probePool, concurrency);
+ } catch (Exception e) {
+ LOGGER.warn("Broker warmup failed; proceeding without it", e);
+ return false;
+ } finally {
+ probePool.shutdownNow();
+ }
+ }
+
+ /// Stage 1: local, no-network warmup. The network probe goes through
[QueryRouter] and stops at the
+ /// gathered DataTables, so it never exercises the SQL compile path's cold
start nor the response build +
+ /// JSON serialization the first real query pays. This compiles a throwaway
query and serializes a small
+ /// synthetic [BrokerResponseNative], warming both. Never throws.
+ private void warmUpLocal() {
+ try {
+ CalciteSqlCompiler.compileToBrokerRequest("SELECT 1");
+ BrokerResponseNative response = new BrokerResponseNative();
+ response.setResultTable(new ResultTable(
+ new DataSchema(new String[]{"warmup"}, new
DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.LONG}),
+ Collections.singletonList(new Object[]{1L})));
+ response.setNumDocsScanned(1);
+ response.toJsonString();
+ } catch (Exception e) {
+ LOGGER.debug("Local (stage 1) warmup failed; continuing", e);
+ }
+ }
+
+ /// Fires the given probe tasks concurrently on the shared pool and returns
the latencies (ms) of the
+ /// probes that completed successfully. Empty means the whole round was
unproductive (nothing routable
+ /// yet / all failed); a failed probe simply does not contribute.
+ private List<Long> runConcurrentRound(List<Callable<Long>> tasks,
ExecutorService pool) {
+ List<Future<Long>> futures = new ArrayList<>(tasks.size());
+ for (Callable<Long> task : tasks) {
+ futures.add(pool.submit(task));
+ }
+ List<Long> latencies = new ArrayList<>(futures.size());
+ for (Future<Long> future : futures) {
+ try {
+ Long elapsedMs = future.get(WARMUP_PROBE_TIMEOUT_MS + 1000L,
TimeUnit.MILLISECONDS);
+ if (elapsedMs != null && elapsedMs >= 0) {
+ latencies.add(elapsedMs);
+ }
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ return latencies;
+ } catch (Exception e) {
+ // Individual probe failed/timed out; it just does not count toward
the round. Cancel it so a probe
+ // that outlived the round's get()-timeout does not keep running
behind the next round (no-op if it
+ // already completed).
+ future.cancel(true);
+ }
+ }
+ return latencies;
+ }
+
+ /// Sleeps [#WARMUP_FAILURE_BACKOFF_MS] after an unproductive probe round.
Returns `false` if interrupted
+ /// (shutdown) so the caller stops promptly and readiness opens.
+ private boolean warmupBackoff() {
+ try {
+ Thread.sleep(WARMUP_FAILURE_BACKOFF_MS);
+ return true;
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ return false;
+ }
+ }
+
+ /// Default probe path: probe `SELECT * FROM "<t>" LIMIT 1` over the
configured tables or a greedy
+ /// set-cover of tables covering every routable server. Exits once at least
`minIterations` probes have
+ /// completed successfully -- a depth floor that guarantees enough
invocations to drive the query path's
+ /// JIT to its top tier -- OR the budget expires. A latency target is
deliberately not used: 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.
+ private boolean warmUpWithAutoSelect(BrokerWarmupConfig config, long
deadlineMs, ExecutorService pool,
+ int concurrency) {
+ long successfulProbes = 0;
+ int rounds = 0;
+ // Select the probe tables once and reuse them across rounds. Only
re-select while the result is empty
+ // (routing not populated yet); recomputing the greedy set-cover every
round would repeat O(tables) work
+ // up to minIterations times on a large-table tenant, for a result that
does not change once non-empty.
+ List<String> tables = List.of();
+ boolean coverageReported = false;
+ // Monotonic across rounds so the round-robin actually advances through
every covered table even at
+ // concurrency 1; resetting per round would probe only the first
`concurrency` tables forever.
+ int probeSeq = 0;
+ while (System.currentTimeMillis() < deadlineMs &&
!Thread.currentThread().isInterrupted()) {
+ if (tables.isEmpty()) {
+ tables = config.hasCustomTables() ?
filterRoutableTables(_routingManager, config.tables())
+ : selectProbeTables(_routingManager, config.maxTables());
+ }
+ if (tables.isEmpty()) {
+ // Converged but nothing routable yet: back off and retry rather than
declaring the broker warm --
+ // an empty routing table here would otherwise make the gate a no-op
precisely on a cold broker.
+ if (!warmupBackoff()) {
+ return false;
+ }
+ continue;
+ }
+ if (!coverageReported) {
+ // Report once, now that the probe tables are settled, whether they
route to every server.
+ coverageReported = true;
+ reportServerCoverage(_routingManager, _brokerMetrics,
config.hasCustomTables(), tables);
+ }
+ long remainingMs = deadlineMs - System.currentTimeMillis();
+ if (remainingMs <= 0) {
+ break;
+ }
+ rounds++;
+ long probeTimeoutMs = Math.min(remainingMs, WARMUP_PROBE_TIMEOUT_MS);
+ // A batch of `concurrency` probes, round-robin over the covered tables
(advancing across rounds via
+ // probeSeq): exercises every server (coverage) AND real concurrency
(contention) at once.
+ List<Callable<Long>> tasks = new ArrayList<>(concurrency);
+ for (String tableNameWithType : roundRobinBatch(tables, probeSeq,
concurrency)) {
+ tasks.add(() -> probe(compileProbe(tableNameWithType),
probeTimeoutMs));
+ }
+ probeSeq += concurrency;
+ List<Long> latencies = runConcurrentRound(tasks, pool);
+ if (latencies.isEmpty()) {
+ if (!warmupBackoff()) {
+ return false;
+ }
+ continue;
+ }
+ successfulProbes += latencies.size();
+ if (successfulProbes >= config.minIterations()) {
+ LOGGER.info("Broker warmup completed after {} round(s) at concurrency
{}; {} probes (floor {})", rounds,
+ concurrency, successfulProbes, config.minIterations());
+ return true;
+ }
+ }
+ LOGGER.warn("Broker warmup budget expired after {} round(s) at concurrency
{}; {} probes (floor {}). Proceeding "
+ + "to serve traffic.", rounds, concurrency, successfulProbes,
config.minIterations());
+ return false;
+ }
+
+ /// Custom-query probe path: run the operator-supplied queries repeatedly,
all of them every round, and
+ /// exit on the same `minIterations` probe-count floor as the default probe,
or the budget. A probe-count
+ /// floor is a cost-independent warm signal; latency plateauing is not (a
trivial query reaches a stable
+ /// latency while still cold), so it is deliberately not used here. An
expensive query that cannot reach
+ /// the floor within the budget exits on the budget, having warmed as much
as the budget allowed.
+ private boolean warmUpWithCustomQuery(BrokerWarmupConfig config, long
deadlineMs, ExecutorService pool,
+ int concurrency) {
+ // Validate-compile each configured query once so a bad one fails fast
(and never busy-spins), keeping
+ // only those that compile. Each probe below compiles its own fresh
request so nothing mutable is shared
+ // across the concurrent tasks.
+ List<String> queries = new ArrayList<>(config.queries().size());
+ for (String query : config.queries()) {
+ try {
+ CalciteSqlCompiler.compileToBrokerRequest(query);
+ queries.add(query);
+ } catch (Exception e) {
+ LOGGER.warn("Broker warmup query failed to compile; skipping it: {}",
query, e);
+ }
+ }
+ if (queries.isEmpty()) {
+ LOGGER.warn("Broker warmup: no configured query compiled; skipping
warmup");
+ return false;
+ }
+ long successfulProbes = 0;
+ int rounds = 0;
+ while (System.currentTimeMillis() < deadlineMs &&
!Thread.currentThread().isInterrupted()) {
+ long remainingMs = deadlineMs - System.currentTimeMillis();
+ if (remainingMs <= 0) {
+ break;
+ }
+ rounds++;
+ long probeTimeoutMs = Math.min(remainingMs, WARMUP_PROBE_TIMEOUT_MS);
+ // Run every configured query each round (each `concurrency` times, to
warm the concurrency step too).
+ List<Callable<Long>> tasks = new ArrayList<>(queries.size() *
concurrency);
+ for (String query : queries) {
+ for (int i = 0; i < concurrency; i++) {
+ tasks.add(() ->
probe(CalciteSqlCompiler.compileToBrokerRequest(query), probeTimeoutMs));
+ }
+ }
+ List<Long> latencies = runConcurrentRound(tasks, pool);
+ if (latencies.isEmpty()) {
+ if (!warmupBackoff()) {
+ return false;
+ }
+ continue;
+ }
+ successfulProbes += latencies.size();
+ // Same depth floor as the default probe: a probe-count floor is a
cost-independent warm signal (it
+ // drives the query path's JIT to its top tier), whereas latency
plateauing is not -- a trivial query
+ // reaches a stable latency while still cold. An expensive query that
cannot reach the floor within the
+ // budget simply exits on the budget, having warmed as much as the
budget allowed.
+ if (successfulProbes >= config.minIterations()) {
+ LOGGER.info("Broker warmup (custom query) completed after {} round(s)
at concurrency {}; {} probes "
+ + "(floor {})", rounds, concurrency, successfulProbes,
config.minIterations());
+ return true;
+ }
+ }
+ LOGGER.warn("Broker warmup (custom query) budget expired after {} round(s)
at concurrency {}; {} probes "
+ + "(floor {}). Proceeding to serve traffic.", rounds, concurrency,
successfulProbes, config.minIterations());
+ return false;
+ }
+
+ /// Compiles the default probe query for a single physical table.
+ private BrokerRequest compileProbe(String tableNameWithType) {
+ return CalciteSqlCompiler.compileToBrokerRequest("SELECT * FROM \"" +
tableNameWithType + "\" LIMIT 1");
+ }
+
+ /// Issues one probe query through [QueryRouter] and returns its wall-clock
duration in ms, or -1 on
+ /// failure. Builds the route the same way the normal path does, minus
auth/quota/logging.
+ private long probe(BrokerRequest brokerRequest, long timeoutMs) {
+ String tableName = brokerRequest.getQuerySource().getTableName();
+ try {
+ String rawTableName = TableNameBuilder.extractRawTableName(tableName);
+ TableType tableType =
TableNameBuilder.getTableTypeFromTableName(tableName);
+ long requestId = _requestIdGenerator.get();
+ TableRouteInfo routeInfo =
+ _implicitHybridTableRouteProvider.getTableRouteInfo(rawTableName,
_tableCache, _routingManager);
+ if (!(routeInfo instanceof ImplicitHybridTableRouteInfo) ||
!routeInfo.isExists()) {
+ return -1;
+ }
+ // A type-suffixed table (the default probe) targets exactly one type; a
raw name (possible with a
+ // custom query) targets whichever types actually exist. Each leg's
request must carry the
+ // type-suffixed table name -- that is what routing resolves against --
so build a per-leg request
+ // from the base query (as the normal path does) rather than reusing the
raw-named request, which
+ // would route to no server and yield an empty request map.
+ BrokerRequest offlineBrokerRequest = shouldRouteOffline(tableType,
routeInfo)
+ ? typedRequest(brokerRequest, routeInfo.getOfflineTableName()) :
null;
+ BrokerRequest realtimeBrokerRequest = shouldRouteRealtime(tableType,
routeInfo)
+ ? typedRequest(brokerRequest, routeInfo.getRealtimeTableName()) :
null;
+ if (offlineBrokerRequest == null && realtimeBrokerRequest == null) {
+ // Neither leg is routable right now (e.g. a hybrid table probed via
its offline name before the
+ // time boundary exists). submitQuery asserts at least one leg is
non-null, so skip cleanly rather
+ // than build an empty route -- the table is retried on the next round
once its leg appears.
+ return -1;
+ }
+ ImplicitHybridTableRouteInfo hybridRouteInfo =
(ImplicitHybridTableRouteInfo) routeInfo;
+ hybridRouteInfo.setOfflineBrokerRequest(offlineBrokerRequest);
+ hybridRouteInfo.setRealtimeBrokerRequest(realtimeBrokerRequest);
+ _implicitHybridTableRouteProvider.calculateRoutes(routeInfo,
_routingManager, offlineBrokerRequest,
+ realtimeBrokerRequest, requestId);
+ if (!routeInfo.isRouteExists()) {
+ return -1;
+ }
+ long startMs = System.currentTimeMillis();
+ // Scatter/gather (+ DataTable deserialize on receive) is the dominant
cost and does not need a
+ // QueryThreadContext; run it directly so its warmth always counts
toward the probe latency.
+ AsyncQueryResponse response = _queryRouter.submitQuery(requestId,
rawTableName, routeInfo, timeoutMs);
+ Map<ServerRoutingInstance, ServerResponse> finalResponses =
response.getFinalResponses();
+ // Best-effort: also warm the broker reduce path (result discarded) so
the first real query of this
+ // shape does not pay it. Isolated in its own try -- reduceOnDataTable
requires a QueryThreadContext
+ // and could fail for edge cases; a reduce failure must never fail the
probe, since the
+ // scatter/gather/deserialize warmth has already happened.
+ try (QueryThreadContext ignore = QueryThreadContext.open(
+ new QueryExecutionContext(QueryExecutionContext.QueryType.SSE,
requestId, Long.toString(requestId),
+ "warmup", startMs, Long.MAX_VALUE, Long.MAX_VALUE, _brokerId,
_brokerId, ""), _threadAccountant)) {
+ Map<ServerRoutingInstance, DataTable> dataTableMap = new
HashMap<>(finalResponses.size());
+ for (Map.Entry<ServerRoutingInstance, ServerResponse> entry :
finalResponses.entrySet()) {
+ DataTable dataTable = entry.getValue().getDataTable();
+ if (dataTable != null) {
+ dataTableMap.put(entry.getKey(), dataTable);
+ }
+ }
+ if (!dataTableMap.isEmpty()) {
+ _brokerReduceService.reduceOnDataTable(brokerRequest, brokerRequest,
dataTableMap, timeoutMs,
+ _brokerMetrics);
+ }
+ } catch (Exception reduceEx) {
+ LOGGER.debug("Warmup reduce step failed (continuing); scatter/gather
already warmed", reduceEx);
+ }
+ return System.currentTimeMillis() - startMs;
+ } catch (Exception e) {
+ if (e instanceof InterruptedException) {
+ // Shutdown interrupted this probe (via shutdownNow); restore the flag
so the pool worker unwinds.
+ Thread.currentThread().interrupt();
+ }
+ LOGGER.debug("Warmup probe failed for query on table: {}", tableName, e);
+ return -1;
+ }
+ }
+
+ /// Builds a per-leg probe request carrying the type-suffixed table name
that routing resolves against,
+ /// from the base (possibly raw-named) request. Deep-copies so the base
request is not mutated.
+ private static BrokerRequest typedRequest(BrokerRequest base, String
tableNameWithType) {
+ PinotQuery pinotQuery = base.getPinotQuery().deepCopy();
+ pinotQuery.getDataSource().setTableName(tableNameWithType);
+ return CalciteSqlCompiler.convertToBrokerRequest(pinotQuery);
+ }
+
+ /// Whether a probe routes the offline leg: the query targets offline (a
type-suffixed offline name) or is
+ /// untyped (a raw name, `tableType` null), AND an offline table exists.
Static so the offline / realtime /
+ /// hybrid decision is unit-testable against a mocked route without a live
cluster.
+ @VisibleForTesting
+ static boolean shouldRouteOffline(@Nullable TableType tableType,
TableRouteInfo routeInfo) {
+ return tableType != TableType.REALTIME && routeInfo.hasOffline();
+ }
+
+ /// Whether a probe routes the realtime leg: the query targets realtime or
is untyped, AND a realtime
+ /// table exists.
+ @VisibleForTesting
+ static boolean shouldRouteRealtime(@Nullable TableType tableType,
TableRouteInfo routeInfo) {
+ return tableType != TableType.OFFLINE && routeInfo.hasRealtime();
+ }
+
+ /// Keeps only the configured tables that this broker actually routes.
Static for unit testing.
+ @VisibleForTesting
+ static List<String> filterRoutableTables(RoutingManager routingManager,
List<String> tables) {
+ Set<String> routable = routingManager.getRoutableTables();
+ List<String> result = new ArrayList<>(tables.size());
+ for (String table : tables) {
+ if (routable.contains(table)) {
+ result.add(table);
+ } else {
+ LOGGER.warn("Broker warmup: configured table {} is not routable on
this broker; skipping", table);
+ }
+ }
+ return result;
+ }
+
+ /// Picks up to `maxTables` tables that between them cover as many routable
servers as possible.
+ ///
+ /// Greedy set cover rather than "all tables" or a configured list: warmup
needs every broker-to-server
+ /// channel and the whole query path exercised, and coverage delivers
exactly that while degrading
+ /// gracefully on a tenant with thousands of tables. Static so it can be
unit-tested against a mocked
+ /// [RoutingManager] without constructing a broker. When `maxTables` is too
small to span every server,
+ /// the leftover servers stay unprobed; [#reportServerCoverage] surfaces
that as a warning and a metric.
+ @VisibleForTesting
+ static List<String> selectProbeTables(RoutingManager routingManager, int
maxTables) {
+ if (maxTables <= 0) {
+ return List.of();
+ }
+ Set<String> uncoveredServers = routableServers(routingManager);
+ // Sorted so selection is deterministic across brokers and across restarts.
+ List<String> candidates = new
ArrayList<>(routingManager.getRoutableTables());
+ Collections.sort(candidates);
+ List<String> selected = new ArrayList<>(Math.min(maxTables,
candidates.size()));
+ for (String tableNameWithType : candidates) {
+ if (selected.size() >= maxTables || uncoveredServers.isEmpty()) {
+ break;
+ }
+ Set<String> serving =
routingManager.getServingInstances(tableNameWithType);
+ if (serving == null || serving.isEmpty()) {
+ continue;
+ }
+ if (uncoveredServers.removeAll(serving)) {
+ selected.add(tableNameWithType);
+ }
+ }
+ return selected;
+ }
+
+ /// Records [BrokerGauge#BROKER_WARMUP_UNCOVERED_SERVERS] and, when
non-zero, warns with the routable
+ /// servers the chosen probe tables do not route to. For the default
set-cover this means `maxTables` was
+ /// too small to span the server set (raise it); for a custom
`warmup.tables` list it means the list omits
+ /// some servers. Warmup still proceeds: those servers' broker-to-server
channels are opened
+ /// deterministically by startup pre-connect, and the serve-path JIT warmup
drives is server-agnostic, so
+ /// this is an observability signal, not a failure.
+ /// Static (taking its collaborators as parameters) so the gauge emission
and warning are unit-testable
+ /// against a mocked [RoutingManager] and [BrokerMetrics].
+ @VisibleForTesting
+ static void reportServerCoverage(RoutingManager routingManager,
BrokerMetrics brokerMetrics,
+ boolean hasCustomTables, List<String> tables) {
+ int totalServers = routableServers(routingManager).size();
+ Set<String> uncovered = uncoveredRoutableServers(routingManager, tables);
+
brokerMetrics.setValueOfGlobalGauge(BrokerGauge.BROKER_WARMUP_UNCOVERED_SERVERS,
uncovered.size());
+ if (!uncovered.isEmpty()) {
+ LOGGER.warn("Broker warmup probes {} table(s) covering {}/{} routable
server(s); {} left unprobed: {}. {}",
+ tables.size(), totalServers - uncovered.size(), totalServers,
uncovered.size(), uncovered,
+ hasCustomTables
+ ? "Add tables routing to them to
pinot.broker.startup.warmup.tables."
+ : "Raise pinot.broker.startup.warmup.maxTables to span every
server.");
+ }
+ }
+
+ /// Returns the routable servers that none of the given probe tables route
to. Empty means the tables
+ /// span every server this broker can reach. Static so it can be unit-tested
against a mocked
+ /// [RoutingManager].
+ @VisibleForTesting
+ static Set<String> uncoveredRoutableServers(RoutingManager routingManager,
Collection<String> tables) {
+ Set<String> uncovered = routableServers(routingManager);
+ for (String tableNameWithType : tables) {
+ Set<String> serving =
routingManager.getServingInstances(tableNameWithType);
+ if (serving != null) {
+ uncovered.removeAll(serving);
+ }
+ }
+ return uncovered;
+ }
+
+ /// The servers this broker actually routes to: the union of the serving
instances of its routable tables.
+ ///
+ /// Deliberately NOT `getRoutableServerInstanceMap()`, which is every
enabled server in the whole cluster
+ /// (no tenant filter) -- on a multi-tenant cluster that would count other
tenants' servers this broker
+ /// never queries, so the coverage metric would be permanently non-zero and
the "raise maxTables" warning
+ /// would be unactionable. This mirrors how startup pre-connect derives its
channels from routing.
+ @VisibleForTesting
+ static Set<String> routableServers(RoutingManager routingManager) {
+ Set<String> servers = new HashSet<>();
+ for (String tableNameWithType : routingManager.getRoutableTables()) {
+ Set<String> serving =
routingManager.getServingInstances(tableNameWithType);
+ if (serving != null) {
+ servers.addAll(serving);
+ }
+ }
+ return servers;
+ }
+
+ /// Returns the next `concurrency` items to probe, round-robin starting at
`startSeq`. Kept generic and
+ /// static so the rotation -- which must advance across rounds, not reset
each round, or only the first
+ /// `concurrency` items would ever be probed at low concurrency -- is
unit-testable without a live probe.
+ /// [Math#floorMod(int,int)] keeps the index valid even if `startSeq`
overflows to a negative value.
+ @VisibleForTesting
+ static <T> List<T> roundRobinBatch(List<T> items, int startSeq, int
concurrency) {
+ int size = items.size();
+ List<T> batch = new ArrayList<>(concurrency);
+ for (int i = 0; i < concurrency; i++) {
+ batch.add(items.get(Math.floorMod(startSeq + i, size)));
Review Comment:
**Nit — `roundRobinBatch` throws `ArithmeticException` (`/ by zero`) on an
empty list.**
`Math.floorMod(x, 0)` throws. Both call sites check for empty first, so this
is latent rather than live — but it is a `@VisibleForTesting static` helper
whose javadoc invites reuse. Suggest returning `List.of()` when
`items.isEmpty()` (or a `Preconditions.checkArgument`), plus the empty case in
`roundRobinBatchWrapsAndSurvivesOverflow`.
##########
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:
**Tests — a 10s slack on a 2s budget is loose enough to hide the overrun I
flagged in `runConcurrentRound`.**
This is the only test guarding the "budget is a hard ceiling" property, and
it tolerates a 6x overshoot. If the `get()`-timeout issue is fixed, tighten
this to something like `budget + 2s` so a regression actually fails the build.
If it is not fixed, then this slack is load-bearing and the comment should say
why.
##########
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:
**Config — documenting that the config layer mangles commas is honest, but
silently rewriting the operator's SQL is a trap.**
`a, b` becoming `a,b` is harmless for column lists, as the comment says. But
a query with a string literal containing `, ` is silently altered, and then
either fails to compile (WARN + skipped) or — worse — compiles into a
*different* query than the operator wrote. Since custom queries are meant to be
"representative production query shapes", literals containing commas are not
exotic.
Options: read the raw value through an un-split accessor if
`PinotConfiguration` exposes one, or fail loudly at startup when a configured
query contains a comma inside quotes. At minimum, mirror the caveat into the
`CommonConstants` comment — an operator reading the config reference will never
see this javadoc.
--
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]
