Add broker startup pre-connect for broker-to-server channels (SSE)

pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/BaseBrokerStarter.java

@@ -142,6 +142,10 @@ public abstract class BaseBrokerStarter implements ServiceStartable {
   /// desynchronize brokers on shared storage.
   private static final int RESPONSE_STORE_CLEANUP_INITIAL_DELAY_JITTER_DIVISOR = 4;
 
+  /// How often the pre-connect thread re-checks whether Helix has converged. Short enough not to add
+  /// meaningful delay to a fast startup, long enough not to hammer the Helix data accessor.
+  private static final long HELIX_CONVERGENCE_POLL_INTERVAL_MS = 200L;
+
   protected PinotConfiguration _brokerConf;
   protected List<ListenerConfig> _listenerConfigs;
   protected String _clusterName;
@@ -189,6 +193,22 @@ public abstract class BaseBrokerStarter implements ServiceStartable {
   protected BrokerGrpcServer _brokerGrpcServer;
   protected FailureDetector _failureDetector;
   protected ThreadAccountant _threadAccountant;
+  /// The Helix-convergence half of the service-status composite, held so startup pre-connect can wait
+  /// on exactly that signal -- convergence is the point at which routing, and the servers it
+  /// references, first exist.
+  @Nullable
+  private volatile ServiceStatus.ServiceStatusCallback _helixConvergenceCallback;
+  /// The background pre-connect thread, tracked so shutdown can interrupt it.
+  @Nullable
+  private volatile Thread _preConnectThread;
+  /// Whether startup pre-connect is enabled, and its budget. Read once in `start()`.
+  private boolean _preConnectEnabled;
+  private long _preConnectTimeoutMs;
+  /// Gates readiness when pre-connect is enabled: readiness reports STARTING until this is true, so a
+  /// broker is only Ready once it has connected to its servers (or the budget expired). Always flips
+  /// true, even on failure, so a rolling restart is never stalled by a broker held not-ready forever.
+  /// When pre-connect is disabled no gate callback is registered and this is irrelevant.
+  private volatile boolean _preConnectComplete;
 
   @Override
   public void init(PinotConfiguration brokerConf)
@@ -669,8 +689,16 @@ public void start()
         _brokerConf.getProperty(CommonConstants.Groovy.GROOVY_QUERY_STATIC_ANALYZER_CONFIG,
         _brokerConf.getProperty(CommonConstants.Groovy.GROOVY_ALL_STATIC_ANALYZER_CONFIG)));
 
-    // Register the service status handler
+    // Read the pre-connect config before registering the status handler: the handler adds the readiness
+    // gate only when pre-connect is enabled, and the gate must be in place before the handler is
+    // registered so there is no window where readiness is granted un-gated. This is single-stage (SSE)
+    // only; multi-stage uses a separate gRPC transport and is unaffected by this flag.
+    _preConnectEnabled = _brokerConf.getProperty(Broker.CONFIG_OF_BROKER_STARTUP_PRECONNECT_ENABLED,
+        Broker.DEFAULT_BROKER_STARTUP_PRECONNECT_ENABLED);
+    _preConnectTimeoutMs = _brokerConf.getProperty(Broker.CONFIG_OF_BROKER_STARTUP_PRECONNECT_TIMEOUT_MS,
+        Broker.DEFAULT_BROKER_STARTUP_PRECONNECT_TIMEOUT_MS);
     registerServiceStatusHandler();
+    startPreConnect();
 
     _isStarting = false;
     _brokerMetrics.addTimedValue(BrokerTimer.STARTUP_SUCCESS_DURATION_MS,
@@ -852,13 +880,114 @@ private void registerServiceStatusHandler() {
             Broker.DEFAULT_BROKER_MIN_RESOURCE_PERCENT_FOR_START);
 
     LOGGER.info("Registering service status handler");
-    ServiceStatus.setServiceStatusCallback(_instanceId, new ServiceStatus.MultipleCallbackServiceStatusCallback(
-        List.of(
-            new ServiceStatus.LifecycleServiceStatusCallback(this::isStarting, this::isShuttingDown),
-            new ServiceStatus.IdealStateAndCurrentStateMatchServiceStatusCallback(_participantHelixManager,
-                _clusterName, _instanceId, resourcesToMonitor, minResourcePercentForStartup),
-            new ServiceStatus.IdealStateAndExternalViewMatchServiceStatusCallback(_participantHelixManager,
-                _clusterName, _instanceId, resourcesToMonitor, minResourcePercentForStartup))));
+    // The two Helix callbacks are grouped into their own composite so startup pre-connect can wait on
+    // exactly the "Helix has converged" signal. CurrentState only reports ONLINE once the
+    // OFFLINE->ONLINE transition has returned, and that transition is what builds routing -- so
+    // convergence is the precondition for routing entries, and the servers they reference, existing.
+    // Behaviour is unchanged: MultipleCallbackServiceStatusCallback surfaces the first non-GOOD
+    // callback, so nesting the two Helix callbacks reports the same status as listing them flat.
+    _helixConvergenceCallback = new ServiceStatus.MultipleCallbackServiceStatusCallback(List.of(
+        new ServiceStatus.IdealStateAndCurrentStateMatchServiceStatusCallback(_participantHelixManager,
+            _clusterName, _instanceId, resourcesToMonitor, minResourcePercentForStartup),
+        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 (_preConnectEnabled) {
+      // The readiness gate. Reports STARTING (not a new status value) until pre-connect 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. This gates the existing
+      // readiness endpoint (getBrokerHealth -> ServiceStatus), which the Kubernetes startupProbe polls:
+      // the pod stays out of the Service until this reports GOOD, and while the startupProbe is failing
+      // Kubernetes does not run the liveness probe, so a warming pod is never killed. No health-endpoint
+      // or probe change is required.
+      callbacks.add(new ServiceStatus.ServiceStatusCallback() {
+        @Override
+        public ServiceStatus.Status getServiceStatus() {
+          return _preConnectComplete ? ServiceStatus.Status.GOOD : ServiceStatus.Status.STARTING;
+        }
+
+        @Override
+        public String getStatusDescription() {
+          return _preConnectComplete ? ServiceStatus.STATUS_DESCRIPTION_NONE
+              : "Pre-connecting broker-to-server channels";
+        }
+      });
+    }
+    ServiceStatus.setServiceStatusCallback(_instanceId,
+        new ServiceStatus.MultipleCallbackServiceStatusCallback(callbacks));
+  }
+
+  /// Runs startup server pre-connect on a background thread and opens the readiness gate
+  /// ([#_preConnectComplete]) 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 pre-connect throws or is interrupted: readiness held open indefinitely
+  /// would stall a rolling restart, a worse failure than serving a broker whose channels are not yet
+  /// warm. When disabled this is a no-op and readiness behaves exactly as before.
+  private void startPreConnect() {
+    if (!_preConnectEnabled) {
+      return;
+    }
+    _preConnectThread = 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 preConnectStartMs = 0L;
+      try {
+        long threadStartMs = System.currentTimeMillis();
+        awaitHelixConvergence();
+        preConnectStartMs = System.currentTimeMillis();
+        LOGGER.info("Helix converged after {} ms; pre-connecting server channels",
+            preConnectStartMs - threadStartMs);
+        int connected = _brokerRequestHandler.preConnectServers(preConnectStartMs + _preConnectTimeoutMs);
+        LOGGER.info("Startup server pre-connect opened {} channel(s); opening readiness", connected);
+      } catch (InterruptedException e) {
+        // Normal on shutdown; stopPreConnect() interrupts us.
+        Thread.currentThread().interrupt();
+        LOGGER.info("Startup server pre-connect interrupted before completion; opening readiness");
+      } catch (Throwable t) {
+        LOGGER.warn("Startup server pre-connect threw; opening readiness anyway", t);
+      } finally {
+        _preConnectComplete = true;
+        // Record the duration only if convergence was reached, so the metric measures the pre-connect work
+        // itself and never the (unbounded) convergence wait -- e.g. when shutdown interrupts the wait.
+        if (preConnectStartMs > 0L) {
+          _brokerMetrics.addTimedValue(BrokerTimer.STARTUP_PRECONNECT_DURATION_MS,
+              System.currentTimeMillis() - preConnectStartMs, TimeUnit.MILLISECONDS);
+        }
+      }
+    }, "broker-startup-preconnect");
+    _preConnectThread.setDaemon(true);
+    _preConnectThread.start();
+  }
+
+  /// Blocks until the Helix-convergence callbacks report GOOD -- the point at which routing entries and
+  /// the servers they reference exist. Deliberately **unbounded** and interruptible: a broker that never
+  /// converges is never Ready regardless of pre-connect, and shutdown interrupts this thread. Monitors
+  /// `brokerResource` only (partitions in {OFFLINE, ONLINE, DROPPED}); segment states live in the table
+  /// resources that servers monitor and cannot hold this up.
+  private void awaitHelixConvergence()
+      throws InterruptedException {
+    ServiceStatus.ServiceStatusCallback callback = _helixConvergenceCallback;
+    if (callback == null) {
+      return;
+    }
+    while (callback.getServiceStatus() != ServiceStatus.Status.GOOD) {
+      Thread.sleep(HELIX_CONVERGENCE_POLL_INTERVAL_MS);
+    }
+  }
+
+  /// Interrupts an in-flight pre-connect so shutdown never waits on it. Best effort: the thread is a
+  /// daemon and records its metric in a `finally` regardless.
+  private void stopPreConnect() {
+    Thread thread = _preConnectThread;
+    if (thread != null && thread.isAlive()) {
+      LOGGER.info("Interrupting in-flight startup server pre-connect for shutdown");
+      thread.interrupt();
+    }
   }
 
   private String getDefaultBrokerId() {
@@ -893,6 +1022,7 @@ protected boolean updatePortIfNeeded(Map<String, String> instanceConfigSimpleFie
   public void stop() {
     LOGGER.info("Shutting down Pinot broker");
     _isShuttingDown = true;
+    stopPreConnect();
 
     LOGGER.info("Disconnecting participant Helix manager");
     _participantHelixManager.disconnect();

pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BrokerRequestHandler.java

@@ -47,6 +47,16 @@ public interface BrokerRequestHandler {
 
   void shutDown();
 
+  /// Opens broker-to-server channels ahead of traffic so the first real query does not pay the blocking
+  /// connect -- and, when broker-to-server TLS is on, the handshake -- on its critical path. Called once
+  /// at startup after Helix has converged, when `pinot.broker.startup.preconnect.enabled` is set.
+  ///
+  /// Only the single-connection SSE handler opens Netty channels, so the default is a no-op. Returns the
+  /// number of channels connected before `deadlineMs` (an absolute [System#currentTimeMillis] value).
+  default int preConnectServers(long deadlineMs) {
+    return 0;
+  }
+
   BrokerResponse handleRequest(JsonNode request, @Nullable SqlNodeAndOptions sqlNodeAndOptions,
       @Nullable RequesterIdentity requesterIdentity, RequestContext requestContext, @Nullable HttpHeaders httpHeaders)
       throws Exception;

pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BrokerRequestHandlerDelegate.java

@@ -76,6 +76,13 @@ public void start() {
     }
   }
 
+  @Override
+  public int preConnectServers(long deadlineMs) {
+    // Only the single-stage handler owns the broker-to-server Netty channels; the multi-stage (gRPC)
+    // and time-series paths have nothing to pre-connect here.
+    return _singleStageBrokerRequestHandler.preConnectServers(deadlineMs);
+  }
+
   @Override
   public void shutDown() {
     _singleStageBrokerRequestHandler.shutDown();

pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/ServerPreConnector.java

@@ -0,0 +1,135 @@
+/**
+ * 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 com.google.common.annotations.VisibleForTesting;
+import com.google.common.util.concurrent.ThreadFactoryBuilder;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.List;
+import java.util.concurrent.CompletionService;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorCompletionService;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.function.BiPredicate;
+import java.util.function.Supplier;
+import javax.annotation.concurrent.ThreadSafe;
+import org.apache.pinot.core.transport.ServerInstance;
+import org.apache.pinot.spi.config.table.TableType;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/// Opens broker-to-server Netty channels ahead of query traffic, so the first real query does not pay
+/// the blocking `connect()` -- and, when broker-to-server TLS is on, the handshake -- on its critical
+/// path.
+///
+/// `ServerRoutingInstance` identity includes the table type, so OFFLINE and REALTIME are **separate**
+/// channels to the same physical server; both are connected here. Connecting an already-active channel
+/// is a no-op, so this is safe to call more than once.
+///
+/// Bounded on two axes so it can never stall startup: a capped thread pool, and a per-channel wait
+/// clamped to the caller's deadline. A server that is unreachable or itself restarting is logged and
+/// skipped -- the existing lazy-connect path still serves it. This class is stateless and thread-safe.
+///
+/// It takes its dependencies as functions rather than concrete `RoutingManager`/`QueryRouter` types so
+/// the parallelism, budget and failure handling can be unit-tested without a live broker.
+@ThreadSafe
+public class ServerPreConnector {
+  private static final Logger LOGGER = LoggerFactory.getLogger(ServerPreConnector.class);
+
+  /// Cap on the connect thread pool: a large tenant must not spawn a thread per server. Safe to exceed
+  /// the core count even on a 2- or 4-vCPU broker: each task is blocking connect + TLS handshake (mostly
+  /// network wait, with the actual I/O on Netty's event loop), and this runs during startup before any
+  /// query load, so the threads are almost entirely parked rather than contending for CPU.
+  @VisibleForTesting
+  static final int MAX_CONNECT_THREADS = 16;
+
+  private final Supplier<Collection<ServerInstance>> _routableServersSupplier;
+  private final BiPredicate<ServerInstance, TableType> _connectFn;
+
+  /// @param routableServersSupplier supplies the servers to connect, evaluated once per [#preConnect]
+  ///     call after the caller has ensured routing is built
+  /// @param connectFn opens the channel for one (server, table type) and returns whether it succeeded
+  public ServerPreConnector(Supplier<Collection<ServerInstance>> routableServersSupplier,
+      BiPredicate<ServerInstance, TableType> connectFn) {
+    _routableServersSupplier = routableServersSupplier;
+    _connectFn = connectFn;
+  }
+
+  /// Opens a channel to every routable server, for both table types, in parallel, bounded by
+  /// `deadlineMs` (an absolute [System#currentTimeMillis] value). Returns the number of channels
+  /// successfully connected. Never throws: a channel that fails or times out is logged and skipped.
+  public int preConnect(long deadlineMs) {
+    // Snapshot the routable-server view once. The supplier may return a live map view that another thread
+    // updates during startup; snapshotting keeps the channel count consistent with the tasks actually
+    // submitted below, so we never poll for phantom channels or under-count real ones.
+    List<ServerInstance> servers = new ArrayList<>(_routableServersSupplier.get());
+    if (servers.isEmpty() || System.currentTimeMillis() >= deadlineMs) {
+      return 0;
+    }
+    long startMs = System.currentTimeMillis();
+    int channelCount = servers.size() * TableType.values().length;
+    ExecutorService executor = Executors.newFixedThreadPool(Math.min(channelCount, MAX_CONNECT_THREADS),
+        new ThreadFactoryBuilder().setNameFormat("broker-preconnect-%d").setDaemon(true).build());
+    // A completion service hands channels back in the order they finish, not the order submitted, so a
+    // slow or unreachable server never blocks the counting of faster ones ahead of the shared deadline
+    // -- no head-of-line blocking, and no under-count of channels that already connected in parallel.
+    CompletionService<Boolean> completionService = new ExecutorCompletionService<>(executor);
+    int connected = 0;
+    try {
+      for (ServerInstance server : servers) {
+        for (TableType tableType : TableType.values()) {
+          completionService.submit(() -> _connectFn.test(server, tableType));
+        }
+      }
+      for (int i = 0; i < channelCount; i++) {
+        long remainingMs = deadlineMs - System.currentTimeMillis();
+        if (remainingMs <= 0) {
+          break;
+        }
+        try {
+          Future<Boolean> future = completionService.poll(remainingMs, TimeUnit.MILLISECONDS);
+          if (future == null) {
+            // Budget elapsed before the next channel finished; the rest fall back to the lazy path.
+            break;
+          }
+          if (Boolean.TRUE.equals(future.get())) {
+            connected++;
+          }
+        } catch (InterruptedException e) {
+          // Shutdown: stopPreConnect() interrupts us. Restore the flag and stop promptly.
+          Thread.currentThread().interrupt();
+          break;
+        } catch (ExecutionException e) {
+          // A server that is unreachable or itself restarting must not block startup.
+          LOGGER.debug("Pre-connect did not complete for one channel", e);
+        }
+      }
+    } finally {
+      executor.shutdownNow();
+    }
+    LOGGER.info("Broker pre-connected {}/{} channel(s) across {} server(s) in {} ms", connected,
+        channelCount, servers.size(), System.currentTimeMillis() - startMs);
+    return connected;
+  }
+}

pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/SingleConnectionBrokerRequestHandler.java

@@ -113,6 +113,17 @@ public void shutDown() {
     _brokerReduceService.shutDown();
   }
 
+  /// Opens a channel to every routable server, for both table types, taking the blocking connect -- and,
+  /// when broker-to-server TLS is on, the handshake -- off the first real query's critical path. The
+  /// caller guarantees Helix has converged, so [RoutingManager#getRoutableServerInstanceMap] already
+  /// reflects the servers this broker will route to. Never throws; returns the number of channels
+  /// connected before `deadlineMs`.
+  @Override
+  public int preConnectServers(long deadlineMs) {
+    return new ServerPreConnector(() -> _routingManager.getRoutableServerInstanceMap().values(),
+        _queryRouter::connect).preConnect(deadlineMs);
+  }
+
   @Override
   protected BrokerResponseNative processBrokerRequest(long requestId, BrokerRequest originalBrokerRequest,
       BrokerRequest serverBrokerRequest, TableRouteInfo route, long timeoutMs,