You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@pinot.apache.org by xi...@apache.org on 2026/10/20 19:05:28 UTC
(pinot) branch xiangfu0/codex2/ratis-query-dispatch created (now 7452b661c3b)
This is an automated email from the ASF dual-hosted git repository.
xiangfu0 pushed a change to branch xiangfu0/codex2/ratis-query-dispatch
in repository https://gitbox.apache.org/repos/asf/pinot.git
at 7452b661c3b Add bounded native broker query dispatch
This branch includes the following new commits:
new 7452b661c3b Add bounded native broker query dispatch
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
(pinot) 01/01: Add bounded native broker query dispatch
Posted by xi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
xiangfu0 pushed a commit to branch xiangfu0/codex2/ratis-query-dispatch
in repository https://gitbox.apache.org/repos/asf/pinot.git
commit 7452b661c3b7ac3bbefa20fc48c86723d5d98e58
Author: Xiang Fu <xi...@apache.org>
AuthorDate: Tue Oct 20 12:05:13 2026 -0700
Add bounded native broker query dispatch
---
.../broker/helix/NativeBrokerParticipant.java | 51 +-
.../requesthandler/NativeGrpcQueryDispatcher.java | 394 ++++++++++++++
.../broker/routing/NativeBrokerQueryAdmission.java | 51 ++
.../broker/transport/NativeDataTableDecoder.java | 347 ++++++++++++
.../broker/transport/NativeGrpcQueryClient.java | 462 ++++++++++++++++
.../broker/helix/NativeBrokerParticipantTest.java | 8 +-
.../NativeGrpcQueryDispatcherTest.java | 598 +++++++++++++++++++++
.../transport/NativeDataTableDecoderTest.java | 224 ++++++++
.../transport/NativeGrpcQueryClientTest.java | 343 ++++++++++++
pinot-coordination-ratis/README.md | 30 +-
.../core/query/reduce/BrokerReduceService.java | 20 +-
.../pinot/core/transport/grpc/GrpcQueryServer.java | 4 +
.../transport/grpc/NativeGrpcQueryGuardTest.java | 29 +-
13 files changed, 2538 insertions(+), 23 deletions(-)
diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/NativeBrokerParticipant.java b/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/NativeBrokerParticipant.java
index aff48e88640..2d461418a60 100644
--- a/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/NativeBrokerParticipant.java
+++ b/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/NativeBrokerParticipant.java
@@ -33,6 +33,7 @@ import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.locks.ReentrantReadWriteLock;
import javax.annotation.Nullable;
+import org.apache.pinot.broker.routing.NativeBrokerQueryAdmission;
import org.apache.pinot.broker.routing.NativeBrokerRoutingImage;
import org.apache.pinot.common.coordination.MetadataRpcConnection;
import org.apache.pinot.common.coordination.NativeMetadataConfig;
@@ -63,7 +64,8 @@ import static java.util.Objects.requireNonNull;
/// qualified serving settings enable local admission and committed readiness, never a network query resource.
/// The applied revision covers only the committed own-participant/control projection, never table or routing state.
/// Lifecycle operations and the single bounded heartbeat lane are serialized; status reads never perform RPCs.
-final class NativeBrokerParticipant implements AutoCloseable, ServiceStatus.ServiceStatusCallback {
+final class NativeBrokerParticipant implements AutoCloseable, ServiceStatus.ServiceStatusCallback,
+ NativeBrokerQueryAdmission {
private static final Logger LOGGER = LoggerFactory.getLogger(NativeBrokerParticipant.class);
private static final int MAX_ROUTING_TABLES = 128;
private static final int MAX_ROUTING_BYTES = 8 * 1024 * 1024;
@@ -127,7 +129,7 @@ final class NativeBrokerParticipant implements AutoCloseable, ServiceStatus.Serv
}
/// A request-owned immutable plan and fanout, bound to its original finite serving and selected-server leases.
- static final class QueryAdmission {
+ private static final class QueryAdmission implements NativeBrokerQueryAdmission.Admission {
private final NativeBrokerRoutingImage _image;
private final List<NativeBrokerRoutingImage.ServerRequest> _requests;
private final List<PeerLeaseBound> _serverLeases;
@@ -143,15 +145,18 @@ final class NativeBrokerParticipant implements AutoCloseable, ServiceStatus.Serv
_generation = generation;
}
- NativeBrokerRoutingImage image() {
+ @Override
+ public NativeBrokerRoutingImage image() {
return _image;
}
- List<NativeBrokerRoutingImage.ServerRequest> requests() {
+ @Override
+ public List<NativeBrokerRoutingImage.ServerRequest> requests() {
return _requests;
}
- NativeBrokerRoutingImage.PlanningMetadata planningMetadata(boolean ignoreCase) {
+ @Override
+ public NativeBrokerRoutingImage.PlanningMetadata planningMetadata(boolean ignoreCase) {
return _image.planningMetadata(ignoreCase);
}
}
@@ -493,7 +498,8 @@ final class NativeBrokerParticipant implements AutoCloseable, ServiceStatus.Serv
}
/// Originate once at authenticated request entry, before planning. Retries retain the returned deadline.
- NativeQueryDeadline originateDeadline(long requestedMillis) {
+ @Override
+ public NativeQueryDeadline originateDeadline(long requestedMillis) {
_admission.readLock().lock();
try {
serving();
@@ -503,7 +509,8 @@ final class NativeBrokerParticipant implements AutoCloseable, ServiceStatus.Serv
}
}
- QueryAdmission admitQuery(String table, UUID queryId, UUID attemptId, NativeQueryDeadline deadline) {
+ @Override
+ public Admission admitQuery(String table, UUID queryId, UUID attemptId, NativeQueryDeadline deadline) {
_admission.readLock().lock();
try {
ServingPermitGuard guard = serving();
@@ -527,10 +534,11 @@ final class NativeBrokerParticipant implements AutoCloseable, ServiceStatus.Serv
}
}
- void checkQueryDelivery(QueryAdmission admission) {
+ @Override
+ public void checkQueryDelivery(Admission admission) {
_admission.readLock().lock();
try {
- validateDelivery(admission);
+ validateDelivery(ownedAdmission(admission));
} finally {
_admission.readLock().unlock();
}
@@ -538,17 +546,38 @@ final class NativeBrokerParticipant implements AutoCloseable, ServiceStatus.Serv
/// The callback must be bounded and must not call lifecycle methods or acquire membership/routing monitors.
/// It executes under the admission read lock so a committed local revocation cannot race check-then-delivery.
- void deliverQuery(QueryAdmission admission, Runnable callback) {
+ @Override
+ public void deliverQuery(Admission admission, Runnable callback) {
requireNonNull(callback);
_admission.readLock().lock();
try {
- validateDelivery(admission);
+ validateDelivery(ownedAdmission(admission));
callback.run();
} finally {
_admission.readLock().unlock();
}
}
+ @Override
+ public long remainingMillis(Admission admission) {
+ _admission.readLock().lock();
+ try {
+ QueryAdmission owned = ownedAdmission(admission);
+ validateDelivery(owned);
+ long remaining = requireNonNull(_servingSettings).queryDeadlines()
+ .remainingMillis(owned._requests.getFirst().context().deadline());
+ validateDelivery(owned);
+ return remaining;
+ } finally {
+ _admission.readLock().unlock();
+ }
+ }
+
+ private static QueryAdmission ownedAdmission(Admission admission) {
+ checkArgument(admission instanceof QueryAdmission, "Admission was not issued by a native broker participant");
+ return (QueryAdmission) admission;
+ }
+
private void validateDelivery(QueryAdmission admission) {
requireNonNull(admission);
ServingPermitGuard guard = serving();
diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/NativeGrpcQueryDispatcher.java b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/NativeGrpcQueryDispatcher.java
new file mode 100644
index 00000000000..d170b181903
--- /dev/null
+++ b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/NativeGrpcQueryDispatcher.java
@@ -0,0 +1,394 @@
+/**
+ * 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.MoreExecutors;
+import java.io.IOException;
+import java.net.URI;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.ArrayBlockingQueue;
+import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.CancellationException;
+import java.util.concurrent.Semaphore;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.LongSupplier;
+import javax.annotation.Nullable;
+import org.apache.pinot.broker.routing.NativeBrokerQueryAdmission;
+import org.apache.pinot.broker.routing.NativeBrokerQueryAdmission.Admission;
+import org.apache.pinot.broker.routing.NativeBrokerRoutingImage;
+import org.apache.pinot.broker.transport.NativeDataTableDecoder;
+import org.apache.pinot.broker.transport.NativeGrpcQueryClient;
+import org.apache.pinot.common.datatable.DataTable;
+import org.apache.pinot.common.metrics.BrokerMetrics;
+import org.apache.pinot.common.proto.Server;
+import org.apache.pinot.common.request.BrokerRequest;
+import org.apache.pinot.common.response.broker.BrokerResponseNative;
+import org.apache.pinot.common.utils.config.QueryOptionsUtils;
+import org.apache.pinot.common.utils.grpc.ServerGrpcRequestBuilder;
+import org.apache.pinot.core.query.reduce.BrokerReduceService;
+import org.apache.pinot.core.routing.SegmentsToQuery;
+import org.apache.pinot.core.routing.TableRouteInfo;
+import org.apache.pinot.core.transport.ServerInstance;
+import org.apache.pinot.core.transport.ServerRoutingInstance;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.env.PinotConfiguration;
+import org.apache.pinot.spi.utils.CommonConstants.Query.Response;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import static com.google.common.base.Preconditions.checkArgument;
+import static com.google.common.base.Preconditions.checkState;
+import static java.util.Objects.requireNonNull;
+
+
+/// Dispatches one admitted OFFLINE query using exact captured routes and the ordinary Pinot DataTable reducer.
+/// Every required destination must return one complete, successful non-streaming response and successful EOF.
+/// Failure cancels the entire attempt. This infrastructure does not register a listener or authorize final HTTP
+/// output: its caller must retain the same admission through serialization and guarded network delivery.
+public final class NativeGrpcQueryDispatcher implements AutoCloseable {
+ private static final Logger LOGGER = LoggerFactory.getLogger(NativeGrpcQueryDispatcher.class);
+ private final String _brokerId;
+ private final NativeBrokerQueryAdmission _authority;
+ private final Transport _transport;
+ @Nullable
+ private final NativeGrpcQueryClient _ownedClient;
+ private final Limits _limits;
+ private final BrokerReduceService _reducer;
+ private final Semaphore _slots;
+ private final Set<Attempt> _attempts = new HashSet<>();
+ private final AtomicBoolean _reducerStopped = new AtomicBoolean();
+ private boolean _closed;
+
+ /// Bounds aggregate work, including cancelled receivers that have not yet exited.
+ public record Limits(int maxConcurrentQueries, int maxFanout, long maxResponseBytes) {
+ public Limits {
+ checkArgument(maxConcurrentQueries > 0 && maxConcurrentQueries <= 128
+ && maxFanout > 0 && maxFanout <= 128 && maxResponseBytes > 0
+ && maxResponseBytes <= 256L * 1024 * 1024, "Invalid native query dispatch limits");
+ }
+ }
+
+ @FunctionalInterface
+ interface Transport {
+ NativeGrpcQueryClient.Call submit(URI endpoint, Server.ServerRequest request, LongSupplier remainingMillis);
+ }
+
+ public NativeGrpcQueryDispatcher(PinotConfiguration config, String brokerId, NativeBrokerQueryAdmission authority,
+ NativeGrpcQueryClient client, Limits limits) {
+ this(config, brokerId, authority, client::submit, limits, client);
+ }
+
+ @VisibleForTesting
+ NativeGrpcQueryDispatcher(PinotConfiguration config, String brokerId, NativeBrokerQueryAdmission authority,
+ Transport transport, Limits limits) {
+ this(config, brokerId, authority, transport, limits, null);
+ }
+
+ private NativeGrpcQueryDispatcher(PinotConfiguration config, String brokerId, NativeBrokerQueryAdmission authority,
+ Transport transport, Limits limits, @Nullable NativeGrpcQueryClient ownedClient) {
+ _brokerId = requireNonNull(brokerId);
+ _authority = requireNonNull(authority);
+ _transport = requireNonNull(transport);
+ _ownedClient = ownedClient;
+ _limits = requireNonNull(limits);
+ _slots = new Semaphore(limits.maxConcurrentQueries());
+ _reducer = new BrokerReduceService(config);
+ }
+
+ /// The caller owns the normal QueryThreadContext and the original admission for the whole request lifetime.
+ public BrokerResponseNative dispatch(long requestId, String correlationId, BrokerRequest original,
+ BrokerRequest server, TableRouteInfo route, Admission admission, boolean trace) throws Exception {
+ Attempt attempt = begin();
+ try {
+ _authority.checkQueryDelivery(admission);
+ validateOptions(original);
+ List<Target> targets = requests(requestId, correlationId, server, route, admission, trace);
+ AtomicLong receivedBytes = new AtomicLong();
+ BlockingQueue<Received> completed = new ArrayBlockingQueue<>(targets.size());
+ for (Target target : targets) {
+ attempt.checkOpen();
+ _authority.checkQueryDelivery(admission);
+ NativeGrpcQueryClient.Call call =
+ _transport.submit(target.endpoint(), target.request(), () -> _authority.remainingMillis(admission));
+ attempt.add(call);
+ attempt.receive(() -> {
+ try {
+ DataTable table = receive(requestId, target, call, admission, attempt, receivedBytes);
+ _authority.deliverQuery(admission, () -> {
+ attempt.checkOpen();
+ completed.add(new Received(target.server(), table, null));
+ });
+ } catch (Exception e) {
+ completed.add(new Received(target.server(), null, e));
+ attempt.cancel();
+ }
+ });
+ }
+ Map<ServerRoutingInstance, DataTable> tables = new HashMap<>();
+ for (int i = 0; i < targets.size(); i++) {
+ // Anchor before reading the qualified budget so a pause cannot restart the receive timeout.
+ long anchor = System.nanoTime();
+ long remainingNanos = TimeUnit.MILLISECONDS.toNanos(_authority.remainingMillis(admission));
+ long waitNanos = remainingNanos - Math.max(0, System.nanoTime() - anchor);
+ Received received = completed.poll(Math.max(0, waitNanos), TimeUnit.NANOSECONDS);
+ if (received == null) {
+ throw new TimeoutException("Native query fanout did not complete within its original deadline");
+ }
+ if (received.failure() != null) {
+ throw received.failure();
+ }
+ _authority.deliverQuery(admission, () -> {
+ attempt.checkOpen();
+ checkState(tables.put(received.server(), received.table()) == null, "Repeated native server result");
+ });
+ }
+ attempt.checkOpen();
+ long remaining = _authority.remainingMillis(admission);
+ BrokerResponseNative result;
+ // Keep reduction, including group-by subtasks, on this attempt's stack. The legacy reducer's asynchronous
+ // timeout/cancel does not wait for worker exit and cannot determine when native query capacity is reusable.
+ try (var executor = MoreExecutors.newDirectExecutorService()) {
+ result = _reducer.reduceOnDataTable(original, server, tables, remaining, BrokerMetrics.get(), executor);
+ }
+ checkState(!result.isPartialResult(), "Native query reduction did not produce a complete result");
+ result.setNumServersQueried(targets.size());
+ result.setNumServersResponded(targets.size());
+ _authority.deliverQuery(admission, attempt::checkOpen);
+ return result;
+ } finally {
+ attempt.finish();
+ }
+ }
+
+ private List<Target> requests(long requestId, String correlationId, BrokerRequest server, TableRouteInfo route,
+ Admission admission, boolean trace) {
+ NativeBrokerRoutingImage image = requireNonNull(admission.image());
+ checkArgument(image.tableName().equals(route.getOfflineTableName()) && route.getRealtimeTableName() == null
+ && route.getRealtimeBrokerRequest() == null && route.getRealtimeRoutingTable() == null
+ && server.equals(route.getOfflineBrokerRequest()) && route.getUnavailableSegments().isEmpty()
+ && !route.isDisabled() && route.getNumPrunedSegmentsTotal() == 0
+ && image.tableName().equals(server.getPinotQuery().getDataSource().getTableName()),
+ "Native dispatch requires the complete captured physical OFFLINE route");
+ validateOptions(server);
+ List<NativeBrokerRoutingImage.ServerRequest> admitted = List.copyOf(admission.requests());
+ checkArgument(!admitted.isEmpty() && admitted.size() <= _limits.maxFanout(), "Invalid native query fanout");
+ var first = admitted.getFirst().context();
+ List<NativeBrokerRoutingImage.ServerRequest> expected =
+ image.requests(first.queryId(), first.attemptId(), requireNonNull(first.deadline()));
+ checkArgument(expected.size() == admitted.size(), "Native admission has incomplete segment coverage");
+ Map<ServerInstance, SegmentsToQuery> routes = requireNonNull(route.getOfflineRoutingTable());
+ checkArgument(routes.size() == expected.size(), "Native route dropped or added a destination");
+ Map<String, Map.Entry<ServerInstance, SegmentsToQuery>> byId = new HashMap<>();
+ routes.entrySet().forEach(entry -> checkArgument(
+ byId.put(entry.getKey().getInstanceId(), entry) == null, "Repeated native routing server"));
+ List<Target> targets = new ArrayList<>();
+ for (int i = 0; i < expected.size(); i++) {
+ var target = expected.get(i);
+ var context = target.context();
+ checkArgument(target.endpoint().equals(admitted.get(i).endpoint())
+ && context.toProto().equals(admitted.get(i).context().toProto())
+ && context.broker().participantId().equals(_brokerId), "Native identity envelope differs from admission");
+ var entry = requireNonNull(byId.remove(context.server().participantId()), "Missing admitted destination");
+ ServerInstance instance = entry.getKey();
+ SegmentsToQuery segments = entry.getValue();
+ checkArgument(target.endpoint().getHost().equals(instance.getHostname())
+ && target.endpoint().getPort() == instance.getGrpcPort()
+ && segments.getOptionalSegments().isEmpty()
+ && segments.getSegments().size() == context.requiredNames().size()
+ && new HashSet<>(segments.getSegments()).equals(new HashSet<>(context.requiredNames())),
+ "Native routing endpoint or segment coverage differs from admission");
+ Server.ServerRequest wire = new ServerGrpcRequestBuilder().setRequestId(requestId).setCid(correlationId)
+ .setBrokerId(_brokerId).setEnableTrace(trace).setEnableStreaming(false).setBrokerRequest(server)
+ .setSegments(context.requiredNames()).build().toBuilder().setNativeQueryContext(context.toProto()).build();
+ targets.add(new Target(target.endpoint(),
+ new ServerRoutingInstance(instance.getInstanceId(), instance.getHostname(), instance.getGrpcPort(),
+ TableType.OFFLINE, true), wire));
+ }
+ return List.copyOf(targets);
+ }
+
+ private static void validateOptions(BrokerRequest request) {
+ Map<String, String> options = request.getPinotQuery().getQueryOptions();
+ checkArgument(options == null || !QueryOptionsUtils.isIgnoreMissingSegments(options),
+ "Native dispatch cannot ignore required missing segments");
+ checkArgument(QueryOptionsUtils.getTableSampler(options) == null,
+ "Native dispatch cannot sample required segment coverage");
+ }
+
+ private DataTable receive(long requestId, Target target, NativeGrpcQueryClient.Call call, Admission admission,
+ Attempt attempt, AtomicLong receivedBytes) throws IOException {
+ Iterator<Server.ServerResponse> responses = call.responses();
+ checkReceiving(admission, attempt);
+ checkState(responses.hasNext(), "Native server returned no DataTable");
+ checkReceiving(admission, attempt);
+ Server.ServerResponse response = responses.next();
+ checkReceiving(admission, attempt);
+ checkState(Response.ResponseType.NON_STREAMING.equals(
+ response.getMetadataOrDefault(Response.MetadataKeys.RESPONSE_TYPE, "")),
+ "Native server returned an unexpected response type");
+ checkState(receivedBytes.addAndGet(response.getSerializedSize()) <= _limits.maxResponseBytes(),
+ "Native query response exceeds its aggregate byte bound");
+ DataTable table = NativeDataTableDecoder.decode(response.getPayload().asReadOnlyByteBuffer());
+ checkReceiving(admission, attempt);
+ checkState(table.getExceptions().isEmpty() && table.getDataSchema() != null
+ && Long.toString(requestId).equals(table.getMetadata().get(DataTable.MetadataKey.REQUEST_ID.getName()))
+ && Integer.toString(target.request().getSegmentsCount()).equals(
+ table.getMetadata().get(DataTable.MetadataKey.NUM_SEGMENTS_QUERIED.getName()))
+ && !Boolean.parseBoolean(table.getMetadata().get(DataTable.MetadataKey.INCOMPLETE_MERGE.getName())),
+ "Native server result is incomplete or does not match the required request");
+ checkState(!responses.hasNext(), "Native server returned multiple DataTables");
+ checkReceiving(admission, attempt);
+ return table;
+ }
+
+ private void checkReceiving(Admission admission, Attempt attempt) {
+ attempt.checkOpen();
+ _authority.checkQueryDelivery(admission);
+ }
+
+ private synchronized Attempt begin() {
+ checkState(!_closed, "Native dispatcher is closed");
+ checkState(_slots.tryAcquire(), "Native query dispatch capacity exhausted");
+ Attempt attempt = new Attempt();
+ _attempts.add(attempt);
+ return attempt;
+ }
+
+ private void release(Attempt attempt) {
+ boolean stop;
+ synchronized (this) {
+ checkState(_attempts.remove(attempt), "Native query attempt was already released");
+ _slots.release();
+ stop = _closed && _attempts.isEmpty();
+ }
+ if (stop) {
+ stopReducer();
+ }
+ }
+
+ @Override
+ public void close() {
+ List<Attempt> attempts;
+ synchronized (this) {
+ _closed = true;
+ attempts = List.copyOf(_attempts);
+ }
+ attempts.forEach(Attempt::cancel);
+ if (_ownedClient != null) {
+ _ownedClient.close();
+ }
+ if (attempts.isEmpty()) {
+ stopReducer();
+ }
+ }
+
+ private void stopReducer() {
+ if (_reducerStopped.compareAndSet(false, true)) {
+ _reducer.shutDown();
+ }
+ }
+
+ private record Target(URI endpoint, ServerRoutingInstance server, Server.ServerRequest request) {
+ }
+
+ private record Received(ServerRoutingInstance server, @Nullable DataTable table, @Nullable Exception failure) {
+ }
+
+ /// Keeps query capacity until both dispatch and every receiver have exited, even after cancellation.
+ private final class Attempt {
+ private final List<NativeGrpcQueryClient.Call> _calls = new ArrayList<>();
+ private final Set<Thread> _receivers = new HashSet<>();
+ private volatile boolean _cancelled;
+ private boolean _finished;
+
+ private void checkOpen() {
+ if (_cancelled) {
+ throw new CancellationException("Native query attempt cancelled");
+ }
+ }
+
+ private synchronized void add(NativeGrpcQueryClient.Call call) {
+ _calls.add(call);
+ if (_cancelled) {
+ call.close();
+ checkOpen();
+ }
+ }
+
+ private synchronized void receive(Runnable receive) {
+ checkOpen();
+ Thread thread = Thread.ofVirtual().name("native-query-receiver").unstarted(() -> {
+ try {
+ checkOpen();
+ receive.run();
+ } finally {
+ boolean release;
+ synchronized (this) {
+ _receivers.remove(Thread.currentThread());
+ release = _finished && _receivers.isEmpty();
+ }
+ if (release) {
+ NativeGrpcQueryDispatcher.this.release(this);
+ }
+ }
+ });
+ _receivers.add(thread);
+ try {
+ thread.start();
+ } catch (RuntimeException | Error e) {
+ _receivers.remove(thread);
+ throw e;
+ }
+ }
+
+ private synchronized void cancel() {
+ _cancelled = true;
+ for (NativeGrpcQueryClient.Call call : _calls) {
+ try {
+ call.close();
+ } catch (RuntimeException e) {
+ LOGGER.warn("Could not close a cancelled native query call", e);
+ }
+ }
+ _receivers.forEach(Thread::interrupt);
+ }
+
+ private void finish() {
+ cancel();
+ boolean release;
+ synchronized (this) {
+ _finished = true;
+ release = _receivers.isEmpty();
+ }
+ if (release) {
+ NativeGrpcQueryDispatcher.this.release(this);
+ }
+ }
+ }
+}
diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/NativeBrokerQueryAdmission.java b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/NativeBrokerQueryAdmission.java
new file mode 100644
index 00000000000..6883ffabed5
--- /dev/null
+++ b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/NativeBrokerQueryAdmission.java
@@ -0,0 +1,51 @@
+/**
+ * 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.routing;
+
+import java.util.List;
+import java.util.UUID;
+import org.apache.pinot.common.request.context.NativeQueryDeadline;
+
+
+/// Broker-local query authority shared by planning, transport and delivery. Implementations bind opaque admissions
+/// to one process registration and perform only local, bounded checks; these methods must never initiate metadata RPCs.
+/// Callers retain the originating deadline and admission through final delivery instead of acquiring fresh authority.
+public interface NativeBrokerQueryAdmission {
+ NativeQueryDeadline originateDeadline(long requestedMillis);
+
+ Admission admitQuery(String table, UUID queryId, UUID attemptId, NativeQueryDeadline deadline);
+
+ void checkQueryDelivery(Admission admission);
+
+ /// Runs a bounded, nonblocking callback atomically with respect to local authority invalidation.
+ void deliverQuery(Admission admission, Runnable callback);
+
+ /// Remaining shared deadline budget, after checking the original local permit and selected-server leases.
+ long remainingMillis(Admission admission);
+
+ /// Immutable captured planning inputs and exact destinations. Only the issuing authority can validate this token;
+ /// constructing another implementation of this interface does not create serving authority.
+ interface Admission {
+ NativeBrokerRoutingImage image();
+
+ List<NativeBrokerRoutingImage.ServerRequest> requests();
+
+ NativeBrokerRoutingImage.PlanningMetadata planningMetadata(boolean ignoreCase);
+ }
+}
diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/transport/NativeDataTableDecoder.java b/pinot-broker/src/main/java/org/apache/pinot/broker/transport/NativeDataTableDecoder.java
new file mode 100644
index 00000000000..c580de039f7
--- /dev/null
+++ b/pinot-broker/src/main/java/org/apache/pinot/broker/transport/NativeDataTableDecoder.java
@@ -0,0 +1,347 @@
+/**
+ * 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.transport;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
+import java.util.HashSet;
+import java.util.Set;
+import org.apache.pinot.common.datatable.DataTable;
+import org.apache.pinot.common.datatable.DataTable.MetadataKey;
+import org.apache.pinot.common.datatable.DataTableFactory;
+import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
+
+
+/// Validates the complete V4 frame before invoking Pinot's ordinary decoder. Native responses have independent
+/// bounds on sections, nested allocation lengths, rows, columns and referenced variable data. MAP, OBJECT and
+/// UNKNOWN results are unsupported because their deferred codecs have their own allocation formats. This does
+/// not change the legacy decoder or establish a heap bound for arbitrary query reduction or result formatting.
+/// The input must remain immutable for the duration of this call; its position and limit are preserved.
+public final class NativeDataTableDecoder {
+ private static final int HEADER_BYTES = 13 * Integer.BYTES;
+ private static final int MAX_PAYLOAD_BYTES = 16 * 1024 * 1024;
+ private static final int MAX_ROWS = 1_000_000;
+ private static final int MAX_COLUMNS = 1024;
+ private static final int MAX_DICTIONARY_ENTRIES = 262_144;
+ private static final int MAX_VALUE_BYTES = 1024 * 1024;
+ private static final int MAX_METADATA_BYTES = 64 * 1024;
+ private static final int MAX_DECIMAL_BYTES = 4096;
+ private static final long MAX_REFERENCED_BYTES = 64L * 1024 * 1024;
+
+ private NativeDataTableDecoder() {
+ }
+
+ public static DataTable decode(ByteBuffer payload) throws IOException {
+ try {
+ ByteBuffer input = payload.slice().order(ByteOrder.BIG_ENDIAN);
+ require(input.remaining() >= HEADER_BYTES + Integer.BYTES && input.remaining() <= MAX_PAYLOAD_BYTES);
+ require(input.getInt() == DataTableFactory.VERSION_4);
+ int rows = bounded(input.getInt(), MAX_ROWS);
+ int columns = bounded(input.getInt(), MAX_COLUMNS);
+ ByteBuffer[] sections = new ByteBuffer[5];
+ int next = HEADER_BYTES;
+ for (int i = 0; i < sections.length; i++) {
+ int offset = input.getInt();
+ int length = input.getInt();
+ require(offset == next && length >= 0 && length <= input.limit() - next);
+ sections[i] = slice(input, offset, length);
+ next += length;
+ }
+ ByteBuffer tail = slice(input, next, input.limit() - next);
+ int metadataLength = bounded(tail.getInt(), MAX_METADATA_BYTES);
+ require(metadataLength == tail.remaining());
+ validateExceptions(sections[0]);
+ int dictionarySize = validateDictionary(sections[1]);
+ ColumnDataType[] types = validateSchema(sections[2], columns);
+ validateMetadata(tail);
+ validateValues(rows, types, dictionarySize, sections[3], sections[4]);
+ return DataTableFactory.getDataTable(payload.slice().order(ByteOrder.BIG_ENDIAN));
+ } catch (RuntimeException e) {
+ throw new IOException("Invalid or unsupported native V4 DataTable", e);
+ }
+ }
+
+ private static void validateExceptions(ByteBuffer input) throws IOException {
+ if (!input.hasRemaining()) {
+ return;
+ }
+ int count = count(input, 256, 8);
+ Set<Integer> codes = new HashSet<>();
+ for (int i = 0; i < count; i++) {
+ require(codes.add(input.getInt()));
+ skipString(input, MAX_METADATA_BYTES);
+ }
+ require(!input.hasRemaining());
+ }
+
+ private static int validateDictionary(ByteBuffer input) throws IOException {
+ if (!input.hasRemaining()) {
+ return 0;
+ }
+ int count = count(input, MAX_DICTIONARY_ENTRIES, Integer.BYTES);
+ for (int i = 0; i < count; i++) {
+ skipString(input, MAX_VALUE_BYTES);
+ }
+ require(!input.hasRemaining());
+ return count;
+ }
+
+ private static ColumnDataType[] validateSchema(ByteBuffer input, int columns) throws IOException {
+ if (!input.hasRemaining()) {
+ require(columns == 0);
+ return new ColumnDataType[0];
+ }
+ int count = count(input, MAX_COLUMNS, 2 * Integer.BYTES);
+ require(count == columns);
+ for (int i = 0; i < count; i++) {
+ skipString(input, 16 * 1024);
+ }
+ ColumnDataType[] types = new ColumnDataType[count];
+ for (int i = 0; i < count; i++) {
+ int length = bounded(input.getInt(), 64);
+ require(length <= input.remaining());
+ byte[] name = new byte[length];
+ input.get(name);
+ ColumnDataType type = ColumnDataType.valueOf(new String(name, UTF_8)).getStoredType();
+ require(type != ColumnDataType.OBJECT && type != ColumnDataType.UNKNOWN && type != ColumnDataType.MAP);
+ types[i] = type;
+ }
+ require(!input.hasRemaining());
+ return types;
+ }
+
+ private static void validateMetadata(ByteBuffer input) throws IOException {
+ if (!input.hasRemaining()) {
+ return;
+ }
+ int count = count(input, MetadataKey.values().length, Integer.BYTES + 1);
+ Set<MetadataKey> keys = new HashSet<>();
+ for (int i = 0; i < count; i++) {
+ MetadataKey key = MetadataKey.getById(input.getInt());
+ require(key != null && keys.add(key));
+ switch (key.getValueType()) {
+ case INT -> skip(input, Integer.BYTES);
+ case LONG -> skip(input, Long.BYTES);
+ case BOOLEAN -> {
+ byte value = input.get();
+ require(value == 0 || value == 1);
+ }
+ case STRING -> skipString(input, MAX_METADATA_BYTES);
+ default -> throw new IOException("Unsupported native metadata value");
+ }
+ }
+ require(!input.hasRemaining());
+ }
+
+ private static void validateValues(int rows, ColumnDataType[] types, int dictionarySize, ByteBuffer fixed,
+ ByteBuffer variable) throws IOException {
+ int rowBytes = 0;
+ for (ColumnDataType type : types) {
+ rowBytes += width(type);
+ }
+ long dataBytes = (long) rows * rowBytes;
+ require((rows == 0 || rowBytes > 0) && dataBytes <= fixed.limit());
+ long bitmapBytes = fixed.limit() - dataBytes;
+ require(bitmapBytes == 0 || bitmapBytes == (long) types.length * 2 * Integer.BYTES);
+ Budget budget = new Budget();
+ for (int row = 0; row < rows; row++) {
+ for (ColumnDataType type : types) {
+ switch (type) {
+ case INT, FLOAT -> skip(fixed, Integer.BYTES);
+ case LONG, DOUBLE -> skip(fixed, Long.BYTES);
+ case STRING -> dictionaryId(fixed.getInt(), dictionarySize);
+ default -> validateVariable(type, variable, fixed.getInt(), fixed.getInt(), dictionarySize, budget);
+ }
+ }
+ }
+ if (bitmapBytes != 0) {
+ for (int i = 0; i < types.length; i++) {
+ int offset = fixed.getInt();
+ int length = bounded(fixed.getInt(), 256 * 1024);
+ ByteBuffer bitmap = slice(variable, offset, length);
+ budget.add(length);
+ if (length != 0) {
+ validateBitmap(bitmap, rows, budget);
+ }
+ }
+ }
+ require(!fixed.hasRemaining());
+ }
+
+ private static void validateVariable(ColumnDataType type, ByteBuffer variable, int offset, int length,
+ int dictionarySize, Budget budget) throws IOException {
+ bounded(length, MAX_VALUE_BYTES);
+ int elementBytes = switch (type) {
+ case INT_ARRAY, FLOAT_ARRAY, STRING_ARRAY -> Integer.BYTES;
+ case LONG_ARRAY, DOUBLE_ARRAY -> Long.BYTES;
+ default -> 1;
+ };
+ if (type == ColumnDataType.BYTES || type == ColumnDataType.BIG_DECIMAL) {
+ ByteBuffer value = slice(variable, offset, length);
+ budget.add(length);
+ if (type == ColumnDataType.BIG_DECIMAL) {
+ validateDecimal(value.remaining());
+ }
+ return;
+ }
+ require(type == ColumnDataType.INT_ARRAY || type == ColumnDataType.LONG_ARRAY
+ || type == ColumnDataType.FLOAT_ARRAY || type == ColumnDataType.DOUBLE_ARRAY
+ || type == ColumnDataType.STRING_ARRAY || type == ColumnDataType.BYTES_ARRAY
+ || type == ColumnDataType.BIG_DECIMAL_ARRAY);
+ bounded(length, MAX_ROWS);
+ budget.add((long) length * 16);
+ if (type == ColumnDataType.BYTES_ARRAY || type == ColumnDataType.BIG_DECIMAL_ARRAY) {
+ ByteBuffer values = slice(variable, offset, variable.limit() - offset);
+ require(length <= values.remaining() / Integer.BYTES);
+ for (int i = 0; i < length; i++) {
+ int size = bounded(values.getInt(), MAX_VALUE_BYTES);
+ if (type == ColumnDataType.BIG_DECIMAL_ARRAY) {
+ validateDecimal(size);
+ }
+ skip(values, size);
+ budget.add(size);
+ }
+ } else {
+ ByteBuffer values = slice(variable, offset, Math.toIntExact((long) length * elementBytes));
+ if (type == ColumnDataType.STRING_ARRAY) {
+ for (int i = 0; i < length; i++) {
+ dictionaryId(values.getInt(), dictionarySize);
+ }
+ }
+ }
+ }
+
+ /// Check portable Roaring containers before the ordinary DataTable null-bitmap reader can allocate from them.
+ private static void validateBitmap(ByteBuffer input, int rows, Budget budget) throws IOException {
+ input.order(ByteOrder.LITTLE_ENDIAN);
+ int cookie = input.getInt();
+ boolean runs = (cookie & 0xffff) == 12347;
+ require(runs || cookie == 12346);
+ int containers = runs ? (cookie >>> 16) + 1 : input.getInt();
+ bounded(containers, (rows + 65535) / 65536);
+ byte[] runFlags = new byte[runs ? (containers + 7) / 8 : 0];
+ input.get(runFlags);
+ require(containers <= input.remaining() / 4);
+ int[] keys = new int[containers];
+ int[] cardinalities = new int[containers];
+ for (int i = 0; i < containers; i++) {
+ keys[i] = Short.toUnsignedInt(input.getShort());
+ cardinalities[i] = Short.toUnsignedInt(input.getShort()) + 1;
+ require((i == 0 || keys[i] > keys[i - 1]) && ((long) keys[i] << 16) < rows);
+ budget.add((long) cardinalities[i] * Integer.BYTES);
+ }
+ int[] offsets = new int[containers];
+ boolean offsetsPresent = !runs || containers >= 4;
+ if (offsetsPresent) {
+ for (int i = 0; i < containers; i++) {
+ offsets[i] = input.getInt();
+ }
+ }
+ for (int i = 0; i < containers; i++) {
+ require(!offsetsPresent || offsets[i] == input.position());
+ int actualCardinality = 0;
+ int maximum = -1;
+ if (runs && (runFlags[i / 8] & (1 << (i % 8))) != 0) {
+ int count = Short.toUnsignedInt(input.getShort());
+ require(count <= input.remaining() / 4);
+ for (int j = 0; j < count; j++) {
+ int start = Short.toUnsignedInt(input.getShort());
+ int length = Short.toUnsignedInt(input.getShort());
+ require(start > maximum && start + length <= 65535);
+ maximum = start + length;
+ actualCardinality += length + 1;
+ }
+ } else if (cardinalities[i] > 4096) {
+ require(input.remaining() >= 8192);
+ for (int j = 0; j < 1024; j++) {
+ long word = input.getLong();
+ actualCardinality += Long.bitCount(word);
+ if (word != 0) {
+ maximum = j * 64 + 63 - Long.numberOfLeadingZeros(word);
+ }
+ }
+ } else {
+ require(cardinalities[i] <= input.remaining() / Short.BYTES);
+ for (int j = 0; j < cardinalities[i]; j++) {
+ int value = Short.toUnsignedInt(input.getShort());
+ require(value > maximum);
+ maximum = value;
+ }
+ actualCardinality = cardinalities[i];
+ }
+ require(actualCardinality == cardinalities[i] && (((long) keys[i] << 16) + maximum) < rows);
+ }
+ require(!input.hasRemaining());
+ }
+
+ private static void validateDecimal(int length) throws IOException {
+ require(length >= 3 && length <= MAX_DECIMAL_BYTES);
+ }
+
+ private static int width(ColumnDataType type) {
+ return type == ColumnDataType.INT || type == ColumnDataType.FLOAT || type == ColumnDataType.STRING ? 4 : 8;
+ }
+
+ private static void dictionaryId(int id, int size) throws IOException {
+ require(id >= 0 && id < size);
+ }
+
+ private static int count(ByteBuffer input, int maximum, int minimumEntryBytes) throws IOException {
+ int count = bounded(input.getInt(), maximum);
+ require(count <= input.remaining() / minimumEntryBytes);
+ return count;
+ }
+
+ private static void skipString(ByteBuffer input, int maximum) throws IOException {
+ skip(input, bounded(input.getInt(), maximum));
+ }
+
+ private static void skip(ByteBuffer input, int length) throws IOException {
+ require(length >= 0 && length <= input.remaining());
+ input.position(input.position() + length);
+ }
+
+ private static ByteBuffer slice(ByteBuffer input, int offset, int length) throws IOException {
+ require(offset >= 0 && offset <= input.limit() && length >= 0 && length <= input.limit() - offset);
+ return input.slice(offset, length).order(ByteOrder.BIG_ENDIAN);
+ }
+
+ private static int bounded(int value, int maximum) throws IOException {
+ require(value >= 0 && value <= maximum);
+ return value;
+ }
+
+ private static void require(boolean valid) throws IOException {
+ if (!valid) {
+ throw new IOException("Invalid or unsupported native V4 DataTable");
+ }
+ }
+
+ private static final class Budget {
+ private long _bytes;
+
+ private void add(long bytes) throws IOException {
+ require(bytes >= 0 && bytes <= MAX_REFERENCED_BYTES - _bytes);
+ _bytes += bytes;
+ }
+ }
+}
diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/transport/NativeGrpcQueryClient.java b/pinot-broker/src/main/java/org/apache/pinot/broker/transport/NativeGrpcQueryClient.java
new file mode 100644
index 00000000000..3937cb776fb
--- /dev/null
+++ b/pinot-broker/src/main/java/org/apache/pinot/broker/transport/NativeGrpcQueryClient.java
@@ -0,0 +1,462 @@
+/**
+ * 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.transport;
+
+import com.google.common.annotations.VisibleForTesting;
+import io.grpc.CallCredentials;
+import io.grpc.CallOptions;
+import io.grpc.ClientCall;
+import io.grpc.Deadline;
+import io.grpc.ForwardingClientCall.SimpleForwardingClientCall;
+import io.grpc.ForwardingClientCallListener.SimpleForwardingClientCallListener;
+import io.grpc.Grpc;
+import io.grpc.ManagedChannel;
+import io.grpc.Metadata;
+import io.grpc.Status;
+import io.grpc.StatusRuntimeException;
+import io.grpc.netty.shaded.io.grpc.netty.GrpcSslContexts;
+import io.grpc.netty.shaded.io.grpc.netty.NettyChannelBuilder;
+import io.grpc.netty.shaded.io.netty.handler.ssl.SslContext;
+import io.grpc.netty.shaded.io.netty.handler.ssl.SslContextBuilder;
+import io.grpc.netty.shaded.io.netty.handler.ssl.SslProvider;
+import io.grpc.stub.ClientCalls;
+import java.io.ByteArrayInputStream;
+import java.io.IOException;
+import java.net.URI;
+import java.net.URISyntaxException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.security.cert.Certificate;
+import java.security.cert.X509Certificate;
+import java.time.Clock;
+import java.util.Arrays;
+import java.util.Date;
+import java.util.Iterator;
+import java.util.Locale;
+import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.Executor;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.Semaphore;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.LongSupplier;
+import javax.annotation.Nullable;
+import javax.net.ssl.SSLSession;
+import org.apache.hc.client5.http.ssl.DefaultHostnameVerifier;
+import org.apache.pinot.common.proto.PinotQueryServerGrpc;
+import org.apache.pinot.common.proto.Server.ServerRequest;
+import org.apache.pinot.common.proto.Server.ServerResponse;
+import org.apache.pinot.common.request.context.NativeQueryContext;
+
+import static com.google.common.base.Preconditions.checkArgument;
+import static java.util.Objects.requireNonNull;
+
+
+/// Strict native query transport with fixed, explicitly scoped PEM credentials and one owned channel per call.
+/// There is no plaintext, authority override, credential renewal, retry, or global TLS/channel cache. Per-call
+/// channels trade connection reuse for bounded ownership in this experimental path. Closing a call cancels its
+/// RPC; closing the client cancels every outstanding call. Callers still enforce query admission and response
+/// identities: transport authentication alone does not authorize a query or establish complete query coverage.
+public final class NativeGrpcQueryClient implements AutoCloseable {
+ private static final int MAX_CREDENTIAL_BYTES = 1024 * 1024;
+ private final SslContext _tls;
+ private final Limits _limits;
+ private final Clock _certificateClock;
+ private final Semaphore _slots;
+ private final Object _lifecycle = new Object();
+ private final Set<OwnedCall> _calls = ConcurrentHashMap.newKeySet();
+ private final Set<ManagedChannel> _channels = ConcurrentHashMap.newKeySet();
+ private final ScheduledExecutorService _reaper = Executors.newSingleThreadScheduledExecutor(
+ runnable -> new Thread(runnable, "native-query-channel-cleanup"));
+ private final AtomicBoolean _closed = new AtomicBoolean();
+
+ /// Hard bounds apply independently to each call; the active channel count cannot exceed maxConcurrentCalls.
+ public record Limits(int maxConcurrentCalls, int maxRequestBytes, int maxFrameBytes, int maxFrames,
+ long maxResponseBytes, long maxCallMillis) {
+ public Limits {
+ checkArgument(maxConcurrentCalls > 0 && maxConcurrentCalls <= 128
+ && maxRequestBytes > 0 && maxRequestBytes <= 8 * 1024 * 1024
+ && maxFrameBytes > 0 && maxFrameBytes <= 16 * 1024 * 1024
+ && maxFrames > 0 && maxFrames <= 65_536 && maxResponseBytes >= maxFrameBytes
+ && maxResponseBytes <= 256L * 1024 * 1024 && maxCallMillis > 0 && maxCallMillis <= 300_000,
+ "Invalid bounded native query transport configuration");
+ }
+ }
+
+ /// An owned, cancellable server stream. The iterator has one consumer; cancellation may run on another thread.
+ public interface Call extends AutoCloseable {
+ Iterator<ServerResponse> responses();
+
+ void cancel();
+
+ @Override
+ void close();
+ }
+
+ public NativeGrpcQueryClient(Path certificateChain, Path privateKey, Path trustCertificates, Limits limits)
+ throws IOException {
+ this(certificateChain, privateKey, trustCertificates, limits, Clock.systemUTC());
+ }
+
+ @VisibleForTesting
+ NativeGrpcQueryClient(Path certificateChain, Path privateKey, Path trustCertificates, Limits limits,
+ Clock certificateClock) throws IOException {
+ _limits = requireNonNull(limits);
+ _certificateClock = requireNonNull(certificateClock);
+ _slots = new Semaphore(limits.maxConcurrentCalls());
+ byte[] certificate = readCredential(certificateChain);
+ byte[] key = readCredential(privateKey);
+ try {
+ byte[] trust = readCredential(trustCertificates);
+ _tls = GrpcSslContexts.configure(SslContextBuilder.forClient(), SslProvider.JDK)
+ .keyManager(new ByteArrayInputStream(certificate), new ByteArrayInputStream(key))
+ .trustManager(new ByteArrayInputStream(trust)).build();
+ } finally {
+ Arrays.fill(key, (byte) 0);
+ }
+ }
+
+ /// The supplier must return the remaining original qualified query budget, not a fresh timeout. It must be local
+ /// and nonblocking; it is checked before connection/request admission and on both sides of iterator delivery.
+ public Call submit(URI endpoint, ServerRequest request, LongSupplier remainingMillis) {
+ validateEndpoint(endpoint);
+ requireNonNull(request);
+ requireNonNull(remainingMillis);
+ checkArgument(request.hasNativeQueryContext()
+ && NativeQueryContext.fromProto(request.getNativeQueryContext()).deadline() != null,
+ "Native query transport requires an identity envelope and original deadline");
+ if (request.getSerializedSize() > _limits.maxRequestBytes()) {
+ throw Status.RESOURCE_EXHAUSTED.withDescription("Native query request exceeds its byte bound")
+ .asRuntimeException();
+ }
+ if (_closed.get()) {
+ throw Status.CANCELLED.withDescription("Native query client is closed").asRuntimeException();
+ }
+ Deadline anchor = Deadline.after(0, TimeUnit.NANOSECONDS);
+ long remaining = remaining(remainingMillis);
+ Deadline deadline = anchor.offset(Math.min(remaining, _limits.maxCallMillis()), TimeUnit.MILLISECONDS);
+ if (!_slots.tryAcquire()) {
+ throw Status.RESOURCE_EXHAUSTED.withDescription("Native query call/channel limit reached").asRuntimeException();
+ }
+ ManagedChannel channel = null;
+ OwnedCall call = null;
+ try {
+ synchronized (_lifecycle) {
+ if (_closed.get()) {
+ throw Status.CANCELLED.withDescription("Native query client is closed").asRuntimeException();
+ }
+ channel = NettyChannelBuilder.forAddress(endpoint.getHost(), endpoint.getPort()).sslContext(_tls)
+ .maxInboundMessageSize(_limits.maxFrameBytes()).maxInboundMetadataSize(16 * 1024)
+ .disableRetry().build();
+ _channels.add(channel);
+ call = new OwnedCall(endpoint, channel, remainingMillis, deadline);
+ _calls.add(call);
+ }
+ call.start(request);
+ return call;
+ } catch (RuntimeException | Error e) {
+ if (call != null) {
+ call.cancel();
+ } else {
+ if (channel != null) {
+ retire(channel, null);
+ } else {
+ _slots.release();
+ }
+ }
+ throw e;
+ }
+ }
+
+ @Override
+ public void close() {
+ synchronized (_lifecycle) {
+ _closed.set(true);
+ }
+ _calls.forEach(OwnedCall::cancel);
+ if (_channels.isEmpty()) {
+ _reaper.shutdown();
+ }
+ }
+
+ @VisibleForTesting
+ int channelCount() {
+ return _channels.size();
+ }
+
+ private void retire(ManagedChannel channel, @Nullable OwnedCall call) {
+ channel.shutdownNow();
+ _reaper.execute(() -> reap(channel, call));
+ }
+
+ private void reap(ManagedChannel channel, @Nullable OwnedCall call) {
+ if (!channel.isTerminated()) {
+ _reaper.schedule(() -> reap(channel, call), 10, TimeUnit.MILLISECONDS);
+ return;
+ }
+ if (call != null) {
+ _calls.remove(call);
+ }
+ _slots.release();
+ _channels.remove(channel);
+ if (_closed.get() && _channels.isEmpty()) {
+ _reaper.shutdown();
+ }
+ }
+
+ private final class OwnedCall implements Call {
+ private final URI _endpoint;
+ private final ManagedChannel _channel;
+ private final LongSupplier _remainingMillis;
+ private final Deadline _deadline;
+ private final AtomicBoolean _released = new AtomicBoolean();
+ private final AtomicReference<StatusRuntimeException> _failure = new AtomicReference<>();
+ private final ClientCall<ServerRequest, ServerResponse> _rpc;
+ @Nullable
+ private volatile SSLSession _session;
+ @Nullable
+ private Iterator<ServerResponse> _responses;
+
+ private OwnedCall(URI endpoint, ManagedChannel channel, LongSupplier remainingMillis, Deadline deadline) {
+ _endpoint = endpoint;
+ _channel = channel;
+ _remainingMillis = remainingMillis;
+ _deadline = deadline;
+ _rpc = channel.newCall(PinotQueryServerGrpc.getSubmitMethod(), CallOptions.DEFAULT.withDeadline(deadline)
+ .withCallCredentials(new PeerCredentials()));
+ }
+
+ private void start(ServerRequest request) {
+ checkBudget();
+ ClientCall<ServerRequest, ServerResponse> guarded =
+ new SimpleForwardingClientCall<>(_rpc) {
+ @Override
+ public void start(Listener<ServerResponse> listener, Metadata headers) {
+ super.start(new SimpleForwardingClientCallListener<>(listener) {
+ private int _frames;
+ private long _bytes;
+
+ @Override
+ public void onMessage(ServerResponse response) {
+ try {
+ checkDelivery();
+ int bytes = response.getSerializedSize();
+ if (++_frames > _limits.maxFrames() || bytes > _limits.maxFrameBytes()
+ || bytes > _limits.maxResponseBytes() - _bytes) {
+ throw Status.RESOURCE_EXHAUSTED.withDescription("Native query response exceeds its bounds")
+ .asRuntimeException();
+ }
+ _bytes += bytes;
+ super.onMessage(response);
+ } catch (RuntimeException e) {
+ fail(e);
+ }
+ }
+
+ @Override
+ public void onClose(Status status, Metadata trailers) {
+ try {
+ if (status.isOk()) {
+ try {
+ checkDelivery();
+ } catch (RuntimeException e) {
+ _failure.compareAndSet(null, statusException(e));
+ }
+ }
+ StatusRuntimeException failure = _failure.get();
+ super.onClose(failure == null ? status : failure.getStatus(), trailers);
+ } finally {
+ release();
+ }
+ }
+ }, headers);
+ }
+ };
+ Iterator<ServerResponse> iterator = ClientCalls.blockingServerStreamingCall(guarded, request);
+ _responses = new Iterator<>() {
+ @Override
+ public boolean hasNext() {
+ try {
+ checkWaiting();
+ boolean next = iterator.hasNext();
+ checkDelivery();
+ return next;
+ } catch (RuntimeException e) {
+ fail(e);
+ throw e;
+ }
+ }
+
+ @Override
+ public ServerResponse next() {
+ try {
+ checkWaiting();
+ ServerResponse response = iterator.next();
+ checkDelivery();
+ return response;
+ } catch (RuntimeException e) {
+ fail(e);
+ throw e;
+ }
+ }
+ };
+ }
+
+ @Override
+ public Iterator<ServerResponse> responses() {
+ return requireNonNull(_responses, "Native query call has not started");
+ }
+
+ @Override
+ public void cancel() {
+ fail(Status.CANCELLED.withDescription("Native query call was cancelled").asRuntimeException());
+ }
+
+ @Override
+ public void close() {
+ cancel();
+ }
+
+ private void fail(RuntimeException failure) {
+ _failure.compareAndSet(null, statusException(failure));
+ _rpc.cancel("Native query call authority or transport failed", failure);
+ release();
+ }
+
+ private void release() {
+ if (_released.compareAndSet(false, true)) {
+ retire(_channel, this);
+ }
+ }
+
+ private void checkBudget() {
+ StatusRuntimeException failure = _failure.get();
+ if (failure != null) {
+ throw failure;
+ }
+ if (_closed.get()) {
+ throw Status.CANCELLED.withDescription("Native query client is closed").asRuntimeException();
+ }
+ remaining(_remainingMillis);
+ if (_deadline.isExpired()) {
+ throw Status.DEADLINE_EXCEEDED.withDescription("Native query call deadline expired").asRuntimeException();
+ }
+ }
+
+ private void checkDelivery() {
+ checkBudget();
+ validatePeer(requireNonNull(_session, "Native query transport has no authenticated peer"), _endpoint.getHost());
+ }
+
+ private void checkWaiting() {
+ checkBudget();
+ // Prior to receiving the first frame, the iterator can wait for negotiation. Request admission itself is
+ // blocked by PeerCredentials until TLS identity has been verified.
+ SSLSession session = _session;
+ if (session != null) {
+ validatePeer(session, _endpoint.getHost());
+ }
+ }
+
+ private final class PeerCredentials extends CallCredentials {
+ @Override
+ public void applyRequestMetadata(RequestInfo info, Executor executor, MetadataApplier applier) {
+ try {
+ checkBudget();
+ } catch (RuntimeException e) {
+ applier.fail(statusException(e).getStatus());
+ return;
+ }
+ try {
+ SSLSession session = requireNonNull(info.getTransportAttrs().get(Grpc.TRANSPORT_ATTR_SSL_SESSION));
+ validatePeer(session, _endpoint.getHost());
+ _session = session;
+ applier.apply(new Metadata());
+ } catch (Exception e) {
+ applier.fail(Status.UNAUTHENTICATED.withDescription("Native query requires a current verified TLS peer")
+ .withCause(e));
+ }
+ }
+ }
+ }
+
+ private void validatePeer(SSLSession session, String hostname) {
+ try {
+ Certificate[] chain = session.getPeerCertificates();
+ checkArgument(chain.length > 0 && chain[0] instanceof X509Certificate, "Missing TLS server identity");
+ Date now = Date.from(_certificateClock.instant());
+ for (Certificate certificate : chain) {
+ ((X509Certificate) certificate).checkValidity(now);
+ }
+ Certificate[] own = session.getLocalCertificates();
+ checkArgument(own != null && own.length > 0, "Missing TLS client identity");
+ for (Certificate certificate : own) {
+ ((X509Certificate) certificate).checkValidity(now);
+ }
+ new DefaultHostnameVerifier().verify(hostname, (X509Certificate) chain[0]);
+ } catch (Exception e) {
+ throw Status.UNAUTHENTICATED.withDescription("Native query TLS identity is absent, expired or mismatched")
+ .withCause(e).asRuntimeException();
+ }
+ }
+
+ private static long remaining(LongSupplier supplier) {
+ long remaining = supplier.getAsLong();
+ if (remaining <= 0) {
+ throw Status.DEADLINE_EXCEEDED.withDescription("Original native query deadline expired").asRuntimeException();
+ }
+ return remaining;
+ }
+
+ private static StatusRuntimeException statusException(RuntimeException failure) {
+ return failure instanceof StatusRuntimeException status ? status
+ : Status.FAILED_PRECONDITION.withDescription("Native query call validation failed").withCause(failure)
+ .asRuntimeException();
+ }
+
+ private static byte[] readCredential(Path path) throws IOException {
+ try (var input = Files.newInputStream(requireNonNull(path))) {
+ byte[] value = input.readNBytes(MAX_CREDENTIAL_BYTES + 1);
+ if (value.length == 0 || value.length > MAX_CREDENTIAL_BYTES) {
+ throw new IOException("Native TLS credential exceeds its byte bound");
+ }
+ return value;
+ }
+ }
+
+ private static void validateEndpoint(URI endpoint) {
+ requireNonNull(endpoint);
+ checkArgument(endpoint.toString().length() <= 2048 && "grpcs".equals(endpoint.getScheme())
+ && endpoint.getHost() != null && endpoint.getHost().equals(endpoint.getHost().toLowerCase(Locale.ROOT))
+ && endpoint.getPort() > 0 && endpoint.getPort() <= 65_535 && endpoint.getRawUserInfo() == null
+ && (endpoint.getRawPath() == null || endpoint.getRawPath().isEmpty())
+ && endpoint.getRawQuery() == null && endpoint.getRawFragment() == null,
+ "Native query endpoint must be a canonical grpcs authority");
+ try {
+ checkArgument(endpoint.toString().equals(
+ new URI("grpcs", null, endpoint.getHost(), endpoint.getPort(), null, null, null).toASCIIString()),
+ "Native query endpoint is not canonical");
+ } catch (URISyntaxException e) {
+ throw new IllegalArgumentException("Invalid native query endpoint", e);
+ }
+ }
+}
diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/broker/helix/NativeBrokerParticipantTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/broker/helix/NativeBrokerParticipantTest.java
index 7ddff4b2250..b0e3779d371 100644
--- a/pinot-broker/src/test/java/org/apache/pinot/broker/broker/helix/NativeBrokerParticipantTest.java
+++ b/pinot-broker/src/test/java/org/apache/pinot/broker/broker/helix/NativeBrokerParticipantTest.java
@@ -35,6 +35,7 @@ import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;
+import org.apache.pinot.broker.routing.NativeBrokerQueryAdmission.Admission;
import org.apache.pinot.common.coordination.MetadataRpcConnection;
import org.apache.pinot.common.coordination.NativeMetadataConfig;
import org.apache.pinot.common.coordination.QueryDeadlineGuard;
@@ -277,6 +278,9 @@ public class NativeBrokerParticipantTest {
fixture._participant.heartbeat();
fixture._participant.refreshRouting();
var original = fixture.admit();
+ long originalBudget = fixture._participant.remainingMillis(original);
+ assertTrue(originalBudget > 0 && originalBudget <= 300);
+ assertThrows(IllegalArgumentException.class, () -> fixture._participant.remainingMillis(mock(Admission.class)));
assertEquals(original.requests().size(), 1);
assertEquals(original.requests().getFirst().context().requiredNames(), List.of("segment"));
assertEquals(original.planningMetadata(false).tableCache().getSchema("orders").getSchemaName(), "orders");
@@ -289,10 +293,12 @@ public class NativeBrokerParticipantTest {
fixture._participant.heartbeat();
fixture._participant.refreshRouting();
fixture._participant.checkQueryDelivery(original);
+ assertTrue(fixture._participant.remainingMillis(original) < originalBudget);
var renewed = fixture.admit();
fixture._millis.set(201);
assertThrows(IllegalStateException.class,
() -> fixture._participant.deliverQuery(original, delivered::incrementAndGet));
+ assertThrows(IllegalStateException.class, () -> fixture._participant.remainingMillis(original));
assertEquals(delivered.get(), 1, "Rejected delivery must never invoke the callback");
fixture._participant.deliverQuery(renewed, delivered::incrementAndGet);
assertEquals(delivered.get(), 2);
@@ -768,7 +774,7 @@ public class NativeBrokerParticipantTest {
_participant.refreshRouting();
}
- private NativeBrokerParticipant.QueryAdmission admit() {
+ private Admission admit() {
return _participant.admitQuery("orders_OFFLINE", UUID.randomUUID(), UUID.randomUUID(),
_participant.originateDeadline(300));
}
diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/NativeGrpcQueryDispatcherTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/NativeGrpcQueryDispatcherTest.java
new file mode 100644
index 00000000000..04df546e41e
--- /dev/null
+++ b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/NativeGrpcQueryDispatcherTest.java
@@ -0,0 +1,598 @@
+/**
+ * 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.protobuf.ByteString;
+import io.grpc.Status;
+import java.io.IOException;
+import java.net.URI;
+import java.util.ArrayList;
+import java.util.Iterator;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.UUID;
+import java.util.concurrent.CancellationException;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.LongSupplier;
+import org.apache.pinot.broker.routing.NativeBrokerQueryAdmission;
+import org.apache.pinot.broker.routing.NativeBrokerQueryAdmission.Admission;
+import org.apache.pinot.broker.routing.NativeBrokerRoutingImage;
+import org.apache.pinot.broker.transport.NativeGrpcQueryClient;
+import org.apache.pinot.common.datatable.DataTable;
+import org.apache.pinot.common.proto.Server;
+import org.apache.pinot.common.request.BrokerRequest;
+import org.apache.pinot.common.request.context.NativeQueryContext;
+import org.apache.pinot.common.request.context.NativeQueryDeadline;
+import org.apache.pinot.common.response.broker.BrokerResponseNative;
+import org.apache.pinot.common.utils.DataSchema;
+import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
+import org.apache.pinot.coordination.ratis.rpc.MetadataRpc;
+import org.apache.pinot.core.common.datatable.DataTableBuilder;
+import org.apache.pinot.core.common.datatable.DataTableBuilderFactory;
+import org.apache.pinot.core.operator.blocks.results.ResultsBlockUtils;
+import org.apache.pinot.core.operator.streaming.StreamingResponseUtils;
+import org.apache.pinot.core.query.request.context.utils.QueryContextConverterUtils;
+import org.apache.pinot.core.routing.ImplicitHybridTableRouteInfo;
+import org.apache.pinot.core.routing.ImplicitHybridTableRouteProvider;
+import org.apache.pinot.core.routing.TableRouteInfo;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.data.FieldSpec.DataType;
+import org.apache.pinot.spi.data.Schema;
+import org.apache.pinot.spi.env.PinotConfiguration;
+import org.apache.pinot.spi.exception.QueryErrorCode;
+import org.apache.pinot.spi.query.QueryThreadContext;
+import org.apache.pinot.spi.utils.CommonConstants.Broker.Request.QueryOptionKey;
+import org.apache.pinot.spi.utils.CommonConstants.Query.Request;
+import org.apache.pinot.spi.utils.CommonConstants.Query.Response;
+import org.apache.pinot.spi.utils.JsonUtils;
+import org.apache.pinot.spi.utils.builder.TableConfigBuilder;
+import org.apache.pinot.sql.parsers.CalciteSqlCompiler;
+import org.testng.annotations.Test;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertThrows;
+import static org.testng.Assert.assertTrue;
+import static org.testng.Assert.expectThrows;
+
+
+/// Exercises the ordinary reducer with real DataTables and exact native fanout, including failures after partial work.
+public class NativeGrpcQueryDispatcherTest {
+ private static final long REQUEST_ID = 42;
+ private static final String BROKER_ID = "Broker_native";
+ private static final String TABLE = "orders_OFFLINE";
+ private static final String SELECTION = "SELECT id FROM orders_OFFLINE ORDER BY id DESC LIMIT 3";
+
+ @Test
+ public void testTwoServerSelectionAndAggregationUseNormalReducer() throws Exception {
+ // Server selection results are already sorted in query order for the broker's normal N-way merge.
+ try (Fixture fixture = new Fixture(List.of(selection(4, 1)), List.of(selection(3, 2)))) {
+ BrokerResponseNative result = fixture.dispatch(SELECTION);
+ assertEquals(result.getResultTable().getRows().stream().map(row -> row[0]).toList(), List.of(4, 3, 2));
+ assertTrue(result.getExceptions().isEmpty());
+ assertEquals(result.getNumSegmentsQueried(), 2L);
+ assertEquals(result.getNumServersQueried(), 2);
+ assertEquals(result.getNumServersResponded(), 2);
+ fixture.assertRequests();
+ fixture.assertClosed();
+ }
+ try (Fixture fixture = new Fixture(List.of(aggregate(5, 2)), List.of(aggregate(5, 2)))) {
+ BrokerResponseNative result = fixture.dispatch("SELECT SUM(id), COUNT(*) FROM orders_OFFLINE");
+ assertEquals(result.getResultTable().getRows().size(), 1);
+ assertEquals(result.getResultTable().getRows().getFirst(), new Object[]{10D, 4L});
+ assertTrue(result.getExceptions().isEmpty());
+ assertEquals(result.getNumSegmentsQueried(), 2L);
+ fixture.assertClosed();
+ }
+ try (Fixture fixture = new Fixture(List.of(groupedSum(1, 2)), List.of(groupedSum(1, 3)))) {
+ BrokerResponseNative result =
+ fixture.dispatch("SELECT id, SUM(id) FROM orders_OFFLINE GROUP BY id ORDER BY id");
+ assertEquals(result.getResultTable().getRows().size(), 3);
+ assertEquals(result.getResultTable().getRows().get(0), new Object[]{1, 2D});
+ assertEquals(result.getResultTable().getRows().get(1), new Object[]{2, 2D});
+ assertEquals(result.getResultTable().getRows().get(2), new Object[]{3, 3D});
+ assertTrue(result.getExceptions().isEmpty());
+ assertEquals(result.getNumSegmentsQueried(), 2L);
+ fixture.assertClosed();
+ }
+ var query = QueryContextConverterUtils.getQueryContext(
+ CalciteSqlCompiler.compileToBrokerRequest(SELECTION).getPinotQuery());
+ // The real all-pruned server path returns an empty STRING schema, even when another server returns INT rows.
+ DataTable empty = metadata(ResultsBlockUtils.buildEmptyQueryResults(query).getDataTable(), 0);
+ assertEquals(empty.getDataSchema().getColumnDataType(0), ColumnDataType.STRING);
+ try (Fixture fixture = new Fixture(
+ List.of(StreamingResponseUtils.getNonStreamingResponse(empty)), List.of(selection(2)))) {
+ BrokerResponseNative result = fixture.dispatch(SELECTION);
+ assertEquals(result.getResultTable().getRows().stream().map(row -> row[0]).toList(), List.of(2));
+ assertTrue(result.getExceptions().isEmpty());
+ assertEquals(result.getNumSegmentsQueried(), 2L);
+ fixture.assertClosed();
+ }
+ }
+
+ @Test(timeOut = 10_000)
+ public void testSubmissionAndReceiveFailureCancelEveryStartedCall() throws Exception {
+ try (Fixture fixture = new Fixture(List.of(selection(1)), List.of(selection(2)))) {
+ fixture._failSubmission = 2;
+ assertThrows(Exception.class, () -> fixture.dispatch(SELECTION));
+ assertEquals(fixture._submissions.get(), 2);
+ assertTrue(fixture._calls.getFirst()._cancelled.get(), "A failed second send must cancel the first RPC");
+ assertTrue(fixture._calls.getFirst()._closed.get());
+ }
+ try (Fixture fixture = new Fixture(List.of(selection(1)), List.of(selection(2)));
+ var executor = Executors.newSingleThreadExecutor()) {
+ CountDownLatch blocked = new CountDownLatch(1);
+ ControlledCall first = fixture._calls.getFirst();
+ ControlledCall second = fixture._calls.get(1);
+ first._beforeRead = () -> {
+ blocked.countDown();
+ await(first._cancelledSignal);
+ throw new CancellationException("Cancelled blocked peer");
+ };
+ second._beforeRead = () -> {
+ await(blocked);
+ throw Status.UNAVAILABLE.withDescription("Server response failed").asRuntimeException();
+ };
+ var dispatched = executor.submit(() -> {
+ assertThrows(Exception.class, () -> fixture.dispatch(SELECTION));
+ return null;
+ });
+ try {
+ assertTrue(blocked.await(1, TimeUnit.SECONDS));
+ dispatched.get(2, TimeUnit.SECONDS);
+ fixture.assertCancelled();
+ } finally {
+ fixture._calls.forEach(ControlledCall::cancel);
+ }
+ }
+ }
+
+ @Test
+ public void testCompleteSingleResponseAndExactCoverageAreRequired() throws Exception {
+ Server.ServerResponse valid = selection(1);
+ DataTable missing = selectionTable(2);
+ missing.getMetadata().put(DataTable.MetadataKey.NUM_SEGMENTS_QUERIED.getName(), "0");
+ DataTable failed = selectionTable(2);
+ failed.addException(QueryErrorCode.SERVER_SEGMENT_MISSING, "Required segment disappeared");
+ DataTable wrongRequest = selectionTable(2);
+ wrongRequest.getMetadata().put(DataTable.MetadataKey.REQUEST_ID.getName(), "41");
+ List<List<Server.ServerResponse>> malformed = List.of(
+ List.of(),
+ List.of(valid, valid),
+ List.of(valid.toBuilder().putMetadata(Response.MetadataKeys.RESPONSE_TYPE, Response.ResponseType.DATA).build()),
+ List.of(StreamingResponseUtils.getNonStreamingResponse(missing)),
+ List.of(StreamingResponseUtils.getNonStreamingResponse(failed)),
+ List.of(StreamingResponseUtils.getNonStreamingResponse(wrongRequest)),
+ List.of(StreamingResponseUtils.getNonStreamingResponse(
+ metadata(DataTableBuilderFactory.getEmptyDataTable(), 0))));
+ for (List<Server.ServerResponse> responses : malformed) {
+ try (Fixture fixture = new Fixture(List.of(valid), responses)) {
+ assertThrows(Exception.class, () -> fixture.dispatch(SELECTION));
+ fixture.assertCancelled();
+ }
+ }
+ try (Fixture fixture = new Fixture(List.of(valid), List.of(valid))) {
+ fixture._calls.get(1)._onEnd = () -> {
+ throw Status.UNAVAILABLE.withDescription("Failure after the full DataTable").asRuntimeException();
+ };
+ assertThrows(Exception.class, () -> fixture.dispatch(SELECTION));
+ fixture.assertCancelled();
+ }
+ }
+
+ @Test
+ public void testRevocationDuringReadAndCompletionCannotExposeSuccess() throws Exception {
+ for (boolean revokeAtEnd : List.of(false, true)) {
+ try (Fixture fixture = new Fixture(List.of(selection(1)), List.of(selection(2)))) {
+ if (revokeAtEnd) {
+ fixture._calls.get(1)._onEnd = () -> fixture._revoked.set(true);
+ } else {
+ fixture._calls.get(1)._onNext = () -> fixture._revoked.set(true);
+ }
+ assertThrows(Exception.class, () -> fixture.dispatch(SELECTION));
+ fixture.assertCancelled();
+ }
+ }
+ }
+
+ @Test
+ public void testRouteOrEnvelopeCannotDropAdmittedSegments() throws Exception {
+ try (Fixture fixture = new Fixture(List.of(selection(1)), List.of(selection(2)))) {
+ BrokerRequest request = CalciteSqlCompiler.compileToBrokerRequest(SELECTION);
+ ImplicitHybridTableRouteInfo route = (ImplicitHybridTableRouteInfo) fixture.route(request);
+ Map<?, ?> original = route.getOfflineRoutingTable();
+ var reduced = new LinkedHashMap<>(route.getOfflineRoutingTable());
+ reduced.remove(reduced.keySet().iterator().next());
+ route.setOfflineRoutingTable(reduced);
+ assertEquals(original.size(), 2);
+ try (QueryThreadContext ignored = QueryThreadContext.openForSseTest()) {
+ assertThrows(Exception.class, () -> fixture._dispatcher.dispatch(REQUEST_ID, "query-42", request, request,
+ route, fixture._admission, false));
+ }
+ assertEquals(fixture._submissions.get(), 0);
+ }
+ try (Fixture fixture = new Fixture(List.of(selection(1)), List.of(selection(2)))) {
+ when(fixture._admission.requests()).thenReturn(List.of(fixture._requests.getFirst()));
+ assertThrows(Exception.class, () -> fixture.dispatch(SELECTION));
+ assertEquals(fixture._submissions.get(), 0);
+ }
+ try (Fixture fixture = new Fixture(List.of(selection(1)), List.of(selection(2)))) {
+ BrokerRequest request = CalciteSqlCompiler.compileToBrokerRequest(SELECTION);
+ request.getPinotQuery().putToQueryOptions(QueryOptionKey.IGNORE_MISSING_SEGMENTS, "true");
+ assertThrows(Exception.class, () -> fixture.dispatch(request));
+ assertEquals(fixture._submissions.get(), 0, "Required coverage cannot be weakened by a legacy query option");
+ }
+ }
+
+ @Test(timeOut = 10_000)
+ public void testCancelledReceiverRetainsCapacityAndAggregateResponseBytesAreBounded() throws Exception {
+ try (Fixture fixture = new Fixture(List.of(selection(1)), List.of(selection(2)),
+ new NativeGrpcQueryDispatcher.Limits(1, 4, 1024 * 1024));
+ var executor = Executors.newSingleThreadExecutor()) {
+ CountDownLatch entered = new CountDownLatch(1);
+ CountDownLatch release = new CountDownLatch(1);
+ AtomicReference<Thread> receiver = new AtomicReference<>();
+ fixture._calls.getFirst()._beforeRead = () -> {
+ receiver.set(Thread.currentThread());
+ entered.countDown();
+ boolean interrupted = false;
+ try {
+ // Cancellation has reached the transport, but its receiver still owns resources until this is released.
+ while (release.getCount() != 0) {
+ try {
+ release.await();
+ } catch (InterruptedException e) {
+ interrupted = true;
+ }
+ }
+ } finally {
+ if (interrupted) {
+ Thread.currentThread().interrupt();
+ }
+ }
+ throw new CancellationException("Slow receiver has finally stopped");
+ };
+ fixture._calls.get(1)._beforeRead = () -> {
+ await(entered);
+ throw Status.UNAVAILABLE.withDescription("Other receiver failed").asRuntimeException();
+ };
+ var dispatched = executor.submit(() -> {
+ assertThrows(Exception.class, () -> fixture.dispatch(SELECTION));
+ return null;
+ });
+ try {
+ assertTrue(entered.await(1, TimeUnit.SECONDS));
+ dispatched.get(2, TimeUnit.SECONDS);
+ fixture.assertCancelled();
+ assertTrue(receiver.get().isAlive(), "The cancelled receiver deliberately still holds resources");
+ IllegalStateException full = expectThrows(IllegalStateException.class, () -> fixture.dispatch(SELECTION));
+ assertTrue(full.getMessage().contains("capacity"), "Cancellation must not prematurely free query capacity");
+ assertEquals(fixture._submissions.get(), 2);
+ fixture._dispatcher.close();
+ IllegalStateException closed = expectThrows(IllegalStateException.class, () -> fixture.dispatch(SELECTION));
+ assertTrue(closed.getMessage().contains("closed"));
+ assertEquals(fixture._submissions.get(), 2, "Closing must prevent any new transport submission");
+ } finally {
+ release.countDown();
+ Thread thread = receiver.get();
+ if (thread != null) {
+ thread.join(2000);
+ assertFalse(thread.isAlive(), "Release must allow the cancelled receiver to finish cleanup");
+ }
+ }
+ assertThrows(IllegalStateException.class, () -> fixture.dispatch(SELECTION));
+ }
+ Server.ServerResponse first = selection(1);
+ Server.ServerResponse second = selection(2);
+ long combinedLimit = first.getSerializedSize() + second.getSerializedSize() - 1L;
+ try (Fixture fixture = new Fixture(List.of(first), List.of(second),
+ new NativeGrpcQueryDispatcher.Limits(1, 4, combinedLimit))) {
+ // Each complete server response fits alone; admitting both must still enforce the aggregate query limit.
+ assertTrue(first.getSerializedSize() <= combinedLimit && second.getSerializedSize() <= combinedLimit);
+ assertThrows(Exception.class, () -> fixture.dispatch(SELECTION));
+ assertEquals(fixture._submissions.get(), 2);
+ fixture.assertCancelled();
+ }
+ }
+
+ private static Server.ServerResponse selection(int... values) throws IOException {
+ return StreamingResponseUtils.getNonStreamingResponse(selectionTable(values));
+ }
+
+ private static DataTable selectionTable(int... values) throws IOException {
+ DataTableBuilder builder = DataTableBuilderFactory.getDataTableBuilder(
+ new DataSchema(new String[]{"id"}, new ColumnDataType[]{ColumnDataType.INT}));
+ for (int value : values) {
+ builder.startRow();
+ builder.setColumn(0, value);
+ builder.finishRow();
+ }
+ return metadata(builder.build(), values.length);
+ }
+
+ private static Server.ServerResponse aggregate(double sum, long count) throws IOException {
+ DataTableBuilder builder = DataTableBuilderFactory.getDataTableBuilder(new DataSchema(
+ new String[]{"sum(id)", "count(*)"}, new ColumnDataType[]{ColumnDataType.DOUBLE, ColumnDataType.LONG}));
+ builder.startRow();
+ builder.setColumn(0, sum);
+ builder.setColumn(1, count);
+ builder.finishRow();
+ return StreamingResponseUtils.getNonStreamingResponse(metadata(builder.build(), count));
+ }
+
+ private static Server.ServerResponse groupedSum(int... ids) throws IOException {
+ DataTableBuilder builder = DataTableBuilderFactory.getDataTableBuilder(
+ new DataSchema(new String[]{"id", "sum(id)"}, new ColumnDataType[]{ColumnDataType.INT, ColumnDataType.DOUBLE}));
+ for (int id : ids) {
+ builder.startRow();
+ builder.setColumn(0, id);
+ builder.setColumn(1, (double) id);
+ builder.finishRow();
+ }
+ return StreamingResponseUtils.getNonStreamingResponse(metadata(builder.build(), ids.length));
+ }
+
+ private static DataTable metadata(DataTable table, long docs) {
+ table.getMetadata().put(DataTable.MetadataKey.REQUEST_ID.getName(), Long.toString(REQUEST_ID));
+ table.getMetadata().put(DataTable.MetadataKey.NUM_SEGMENTS_QUERIED.getName(), "1");
+ table.getMetadata().put(DataTable.MetadataKey.NUM_SEGMENTS_PROCESSED.getName(), "1");
+ table.getMetadata().put(DataTable.MetadataKey.NUM_SEGMENTS_MATCHED.getName(), "1");
+ table.getMetadata().put(DataTable.MetadataKey.NUM_DOCS_SCANNED.getName(), Long.toString(docs));
+ return table;
+ }
+
+ private static void await(CountDownLatch latch) {
+ try {
+ assertTrue(latch.await(3, TimeUnit.SECONDS), "Expected concurrent dispatch/cancellation did not arrive");
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new IllegalStateException(e);
+ }
+ }
+
+ private static final class ControlledCall implements NativeGrpcQueryClient.Call {
+ private final List<Server.ServerResponse> _frames;
+ private final AtomicBoolean _cancelled = new AtomicBoolean();
+ private final AtomicBoolean _closed = new AtomicBoolean();
+ private final CountDownLatch _cancelledSignal = new CountDownLatch(1);
+ private Runnable _beforeRead = () -> { };
+ private Runnable _onNext = () -> { };
+ private Runnable _onEnd = () -> { };
+
+ private ControlledCall(List<Server.ServerResponse> frames) {
+ _frames = List.copyOf(frames);
+ }
+
+ @Override
+ public Iterator<Server.ServerResponse> responses() {
+ return new Iterator<>() {
+ private int _index;
+
+ @Override
+ public boolean hasNext() {
+ if (_cancelled.get()) {
+ throw new CancellationException("RPC cancelled");
+ }
+ _beforeRead.run();
+ if (_index == _frames.size()) {
+ _onEnd.run();
+ return false;
+ }
+ return true;
+ }
+
+ @Override
+ public Server.ServerResponse next() {
+ _onNext.run();
+ return _frames.get(_index++);
+ }
+ };
+ }
+
+ @Override
+ public void cancel() {
+ _cancelled.set(true);
+ _cancelledSignal.countDown();
+ }
+
+ @Override
+ public void close() {
+ _closed.set(true);
+ cancel();
+ }
+ }
+
+ private static final class Fixture implements AutoCloseable {
+ private final NativeBrokerQueryAdmission _authority = mock(NativeBrokerQueryAdmission.class);
+ private final Admission _admission = mock(Admission.class);
+ private final AtomicBoolean _revoked = new AtomicBoolean();
+ private final AtomicInteger _submissions = new AtomicInteger();
+ private final List<ControlledCall> _calls;
+ private final List<Server.ServerRequest> _sent = new ArrayList<>();
+ private final NativeBrokerRoutingImage _image = image();
+ private final List<NativeBrokerRoutingImage.ServerRequest> _requests = _image.requests(
+ UUID.randomUUID(), UUID.randomUUID(), new NativeQueryDeadline("test-clock", 1000, 11000));
+ private final NativeGrpcQueryDispatcher _dispatcher;
+ private int _failSubmission;
+
+ private Fixture(List<Server.ServerResponse> first, List<Server.ServerResponse> second) throws IOException {
+ this(first, second, new NativeGrpcQueryDispatcher.Limits(2, 4, 1024 * 1024));
+ }
+
+ private Fixture(List<Server.ServerResponse> first, List<Server.ServerResponse> second,
+ NativeGrpcQueryDispatcher.Limits limits) throws IOException {
+ _calls = List.of(new ControlledCall(first), new ControlledCall(second));
+ when(_admission.image()).thenReturn(_image);
+ when(_admission.requests()).thenReturn(_requests);
+ when(_admission.planningMetadata(false)).thenAnswer(ignored -> _image.planningMetadata(false));
+ doAnswer(ignored -> {
+ checkAuthority();
+ return null;
+ }).when(_authority).checkQueryDelivery(_admission);
+ doAnswer(invocation -> {
+ checkAuthority();
+ invocation.<Runnable>getArgument(1).run();
+ return null;
+ }).when(_authority).deliverQuery(any(), any());
+ when(_authority.remainingMillis(_admission)).thenAnswer(ignored -> {
+ checkAuthority();
+ return 10_000L;
+ });
+ _dispatcher = new NativeGrpcQueryDispatcher(new PinotConfiguration(), BROKER_ID, _authority, this::submit,
+ limits);
+ }
+
+ private NativeGrpcQueryClient.Call submit(URI endpoint, Server.ServerRequest request, LongSupplier remaining) {
+ assertTrue(remaining.getAsLong() > 0);
+ int submission = _submissions.incrementAndGet();
+ if (submission == _failSubmission) {
+ throw Status.UNAVAILABLE.withDescription("Connection failed before submission").asRuntimeException();
+ }
+ int index = _requests.stream().map(NativeBrokerRoutingImage.ServerRequest::endpoint).toList().indexOf(endpoint);
+ assertTrue(index >= 0, "Dispatch must use the captured approved endpoint");
+ _sent.add(request);
+ return _calls.get(index);
+ }
+
+ private void checkAuthority() {
+ if (_revoked.get()) {
+ throw new IllegalStateException("Captured broker authority revoked");
+ }
+ }
+
+ private BrokerResponseNative dispatch(String sql) throws Exception {
+ return dispatch(CalciteSqlCompiler.compileToBrokerRequest(sql));
+ }
+
+ private BrokerResponseNative dispatch(BrokerRequest request) throws Exception {
+ try (QueryThreadContext ignored = QueryThreadContext.openForSseTest()) {
+ return _dispatcher.dispatch(REQUEST_ID, "query-42", request, request, route(request), _admission, false);
+ }
+ }
+
+ private TableRouteInfo route(BrokerRequest request) {
+ var planning = _image.planningMetadata(false);
+ var provider = new ImplicitHybridTableRouteProvider();
+ var route = provider.getTableRouteInfo(TABLE, planning.tableCache(), planning.routingManager());
+ provider.calculateRoutes(route, planning.routingManager(), request, null, REQUEST_ID);
+ return route;
+ }
+
+ private void assertRequests() {
+ assertEquals(_sent.size(), 2);
+ for (Server.ServerRequest request : _sent) {
+ NativeQueryContext context = NativeQueryContext.fromProto(request.getNativeQueryContext());
+ var expected = _requests.stream().filter(target ->
+ target.context().server().equals(context.server())).findFirst().orElseThrow();
+ assertEquals(context.toProto(), expected.context().toProto());
+ assertEquals(request.getSegmentsList(), expected.context().requiredNames());
+ assertEquals(request.getMetadataOrThrow(Request.MetadataKeys.ENABLE_STREAMING), "false");
+ assertEquals(request.getMetadataOrThrow(Request.MetadataKeys.REQUEST_ID), Long.toString(REQUEST_ID));
+ assertEquals(request.getMetadataOrThrow(Request.MetadataKeys.CORRELATION_ID), "query-42");
+ assertEquals(request.getMetadataOrThrow(Request.MetadataKeys.BROKER_ID), BROKER_ID);
+ assertFalse(request.getPayload().isEmpty(), "Use the compiled ordinary BrokerRequest");
+ }
+ }
+
+ private void assertClosed() {
+ _calls.forEach(call -> assertTrue(call._closed.get(), "Each started RPC must be closed"));
+ }
+
+ private void assertCancelled() {
+ _calls.forEach(call -> assertTrue(call._cancelled.get(), "A failed query must cancel every started RPC"));
+ assertClosed();
+ }
+
+ @Override
+ public void close() {
+ _dispatcher.close();
+ }
+ }
+
+ private static NativeBrokerRoutingImage image() throws IOException {
+ String history = UUID.randomUUID().toString();
+ var broker = token(history, BROKER_ID, MetadataRpc.ParticipantRole.BROKER);
+ var table = object(history);
+ var schema = object(history);
+ var revision = MetadataRpc.Revision.newBuilder().setHistoryId(history).setIndex(20);
+ var applied = MetadataRpc.Revision.newBuilder().setHistoryId(history).setIndex(1);
+ var control = MetadataRpc.MembershipControl.newBuilder().setControlEpoch(MetadataRpc.ControlToken.newBuilder()
+ .setHistoryId(history).setRaftTerm(1).setBarrierIndex(1).setLeaderId("controller")
+ .setActivationId(UUID.randomUUID().toString())).setClockBarrierIndex(1).setMinimumTickDelayMillis(100)
+ .setMaxConfiguredPermitLifetimeMillis(200).setPolicy(MetadataRpc.MembershipPolicy.newBuilder().setVersion(1)
+ .setMinTickIntervalMillis(100).setLeaseRounds(3).setPermitLifetimeMillis(200)
+ .setMaxQueryLifetimeMillis(300).setMaxRevisionLag(100));
+ var tableConfig = new TableConfigBuilder(TableType.OFFLINE).setTableName("orders").build();
+ var schemaConfig = new Schema.SchemaBuilder().setSchemaName("orders")
+ .addSingleValueDimension("id", DataType.INT).build();
+ var routing = MetadataRpc.TableRouting.newBuilder().setTableName(TABLE).setTable(table)
+ .setSchemaName("orders").setSchema(schema)
+ .setTableConfig(record(table, TABLE, MetadataRpc.RecordKind.TABLE_CONFIG,
+ MetadataRpc.CatalogValue.newBuilder().setTableConfigJson(
+ ByteString.copyFrom(JsonUtils.objectToBytes(tableConfig))).build()))
+ .setSchemaRecord(record(schema, "orders", MetadataRpc.RecordKind.SCHEMA,
+ MetadataRpc.CatalogValue.newBuilder().setSchemaJson(
+ ByteString.copyFrom(JsonUtils.objectToBytes(schemaConfig))).build()));
+ for (int i = 0; i < 2; i++) {
+ String endpoint = "grpcs://server-" + i + ":8090";
+ var server = token(history, "Server_" + i, MetadataRpc.ParticipantRole.SERVER);
+ routing.addSegments(MetadataRpc.SegmentRouting.newBuilder().setName("segment-" + i)
+ .setContent(object(history)).setContentSha256("a".repeat(64))
+ .addReplicas(MetadataRpc.ReplicaRouting.newBuilder().setServer(server).setAssignment(object(history))
+ .setAssignmentEpoch(1).setMaterializationId(UUID.randomUUID().toString())
+ .setInstallationId(UUID.randomUUID().toString()).setEndpoint(endpoint)));
+ routing.addServerAuthorities(MetadataRpc.ParticipantState.newBuilder().setParticipant(server)
+ .setReady(true).setExpiryRound(3).setLastModifiedIndex(2).setAppliedRevision(applied)
+ .addCapabilities(MetadataRpc.ServingCapability.SERVER_QUERY)
+ .putEndpoints("queryGrpc", endpoint).putEndpoints("nativeQueryVersion", "2"));
+ }
+ var snapshot = MetadataRpc.RoutingSnapshot.newBuilder().setStatus(MetadataRpc.ResultStatus.OK)
+ .setRevision(revision).setControl(control).setTable(routing)
+ .setParticipant(MetadataRpc.ParticipantState.newBuilder().setParticipant(broker).setReady(true)
+ .setExpiryRound(3).setLastModifiedIndex(2).setAppliedRevision(applied)
+ .addCapabilities(MetadataRpc.ServingCapability.BROKER_QUERY)).build();
+ return new NativeBrokerRoutingImage(snapshot, broker, TABLE);
+ }
+
+ private static MetadataRpc.ParticipantToken token(String history, String id, MetadataRpc.ParticipantRole role) {
+ return MetadataRpc.ParticipantToken.newBuilder().setHistoryId(history).setParticipantId(id).setRole(role)
+ .setIncarnation(UUID.randomUUID().toString()).setRegistrationEpoch(1).build();
+ }
+
+ private static MetadataRpc.ObjectVersion object(String history) {
+ return MetadataRpc.ObjectVersion.newBuilder().setHistoryId(history).setObjectId(UUID.randomUUID().toString())
+ .setVersion(1).build();
+ }
+
+ private static MetadataRpc.CatalogRecord record(MetadataRpc.ObjectVersion identity, String name,
+ MetadataRpc.RecordKind kind, MetadataRpc.CatalogValue value) {
+ return MetadataRpc.CatalogRecord.newBuilder().setKind(kind).setName(name).setObjectId(identity.getObjectId())
+ .setVersion(1).setSchemaVersion(1).setLastModifiedIndex(1).setValue(value).build();
+ }
+}
diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/transport/NativeDataTableDecoderTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/transport/NativeDataTableDecoderTest.java
new file mode 100644
index 00000000000..142432ea1ec
--- /dev/null
+++ b/pinot-broker/src/test/java/org/apache/pinot/broker/transport/NativeDataTableDecoderTest.java
@@ -0,0 +1,224 @@
+/**
+ * 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.transport;
+
+import java.io.IOException;
+import java.math.BigDecimal;
+import java.nio.ByteBuffer;
+import java.util.List;
+import org.apache.pinot.common.datatable.DataTable;
+import org.apache.pinot.common.datatable.DataTable.MetadataKey;
+import org.apache.pinot.common.utils.DataSchema;
+import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
+import org.apache.pinot.core.common.datatable.DataTableBuilder;
+import org.apache.pinot.core.common.datatable.DataTableBuilderFactory;
+import org.apache.pinot.spi.utils.ByteArray;
+import org.roaringbitmap.RoaringBitmap;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertThrows;
+
+
+/// Uses ordinary V4 writers and small malformed frames to verify native checks precede allocation and reduction.
+public class NativeDataTableDecoderTest {
+ @Test
+ public void testOrdinaryValuesAndEmptyTable() throws IOException {
+ ColumnDataType[] types = {ColumnDataType.INT, ColumnDataType.LONG, ColumnDataType.FLOAT, ColumnDataType.DOUBLE,
+ ColumnDataType.STRING, ColumnDataType.BYTES, ColumnDataType.BIG_DECIMAL, ColumnDataType.INT_ARRAY,
+ ColumnDataType.LONG_ARRAY, ColumnDataType.FLOAT_ARRAY, ColumnDataType.DOUBLE_ARRAY,
+ ColumnDataType.STRING_ARRAY, ColumnDataType.BYTES_ARRAY, ColumnDataType.BIG_DECIMAL_ARRAY};
+ String[] names = new String[types.length];
+ for (int i = 0; i < names.length; i++) {
+ names[i] = "column" + i;
+ }
+ DataSchema schema = new DataSchema(names, types);
+ DataTableBuilder builder = DataTableBuilderFactory.getDataTableBuilder(schema);
+ builder.startRow();
+ builder.setColumn(0, 7);
+ builder.setColumn(1, 8L);
+ builder.setColumn(2, 1.5F);
+ builder.setColumn(3, 2.5D);
+ builder.setColumn(4, "hello");
+ builder.setColumn(5, new ByteArray(new byte[]{1, 2}));
+ builder.setColumn(6, new BigDecimal("1.25"));
+ builder.setColumn(7, new int[]{1, 2});
+ builder.setColumn(8, new long[]{3, 4});
+ builder.setColumn(9, new float[]{1.5F});
+ builder.setColumn(10, new double[]{2.5D});
+ builder.setColumn(11, new String[]{"hello", "world"});
+ builder.setColumn(12, new ByteArray[]{new ByteArray(new byte[]{3}), new ByteArray(new byte[0])});
+ builder.setColumn(13, new BigDecimal[]{new BigDecimal("2.75")});
+ builder.finishRow();
+ DataTable original = builder.build();
+ original.getMetadata().put(MetadataKey.REQUEST_ID.getName(), "42");
+ original.getMetadata().put(MetadataKey.TRACE_INFO.getName(), "trace");
+ ByteBuffer input = ByteBuffer.wrap(original.toBytes());
+ DataTable decoded = NativeDataTableDecoder.decode(input);
+ assertEquals(input.position(), 0);
+ assertEquals(decoded.getDataSchema(), schema);
+ assertEquals(decoded.getNumberOfRows(), 1);
+ assertEquals(decoded.getInt(0, 0), 7);
+ assertEquals(decoded.getLong(0, 1), 8L);
+ assertEquals(decoded.getFloat(0, 2), 1.5F);
+ assertEquals(decoded.getDouble(0, 3), 2.5D);
+ assertEquals(decoded.getString(0, 4), "hello");
+ assertEquals(decoded.getBytes(0, 5), new ByteArray(new byte[]{1, 2}));
+ assertEquals(decoded.getBigDecimal(0, 6), new BigDecimal("1.25"));
+ assertEquals(decoded.getIntArray(0, 7), new int[]{1, 2});
+ assertEquals(decoded.getLongArray(0, 8), new long[]{3, 4});
+ assertEquals(decoded.getFloatArray(0, 9), new float[]{1.5F});
+ assertEquals(decoded.getDoubleArray(0, 10), new double[]{2.5D});
+ assertEquals(decoded.getStringArray(0, 11), new String[]{"hello", "world"});
+ assertEquals(decoded.getBytesArray(0, 12),
+ new ByteArray[]{new ByteArray(new byte[]{3}), new ByteArray(new byte[0])});
+ assertEquals(decoded.getBigDecimalArray(0, 13), new BigDecimal[]{new BigDecimal("2.75")});
+ assertEquals(decoded.getMetadata(), original.getMetadata());
+ DataTable empty = DataTableBuilderFactory.getDataTableBuilder(schema).build();
+ assertEquals(NativeDataTableDecoder.decode(ByteBuffer.wrap(empty.toBytes())).getNumberOfRows(), 0);
+ }
+
+ @Test
+ public void testNullBitmapArrayBitmapAndRunContainers() throws IOException {
+ RoaringBitmap array = RoaringBitmap.bitmapOf(0, 70_000);
+ RoaringBitmap bitmap = new RoaringBitmap();
+ for (int i = 0; i < 10_000; i += 2) {
+ bitmap.add(i);
+ }
+ RoaringBitmap run = new RoaringBitmap();
+ run.add(0L, 5000L);
+ run.runOptimize();
+ for (RoaringBitmap nulls : List.of(array, bitmap, run)) {
+ int rows = nulls.last() + 1;
+ DataTableBuilder builder = DataTableBuilderFactory.getDataTableBuilder(
+ new DataSchema(new String[]{"id"}, new ColumnDataType[]{ColumnDataType.INT}));
+ for (int i = 0; i < rows; i++) {
+ builder.startRow();
+ builder.setColumn(0, i);
+ builder.finishRow();
+ }
+ builder.setNullRowIds(nulls);
+ byte[] bytes = builder.build().toBytes();
+ DataTable decoded = NativeDataTableDecoder.decode(ByteBuffer.wrap(bytes));
+ assertEquals(decoded.getNullRowIds(0), nulls);
+ int fixed = ByteBuffer.wrap(bytes).getInt(9 * Integer.BYTES);
+ rejectInt(bytes, fixed + rows * Integer.BYTES + Integer.BYTES, Integer.MAX_VALUE);
+ int variable = ByteBuffer.wrap(bytes).getInt(11 * Integer.BYTES);
+ rejectInt(bytes, variable, 0);
+ }
+ }
+
+ @Test
+ public void testRejectAllocationLengthsBeforeDecoding() throws IOException {
+ DataTableBuilder builder = DataTableBuilderFactory.getDataTableBuilder(
+ new DataSchema(new String[]{"value"}, new ColumnDataType[]{ColumnDataType.STRING}));
+ builder.startRow();
+ builder.setColumn(0, "small");
+ builder.finishRow();
+ DataTable table = builder.build();
+ table.getMetadata().put(MetadataKey.TRACE_INFO.getName(), "small");
+ byte[] bytes = table.toBytes();
+ ByteBuffer header = ByteBuffer.wrap(bytes);
+ for (int field : new int[]{0, 1, 2, 4, 6, 8, 10, 12}) {
+ rejectInt(bytes, field * Integer.BYTES, Integer.MAX_VALUE);
+ rejectInt(bytes, field * Integer.BYTES, -1);
+ }
+ int dictionary = header.getInt(5 * Integer.BYTES);
+ rejectInt(bytes, dictionary, Integer.MAX_VALUE);
+ rejectInt(bytes, dictionary + Integer.BYTES, Integer.MAX_VALUE);
+ int schema = header.getInt(7 * Integer.BYTES);
+ rejectInt(bytes, schema, Integer.MAX_VALUE);
+ rejectInt(bytes, schema + Integer.BYTES, Integer.MAX_VALUE);
+ int fixed = header.getInt(9 * Integer.BYTES);
+ rejectInt(bytes, fixed, Integer.MAX_VALUE);
+ int metadata = header.getInt(11 * Integer.BYTES) + header.getInt(12 * Integer.BYTES);
+ rejectInt(bytes, metadata, Integer.MAX_VALUE);
+ rejectInt(bytes, metadata + Integer.BYTES, Integer.MAX_VALUE);
+ rejectInt(bytes, metadata + 2 * Integer.BYTES, Integer.MAX_VALUE);
+ ByteBuffer metadataEntries = ByteBuffer.wrap(bytes);
+ metadataEntries.position(metadata + 2 * Integer.BYTES);
+ while (metadataEntries.hasRemaining()) {
+ MetadataKey key = MetadataKey.getById(metadataEntries.getInt());
+ if (key == MetadataKey.TRACE_INFO) {
+ rejectInt(bytes, metadataEntries.position(), Integer.MAX_VALUE);
+ break;
+ }
+ assertEquals(key.getValueType(), DataTable.MetadataValueType.LONG);
+ metadataEntries.getLong();
+ }
+ for (int size : new int[]{0, 4, 51, bytes.length - 1}) {
+ assertThrows(IOException.class, () -> NativeDataTableDecoder.decode(ByteBuffer.wrap(bytes, 0, size)));
+ }
+ }
+
+ @Test
+ public void testRejectVariableReferencesAndUnsupportedCodecs() throws IOException {
+ for (ColumnDataType type : List.of(ColumnDataType.INT_ARRAY, ColumnDataType.BYTES_ARRAY,
+ ColumnDataType.BIG_DECIMAL_ARRAY)) {
+ DataTableBuilder builder = DataTableBuilderFactory.getDataTableBuilder(
+ new DataSchema(new String[]{"value"}, new ColumnDataType[]{type}));
+ builder.startRow();
+ if (type == ColumnDataType.INT_ARRAY) {
+ builder.setColumn(0, new int[]{1});
+ } else if (type == ColumnDataType.BYTES_ARRAY) {
+ builder.setColumn(0, new ByteArray[]{new ByteArray(new byte[]{1})});
+ } else {
+ builder.setColumn(0, new BigDecimal[]{BigDecimal.ONE});
+ }
+ builder.finishRow();
+ byte[] bytes = builder.build().toBytes();
+ ByteBuffer header = ByteBuffer.wrap(bytes);
+ int fixed = header.getInt(9 * Integer.BYTES);
+ rejectInt(bytes, fixed, -1);
+ rejectInt(bytes, fixed + Integer.BYTES, Integer.MAX_VALUE);
+ if (type != ColumnDataType.INT_ARRAY) {
+ rejectInt(bytes, header.getInt(11 * Integer.BYTES), Integer.MAX_VALUE);
+ }
+ }
+ for (ColumnDataType type : List.of(ColumnDataType.MAP, ColumnDataType.OBJECT, ColumnDataType.UNKNOWN)) {
+ DataTable table = DataTableBuilderFactory.getDataTableBuilder(
+ new DataSchema(new String[]{"value"}, new ColumnDataType[]{type})).build();
+ assertThrows(IOException.class, () -> NativeDataTableDecoder.decode(ByteBuffer.wrap(table.toBytes())));
+ }
+ DataSchema schema = new DataSchema(new String[]{"value"}, new ColumnDataType[]{ColumnDataType.BYTES});
+ ByteBuffer fixed = ByteBuffer.allocate(65 * 8);
+ for (int i = 0; i < 65; i++) {
+ fixed.putInt(0).putInt(1024 * 1024);
+ }
+ byte[][] sections = {new byte[4], new byte[4], schema.toBytes(), fixed.array(), new byte[1024 * 1024]};
+ int next = 13 * Integer.BYTES;
+ ByteBuffer input = ByteBuffer.allocate(next + 8 + List.of(sections).stream().mapToInt(s -> s.length).sum());
+ input.putInt(4).putInt(65).putInt(1);
+ for (byte[] section : sections) {
+ input.putInt(next).putInt(section.length);
+ next += section.length;
+ }
+ for (byte[] section : sections) {
+ input.put(section);
+ }
+ input.putInt(4).putInt(0).flip();
+ assertThrows(IOException.class, () -> NativeDataTableDecoder.decode(input));
+ }
+
+ private static void rejectInt(byte[] original, int offset, int value) {
+ byte[] malformed = original.clone();
+ ByteBuffer.wrap(malformed).putInt(offset, value);
+ assertThrows(IOException.class, () -> NativeDataTableDecoder.decode(ByteBuffer.wrap(malformed)));
+ }
+}
diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/transport/NativeGrpcQueryClientTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/transport/NativeGrpcQueryClientTest.java
new file mode 100644
index 00000000000..d7f71d9b291
--- /dev/null
+++ b/pinot-broker/src/test/java/org/apache/pinot/broker/transport/NativeGrpcQueryClientTest.java
@@ -0,0 +1,343 @@
+/**
+ * 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.transport;
+
+import com.google.protobuf.ByteString;
+import io.grpc.Server;
+import io.grpc.Status;
+import io.grpc.StatusRuntimeException;
+import io.grpc.netty.shaded.io.grpc.netty.GrpcSslContexts;
+import io.grpc.netty.shaded.io.grpc.netty.NettyServerBuilder;
+import io.grpc.netty.shaded.io.netty.handler.ssl.ClientAuth;
+import io.grpc.stub.ServerCallStreamObserver;
+import io.grpc.stub.StreamObserver;
+import java.io.IOException;
+import java.net.InetSocketAddress;
+import java.net.URI;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.security.KeyStore;
+import java.security.cert.X509Certificate;
+import java.time.Clock;
+import java.time.Instant;
+import java.time.ZoneId;
+import java.time.ZoneOffset;
+import java.util.Base64;
+import java.util.List;
+import java.util.UUID;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Consumer;
+import javax.annotation.Nullable;
+import org.apache.commons.io.FileUtils;
+import org.apache.pinot.common.config.provider.NativeOfflineSegmentAction.ObjectVersion;
+import org.apache.pinot.common.config.provider.NativeOfflineSegmentAction.Participant;
+import org.apache.pinot.common.config.provider.NativeOfflineSegmentAction.Role;
+import org.apache.pinot.common.proto.PinotQueryServerGrpc;
+import org.apache.pinot.common.proto.Server.ServerRequest;
+import org.apache.pinot.common.proto.Server.ServerResponse;
+import org.apache.pinot.common.request.context.NativeQueryContext;
+import org.apache.pinot.common.request.context.NativeQueryDeadline;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertThrows;
+import static org.testng.Assert.assertTrue;
+import static org.testng.Assert.expectThrows;
+
+
+/// Exercises real loopback TLS, live certificate revalidation, bounded channels and cancellable query streams.
+public class NativeGrpcQueryClientTest {
+ private static final NativeGrpcQueryClient.Limits LIMITS =
+ new NativeGrpcQueryClient.Limits(4, 64 * 1024, 4096, 8, 32 * 1024, 5000);
+ private static final ServerResponse FRAME =
+ ServerResponse.newBuilder().setPayload(ByteString.copyFromUtf8("response")).build();
+
+ @Test(timeOut = 30_000)
+ public void testStrictTlsTrustHostnameAndNoPlaintextFallback() throws Exception {
+ try (TlsFixture tls = new TlsFixture(); TlsFixture unrelated = new TlsFixture();
+ var server = new QueryServer(tls, observer -> {
+ observer.onNext(FRAME);
+ observer.onCompleted();
+ });
+ var client = tls.client(LIMITS)) {
+ try (var call = client.submit(server.endpoint(), request(), () -> 5000)) {
+ assertTrue(call.responses().hasNext());
+ assertEquals(call.responses().next(), FRAME);
+ assertFalse(call.responses().hasNext());
+ }
+ assertEquals(server._requests.get(), 1);
+ URI wrongHostname = URI.create("grpcs://127.0.0.1:" + server._server.getPort());
+ try (var call = client.submit(wrongHostname, request(), () -> 5000)) {
+ assertThrows(StatusRuntimeException.class, () -> call.responses().hasNext());
+ }
+ try (var untrusted = new NativeGrpcQueryClient(tls._certificate, tls._key, unrelated._certificate, LIMITS);
+ var call = untrusted.submit(server.endpoint(), request(), () -> 5000)) {
+ assertThrows(StatusRuntimeException.class, () -> call.responses().hasNext());
+ }
+ assertEquals(server._requests.get(), 1, "Invalid TLS peers must not reach the query service");
+ for (String invalid : List.of("grpc://localhost:8090", "grpcs://LOCALHOST:8090",
+ "grpcs://localhost:8090/", "grpcs://user@localhost:8090")) {
+ assertThrows(IllegalArgumentException.class, () -> client.submit(URI.create(invalid), request(), () -> 5000));
+ }
+ assertThrows(IllegalArgumentException.class,
+ () -> client.submit(server.endpoint(), request().toBuilder().clearNativeQueryContext().build(), () -> 5000));
+ assertEquals(expectThrows(StatusRuntimeException.class, () -> client.submit(server.endpoint(),
+ request().toBuilder().setSql("x".repeat(65 * 1024)).build(), () -> 5000)).getStatus().getCode(),
+ Status.Code.RESOURCE_EXHAUSTED);
+ try (var plaintext = new QueryServer(null, observer -> observer.onCompleted());
+ var call = client.submit(plaintext.endpoint(), request(), () -> 5000)) {
+ assertThrows(StatusRuntimeException.class, () -> call.responses().hasNext());
+ assertEquals(plaintext._requests.get(), 0);
+ }
+ client.close();
+ awaitClosed(client);
+ }
+ }
+
+ @Test(timeOut = 30_000)
+ public void testCertificateExpiryAfterHandshakeRejectsFramesAndNextRequest() throws Exception {
+ try (TlsFixture tls = new TlsFixture()) {
+ MutableClock clock = new MutableClock();
+ AtomicReference<StreamObserver<ServerResponse>> producer = new AtomicReference<>();
+ try (var server = new QueryServer(tls, observer -> {
+ producer.set(observer);
+ observer.onNext(FRAME);
+ }); var client = new NativeGrpcQueryClient(tls._certificate, tls._key, tls._certificate, LIMITS, clock)) {
+ try (var call = client.submit(server.endpoint(), request(), () -> 5000)) {
+ assertEquals(call.responses().next(), FRAME, "A real TLS handshake and authenticated request must succeed");
+ clock._now.set(tls._identity.getNotAfter().toInstant().plusSeconds(1));
+ producer.get().onNext(FRAME);
+ assertEquals(expectThrows(StatusRuntimeException.class, () -> call.responses().hasNext())
+ .getStatus().getCode(), Status.Code.UNAUTHENTICATED);
+ }
+ try (var next = client.submit(server.endpoint(), request(), () -> 5000)) {
+ assertThrows(StatusRuntimeException.class, () -> next.responses().hasNext());
+ }
+ assertEquals(server._requests.get(), 1,
+ "Call credentials must reject expired identity before another request, independently of TLS resumption");
+ client.close();
+ awaitClosed(client);
+ }
+ }
+ }
+
+ @Test(timeOut = 30_000)
+ public void testCancellationDeadlinesAndTransportBoundsReleaseOwnedChannels() throws Exception {
+ try (TlsFixture tls = new TlsFixture(); var executor = Executors.newSingleThreadExecutor()) {
+ CountDownLatch entered = new CountDownLatch(1);
+ CountDownLatch cancelled = new CountDownLatch(1);
+ var limits = new NativeGrpcQueryClient.Limits(1, 64 * 1024, 1024, 1, 1024, 5000);
+ try (var server = new QueryServer(tls, observer -> {
+ ((ServerCallStreamObserver<ServerResponse>) observer).setOnCancelHandler(cancelled::countDown);
+ entered.countDown();
+ }); var client = tls.client(limits)) {
+ var call = client.submit(server.endpoint(), request(), () -> 5000);
+ var waiting = executor.submit(
+ () -> expectThrows(StatusRuntimeException.class, () -> call.responses().hasNext()));
+ assertTrue(entered.await(5, TimeUnit.SECONDS));
+ assertEquals(client.channelCount(), 1);
+ assertEquals(expectThrows(StatusRuntimeException.class,
+ () -> client.submit(server.endpoint(), request(), () -> 5000)).getStatus().getCode(),
+ Status.Code.RESOURCE_EXHAUSTED);
+ call.cancel();
+ assertEquals(waiting.get(5, TimeUnit.SECONDS).getStatus().getCode(), Status.Code.CANCELLED);
+ assertTrue(cancelled.await(5, TimeUnit.SECONDS));
+ awaitClosed(client);
+ assertEquals(client.channelCount(), 0, "Slots are released only after the owned channel terminates");
+ AtomicLong remaining = new AtomicLong(5000);
+ try (var expired = client.submit(server.endpoint(), request(), remaining::get)) {
+ remaining.set(0);
+ assertEquals(expectThrows(StatusRuntimeException.class, () -> expired.responses().hasNext())
+ .getStatus().getCode(), Status.Code.DEADLINE_EXCEEDED);
+ }
+ awaitClosed(client);
+ try (var expired = client.submit(server.endpoint(), request(), () -> 50)) {
+ assertEquals(expectThrows(StatusRuntimeException.class, () -> expired.responses().hasNext())
+ .getStatus().getCode(), Status.Code.DEADLINE_EXCEEDED);
+ }
+ awaitClosed(client);
+ AtomicInteger observations = new AtomicInteger();
+ assertEquals(expectThrows(StatusRuntimeException.class,
+ () -> client.submit(server.endpoint(), request(), () -> {
+ if (observations.getAndIncrement() == 0) {
+ try {
+ Thread.sleep(100);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new IllegalStateException(e);
+ }
+ return 50;
+ }
+ return 5000;
+ })).getStatus().getCode(), Status.Code.DEADLINE_EXCEEDED,
+ "A pause inside the remaining-budget supplier must consume the anchored RPC deadline");
+ awaitClosed(client);
+ var stopped = client.submit(server.endpoint(), request(), () -> 5000);
+ client.close();
+ assertThrows(StatusRuntimeException.class, () -> stopped.responses().hasNext());
+ assertThrows(StatusRuntimeException.class, () -> client.submit(server.endpoint(), request(), () -> 5000));
+ awaitClosed(client);
+ }
+ for (int limit = 0; limit < 3; limit++) {
+ int bound = limit;
+ var responseLimits = new NativeGrpcQueryClient.Limits(1, 64 * 1024, 1024, bound == 2 ? 8 : 1, 1024, 5000);
+ try (var server = new QueryServer(tls, observer -> {
+ ServerResponse frame = bound == 0 ? FRAME
+ : FRAME.toBuilder().setPayload(ByteString.copyFrom(new byte[bound == 1 ? 2048 : 700])).build();
+ observer.onNext(frame);
+ if (bound != 1) {
+ observer.onNext(frame);
+ }
+ observer.onCompleted();
+ }); var client = tls.client(responseLimits);
+ var call = client.submit(server.endpoint(), request(), () -> 5000)) {
+ assertEquals(expectThrows(StatusRuntimeException.class, () -> {
+ while (call.responses().hasNext()) {
+ call.responses().next();
+ }
+ }).getStatus().getCode(), Status.Code.RESOURCE_EXHAUSTED);
+ }
+ }
+ }
+ }
+
+ private static void awaitClosed(NativeGrpcQueryClient client) throws InterruptedException {
+ long until = System.nanoTime() + TimeUnit.SECONDS.toNanos(5);
+ while (client.channelCount() != 0 && until - System.nanoTime() > 0) {
+ Thread.sleep(10);
+ }
+ assertEquals(client.channelCount(), 0, "Owned query channels did not terminate");
+ }
+
+ private static ServerRequest request() {
+ NativeQueryContext context = new NativeQueryContext(1, UUID.randomUUID(),
+ new NativeQueryContext.ProcessIdentity("broker", UUID.randomUUID(), 1),
+ new Participant("server", Role.SERVER, UUID.randomUUID(), 1), UUID.randomUUID(), UUID.randomUUID(),
+ "tbl_OFFLINE", object(), "tbl", object(),
+ List.of(new NativeQueryContext.Segment("segment", object(), "a".repeat(64), object(), UUID.randomUUID())),
+ List.of()).withDeadline(new NativeQueryDeadline("test-clock", 10, 100));
+ return ServerRequest.newBuilder().setSql("SELECT count(*) FROM tbl_OFFLINE").addSegments("segment")
+ .setNativeQueryContext(context.toProto()).build();
+ }
+
+ private static ObjectVersion object() {
+ return new ObjectVersion(UUID.randomUUID(), 1);
+ }
+
+ private static final class QueryServer implements AutoCloseable {
+ private final AtomicInteger _requests = new AtomicInteger();
+ private final Server _server;
+
+ private QueryServer(@Nullable TlsFixture tls, Consumer<StreamObserver<ServerResponse>> responder)
+ throws IOException {
+ var builder = NettyServerBuilder.forAddress(new InetSocketAddress("127.0.0.1", 0));
+ if (tls != null) {
+ builder.sslContext(GrpcSslContexts.forServer(tls._certificate.toFile(), tls._key.toFile())
+ .trustManager(tls._certificate.toFile()).clientAuth(ClientAuth.REQUIRE).build());
+ }
+ _server = builder.addService(new PinotQueryServerGrpc.PinotQueryServerImplBase() {
+ @Override
+ public void submit(ServerRequest request, StreamObserver<ServerResponse> observer) {
+ _requests.incrementAndGet();
+ responder.accept(observer);
+ }
+ }).build().start();
+ }
+
+ private URI endpoint() {
+ return URI.create("grpcs://localhost:" + _server.getPort());
+ }
+
+ @Override
+ public void close() throws InterruptedException {
+ _server.shutdownNow();
+ assertTrue(_server.awaitTermination(5, TimeUnit.SECONDS));
+ }
+ }
+
+ private static final class MutableClock extends Clock {
+ private final AtomicReference<Instant> _now = new AtomicReference<>(Instant.now());
+
+ @Override
+ public ZoneId getZone() {
+ return ZoneOffset.UTC;
+ }
+
+ @Override
+ public Clock withZone(ZoneId zone) {
+ return this;
+ }
+
+ @Override
+ public Instant instant() {
+ return _now.get();
+ }
+ }
+
+ private static final class TlsFixture implements AutoCloseable {
+ private final Path _directory = Files.createTempDirectory("native-grpc-client-");
+ private final Path _certificate = _directory.resolve("identity.pem");
+ private final Path _key = _directory.resolve("identity-key.pem");
+ private final X509Certificate _identity;
+
+ private TlsFixture() throws Exception {
+ Path file = _directory.resolve("identity.p12");
+ Process keytool = new ProcessBuilder(Path.of(System.getProperty("java.home"), "bin", "keytool").toString(),
+ "-genkeypair", "-alias", "test", "-keyalg", "EC", "-groupname", "secp256r1", "-dname", "CN=localhost",
+ "-ext", "SAN=DNS:localhost", "-ext", "EKU=serverAuth,clientAuth", "-validity", "2", "-storetype", "PKCS12",
+ "-keystore", file.toString(), "-storepass", "test-password", "-keypass", "test-password", "-noprompt")
+ .redirectErrorStream(true).start();
+ if (!keytool.waitFor(20, TimeUnit.SECONDS)) {
+ keytool.destroyForcibly();
+ throw new IOException("TLS test certificate creation timed out");
+ }
+ assertEquals(keytool.exitValue(), 0,
+ new String(keytool.getInputStream().readAllBytes(), StandardCharsets.UTF_8));
+ KeyStore store = KeyStore.getInstance("PKCS12");
+ try (var input = Files.newInputStream(file)) {
+ store.load(input, "test-password".toCharArray());
+ }
+ _identity = (X509Certificate) store.getCertificate("test");
+ pem(_certificate, "CERTIFICATE", _identity.getEncoded());
+ pem(_key, "PRIVATE KEY", store.getKey("test", "test-password".toCharArray()).getEncoded());
+ }
+
+ private NativeGrpcQueryClient client(NativeGrpcQueryClient.Limits limits) throws IOException {
+ return new NativeGrpcQueryClient(_certificate, _key, _certificate, limits);
+ }
+
+ private static void pem(Path file, String type, byte[] value) throws IOException {
+ Files.writeString(file, "-----BEGIN " + type + "-----\n"
+ + Base64.getMimeEncoder(64, new byte[]{'\n'}).encodeToString(value) + "\n-----END " + type + "-----\n",
+ StandardCharsets.US_ASCII);
+ }
+
+ @Override
+ public void close() throws IOException {
+ FileUtils.deleteDirectory(_directory.toFile());
+ }
+ }
+}
diff --git a/pinot-coordination-ratis/README.md b/pinot-coordination-ratis/README.md
index 88214c053d0..c870d27c815 100644
--- a/pinot-coordination-ratis/README.md
+++ b/pinot-coordination-ratis/README.md
@@ -459,9 +459,37 @@ Admission and delivery use a separate local lock and perform no metadata RPCs. A
under that lock, serialized with local invalidation; it must not call lifecycle operations or acquire their locks.
Shutdown fences locally before waiting for an in-flight membership RPC or unregistering.
-These APIs are not yet connected to authenticated query entry, normal network dispatch or final HTTP serialization.
+These APIs are not yet connected to authenticated query entry or final HTTP serialization.
Synthetic clocks in the focused tests establish lifecycle behavior, not production clock qualification.
+## Native broker dispatch infrastructure
+
+`NativeBrokerQueryAdmission` exposes the participant's existing local admission and delivery checks without exposing
+its lifecycle locks. `NativeGrpcQueryDispatcher` binds the normal OFFLINE route to every exact admitted destination
+and identity envelope before sending requests. It rejects dropped or optional segments, sampling and
+`ignoreMissingSegments`. Every server must provide one complete non-streaming DataTable with matching request ID,
+required-segment count and successful RPC completion. Missing, repeated, malformed or failed responses cancel the
+whole attempt. The existing Pinot reducer handles selection, ordering, aggregation and group-by; any partial reducer
+result fails the attempt. Reduction runs synchronously on the dispatching thread, including group-by subtasks, so
+timed-out reducer workers cannot accumulate after query capacity is released. This currently trades parallel reduction
+for explicit attempt ownership. Empty result rows remain valid; empty-table routing remains unsupported.
+
+Native response decoding accepts V4 with bounded section lengths, nested counts, variable data and null bitmaps.
+It supports scalar and array representations but explicitly rejects `OBJECT`, `MAP` and `UNKNOWN` columns; functions
+requiring those intermediate formats remain unsupported. These decoding limits do not establish a total heap bound
+for arbitrary reducers or result formatting.
+
+The transport requires fixed, explicitly scoped mutual-TLS credentials and canonical approved `grpcs` endpoints,
+checks the negotiated certificate chain, hostname and validity, and has no plaintext fallback or retries. It bounds
+request/frame/response bytes, call duration and active channels; a cancelled channel retains capacity until termination.
+Each call currently owns its channel, accepting handshake overhead instead of introducing another shared TLS cache.
+Dispatch bounds concurrent queries, fanout and aggregate bytes; cancelled receivers retain capacity until they exit.
+All fanout requests carry the original deadline, and acceptance checks retain the original finite admission.
+
+This path is not registered with the normal request handler or a native listener. Authenticated request entry, quotas,
+qualified production clocks and final HTTP output fencing still need integration before query serving can be enabled.
+Helix/ZK remains the default, and native query listeners remain disabled.
+
## Authenticated native query admission
The optional `principals` object in `pinot.coordination.ratis.broker.routing` maps lowercase SHA-256 leaf-certificate
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/BrokerReduceService.java b/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/BrokerReduceService.java
index 05021e4bf0e..7c92432babe 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/BrokerReduceService.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/BrokerReduceService.java
@@ -23,6 +23,7 @@ import java.util.Arrays;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
+import java.util.concurrent.ExecutorService;
import javax.annotation.Nullable;
import javax.annotation.concurrent.ThreadSafe;
import org.apache.pinot.common.datatable.DataTable;
@@ -62,6 +63,16 @@ public class BrokerReduceService extends BaseReduceService {
/// [org.apache.pinot.spi.query.QueryThreadContext] must already be set up before calling this method.
public BrokerResponseNative reduceOnDataTable(BrokerRequest brokerRequest, BrokerRequest serverBrokerRequest,
Map<ServerRoutingInstance, DataTable> dataTableMap, long reduceTimeOutMs, BrokerMetrics brokerMetrics) {
+ return reduceOnDataTable(brokerRequest, serverBrokerRequest, dataTableMap, reduceTimeOutMs, brokerMetrics,
+ _reduceExecutorService);
+ }
+
+ /// Reduces with a caller-owned executor, which also runs any nested group-by finalization work.
+ /// The caller owns its lifetime; this method does not shut it down. [org.apache.pinot.spi.query.QueryThreadContext]
+ /// must already be set up, and an asynchronous executor must propagate that context.
+ public BrokerResponseNative reduceOnDataTable(BrokerRequest brokerRequest, BrokerRequest serverBrokerRequest,
+ Map<ServerRoutingInstance, DataTable> dataTableMap, long reduceTimeOutMs, BrokerMetrics brokerMetrics,
+ ExecutorService executorService) {
if (dataTableMap.isEmpty()) {
// Empty response.
return BrokerResponseNative.empty();
@@ -112,7 +123,7 @@ public class BrokerReduceService extends BaseReduceService {
QueryContext serverQueryContext = QueryContextConverterUtils.getQueryContext(serverBrokerRequest.getPinotQuery());
DataTableReducer dataTableReducer = ResultReducerFactory.getResultReducer(serverQueryContext);
- DataTableReducerContext reducerContext = createReducerContext(queryOptions, reduceTimeOutMs);
+ DataTableReducerContext reducerContext = createReducerContext(queryOptions, reduceTimeOutMs, executorService);
try {
dataTableReducer.reduceAndSetResults(rawTableName, cachedDataSchema, dataTableMap, brokerResponseNative,
@@ -228,7 +239,8 @@ public class BrokerReduceService extends BaseReduceService {
QueryContext serverQueryContext = QueryContextConverterUtils.getQueryContext(serverBrokerRequest.getPinotQuery());
DataTableReducer dataTableReducer = ResultReducerFactory.getResultReducer(serverQueryContext);
DataTableReducerContext reducerContext =
- createReducerContext(serverBrokerRequest.getPinotQuery().getQueryOptions(), reduceTimeOutMs);
+ createReducerContext(serverBrokerRequest.getPinotQuery().getQueryOptions(), reduceTimeOutMs,
+ _reduceExecutorService);
DataTable merged = dataTableReducer.mergeDataTablesOnly(rawTableName, cachedDataSchema, dataTableMap,
reducerContext, brokerMetrics);
@@ -303,7 +315,7 @@ public class BrokerReduceService extends BaseReduceService {
/// Resolves the group-by trim parameters (query option overrides, else broker defaults) and builds the
/// [DataTableReducerContext]
private DataTableReducerContext createReducerContext(@Nullable Map<String, String> queryOptions,
- long reduceTimeOutMs) {
+ long reduceTimeOutMs, ExecutorService executorService) {
Integer minGroupTrimSizeQueryOption = null;
Integer groupTrimThresholdQueryOption = null;
Integer minInitialIndexedTableCapacityQueryOption = null;
@@ -318,7 +330,7 @@ public class BrokerReduceService extends BaseReduceService {
int minInitialIndexedTableCapacity =
minInitialIndexedTableCapacityQueryOption != null ? minInitialIndexedTableCapacityQueryOption
: _minInitialIndexedTableCapacity;
- return new DataTableReducerContext(_reduceExecutorService, _maxReduceThreadsPerQuery, reduceTimeOutMs,
+ return new DataTableReducerContext(executorService, _maxReduceThreadsPerQuery, reduceTimeOutMs,
groupTrimThreshold, minGroupTrimSize, minInitialIndexedTableCapacity);
}
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/transport/grpc/GrpcQueryServer.java b/pinot-core/src/main/java/org/apache/pinot/core/transport/grpc/GrpcQueryServer.java
index 627aa46df09..7f5bccc02f5 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/transport/grpc/GrpcQueryServer.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/transport/grpc/GrpcQueryServer.java
@@ -359,6 +359,10 @@ public class GrpcQueryServer extends PinotQueryServerGrpc.PinotQueryServerImplBa
ServerResponse serverResponse;
try {
DataTable dataTable = instanceResponse.toDataTable();
+ if (_nativeQueryGate != null) {
+ dataTable.getMetadata().put(DataTable.MetadataKey.REQUEST_ID.getName(),
+ Long.toString(queryRequest.getRequestId()));
+ }
serverResponse = queryRequest.isEnableStreaming() ? StreamingResponseUtils.getMetadataResponse(dataTable)
: StreamingResponseUtils.getNonStreamingResponse(dataTable);
} catch (Exception e) {
diff --git a/pinot-core/src/test/java/org/apache/pinot/core/transport/grpc/NativeGrpcQueryGuardTest.java b/pinot-core/src/test/java/org/apache/pinot/core/transport/grpc/NativeGrpcQueryGuardTest.java
index d0cb6790f44..a95462737cb 100644
--- a/pinot-core/src/test/java/org/apache/pinot/core/transport/grpc/NativeGrpcQueryGuardTest.java
+++ b/pinot-core/src/test/java/org/apache/pinot/core/transport/grpc/NativeGrpcQueryGuardTest.java
@@ -38,6 +38,7 @@ import java.security.MessageDigest;
import java.security.cert.Certificate;
import java.security.cert.CertificateExpiredException;
import java.security.cert.X509Certificate;
+import java.util.HashMap;
import java.util.HexFormat;
import java.util.Iterator;
import java.util.List;
@@ -57,6 +58,7 @@ import org.apache.pinot.common.config.provider.NativeOfflineSegmentAction.Object
import org.apache.pinot.common.config.provider.NativeOfflineSegmentAction.Participant;
import org.apache.pinot.common.config.provider.NativeOfflineSegmentAction.Role;
import org.apache.pinot.common.datatable.DataTable;
+import org.apache.pinot.common.datatable.DataTableFactory;
import org.apache.pinot.common.metrics.ServerMetrics;
import org.apache.pinot.common.proto.PinotQueryServerGrpc;
import org.apache.pinot.common.proto.Server.ServerRequest;
@@ -64,6 +66,10 @@ import org.apache.pinot.common.proto.Server.ServerResponse;
import org.apache.pinot.common.request.context.NativeQueryContext;
import org.apache.pinot.common.request.context.NativeQueryDeadline;
import org.apache.pinot.common.request.context.NativeQueryGate;
+import org.apache.pinot.common.utils.DataSchema;
+import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
+import org.apache.pinot.core.common.datatable.DataTableBuilder;
+import org.apache.pinot.core.common.datatable.DataTableBuilderFactory;
import org.apache.pinot.core.operator.blocks.InstanceResponseBlock;
import org.apache.pinot.core.operator.blocks.results.BaseResultsBlock;
import org.apache.pinot.core.query.executor.QueryExecutor;
@@ -156,7 +162,9 @@ public class NativeGrpcQueryGuardTest {
assertTrue(responses.hasNext());
responses.next();
assertTrue(responses.hasNext());
- responses.next();
+ DataTable terminal = DataTableFactory.getDataTable(responses.next().getPayload().asReadOnlyByteBuffer());
+ assertEquals(terminal.getMetadata().get(DataTable.MetadataKey.REQUEST_ID.getName()), "42",
+ "Native responses must carry the exact admitted request ID");
assertFalse(responses.hasNext());
assertEquals(peer.get(), tls._fingerprint, "Request metadata must not supply the authenticated identity");
assertEquals(callbacks.get(), 3, "Streaming data, final metadata and completion must all be guarded");
@@ -218,7 +226,11 @@ public class NativeGrpcQueryGuardTest {
Iterator<ServerResponse> responses = legacy.submit(request().toBuilder().clearNativeQueryContext().build());
int count = 0;
while (responses.hasNext()) {
- assertNotNull(responses.next());
+ ServerResponse response = responses.next();
+ assertNotNull(response);
+ assertFalse(DataTableFactory.getDataTable(response.getPayload().asReadOnlyByteBuffer())
+ .getMetadata().containsKey(DataTable.MetadataKey.REQUEST_ID.getName()),
+ "Native response stamping must not change the legacy gRPC path");
count++;
}
assertEquals(count, 2);
@@ -242,6 +254,7 @@ public class NativeGrpcQueryGuardTest {
List.of(new NativeQueryContext.Segment("segment", object(), "a".repeat(64), object(), UUID.randomUUID())),
List.of()).withDeadline(new NativeQueryDeadline("test-clock", 10, 100));
return ServerRequest.newBuilder().setSql("SELECT count(*) FROM tbl_OFFLINE").addSegments("segment")
+ .putMetadata(Request.MetadataKeys.REQUEST_ID, "42")
.putMetadata(Request.MetadataKeys.BROKER_ID, "broker")
.putMetadata(Request.MetadataKeys.ENABLE_STREAMING, "true")
.putMetadata("native-query-peer-certificate", "forged-request-identity")
@@ -259,12 +272,16 @@ public class NativeGrpcQueryGuardTest {
private Fixture(TlsFixture tls, NativeQueryGate gate) throws Exception {
ServerMetrics.register(mock(ServerMetrics.class));
- DataTable data = mock(DataTable.class);
- when(data.toBytes()).thenReturn(new byte[]{1, 2});
+ DataTableBuilder builder = DataTableBuilderFactory.getDataTableBuilder(
+ new DataSchema(new String[]{"count(*)"}, new ColumnDataType[]{ColumnDataType.LONG}));
+ builder.startRow();
+ builder.setColumn(0, 1L);
+ builder.finishRow();
+ DataTable data = builder.build();
BaseResultsBlock block = mock(BaseResultsBlock.class);
when(block.getDataTable()).thenReturn(data);
- InstanceResponseBlock result = mock(InstanceResponseBlock.class);
- when(result.toDataTable()).thenReturn(data);
+ when(block.getResultsMetadata()).thenReturn(new HashMap<>());
+ InstanceResponseBlock result = new InstanceResponseBlock(block);
when(_executor.execute(any(), any(), any())).thenAnswer(invocation -> {
ResultsBlockStreamer streamer = invocation.getArgument(2);
streamer.send(block);
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org