jineshparakh commented on code in PR #19500:
URL: https://github.com/apache/pinot/pull/19500#discussion_r3965810661
##########
pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/SingleConnectionBrokerRequestHandler.java:
##########
@@ -117,6 +152,447 @@ public void shutDown() {
_brokerReduceService.shutDown();
}
+ /// Warms this broker's serve path before readiness is granted, in two
stages:
+ /// 1. a local, no-network pass ([#warmUpLocal]) that JIT-warms the
compile and response-serialization
+ /// paths, so they are warm even if no server is reachable; then
+ /// 2. a network probe -- repeatedly running real probe queries until they
have run enough times (the
+ /// `minIterations` depth floor) or the budget expires -- which warms
the scatter/gather/deserialize
+ /// and reduce paths and opens the broker-to-server channels lazily as
it goes.
+ ///
+ /// 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 (probe reduces also use a throwaway metrics
registry, see
+ /// [#WARMUP_NOOP_METRICS]).
+ ///
+ /// Contract for the readiness caller: this **never throws** and **always
returns by `deadlineMs`** (an
+ /// absolute [System#currentTimeMillis] value). Every wait -- each probe and
each `Future#get` -- is
+ /// bounded by the remaining budget, and once it is spent all outstanding
probes are cancelled, so a slow
+ /// or unreachable server can never stall the rolling restart this gates.
The return value (reached the
+ /// floor vs. hit the budget) is for logging and metrics only; the caller
proceeds either way.
+ ///
+ /// Runs on the single `broker-startup-warmup` thread. The probe pool it
creates is local to the call and
+ /// shut down before returning, so there is no cross-call shared mutable
state.
+ @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. Looped
enough to JIT the serialization
+ // path (a single pass leaves it interpreted), but capped well below the
network-probe depth floor and
+ // deadline-guarded so it stays a quick prelude and never eats the budget
the network probe needs.
+ warmUpLocal(Math.min(config.minIterations(), LOCAL_WARMUP_MAX_ITERATIONS),
deadlineMs);
+ int concurrency = Math.max(1, config.concurrency());
+ ExecutorService probePool = Executors.newFixedThreadPool(concurrency,
+ new
ThreadFactoryBuilder().setNameFormat("broker-warmup-probe-%d").setDaemon(true).build());
+ try {
+ return warmUpNetwork(config, deadlineMs, probePool, concurrency);
+ } catch (Exception e) {
+ LOGGER.warn("Broker warmup failed; proceeding without it", e);
+ return false;
+ } finally {
+ probePool.shutdownNow();
+ try {
+ // Wait briefly for interrupted probe workers to unwind before
returning, so a probe cannot outlive
+ // warmUp() and race the request handler being torn down on shutdown.
Bounded, so shutdown never hangs
+ // on a stuck probe.
+ if (!probePool.awaitTermination(PROBE_POOL_SHUTDOWN_WAIT_MS,
TimeUnit.MILLISECONDS)) {
+ LOGGER.debug("Warmup probe pool did not fully terminate within {}
ms; proceeding",
+ PROBE_POOL_SHUTDOWN_WAIT_MS);
+ }
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+
+ /// Stage 1: local, no-network warmup. The network probe goes through
[QueryRouter] and stops at the
+ /// gathered DataTables, so it never exercises the response build + JSON
serialization the first real query
+ /// pays (the compile path it does re-warm, but only via a real table). This
compiles a throwaway query and
+ /// serializes a small synthetic [BrokerResponseNative] `iterations` times
so the serialization path
+ /// reaches JIT rather than staying interpreted after a single pass.
`iterations` is capped by the caller
+ /// (see [#LOCAL_WARMUP_MAX_ITERATIONS]); the loop is guarded by both
`deadlineMs` and the interrupt flag,
+ /// so stage 1 stays a quick prelude and stops promptly when shutdown
interrupts the warmup thread. Never
+ /// throws; never runs the network.
+ private void warmUpLocal(int iterations, long deadlineMs) {
+ try {
+ for (int i = 0; i < iterations && System.currentTimeMillis() < deadlineMs
+ && !Thread.currentThread().isInterrupted(); i++) {
+ 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.
+ ///
+ /// The total wait is bounded by `deadlineMs`, never by a fixed per-probe
constant: each `get()` waits
+ /// only the remaining budget, and once the deadline has passed every
still-outstanding future is
+ /// cancelled instead of waited on. This is what makes the budget a hard
ceiling even when more tasks were
+ /// submitted than the pool has threads (so tasks queue) -- otherwise a
queued straggler could hold the
+ /// readiness gate well past the budget.
+ ///
+ /// Static and package-private so the budget / cancellation / interrupt
semantics can be unit-tested with
+ /// injected probe callables, without a live cluster.
+ @VisibleForTesting
+ static List<Long> runConcurrentRound(List<Callable<Long>> tasks,
ExecutorService pool, long deadlineMs) {
+ 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 (int i = 0; i < futures.size(); i++) {
+ long remainingMs = deadlineMs - System.currentTimeMillis();
+ if (remainingMs <= 0) {
+ // Budget spent: stop waiting and cancel everything still outstanding
so no probe outlives it.
+ cancelAll(futures, i);
+ break;
+ }
+ Future<Long> future = futures.get(i);
+ try {
+ Long elapsedMs = future.get(remainingMs, TimeUnit.MILLISECONDS);
+ if (elapsedMs != null && elapsedMs >= 0) {
+ latencies.add(elapsedMs);
+ }
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ cancelAll(futures, i);
+ return latencies;
+ } catch (Exception e) {
+ // Individual probe failed or overran the budget; it does not count,
and is cancelled so it does not
+ // keep running behind the next round (no-op if it already completed).
+ future.cancel(true);
+ }
+ }
+ return latencies;
+ }
+
+ private static void cancelAll(List<Future<Long>> futures, int from) {
+ for (int j = from; j < futures.size(); j++) {
+ futures.get(j).cancel(true);
+ }
+ }
+
+ /// 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;
+ }
+ }
+
+ /// The network warmup: probe the static `SELECT * FROM "<t>" LIMIT 1` over
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 warmUpNetwork(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 = selectProbeTables(_routingManager);
+ }
+ 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, tables);
+ }
+ if (System.currentTimeMillis() >= deadlineMs) {
+ break;
+ }
+ rounds++;
+ // 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. Each task
+ // derives its own timeout from the deadline when it actually starts
(see probeWithinDeadline), so a
+ // task that queued behind the pool cannot overrun the budget.
+ List<Callable<Long>> tasks = new ArrayList<>(concurrency);
+ for (String tableNameWithType : roundRobinBatch(tables, probeSeq,
concurrency)) {
+ tasks.add(() -> probeWithinDeadline(compileProbe(tableNameWithType),
deadlineMs));
+ }
+ probeSeq += concurrency;
+ List<Long> latencies = runConcurrentRound(tasks, pool, deadlineMs);
+ 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;
+ }
+ }
+ logBudgetExpiry(rounds, concurrency, successfulProbes,
config.minIterations());
+ return false;
+ }
+
+ /// Logs warmup budget expiry: WARN only when nothing warmed
(`successfulProbes == 0`), INFO otherwise.
+ /// Expiring after warming some probes is the expected steady-state exit for
a probe too expensive to reach
+ /// the floor within the budget -- it warmed as much as the budget allowed
-- so it is not warning-worthy.
+ private static void logBudgetExpiry(int rounds, int concurrency, long
successfulProbes, int floor) {
+ if (successfulProbes == 0) {
+ LOGGER.warn("Broker warmup budget expired after {} round(s) at
concurrency {}; {} probes (floor {}). "
+ + "Proceeding to serve traffic (nothing warmed).", rounds,
concurrency, successfulProbes, floor);
+ } else {
+ LOGGER.info("Broker warmup budget expired after {} round(s) at
concurrency {}; {} probes (floor {}). "
+ + "Proceeding to serve traffic.", rounds, concurrency,
successfulProbes, floor);
+ }
+ }
+
+ /// Compiles the default probe query for a single physical table.
+ private BrokerRequest compileProbe(String tableNameWithType) {
+ return CalciteSqlCompiler.compileToBrokerRequest("SELECT * FROM \"" +
tableNameWithType + "\" LIMIT 1");
+ }
+
+ /// Runs one probe with a timeout derived from the remaining budget **at the
moment the task starts**,
+ /// never a fixed per-probe constant. When more tasks are submitted than the
pool has threads they queue,
+ /// and a queued task may not start until much of the budget is already
gone; deriving the timeout here
+ /// (rather than when the round was built) is what keeps the budget a hard
ceiling. Returns -1 (uncounted)
+ /// if the budget is already spent when the task starts, so no probe is
fired past the deadline.
+ private long probeWithinDeadline(BrokerRequest brokerRequest, long
deadlineMs) {
+ long timeoutMs = Math.min(deadlineMs - System.currentTimeMillis(),
WARMUP_PROBE_TIMEOUT_MS);
+ return timeoutMs <= 0 ? -1 : probe(brokerRequest, timeoutMs);
+ }
+
+ /// 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;
+ }
+ // The probe table is always type-suffixed (the set-cover picks names
straight from getRoutableTables),
+ // so it targets exactly one type. 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). shouldRouteOffline/Realtime still handle the general case
for robustness.
+ 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;
+ }
+ // calculateRoutes sets the per-leg broker requests on the routeInfo
itself (nulling a leg whose
+ // routing table turns out empty), so we pass them as arguments and do
not pre-set them here.
+ _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.
+ // Side effect worth naming: submitQuery records
adaptive-server-selector stats (AsyncQueryResponse ->
+ // ServerRoutingStatsManager), so probes seed each server's latency EMA
before real traffic. A probe
+ // that gets a real response seeds the EMA with that server's (cold)
response latency; a probe that
+ // times out or errors seeds it with the full timeout (see
AsyncQueryResponse#getFinalResponses). On a
+ // healthy cluster every probe responds, so the seeds are comparable
across servers and the EMA decays
+ // to true warm latencies within a few real requests -- a small,
self-correcting bias. The one case
+ // that biases RELATIVE ordering is a server slow enough to time out
probes while its peers respond:
+ // it is seeded high and the selector routes less to it at first. That
is acceptable (it steers early
+ // traffic away from a genuinely slow server and self-corrects), and
still strictly better than the
+ // selector starting with no per-server history at all.
+ 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()) {
+ // Discard-only warmth: route metrics to the throwaway registry so
the probe does not pollute the
+ // broker's real counters (see WARMUP_NOOP_METRICS).
+ _brokerReduceService.reduceOnDataTable(brokerRequest, brokerRequest,
dataTableMap, timeoutMs,
+ WARMUP_NOOP_METRICS);
+ }
+ } 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();
+ }
+
+ /// Picks the fewest tables that between them cover **every** routable
server -- the one, only behavior:
+ /// warmup always covers the whole fleet, with no cap to tune.
+ ///
+ /// 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. The loop stops as soon
+ /// as every server is covered, so it self-limits to at most one table per
server (never the full table
+ /// list) even on a tenant with thousands of tables. Because the coverage
universe is exactly the servers
+ /// some routable table serves, this always achieves full coverage. Static
so it can be unit-tested against
+ /// a mocked [RoutingManager] without constructing a broker.
+ @VisibleForTesting
+ static List<String> selectProbeTables(RoutingManager routingManager) {
+ 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);
+ // Capacity bounded by the server count -- the true ceiling on how many
tables get selected.
+ List<String> selected = new ArrayList<>(Math.min(uncoveredServers.size(),
candidates.size()));
+ for (String tableNameWithType : candidates) {
+ if (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#STARTUP_WARMUP_UNCOVERED_SERVERS] -- the routable
servers the chosen probe tables
+ /// do not route to. Since [#selectProbeTables] always spans every routable
server, this is normally 0; it
+ /// is emitted as an observability signal confirming full coverage (and
would flag a coverage regression).
+ /// Static (taking its collaborators as parameters) so the gauge emission is
unit-testable against a mocked
+ /// [RoutingManager] and [BrokerMetrics].
+ @VisibleForTesting
+ static void reportServerCoverage(RoutingManager routingManager,
BrokerMetrics brokerMetrics,
Review Comment:
I think we should have this gauge as we can either hit the buget or the
iteration count without warming up everything. It'll serve as a good metric to
tune the defaults.
--
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]
