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 <[email protected]> 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: [email protected] For additional commands, e-mail: [email protected]
