hvanhovell commented on code in PR #55683: URL: https://github.com/apache/spark/pull/55683#discussion_r3344073664
########## udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/grpc/GrpcWorkerSession.scala: ########## @@ -0,0 +1,773 @@ +/* + * 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.spark.udf.worker.core.grpc + +import java.util.concurrent.{CountDownLatch, LinkedBlockingQueue, TimeoutException, TimeUnit} +import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference} + +import scala.util.control.NonFatal + +import com.google.protobuf.ByteString +import io.grpc.{ConnectivityState, ManagedChannel} +import io.grpc.stub.StreamObserver + +import org.apache.spark.annotation.Experimental +import org.apache.spark.udf.worker.{Cancel, DataRequest, ExecutionError, Finish, Init, + UdfControlRequest, UdfControlResponse, UdfRequest, UdfResponse, UdfWorkerGrpc} +import org.apache.spark.udf.worker.core.{WorkerHandle, WorkerLogger, WorkerSession} +import org.apache.spark.udf.worker.core.grpc.GrpcWorkerSession._ + +/** + * :: Experimental :: + * gRPC implementation of [[WorkerSession]] for the `UdfWorker.Execute` + * bidirectional RPC. + * + * Drives one bidirectional `Execute` stream against the worker per the + * ordering invariants documented in `udf_protocol.proto`: + * {{{ + * Engine -> Worker: Init -> (DataRequest)* -> Finish (Cancel)? + * | Cancel + * Worker -> Engine: InitResponse -> (DataResponse)* -> (ErrorResponse)? -> (FinishResponse | CancelResponse) + * }}} + * + * Knows nothing about how the worker was provisioned (locally spawned, + * indirectly looked up, ...) -- the dispatcher constructs this with a + * [[WorkerHandle]] and channel; the base [[WorkerSession]] handles + * dispatcher-side cleanup on close. + * + * Threading: + * - [[doInit]] is synchronous: sends `Init` and blocks on `InitResponse`. + * - [[doProcess]] returns an iterator. Input batches are forwarded inline + * (the iterator's `next()` thread also sends `DataRequest`). Output + * batches arrive via the response observer (gRPC callback thread) and + * are consumed by the same iterator. A terminator (`FinishResponse`, + * `CancelResponse`, `ErrorResponse`, gRPC stream error) is published + * once. + * - [[cancel]] is thread-safe and idempotent. + * - [[doClose]] half-closes the request stream after the terminator arrives. + * + * TODO [SPARK-55278]: this class does not yet implement payload chunking; + * the entire [[Init.udf]] payload is sent inline. Chunking will be added + * when a UDF payload large enough to exceed gRPC's default message size + * limit is introduced. + * + * @param workerHandle dispatcher-side handle for releasing the worker on + * [[close]] (see [[WorkerSession]]). + * @param channel built and owned by the caller (typically a + * [[GrpcWorkerChannel]]). Not closed here -- the + * dispatcher tears it down via [[WorkerHandle]]. + * @param logger diagnostics. Defaults to [[WorkerLogger.NoOp]]. + * @param initResponseTimeoutMs upper bound on the wait for `InitResponse` + * after [[doInit]] sends `Init`. + * @param terminalTimeoutMs upper bound on the wait for a stream + * terminator (`FinishResponse`, + * `CancelResponse`, or `ErrorResponse`). + * Each output-queue poll resets this wait; + * see [[doProcess]] / `ProcessIterator`. + */ +@Experimental +class GrpcWorkerSession( + workerHandle: WorkerHandle, + channel: ManagedChannel, + logger: WorkerLogger = WorkerLogger.NoOp, + initResponseTimeoutMs: Long = DEFAULT_INIT_RESPONSE_TIMEOUT_MS, + terminalTimeoutMs: Long = DEFAULT_TERMINAL_TIMEOUT_MS) + extends WorkerSession(workerHandle, logger) { + + require(channel != null, "channel is required") + + private val asyncStub = UdfWorkerGrpc.newStub(channel) + + // Output batches received from the worker, drained by the iterator returned + // from process(). Intentionally unbounded -- the alternative (a bounded + // queue that the gRPC callback thread blocks on when full) would stall the + // Netty event loop and back-pressure the entire channel, including control + // messages, which makes cancel / terminator delivery latency-sensitive. + // + // Today this is safe because gRPC's HTTP/2 flow control bounds in-flight + // bytes on the wire and the engine consumer typically drains promptly; the + // queue holds at most a small number of decoded batches. + // + // Future work -- protocol-level back-pressure: a slow downstream consumer + // (e.g. a Spark operator that stalls) can in principle let the queue grow, + // because flow control only gates the wire path, not the in-heap queue. + // The intended fix is at the protocol layer, not the queue: extend the + // proto so the engine can ask the worker to hold UDF execution while the + // downstream is behind (e.g. periodic data-request acks with a sequence + // number, or a dedicated pause/resume control message). That keeps the + // gRPC receive path non-blocking while still bounding driver memory. + // We deliberately do NOT bound this queue: blocking the gRPC receiving + // thread is worse than the current unbounded-but-fast behaviour. A proto + // change is out of scope here. + private val outputQueue = new LinkedBlockingQueue[QueueItem]() + + // Latch fired when `InitResponse` (success or error) or a transport error + // arrives. init() blocks on this; until it fires we have no proof the + // worker actually accepted the session. + private val initLatch = new CountDownLatch(1) + private val initError = new AtomicReference[Option[ExecutionError]](None) + + // Holds the terminal outcome of the stream. Set exactly once by whichever + // response observer event reaches a terminator first. + private val terminal = new AtomicReference[Option[Terminal]](None) + private val terminalLatch = new CountDownLatch(1) + + // Captures an ErrorResponse encountered during the data phase so that + // the CancelResponse handler can attribute the failure to the original + // user / worker / protocol error rather than reporting a bare "Cancelled". + private val executionError = new AtomicReference[Option[ExecutionError]](None) + + // True once further engine-to-worker writes are invalid (terminal arrived + // from the worker, transport error, or last outgoing write failed). Guards + // [[sendRequest]] / [[sendCancelInternal]] against post-terminal writes. + // Note: this is distinct from [[halfClosed]] -- "no more writes are valid" + // and "we have called requestObserver.onCompleted()" are independent. + private val outgoingClosed = new AtomicBoolean(false) + + // True once we have invoked requestObserver.onCompleted() to half-close + // the request side. Set exactly once by [[close]]. + private val halfClosed = new AtomicBoolean(false) + + // Tracks whether cancel() has been invoked. ProcessIterator uses this to + // honour the "Cancel MUST NOT precede Finish" invariant; sendRequest + // re-reads it inside [[requestLock]] so a Data/Finish never races past a + // concurrent Cancel on the wire. + private val cancelRequested = new AtomicBoolean(false) + + // Latched on InitResponse success. + private val initSucceeded = new AtomicBoolean(false) + + // gRPC requires serialized writes to a request StreamObserver. + private val requestLock = new Object + + // Initialised in init() -- before that, cancel() / close() are no-ops on the + // request side, which is exactly the contract the wrapping WorkerSession + // expects. + @volatile private var requestObserver: StreamObserver[UdfRequest] = _ + + private val responseObserver: StreamObserver[UdfResponse] = new StreamObserver[UdfResponse] { + override def onNext(response: UdfResponse): Unit = { + response.getResponseCase match { + case UdfResponse.ResponseCase.DATA => + outputQueue.put(QueueItem.Batch(response.getData.getData.toByteArray)) + + case UdfResponse.ResponseCase.CONTROL => + handleControl(response.getControl) + + case other => + completeTerminal(Terminal.TransportError(new IllegalStateException( + s"unexpected response oneof: $other"))) + } + } + + override def onError(t: Throwable): Unit = { + // Transport-level failure: the stream is dead. No further writes are + // possible. Unblock any thread waiting on InitResponse so init() + // surfaces the transport cause instead of timing out after + // initResponseTimeoutMs with a misleading "timed out" message. + outgoingClosed.set(true) + initLatch.countDown() + completeTerminal(Terminal.TransportError(t)) + } + + override def onCompleted(): Unit = { + // Worker half-closed its side without sending a terminator (FinishResponse + // / CancelResponse). Treat as transport error so the engine sees a + // failure, not a silent end-of-stream. + if (terminal.get().isEmpty) { + completeTerminal(Terminal.TransportError(new IllegalStateException( + "worker response stream closed without a terminator"))) + } + // Defensive: if onCompleted reached us before InitResponse, doInit is + // still blocked on initLatch and would otherwise time out. + initLatch.countDown() + } + } + + private def handleControl(ctrl: UdfControlResponse): Unit = ctrl.getControlCase match { + case UdfControlResponse.ControlCase.INIT => + val resp = ctrl.getInit + if (resp.hasError) { + initError.set(Some(resp.getError)) + initLatch.countDown() + sendCancelInternal("init failed") + } else { + initSucceeded.set(true) + initLatch.countDown() + } + + case UdfControlResponse.ControlCase.ERROR => + val err = ctrl.getError.getError + executionError.compareAndSet(None, Some(err)) + // ERROR before InitResponse means init will never succeed; surface it + // to doInit immediately instead of waiting for the eventual CancelResponse. + // countDown is idempotent so this is safe in the normal data-phase path too. + initLatch.countDown() + // ERROR before InitResponse is special: under directExecutor or with + // a fast worker, sendCancelInternal may run while requestObserver is + // still null (we're inside stream.onNext(initRequest) and have not + // yet published the field), so no Cancel reaches the wire and no + // CancelResponse will follow to settle the terminal. Settle it here + // so doInit doesn't return as if init succeeded. + // completeTerminal is idempotent, so the data-phase ERROR -> Cancel -> + // CancelResponse path is unaffected. + if (!initSucceeded.get()) { + completeTerminal(Terminal.ExecutionFailed(err)) + } + sendCancelInternal("aborting after ErrorResponse") + + case UdfControlResponse.ControlCase.FINISH => + val resp = ctrl.getFinish + val maybeFinishErr = + if (resp.hasError) Some(resp.getError) else None + val priorErr = executionError.get() + val combined = priorErr.orElse(maybeFinishErr) + outgoingClosed.set(true) + combined match { + case Some(err) => completeTerminal(Terminal.ExecutionFailed(err)) + case None => completeTerminal(Terminal.FinishOk) + } + // Defensive: FINISH before InitResponse is a worker protocol bug, but + // we should fail init fast rather than hang the 30s init timeout. + initLatch.countDown() + + case UdfControlResponse.ControlCase.CANCEL => + val resp = ctrl.getCancel + val cancelErr = if (resp.hasError) Some(resp.getError) else None + val priorErr = executionError.get() + val effective = priorErr + .orElse(initError.get()) + .orElse(cancelErr) + outgoingClosed.set(true) + effective match { + case Some(err) => completeTerminal(Terminal.ExecutionFailed(err)) + case None => completeTerminal(Terminal.Cancelled) + } + // Defensive: CANCEL before InitResponse unblocks doInit so it can + // surface the cancellation instead of timing out. + initLatch.countDown() + + case UdfControlResponse.ControlCase.CONTROL_NOT_SET => + completeTerminal(Terminal.TransportError(new IllegalStateException( + "empty UdfControlResponse oneof"))) + initLatch.countDown() + } + + private def completeTerminal(t: Terminal): Unit = { + if (terminal.compareAndSet(None, Some(t))) { + // Wake any iterator blocked on outputQueue.take() / poll(). + outputQueue.put(QueueItem.EndOfStream) + terminalLatch.countDown() + } + } + + // ---- WorkerSession hooks ------------------------------------------------ + + override protected def doInit(message: Init): Unit = { + // Fail fast if the channel is already shut down. Without this check, + // asyncStub.execute(...) would still succeed and the failure would + // surface ~initResponseTimeoutMs later as a misleading "InitResponse + // timed out" error. + if (channel.getState(false) == ConnectivityState.SHUTDOWN) { + val ex = new IllegalStateException("gRPC channel is shut down") + completeTerminal(Terminal.TransportError(ex)) + throw new GrpcWorkerSessionException("UDF worker channel is closed", ex) + } + // Open the stream as a local first; only publish `requestObserver` + // AFTER the Init has been put on the wire. `cancel()` reads + // `requestObserver` outside [[requestLock]] and returns early when + // it is null, so this ordering prevents a concurrent cancel from + // sneaking a Cancel ahead of Init. + val stream = asyncStub.execute(responseObserver) + val initRequest = UdfRequest.newBuilder() + .setControl(UdfControlRequest.newBuilder().setInit(message).build()) + .build() + try { + requestLock.synchronized { + stream.onNext(initRequest) + requestObserver = stream + } + } catch { + case NonFatal(e) => + // Expose the stream so a subsequent close() can still attempt a + // best-effort half-close; the terminal already reflects the failure. + requestObserver = stream + outgoingClosed.set(true) + completeTerminal(Terminal.TransportError(e)) + throw e + } + + val responded = try { + initLatch.await(initResponseTimeoutMs, TimeUnit.MILLISECONDS) + } catch { + case _: InterruptedException => + Thread.currentThread().interrupt() + sendCancelInternal("interrupted during init") + // Make sure close() does not block waiting for a terminator that + // will never arrive on this thread's behalf. + completeTerminal(Terminal.TransportError( + new InterruptedException("interrupted while waiting for InitResponse"))) + throw new InterruptedException("interrupted while waiting for InitResponse") + } + + if (!responded) { + sendCancelInternal("InitResponse timed out") + // Settle the terminal so a subsequent close() does not stall for a + // second `terminalTimeoutMs` waiting for a worker that already missed + // its init deadline. + completeTerminal(Terminal.TransportError(new TimeoutException( + s"timed out waiting for InitResponse after ${initResponseTimeoutMs}ms"))) + throw new IllegalStateException( + s"timed out waiting for InitResponse after ${initResponseTimeoutMs}ms") + } + + initError.get() match { + case Some(err) => + // The protocol requires the engine to send Cancel after an init + // error and the worker to respond with CancelResponse. Drain it + // before throwing so we don't leave a dangling stream. + awaitTerminal() + throw new GrpcWorkerSessionException( + s"UDF worker init failed: ${describeError(err)}", err) + case None => + // No InitResponse was observed but the latch fired. The defensive + // initLatch.countDown() in handleControl / onError / onCompleted + // means we get here when the worker terminated the stream before + // sending InitResponse. Surface that as an init failure rather than + // letting the caller proceed as if init succeeded. + terminal.get() match { + case Some(Terminal.TransportError(cause)) => + throw new GrpcWorkerSessionException( + "UDF worker stream failed during init", cause) + case Some(Terminal.ExecutionFailed(err)) => + throw new GrpcWorkerSessionException( + s"UDF worker reported an error before init completed: " + + describeError(err), err) + case Some(Terminal.Cancelled) => + throw new GrpcWorkerSessionException( + "UDF worker stream was cancelled before init completed") + case Some(Terminal.FinishOk) => + throw new GrpcWorkerSessionException( + "UDF worker finished before init completed") + case None => + // success: InitResponse arrived and was a success + } + } + } + + override protected def doProcess( + input: Iterator[Array[Byte]]): Iterator[Array[Byte]] = { + require(initSucceeded.get(), + "process called without a successful init " + + "(WorkerSession lifecycle guard should normally prevent this)") + new ProcessIterator(input) + } + + /** Thread-safe; idempotent. See [[WorkerSession#cancel]] for the contract. */ + override def cancel(): Unit = sendCancelInternal("session.cancel()") + + /** + * The worker is salvageable when the session ended cleanly + * ([[Terminal.FinishOk]]) or via cooperative engine-initiated + * cancellation. Any other terminal -- transport error, execution + * failure, ErrorResponse -- leaves the worker in an unknown state and + * the wrapping [[org.apache.spark.udf.worker.core.WorkerSession]] must + * mark `workerHandle` invalid so a future pool will not recycle the + * worker. + * + * UDF-level resource cleanup (anything the UDF allocated while + * processing partial input) is the UDF's responsibility via its + * cleanup callback; it is not a salvageability concern at the protocol + * layer. + */ + override protected def isWorkerSalvageable: Boolean = terminal.get() match { + case Some(Terminal.FinishOk) => true + // Engine-initiated cancel ran cooperatively; the worker is in a clean + // post-cancel state and is safe to reuse. + case Some(Terminal.Cancelled) => true + case Some(Terminal.ExecutionFailed(_)) => false + case Some(Terminal.TransportError(_)) => false + case None => + // close() always awaits a terminator; reaching here means the worker + // failed to settle within the timeout, which counts as unsafe. + false + } + + override protected def doClose(): Unit = { + if (requestObserver == null) return + // If the terminator hasn't arrived yet (e.g. cancel() invoked, worker + // mid-cleanup), wait briefly. Bound the wait so a misbehaving worker + // cannot pin close() forever. + if (terminal.get().isEmpty) { + try { + terminalLatch.await(terminalTimeoutMs, TimeUnit.MILLISECONDS) + } catch { + case _: InterruptedException => Thread.currentThread().interrupt() + } + } + // If the worker still has not settled (timeout or interrupt above), + // record a terminal so isWorkerSalvageable returns a definite answer + // and any other thread reading `terminal` sees a stable value. + if (terminal.get().isEmpty) { + completeTerminal(Terminal.TransportError(new TimeoutException( + s"timed out waiting for stream terminator after ${terminalTimeoutMs}ms"))) + } + // Half-close the request side regardless of whether we sent Cancel or + // Finish earlier -- the worker's onCompleted handler relies on this to + // tear down the per-stream state cleanly. Independent from + // [[outgoingClosed]], which gates writes: a session may legitimately + // have written its last message (Cancel/Finish) and still need to + // half-close. + if (halfClosed.compareAndSet(false, true)) { + try { + requestLock.synchronized { requestObserver.onCompleted() } + } catch { + case NonFatal(e) => + // The stream may already be torn down (transport error, server + // half-close ahead of us). Either way, onCompleted is best-effort. + logger.debug("Error half-closing UDF Execute stream", e) + } + } + } + + // ---- Internal request helpers --------------------------------------------- + + /** + * Sends a Data or Finish request to the worker. Two invariants are + * checked inside [[requestLock]]: + * - [[outgoingClosed]]: writes are unsafe (transport dead or terminal + * received). Throws, so the caller's terminal/exception path runs. + * - [[cancelRequested]]: a Cancel has been (or is about to be) sent. + * Silently no-ops so a Data/Finish never appears on the wire after + * Cancel; the caller's iterator falls through to the terminator wait. + * + * Init is NOT sent through this helper -- [[doInit]] writes Init + * directly under [[requestLock]] before publishing + * [[requestObserver]], so there is no path through which a concurrent + * cancel could send Cancel before Init. + */ + private def sendRequest(req: UdfRequest): Unit = + requestLock.synchronized { + if (outgoingClosed.get()) { + throw new IllegalStateException( + "cannot send request: UDF Execute stream is already closed") + } + if (cancelRequested.get()) { + // A concurrent cancel() either has been or is about to be flushed + // through this lock; suppressing the write preserves the proto + // ordering invariant (no Data/Finish after Cancel). + return + } + requestObserver.onNext(req) + } + + /** + * Sends a `Cancel` control message. Idempotent across ALL call sites + * (the public `cancel()`, in-band cancels from `handleControl`'s INIT + * error / ERROR paths, and engine-internal cancels from the iterator) + * via [[cancelRequested]]: only the first caller actually puts a + * `Cancel` on the wire. Setting [[cancelRequested]] also blocks + * subsequent Data/Finish writes from [[ProcessIterator]] (see + * [[sendRequest]]), preserving the proto invariant that nothing + * follows `Cancel` on the engine-to-worker side. + */ + private def sendCancelInternal(reason: String): Unit = { + if (!cancelRequested.compareAndSet(false, true)) return + if (requestObserver == null) return + if (outgoingClosed.get()) return + try { + requestLock.synchronized { + if (outgoingClosed.get()) return + requestObserver.onNext(UdfRequest.newBuilder() + .setControl(UdfControlRequest.newBuilder() + .setCancel(Cancel.newBuilder().setReason(reason).build()) + .build()) + .build()) + } + } catch { + case NonFatal(e) => + logger.debug(s"Cancel send failed (stream may already be torn down): ${e.getMessage}") + outgoingClosed.set(true) + completeTerminal(Terminal.TransportError(e)) + } + } + + private def awaitTerminal(): Unit = { + if (terminal.get().isDefined) return + val ok = try { + terminalLatch.await(terminalTimeoutMs, TimeUnit.MILLISECONDS) + } catch { + case _: InterruptedException => + Thread.currentThread().interrupt() + false + } + if (!ok) { + completeTerminal(Terminal.TransportError(new IllegalStateException( + s"timed out waiting for stream terminator after ${terminalTimeoutMs}ms"))) + } + } + + // ---- ProcessIterator ------------------------------------------------------ + + /** + * Iterator returned by [[doProcess]]. Drives the data phase of the + * stream end-to-end: each call to `hasNext` / `next` may send a + * `DataRequest` or `Finish` to the worker, and reads result batches + * out of [[outputQueue]] as the worker emits them. + * + * The iterator is single-threaded with respect to the engine, but + * coexists with the gRPC callback thread (which enqueues responses + * and may set [[terminal]] / [[outgoingClosed]]) and with any thread + * that calls [[cancel]]. + */ + private class ProcessIterator(input: Iterator[Array[Byte]]) + extends Iterator[Array[Byte]] { + + private val finishedSending = new AtomicBoolean(false) + // Latched once the terminator sentinel has been observed. Without this, + // a second hasNext() after the iterator naturally exhausts would re-enter + // advance(), fall through all four branches, and block branch 4 for the + // full terminalTimeoutMs before returning. Callers that probe hasNext + // an extra time (iterator.size, instrumentation wrappers) would hang. + private val exhausted = new AtomicBoolean(false) + @volatile private var prefetched: Array[Byte] = _ + + // Pre-init cancel fast-fail: if cancel() ran before doInit published the + // requestObserver, no Cancel was sent on the wire and the worker may + // never emit a terminator. Trip the terminal explicitly so the first + // hasNext / next surfaces a cancellation instead of blocking for + // terminalTimeoutMs in branch 4. + if (cancelRequested.get() && terminal.get().isEmpty) { + completeTerminal(Terminal.Cancelled) + } + + override def hasNext: Boolean = { + if (prefetched ne null) return true + advance() + prefetched ne null + } + + override def next(): Array[Byte] = { + if (prefetched eq null) advance() + val out = prefetched + if (out eq null) { + throw new NoSuchElementException("ProcessIterator exhausted") + } + prefetched = null + out + } + + /** + * Four-branch state machine in a single loop. Each iteration tries + * the branches in order; branch 1 always runs first so any output + * already produced is delivered immediately. + * + * 1. '''Drain output already available.''' A `DataResponse` is in + * [[outputQueue]] (or the terminal sentinel). Pull it out and + * either return the batch or surface the terminator. + * + * 2. '''Send the next input batch.''' Input has more data and the + * session is not cancelled. Pull the next batch from the user's + * iterator and ship it as a `DataRequest`. After this the loop + * continues -- branch 1 may now find new output, or we fall + * through to branch 2 again for the next batch. + * + * 3. '''Send Finish.''' Input is exhausted and the session is not + * cancelled. Send `Finish` exactly once (CAS on + * [[finishedSending]]). After this, branches 2 and 3 are both + * short-circuited and the loop will reach branch 4 to await + * trailing output and the `FinishResponse` terminator. + * + * 4. '''Block waiting for late output.''' Nothing queued, nothing + * more to send. Reached only after `Finish` has been sent (or + * cancellation requested); while input is still being fed the + * loop just alternates between drain and send. + * + * '''Per-poll timeout.''' Each poll waits up to + * [[terminalTimeoutMs]] and every event from the worker resets + * the wait. The worker-side contract is "emit at least one + * event every [[terminalTimeoutMs]] after Finish", not + * "complete the UDF within [[terminalTimeoutMs]]". For + * reduce-style UDFs the timeout window typically only covers + * tail processing: the worker can fold each in-flight batch as + * it arrives, so by the time Finish lands most of the work is + * already done. + * + * '''Liveness via empty response.''' A worker that expects a + * long post-Finish silence MAY emit an empty `DataResponse` + * purely to reset the per-poll wait. The engine surfaces it to + * the caller, so the batch should be one the caller recognises + * as a heartbeat (e.g. an empty Arrow IPC batch with zero + * rows). + */ + private def advance(): Unit = { + // After the terminator has been observed once, further hasNext() / + // next() calls must return immediately. Without this short-circuit + // a probing caller would re-enter branch 4 and block on outputQueue + // for the full terminalTimeoutMs. + if (exhausted.get()) return + while (prefetched eq null) { + // (1) Drain anything the worker has already produced. + outputQueue.poll() match { + case null => // queue empty, fall through to send/wait branches below + case QueueItem.EndOfStream => + exhausted.set(true) + throwIfTerminalError() + return + case QueueItem.Batch(b) => + prefetched = b + return + } + + // (2) Send next input batch when input has data and session is live. + if (!cancelRequested.get() && !finishedSending.get() && input.hasNext) { + val batch = try input.next() catch { Review Comment: This sends a singular batch per iteration. Shouldn't we try to fill the send queue to a certain level? -- 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]
