gortiz commented on code in PR #18458:
URL: https://github.com/apache/pinot/pull/18458#discussion_r3356942870
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/service/dispatch/DispatchClient.java:
##########
@@ -130,6 +133,38 @@ public void submit(Worker.QueryRequest request,
QueryServerInstance virtualServe
_dispatchStub.withDeadline(deadline).submit(request, new
LastValueDispatchObserver<>(virtualServer, callback));
}
+ /**
+ * Opens a {@code SubmitWithStream} bidi RPC for one server, sends the
initial {@code submit}, and registers the
+ * resulting {@link StreamingDispatchObserver} with {@code session} for
cancel fan-out and {@code OpChainComplete}
+ * accumulation.
+ *
+ * <p>The submit-ack callback is invoked exactly once: with the {@link
Worker.QueryResponse} on the first
+ * {@code submit_ack} from the server, or with a non-null {@link Throwable}
if the stream errors before the ack
+ * arrives.
+ *
+ * @param request the plan submission
+ * @param virtualServer server identity (used in callbacks for
routing decisions on failure)
+ * @param deadline gRPC deadline for the call
+ * @param session broker-side streaming session — the returned
observer registers itself here
+ * @param expectedOpChainsForThisServer number of opchains this server is
expected to report; used to drain the
+ * session latch correctly when the stream
errors before all opchains have responded
+ * @param ackCallback receives the submit-ack response or a
failure throwable
+ * @return the observer, also exposed as
+ * {@link
org.apache.pinot.query.service.dispatch.streaming.StreamingServerHandle} on the
session for
+ * cancel fan-out
+ */
+ public StreamingDispatchObserver submitWithStream(Worker.QueryRequest
request, QueryServerInstance virtualServer,
+ Deadline deadline, StreamingQuerySession session, int
expectedOpChainsForThisServer,
+ BiConsumer<Worker.QueryResponse, Throwable> ackCallback) {
+ StreamingDispatchObserver observer = new
StreamingDispatchObserver(virtualServer, session,
+ expectedOpChainsForThisServer, ackCallback);
+ StreamObserver<Worker.BrokerToServer> outbound =
_dispatchStub.withDeadline(deadline).submitWithStream(observer);
Review Comment:
You're right that there is no automatic fallback — that is intentional
(matching the `submitAndReduceWithStream` doc): stream mode requires every
server to implement `SubmitWithStream`. To make the rolling-upgrade hazard fail
*loud* instead of as an opaque query failure, 1476e498ba detects gRPC
`UNIMPLEMENTED` in `StreamingDispatchObserver.onError` and logs a clear message
telling the operator to disable stream stats (query option /
`pinot.broker.mse.stream.stats`) until all servers are upgraded. I'll also fix
the PR description's compat table to drop the "fallback" wording and match the
code.
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/service/dispatch/QueryDispatcher.java:
##########
@@ -243,19 +459,52 @@ private QueryResult tryRecover(long requestId,
Set<QueryServerInstance> servers,
} else if (ex instanceof QueryException) {
errorCode = ((QueryException) ex).getErrorCode();
} else {
- // in case of unknown exceptions, the exception will be rethrown, so we
don't need stats
cancel(requestId, servers);
throw ex;
}
- // in case of known exceptions (timeout or query exception), we need can
build here the erroneous QueryResult
- // that include the stats.
- LOGGER.warn("Query failed with a known exception. Trying to cancel the
other opchains");
- MultiStageQueryStats stats = cancelWithStats(requestId, servers);
- if (stats == null) {
+ LOGGER.warn("Query failed with a known exception. Cancelling remaining
opchains.");
+ cancel(requestId, servers);
+ QueryProcessingException processingException = new
QueryProcessingException(errorCode, ex.getMessage());
+ return new QueryResult(processingException,
MultiStageQueryStats.emptyStats(0), 0L);
Review Comment:
Acknowledged — this is a default-path behavior change (legacy `tryRecover`
returned partial stats via `cancelWithStats`; it now returns empty stats on a
known dispatch exception, trading per-failure cancel fan-out for lower
error-path load). I reworded the `QueryRunner.cancel` Javadoc in 1476e498ba so
it no longer reads as if the consumer was already gone — it makes clear this
change removes both the synchronous collection and its only consumer
(`tryRecover`). I'll add a release-note line to the PR description for the
behavior change.
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/service/dispatch/streaming/StreamingQuerySession.java:
##########
@@ -0,0 +1,317 @@
+/**
+ * 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.query.service.dispatch.streaming;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.LinkedHashSet;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.locks.ReentrantLock;
+import javax.annotation.Nullable;
+import org.apache.pinot.common.proto.Worker;
+import org.apache.pinot.query.runtime.plan.MultiStageStatsTreeDecoder;
+import org.apache.pinot.query.runtime.plan.StageStatsTreeNode;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * Broker-side per-query state for the {@code SubmitWithStream} dispatch path.
Owns the per-stage tree accumulator,
+ * the outstanding-opchain count, the per-stage coverage counters, and the set
of open server streams.
+ *
+ * <p>Concurrency model — all mutating methods acquire the per-session lock,
so the accumulator and counters need no
+ * additional internal synchronization. gRPC client {@code onNext} callbacks
land on I/O threads and call into this
+ * session directly; the work per call is short (decode + merge + decrement)
so doing it on the I/O thread is fine.
+ *
+ * <p>Completion semantics — {@link #awaitCompletion(long, TimeUnit)} returns
{@code true} as soon as every expected
+ * opchain has reported (early completion), and {@code false} if the timeout
fires first. The dispatcher should call
+ * it <strong>only after</strong> the broker receiving mailbox has finished,
so that a successful return means both
+ * "data done" and "stats fully accounted for". When it returns {@code false}
the per-stage coverage exposes which
+ * stages are missing.
+ */
+public class StreamingQuerySession {
+ private static final Logger LOGGER =
LoggerFactory.getLogger(StreamingQuerySession.class);
+
+ private final long _requestId;
+ private final int _expectedOpChains;
+ private final CountDownLatch _completionLatch;
+ private final ReentrantLock _lock = new ReentrantLock();
+
+ /** Per-stage merged accumulator. Mutated under {@link #_lock}. */
+ private final Map<Integer, StageStatsTreeNode> _stageAccumulator = new
HashMap<>();
+ /** Per-stage count of opchains that have responded successfully and merged
cleanly. */
+ private final Map<Integer, Integer> _respondedByStage = new HashMap<>();
+ /** Per-stage count of opchains that responded but the broker couldn't merge
their payload. */
+ private final Map<Integer, Integer> _mergeFailedByStage = new HashMap<>();
+
+ /** Set of open server streams. Iteration order is insertion order so cancel
fan-out is deterministic. */
+ private final Set<StreamingServerHandle> _openStreams = new
LinkedHashSet<>();
+
+ /** True after the first peer error (success=false OpChainComplete or stream
onError). Used to trigger fan-out
+ * cancel idempotently. */
+ private boolean _peerErrorObserved = false;
+
+ public StreamingQuerySession(long requestId, int expectedOpChains) {
+ _requestId = requestId;
+ _expectedOpChains = expectedOpChains;
+ _completionLatch = new CountDownLatch(expectedOpChains);
+ }
+
+ public long getRequestId() {
+ return _requestId;
+ }
+
+ public int getExpectedOpChains() {
+ return _expectedOpChains;
+ }
+
+ /**
+ * Registers an open server stream so the session can iterate them later for
fan-out cancel. Must be called by the
+ * dispatcher when the {@code SubmitWithStream} call is opened.
+ */
+ public void registerStream(StreamingServerHandle stream) {
+ _lock.lock();
+ try {
+ _openStreams.add(stream);
+ } finally {
+ _lock.unlock();
+ }
+ }
+
+ /**
+ * Removes a stream from the open-streams set. Called when the server emits
{@code ServerDone} (clean close) or the
+ * stream errors. Idempotent.
+ */
+ public void unregisterStream(StreamingServerHandle stream) {
+ _lock.lock();
+ try {
+ _openStreams.remove(stream);
+ } finally {
+ _lock.unlock();
+ }
+ }
+
+ /**
+ * Records an {@link Worker.OpChainComplete} message decoded from a server
stream. Decrements the outstanding count
+ * and merges the contained tree into the per-stage accumulator (or marks
the stage {@code mergeFailed} on a shape
+ * mismatch / decode failure). Also records {@code success=false} reports as
peer errors so fan-out cancel can fire.
+ */
+ public void recordOpChainComplete(Worker.OpChainComplete message) {
+ int stageId = message.getStageId();
+ boolean isSuccess = message.getSuccess();
+ Worker.MultiStageStatsTree statsTree = message.getStats();
+
+ boolean shouldFanOutCancel = false;
+ _lock.lock();
+ try {
+ if (!isSuccess) {
+ if (!_peerErrorObserved) {
+ _peerErrorObserved = true;
+ shouldFanOutCancel = true;
Review Comment:
The fan-out is now debounced to once per query for the first peer error
(`_peerErrorObserved`), so a flaky server no longer triggers N-1 cancels per
affected query. For observability, 1476e498ba adds the
`MSE_STREAM_STATS_CANCEL_FANOUTS` broker meter (counts non-empty cancel
broadcasts) so the storm condition is visible in prod. Distinguishing a
server-returned error from a pure transport close (so transport blips don't
cancel peers) is a reasonable refinement I'd like to do as a follow-up rather
than in this PR.
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/service/server/QueryServer.java:
##########
@@ -450,6 +454,255 @@ private void
submitTimeSeriesInternal(Worker.TimeSeriesQueryRequest request,
}
}
+ /// Stream-mode submission handler. The broker keeps the stream open for the
query lifetime; the server replies
+ /// with a {@code submit_ack} as the first message and (in subsequent
commits) per-opchain
+ /// {@link Worker.OpChainComplete} messages followed by a final {@link
Worker.ServerDone}.
+ ///
+ /// This skeleton wires up the gRPC mechanics + plan submission via the
existing submission path. It does NOT yet
+ /// emit OpChainComplete / ServerDone — those need a per-opchain completion
hook on
+ /// {@link org.apache.pinot.query.runtime.executor.OpChainSchedulerService},
which is layered on next.
+ /// Cancel still routes through the existing unary {@link
#cancel(Worker.CancelRequest, StreamObserver)} RPC; broker
+ /// stream-close also triggers a cancel here.
+ @Override
+ public StreamObserver<Worker.BrokerToServer> submitWithStream(
+ StreamObserver<Worker.ServerToBroker> responseObserver) {
+ return new SubmitWithStreamObserver(responseObserver);
+ }
+
+ /// Per-query state for an open {@code SubmitWithStream} call. Owns the
response stream and serialises every
+ /// {@code onNext} call on it via a {@code synchronized} block — gRPC
requires {@code StreamObserver.onNext} to be
+ /// called serially.
+ ///
+ /// Tracks the expected number of opchains for the request (sum of
WorkerMetadata across all stages). An
+ /// {@link OpChainCompletionListener} registered with {@link
QueryRunner#registerOpChainCompletionListener}
+ /// fires once per opchain finishing, encodes its stats via {@link
MultiStageStatsTreeEncoder}, and emits an
+ /// {@link Worker.OpChainComplete} on the response stream. When the
per-request completed-count reaches the
+ /// expected total, {@link Worker.ServerDone} is emitted and the stream is
closed.
+ ///
+ /// All blocking work (plan deserialization, opchain construction) runs on
+ /// {@link QueryServer#_submissionExecutorService}.
+ private final class SubmitWithStreamObserver implements
StreamObserver<Worker.BrokerToServer> {
+ private final StreamObserver<Worker.ServerToBroker> _responseObserver;
+ /// Serialises onNext calls on the response stream and guards mutable
session state.
+ private final Object _streamLock = new Object();
+ /// True once we've received the first {@code submit} and dispatched it.
+ private final AtomicBoolean _submitted = new AtomicBoolean(false);
+ /// True once we've completed the response stream (success or error).
Idempotent guard.
+ private final AtomicBoolean _completed = new AtomicBoolean(false);
+ /// Number of opchains we expect to report for this request — set after we
deserialize the plan.
+ private final AtomicInteger _expectedOpChains = new AtomicInteger(-1);
+ /// Number of opchains that have reported so far via the completion
listener.
+ private final AtomicInteger _completedOpChains = new AtomicInteger(0);
+ /// Set once we successfully parse the request id from the submit
metadata. Used by cancel-via-stream.
+ private volatile long _requestId = -1;
+
+ SubmitWithStreamObserver(StreamObserver<Worker.ServerToBroker>
responseObserver) {
+ _responseObserver = responseObserver;
+ }
+
+ @Override
+ public void onNext(Worker.BrokerToServer message) {
+ switch (message.getPayloadCase()) {
+ case SUBMIT:
+ handleSubmit(message.getSubmit());
+ break;
+ case CANCEL:
+ handleCancel(message.getCancel());
+ break;
+ case PAYLOAD_NOT_SET:
+ default:
+ sendErrorAndComplete("Unexpected BrokerToServer payload: " +
message.getPayloadCase());
+ break;
+ }
+ }
+
+ @Override
+ public void onError(Throwable t) {
+ // Broker-side stream error / disconnect. Treat like a cancel and clean
up; do not reply on the response stream
+ // (the underlying transport is gone).
+ LOGGER.warn("SubmitWithStream stream error for request {}: {}",
_requestId, t.getMessage());
+ _completed.set(true);
+ cleanupListener();
+ cancelIfSubmitted();
+ }
+
+ @Override
+ public void onCompleted() {
+ // Broker has half-closed (no more inbound messages). The server stream
stays open until all opchains have
+ // reported via the completion listener — it's the listener's job to
emit ServerDone and complete the stream.
+ // If the broker half-closes before the server is done, that's OK; we
keep emitting on the response stream
+ // until our own completion criterion is met.
+ }
+
+ private void handleSubmit(Worker.QueryRequest request) {
+ if (!_submitted.compareAndSet(false, true)) {
+ sendErrorAndComplete("Multiple submit messages on the same stream are
not allowed");
+ return;
+ }
+ ServerMetrics.get().addMeteredGlobalValue(ServerMeter.MSE_QUERIES, 1L);
+ Map<String, String> deserializedMetadata;
+ try {
+ deserializedMetadata =
QueryPlanSerDeUtils.fromProtoProperties(request.getMetadata());
+ } catch (Exception e) {
+ LOGGER.error("Caught exception while deserializing request metadata",
e);
+ sendErrorAndComplete("Caught exception while deserializing request
metadata: " + e.getMessage());
+ return;
+ }
+ // Override the cluster-level _sendStats decision for this request:
stats travel out-of-band on the bidi stream
+ // (via the OpChainCompletionListener), so we suppress the mailbox-side
stats path. The override is read by
+ // QueryRunner.effectiveSendStats(...).
+ Map<String, String> reqMetadata = new HashMap<>(deserializedMetadata);
+
reqMetadata.put(CommonConstants.MultiStageQueryRunner.KEY_OF_STATS_REPORTING_MODE,
+ CommonConstants.MultiStageQueryRunner.STATS_REPORTING_MODE_STREAM);
+ try {
+ _requestId = Long.parseLong(reqMetadata.get(MetadataKeys.REQUEST_ID));
+ } catch (Exception ignored) {
+ // _requestId stays at -1; cancel-on-stream-close will just be a no-op.
+ }
+ // Count how many opchains will run on this server: sum of
WorkerMetadata across all stage plans.
+ int opChainCount = 0;
+ for (Worker.StagePlan stagePlan : request.getStagePlanList()) {
+ opChainCount += stagePlan.getStageMetadata().getWorkerMetadataCount();
+ }
+ final int expected = opChainCount;
+ _expectedOpChains.set(expected);
+
+ // Register the per-request completion listener BEFORE submitting.
Otherwise short opchains could finish before
+ // we've registered and we'd miss their events.
+ if (_requestId >= 0 && expected > 0) {
+ _queryRunner.registerOpChainCompletionListener(_requestId,
this::onOpChainComplete);
+ }
+
+ long timeoutMs =
Long.parseLong(reqMetadata.get(QueryOptionKey.TIMEOUT_MS));
+ CompletableFuture.runAsync(() -> submitInternal(request, reqMetadata),
_submissionExecutorService)
+ .orTimeout(timeoutMs, TimeUnit.MILLISECONDS)
+ .whenComplete((result, error) -> {
+ if (error != null) {
+ LOGGER.error("Caught exception while submitting request: {}",
_requestId, error);
+ sendSubmitAck(buildErrorResponse("Caught exception while
submitting request: " + error.getMessage()));
+ // Submission failed — no opchains will run, so emit ServerDone
immediately and clean up.
+ cleanupListener();
+ sendDoneAndComplete();
+ } else {
+ sendSubmitAck(buildOkResponse());
+ // If for some reason expected was 0 (empty plan), close the
stream now.
+ if (expected == 0) {
+ cleanupListener();
+ sendDoneAndComplete();
+ }
+ }
+ });
+ }
+
+ /**
+ * Fires once per opchain on this server completing. Encodes the stats
into a {@link Worker.MultiStageStatsTree},
+ * emits an {@link Worker.OpChainComplete}, and emits {@link
Worker.ServerDone} once all expected opchains have
+ * reported.
+ */
+ private void onOpChainComplete(OpChainId opChainId, MultiStageOperator
rootOperator,
+ @Nullable MultiStageQueryStats stats, OpChainExecutionContext context,
@Nullable Throwable error) {
+ Worker.OpChainComplete.Builder builder =
Worker.OpChainComplete.newBuilder()
+ .setStageId(opChainId.getStageId())
+ .setWorkerId(opChainId.getVirtualServerId())
+ .setSuccess(error == null);
+ if (error != null) {
+ builder.setErrorMsg(error.getMessage() == null ?
error.getClass().getSimpleName() : error.getMessage());
+ }
+ if (stats != null) {
+ try {
+ builder.setStats(MultiStageStatsTreeEncoder.encode(rootOperator,
stats, context));
Review Comment:
Re the encoder limitation noted here: the pipeline-breaker "follow-up" that
`MultiStageStatsTreeEncoder`'s Javadoc referred to is now implemented in
2965797a77 (the encoder grafts the folded pipeline-breaker subtree so a
leaf-with-PB stage encodes correctly). The all-or-nothing encoder contract and
the `onNext`/increment liveness fix discussed above still stand.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
