jineshparakh commented on code in PR #19500:
URL: https://github.com/apache/pinot/pull/19500#discussion_r3967076017


##########
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:
   Follow-up done — `runConcurrentRound` now distinguishes `ExecutionException` 
(a probe that *threw*) from a timeout and logs the first one of the whole run 
at WARN (a per-run latch, so no per-round spam), so a consistently broken probe 
is diagnosable without a debug build.



##########
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:
   Follow-up done — added a clause to the PR description: `submitQuery` 
deliberately seeds the adaptive server selector's routing stats 
(`ServerRoutingStatsManager`), so "pollutes no customer-facing surface" covers 
metrics/auth/quota/log, not routing stats.



##########
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:
   Follow-up done — corrected the wording. `min(minIterations, 
LOCAL_WARMUP_MAX_ITERATIONS)` equals the floor at the defaults, so it now reads 
"capped at 2000 (the cap only bites if minIterations is raised above 2000)" 
instead of "well below the floor".



-- 
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]

Reply via email to