You are viewing a plain text version of this content. The canonical link for it is here.
Posted to reviews@spark.apache.org by "haiyangsun-db (via GitHub)" <gi...@apache.org> on 2026/06/23 13:13:36 UTC
[PR] grpc session. [spark]
haiyangsun-db opened a new pull request, #56702:
URL: https://github.com/apache/spark/pull/56702
<!--
Thanks for sending a pull request! Here are some tips for you:
1. If this is your first time, please read our contributor guidelines: https://spark.apache.org/contributing.html
2. Ensure you have added or run the appropriate tests for your PR: https://spark.apache.org/developer-tools.html
3. If the PR is unfinished, add '[WIP]' in your PR title, e.g., '[WIP][SPARK-XXXX] Your PR title ...'.
4. Be sure to keep the PR description updated to reflect all changes.
5. Please write your PR title to summarize what this PR proposes.
6. If possible, provide a concise example to reproduce the issue for a faster review.
7. If you want to add a new configuration, please read the guideline first for naming configurations in
'core/src/main/scala/org/apache/spark/internal/config/ConfigEntry.scala'.
8. If you want to add or modify an error type or message, please read the guideline first in
'common/utils/src/main/resources/error/README.md'.
-->
### What changes were proposed in this pull request?
<!--
Please clarify what changes you are proposing. The purpose of this section is to outline the changes and how this PR fixes the issue.
If possible, please consider writing useful notes for better and faster reviews in your PR. See the examples below.
1. If you refactor some codes with changing classes, showing the class hierarchy will help reviewers.
2. If you fix some SQL features, you can provide some references of other DBMSes.
3. If there is design documentation, please add the link.
4. If there is a discussion in the mailing list, please add the link.
-->
### Why are the changes needed?
<!--
Please clarify why the changes are needed. For instance,
1. If you propose a new API, clarify the use case for a new API.
2. If you fix a bug, you can clarify why it is a bug.
-->
### Does this PR introduce _any_ user-facing change?
<!--
Note that it means *any* user-facing change including all aspects such as new features, bug fixes, or other behavior changes. Documentation-only updates are not considered user-facing changes.
If yes, please clarify the previous behavior and the change this PR proposes - provide the console output, description and/or an example to show the behavior difference if possible.
If possible, please also clarify if this is a user-facing change compared to the released Spark versions or within the unreleased branches such as master.
If no, write 'No'.
-->
### How was this patch tested?
<!--
If tests were added, say they were added here. Please make sure to add some test cases that check the changes thoroughly including negative and positive cases if possible.
If it was tested in a way different from regular unit tests, please clarify how you tested step by step, ideally copy and paste-able, so that other reviewers can test and check, and descendants can verify in the future.
If tests were not added, please describe why they were not added and/or why it was difficult to add.
If benchmark tests were added, please run the benchmarks in GitHub Actions for the consistent environment, and the instructions could accord to: https://spark.apache.org/developer-tools.html#github-workflow-benchmarks.
-->
### Was this patch authored or co-authored using generative AI tooling?
<!--
If generative AI tooling has been used in the process of authoring this patch, please include the
phrase: 'Generated-by: ' followed by the name of the tool and its version.
If no, write 'No'.
Please refer to the [ASF Generative Tooling Guidance](https://www.apache.org/legal/generative-tooling.html) for details.
-->
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57640] Implement gRPC WorkerSession and test it against an echo server. [spark]
Posted by "haiyangsun-db (via GitHub)" <gi...@apache.org>.
haiyangsun-db commented on code in PR #56702:
URL: https://github.com/apache/spark/pull/56702#discussion_r3666321486
##########
udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/GrpcWorkerSession.scala:
##########
@@ -0,0 +1,933 @@
+/*
+ * 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.grpc
+
+import java.util.concurrent.{CountDownLatch, LinkedBlockingQueue, TimeoutException, TimeUnit}
+import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference}
+
+import scala.util.control.NonFatal
+
+import io.grpc.{ConnectivityState, ManagedChannel}
+import io.grpc.stub.StreamObserver
+
+import org.apache.spark.annotation.Experimental
+import org.apache.spark.udf.worker.{Cancel, CancelResponse, DataRequest, DataResponse,
+ ExecutionError, Finish, FinishResponse, Init, InitResponse, UdfControlRequest,
+ UdfControlResponse, UdfRequest, UdfResponse, UdfWorkerGrpc}
+import org.apache.spark.udf.worker.core.{Termination, WorkerHandle, WorkerLogger, WorkerSession}
+import org.apache.spark.udf.worker.core.WorkerSession.SessionState
+import org.apache.spark.udf.worker.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_message.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.
+ *
+ * '''Driving model.''' Consumption-driven (Volcano / pull): the thread that
+ * consumes the [[doProcess]] result iterator is the one that pulls input and
+ * sends each `DataRequest`; the gRPC callback thread only receives output. It is
+ * pull-driven but not one-input-per-output -- `advance` sends the next input
+ * whenever the output queue is momentarily empty, so under async delivery it may
+ * push several input batches before any output is read (bounded by HTTP/2 flow
+ * control).
+ *
+ * '''State machine.''' This class does not keep its own state machine: it
+ * drives the single [[WorkerSession.SessionState]] owned by the base. The base
+ * advances `Created -> Initializing` (in `init`) and `Initialized -> Streaming`
+ * (in `process`); this class advances the protocol-event edges through
+ * [[compareAndSetState]] / [[completeTerminal]] as it exchanges messages:
+ * {{{
+ * Initializing --(InitResponse ok)--> Initialized [handleControl]
+ * Streaming ----(Finish written)----> Finishing [ProcessIterator]
+ * <any non-terminal> --(Cancel written)--> Cancelling [sendCancelInternal]
+ * <any non-terminal> --(terminator/error)--> terminal [completeTerminal]
+ * }}}
+ * The two clean terminals carry the worker's `FinishResponse` / `CancelResponse`
+ * (metrics + finish/cancel callback `data`/`error`) so [[close]] can return
+ * them. The only flag kept outside the machine is [[cancelRequested]]: a
+ * cancellation can be requested before the stream exists (so it cannot be a
+ * state transition yet), and it must both fast-fail the result iterator and
+ * suppress any in-flight Data/Finish.
+ *
+ * Threading:
+ * - [[doInit]] is synchronous: sends `Init` and blocks on `InitResponse`,
+ * returning it.
+ * - [[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.
+ * - [[doClose]] is thread-safe and idempotent: it settles + returns the
+ * terminator (cancelling in-flight work if the stream had not finished)
+ * and half-closes the request side.
+ *
+ * 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 a gRPC channel built and owned by the caller (the
+ * dispatcher). 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 from the worker, drained by the process() iterator.
+ // Intentionally unbounded: a bounded queue would block the gRPC callback
+ // (Netty event-loop) thread when full, stalling terminator/control delivery on
+ // the whole channel. HTTP/2 flow control bounds the wire and the consumer
+ // normally drains promptly; a stalled downstream can still grow it, but the fix
+ // is protocol-level back-pressure (out of scope), not bounding the queue.
+ //
+ // TODO [SPARK-57324]: expose queue depth as a metric (early warning for a
+ // stalled consumer).
+ private val outputQueue = new LinkedBlockingQueue[QueueItem]()
+
+ // The worker's `InitResponse` (success or error) coupled with the latch that
+ // init() blocks on. Fires when the InitResponse arrives (`complete`) or when a
+ // terminal settles first without one -- transport error, half-close, or a
+ // premature terminator (`signalWithoutValue`). Until it fires we have no proof
+ // the worker accepted the session.
+ //
+ // Settle-before-release rule (referenced from every callback that both settles
+ // a terminal and fires this latch): settle the terminal FIRST, then complete /
+ // signal. init() blocks on the latch, and only latch await/release establish a
+ // happens-before edge, so a woken init() is guaranteed to observe the terminal
+ // rather than a transient state. OneShotValue keeps that publish-then-release
+ // in one place instead of every caller remembering to count down after setting
+ // the reference.
+ private val initValue = new OneShotValue[InitResponse]
+
+ // Fired when the session reaches a terminal [[SessionState]]. doClose() and
+ // the init-error path block on this to drain the terminator.
+ private val terminalLatch = new CountDownLatch(1)
+
+ // Captures an ErrorResponse encountered during the data phase so that
+ // the CancelResponse terminator 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)
+
+ // Cancellation INTENT -- distinct from the `Cancelling` state, which is reached
+ // only once a Cancel is actually written to the wire ([[sendCancelInternal]]).
+ // Intent is set first and can outrun (or never reach) that write, so
+ // `cancelRequested` does NOT imply state `Cancelling`. Kept outside the
+ // [[SessionState]] machine because a cancel can be requested before the stream
+ // exists (pre-init), where there is no wire transition to make yet. Used to (a)
+ // make cancel idempotent across all call sites, (b) fast-fail ProcessIterator
+ // on a pre-init cancel, and (c) suppress any Data/Finish that would otherwise
+ // race a Cancel onto the wire (re-read inside [[requestLock]]).
+ private val cancelRequested = new AtomicBoolean(false)
+
+ // gRPC requires serialized writes to a request StreamObserver.
+ private val requestLock = new Object
+
+ // Initialised in init() -- before that, close() is a no-op 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 =>
+ // A DataResponse before InitResponse violates the protocol (InitResponse
+ // must precede any DataResponse): fast-fail rather than enqueue it and let
+ // init() block to its timeout. Settle-before-release (see initValue).
+ if (!initResolved) {
+ Transitions.transportFailed(new IllegalStateException(
+ "worker sent a DataResponse before InitResponse"))
+ initValue.signalWithoutValue()
+ } else {
+ outputQueue.put(QueueItem.Batch(response.getData))
+ }
+
+ case UdfResponse.ResponseCase.CONTROL =>
+ handleControl(response.getControl)
+
+ case other =>
+ // A malformed response (empty / unknown oneof) is a terminal transport
+ // failure; fast-fail init the same way. Settle-before-release (see initValue).
+ Transitions.transportFailed(new IllegalStateException(
+ s"unexpected response oneof: $other"))
+ initValue.signalWithoutValue()
+ }
+ }
+
+ override def onError(t: Throwable): Unit = {
+ // Transport-level failure: the stream is dead, no further writes possible.
+ // Settle-before-release (see initValue) so init() surfaces the transport
+ // cause instead of the initResponseTimeoutMs "timed out" error.
+ Transitions.transportFailed(t)
+ initValue.signalWithoutValue()
+ }
+
+ 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 (!currentState.isTerminal) {
+ Transitions.transportFailed(new IllegalStateException(
+ "worker response stream closed without a terminator"))
+ }
+ // Defensive: if onCompleted reached us before InitResponse, doInit is
+ // still blocked on initValue and would otherwise time out.
+ initValue.signalWithoutValue()
+ }
+ }
+
+ /**
+ * Wakes the result iterator (blocked on [[outputQueue]]) and any thread
+ * waiting on [[terminalLatch]] when the base settles a terminal. Invoked once,
+ * by the caller that wins [[completeTerminal]].
+ */
+ override protected def onTerminalSettled(termination: Termination): Unit = {
+ outputQueue.put(QueueItem.EndOfStream)
+ terminalLatch.countDown()
+ }
+
+ private def handleControl(ctrl: UdfControlResponse): Unit = ctrl.getControlCase match {
+ case UdfControlResponse.ControlCase.INIT =>
+ val resp = ctrl.getInit
+ if (resp.hasError) {
+ // Attribute a subsequent close()'s terminator to the init error rather
+ // than a bare Cancelled, mirroring the data-phase ERROR branch.
+ executionError.compareAndSet(None, Some(resp.getError))
+ // Settle Failed(err) before the Cancel below (not after its CancelResponse):
+ // the sticky terminal makes close() report the init error, and the Cancel
+ // becomes best-effort worker cleanup -- so if it can't be written (a
+ // directExecutor worker delivered this reentrantly, before doInit published
+ // requestObserver), doInit still doesn't wait terminalTimeoutMs for a
+ // CancelResponse that can never arrive. Settle-before-release (see initValue).
+ Transitions.failed(resp.getError)
+ initValue.complete(resp)
+ sendCancelInternal(() => cancelWithReason("init failed"))
+ } else {
+ // InitResponse OK. Only advance from Initializing so a terminal that raced
+ // in (e.g. a transport error) still wins; process() then opens the data
+ // phase. Settle-before-release (see initValue): publish after the CAS.
+ Transitions.initAccepted()
+ initValue.complete(resp)
+ }
+
+ case UdfControlResponse.ControlCase.ERROR =>
+ val err = ctrl.getError.getError
+ executionError.compareAndSet(None, Some(err))
+ // Before init resolves, settle Failed here rather than rely on the Cancel
+ // below: sendCancelInternal may find requestObserver still null (reentrant
+ // delivery inside stream.onNext, before doInit published it), so no Cancel --
+ // hence no CancelResponse -- would settle the terminal, and doInit would
+ // return as if init succeeded. completeTerminal is idempotent, so the
+ // data-phase ERROR -> Cancel -> CancelResponse path is unaffected.
+ if (!initResolved) {
+ Transitions.failed(err)
+ }
+ // Settle-before-release (see initValue). No value to publish -- ERROR is not
+ // an InitResponse; on the data-phase path the latch already fired in init.
+ initValue.signalWithoutValue()
+ sendCancelInternal(() => cancelWithReason("aborting after ErrorResponse"))
+
+ case UdfControlResponse.ControlCase.FINISH =>
+ // The FinishResponse carries metrics + the finish-callback data/error.
+ // Keep it on the terminal so close() can return it; the iterator inspects
+ // its error field to decide whether to throw.
+ Transitions.finished(ctrl.getFinish)
+ // Defensive: FINISH before InitResponse is a worker protocol bug, but
+ // we should fail init fast rather than hang the 30s init timeout.
+ initValue.signalWithoutValue()
+
+ case UdfControlResponse.ControlCase.CANCEL =>
+ // The CancelResponse carries metrics + the cancel-callback error. Keep it
+ // on the terminal so close() can return it; any prior ErrorResponse is
+ // tracked in executionError and surfaced by the iterator.
+ Transitions.cancelled(ctrl.getCancel)
+ // Defensive: CANCEL before InitResponse unblocks doInit so it can
+ // surface the cancellation instead of timing out.
+ initValue.signalWithoutValue()
+
+ case UdfControlResponse.ControlCase.CONTROL_NOT_SET =>
+ Transitions.transportFailed(new IllegalStateException(
+ "empty UdfControlResponse oneof"))
+ initValue.signalWithoutValue()
+ }
+
+ /**
+ * True once init is no longer pending -- i.e. the stream is past `Initializing`.
+ * Not "init succeeded": a terminal (including a failure) also counts as resolved.
+ */
+ private def initResolved: Boolean = currentState match {
+ case SessionState.Created | SessionState.Initializing => false
+ case _ => true
+ }
+
+ private def cancelWithReason(reason: String): Cancel =
+ Cancel.newBuilder().setReason(reason).build()
+
+ /**
+ * The protocol transition graph in one place: names for the edges, not a
+ * second source of truth. Every edge acts on the single
+ * [[WorkerSession.SessionState]] owned by the base via [[compareAndSetState]]
+ * (non-terminal) or [[completeTerminal]] (terminal), so the base's CAS is the
+ * only synchronization and a terminal that arrived first always wins (the
+ * non-terminal CASes fail against it; [[completeTerminal]] is first-wins).
+ *
+ * Edges this class drives -- edge, method, then driver site(s) / thread (the
+ * base drives the API-call edges: `Created -> Initializing` in `init`,
+ * `Initialized -> Streaming` in `process`):
+ * {{{
+ * Initializing -> Initialized initAccepted handleControl INIT-ok [gRPC cb]
+ * Streaming -> Finishing finishSent advance branch 3 [engine]
+ * non-terminal -> Cancelling cancelSentFrom sendCancelInternal, iff a Cancel
+ * reaches the wire [gRPC cb | engine | init | close]
+ * non-terminal -> Terminal finished/cancelled/failed/transportFailed
+ * handleControl / doInit / doClose / advance / onError
+ * }}}
+ * `Cancelling` is reached only if a Cancel is actually written; a pre-stream or
+ * raced cancel goes straight to a `Cancelled`/`TransportFailed` terminal (see
+ * [[sendCancelInternal]], [[doClose]], `ProcessIterator`) or nowhere -- so
+ * `cancelRequested` (intent) does not imply state `Cancelling`.
+ */
+ private object Transitions {
+ /** `InitResponse` OK: `Initializing -> Initialized`. */
+ def initAccepted(): Boolean =
+ compareAndSetState(SessionState.Initializing, SessionState.Initialized)
+
+ /** Input exhausted and `Finish` written: `Streaming -> Finishing` (once). */
+ def finishSent(): Boolean =
+ compareAndSetState(SessionState.Streaming, SessionState.Finishing)
+
+ /** `Cancel` written: `cur -> Cancelling`, from any non-terminal `cur`. */
+ def cancelSentFrom(cur: SessionState): Boolean =
+ !cur.isTerminal && compareAndSetState(cur, SessionState.Cancelling)
+
+ /** Clean terminal carrying the worker's `FinishResponse`. */
+ def finished(response: FinishResponse): Boolean =
+ completeTerminal(Termination.Finished(response))
+
+ /** Clean terminal carrying the worker's `CancelResponse`. */
+ def cancelled(response: CancelResponse): Boolean =
+ completeTerminal(Termination.Cancelled(response))
+
+ /** Failure terminal carrying a structured [[ExecutionError]]. */
+ def failed(error: ExecutionError): Boolean =
+ completeTerminal(Termination.Failed(error))
+
+ /** Failure terminal carrying a transport-level cause. */
+ def transportFailed(cause: Throwable): Boolean =
+ completeTerminal(Termination.TransportFailed(cause))
+ }
+
+ // ---- WorkerSession hooks ------------------------------------------------
+
+ override protected def doInit(message: Init): InitResponse = {
+ // 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")
+ Transitions.transportFailed(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. close()'s 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 {
+ // Reentrant delivery: a directExecutor worker (one whose gRPC callbacks run
+ // synchronously on the caller's thread) can deliver InitResponse -- and thus
+ // run responseObserver/handleControl -- from *inside* this stream.onNext,
+ // before requestObserver below is published. The base advanced
+ // Created -> Initializing before calling doInit, so that reentrant
+ // InitResponse still finds Initializing and advances to Initialized; the
+ // handleControl error paths guard the still-null requestObserver. This is
+ // the scenario the "before doInit published requestObserver" comments mean.
+ 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
+ Transitions.transportFailed(e)
+ // Surface as GrpcWorkerSessionException so the engine integration layer
+ // (which catches that type and wraps it) sees a uniform init-failure
+ // exception rather than the raw transport error.
+ throw new GrpcWorkerSessionException("UDF worker stream failed during init", e)
+ }
+
+ val responded = try {
+ initValue.await(initResponseTimeoutMs)
+ } catch {
+ case _: InterruptedException =>
+ Thread.currentThread().interrupt()
+ sendCancelInternal(() => cancelWithReason("interrupted during init"))
+ // Make sure close() does not block waiting for a terminator that
+ // will never arrive on this thread's behalf.
+ Transitions.transportFailed(
+ new InterruptedException("interrupted while waiting for InitResponse"))
+ throw new InterruptedException("interrupted while waiting for InitResponse")
Review Comment:
it's a first-class Interrupted outcome that cooperatively cancels (bounded), keeps the worker reusable, and no longer masquerades as TransportFailed
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57640][UDF] Implement gRPC WorkerSession and test it against an echo server [spark]
Posted by "Yicong-Huang (via GitHub)" <gi...@apache.org>.
Yicong-Huang closed pull request #56702: [SPARK-57640][UDF] Implement gRPC WorkerSession and test it against an echo server
URL: https://github.com/apache/spark/pull/56702
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57640] Implement gRPC WorkerSession and test it against an echo server. [spark]
Posted by "haiyangsun-db (via GitHub)" <gi...@apache.org>.
haiyangsun-db commented on code in PR #56702:
URL: https://github.com/apache/spark/pull/56702#discussion_r3665868834
##########
udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/GrpcWorkerSession.scala:
##########
@@ -0,0 +1,933 @@
+/*
+ * 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.grpc
+
+import java.util.concurrent.{CountDownLatch, LinkedBlockingQueue, TimeoutException, TimeUnit}
+import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference}
+
+import scala.util.control.NonFatal
+
+import io.grpc.{ConnectivityState, ManagedChannel}
+import io.grpc.stub.StreamObserver
+
+import org.apache.spark.annotation.Experimental
+import org.apache.spark.udf.worker.{Cancel, CancelResponse, DataRequest, DataResponse,
+ ExecutionError, Finish, FinishResponse, Init, InitResponse, UdfControlRequest,
+ UdfControlResponse, UdfRequest, UdfResponse, UdfWorkerGrpc}
+import org.apache.spark.udf.worker.core.{Termination, WorkerHandle, WorkerLogger, WorkerSession}
+import org.apache.spark.udf.worker.core.WorkerSession.SessionState
+import org.apache.spark.udf.worker.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_message.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.
+ *
+ * '''Driving model.''' Consumption-driven (Volcano / pull): the thread that
+ * consumes the [[doProcess]] result iterator is the one that pulls input and
+ * sends each `DataRequest`; the gRPC callback thread only receives output. It is
+ * pull-driven but not one-input-per-output -- `advance` sends the next input
+ * whenever the output queue is momentarily empty, so under async delivery it may
+ * push several input batches before any output is read (bounded by HTTP/2 flow
+ * control).
+ *
+ * '''State machine.''' This class does not keep its own state machine: it
+ * drives the single [[WorkerSession.SessionState]] owned by the base. The base
+ * advances `Created -> Initializing` (in `init`) and `Initialized -> Streaming`
+ * (in `process`); this class advances the protocol-event edges through
+ * [[compareAndSetState]] / [[completeTerminal]] as it exchanges messages:
+ * {{{
+ * Initializing --(InitResponse ok)--> Initialized [handleControl]
+ * Streaming ----(Finish written)----> Finishing [ProcessIterator]
+ * <any non-terminal> --(Cancel written)--> Cancelling [sendCancelInternal]
+ * <any non-terminal> --(terminator/error)--> terminal [completeTerminal]
+ * }}}
+ * The two clean terminals carry the worker's `FinishResponse` / `CancelResponse`
+ * (metrics + finish/cancel callback `data`/`error`) so [[close]] can return
+ * them. The only flag kept outside the machine is [[cancelRequested]]: a
+ * cancellation can be requested before the stream exists (so it cannot be a
+ * state transition yet), and it must both fast-fail the result iterator and
+ * suppress any in-flight Data/Finish.
+ *
+ * Threading:
+ * - [[doInit]] is synchronous: sends `Init` and blocks on `InitResponse`,
+ * returning it.
+ * - [[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.
+ * - [[doClose]] is thread-safe and idempotent: it settles + returns the
+ * terminator (cancelling in-flight work if the stream had not finished)
+ * and half-closes the request side.
+ *
+ * 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 a gRPC channel built and owned by the caller (the
+ * dispatcher). 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 from the worker, drained by the process() iterator.
+ // Intentionally unbounded: a bounded queue would block the gRPC callback
+ // (Netty event-loop) thread when full, stalling terminator/control delivery on
+ // the whole channel. HTTP/2 flow control bounds the wire and the consumer
+ // normally drains promptly; a stalled downstream can still grow it, but the fix
+ // is protocol-level back-pressure (out of scope), not bounding the queue.
+ //
+ // TODO [SPARK-57324]: expose queue depth as a metric (early warning for a
+ // stalled consumer).
+ private val outputQueue = new LinkedBlockingQueue[QueueItem]()
+
+ // The worker's `InitResponse` (success or error) coupled with the latch that
+ // init() blocks on. Fires when the InitResponse arrives (`complete`) or when a
+ // terminal settles first without one -- transport error, half-close, or a
+ // premature terminator (`signalWithoutValue`). Until it fires we have no proof
+ // the worker accepted the session.
+ //
+ // Settle-before-release rule (referenced from every callback that both settles
+ // a terminal and fires this latch): settle the terminal FIRST, then complete /
+ // signal. init() blocks on the latch, and only latch await/release establish a
+ // happens-before edge, so a woken init() is guaranteed to observe the terminal
+ // rather than a transient state. OneShotValue keeps that publish-then-release
+ // in one place instead of every caller remembering to count down after setting
+ // the reference.
+ private val initValue = new OneShotValue[InitResponse]
+
+ // Fired when the session reaches a terminal [[SessionState]]. doClose() and
+ // the init-error path block on this to drain the terminator.
+ private val terminalLatch = new CountDownLatch(1)
+
+ // Captures an ErrorResponse encountered during the data phase so that
+ // the CancelResponse terminator 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)
+
+ // Cancellation INTENT -- distinct from the `Cancelling` state, which is reached
+ // only once a Cancel is actually written to the wire ([[sendCancelInternal]]).
+ // Intent is set first and can outrun (or never reach) that write, so
+ // `cancelRequested` does NOT imply state `Cancelling`. Kept outside the
+ // [[SessionState]] machine because a cancel can be requested before the stream
+ // exists (pre-init), where there is no wire transition to make yet. Used to (a)
+ // make cancel idempotent across all call sites, (b) fast-fail ProcessIterator
+ // on a pre-init cancel, and (c) suppress any Data/Finish that would otherwise
+ // race a Cancel onto the wire (re-read inside [[requestLock]]).
+ private val cancelRequested = new AtomicBoolean(false)
+
+ // gRPC requires serialized writes to a request StreamObserver.
+ private val requestLock = new Object
+
+ // Initialised in init() -- before that, close() is a no-op 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 =>
+ // A DataResponse before InitResponse violates the protocol (InitResponse
+ // must precede any DataResponse): fast-fail rather than enqueue it and let
+ // init() block to its timeout. Settle-before-release (see initValue).
+ if (!initResolved) {
+ Transitions.transportFailed(new IllegalStateException(
+ "worker sent a DataResponse before InitResponse"))
+ initValue.signalWithoutValue()
+ } else {
+ outputQueue.put(QueueItem.Batch(response.getData))
+ }
+
+ case UdfResponse.ResponseCase.CONTROL =>
+ handleControl(response.getControl)
+
+ case other =>
+ // A malformed response (empty / unknown oneof) is a terminal transport
+ // failure; fast-fail init the same way. Settle-before-release (see initValue).
+ Transitions.transportFailed(new IllegalStateException(
+ s"unexpected response oneof: $other"))
+ initValue.signalWithoutValue()
+ }
+ }
+
+ override def onError(t: Throwable): Unit = {
+ // Transport-level failure: the stream is dead, no further writes possible.
+ // Settle-before-release (see initValue) so init() surfaces the transport
+ // cause instead of the initResponseTimeoutMs "timed out" error.
+ Transitions.transportFailed(t)
+ initValue.signalWithoutValue()
+ }
+
+ 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 (!currentState.isTerminal) {
+ Transitions.transportFailed(new IllegalStateException(
+ "worker response stream closed without a terminator"))
+ }
+ // Defensive: if onCompleted reached us before InitResponse, doInit is
+ // still blocked on initValue and would otherwise time out.
+ initValue.signalWithoutValue()
+ }
+ }
+
+ /**
+ * Wakes the result iterator (blocked on [[outputQueue]]) and any thread
+ * waiting on [[terminalLatch]] when the base settles a terminal. Invoked once,
+ * by the caller that wins [[completeTerminal]].
+ */
+ override protected def onTerminalSettled(termination: Termination): Unit = {
+ outputQueue.put(QueueItem.EndOfStream)
+ terminalLatch.countDown()
+ }
+
+ private def handleControl(ctrl: UdfControlResponse): Unit = ctrl.getControlCase match {
+ case UdfControlResponse.ControlCase.INIT =>
+ val resp = ctrl.getInit
+ if (resp.hasError) {
+ // Attribute a subsequent close()'s terminator to the init error rather
+ // than a bare Cancelled, mirroring the data-phase ERROR branch.
+ executionError.compareAndSet(None, Some(resp.getError))
+ // Settle Failed(err) before the Cancel below (not after its CancelResponse):
+ // the sticky terminal makes close() report the init error, and the Cancel
+ // becomes best-effort worker cleanup -- so if it can't be written (a
+ // directExecutor worker delivered this reentrantly, before doInit published
+ // requestObserver), doInit still doesn't wait terminalTimeoutMs for a
+ // CancelResponse that can never arrive. Settle-before-release (see initValue).
+ Transitions.failed(resp.getError)
+ initValue.complete(resp)
+ sendCancelInternal(() => cancelWithReason("init failed"))
+ } else {
+ // InitResponse OK. Only advance from Initializing so a terminal that raced
+ // in (e.g. a transport error) still wins; process() then opens the data
+ // phase. Settle-before-release (see initValue): publish after the CAS.
+ Transitions.initAccepted()
+ initValue.complete(resp)
+ }
+
+ case UdfControlResponse.ControlCase.ERROR =>
+ val err = ctrl.getError.getError
+ executionError.compareAndSet(None, Some(err))
+ // Before init resolves, settle Failed here rather than rely on the Cancel
+ // below: sendCancelInternal may find requestObserver still null (reentrant
+ // delivery inside stream.onNext, before doInit published it), so no Cancel --
+ // hence no CancelResponse -- would settle the terminal, and doInit would
+ // return as if init succeeded. completeTerminal is idempotent, so the
+ // data-phase ERROR -> Cancel -> CancelResponse path is unaffected.
+ if (!initResolved) {
+ Transitions.failed(err)
+ }
+ // Settle-before-release (see initValue). No value to publish -- ERROR is not
+ // an InitResponse; on the data-phase path the latch already fired in init.
+ initValue.signalWithoutValue()
+ sendCancelInternal(() => cancelWithReason("aborting after ErrorResponse"))
+
+ case UdfControlResponse.ControlCase.FINISH =>
+ // The FinishResponse carries metrics + the finish-callback data/error.
+ // Keep it on the terminal so close() can return it; the iterator inspects
+ // its error field to decide whether to throw.
+ Transitions.finished(ctrl.getFinish)
+ // Defensive: FINISH before InitResponse is a worker protocol bug, but
+ // we should fail init fast rather than hang the 30s init timeout.
+ initValue.signalWithoutValue()
+
+ case UdfControlResponse.ControlCase.CANCEL =>
+ // The CancelResponse carries metrics + the cancel-callback error. Keep it
+ // on the terminal so close() can return it; any prior ErrorResponse is
+ // tracked in executionError and surfaced by the iterator.
+ Transitions.cancelled(ctrl.getCancel)
+ // Defensive: CANCEL before InitResponse unblocks doInit so it can
+ // surface the cancellation instead of timing out.
+ initValue.signalWithoutValue()
+
+ case UdfControlResponse.ControlCase.CONTROL_NOT_SET =>
+ Transitions.transportFailed(new IllegalStateException(
+ "empty UdfControlResponse oneof"))
+ initValue.signalWithoutValue()
+ }
+
+ /**
+ * True once init is no longer pending -- i.e. the stream is past `Initializing`.
+ * Not "init succeeded": a terminal (including a failure) also counts as resolved.
+ */
+ private def initResolved: Boolean = currentState match {
+ case SessionState.Created | SessionState.Initializing => false
+ case _ => true
+ }
+
+ private def cancelWithReason(reason: String): Cancel =
+ Cancel.newBuilder().setReason(reason).build()
+
+ /**
+ * The protocol transition graph in one place: names for the edges, not a
+ * second source of truth. Every edge acts on the single
+ * [[WorkerSession.SessionState]] owned by the base via [[compareAndSetState]]
+ * (non-terminal) or [[completeTerminal]] (terminal), so the base's CAS is the
+ * only synchronization and a terminal that arrived first always wins (the
+ * non-terminal CASes fail against it; [[completeTerminal]] is first-wins).
+ *
+ * Edges this class drives -- edge, method, then driver site(s) / thread (the
+ * base drives the API-call edges: `Created -> Initializing` in `init`,
+ * `Initialized -> Streaming` in `process`):
+ * {{{
+ * Initializing -> Initialized initAccepted handleControl INIT-ok [gRPC cb]
+ * Streaming -> Finishing finishSent advance branch 3 [engine]
+ * non-terminal -> Cancelling cancelSentFrom sendCancelInternal, iff a Cancel
+ * reaches the wire [gRPC cb | engine | init | close]
+ * non-terminal -> Terminal finished/cancelled/failed/transportFailed
+ * handleControl / doInit / doClose / advance / onError
+ * }}}
+ * `Cancelling` is reached only if a Cancel is actually written; a pre-stream or
+ * raced cancel goes straight to a `Cancelled`/`TransportFailed` terminal (see
+ * [[sendCancelInternal]], [[doClose]], `ProcessIterator`) or nowhere -- so
+ * `cancelRequested` (intent) does not imply state `Cancelling`.
+ */
+ private object Transitions {
+ /** `InitResponse` OK: `Initializing -> Initialized`. */
+ def initAccepted(): Boolean =
+ compareAndSetState(SessionState.Initializing, SessionState.Initialized)
+
+ /** Input exhausted and `Finish` written: `Streaming -> Finishing` (once). */
+ def finishSent(): Boolean =
+ compareAndSetState(SessionState.Streaming, SessionState.Finishing)
+
+ /** `Cancel` written: `cur -> Cancelling`, from any non-terminal `cur`. */
+ def cancelSentFrom(cur: SessionState): Boolean =
+ !cur.isTerminal && compareAndSetState(cur, SessionState.Cancelling)
+
+ /** Clean terminal carrying the worker's `FinishResponse`. */
+ def finished(response: FinishResponse): Boolean =
+ completeTerminal(Termination.Finished(response))
+
+ /** Clean terminal carrying the worker's `CancelResponse`. */
+ def cancelled(response: CancelResponse): Boolean =
+ completeTerminal(Termination.Cancelled(response))
+
+ /** Failure terminal carrying a structured [[ExecutionError]]. */
+ def failed(error: ExecutionError): Boolean =
+ completeTerminal(Termination.Failed(error))
+
+ /** Failure terminal carrying a transport-level cause. */
+ def transportFailed(cause: Throwable): Boolean =
+ completeTerminal(Termination.TransportFailed(cause))
+ }
+
+ // ---- WorkerSession hooks ------------------------------------------------
+
+ override protected def doInit(message: Init): InitResponse = {
+ // 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")
+ Transitions.transportFailed(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. close()'s 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 {
+ // Reentrant delivery: a directExecutor worker (one whose gRPC callbacks run
+ // synchronously on the caller's thread) can deliver InitResponse -- and thus
+ // run responseObserver/handleControl -- from *inside* this stream.onNext,
+ // before requestObserver below is published. The base advanced
+ // Created -> Initializing before calling doInit, so that reentrant
+ // InitResponse still finds Initializing and advances to Initialized; the
+ // handleControl error paths guard the still-null requestObserver. This is
+ // the scenario the "before doInit published requestObserver" comments mean.
+ 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
+ Transitions.transportFailed(e)
+ // Surface as GrpcWorkerSessionException so the engine integration layer
+ // (which catches that type and wraps it) sees a uniform init-failure
+ // exception rather than the raw transport error.
+ throw new GrpcWorkerSessionException("UDF worker stream failed during init", e)
+ }
+
+ val responded = try {
+ initValue.await(initResponseTimeoutMs)
+ } catch {
+ case _: InterruptedException =>
+ Thread.currentThread().interrupt()
+ sendCancelInternal(() => cancelWithReason("interrupted during init"))
+ // Make sure close() does not block waiting for a terminator that
+ // will never arrive on this thread's behalf.
+ Transitions.transportFailed(
+ new InterruptedException("interrupted while waiting for InitResponse"))
+ throw new InterruptedException("interrupted while waiting for InitResponse")
+ }
+
+ if (!responded) {
+ sendCancelInternal(() => cancelWithReason("InitResponse timed out"))
+ val timeout = new TimeoutException(
+ s"timed out waiting for InitResponse after ${initResponseTimeoutMs}ms")
+ // Settle the terminal so a subsequent close() does not stall for a
+ // second `terminalTimeoutMs` waiting for a worker that already missed
+ // its init deadline.
+ Transitions.transportFailed(timeout)
+ // Surface as GrpcWorkerSessionException (carrying the timeout cause) so
+ // the engine integration layer that catches that type can wrap it.
+ throw new GrpcWorkerSessionException(
+ s"timed out waiting for InitResponse after ${initResponseTimeoutMs}ms", timeout)
+ }
+
+ initValue.get match {
+ case Some(resp) if resp.hasError =>
+ // 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(resp.getError)}", resp.getError)
+ case Some(resp) =>
+ resp
+ case None =>
+ // No InitResponse arrived but the latch fired. The defensive
+ // initValue.signalWithoutValue() 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.
+ currentState match {
+ case SessionState.Terminal(Termination.TransportFailed(cause)) =>
+ throw new GrpcWorkerSessionException(
+ "UDF worker stream failed during init", cause)
+ case SessionState.Terminal(Termination.Failed(err)) =>
+ throw new GrpcWorkerSessionException(
+ s"UDF worker reported an error before init completed: " +
+ describeError(err), err)
+ case SessionState.Terminal(Termination.Cancelled(_)) =>
+ throw new GrpcWorkerSessionException(
+ "UDF worker stream was cancelled before init completed")
+ case SessionState.Terminal(Termination.Finished(_)) =>
+ throw new GrpcWorkerSessionException(
+ "UDF worker finished before init completed")
+ case other =>
+ throw new IllegalStateException(
+ s"init latch fired without an InitResponse or terminal: $other")
+ }
+ }
+ }
+
+ override protected def doProcess(
+ input: Iterator[DataRequest],
+ finish: () => Finish): Iterator[DataResponse] = {
+ // Init success is guaranteed by the base [[WorkerSession]] lifecycle: if
+ // doInit had failed it would have thrown and process() would never run.
+ new ProcessIterator(input, finish)
+ }
+
+ override protected def doClose(cancel: () => Cancel): Termination = {
+ if (requestObserver == null) {
+ // init() never put a stream on the wire (closed before/around init, or
+ // init threw before publishing). There is no protocol terminator; treat
+ // the session as cancelled-before-start. The base WorkerSession still
+ // releases the worker handle, so the worker is torn down.
+ Transitions.cancelled(CancelResponse.getDefaultInstance)
+ return Termination.Cancelled(CancelResponse.getDefaultInstance)
Review Comment:
There's no Transitions.failed anymore - transportFailed and the terminal outcomes are now distinct by design
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57640] Implement gRPC WorkerSession and test it against an echo server. [spark]
Posted by "haiyangsun-db (via GitHub)" <gi...@apache.org>.
haiyangsun-db commented on code in PR #56702:
URL: https://github.com/apache/spark/pull/56702#discussion_r3707824036
##########
udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/GrpcWorkerSession.scala:
##########
@@ -0,0 +1,1177 @@
+/*
+ * 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.grpc
+
+import java.util.Objects
+import java.util.concurrent.{CountDownLatch, LinkedBlockingQueue, TimeoutException, TimeUnit}
+import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference}
+
+import scala.util.control.NonFatal
+
+import io.grpc.{ConnectivityState, ManagedChannel}
+import io.grpc.stub.StreamObserver
+
+import org.apache.spark.annotation.Experimental
+import org.apache.spark.udf.worker.{Cancel, CancelResponse, DataRequest, DataResponse,
+ ExecutionError, Finish, FinishResponse, Init, InitResponse, UdfControlRequest,
+ UdfControlResponse, UdfPayload, UdfRequest, UdfResponse, UdfWorkerGrpc}
+import org.apache.spark.udf.worker.core.{Termination, WorkerHandle, WorkerLogger, WorkerSession}
+import org.apache.spark.udf.worker.core.WorkerSession.SessionState
+import org.apache.spark.udf.worker.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_message.proto` (`PayloadChunk*`
+ * between Init and InitResponse omitted here; chunking is not yet
+ * implemented -- see the TODO below):
+ * {{{
+ * Engine -> Worker: Init -> (DataRequest)* -> Finish (Cancel)?
+ * | Cancel
+ * Worker -> Engine: InitResponse -> (DataResponse)* ->
+ * (ErrorResponse)? -> (FinishResponse | CancelResponse)
+ * }}}
+ *
+ * Knows nothing about how the worker was provisioned -- the dispatcher
+ * constructs this with a [[WorkerHandle]] and channel; the base
+ * [[WorkerSession]] handles dispatcher-side cleanup on close.
+ *
+ * '''Driving model.''' Consumption-driven (Volcano / pull): the thread that
+ * consumes the [[doProcess]] result iterator is the one that pulls input and
+ * sends each `DataRequest`; the gRPC callback thread only receives output. It is
+ * pull-driven but not one-input-per-output -- `advance` sends the next input
+ * whenever the output queue is momentarily empty, so under async delivery it may
+ * push several input batches before any output is read. HTTP/2 flow control
+ * bounds wire traffic, but application-level buffering is intentionally left to
+ * a follow-up (see the TODO on [[outputQueue]]).
+ *
+ * '''State machine.''' This class does not keep its own state machine: it
+ * drives the single [[WorkerSession.SessionState]] owned by the base. The base
+ * advances `Created -> Initializing` (in `init`) and `Initialized -> Streaming`
+ * (in `process`); this class advances the protocol-event edges through
+ * [[compareAndSetState]] / [[completeTerminal]] as it exchanges messages:
+ * {{{
+ * Initializing --(InitResponse ok)--> Initialized [handleControl]
+ * Streaming ----(input exhausted)-------> Finishing [ProcessIterator]
+ * <any non-terminal> --(Cancel send starts)--> Cancelling [sendCancelInternal]
+ * <any non-terminal> --(terminator/error)--> terminal [completeTerminal]
+ * }}}
+ * The two clean terminals carry the worker's `FinishResponse` / `CancelResponse`
+ * (metrics + finish/cancel callback `data`/`error`) so [[close]] can return
+ * them. Cancellation intent is deliberately tracked outside the machine in
+ * [[cancelRequested]], which makes the wire write idempotent and suppresses any
+ * in-flight Data/Finish once cancellation begins.
+ *
+ * Threading:
+ * - [[doInit]] is synchronous: sends `Init` and blocks on `InitResponse`,
+ * returning it.
+ * - [[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.
+ * - [[doClose]] is thread-safe and idempotent: it settles + returns the
+ * terminator (cancelling in-flight work if the stream had not finished)
+ * and terminates the request side.
+ *
+ * TODO [SPARK-55278]: this class does not yet implement payload chunking;
+ * the entire [[UdfPayload]] 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 a gRPC channel built and owned by the caller (the
+ * dispatcher). 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`.
+ * @param interruptCancelTimeoutMs upper bound on the wait for a `CancelResponse`
+ * after an interrupt (cancelled query / killed
+ * task) sends `Cancel`. Short by design; on expiry
+ * the session settles `Interrupted`. See
+ * `handleInterrupt`.
+ */
+@Experimental
+class GrpcWorkerSession(
+ workerHandle: WorkerHandle,
+ channel: ManagedChannel,
+ logger: WorkerLogger = WorkerLogger.NoOp,
+ initResponseTimeoutMs: Long = DEFAULT_INIT_RESPONSE_TIMEOUT_MS,
+ terminalTimeoutMs: Long = DEFAULT_TERMINAL_TIMEOUT_MS,
+ interruptCancelTimeoutMs: Long = DEFAULT_INTERRUPT_CANCEL_TIMEOUT_MS)
+ extends WorkerSession(workerHandle, logger) {
+
+ require(channel != null, "channel is required")
+
+ private val asyncStub = UdfWorkerGrpc.newStub(channel)
+
+ // Output batches from the worker, drained by the process() iterator.
+ // Intentionally unbounded in this first implementation: a bounded queue would
+ // block the gRPC callback (Netty event-loop) thread when full, stalling
+ // terminator/control delivery on the whole channel. The consumer normally
+ // drains promptly, but HTTP/2 flow control alone does not bound this application
+ // queue or gRPC's asynchronous send buffers.
+ //
+ // TODO [SPARK-55278]: add application-level gRPC flow control in a follow-up,
+ // using ClientCallStreamObserver readiness for requests and manual inbound
+ // demand for responses, before wiring this transport into a production path.
+ // TODO [SPARK-57324]: expose queue depth as a metric (early warning for a
+ // stalled consumer).
+ private val outputQueue = new LinkedBlockingQueue[QueueItem]()
+
+ // This value couples the worker's `InitResponse` (success or error) with the
+ // latch on which init() blocks. The latch fires when the InitResponse arrives
+ // (`complete`), when a pre-init ErrorResponse arrives, or when a terminal
+ // settles first without an InitResponse (`signalWithoutValue`). Until it fires
+ // we have no proof the worker accepted the session.
+ //
+ // Settle-before-release rule (referenced from every callback that both settles
+ // a terminal and fires this latch): settle the terminal FIRST, then complete /
+ // signal. init() blocks on the latch, and only a latch await/release pair
+ // establishes a happens-before edge, so a woken init() is guaranteed to
+ // observe the terminal rather than a transient state. OneShotValue keeps that
+ // publish-then-release in one place instead of every caller remembering to
+ // count down after setting the reference.
+ private val initValue = new OneShotValue[InitResponse]
+
+ // Fired when the session reaches a terminal [[SessionState]]. doClose() and
+ // the init-error path block on this to drain the terminator.
+ private val terminalLatch = new CountDownLatch(1)
+
+ // Captures an ErrorResponse encountered during the data phase so that
+ // the CancelResponse terminator 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 immediately before Finish is handed to the request observer. It is set
+ // before onNext so a reentrant callback may deliver FinishResponse while the
+ // Finish request is still being written.
+ private val finishSendStarted = new AtomicBoolean(false)
+
+ // Cancellation INTENT -- distinct from the `Cancelling` state, which is reached
+ // once a Cancel send is attempted under [[requestLock]] ([[sendCancelInternal]]).
+ // Intent is set first and can outrun (or never reach) that attempt, so
+ // `cancelRequested` does NOT imply state `Cancelling`. It is kept outside the
+ // [[SessionState]] machine to (a) make cancellation idempotent across all call
+ // sites and (b) suppress any Data/Finish that would otherwise race a Cancel
+ // onto the wire (re-read inside [[requestLock]]).
+ private val cancelRequested = new AtomicBoolean(false)
+
+ // gRPC requires serialized writes to a request StreamObserver.
+ private val requestLock = new Object
+
+ // Initialized in init() -- before that, close() is a no-op on the request
+ // side, which is exactly the contract the wrapping WorkerSession expects.
+ @volatile private var requestObserver: StreamObserver[UdfRequest] = _
+
+ // True after the request side has been half-closed, aborted, or observed the
+ // transport closing. It prevents this session from issuing duplicate terminal
+ // calls, including following onError with onCompleted from close().
+ private val requestSideTerminated = new AtomicBoolean(false)
+
+ private val responseObserver: StreamObserver[UdfResponse] = new StreamObserver[UdfResponse] {
+ override def onNext(response: UdfResponse): Unit = {
+ response.getResponseCase match {
+ case UdfResponse.ResponseCase.DATA =>
+ // A DataResponse before InitResponse violates the protocol (InitResponse
+ // must precede any DataResponse): fast-fail rather than enqueue it and let
+ // init() block to its timeout. Settling the terminal wakes a blocked
+ // init() via onTerminalSettled.
+ if (!initResolved) {
+ Transitions.transportFailed(new IllegalStateException(
+ "worker sent a DataResponse before InitResponse"))
+ } else {
+ outputQueue.put(QueueItem.Batch(response.getData))
+ }
+
+ case UdfResponse.ResponseCase.CONTROL =>
+ handleControl(response.getControl)
+
+ case other =>
+ // A malformed response (empty / unknown oneof) is a terminal transport
+ // failure; fast-fail init the same way.
+ Transitions.transportFailed(new IllegalStateException(
+ s"unexpected response oneof: $other"))
+ }
+ }
+
+ override def onError(t: Throwable): Unit = {
+ // Transport-level failure: the stream is dead, no further writes possible.
+ // Settling the terminal wakes a blocked init() (via onTerminalSettled) so it
+ // surfaces the transport cause instead of the initResponseTimeoutMs error.
+ requestSideTerminated.set(true)
+ Transitions.transportFailed(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. Settling the terminal wakes a blocked
+ // init() via onTerminalSettled.
+ requestSideTerminated.set(true)
+ if (!currentState.isTerminal) {
+ Transitions.transportFailed(new IllegalStateException(
+ "worker response stream closed without a terminator"))
+ }
+ }
+ }
+
+ /**
+ * Wakes everything that can be blocked when the base settles a terminal:
+ * the result iterator (on [[outputQueue]]), a thread on [[terminalLatch]]
+ * (close), and a thread still blocked in [[doInit]] on [[initValue]]. Invoked
+ * once, by the caller that wins [[completeTerminal]], so this is the '''single'''
+ * place a settled terminal wakes a blocked [[doInit]] -- callers that settle a
+ * terminal (any response-callback failure/terminator branch, the timeout paths,
+ * or close()) do NOT signal [[initValue]] themselves; settling is enough.
+ * Settle-before-release (see [[initValue]]) holds because this runs after the
+ * terminal CAS in [[completeTerminal]]. The only direct [[initValue]] signals
+ * left are the non-terminal init paths ([[handleControl]]'s INIT-ok / INIT-error
+ * / pre-init-ERROR branches), which must wake [[doInit]] '''without''' settling a
+ * terminal. INIT-ok first advances the session to `Initialized`; the error
+ * branches remain `Initializing` only until [[doInit]] sends Cancel and throws.
+ */
+ override protected def onTerminalSettled(termination: Termination): Unit = {
+ outputQueue.put(QueueItem.EndOfStream)
+ initValue.signalWithoutValue()
+ terminalLatch.countDown()
+ }
+
+ private def handleControl(ctrl: UdfControlResponse): Unit = ctrl.getControlCase match {
+ case UdfControlResponse.ControlCase.INIT =>
+ val resp = ctrl.getInit
+ if (resp.hasError) {
+ // Record the error and publish the InitResponse; do NOT settle a terminal
+ // or send Cancel here. The proto requires the engine to Cancel after an
+ // init error (udf_message.proto); doInit owns that synchronous cleanup
+ // after the initial onNext returns, then drains the CancelResponse before
+ // throwing. Settle-before-release (see initValue): there is no terminal to
+ // settle first here, just publish.
+ executionError.compareAndSet(None, Some(resp.getError))
+ initValue.complete(resp)
+ } else {
+ // InitResponse OK. Only advance from Initializing so a terminal that raced
+ // in (e.g. a transport error) still wins; process() then opens the data
+ // phase. Settle-before-release (see initValue): publish after the CAS.
+ Transitions.initAccepted()
+ initValue.complete(resp)
+ }
+
+ case UdfControlResponse.ControlCase.ERROR =>
+ val err = ctrl.getError.getError
+ executionError.compareAndSet(None, Some(err))
+ if (!initResolved) {
+ // Pre-init ErrorResponse. Leave the session in Initializing and just wake
+ // init(): the recorded executionError tells doInit the worker failed, and
+ // doInit sends the Cancel and drains the CancelResponse after the initial
+ // onNext returns, just like the INIT-error branch.
+ initValue.signalWithoutValue()
+ } else {
+ // Data-phase ErrorResponse: requestObserver is published, so cancel here.
+ // Cancel -> CancelResponse settles the terminal; the iterator surfaces the
+ // recorded executionError. The init latch already fired in init(), so no
+ // signalWithoutValue is needed.
+ sendCancelInternal(() => cancelWithReason("aborting after ErrorResponse"))
+ }
+
+ case UdfControlResponse.ControlCase.FINISH =>
+ // The FinishResponse carries metrics + the finish-callback data/error.
+ // Keep it on the terminal so close() can return it; the iterator inspects
+ // its error field to decide whether to throw. A FinishResponse is valid only
+ // after the engine has started sending Finish; accepting one earlier could
+ // silently truncate input and report a clean result.
+ if (finishSendStarted.get()) {
+ Transitions.finished(ctrl.getFinish)
+ } else {
+ Transitions.transportFailed(new IllegalStateException(
+ "worker sent FinishResponse before the engine sent Finish"))
+ }
+
+ case UdfControlResponse.ControlCase.CANCEL =>
+ // The CancelResponse carries metrics + the cancel-callback error. Keep it
+ // on the terminal so close() can return it; any prior ErrorResponse is
+ // tracked in executionError and surfaced by the iterator. Settling the
+ // terminal wakes a blocked init() (a CANCEL before InitResponse) via
+ // onTerminalSettled.
+ Transitions.cancelled(ctrl.getCancel)
+
+ case UdfControlResponse.ControlCase.CONTROL_NOT_SET =>
+ // Settling the terminal wakes a blocked init() via onTerminalSettled.
+ Transitions.transportFailed(new IllegalStateException(
+ "empty UdfControlResponse oneof"))
+ }
+
+ /**
+ * True once init is no longer pending -- i.e. the stream is past `Initializing`.
+ * Not "init succeeded": a terminal (including a failure) also counts as resolved.
+ */
+ private def initResolved: Boolean = currentState match {
+ case SessionState.Created | SessionState.Initializing => false
+ case _ => true
+ }
+
+ private def cancelWithReason(reason: String): Cancel =
+ Cancel.newBuilder().setReason(reason).build()
+
+ /**
+ * The protocol transition graph in one place: names for the edges, not a
+ * second source of truth. Every edge acts on the single
+ * [[WorkerSession.SessionState]] owned by the base via [[compareAndSetState]]
+ * (non-terminal) or [[completeTerminal]] (terminal), so the base's CAS is the
+ * only synchronization and a terminal that arrived first always wins (the
+ * non-terminal CASes fail against it; [[completeTerminal]] is first-wins).
+ *
+ * Edges this class drives -- edge, method, then driver site(s) / thread (the
+ * base drives the API-call edges: `Created -> Initializing` in `init`,
+ * `Initialized -> Streaming` in `process`):
+ * {{{
+ * Initializing -> Initialized initAccepted handleControl INIT-ok [gRPC cb]
+ * Streaming -> Finishing beginFinish advance branch 3 [engine]
+ * non-terminal -> Cancelling beginCancelFrom sendCancelInternal, immediately
+ * before the Cancel send [gRPC cb | engine | init | close]
+ * non-terminal -> Terminal finished/cancelled/transportFailed/interrupted
+ * handleControl / doInit / doClose / advance / onError
+ * }}}
+ * `Cancelling` is reached only when a Cancel send starts; a pre-stream or
+ * raced cancel goes straight to a `Cancelled`/`TransportFailed` terminal (see
+ * [[sendCancelInternal]], [[doClose]], `ProcessIterator`) or nowhere -- so
+ * `cancelRequested` (intent) does not imply state `Cancelling`.
+ */
+ private object Transitions {
+ /** `InitResponse` OK: `Initializing -> Initialized`. */
+ def initAccepted(): Boolean =
+ compareAndSetState(SessionState.Initializing, SessionState.Initialized)
+
+ /** Input exhausted: `Streaming -> Finishing` (once). */
+ def beginFinish(): Boolean =
+ compareAndSetState(SessionState.Streaming, SessionState.Finishing)
+
+ /** `Cancel` send starts: `cur -> Cancelling`, from any non-terminal `cur`. */
+ def beginCancelFrom(cur: SessionState): Boolean =
+ !cur.isTerminal && compareAndSetState(cur, SessionState.Cancelling)
+
+ /** Clean terminal carrying the worker's `FinishResponse`. */
+ def finished(response: FinishResponse): Boolean =
+ completeTerminal(Termination.Finished(response))
+
+ /** Clean terminal carrying the worker's `CancelResponse`. */
+ def cancelled(response: CancelResponse): Boolean =
+ completeTerminal(Termination.Cancelled(response))
+
+ /** Failure terminal carrying a transport-level cause. */
+ def transportFailed(cause: Throwable): Boolean =
+ completeTerminal(Termination.TransportFailed(cause))
+
+ /**
+ * Terminal for an engine-thread interrupt (e.g. a Spark task kill). Distinct
+ * from [[transportFailed]] because the cause is an engine-side interrupt, but
+ * still unsalvageable without a worker acknowledgement. Settled by
+ * [[handleInterrupt]] only when the brief post-Cancel drain expires without a
+ * `CancelResponse`; if the worker acks in time the session settles a
+ * cooperative [[cancelled]] instead.
+ */
+ def interrupted(cause: Throwable): Boolean =
+ completeTerminal(Termination.Interrupted(cause))
+ }
+
+ // ---- WorkerSession hooks ------------------------------------------------
+
+ override protected def doInit(message: Init): InitResponse = {
+ // Construct the request before opening the RPC. A malformed Init must not
+ // leave an otherwise-unused Execute stream waiting for its first request.
+ val initRequest = UdfRequest.newBuilder()
+ .setControl(UdfControlRequest.newBuilder().setInit(message).build())
+ .build()
+
+ // 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")
+ Transitions.transportFailed(ex)
+ throw new GrpcWorkerSessionException("UDF worker channel is closed", ex)
+ }
+ try {
+ requestLock.synchronized {
+ // Serialize stream creation, publication, and Init with close/cancel. A
+ // close that wins the lock first prevents the RPC from opening; one that
+ // loses cannot release the worker until the initial write returns.
+ if (!currentState.isTerminal) {
+ val stream = asyncStub.execute(responseObserver)
+ requestObserver = stream
+ // Stream creation may invoke a response callback synchronously. Do not
+ // send Init if that callback has already settled a terminal.
+ if (!currentState.isTerminal) {
+ // With directExecutor, gRPC callbacks run synchronously on the caller's
+ // thread, so InitResponse can reach responseObserver/handleControl from
+ // *inside* this stream.onNext.
+ // requestObserver is already published, so an immediate ErrorResponse
+ // can write its required Cancel reentrantly without losing the request.
+ sendOnNext(initRequest)
+ }
+ }
+ }
+ } catch {
+ case NonFatal(e) =>
+ Transitions.transportFailed(e)
+ // Surface as GrpcWorkerSessionException so the engine integration layer
+ // (which catches that type and wraps it) sees a uniform init-failure
+ // exception rather than the raw transport error.
+ throw new GrpcWorkerSessionException("UDF worker stream failed during init", e)
+ }
+
+ try {
+ initValue.await(initResponseTimeoutMs)
+ } catch {
+ case _: InterruptedException =>
+ // Interrupt (cancelled query / killed task) while awaiting InitResponse:
+ // cooperatively Cancel with a bounded drain, settling Cancelled if the
+ // worker acks in time, else Interrupted. See handleInterrupt.
+ handleInterrupt("waiting for InitResponse")
+ throw new InterruptedException("interrupted while waiting for InitResponse")
+ case e: TimeoutException =>
+ sendCancelInternal(() => cancelWithReason("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.
+ Transitions.transportFailed(e)
+ // Surface as GrpcWorkerSessionException (carrying the timeout cause) so
+ // the engine integration layer that catches that type can wrap it.
+ throw new GrpcWorkerSessionException(
+ s"timed out waiting for InitResponse after ${initResponseTimeoutMs}ms", e)
+ }
+
+ initValue.get match {
+ case Some(resp) if resp.hasError =>
+ failInitWithError(resp.getError,
+ s"UDF worker init failed: ${describeError(resp.getError)}")
+ case Some(resp) =>
+ executionError.get() match {
+ case Some(err) =>
+ // A reentrant callback can report an ErrorResponse immediately after
+ // InitResponse and before the initial onNext returns. The callback
+ // has already sent Cancel; drain its response
+ // and surface the original error instead of returning init success
+ // for a session that is already cancelling or terminal.
+ failInitWithError(err,
+ s"UDF worker reported an error as init completed: ${describeError(err)}")
+ case None if currentState != SessionState.Initialized =>
+ // Likewise, an immediate protocol failure or concurrent close may
+ // have moved the session out of Initialized after publishing the
+ // InitResponse. A normal return would violate init()'s contract and
+ // leave process() to fail later with only an ordering error.
+ failInitFromCurrentState()
+ case None =>
+ resp
+ }
+ case None =>
+ // No InitResponse arrived but the latch fired.
+ //
+ // A pre-init ErrorResponse leaves the session in Initializing with the
+ // error recorded in executionError (see handleControl ERROR branch): the
+ // engine must now Cancel and drain the CancelResponse, which failInit does.
+ executionError.get() match {
+ case Some(err) if !currentState.isTerminal =>
+ failInitWithError(err,
+ s"UDF worker reported an error before init completed: ${describeError(err)}")
+ case _ =>
+ // Otherwise a terminal already settled -- the worker terminated the
+ // stream before sending InitResponse (transport error, half-close, or
+ // a premature FinishResponse/CancelResponse). Surface it as an init
+ // failure rather than letting the caller proceed as if init succeeded.
+ failInitFromCurrentState()
+ }
+ }
+ }
+
+ /** Surfaces the terminal that prevented init from completing normally. */
+ private def failInitFromCurrentState(): Nothing = {
+ executionError.get() match {
+ case Some(err) =>
+ failInitWithError(err,
+ s"UDF worker reported an error as init completed: ${describeError(err)}")
+ case None => ()
+ }
+ // Cancelling can be observed after a concurrent close or a reentrant error
+ // before its CancelResponse arrives. Drain it so the exception reflects the
+ // stable terminal rather than a transient state.
+ if (currentState == SessionState.Cancelling) {
+ awaitTerminal()
+ }
+ currentState match {
+ case SessionState.Terminal(Termination.TransportFailed(cause)) =>
+ throw new GrpcWorkerSessionException("UDF worker stream failed during init", cause)
+ case SessionState.Terminal(Termination.Failed(err)) =>
+ throw new GrpcWorkerSessionException(
+ s"UDF worker reported an error before init completed: ${describeError(err)}", err)
+ case SessionState.Terminal(Termination.Cancelled(_)) =>
+ throw new GrpcWorkerSessionException(
+ "UDF worker stream was cancelled before init completed")
+ case SessionState.Terminal(Termination.Interrupted(cause)) =>
+ throw new GrpcWorkerSessionException("UDF worker init was interrupted", cause)
+ case SessionState.Terminal(Termination.Finished(_)) =>
+ throw new GrpcWorkerSessionException("UDF worker finished before init completed")
+ case other =>
+ throw new IllegalStateException(
+ s"init completed without an accepted session or terminal: $other")
+ }
+ }
+
+ /**
+ * Fails init when the worker reports an error before init returns: an
+ * `InitResponse` carrying an error, a pre-init `ErrorResponse`, or an immediate
+ * post-init `ErrorResponse`. The protocol requires the engine to send `Cancel`
+ * and the worker to reply with `CancelResponse` (udf_message.proto); we send it,
+ * drain the terminator so no stream is left dangling, and throw the structured
+ * error. A responsive worker settles `Cancelled`; a failed drain settles
+ * `TransportFailed`.
+ */
+ private def failInitWithError(err: ExecutionError, message: String): Nothing = {
+ sendCancelInternal(() => cancelWithReason("init failed"))
+ awaitTerminal()
+ throw new GrpcWorkerSessionException(message, err)
+ }
+
+ override protected def doProcess(
+ input: Iterator[DataRequest],
+ finish: () => Finish): Iterator[DataResponse] = {
+ // Init success is guaranteed by the base [[WorkerSession]] lifecycle: if
+ // doInit had failed it would have thrown and process() would never run.
+ new ProcessIterator(input, finish)
+ }
+
+ override protected def doClose(cancel: () => Cancel): Termination = {
+ // Coordinate the no-stream decision with doInit's publication + Init write.
+ // If close wins this lock, doInit observes the terminal and never sends Init;
+ // if doInit wins, close observes the published observer and sends Cancel only
+ // after the Init send has returned.
+ val noPublishedStream = requestLock.synchronized {
+ if (requestObserver == null) {
+ // init() never put a stream on the wire (closed before/around init, or
+ // init threw before publishing). There is no protocol terminator; if no
+ // terminal has settled yet, treat the session as cancelled-before-start.
+ // A terminal may already be settled here (e.g. the channel-shutdown
+ // TransportFailed in doInit also leaves requestObserver null); return the
+ // settled terminal as-is rather than a bare Cancelled that disagrees with
+ // the state. The base WorkerSession still releases the worker handle.
+ if (!currentState.isTerminal) {
+ Transitions.cancelled(CancelResponse.getDefaultInstance)
+ }
+ true
+ } else {
+ false
+ }
+ }
+ if (noPublishedStream) {
+ return settledTermination
+ }
+ // If the stream has not finished on its own, cancel anything in flight so
+ // the worker can clean up, then wait for the terminator. sendCancelInternal
+ // evaluates the cancel thunk only when it attempts a Cancel, and at
+ // most once across all callers.
+ if (!currentState.isTerminal) {
+ sendCancelInternal(cancel)
+ try {
+ terminalLatch.await(terminalTimeoutMs, TimeUnit.MILLISECONDS)
+ } catch {
+ case _: InterruptedException =>
+ // Interrupted mid-close: settle Interrupted via the
+ // shared handler rather than falling through to the TransportFailed
+ // guard below. The Cancel was already attempted, so handleInterrupt's
+ // sendCancelInternal is a no-op and it just runs the bounded drain.
+ //
+ // TODO [SPARK-57640]: close() is a finalizer often already on the task
+ // kill path, yet this makes an interrupted close block up to another
+ // interruptCancelTimeoutMs (chasing a clean, salvageable Cancelled
+ // rather than immediately unwinding to Interrupted). Current choice:
+ // spend the bounded wait to obtain proof the worker is recyclable; it
+ // is small next to terminalTimeoutMs. Revisit if killed-task teardown
+ // latency (especially many sessions torn down at once) makes an
+ // immediate unwind preferable -- e.g. skip this second drain in the
+ // close() path, since close already gave the worker its terminalTimeoutMs window.
+ handleInterrupt("closing the session")
+ }
+ }
+ // If the worker still has not settled (timeout above, or an interrupt whose
+ // bounded drain did not settle a terminal), record a terminal so
+ // isWorkerSalvageable returns a definite answer and any other thread reading
+ // the state sees a stable value.
+ if (!currentState.isTerminal) {
+ Transitions.transportFailed(new TimeoutException(
+ s"timed out waiting for stream terminator after ${terminalTimeoutMs}ms"))
+ }
+ // Close the request side according to the settled outcome. Clean protocol
+ // terminators are half-closed; failures abort the RPC so close() never follows
+ // an onError with onCompleted.
+ currentState match {
+ case SessionState.Terminal(Termination.Finished(_)) |
+ SessionState.Terminal(Termination.Cancelled(_)) =>
+ completeRequestSide()
+ case SessionState.Terminal(Termination.Failed(error)) =>
+ abortRequestSide(new GrpcWorkerSessionException(
+ s"UDF execution failed: ${describeError(error)}", error))
+ case SessionState.Terminal(Termination.TransportFailed(cause)) =>
+ abortRequestSide(cause)
+ case SessionState.Terminal(Termination.Interrupted(cause)) =>
+ abortRequestSide(cause)
+ case other =>
+ logger.debug(s"UDF Execute stream closed without a terminal outcome: $other")
+ }
+ // Derive the Termination from the settled terminal. close() is the
+ // finalizer/cleanup path and must NOT re-throw: a UDF / data-phase error is
+ // already surfaced while the result iterator is consumed, and an init error
+ // is surfaced from init(). settledTermination returns the settled terminal
+ // as-is: the response proto for the clean terminators, and the failure
+ // terminals (Failed / TransportFailed / Interrupted) carrying their cause --
+ // not a bare Cancelled.
+ settledTermination
+ }
+
+ // ---- Internal request helpers ---------------------------------------------
+
+ /**
+ * Writes one request and aborts the request side if the observer rejects it.
+ * Callers hold [[requestLock]], so this also serializes the compensating
+ * `onError` with every other request-side operation.
+ */
+ private def sendOnNext(req: UdfRequest): Unit = {
+ try {
+ requestObserver.onNext(req)
+ } catch {
+ case NonFatal(e) =>
+ abortRequestSide(e)
+ throw e
+ }
+ }
+
+ /** Aborts the request side at most once. Safe to call while holding [[requestLock]]. */
+ private def abortRequestSide(cause: Throwable): Unit = requestLock.synchronized {
+ if (requestObserver != null && requestSideTerminated.compareAndSet(false, true)) {
+ try {
+ requestObserver.onError(cause)
+ } catch {
+ case NonFatal(e) => logger.debug("Error aborting UDF Execute stream", e)
+ }
+ }
+ }
+
+ /** Half-closes the request side at most once. */
+ private def completeRequestSide(): Unit = requestLock.synchronized {
+ if (requestObserver != null && requestSideTerminated.compareAndSet(false, true)) {
+ try {
+ requestObserver.onCompleted()
+ } catch {
+ case NonFatal(e) => logger.debug("Error half-closing UDF Execute stream", e)
+ }
+ }
+ }
+
+ /**
+ * Sends a Data or Finish request to the worker. Three invariants are
+ * checked inside [[requestLock]]:
+ * - terminal state: writes are unsafe (transport dead or terminal
+ * received). Throws, so the caller's terminal/exception path runs.
+ * - request-side termination: no write may follow `onError` / `onCompleted`.
+ * Throws so the caller settles or surfaces the failure.
+ * - [[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]] publishes
+ * [[requestObserver]] and writes Init together under [[requestLock]], so a
+ * concurrent cancel observes the observer but cannot acquire the lock and send
+ * Cancel until after Init.
+ */
+ private def sendRequest(req: UdfRequest): Unit =
+ requestLock.synchronized {
+ if (currentState.isTerminal) {
+ throw new IllegalStateException(
+ "cannot send request: UDF Execute stream is already closed")
+ }
+ if (requestSideTerminated.get()) {
+ throw new IllegalStateException(
+ "cannot send request: UDF Execute request stream is already closed")
+ }
+ // Suppress the write when a cancel has been (or is about to be) flushed
+ // through this lock; that preserves the proto ordering invariant (no
+ // Data/Finish after Cancel). Test the flag rather than `return`-ing from
+ // this by-name `synchronized` body: a non-local return compiles to a thrown
+ // NonLocalReturnControl, which is brittle (also relied on in sendCancelInternal).
+ if (!cancelRequested.get()) {
+ if (req.hasControl && req.getControl.hasFinish) {
+ // Set before onNext: directExecutor delivery may return FinishResponse
+ // reentrantly from inside this call.
+ finishSendStarted.set(true)
+ }
+ sendOnNext(req)
+ }
+ }
+
+ /**
+ * Sends a `Cancel` control message, returning `true` iff the request observer
+ * accepts the `Cancel`. Idempotent across ALL call sites (close()'s
+ * cancel, in-band cancels from `handleControl`'s INIT error / ERROR paths,
+ * and engine-internal cancels from the iterator) via [[cancelRequested]]:
+ * only the first caller attempts a `Cancel`, and the `cancel` thunk is evaluated
+ * only by that caller and only when the send is about to start --
+ * so a side-effecting thunk (e.g. one carrying a client cancel callback) runs
+ * at most once. 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(cancel: () => Cancel): Boolean = {
+ // Do not consume the one-shot cancellation intent until there is an observer
+ // that can carry it. doInit publishes the observer before sending Init while
+ // holding requestLock, so every response to Init sees a non-null observer.
+ if (requestObserver == null) return false
+ if (!cancelRequested.compareAndSet(false, true)) return false
+ if (currentState.isTerminal) return false
+ try {
+ requestLock.synchronized {
+ // A terminal that arrived first must win; bail without writing. Compute
+ // the block's value rather than `return` (see sendRequest for why non-local
+ // return is avoided -- here it would also escape the NonFatal catch below).
+ val cur = currentState
+ if (cur.isTerminal || requestSideTerminated.get()) {
+ false
+ } else {
+ val request = UdfRequest.newBuilder()
+ .setControl(UdfControlRequest.newBuilder().setCancel(cancel()).build())
+ .build()
+ // Record that a Cancel send has started by advancing to Cancelling.
+ Transitions.beginCancelFrom(cur)
+ sendOnNext(request)
+ true
+ }
+ }
+ } catch {
+ case NonFatal(e) =>
+ logger.debug(s"Cancel send failed (stream may already be torn down): ${e.getMessage}")
+ Transitions.transportFailed(e)
+ false
+ }
+ }
+
+ private def awaitTerminal(): Unit = {
+ if (currentState.isTerminal) return
+ try {
+ if (!terminalLatch.await(terminalTimeoutMs, TimeUnit.MILLISECONDS)) {
+ Transitions.transportFailed(new TimeoutException(
+ s"timed out waiting for stream terminator after ${terminalTimeoutMs}ms"))
+ }
+ } catch {
+ case _: InterruptedException => handleInterrupt("waiting for stream terminator")
+ }
+ }
+
+ /**
+ * Handles an [[InterruptedException]] observed while blocked on a worker event
+ * (init, terminal drain, close, or the result-iterator poll). An interrupt is
+ * an engine-side event -- typically a cancelled query / killed task -- not a
+ * worker fault, so we cancel cooperatively and keep the worker salvageable only
+ * when it acks:
+ *
+ * 1. Send a best-effort `Cancel` (idempotent across call sites).
+ * 2. Wait up to [[interruptCancelTimeoutMs]] -- short, so the interrupted
+ * thread unwinds promptly -- for the worker's `CancelResponse`. A healthy
+ * worker acks in that window and the callback settles a clean `Cancelled`
+ * terminal (worker salvageable, response drained).
+ * 3. If the ack does not arrive in time, settle `Interrupted`. Without a
+ * terminator or liveness proof the worker is not salvageable and must not
+ * be returned to a reuse pool.
+ *
+ * The bounded wait deliberately runs '''before''' the thread's interrupt flag
+ * is restored: [[InterruptedException]] clears the flag on throw, so re-setting
+ * it first would make the wait below throw immediately and defeat the drain.
+ * The flag is restored at the end so the caller's unwind still observes the
+ * interrupt.
+ *
+ * TODO [SPARK-57640]: add a worker liveness/heartbeat check so a future version
+ * may prove that an interrupted worker which missed the short ack window is
+ * nevertheless safe to recycle.
+ */
+ private def handleInterrupt(waitContext: String): Unit = {
+ // NB: flag is currently clear (InterruptedException cleared it); do not
+ // restore it until after the bounded drain below.
+ if (!currentState.isTerminal) {
+ sendCancelInternal(() => cancelWithReason(s"interrupted while $waitContext"))
+ val acked = try {
+ terminalLatch.await(interruptCancelTimeoutMs, TimeUnit.MILLISECONDS)
+ } catch {
+ case _: InterruptedException =>
+ // A second interrupt during the drain: give up the wait immediately.
+ false
+ }
+ if (!acked && !currentState.isTerminal) {
+ // No CancelResponse in time: settle Interrupted rather than
+ // TransportFailed, and stop a later close() from blocking on a terminator
+ // that will not be drained on this thread's behalf. Interrupted remains
+ // unsalvageable until a separate liveness check can prove otherwise.
+ Transitions.interrupted(
+ new InterruptedException(s"interrupted while $waitContext"))
+ }
+ }
+ Thread.currentThread().interrupt()
+ }
+
+ // ---- 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 settle the state) and with any thread that finalizes via
+ * [[close]]. It drives the `Streaming -> Finishing` transition (sending
+ * `Finish` once input is exhausted, built lazily from the `finish` thunk).
+ */
+ private class ProcessIterator(input: Iterator[DataRequest], finish: () => Finish)
+ extends Iterator[DataResponse] {
+
+ // 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 branches, and block branch 4 for the
+ // full terminalTimeoutMs before returning. Callers that probe hasNext
+ // an extra time (iterator.size, instrumentation wrappers) would hang.
+ // Iterator-local and only touched by the single engine thread.
+ private val exhausted = new AtomicBoolean(false)
+ @volatile private var prefetched: DataResponse = _
+
+ /**
+ * Runs a caller-supplied thunk (`input.hasNext`, `input.next()`, or the
+ * `finish` builder) that may throw, and on failure Cancels the in-flight
+ * stream before rethrowing. Without the Cancel a failure would leave the
+ * worker awaiting more input with no terminator owed on the wire; with it the
+ * worker tears down and the exception still propagates to the engine. Setting
+ * `cancelRequested` also suppresses any further Data/Finish this loop might
+ * otherwise attempt (see [[sendRequest]]).
+ */
+ private def cancelOnThrow[T](reason: String)(op: => T): T =
+ try op catch {
+ case NonFatal(e) =>
+ sendCancelInternal(() => cancelWithReason(reason))
+ throw e
+ }
+
+ override def hasNext: Boolean = {
+ if (prefetched ne null) return true
+ advance()
+ prefetched ne null
+ }
+
+ override def next(): DataResponse = {
+ if (prefetched eq null) advance()
+ val out = prefetched
+ if (out eq null) {
+ throw new NoSuchElementException("ProcessIterator exhausted")
+ }
+ prefetched = null
+ out
+ }
+
+ /**
+ * Fills [[prefetched]] with the next output batch, or leaves it null at the
+ * terminator (which also sets [[exhausted]], so `hasNext` reads null as "done").
+ * Each loop iteration tries in order: (1) drain queued output -- a batch (fill
+ * and return) or the terminator sentinel (return); (2) while `Streaming`, send
+ * the next input batch and loop; (3) once input is exhausted, send `Finish`
+ * (CAS `Streaming -> Finishing`, once) and loop; (4) otherwise block for late
+ * output or the terminator, then return. Branches 2-3 loop; 1 and 4 return.
+ *
+ * '''Per-poll timeout.''' Branch 4 waits up to [[terminalTimeoutMs]] per
+ * poll, reset by every worker event -- the contract is "emit at least one
+ * event every [[terminalTimeoutMs]] after Finish", not "finish the UDF within
+ * [[terminalTimeoutMs]]". A worker expecting a long post-Finish silence MAY
+ * emit an empty `DataResponse` as a heartbeat to reset the wait; it is
+ * surfaced to the caller, so it should be a batch the caller recognises as
+ * empty (e.g. a zero-row Arrow batch).
+ */
+ private def advance(): Unit = {
+ if (exhausted.get()) return // terminator already seen; never re-block
+ 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 while the stream is open for data.
+ // `input.hasNext` may itself fetch/compute the next element (many Spark
+ // iterators prefetch), so it can throw just like `input.next()`; both go
+ // through cancelOnThrow so a failing upstream Cancels the stream instead
+ // of stranding the worker waiting for input that will never arrive.
+ if (!cancelRequested.get() && currentState == SessionState.Streaming &&
+ cancelOnThrow("input iterator failed")(input.hasNext)) {
+ val request = cancelOnThrow("input iterator failed") {
+ val batch = Objects.requireNonNull(
+ input.next(), "input iterator returned null")
+ UdfRequest.newBuilder().setData(batch).build()
+ }
+ if (!sendOrEndOnRacedTerminal(request)) return
+ } else if (!cancelRequested.get() && Transitions.beginFinish()) {
+ // (3) No more input; send Finish exactly once (unless cancelled). The
+ // `finish` thunk is caller-supplied (it may run a finish callback), so
+ // it can throw; cancelOnThrow Cancels the stream before rethrowing. The
+ // Streaming -> Finishing CAS has already run, but we have not written
+ // Finish yet, so the Cancel is the only engine-to-worker message that
+ // reaches the wire -- the proto "nothing after Cancel" invariant holds.
+ val request = cancelOnThrow("finish callback failed") {
+ val finishMsg = Objects.requireNonNull(
+ finish(), "finish callback returned null")
+ UdfRequest.newBuilder()
+ .setControl(UdfControlRequest.newBuilder().setFinish(finishMsg).build())
+ .build()
+ }
+ if (!sendOrEndOnRacedTerminal(request)) {
+ return
+ }
+ } else {
+ // (4) Block for late output or the terminator. See class doc
+ // above for the per-poll-vs-total-session timeout semantics.
+ val item = try {
+ outputQueue.poll(terminalTimeoutMs, TimeUnit.MILLISECONDS)
+ } catch {
+ case _: InterruptedException =>
+ // Interrupt (cancelled query / killed task) while reading results:
+ // cooperatively Cancel with a bounded drain, settling Cancelled if
+ // the worker acks in time, else Interrupted. See handleInterrupt.
+ handleInterrupt("reading UDF result")
+ exhausted.set(true)
+ throw new InterruptedException("interrupted while reading UDF result")
+ }
+ item match {
+ case null =>
+ Transitions.transportFailed(new IllegalStateException(
+ s"timed out waiting for UDF output after ${terminalTimeoutMs}ms"))
+ exhausted.set(true)
+ throwIfTerminalError()
+ return
+ case QueueItem.EndOfStream =>
+ exhausted.set(true)
+ throwIfTerminalError()
+ return
+ case QueueItem.Batch(b) =>
+ prefetched = b
+ return
+ }
+ }
+ }
+ }
+
+ /**
+ * Sends one data-phase request (`DataRequest` / `Finish`), recovering from a
+ * terminator that raced the write. Returns `true` to keep looping (sent, or
+ * suppressed by a pending cancel), `false` if a terminator settled and the
+ * iterator should end (caller `return`s).
+ *
+ * On a write failure a terminator has usually already settled (the worker
+ * finished/failed early) and the write only failed because the stream is
+ * closed: record a transport terminal if none is set ([[completeTerminal]] is
+ * a no-op once settled), then [[throwIfTerminalError]] -- which throws for an
+ * error/transport terminal, or returns for a clean `Finished` that raced the
+ * send, in which case the benign "stream closed" error is dropped.
+ */
+ private def sendOrEndOnRacedTerminal(req: UdfRequest): Boolean = {
+ try {
+ sendRequest(req)
+ true
+ } catch {
+ case NonFatal(e) =>
+ if (!currentState.isTerminal) {
+ Transitions.transportFailed(e)
+ }
+ exhausted.set(true)
+ throwIfTerminalError()
+ false
+ }
+ }
+
+ /**
+ * Surfaces a failed terminator as an exception when the result iterator is
+ * drained. A data-phase [[ExecutionError]] (captured in [[executionError]])
+ * takes precedence over the terminator's own error, then the finish/cancel
+ * callback error carried on the terminator, then a bare cancellation.
+ */
+ private def throwIfTerminalError(): Unit = currentState match {
+ case SessionState.Terminal(Termination.Finished(response)) =>
+ responseError(response.hasError, response.getError, executionError.get())
+ .foreach(err => throw new GrpcWorkerSessionException(
+ s"UDF execution failed: ${describeError(err)}", err))
+ case SessionState.Terminal(Termination.Cancelled(response)) =>
+ responseError(response.hasError, response.getError, executionError.get()) match {
+ case Some(err) =>
+ throw new GrpcWorkerSessionException(
+ s"UDF execution failed: ${describeError(err)}", err)
+ case None =>
+ throw new GrpcWorkerSessionException("UDF execution was cancelled")
+ }
+ case SessionState.Terminal(Termination.Failed(err)) =>
+ throw new GrpcWorkerSessionException(
+ s"UDF execution failed: ${describeError(err)}", err)
+ case SessionState.Terminal(Termination.TransportFailed(t)) =>
+ throw new GrpcWorkerSessionException("UDF worker stream failed", t)
+ case SessionState.Terminal(Termination.Interrupted(t)) =>
+ throw new GrpcWorkerSessionException("UDF execution was interrupted", t)
+ case other =>
+ throw new IllegalStateException(s"terminator sentinel without terminal: $other")
+ }
+
+ /** Picks the error to surface: prior data-phase error, else the terminator's. */
+ private def responseError(
+ hasError: Boolean,
+ error: ExecutionError,
+ priorError: Option[ExecutionError]): Option[ExecutionError] =
+ priorError.orElse(if (hasError) Some(error) else None)
+ }
+}
+
+object GrpcWorkerSession {
+ /** Upper bound on the wait for `InitResponse`. */
+ val DEFAULT_INIT_RESPONSE_TIMEOUT_MS: Long = 30000L
+
+ /** Upper bound on the wait for `FinishResponse` / `CancelResponse`. */
+ val DEFAULT_TERMINAL_TIMEOUT_MS: Long = 30000L
+
+ /**
+ * Upper bound on the wait for a `CancelResponse` after an interrupt (e.g. a
+ * cancelled query / killed task) sends `Cancel`. Much shorter than
+ * [[DEFAULT_TERMINAL_TIMEOUT_MS]]: the interrupted thread must unwind
+ * promptly, so a healthy worker gets a brief window to ack the Cancel (clean
+ * `Cancelled` terminal, worker salvageable) before the session falls back to
+ * an unsalvageable `Interrupted` terminal. See `handleInterrupt`.
+ */
+ val DEFAULT_INTERRUPT_CANCEL_TIMEOUT_MS: Long = 2000L
+
+ // Distinguishes a (possibly empty) data batch from end-of-stream, and makes
+ // the iterator's match exhaustive.
+ private sealed trait QueueItem
+ private object QueueItem {
+ final case class Batch(response: DataResponse) extends QueueItem
+ case object EndOfStream extends QueueItem
+ }
+
+ /**
+ * A write-once value paired with the latch a waiter blocks on, so the "publish
+ * the value, then release the waiter" ordering lives in one place instead of
+ * every call site having to remember to count the latch down after setting the
+ * reference. The value is set at most once ([[complete]]); the latch can also
+ * be released without a value ([[signalWithoutValue]]) when init fails through
+ * a pre-init error or terminal. All reads/writes go through the latch, so a
+ * waiter released by [[await]] has a happens-before edge to the [[complete]]
+ * that set the value.
+ */
+ private final class OneShotValue[A] {
+ private val latch = new CountDownLatch(1)
+ private val ref = new AtomicReference[Option[A]](None)
+ private val completed = new AtomicBoolean(false)
+
+ /** Publishes the value (first writer wins) and releases any waiter. */
+ def complete(value: A): Unit = {
+ if (completed.compareAndSet(false, true)) {
+ ref.set(Some(value))
+ latch.countDown()
+ }
+ }
+
+ /** Completes without a value (init failed through another event) and releases the waiter. */
+ def signalWithoutValue(): Unit = {
+ if (completed.compareAndSet(false, true)) {
+ latch.countDown()
+ }
+ }
+
+ /**
+ * Blocks up to `timeoutMs` for a release, throwing [[TimeoutException]] if
+ * none arrives so the caller handles the timeout on the exception path rather
+ * than by inspecting a return value.
+ */
+ def await(timeoutMs: Long): Unit =
+ if (!latch.await(timeoutMs, TimeUnit.MILLISECONDS)) {
+ throw new TimeoutException(
+ s"timed out waiting for InitResponse after ${timeoutMs}ms")
Review Comment:
fixed
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57640] Implement gRPC WorkerSession and test it against an echo server. [spark]
Posted by "haiyangsun-db (via GitHub)" <gi...@apache.org>.
haiyangsun-db commented on code in PR #56702:
URL: https://github.com/apache/spark/pull/56702#discussion_r3665860704
##########
udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/GrpcWorkerSession.scala:
##########
@@ -0,0 +1,933 @@
+/*
+ * 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.grpc
+
+import java.util.concurrent.{CountDownLatch, LinkedBlockingQueue, TimeoutException, TimeUnit}
+import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference}
+
+import scala.util.control.NonFatal
+
+import io.grpc.{ConnectivityState, ManagedChannel}
+import io.grpc.stub.StreamObserver
+
+import org.apache.spark.annotation.Experimental
+import org.apache.spark.udf.worker.{Cancel, CancelResponse, DataRequest, DataResponse,
+ ExecutionError, Finish, FinishResponse, Init, InitResponse, UdfControlRequest,
+ UdfControlResponse, UdfRequest, UdfResponse, UdfWorkerGrpc}
+import org.apache.spark.udf.worker.core.{Termination, WorkerHandle, WorkerLogger, WorkerSession}
+import org.apache.spark.udf.worker.core.WorkerSession.SessionState
+import org.apache.spark.udf.worker.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_message.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.
+ *
+ * '''Driving model.''' Consumption-driven (Volcano / pull): the thread that
+ * consumes the [[doProcess]] result iterator is the one that pulls input and
+ * sends each `DataRequest`; the gRPC callback thread only receives output. It is
+ * pull-driven but not one-input-per-output -- `advance` sends the next input
+ * whenever the output queue is momentarily empty, so under async delivery it may
+ * push several input batches before any output is read (bounded by HTTP/2 flow
+ * control).
+ *
+ * '''State machine.''' This class does not keep its own state machine: it
+ * drives the single [[WorkerSession.SessionState]] owned by the base. The base
+ * advances `Created -> Initializing` (in `init`) and `Initialized -> Streaming`
+ * (in `process`); this class advances the protocol-event edges through
+ * [[compareAndSetState]] / [[completeTerminal]] as it exchanges messages:
+ * {{{
+ * Initializing --(InitResponse ok)--> Initialized [handleControl]
+ * Streaming ----(Finish written)----> Finishing [ProcessIterator]
+ * <any non-terminal> --(Cancel written)--> Cancelling [sendCancelInternal]
+ * <any non-terminal> --(terminator/error)--> terminal [completeTerminal]
+ * }}}
+ * The two clean terminals carry the worker's `FinishResponse` / `CancelResponse`
+ * (metrics + finish/cancel callback `data`/`error`) so [[close]] can return
+ * them. The only flag kept outside the machine is [[cancelRequested]]: a
+ * cancellation can be requested before the stream exists (so it cannot be a
+ * state transition yet), and it must both fast-fail the result iterator and
+ * suppress any in-flight Data/Finish.
+ *
+ * Threading:
+ * - [[doInit]] is synchronous: sends `Init` and blocks on `InitResponse`,
+ * returning it.
+ * - [[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.
+ * - [[doClose]] is thread-safe and idempotent: it settles + returns the
+ * terminator (cancelling in-flight work if the stream had not finished)
+ * and half-closes the request side.
+ *
+ * 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 a gRPC channel built and owned by the caller (the
+ * dispatcher). 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 from the worker, drained by the process() iterator.
+ // Intentionally unbounded: a bounded queue would block the gRPC callback
+ // (Netty event-loop) thread when full, stalling terminator/control delivery on
+ // the whole channel. HTTP/2 flow control bounds the wire and the consumer
+ // normally drains promptly; a stalled downstream can still grow it, but the fix
+ // is protocol-level back-pressure (out of scope), not bounding the queue.
+ //
+ // TODO [SPARK-57324]: expose queue depth as a metric (early warning for a
+ // stalled consumer).
+ private val outputQueue = new LinkedBlockingQueue[QueueItem]()
+
+ // The worker's `InitResponse` (success or error) coupled with the latch that
+ // init() blocks on. Fires when the InitResponse arrives (`complete`) or when a
+ // terminal settles first without one -- transport error, half-close, or a
+ // premature terminator (`signalWithoutValue`). Until it fires we have no proof
+ // the worker accepted the session.
+ //
+ // Settle-before-release rule (referenced from every callback that both settles
+ // a terminal and fires this latch): settle the terminal FIRST, then complete /
+ // signal. init() blocks on the latch, and only latch await/release establish a
+ // happens-before edge, so a woken init() is guaranteed to observe the terminal
+ // rather than a transient state. OneShotValue keeps that publish-then-release
+ // in one place instead of every caller remembering to count down after setting
+ // the reference.
+ private val initValue = new OneShotValue[InitResponse]
+
+ // Fired when the session reaches a terminal [[SessionState]]. doClose() and
+ // the init-error path block on this to drain the terminator.
+ private val terminalLatch = new CountDownLatch(1)
+
+ // Captures an ErrorResponse encountered during the data phase so that
+ // the CancelResponse terminator 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)
+
+ // Cancellation INTENT -- distinct from the `Cancelling` state, which is reached
+ // only once a Cancel is actually written to the wire ([[sendCancelInternal]]).
+ // Intent is set first and can outrun (or never reach) that write, so
+ // `cancelRequested` does NOT imply state `Cancelling`. Kept outside the
+ // [[SessionState]] machine because a cancel can be requested before the stream
+ // exists (pre-init), where there is no wire transition to make yet. Used to (a)
+ // make cancel idempotent across all call sites, (b) fast-fail ProcessIterator
+ // on a pre-init cancel, and (c) suppress any Data/Finish that would otherwise
+ // race a Cancel onto the wire (re-read inside [[requestLock]]).
+ private val cancelRequested = new AtomicBoolean(false)
+
+ // gRPC requires serialized writes to a request StreamObserver.
+ private val requestLock = new Object
+
+ // Initialised in init() -- before that, close() is a no-op 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 =>
+ // A DataResponse before InitResponse violates the protocol (InitResponse
+ // must precede any DataResponse): fast-fail rather than enqueue it and let
+ // init() block to its timeout. Settle-before-release (see initValue).
+ if (!initResolved) {
+ Transitions.transportFailed(new IllegalStateException(
+ "worker sent a DataResponse before InitResponse"))
+ initValue.signalWithoutValue()
+ } else {
+ outputQueue.put(QueueItem.Batch(response.getData))
+ }
+
+ case UdfResponse.ResponseCase.CONTROL =>
+ handleControl(response.getControl)
+
+ case other =>
+ // A malformed response (empty / unknown oneof) is a terminal transport
+ // failure; fast-fail init the same way. Settle-before-release (see initValue).
+ Transitions.transportFailed(new IllegalStateException(
+ s"unexpected response oneof: $other"))
+ initValue.signalWithoutValue()
+ }
+ }
+
+ override def onError(t: Throwable): Unit = {
+ // Transport-level failure: the stream is dead, no further writes possible.
+ // Settle-before-release (see initValue) so init() surfaces the transport
+ // cause instead of the initResponseTimeoutMs "timed out" error.
+ Transitions.transportFailed(t)
+ initValue.signalWithoutValue()
+ }
+
+ 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 (!currentState.isTerminal) {
+ Transitions.transportFailed(new IllegalStateException(
+ "worker response stream closed without a terminator"))
+ }
+ // Defensive: if onCompleted reached us before InitResponse, doInit is
+ // still blocked on initValue and would otherwise time out.
+ initValue.signalWithoutValue()
+ }
+ }
+
+ /**
+ * Wakes the result iterator (blocked on [[outputQueue]]) and any thread
+ * waiting on [[terminalLatch]] when the base settles a terminal. Invoked once,
+ * by the caller that wins [[completeTerminal]].
+ */
+ override protected def onTerminalSettled(termination: Termination): Unit = {
+ outputQueue.put(QueueItem.EndOfStream)
+ terminalLatch.countDown()
+ }
+
+ private def handleControl(ctrl: UdfControlResponse): Unit = ctrl.getControlCase match {
+ case UdfControlResponse.ControlCase.INIT =>
+ val resp = ctrl.getInit
+ if (resp.hasError) {
+ // Attribute a subsequent close()'s terminator to the init error rather
+ // than a bare Cancelled, mirroring the data-phase ERROR branch.
+ executionError.compareAndSet(None, Some(resp.getError))
+ // Settle Failed(err) before the Cancel below (not after its CancelResponse):
+ // the sticky terminal makes close() report the init error, and the Cancel
+ // becomes best-effort worker cleanup -- so if it can't be written (a
+ // directExecutor worker delivered this reentrantly, before doInit published
+ // requestObserver), doInit still doesn't wait terminalTimeoutMs for a
+ // CancelResponse that can never arrive. Settle-before-release (see initValue).
+ Transitions.failed(resp.getError)
Review Comment:
INIT-error branch now executionError.compareAndSet(...) + initValue.complete(resp), leaves state Initializing; doInit→failInitWithError sends Cancel then drains (L506-508). No early failed.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57640] Implement gRPC WorkerSession and test it against an echo server. [spark]
Posted by "haiyangsun-db (via GitHub)" <gi...@apache.org>.
haiyangsun-db commented on PR #56702:
URL: https://github.com/apache/spark/pull/56702#issuecomment-5107090751
@Yicong-Huang
Thanks for the thorough review -- this was really helpful. I've pushed a series of commits addressing everything; recap below, with the design discussion on high-level point 1 at the end.
**High-level point 2 (Cancel-after-init-error should go through Cancelling, not settle Failed early):** done. The init-error path no longer settles `Failed` in the response callback; it records the error and leaves the session `Initializing`, and `doInit` (where `requestObserver` is guaranteed published) sends the `Cancel` and drains the `CancelResponse`, so the terminal settles `Cancelled`. `Transitions.failed` was removed entirely.
**Inline comments:** all addressed in the threads (recap so it's in one place):
- INIT-error no longer settles a terminal early -> the wire `Cancel` is now actually written and its `CancelResponse` drained.
- `awaitTerminal()` genuinely blocks for the `CancelResponse` again (it was a no-op only because of the early-settle bug above).
- `doClose`'s no-stream path only settles `Cancelled` when not already terminal, then returns `settledTermination` (no more fresh `Cancelled` disagreeing with a settled state).
- `OneShotValue.await` now returns `Unit` and throws `TimeoutException` instead of returning a boolean.
- Interrupt handling is now a first-class `Interrupted` termination (see that thread): a best-effort `Cancel` with a short bounded wait for the ack -- a clean, salvageable `Cancelled` if the worker responds in time, otherwise `Interrupted` (also salvageable, since the interrupt is an engine-side event, not a worker fault).
**High-level point 1 (init-handshake complexity):**
You're right that the signaling was the main source of noise, and it turned out to be mostly accidental. `onTerminalSettled` is now the single place a settled terminal wakes a blocked `init()`, so every branch that settles a terminal (the `responseObserver` failure/terminator cases, the timeout paths, `close()`) no longer signals the init waiter itself -- it just settles. That removed 7 redundant `signalWithoutValue()` calls; the only direct init signals left are the three genuinely non-terminal paths (`INIT`-ok, `INIT`-error, pre-init `ERROR`), which must wake `init()` *without* settling a terminal so the session stays `Initializing` and `process()` can still run. That's the irreducible core of the handshake.
On fully isolating init behind a one-shot handler and switching to a data-phase handler once it resolves: I looked at this and decided against it. gRPC gives us a single `StreamObserver` for the stream's whole life, so "switching" means either a swappable `@volatile var handler` or an `if (initResolved) ...` at the top of `onNext`. The former adds a mutable concurrency surface and a swap-vs-reentrant-delivery race -- exactly the trickiest part of this class -- and the latter just relocates the two `initResolved` checks that already exist. Neither removes the essential constraint (reentrant directExecutor delivery forcing the deferred `Cancel`), so I don't think the extra indirection pays for itself. Happy to revisit if you feel strongly.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57640] Implement gRPC WorkerSession and test it against an echo server. [spark]
Posted by "haiyangsun-db (via GitHub)" <gi...@apache.org>.
haiyangsun-db commented on code in PR #56702:
URL: https://github.com/apache/spark/pull/56702#discussion_r3665862323
##########
udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/GrpcWorkerSession.scala:
##########
@@ -0,0 +1,933 @@
+/*
+ * 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.grpc
+
+import java.util.concurrent.{CountDownLatch, LinkedBlockingQueue, TimeoutException, TimeUnit}
+import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference}
+
+import scala.util.control.NonFatal
+
+import io.grpc.{ConnectivityState, ManagedChannel}
+import io.grpc.stub.StreamObserver
+
+import org.apache.spark.annotation.Experimental
+import org.apache.spark.udf.worker.{Cancel, CancelResponse, DataRequest, DataResponse,
+ ExecutionError, Finish, FinishResponse, Init, InitResponse, UdfControlRequest,
+ UdfControlResponse, UdfRequest, UdfResponse, UdfWorkerGrpc}
+import org.apache.spark.udf.worker.core.{Termination, WorkerHandle, WorkerLogger, WorkerSession}
+import org.apache.spark.udf.worker.core.WorkerSession.SessionState
+import org.apache.spark.udf.worker.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_message.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.
+ *
+ * '''Driving model.''' Consumption-driven (Volcano / pull): the thread that
+ * consumes the [[doProcess]] result iterator is the one that pulls input and
+ * sends each `DataRequest`; the gRPC callback thread only receives output. It is
+ * pull-driven but not one-input-per-output -- `advance` sends the next input
+ * whenever the output queue is momentarily empty, so under async delivery it may
+ * push several input batches before any output is read (bounded by HTTP/2 flow
+ * control).
+ *
+ * '''State machine.''' This class does not keep its own state machine: it
+ * drives the single [[WorkerSession.SessionState]] owned by the base. The base
+ * advances `Created -> Initializing` (in `init`) and `Initialized -> Streaming`
+ * (in `process`); this class advances the protocol-event edges through
+ * [[compareAndSetState]] / [[completeTerminal]] as it exchanges messages:
+ * {{{
+ * Initializing --(InitResponse ok)--> Initialized [handleControl]
+ * Streaming ----(Finish written)----> Finishing [ProcessIterator]
+ * <any non-terminal> --(Cancel written)--> Cancelling [sendCancelInternal]
+ * <any non-terminal> --(terminator/error)--> terminal [completeTerminal]
+ * }}}
+ * The two clean terminals carry the worker's `FinishResponse` / `CancelResponse`
+ * (metrics + finish/cancel callback `data`/`error`) so [[close]] can return
+ * them. The only flag kept outside the machine is [[cancelRequested]]: a
+ * cancellation can be requested before the stream exists (so it cannot be a
+ * state transition yet), and it must both fast-fail the result iterator and
+ * suppress any in-flight Data/Finish.
+ *
+ * Threading:
+ * - [[doInit]] is synchronous: sends `Init` and blocks on `InitResponse`,
+ * returning it.
+ * - [[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.
+ * - [[doClose]] is thread-safe and idempotent: it settles + returns the
+ * terminator (cancelling in-flight work if the stream had not finished)
+ * and half-closes the request side.
+ *
+ * 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 a gRPC channel built and owned by the caller (the
+ * dispatcher). 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 from the worker, drained by the process() iterator.
+ // Intentionally unbounded: a bounded queue would block the gRPC callback
+ // (Netty event-loop) thread when full, stalling terminator/control delivery on
+ // the whole channel. HTTP/2 flow control bounds the wire and the consumer
+ // normally drains promptly; a stalled downstream can still grow it, but the fix
+ // is protocol-level back-pressure (out of scope), not bounding the queue.
+ //
+ // TODO [SPARK-57324]: expose queue depth as a metric (early warning for a
+ // stalled consumer).
+ private val outputQueue = new LinkedBlockingQueue[QueueItem]()
+
+ // The worker's `InitResponse` (success or error) coupled with the latch that
+ // init() blocks on. Fires when the InitResponse arrives (`complete`) or when a
+ // terminal settles first without one -- transport error, half-close, or a
+ // premature terminator (`signalWithoutValue`). Until it fires we have no proof
+ // the worker accepted the session.
+ //
+ // Settle-before-release rule (referenced from every callback that both settles
+ // a terminal and fires this latch): settle the terminal FIRST, then complete /
+ // signal. init() blocks on the latch, and only latch await/release establish a
+ // happens-before edge, so a woken init() is guaranteed to observe the terminal
+ // rather than a transient state. OneShotValue keeps that publish-then-release
+ // in one place instead of every caller remembering to count down after setting
+ // the reference.
+ private val initValue = new OneShotValue[InitResponse]
+
+ // Fired when the session reaches a terminal [[SessionState]]. doClose() and
+ // the init-error path block on this to drain the terminator.
+ private val terminalLatch = new CountDownLatch(1)
+
+ // Captures an ErrorResponse encountered during the data phase so that
+ // the CancelResponse terminator 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)
+
+ // Cancellation INTENT -- distinct from the `Cancelling` state, which is reached
+ // only once a Cancel is actually written to the wire ([[sendCancelInternal]]).
+ // Intent is set first and can outrun (or never reach) that write, so
+ // `cancelRequested` does NOT imply state `Cancelling`. Kept outside the
+ // [[SessionState]] machine because a cancel can be requested before the stream
+ // exists (pre-init), where there is no wire transition to make yet. Used to (a)
+ // make cancel idempotent across all call sites, (b) fast-fail ProcessIterator
+ // on a pre-init cancel, and (c) suppress any Data/Finish that would otherwise
+ // race a Cancel onto the wire (re-read inside [[requestLock]]).
+ private val cancelRequested = new AtomicBoolean(false)
+
+ // gRPC requires serialized writes to a request StreamObserver.
+ private val requestLock = new Object
+
+ // Initialised in init() -- before that, close() is a no-op 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 =>
+ // A DataResponse before InitResponse violates the protocol (InitResponse
+ // must precede any DataResponse): fast-fail rather than enqueue it and let
+ // init() block to its timeout. Settle-before-release (see initValue).
+ if (!initResolved) {
+ Transitions.transportFailed(new IllegalStateException(
+ "worker sent a DataResponse before InitResponse"))
+ initValue.signalWithoutValue()
+ } else {
+ outputQueue.put(QueueItem.Batch(response.getData))
+ }
+
+ case UdfResponse.ResponseCase.CONTROL =>
+ handleControl(response.getControl)
+
+ case other =>
+ // A malformed response (empty / unknown oneof) is a terminal transport
+ // failure; fast-fail init the same way. Settle-before-release (see initValue).
+ Transitions.transportFailed(new IllegalStateException(
+ s"unexpected response oneof: $other"))
+ initValue.signalWithoutValue()
+ }
+ }
+
+ override def onError(t: Throwable): Unit = {
+ // Transport-level failure: the stream is dead, no further writes possible.
+ // Settle-before-release (see initValue) so init() surfaces the transport
+ // cause instead of the initResponseTimeoutMs "timed out" error.
+ Transitions.transportFailed(t)
+ initValue.signalWithoutValue()
+ }
+
+ 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 (!currentState.isTerminal) {
+ Transitions.transportFailed(new IllegalStateException(
+ "worker response stream closed without a terminator"))
+ }
+ // Defensive: if onCompleted reached us before InitResponse, doInit is
+ // still blocked on initValue and would otherwise time out.
+ initValue.signalWithoutValue()
+ }
+ }
+
+ /**
+ * Wakes the result iterator (blocked on [[outputQueue]]) and any thread
+ * waiting on [[terminalLatch]] when the base settles a terminal. Invoked once,
+ * by the caller that wins [[completeTerminal]].
+ */
+ override protected def onTerminalSettled(termination: Termination): Unit = {
+ outputQueue.put(QueueItem.EndOfStream)
+ terminalLatch.countDown()
+ }
+
+ private def handleControl(ctrl: UdfControlResponse): Unit = ctrl.getControlCase match {
+ case UdfControlResponse.ControlCase.INIT =>
+ val resp = ctrl.getInit
+ if (resp.hasError) {
+ // Attribute a subsequent close()'s terminator to the init error rather
+ // than a bare Cancelled, mirroring the data-phase ERROR branch.
+ executionError.compareAndSet(None, Some(resp.getError))
+ // Settle Failed(err) before the Cancel below (not after its CancelResponse):
+ // the sticky terminal makes close() report the init error, and the Cancel
+ // becomes best-effort worker cleanup -- so if it can't be written (a
+ // directExecutor worker delivered this reentrantly, before doInit published
+ // requestObserver), doInit still doesn't wait terminalTimeoutMs for a
+ // CancelResponse that can never arrive. Settle-before-release (see initValue).
+ Transitions.failed(resp.getError)
+ initValue.complete(resp)
+ sendCancelInternal(() => cancelWithReason("init failed"))
+ } else {
+ // InitResponse OK. Only advance from Initializing so a terminal that raced
+ // in (e.g. a transport error) still wins; process() then opens the data
+ // phase. Settle-before-release (see initValue): publish after the CAS.
+ Transitions.initAccepted()
+ initValue.complete(resp)
+ }
+
+ case UdfControlResponse.ControlCase.ERROR =>
+ val err = ctrl.getError.getError
+ executionError.compareAndSet(None, Some(err))
+ // Before init resolves, settle Failed here rather than rely on the Cancel
+ // below: sendCancelInternal may find requestObserver still null (reentrant
+ // delivery inside stream.onNext, before doInit published it), so no Cancel --
+ // hence no CancelResponse -- would settle the terminal, and doInit would
+ // return as if init succeeded. completeTerminal is idempotent, so the
+ // data-phase ERROR -> Cancel -> CancelResponse path is unaffected.
+ if (!initResolved) {
+ Transitions.failed(err)
+ }
+ // Settle-before-release (see initValue). No value to publish -- ERROR is not
+ // an InitResponse; on the data-phase path the latch already fired in init.
+ initValue.signalWithoutValue()
+ sendCancelInternal(() => cancelWithReason("aborting after ErrorResponse"))
+
+ case UdfControlResponse.ControlCase.FINISH =>
+ // The FinishResponse carries metrics + the finish-callback data/error.
+ // Keep it on the terminal so close() can return it; the iterator inspects
+ // its error field to decide whether to throw.
+ Transitions.finished(ctrl.getFinish)
+ // Defensive: FINISH before InitResponse is a worker protocol bug, but
+ // we should fail init fast rather than hang the 30s init timeout.
+ initValue.signalWithoutValue()
+
+ case UdfControlResponse.ControlCase.CANCEL =>
+ // The CancelResponse carries metrics + the cancel-callback error. Keep it
+ // on the terminal so close() can return it; any prior ErrorResponse is
+ // tracked in executionError and surfaced by the iterator.
+ Transitions.cancelled(ctrl.getCancel)
+ // Defensive: CANCEL before InitResponse unblocks doInit so it can
+ // surface the cancellation instead of timing out.
+ initValue.signalWithoutValue()
+
+ case UdfControlResponse.ControlCase.CONTROL_NOT_SET =>
+ Transitions.transportFailed(new IllegalStateException(
+ "empty UdfControlResponse oneof"))
+ initValue.signalWithoutValue()
+ }
+
+ /**
+ * True once init is no longer pending -- i.e. the stream is past `Initializing`.
+ * Not "init succeeded": a terminal (including a failure) also counts as resolved.
+ */
+ private def initResolved: Boolean = currentState match {
+ case SessionState.Created | SessionState.Initializing => false
+ case _ => true
+ }
+
+ private def cancelWithReason(reason: String): Cancel =
+ Cancel.newBuilder().setReason(reason).build()
+
+ /**
+ * The protocol transition graph in one place: names for the edges, not a
+ * second source of truth. Every edge acts on the single
+ * [[WorkerSession.SessionState]] owned by the base via [[compareAndSetState]]
+ * (non-terminal) or [[completeTerminal]] (terminal), so the base's CAS is the
+ * only synchronization and a terminal that arrived first always wins (the
+ * non-terminal CASes fail against it; [[completeTerminal]] is first-wins).
+ *
+ * Edges this class drives -- edge, method, then driver site(s) / thread (the
+ * base drives the API-call edges: `Created -> Initializing` in `init`,
+ * `Initialized -> Streaming` in `process`):
+ * {{{
+ * Initializing -> Initialized initAccepted handleControl INIT-ok [gRPC cb]
+ * Streaming -> Finishing finishSent advance branch 3 [engine]
+ * non-terminal -> Cancelling cancelSentFrom sendCancelInternal, iff a Cancel
+ * reaches the wire [gRPC cb | engine | init | close]
+ * non-terminal -> Terminal finished/cancelled/failed/transportFailed
+ * handleControl / doInit / doClose / advance / onError
+ * }}}
+ * `Cancelling` is reached only if a Cancel is actually written; a pre-stream or
+ * raced cancel goes straight to a `Cancelled`/`TransportFailed` terminal (see
+ * [[sendCancelInternal]], [[doClose]], `ProcessIterator`) or nowhere -- so
+ * `cancelRequested` (intent) does not imply state `Cancelling`.
+ */
+ private object Transitions {
+ /** `InitResponse` OK: `Initializing -> Initialized`. */
+ def initAccepted(): Boolean =
+ compareAndSetState(SessionState.Initializing, SessionState.Initialized)
+
+ /** Input exhausted and `Finish` written: `Streaming -> Finishing` (once). */
+ def finishSent(): Boolean =
+ compareAndSetState(SessionState.Streaming, SessionState.Finishing)
+
+ /** `Cancel` written: `cur -> Cancelling`, from any non-terminal `cur`. */
+ def cancelSentFrom(cur: SessionState): Boolean =
+ !cur.isTerminal && compareAndSetState(cur, SessionState.Cancelling)
+
+ /** Clean terminal carrying the worker's `FinishResponse`. */
+ def finished(response: FinishResponse): Boolean =
+ completeTerminal(Termination.Finished(response))
+
+ /** Clean terminal carrying the worker's `CancelResponse`. */
+ def cancelled(response: CancelResponse): Boolean =
+ completeTerminal(Termination.Cancelled(response))
+
+ /** Failure terminal carrying a structured [[ExecutionError]]. */
+ def failed(error: ExecutionError): Boolean =
+ completeTerminal(Termination.Failed(error))
+
+ /** Failure terminal carrying a transport-level cause. */
+ def transportFailed(cause: Throwable): Boolean =
+ completeTerminal(Termination.TransportFailed(cause))
+ }
+
+ // ---- WorkerSession hooks ------------------------------------------------
+
+ override protected def doInit(message: Init): InitResponse = {
+ // 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")
+ Transitions.transportFailed(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. close()'s 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 {
+ // Reentrant delivery: a directExecutor worker (one whose gRPC callbacks run
+ // synchronously on the caller's thread) can deliver InitResponse -- and thus
+ // run responseObserver/handleControl -- from *inside* this stream.onNext,
+ // before requestObserver below is published. The base advanced
+ // Created -> Initializing before calling doInit, so that reentrant
+ // InitResponse still finds Initializing and advances to Initialized; the
+ // handleControl error paths guard the still-null requestObserver. This is
+ // the scenario the "before doInit published requestObserver" comments mean.
+ 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
+ Transitions.transportFailed(e)
+ // Surface as GrpcWorkerSessionException so the engine integration layer
+ // (which catches that type and wraps it) sees a uniform init-failure
+ // exception rather than the raw transport error.
+ throw new GrpcWorkerSessionException("UDF worker stream failed during init", e)
+ }
+
+ val responded = try {
+ initValue.await(initResponseTimeoutMs)
+ } catch {
+ case _: InterruptedException =>
+ Thread.currentThread().interrupt()
+ sendCancelInternal(() => cancelWithReason("interrupted during init"))
+ // Make sure close() does not block waiting for a terminator that
+ // will never arrive on this thread's behalf.
+ Transitions.transportFailed(
+ new InterruptedException("interrupted while waiting for InitResponse"))
+ throw new InterruptedException("interrupted while waiting for InitResponse")
+ }
+
+ if (!responded) {
+ sendCancelInternal(() => cancelWithReason("InitResponse timed out"))
+ val timeout = new TimeoutException(
+ s"timed out waiting for InitResponse after ${initResponseTimeoutMs}ms")
+ // Settle the terminal so a subsequent close() does not stall for a
+ // second `terminalTimeoutMs` waiting for a worker that already missed
+ // its init deadline.
+ Transitions.transportFailed(timeout)
+ // Surface as GrpcWorkerSessionException (carrying the timeout cause) so
+ // the engine integration layer that catches that type can wrap it.
+ throw new GrpcWorkerSessionException(
+ s"timed out waiting for InitResponse after ${initResponseTimeoutMs}ms", timeout)
+ }
+
+ initValue.get match {
+ case Some(resp) if resp.hasError =>
+ // 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()
Review Comment:
Since no terminal is settled early anymore, awaitTerminal() at L507 now genuinely blocks for the CancelResponse. The comment is accurate again.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57640] Implement gRPC WorkerSession and test it against an echo server. [spark]
Posted by "haiyangsun-db (via GitHub)" <gi...@apache.org>.
haiyangsun-db commented on code in PR #56702:
URL: https://github.com/apache/spark/pull/56702#discussion_r3665853092
##########
udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/GrpcWorkerSession.scala:
##########
@@ -0,0 +1,933 @@
+/*
+ * 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.grpc
+
+import java.util.concurrent.{CountDownLatch, LinkedBlockingQueue, TimeoutException, TimeUnit}
+import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference}
+
+import scala.util.control.NonFatal
+
+import io.grpc.{ConnectivityState, ManagedChannel}
+import io.grpc.stub.StreamObserver
+
+import org.apache.spark.annotation.Experimental
+import org.apache.spark.udf.worker.{Cancel, CancelResponse, DataRequest, DataResponse,
+ ExecutionError, Finish, FinishResponse, Init, InitResponse, UdfControlRequest,
+ UdfControlResponse, UdfRequest, UdfResponse, UdfWorkerGrpc}
+import org.apache.spark.udf.worker.core.{Termination, WorkerHandle, WorkerLogger, WorkerSession}
+import org.apache.spark.udf.worker.core.WorkerSession.SessionState
+import org.apache.spark.udf.worker.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_message.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.
+ *
+ * '''Driving model.''' Consumption-driven (Volcano / pull): the thread that
+ * consumes the [[doProcess]] result iterator is the one that pulls input and
+ * sends each `DataRequest`; the gRPC callback thread only receives output. It is
+ * pull-driven but not one-input-per-output -- `advance` sends the next input
+ * whenever the output queue is momentarily empty, so under async delivery it may
+ * push several input batches before any output is read (bounded by HTTP/2 flow
+ * control).
+ *
+ * '''State machine.''' This class does not keep its own state machine: it
+ * drives the single [[WorkerSession.SessionState]] owned by the base. The base
+ * advances `Created -> Initializing` (in `init`) and `Initialized -> Streaming`
+ * (in `process`); this class advances the protocol-event edges through
+ * [[compareAndSetState]] / [[completeTerminal]] as it exchanges messages:
+ * {{{
+ * Initializing --(InitResponse ok)--> Initialized [handleControl]
+ * Streaming ----(Finish written)----> Finishing [ProcessIterator]
+ * <any non-terminal> --(Cancel written)--> Cancelling [sendCancelInternal]
+ * <any non-terminal> --(terminator/error)--> terminal [completeTerminal]
+ * }}}
+ * The two clean terminals carry the worker's `FinishResponse` / `CancelResponse`
+ * (metrics + finish/cancel callback `data`/`error`) so [[close]] can return
+ * them. The only flag kept outside the machine is [[cancelRequested]]: a
+ * cancellation can be requested before the stream exists (so it cannot be a
+ * state transition yet), and it must both fast-fail the result iterator and
+ * suppress any in-flight Data/Finish.
+ *
+ * Threading:
+ * - [[doInit]] is synchronous: sends `Init` and blocks on `InitResponse`,
+ * returning it.
+ * - [[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.
+ * - [[doClose]] is thread-safe and idempotent: it settles + returns the
+ * terminator (cancelling in-flight work if the stream had not finished)
+ * and half-closes the request side.
+ *
+ * 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 a gRPC channel built and owned by the caller (the
+ * dispatcher). 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 from the worker, drained by the process() iterator.
+ // Intentionally unbounded: a bounded queue would block the gRPC callback
+ // (Netty event-loop) thread when full, stalling terminator/control delivery on
+ // the whole channel. HTTP/2 flow control bounds the wire and the consumer
+ // normally drains promptly; a stalled downstream can still grow it, but the fix
+ // is protocol-level back-pressure (out of scope), not bounding the queue.
+ //
+ // TODO [SPARK-57324]: expose queue depth as a metric (early warning for a
+ // stalled consumer).
+ private val outputQueue = new LinkedBlockingQueue[QueueItem]()
+
+ // The worker's `InitResponse` (success or error) coupled with the latch that
+ // init() blocks on. Fires when the InitResponse arrives (`complete`) or when a
+ // terminal settles first without one -- transport error, half-close, or a
+ // premature terminator (`signalWithoutValue`). Until it fires we have no proof
+ // the worker accepted the session.
+ //
+ // Settle-before-release rule (referenced from every callback that both settles
+ // a terminal and fires this latch): settle the terminal FIRST, then complete /
+ // signal. init() blocks on the latch, and only latch await/release establish a
+ // happens-before edge, so a woken init() is guaranteed to observe the terminal
+ // rather than a transient state. OneShotValue keeps that publish-then-release
+ // in one place instead of every caller remembering to count down after setting
+ // the reference.
+ private val initValue = new OneShotValue[InitResponse]
+
+ // Fired when the session reaches a terminal [[SessionState]]. doClose() and
+ // the init-error path block on this to drain the terminator.
+ private val terminalLatch = new CountDownLatch(1)
+
+ // Captures an ErrorResponse encountered during the data phase so that
+ // the CancelResponse terminator 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)
+
+ // Cancellation INTENT -- distinct from the `Cancelling` state, which is reached
+ // only once a Cancel is actually written to the wire ([[sendCancelInternal]]).
+ // Intent is set first and can outrun (or never reach) that write, so
+ // `cancelRequested` does NOT imply state `Cancelling`. Kept outside the
+ // [[SessionState]] machine because a cancel can be requested before the stream
+ // exists (pre-init), where there is no wire transition to make yet. Used to (a)
+ // make cancel idempotent across all call sites, (b) fast-fail ProcessIterator
+ // on a pre-init cancel, and (c) suppress any Data/Finish that would otherwise
+ // race a Cancel onto the wire (re-read inside [[requestLock]]).
+ private val cancelRequested = new AtomicBoolean(false)
+
+ // gRPC requires serialized writes to a request StreamObserver.
+ private val requestLock = new Object
+
+ // Initialised in init() -- before that, close() is a no-op 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 =>
+ // A DataResponse before InitResponse violates the protocol (InitResponse
+ // must precede any DataResponse): fast-fail rather than enqueue it and let
+ // init() block to its timeout. Settle-before-release (see initValue).
+ if (!initResolved) {
+ Transitions.transportFailed(new IllegalStateException(
+ "worker sent a DataResponse before InitResponse"))
+ initValue.signalWithoutValue()
+ } else {
+ outputQueue.put(QueueItem.Batch(response.getData))
+ }
+
+ case UdfResponse.ResponseCase.CONTROL =>
+ handleControl(response.getControl)
+
+ case other =>
+ // A malformed response (empty / unknown oneof) is a terminal transport
+ // failure; fast-fail init the same way. Settle-before-release (see initValue).
+ Transitions.transportFailed(new IllegalStateException(
+ s"unexpected response oneof: $other"))
+ initValue.signalWithoutValue()
+ }
+ }
+
+ override def onError(t: Throwable): Unit = {
+ // Transport-level failure: the stream is dead, no further writes possible.
+ // Settle-before-release (see initValue) so init() surfaces the transport
+ // cause instead of the initResponseTimeoutMs "timed out" error.
+ Transitions.transportFailed(t)
+ initValue.signalWithoutValue()
+ }
+
+ 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 (!currentState.isTerminal) {
+ Transitions.transportFailed(new IllegalStateException(
+ "worker response stream closed without a terminator"))
+ }
+ // Defensive: if onCompleted reached us before InitResponse, doInit is
+ // still blocked on initValue and would otherwise time out.
+ initValue.signalWithoutValue()
+ }
+ }
+
+ /**
+ * Wakes the result iterator (blocked on [[outputQueue]]) and any thread
+ * waiting on [[terminalLatch]] when the base settles a terminal. Invoked once,
+ * by the caller that wins [[completeTerminal]].
+ */
+ override protected def onTerminalSettled(termination: Termination): Unit = {
+ outputQueue.put(QueueItem.EndOfStream)
+ terminalLatch.countDown()
+ }
+
+ private def handleControl(ctrl: UdfControlResponse): Unit = ctrl.getControlCase match {
+ case UdfControlResponse.ControlCase.INIT =>
+ val resp = ctrl.getInit
+ if (resp.hasError) {
+ // Attribute a subsequent close()'s terminator to the init error rather
+ // than a bare Cancelled, mirroring the data-phase ERROR branch.
+ executionError.compareAndSet(None, Some(resp.getError))
+ // Settle Failed(err) before the Cancel below (not after its CancelResponse):
+ // the sticky terminal makes close() report the init error, and the Cancel
+ // becomes best-effort worker cleanup -- so if it can't be written (a
+ // directExecutor worker delivered this reentrantly, before doInit published
+ // requestObserver), doInit still doesn't wait terminalTimeoutMs for a
+ // CancelResponse that can never arrive. Settle-before-release (see initValue).
+ Transitions.failed(resp.getError)
Review Comment:
resolved
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57640] Implement gRPC WorkerSession and test it against an echo server. [spark]
Posted by "haiyangsun-db (via GitHub)" <gi...@apache.org>.
haiyangsun-db commented on code in PR #56702:
URL: https://github.com/apache/spark/pull/56702#discussion_r3666310353
##########
udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/GrpcWorkerSession.scala:
##########
@@ -0,0 +1,933 @@
+/*
+ * 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.grpc
+
+import java.util.concurrent.{CountDownLatch, LinkedBlockingQueue, TimeoutException, TimeUnit}
+import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference}
+
+import scala.util.control.NonFatal
+
+import io.grpc.{ConnectivityState, ManagedChannel}
+import io.grpc.stub.StreamObserver
+
+import org.apache.spark.annotation.Experimental
+import org.apache.spark.udf.worker.{Cancel, CancelResponse, DataRequest, DataResponse,
+ ExecutionError, Finish, FinishResponse, Init, InitResponse, UdfControlRequest,
+ UdfControlResponse, UdfRequest, UdfResponse, UdfWorkerGrpc}
+import org.apache.spark.udf.worker.core.{Termination, WorkerHandle, WorkerLogger, WorkerSession}
+import org.apache.spark.udf.worker.core.WorkerSession.SessionState
+import org.apache.spark.udf.worker.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_message.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.
+ *
+ * '''Driving model.''' Consumption-driven (Volcano / pull): the thread that
+ * consumes the [[doProcess]] result iterator is the one that pulls input and
+ * sends each `DataRequest`; the gRPC callback thread only receives output. It is
+ * pull-driven but not one-input-per-output -- `advance` sends the next input
+ * whenever the output queue is momentarily empty, so under async delivery it may
+ * push several input batches before any output is read (bounded by HTTP/2 flow
+ * control).
+ *
+ * '''State machine.''' This class does not keep its own state machine: it
+ * drives the single [[WorkerSession.SessionState]] owned by the base. The base
+ * advances `Created -> Initializing` (in `init`) and `Initialized -> Streaming`
+ * (in `process`); this class advances the protocol-event edges through
+ * [[compareAndSetState]] / [[completeTerminal]] as it exchanges messages:
+ * {{{
+ * Initializing --(InitResponse ok)--> Initialized [handleControl]
+ * Streaming ----(Finish written)----> Finishing [ProcessIterator]
+ * <any non-terminal> --(Cancel written)--> Cancelling [sendCancelInternal]
+ * <any non-terminal> --(terminator/error)--> terminal [completeTerminal]
+ * }}}
+ * The two clean terminals carry the worker's `FinishResponse` / `CancelResponse`
+ * (metrics + finish/cancel callback `data`/`error`) so [[close]] can return
+ * them. The only flag kept outside the machine is [[cancelRequested]]: a
+ * cancellation can be requested before the stream exists (so it cannot be a
+ * state transition yet), and it must both fast-fail the result iterator and
+ * suppress any in-flight Data/Finish.
+ *
+ * Threading:
+ * - [[doInit]] is synchronous: sends `Init` and blocks on `InitResponse`,
+ * returning it.
+ * - [[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.
+ * - [[doClose]] is thread-safe and idempotent: it settles + returns the
+ * terminator (cancelling in-flight work if the stream had not finished)
+ * and half-closes the request side.
+ *
+ * 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 a gRPC channel built and owned by the caller (the
+ * dispatcher). 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 from the worker, drained by the process() iterator.
+ // Intentionally unbounded: a bounded queue would block the gRPC callback
+ // (Netty event-loop) thread when full, stalling terminator/control delivery on
+ // the whole channel. HTTP/2 flow control bounds the wire and the consumer
+ // normally drains promptly; a stalled downstream can still grow it, but the fix
+ // is protocol-level back-pressure (out of scope), not bounding the queue.
+ //
+ // TODO [SPARK-57324]: expose queue depth as a metric (early warning for a
+ // stalled consumer).
+ private val outputQueue = new LinkedBlockingQueue[QueueItem]()
+
+ // The worker's `InitResponse` (success or error) coupled with the latch that
+ // init() blocks on. Fires when the InitResponse arrives (`complete`) or when a
+ // terminal settles first without one -- transport error, half-close, or a
+ // premature terminator (`signalWithoutValue`). Until it fires we have no proof
+ // the worker accepted the session.
+ //
+ // Settle-before-release rule (referenced from every callback that both settles
+ // a terminal and fires this latch): settle the terminal FIRST, then complete /
+ // signal. init() blocks on the latch, and only latch await/release establish a
+ // happens-before edge, so a woken init() is guaranteed to observe the terminal
+ // rather than a transient state. OneShotValue keeps that publish-then-release
+ // in one place instead of every caller remembering to count down after setting
+ // the reference.
+ private val initValue = new OneShotValue[InitResponse]
+
+ // Fired when the session reaches a terminal [[SessionState]]. doClose() and
+ // the init-error path block on this to drain the terminator.
+ private val terminalLatch = new CountDownLatch(1)
+
+ // Captures an ErrorResponse encountered during the data phase so that
+ // the CancelResponse terminator 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)
+
+ // Cancellation INTENT -- distinct from the `Cancelling` state, which is reached
+ // only once a Cancel is actually written to the wire ([[sendCancelInternal]]).
+ // Intent is set first and can outrun (or never reach) that write, so
+ // `cancelRequested` does NOT imply state `Cancelling`. Kept outside the
+ // [[SessionState]] machine because a cancel can be requested before the stream
+ // exists (pre-init), where there is no wire transition to make yet. Used to (a)
+ // make cancel idempotent across all call sites, (b) fast-fail ProcessIterator
+ // on a pre-init cancel, and (c) suppress any Data/Finish that would otherwise
+ // race a Cancel onto the wire (re-read inside [[requestLock]]).
+ private val cancelRequested = new AtomicBoolean(false)
+
+ // gRPC requires serialized writes to a request StreamObserver.
+ private val requestLock = new Object
+
+ // Initialised in init() -- before that, close() is a no-op 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 =>
+ // A DataResponse before InitResponse violates the protocol (InitResponse
+ // must precede any DataResponse): fast-fail rather than enqueue it and let
+ // init() block to its timeout. Settle-before-release (see initValue).
+ if (!initResolved) {
+ Transitions.transportFailed(new IllegalStateException(
+ "worker sent a DataResponse before InitResponse"))
+ initValue.signalWithoutValue()
+ } else {
+ outputQueue.put(QueueItem.Batch(response.getData))
+ }
+
+ case UdfResponse.ResponseCase.CONTROL =>
+ handleControl(response.getControl)
+
+ case other =>
+ // A malformed response (empty / unknown oneof) is a terminal transport
+ // failure; fast-fail init the same way. Settle-before-release (see initValue).
+ Transitions.transportFailed(new IllegalStateException(
+ s"unexpected response oneof: $other"))
+ initValue.signalWithoutValue()
+ }
+ }
+
+ override def onError(t: Throwable): Unit = {
+ // Transport-level failure: the stream is dead, no further writes possible.
+ // Settle-before-release (see initValue) so init() surfaces the transport
+ // cause instead of the initResponseTimeoutMs "timed out" error.
+ Transitions.transportFailed(t)
+ initValue.signalWithoutValue()
+ }
+
+ 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 (!currentState.isTerminal) {
+ Transitions.transportFailed(new IllegalStateException(
+ "worker response stream closed without a terminator"))
+ }
+ // Defensive: if onCompleted reached us before InitResponse, doInit is
+ // still blocked on initValue and would otherwise time out.
+ initValue.signalWithoutValue()
+ }
+ }
+
+ /**
+ * Wakes the result iterator (blocked on [[outputQueue]]) and any thread
+ * waiting on [[terminalLatch]] when the base settles a terminal. Invoked once,
+ * by the caller that wins [[completeTerminal]].
+ */
+ override protected def onTerminalSettled(termination: Termination): Unit = {
+ outputQueue.put(QueueItem.EndOfStream)
+ terminalLatch.countDown()
+ }
+
+ private def handleControl(ctrl: UdfControlResponse): Unit = ctrl.getControlCase match {
+ case UdfControlResponse.ControlCase.INIT =>
+ val resp = ctrl.getInit
+ if (resp.hasError) {
+ // Attribute a subsequent close()'s terminator to the init error rather
+ // than a bare Cancelled, mirroring the data-phase ERROR branch.
+ executionError.compareAndSet(None, Some(resp.getError))
+ // Settle Failed(err) before the Cancel below (not after its CancelResponse):
+ // the sticky terminal makes close() report the init error, and the Cancel
+ // becomes best-effort worker cleanup -- so if it can't be written (a
+ // directExecutor worker delivered this reentrantly, before doInit published
+ // requestObserver), doInit still doesn't wait terminalTimeoutMs for a
+ // CancelResponse that can never arrive. Settle-before-release (see initValue).
+ Transitions.failed(resp.getError)
+ initValue.complete(resp)
+ sendCancelInternal(() => cancelWithReason("init failed"))
+ } else {
+ // InitResponse OK. Only advance from Initializing so a terminal that raced
+ // in (e.g. a transport error) still wins; process() then opens the data
+ // phase. Settle-before-release (see initValue): publish after the CAS.
+ Transitions.initAccepted()
+ initValue.complete(resp)
+ }
+
+ case UdfControlResponse.ControlCase.ERROR =>
+ val err = ctrl.getError.getError
+ executionError.compareAndSet(None, Some(err))
+ // Before init resolves, settle Failed here rather than rely on the Cancel
+ // below: sendCancelInternal may find requestObserver still null (reentrant
+ // delivery inside stream.onNext, before doInit published it), so no Cancel --
+ // hence no CancelResponse -- would settle the terminal, and doInit would
+ // return as if init succeeded. completeTerminal is idempotent, so the
+ // data-phase ERROR -> Cancel -> CancelResponse path is unaffected.
+ if (!initResolved) {
+ Transitions.failed(err)
+ }
+ // Settle-before-release (see initValue). No value to publish -- ERROR is not
+ // an InitResponse; on the data-phase path the latch already fired in init.
+ initValue.signalWithoutValue()
+ sendCancelInternal(() => cancelWithReason("aborting after ErrorResponse"))
+
+ case UdfControlResponse.ControlCase.FINISH =>
+ // The FinishResponse carries metrics + the finish-callback data/error.
+ // Keep it on the terminal so close() can return it; the iterator inspects
+ // its error field to decide whether to throw.
+ Transitions.finished(ctrl.getFinish)
+ // Defensive: FINISH before InitResponse is a worker protocol bug, but
+ // we should fail init fast rather than hang the 30s init timeout.
+ initValue.signalWithoutValue()
+
+ case UdfControlResponse.ControlCase.CANCEL =>
+ // The CancelResponse carries metrics + the cancel-callback error. Keep it
+ // on the terminal so close() can return it; any prior ErrorResponse is
+ // tracked in executionError and surfaced by the iterator.
+ Transitions.cancelled(ctrl.getCancel)
+ // Defensive: CANCEL before InitResponse unblocks doInit so it can
+ // surface the cancellation instead of timing out.
+ initValue.signalWithoutValue()
+
+ case UdfControlResponse.ControlCase.CONTROL_NOT_SET =>
+ Transitions.transportFailed(new IllegalStateException(
+ "empty UdfControlResponse oneof"))
+ initValue.signalWithoutValue()
+ }
+
+ /**
+ * True once init is no longer pending -- i.e. the stream is past `Initializing`.
+ * Not "init succeeded": a terminal (including a failure) also counts as resolved.
+ */
+ private def initResolved: Boolean = currentState match {
+ case SessionState.Created | SessionState.Initializing => false
+ case _ => true
+ }
+
+ private def cancelWithReason(reason: String): Cancel =
+ Cancel.newBuilder().setReason(reason).build()
+
+ /**
+ * The protocol transition graph in one place: names for the edges, not a
+ * second source of truth. Every edge acts on the single
+ * [[WorkerSession.SessionState]] owned by the base via [[compareAndSetState]]
+ * (non-terminal) or [[completeTerminal]] (terminal), so the base's CAS is the
+ * only synchronization and a terminal that arrived first always wins (the
+ * non-terminal CASes fail against it; [[completeTerminal]] is first-wins).
+ *
+ * Edges this class drives -- edge, method, then driver site(s) / thread (the
+ * base drives the API-call edges: `Created -> Initializing` in `init`,
+ * `Initialized -> Streaming` in `process`):
+ * {{{
+ * Initializing -> Initialized initAccepted handleControl INIT-ok [gRPC cb]
+ * Streaming -> Finishing finishSent advance branch 3 [engine]
+ * non-terminal -> Cancelling cancelSentFrom sendCancelInternal, iff a Cancel
+ * reaches the wire [gRPC cb | engine | init | close]
+ * non-terminal -> Terminal finished/cancelled/failed/transportFailed
+ * handleControl / doInit / doClose / advance / onError
+ * }}}
+ * `Cancelling` is reached only if a Cancel is actually written; a pre-stream or
+ * raced cancel goes straight to a `Cancelled`/`TransportFailed` terminal (see
+ * [[sendCancelInternal]], [[doClose]], `ProcessIterator`) or nowhere -- so
+ * `cancelRequested` (intent) does not imply state `Cancelling`.
+ */
+ private object Transitions {
+ /** `InitResponse` OK: `Initializing -> Initialized`. */
+ def initAccepted(): Boolean =
+ compareAndSetState(SessionState.Initializing, SessionState.Initialized)
+
+ /** Input exhausted and `Finish` written: `Streaming -> Finishing` (once). */
+ def finishSent(): Boolean =
+ compareAndSetState(SessionState.Streaming, SessionState.Finishing)
+
+ /** `Cancel` written: `cur -> Cancelling`, from any non-terminal `cur`. */
+ def cancelSentFrom(cur: SessionState): Boolean =
+ !cur.isTerminal && compareAndSetState(cur, SessionState.Cancelling)
+
+ /** Clean terminal carrying the worker's `FinishResponse`. */
+ def finished(response: FinishResponse): Boolean =
+ completeTerminal(Termination.Finished(response))
+
+ /** Clean terminal carrying the worker's `CancelResponse`. */
+ def cancelled(response: CancelResponse): Boolean =
+ completeTerminal(Termination.Cancelled(response))
+
+ /** Failure terminal carrying a structured [[ExecutionError]]. */
+ def failed(error: ExecutionError): Boolean =
+ completeTerminal(Termination.Failed(error))
+
+ /** Failure terminal carrying a transport-level cause. */
+ def transportFailed(cause: Throwable): Boolean =
+ completeTerminal(Termination.TransportFailed(cause))
+ }
+
+ // ---- WorkerSession hooks ------------------------------------------------
+
+ override protected def doInit(message: Init): InitResponse = {
+ // 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")
+ Transitions.transportFailed(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. close()'s 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 {
+ // Reentrant delivery: a directExecutor worker (one whose gRPC callbacks run
+ // synchronously on the caller's thread) can deliver InitResponse -- and thus
+ // run responseObserver/handleControl -- from *inside* this stream.onNext,
+ // before requestObserver below is published. The base advanced
+ // Created -> Initializing before calling doInit, so that reentrant
+ // InitResponse still finds Initializing and advances to Initialized; the
+ // handleControl error paths guard the still-null requestObserver. This is
+ // the scenario the "before doInit published requestObserver" comments mean.
+ 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
+ Transitions.transportFailed(e)
+ // Surface as GrpcWorkerSessionException so the engine integration layer
+ // (which catches that type and wraps it) sees a uniform init-failure
+ // exception rather than the raw transport error.
+ throw new GrpcWorkerSessionException("UDF worker stream failed during init", e)
+ }
+
+ val responded = try {
+ initValue.await(initResponseTimeoutMs)
+ } catch {
+ case _: InterruptedException =>
+ Thread.currentThread().interrupt()
+ sendCancelInternal(() => cancelWithReason("interrupted during init"))
+ // Make sure close() does not block waiting for a terminator that
+ // will never arrive on this thread's behalf.
+ Transitions.transportFailed(
+ new InterruptedException("interrupted while waiting for InitResponse"))
+ throw new InterruptedException("interrupted while waiting for InitResponse")
Review Comment:
@Yicong-Huang In spark's case, this usually means query cancellation, that can originate from a end-user clicking cancel or interrupt the query, or AQE/Spark engine cancel a stage.
I'm trying to introduce a better and consistent way of handling interrupts here.
On the other hand, I would propose us to move faster with this first version, as we can evolve over time later.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57640] Implement gRPC WorkerSession and test it against an echo server. [spark]
Posted by "Yicong-Huang (via GitHub)" <gi...@apache.org>.
Yicong-Huang commented on code in PR #56702:
URL: https://github.com/apache/spark/pull/56702#discussion_r3707554551
##########
udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/GrpcWorkerSession.scala:
##########
@@ -0,0 +1,816 @@
+/*
+ * 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.grpc
+
+import java.util.concurrent.{CountDownLatch, LinkedBlockingQueue, TimeoutException, TimeUnit}
+import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference}
+
+import scala.util.control.NonFatal
+
+import io.grpc.{ConnectivityState, ManagedChannel}
+import io.grpc.stub.StreamObserver
+
+import org.apache.spark.annotation.Experimental
+import org.apache.spark.udf.worker.{Cancel, CancelResponse, DataRequest, DataResponse,
+ ExecutionError, Finish, Init, InitResponse, UdfControlRequest,
+ UdfControlResponse, UdfRequest, UdfResponse, UdfWorkerGrpc}
+import org.apache.spark.udf.worker.core.{Termination, WorkerHandle, WorkerLogger, WorkerSession}
+import org.apache.spark.udf.worker.core.WorkerSession.SessionState
+import org.apache.spark.udf.worker.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_message.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.
+ *
+ * '''Driving model.''' Consumption-driven (Volcano / pull): the thread that
+ * consumes the [[doProcess]] result iterator is the one that pulls the next
+ * input batch and sends its `DataRequest`. Nothing is sent until the iterator
+ * is consumed, and input flows one batch per output pull. The gRPC callback
+ * thread only receives output.
+ *
+ * '''State machine.''' This class does not keep its own state machine: it
+ * drives the single [[WorkerSession.SessionState]] owned by the base. The base
+ * advances `Created -> Initializing` (in `init`) and `Initialized -> Streaming`
+ * (in `process`); this class advances the protocol-event edges through
+ * [[compareAndSetState]] / [[completeTerminal]] as it exchanges messages:
+ * {{{
+ * Initializing --(InitResponse ok)--> Initialized [handleControl]
+ * Streaming ----(Finish written)----> Finishing [ProcessIterator]
+ * <any non-terminal> --(Cancel written)--> Cancelling [sendCancelInternal]
+ * <any non-terminal> --(terminator/error)--> terminal [completeTerminal]
+ * }}}
+ * The two clean terminals carry the worker's `FinishResponse` / `CancelResponse`
+ * (metrics + finish/cancel callback `data`/`error`) so [[close]] can return
+ * them. The only flag kept outside the machine is [[cancelRequested]]: a
+ * cancellation can be requested before the stream exists (so it cannot be a
+ * state transition yet), and it must both fast-fail the result iterator and
+ * suppress any in-flight Data/Finish.
+ *
+ * Threading:
+ * - [[doInit]] is synchronous: sends `Init` and blocks on `InitResponse`,
+ * returning it.
+ * - [[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.
+ * - [[doClose]] is thread-safe and idempotent: it settles + returns the
+ * terminator (cancelling in-flight work if the stream had not finished)
+ * and half-closes the request side.
+ *
+ * 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 a gRPC channel built and owned by the caller (the
+ * dispatcher). 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 from the worker, drained by the process() iterator.
+ // Intentionally unbounded: a bounded queue would block the gRPC callback
+ // (Netty event-loop) thread when full, stalling control-message and
+ // terminator delivery on the whole channel. HTTP/2 flow control bounds the
+ // wire and the consumer normally drains promptly, so it stays small. A
+ // stalled downstream can still grow it; the real fix is protocol-level
+ // back-pressure (out of scope here), not bounding the queue.
+ //
+ // TODO [SPARK-57324]: expose queue depth as a metric (early warning for a
+ // stalled consumer).
+ 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)
+
+ // The InitResponse the worker sent (success or error), captured so init()
+ // can return it. None until InitResponse arrives.
+ private val initResponse = new AtomicReference[Option[InitResponse]](None)
+
+ // Fired when the session reaches a terminal [[SessionState]]. doClose() and
+ // the init-error path block on this to drain the terminator.
+ private val terminalLatch = new CountDownLatch(1)
+
+ // Captures an ErrorResponse encountered during the data phase so that
+ // the CancelResponse terminator 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)
+
+ // Cancellation intent. Kept outside the [[SessionState]] machine because a
+ // cancel can be requested before the stream exists (pre-init), where there is
+ // no wire transition to make yet -- only an intent to record. Used to (a)
+ // make cancel idempotent across all call sites, (b) fast-fail ProcessIterator
+ // on a pre-init cancel, and (c) suppress any Data/Finish that would otherwise
+ // race a Cancel onto the wire (re-read inside [[requestLock]]).
+ private val cancelRequested = new AtomicBoolean(false)
+
+ // gRPC requires serialized writes to a request StreamObserver.
+ private val requestLock = new Object
+
+ // Initialised in init() -- before that, close() is a no-op 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))
+
+ case UdfResponse.ResponseCase.CONTROL =>
+ handleControl(response.getControl)
+
+ case other =>
+ // A malformed response (empty / unknown oneof) is a terminal
+ // transport failure. Count down initLatch like every other
+ // terminal-settling path here (onError / onCompleted / each
+ // handleControl branch): if this arrives before InitResponse, doInit
+ // must fail fast with this cause rather than block until
+ // initResponseTimeoutMs and report a misleading "timed out" error.
+ completeTerminal(Termination.TransportFailed(new IllegalStateException(
+ s"unexpected response oneof: $other")))
+ initLatch.countDown()
+ }
+ }
+
+ override def onError(t: Throwable): Unit = {
+ // Transport-level failure: the stream is dead. No further writes are
+ // possible (reaching a terminal state closes the write side). Settle the
+ // terminal BEFORE counting down initLatch: init() blocks on that latch,
+ // and only await/countDown establish a happens-before edge, so a thread
+ // woken by the countDown must already be able to observe the terminal.
+ // That lets init() surface the transport cause instead of timing out
+ // after initResponseTimeoutMs with a misleading "timed out" message --
+ // or, if it saw a transient non-terminal state, the defensive "init latch
+ // fired without an InitResponse or terminal" error.
+ completeTerminal(Termination.TransportFailed(t))
+ initLatch.countDown()
+ }
+
+ 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 (!currentState.isTerminal) {
+ completeTerminal(Termination.TransportFailed(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()
+ }
+ }
+
+ /**
+ * Wakes the result iterator (blocked on [[outputQueue]]) and any thread
+ * waiting on [[terminalLatch]] when the base settles a terminal. Invoked once,
+ * by the caller that wins [[completeTerminal]].
+ */
+ override protected def onTerminalSettled(termination: Termination): Unit = {
+ outputQueue.put(QueueItem.EndOfStream)
+ terminalLatch.countDown()
+ }
+
+ private def handleControl(ctrl: UdfControlResponse): Unit = ctrl.getControlCase match {
+ case UdfControlResponse.ControlCase.INIT =>
+ val resp = ctrl.getInit
+ // Capture the InitResponse so init() can return it (or throw on its error).
+ initResponse.set(Some(resp))
+ if (resp.hasError) {
+ initLatch.countDown()
+ // Per the protocol the engine follows an init error with Cancel; send it
+ // if the request stream is up. If the Cancel cannot be written -- e.g. a
+ // directExecutor worker delivered InitResponse reentrantly, before
+ // doInit published requestObserver -- settle the terminal here so doInit
+ // does not wait the full terminalTimeoutMs for a CancelResponse that can
+ // never arrive. (When the Cancel is sent, the CancelResponse settles the
+ // terminal instead.)
+ if (!sendCancelInternal(() => cancelWithReason("init failed")) &&
Review Comment:
Confirmed fixed: the INIT-error branch now records `executionError` before publishing the `InitResponse`, and `doInit` owns the Cancel plus `CancelResponse` drain, so the structured error reaches the caller and the terminal is the proto `Cancelled`.
##########
udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/GrpcWorkerSessionConcurrencySuite.scala:
##########
@@ -0,0 +1,742 @@
+/*
+ * 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.grpc
+
+import java.util.Locale
+import java.util.concurrent.{ConcurrentLinkedQueue, CountDownLatch, TimeUnit}
+import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference}
+
+import scala.jdk.CollectionConverters._
+
+import com.google.protobuf.ByteString
+import io.grpc.{ManagedChannel, Server}
+import io.grpc.inprocess.{InProcessChannelBuilder, InProcessServerBuilder}
+import io.grpc.stub.StreamObserver
+import org.scalatest.BeforeAndAfterEach
+// scalastyle:off funsuite
+import org.scalatest.funsuite.AnyFunSuite
+
+import org.apache.spark.udf.worker.{Cancel, CancelResponse, DataRequest, DataResponse,
+ ErrorResponse, ExecutionError, Finish, FinishResponse, Init, InitResponse, UdfControlResponse,
+ UdfPayload, UdfRequest, UdfResponse, UDFWorkerDataFormat, UdfWorkerGrpc, UserError,
+ WorkerRequest, WorkerResponse}
+import org.apache.spark.udf.worker.core.{Termination, WorkerHandle, WorkerLogger}
+
+/**
+ * Concurrency tests for [[GrpcWorkerSession]] that pin the wire-ordering and
+ * fast-fail invariants under concurrent and worker-misbehavior scenarios:
+ * - Cancel must never appear on the wire before Init.
+ * - A worker terminator (ERROR / FINISH / CANCEL / onCompleted) arriving
+ * before InitResponse must fail [[GrpcWorkerSession#init]] fast, not hang
+ * for `initResponseTimeoutMs`.
+ * - Repeated [[Iterator#hasNext]] after natural iterator exhaustion must
+ * return immediately, not block for `terminalTimeoutMs`.
+ * - Cancel issued before [[GrpcWorkerSession#init]] makes a subsequent
+ * [[GrpcWorkerSession#doProcess]] surface the cancellation immediately,
+ * instead of blocking for `terminalTimeoutMs`.
+ * - Cancel / close concurrent with an in-progress data phase terminate
+ * cleanly without leaks or unbounded hangs.
+ *
+ * Runs entirely in-process: no subprocess, no UDS. Server services are
+ * custom-built per test so we can drive specific worker misbehavior.
+ */
+class GrpcWorkerSessionConcurrencySuite
+ extends AnyFunSuite with BeforeAndAfterEach {
+// scalastyle:on funsuite
+
+ /** Used by tests to keep stale in-flight infra reachable for teardown. */
+ private val openServers = new ConcurrentLinkedQueue[Server]()
+ private val openChannels = new ConcurrentLinkedQueue[ManagedChannel]()
+ private val openSessions = new ConcurrentLinkedQueue[GrpcWorkerSession]()
+
+ override def afterEach(): Unit = {
+ // Shut channels down first. This fires onError on any still-live stream,
+ // which settles the session terminal and counts down the init/terminal
+ // latches. That unblocks both the session.close() below and any worker
+ // thread a failing test left parked on a (deliberately large) timeout, so
+ // teardown never hangs even when a test asserts via assertFinishesWithin.
+ openChannels.asScala.foreach { c =>
+ try c.shutdownNow().awaitTermination(2, TimeUnit.SECONDS) catch { case _: Throwable => () }
+ }
+ openChannels.clear()
+ openSessions.asScala.foreach { s => try s.close(emptyCancel) catch { case _: Throwable => () } }
+ openSessions.clear()
+ openServers.asScala.foreach { s =>
+ try s.shutdownNow().awaitTermination(2, TimeUnit.SECONDS) catch { case _: Throwable => () }
+ }
+ openServers.clear()
+ super.afterEach()
+ }
+
+ // A session timeout large enough that a correct test never reaches it; a
+ // regression that fails to short-circuit blocks here for minutes and is caught
+ // by assertFinishesWithin (below) instead of a flaky `elapsed < timeout` bound.
+ private val NeverReachedTimeoutMs = TimeUnit.MINUTES.toMillis(10)
+
+ /**
+ * Runs `body` on a daemon thread and asserts it finishes within `withinMs`,
+ * rethrowing whatever `body` threw (so an `intercept` inside `body` still
+ * works). Pair with [[NeverReachedTimeoutMs]]: the correct fast path returns
+ * in milliseconds, so `withinMs` (seconds) has an enormous safety margin and
+ * does not flake, while a regression that parks on the timeout never returns
+ * within `withinMs` and fails the assertion.
+ */
+ private def assertFinishesWithin(withinMs: Long, name: String)(body: => Unit): Unit = {
+ val thrown = new AtomicReference[Throwable]()
+ val done = new CountDownLatch(1)
+ val worker = new Thread(() => {
+ try body catch { case t: Throwable => thrown.set(t) } finally done.countDown()
+ }, name)
+ worker.setDaemon(true)
+ worker.start()
+ assert(done.await(withinMs, TimeUnit.MILLISECONDS),
+ s"$name did not finish within ${withinMs}ms; it parked on a session timeout " +
+ "that the fast path should have short-circuited")
+ Option(thrown.get()).foreach(t => throw t)
+ }
+
+ // ---------------------------------------------------------------------------
+ // Helpers
+ // ---------------------------------------------------------------------------
+
+ private class TestWorkerHandle extends WorkerHandle {
+ val invalidated = new AtomicBoolean(false)
Review Comment:
Confirmed: `released`/`invalidated` are now asserted across the clean-finish, TransportFailed, acknowledged-Cancelled and unacknowledged-Interrupted paths, including the idempotent repeat `close()`.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57640] Implement gRPC WorkerSession and test it against an echo server. [spark]
Posted by "haiyangsun-db (via GitHub)" <gi...@apache.org>.
haiyangsun-db commented on code in PR #56702:
URL: https://github.com/apache/spark/pull/56702#discussion_r3549402052
##########
udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/GrpcWorkerSession.scala:
##########
@@ -0,0 +1,816 @@
+/*
+ * 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.grpc
+
+import java.util.concurrent.{CountDownLatch, LinkedBlockingQueue, TimeoutException, TimeUnit}
+import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference}
+
+import scala.util.control.NonFatal
+
+import io.grpc.{ConnectivityState, ManagedChannel}
+import io.grpc.stub.StreamObserver
+
+import org.apache.spark.annotation.Experimental
+import org.apache.spark.udf.worker.{Cancel, CancelResponse, DataRequest, DataResponse,
+ ExecutionError, Finish, Init, InitResponse, UdfControlRequest,
+ UdfControlResponse, UdfRequest, UdfResponse, UdfWorkerGrpc}
+import org.apache.spark.udf.worker.core.{Termination, WorkerHandle, WorkerLogger, WorkerSession}
+import org.apache.spark.udf.worker.core.WorkerSession.SessionState
+import org.apache.spark.udf.worker.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_message.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.
+ *
+ * '''Driving model.''' Consumption-driven (Volcano / pull): the thread that
+ * consumes the [[doProcess]] result iterator is the one that pulls the next
+ * input batch and sends its `DataRequest`. Nothing is sent until the iterator
Review Comment:
yeah, that's correct. i can adjust the comment
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57640] Implement gRPC WorkerSession and test it against an echo server. [spark]
Posted by "Yicong-Huang (via GitHub)" <gi...@apache.org>.
Yicong-Huang commented on code in PR #56702:
URL: https://github.com/apache/spark/pull/56702#discussion_r3707545817
##########
udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/GrpcWorkerSession.scala:
##########
@@ -0,0 +1,1177 @@
+/*
+ * 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.grpc
+
+import java.util.Objects
+import java.util.concurrent.{CountDownLatch, LinkedBlockingQueue, TimeoutException, TimeUnit}
+import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference}
+
+import scala.util.control.NonFatal
+
+import io.grpc.{ConnectivityState, ManagedChannel}
+import io.grpc.stub.StreamObserver
+
+import org.apache.spark.annotation.Experimental
+import org.apache.spark.udf.worker.{Cancel, CancelResponse, DataRequest, DataResponse,
+ ExecutionError, Finish, FinishResponse, Init, InitResponse, UdfControlRequest,
+ UdfControlResponse, UdfPayload, UdfRequest, UdfResponse, UdfWorkerGrpc}
+import org.apache.spark.udf.worker.core.{Termination, WorkerHandle, WorkerLogger, WorkerSession}
+import org.apache.spark.udf.worker.core.WorkerSession.SessionState
+import org.apache.spark.udf.worker.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_message.proto` (`PayloadChunk*`
+ * between Init and InitResponse omitted here; chunking is not yet
+ * implemented -- see the TODO below):
+ * {{{
+ * Engine -> Worker: Init -> (DataRequest)* -> Finish (Cancel)?
+ * | Cancel
+ * Worker -> Engine: InitResponse -> (DataResponse)* ->
+ * (ErrorResponse)? -> (FinishResponse | CancelResponse)
+ * }}}
+ *
+ * Knows nothing about how the worker was provisioned -- the dispatcher
+ * constructs this with a [[WorkerHandle]] and channel; the base
+ * [[WorkerSession]] handles dispatcher-side cleanup on close.
+ *
+ * '''Driving model.''' Consumption-driven (Volcano / pull): the thread that
+ * consumes the [[doProcess]] result iterator is the one that pulls input and
+ * sends each `DataRequest`; the gRPC callback thread only receives output. It is
+ * pull-driven but not one-input-per-output -- `advance` sends the next input
+ * whenever the output queue is momentarily empty, so under async delivery it may
+ * push several input batches before any output is read. HTTP/2 flow control
+ * bounds wire traffic, but application-level buffering is intentionally left to
+ * a follow-up (see the TODO on [[outputQueue]]).
+ *
+ * '''State machine.''' This class does not keep its own state machine: it
+ * drives the single [[WorkerSession.SessionState]] owned by the base. The base
+ * advances `Created -> Initializing` (in `init`) and `Initialized -> Streaming`
+ * (in `process`); this class advances the protocol-event edges through
+ * [[compareAndSetState]] / [[completeTerminal]] as it exchanges messages:
+ * {{{
+ * Initializing --(InitResponse ok)--> Initialized [handleControl]
+ * Streaming ----(input exhausted)-------> Finishing [ProcessIterator]
+ * <any non-terminal> --(Cancel send starts)--> Cancelling [sendCancelInternal]
+ * <any non-terminal> --(terminator/error)--> terminal [completeTerminal]
+ * }}}
+ * The two clean terminals carry the worker's `FinishResponse` / `CancelResponse`
+ * (metrics + finish/cancel callback `data`/`error`) so [[close]] can return
+ * them. Cancellation intent is deliberately tracked outside the machine in
+ * [[cancelRequested]], which makes the wire write idempotent and suppresses any
+ * in-flight Data/Finish once cancellation begins.
+ *
+ * Threading:
+ * - [[doInit]] is synchronous: sends `Init` and blocks on `InitResponse`,
+ * returning it.
+ * - [[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.
+ * - [[doClose]] is thread-safe and idempotent: it settles + returns the
+ * terminator (cancelling in-flight work if the stream had not finished)
+ * and terminates the request side.
+ *
+ * TODO [SPARK-55278]: this class does not yet implement payload chunking;
+ * the entire [[UdfPayload]] 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 a gRPC channel built and owned by the caller (the
+ * dispatcher). 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`.
+ * @param interruptCancelTimeoutMs upper bound on the wait for a `CancelResponse`
+ * after an interrupt (cancelled query / killed
+ * task) sends `Cancel`. Short by design; on expiry
+ * the session settles `Interrupted`. See
+ * `handleInterrupt`.
+ */
+@Experimental
+class GrpcWorkerSession(
+ workerHandle: WorkerHandle,
+ channel: ManagedChannel,
+ logger: WorkerLogger = WorkerLogger.NoOp,
+ initResponseTimeoutMs: Long = DEFAULT_INIT_RESPONSE_TIMEOUT_MS,
+ terminalTimeoutMs: Long = DEFAULT_TERMINAL_TIMEOUT_MS,
+ interruptCancelTimeoutMs: Long = DEFAULT_INTERRUPT_CANCEL_TIMEOUT_MS)
+ extends WorkerSession(workerHandle, logger) {
+
+ require(channel != null, "channel is required")
+
+ private val asyncStub = UdfWorkerGrpc.newStub(channel)
+
+ // Output batches from the worker, drained by the process() iterator.
+ // Intentionally unbounded in this first implementation: a bounded queue would
+ // block the gRPC callback (Netty event-loop) thread when full, stalling
+ // terminator/control delivery on the whole channel. The consumer normally
+ // drains promptly, but HTTP/2 flow control alone does not bound this application
+ // queue or gRPC's asynchronous send buffers.
+ //
+ // TODO [SPARK-55278]: add application-level gRPC flow control in a follow-up,
+ // using ClientCallStreamObserver readiness for requests and manual inbound
+ // demand for responses, before wiring this transport into a production path.
+ // TODO [SPARK-57324]: expose queue depth as a metric (early warning for a
+ // stalled consumer).
+ private val outputQueue = new LinkedBlockingQueue[QueueItem]()
+
+ // This value couples the worker's `InitResponse` (success or error) with the
+ // latch on which init() blocks. The latch fires when the InitResponse arrives
+ // (`complete`), when a pre-init ErrorResponse arrives, or when a terminal
+ // settles first without an InitResponse (`signalWithoutValue`). Until it fires
+ // we have no proof the worker accepted the session.
+ //
+ // Settle-before-release rule (referenced from every callback that both settles
+ // a terminal and fires this latch): settle the terminal FIRST, then complete /
+ // signal. init() blocks on the latch, and only a latch await/release pair
+ // establishes a happens-before edge, so a woken init() is guaranteed to
+ // observe the terminal rather than a transient state. OneShotValue keeps that
+ // publish-then-release in one place instead of every caller remembering to
+ // count down after setting the reference.
+ private val initValue = new OneShotValue[InitResponse]
+
+ // Fired when the session reaches a terminal [[SessionState]]. doClose() and
+ // the init-error path block on this to drain the terminator.
+ private val terminalLatch = new CountDownLatch(1)
+
+ // Captures an ErrorResponse encountered during the data phase so that
+ // the CancelResponse terminator 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 immediately before Finish is handed to the request observer. It is set
+ // before onNext so a reentrant callback may deliver FinishResponse while the
+ // Finish request is still being written.
+ private val finishSendStarted = new AtomicBoolean(false)
+
+ // Cancellation INTENT -- distinct from the `Cancelling` state, which is reached
+ // once a Cancel send is attempted under [[requestLock]] ([[sendCancelInternal]]).
+ // Intent is set first and can outrun (or never reach) that attempt, so
+ // `cancelRequested` does NOT imply state `Cancelling`. It is kept outside the
+ // [[SessionState]] machine to (a) make cancellation idempotent across all call
+ // sites and (b) suppress any Data/Finish that would otherwise race a Cancel
+ // onto the wire (re-read inside [[requestLock]]).
+ private val cancelRequested = new AtomicBoolean(false)
+
+ // gRPC requires serialized writes to a request StreamObserver.
+ private val requestLock = new Object
+
+ // Initialized in init() -- before that, close() is a no-op on the request
+ // side, which is exactly the contract the wrapping WorkerSession expects.
+ @volatile private var requestObserver: StreamObserver[UdfRequest] = _
+
+ // True after the request side has been half-closed, aborted, or observed the
+ // transport closing. It prevents this session from issuing duplicate terminal
+ // calls, including following onError with onCompleted from close().
+ private val requestSideTerminated = new AtomicBoolean(false)
+
+ private val responseObserver: StreamObserver[UdfResponse] = new StreamObserver[UdfResponse] {
+ override def onNext(response: UdfResponse): Unit = {
+ response.getResponseCase match {
+ case UdfResponse.ResponseCase.DATA =>
+ // A DataResponse before InitResponse violates the protocol (InitResponse
+ // must precede any DataResponse): fast-fail rather than enqueue it and let
+ // init() block to its timeout. Settling the terminal wakes a blocked
+ // init() via onTerminalSettled.
+ if (!initResolved) {
+ Transitions.transportFailed(new IllegalStateException(
+ "worker sent a DataResponse before InitResponse"))
+ } else {
+ outputQueue.put(QueueItem.Batch(response.getData))
+ }
+
+ case UdfResponse.ResponseCase.CONTROL =>
+ handleControl(response.getControl)
+
+ case other =>
+ // A malformed response (empty / unknown oneof) is a terminal transport
+ // failure; fast-fail init the same way.
+ Transitions.transportFailed(new IllegalStateException(
+ s"unexpected response oneof: $other"))
+ }
+ }
+
+ override def onError(t: Throwable): Unit = {
+ // Transport-level failure: the stream is dead, no further writes possible.
+ // Settling the terminal wakes a blocked init() (via onTerminalSettled) so it
+ // surfaces the transport cause instead of the initResponseTimeoutMs error.
+ requestSideTerminated.set(true)
+ Transitions.transportFailed(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. Settling the terminal wakes a blocked
+ // init() via onTerminalSettled.
+ requestSideTerminated.set(true)
+ if (!currentState.isTerminal) {
+ Transitions.transportFailed(new IllegalStateException(
+ "worker response stream closed without a terminator"))
+ }
+ }
+ }
+
+ /**
+ * Wakes everything that can be blocked when the base settles a terminal:
+ * the result iterator (on [[outputQueue]]), a thread on [[terminalLatch]]
+ * (close), and a thread still blocked in [[doInit]] on [[initValue]]. Invoked
+ * once, by the caller that wins [[completeTerminal]], so this is the '''single'''
+ * place a settled terminal wakes a blocked [[doInit]] -- callers that settle a
+ * terminal (any response-callback failure/terminator branch, the timeout paths,
+ * or close()) do NOT signal [[initValue]] themselves; settling is enough.
+ * Settle-before-release (see [[initValue]]) holds because this runs after the
+ * terminal CAS in [[completeTerminal]]. The only direct [[initValue]] signals
+ * left are the non-terminal init paths ([[handleControl]]'s INIT-ok / INIT-error
+ * / pre-init-ERROR branches), which must wake [[doInit]] '''without''' settling a
+ * terminal. INIT-ok first advances the session to `Initialized`; the error
+ * branches remain `Initializing` only until [[doInit]] sends Cancel and throws.
+ */
+ override protected def onTerminalSettled(termination: Termination): Unit = {
+ outputQueue.put(QueueItem.EndOfStream)
+ initValue.signalWithoutValue()
+ terminalLatch.countDown()
+ }
+
+ private def handleControl(ctrl: UdfControlResponse): Unit = ctrl.getControlCase match {
+ case UdfControlResponse.ControlCase.INIT =>
+ val resp = ctrl.getInit
+ if (resp.hasError) {
+ // Record the error and publish the InitResponse; do NOT settle a terminal
+ // or send Cancel here. The proto requires the engine to Cancel after an
+ // init error (udf_message.proto); doInit owns that synchronous cleanup
+ // after the initial onNext returns, then drains the CancelResponse before
+ // throwing. Settle-before-release (see initValue): there is no terminal to
+ // settle first here, just publish.
+ executionError.compareAndSet(None, Some(resp.getError))
+ initValue.complete(resp)
+ } else {
+ // InitResponse OK. Only advance from Initializing so a terminal that raced
+ // in (e.g. a transport error) still wins; process() then opens the data
+ // phase. Settle-before-release (see initValue): publish after the CAS.
+ Transitions.initAccepted()
+ initValue.complete(resp)
+ }
+
+ case UdfControlResponse.ControlCase.ERROR =>
+ val err = ctrl.getError.getError
+ executionError.compareAndSet(None, Some(err))
+ if (!initResolved) {
+ // Pre-init ErrorResponse. Leave the session in Initializing and just wake
+ // init(): the recorded executionError tells doInit the worker failed, and
+ // doInit sends the Cancel and drains the CancelResponse after the initial
+ // onNext returns, just like the INIT-error branch.
+ initValue.signalWithoutValue()
+ } else {
+ // Data-phase ErrorResponse: requestObserver is published, so cancel here.
+ // Cancel -> CancelResponse settles the terminal; the iterator surfaces the
+ // recorded executionError. The init latch already fired in init(), so no
+ // signalWithoutValue is needed.
+ sendCancelInternal(() => cancelWithReason("aborting after ErrorResponse"))
+ }
+
+ case UdfControlResponse.ControlCase.FINISH =>
+ // The FinishResponse carries metrics + the finish-callback data/error.
+ // Keep it on the terminal so close() can return it; the iterator inspects
+ // its error field to decide whether to throw. A FinishResponse is valid only
+ // after the engine has started sending Finish; accepting one earlier could
+ // silently truncate input and report a clean result.
+ if (finishSendStarted.get()) {
+ Transitions.finished(ctrl.getFinish)
+ } else {
+ Transitions.transportFailed(new IllegalStateException(
+ "worker sent FinishResponse before the engine sent Finish"))
+ }
+
+ case UdfControlResponse.ControlCase.CANCEL =>
+ // The CancelResponse carries metrics + the cancel-callback error. Keep it
+ // on the terminal so close() can return it; any prior ErrorResponse is
+ // tracked in executionError and surfaced by the iterator. Settling the
+ // terminal wakes a blocked init() (a CANCEL before InitResponse) via
+ // onTerminalSettled.
+ Transitions.cancelled(ctrl.getCancel)
+
+ case UdfControlResponse.ControlCase.CONTROL_NOT_SET =>
+ // Settling the terminal wakes a blocked init() via onTerminalSettled.
+ Transitions.transportFailed(new IllegalStateException(
+ "empty UdfControlResponse oneof"))
+ }
+
+ /**
+ * True once init is no longer pending -- i.e. the stream is past `Initializing`.
+ * Not "init succeeded": a terminal (including a failure) also counts as resolved.
+ */
+ private def initResolved: Boolean = currentState match {
+ case SessionState.Created | SessionState.Initializing => false
+ case _ => true
+ }
+
+ private def cancelWithReason(reason: String): Cancel =
+ Cancel.newBuilder().setReason(reason).build()
+
+ /**
+ * The protocol transition graph in one place: names for the edges, not a
+ * second source of truth. Every edge acts on the single
+ * [[WorkerSession.SessionState]] owned by the base via [[compareAndSetState]]
+ * (non-terminal) or [[completeTerminal]] (terminal), so the base's CAS is the
+ * only synchronization and a terminal that arrived first always wins (the
+ * non-terminal CASes fail against it; [[completeTerminal]] is first-wins).
+ *
+ * Edges this class drives -- edge, method, then driver site(s) / thread (the
+ * base drives the API-call edges: `Created -> Initializing` in `init`,
+ * `Initialized -> Streaming` in `process`):
+ * {{{
+ * Initializing -> Initialized initAccepted handleControl INIT-ok [gRPC cb]
+ * Streaming -> Finishing beginFinish advance branch 3 [engine]
+ * non-terminal -> Cancelling beginCancelFrom sendCancelInternal, immediately
+ * before the Cancel send [gRPC cb | engine | init | close]
+ * non-terminal -> Terminal finished/cancelled/transportFailed/interrupted
+ * handleControl / doInit / doClose / advance / onError
+ * }}}
+ * `Cancelling` is reached only when a Cancel send starts; a pre-stream or
+ * raced cancel goes straight to a `Cancelled`/`TransportFailed` terminal (see
+ * [[sendCancelInternal]], [[doClose]], `ProcessIterator`) or nowhere -- so
+ * `cancelRequested` (intent) does not imply state `Cancelling`.
+ */
+ private object Transitions {
+ /** `InitResponse` OK: `Initializing -> Initialized`. */
+ def initAccepted(): Boolean =
+ compareAndSetState(SessionState.Initializing, SessionState.Initialized)
+
+ /** Input exhausted: `Streaming -> Finishing` (once). */
+ def beginFinish(): Boolean =
+ compareAndSetState(SessionState.Streaming, SessionState.Finishing)
+
+ /** `Cancel` send starts: `cur -> Cancelling`, from any non-terminal `cur`. */
+ def beginCancelFrom(cur: SessionState): Boolean =
+ !cur.isTerminal && compareAndSetState(cur, SessionState.Cancelling)
+
+ /** Clean terminal carrying the worker's `FinishResponse`. */
+ def finished(response: FinishResponse): Boolean =
+ completeTerminal(Termination.Finished(response))
+
+ /** Clean terminal carrying the worker's `CancelResponse`. */
+ def cancelled(response: CancelResponse): Boolean =
+ completeTerminal(Termination.Cancelled(response))
+
+ /** Failure terminal carrying a transport-level cause. */
+ def transportFailed(cause: Throwable): Boolean =
+ completeTerminal(Termination.TransportFailed(cause))
+
+ /**
+ * Terminal for an engine-thread interrupt (e.g. a Spark task kill). Distinct
+ * from [[transportFailed]] because the cause is an engine-side interrupt, but
+ * still unsalvageable without a worker acknowledgement. Settled by
+ * [[handleInterrupt]] only when the brief post-Cancel drain expires without a
+ * `CancelResponse`; if the worker acks in time the session settles a
+ * cooperative [[cancelled]] instead.
+ */
+ def interrupted(cause: Throwable): Boolean =
+ completeTerminal(Termination.Interrupted(cause))
+ }
+
+ // ---- WorkerSession hooks ------------------------------------------------
+
+ override protected def doInit(message: Init): InitResponse = {
+ // Construct the request before opening the RPC. A malformed Init must not
+ // leave an otherwise-unused Execute stream waiting for its first request.
+ val initRequest = UdfRequest.newBuilder()
+ .setControl(UdfControlRequest.newBuilder().setInit(message).build())
+ .build()
+
+ // 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")
+ Transitions.transportFailed(ex)
+ throw new GrpcWorkerSessionException("UDF worker channel is closed", ex)
+ }
+ try {
+ requestLock.synchronized {
+ // Serialize stream creation, publication, and Init with close/cancel. A
+ // close that wins the lock first prevents the RPC from opening; one that
+ // loses cannot release the worker until the initial write returns.
+ if (!currentState.isTerminal) {
+ val stream = asyncStub.execute(responseObserver)
+ requestObserver = stream
+ // Stream creation may invoke a response callback synchronously. Do not
+ // send Init if that callback has already settled a terminal.
+ if (!currentState.isTerminal) {
+ // With directExecutor, gRPC callbacks run synchronously on the caller's
+ // thread, so InitResponse can reach responseObserver/handleControl from
+ // *inside* this stream.onNext.
+ // requestObserver is already published, so an immediate ErrorResponse
+ // can write its required Cancel reentrantly without losing the request.
+ sendOnNext(initRequest)
+ }
+ }
+ }
+ } catch {
+ case NonFatal(e) =>
+ Transitions.transportFailed(e)
+ // Surface as GrpcWorkerSessionException so the engine integration layer
+ // (which catches that type and wraps it) sees a uniform init-failure
+ // exception rather than the raw transport error.
+ throw new GrpcWorkerSessionException("UDF worker stream failed during init", e)
+ }
+
+ try {
+ initValue.await(initResponseTimeoutMs)
+ } catch {
+ case _: InterruptedException =>
+ // Interrupt (cancelled query / killed task) while awaiting InitResponse:
+ // cooperatively Cancel with a bounded drain, settling Cancelled if the
+ // worker acks in time, else Interrupted. See handleInterrupt.
+ handleInterrupt("waiting for InitResponse")
+ throw new InterruptedException("interrupted while waiting for InitResponse")
+ case e: TimeoutException =>
+ sendCancelInternal(() => cancelWithReason("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.
+ Transitions.transportFailed(e)
+ // Surface as GrpcWorkerSessionException (carrying the timeout cause) so
+ // the engine integration layer that catches that type can wrap it.
+ throw new GrpcWorkerSessionException(
+ s"timed out waiting for InitResponse after ${initResponseTimeoutMs}ms", e)
+ }
+
+ initValue.get match {
+ case Some(resp) if resp.hasError =>
+ failInitWithError(resp.getError,
+ s"UDF worker init failed: ${describeError(resp.getError)}")
+ case Some(resp) =>
+ executionError.get() match {
+ case Some(err) =>
+ // A reentrant callback can report an ErrorResponse immediately after
+ // InitResponse and before the initial onNext returns. The callback
+ // has already sent Cancel; drain its response
+ // and surface the original error instead of returning init success
+ // for a session that is already cancelling or terminal.
+ failInitWithError(err,
+ s"UDF worker reported an error as init completed: ${describeError(err)}")
+ case None if currentState != SessionState.Initialized =>
+ // Likewise, an immediate protocol failure or concurrent close may
+ // have moved the session out of Initialized after publishing the
+ // InitResponse. A normal return would violate init()'s contract and
+ // leave process() to fail later with only an ordering error.
+ failInitFromCurrentState()
+ case None =>
+ resp
+ }
+ case None =>
+ // No InitResponse arrived but the latch fired.
+ //
+ // A pre-init ErrorResponse leaves the session in Initializing with the
+ // error recorded in executionError (see handleControl ERROR branch): the
+ // engine must now Cancel and drain the CancelResponse, which failInit does.
+ executionError.get() match {
+ case Some(err) if !currentState.isTerminal =>
+ failInitWithError(err,
+ s"UDF worker reported an error before init completed: ${describeError(err)}")
+ case _ =>
+ // Otherwise a terminal already settled -- the worker terminated the
+ // stream before sending InitResponse (transport error, half-close, or
+ // a premature FinishResponse/CancelResponse). Surface it as an init
+ // failure rather than letting the caller proceed as if init succeeded.
+ failInitFromCurrentState()
+ }
+ }
+ }
+
+ /** Surfaces the terminal that prevented init from completing normally. */
+ private def failInitFromCurrentState(): Nothing = {
+ executionError.get() match {
+ case Some(err) =>
+ failInitWithError(err,
+ s"UDF worker reported an error as init completed: ${describeError(err)}")
+ case None => ()
+ }
+ // Cancelling can be observed after a concurrent close or a reentrant error
+ // before its CancelResponse arrives. Drain it so the exception reflects the
+ // stable terminal rather than a transient state.
+ if (currentState == SessionState.Cancelling) {
+ awaitTerminal()
+ }
+ currentState match {
+ case SessionState.Terminal(Termination.TransportFailed(cause)) =>
+ throw new GrpcWorkerSessionException("UDF worker stream failed during init", cause)
+ case SessionState.Terminal(Termination.Failed(err)) =>
+ throw new GrpcWorkerSessionException(
+ s"UDF worker reported an error before init completed: ${describeError(err)}", err)
+ case SessionState.Terminal(Termination.Cancelled(_)) =>
+ throw new GrpcWorkerSessionException(
+ "UDF worker stream was cancelled before init completed")
+ case SessionState.Terminal(Termination.Interrupted(cause)) =>
+ throw new GrpcWorkerSessionException("UDF worker init was interrupted", cause)
+ case SessionState.Terminal(Termination.Finished(_)) =>
+ throw new GrpcWorkerSessionException("UDF worker finished before init completed")
+ case other =>
+ throw new IllegalStateException(
+ s"init completed without an accepted session or terminal: $other")
+ }
+ }
+
+ /**
+ * Fails init when the worker reports an error before init returns: an
+ * `InitResponse` carrying an error, a pre-init `ErrorResponse`, or an immediate
+ * post-init `ErrorResponse`. The protocol requires the engine to send `Cancel`
+ * and the worker to reply with `CancelResponse` (udf_message.proto); we send it,
+ * drain the terminator so no stream is left dangling, and throw the structured
+ * error. A responsive worker settles `Cancelled`; a failed drain settles
+ * `TransportFailed`.
+ */
+ private def failInitWithError(err: ExecutionError, message: String): Nothing = {
+ sendCancelInternal(() => cancelWithReason("init failed"))
+ awaitTerminal()
+ throw new GrpcWorkerSessionException(message, err)
+ }
+
+ override protected def doProcess(
+ input: Iterator[DataRequest],
+ finish: () => Finish): Iterator[DataResponse] = {
+ // Init success is guaranteed by the base [[WorkerSession]] lifecycle: if
+ // doInit had failed it would have thrown and process() would never run.
+ new ProcessIterator(input, finish)
+ }
+
+ override protected def doClose(cancel: () => Cancel): Termination = {
+ // Coordinate the no-stream decision with doInit's publication + Init write.
+ // If close wins this lock, doInit observes the terminal and never sends Init;
+ // if doInit wins, close observes the published observer and sends Cancel only
+ // after the Init send has returned.
+ val noPublishedStream = requestLock.synchronized {
+ if (requestObserver == null) {
+ // init() never put a stream on the wire (closed before/around init, or
+ // init threw before publishing). There is no protocol terminator; if no
+ // terminal has settled yet, treat the session as cancelled-before-start.
+ // A terminal may already be settled here (e.g. the channel-shutdown
+ // TransportFailed in doInit also leaves requestObserver null); return the
+ // settled terminal as-is rather than a bare Cancelled that disagrees with
+ // the state. The base WorkerSession still releases the worker handle.
+ if (!currentState.isTerminal) {
+ Transitions.cancelled(CancelResponse.getDefaultInstance)
+ }
+ true
+ } else {
+ false
+ }
+ }
+ if (noPublishedStream) {
+ return settledTermination
+ }
+ // If the stream has not finished on its own, cancel anything in flight so
+ // the worker can clean up, then wait for the terminator. sendCancelInternal
+ // evaluates the cancel thunk only when it attempts a Cancel, and at
+ // most once across all callers.
+ if (!currentState.isTerminal) {
+ sendCancelInternal(cancel)
+ try {
+ terminalLatch.await(terminalTimeoutMs, TimeUnit.MILLISECONDS)
+ } catch {
+ case _: InterruptedException =>
+ // Interrupted mid-close: settle Interrupted via the
+ // shared handler rather than falling through to the TransportFailed
+ // guard below. The Cancel was already attempted, so handleInterrupt's
+ // sendCancelInternal is a no-op and it just runs the bounded drain.
+ //
+ // TODO [SPARK-57640]: close() is a finalizer often already on the task
+ // kill path, yet this makes an interrupted close block up to another
+ // interruptCancelTimeoutMs (chasing a clean, salvageable Cancelled
+ // rather than immediately unwinding to Interrupted). Current choice:
+ // spend the bounded wait to obtain proof the worker is recyclable; it
+ // is small next to terminalTimeoutMs. Revisit if killed-task teardown
+ // latency (especially many sessions torn down at once) makes an
+ // immediate unwind preferable -- e.g. skip this second drain in the
+ // close() path, since close already gave the worker its terminalTimeoutMs window.
+ handleInterrupt("closing the session")
+ }
+ }
+ // If the worker still has not settled (timeout above, or an interrupt whose
+ // bounded drain did not settle a terminal), record a terminal so
+ // isWorkerSalvageable returns a definite answer and any other thread reading
+ // the state sees a stable value.
+ if (!currentState.isTerminal) {
+ Transitions.transportFailed(new TimeoutException(
+ s"timed out waiting for stream terminator after ${terminalTimeoutMs}ms"))
+ }
+ // Close the request side according to the settled outcome. Clean protocol
+ // terminators are half-closed; failures abort the RPC so close() never follows
+ // an onError with onCompleted.
+ currentState match {
+ case SessionState.Terminal(Termination.Finished(_)) |
+ SessionState.Terminal(Termination.Cancelled(_)) =>
+ completeRequestSide()
+ case SessionState.Terminal(Termination.Failed(error)) =>
+ abortRequestSide(new GrpcWorkerSessionException(
+ s"UDF execution failed: ${describeError(error)}", error))
+ case SessionState.Terminal(Termination.TransportFailed(cause)) =>
+ abortRequestSide(cause)
+ case SessionState.Terminal(Termination.Interrupted(cause)) =>
+ abortRequestSide(cause)
+ case other =>
+ logger.debug(s"UDF Execute stream closed without a terminal outcome: $other")
+ }
+ // Derive the Termination from the settled terminal. close() is the
+ // finalizer/cleanup path and must NOT re-throw: a UDF / data-phase error is
+ // already surfaced while the result iterator is consumed, and an init error
+ // is surfaced from init(). settledTermination returns the settled terminal
+ // as-is: the response proto for the clean terminators, and the failure
+ // terminals (Failed / TransportFailed / Interrupted) carrying their cause --
+ // not a bare Cancelled.
+ settledTermination
+ }
+
+ // ---- Internal request helpers ---------------------------------------------
+
+ /**
+ * Writes one request and aborts the request side if the observer rejects it.
+ * Callers hold [[requestLock]], so this also serializes the compensating
+ * `onError` with every other request-side operation.
+ */
+ private def sendOnNext(req: UdfRequest): Unit = {
+ try {
+ requestObserver.onNext(req)
+ } catch {
+ case NonFatal(e) =>
+ abortRequestSide(e)
+ throw e
+ }
+ }
+
+ /** Aborts the request side at most once. Safe to call while holding [[requestLock]]. */
+ private def abortRequestSide(cause: Throwable): Unit = requestLock.synchronized {
+ if (requestObserver != null && requestSideTerminated.compareAndSet(false, true)) {
+ try {
+ requestObserver.onError(cause)
+ } catch {
+ case NonFatal(e) => logger.debug("Error aborting UDF Execute stream", e)
+ }
+ }
+ }
+
+ /** Half-closes the request side at most once. */
+ private def completeRequestSide(): Unit = requestLock.synchronized {
+ if (requestObserver != null && requestSideTerminated.compareAndSet(false, true)) {
+ try {
+ requestObserver.onCompleted()
+ } catch {
+ case NonFatal(e) => logger.debug("Error half-closing UDF Execute stream", e)
+ }
+ }
+ }
+
+ /**
+ * Sends a Data or Finish request to the worker. Three invariants are
+ * checked inside [[requestLock]]:
+ * - terminal state: writes are unsafe (transport dead or terminal
+ * received). Throws, so the caller's terminal/exception path runs.
+ * - request-side termination: no write may follow `onError` / `onCompleted`.
+ * Throws so the caller settles or surfaces the failure.
+ * - [[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]] publishes
+ * [[requestObserver]] and writes Init together under [[requestLock]], so a
+ * concurrent cancel observes the observer but cannot acquire the lock and send
+ * Cancel until after Init.
+ */
+ private def sendRequest(req: UdfRequest): Unit =
+ requestLock.synchronized {
+ if (currentState.isTerminal) {
+ throw new IllegalStateException(
+ "cannot send request: UDF Execute stream is already closed")
+ }
+ if (requestSideTerminated.get()) {
+ throw new IllegalStateException(
+ "cannot send request: UDF Execute request stream is already closed")
+ }
+ // Suppress the write when a cancel has been (or is about to be) flushed
+ // through this lock; that preserves the proto ordering invariant (no
+ // Data/Finish after Cancel). Test the flag rather than `return`-ing from
+ // this by-name `synchronized` body: a non-local return compiles to a thrown
+ // NonLocalReturnControl, which is brittle (also relied on in sendCancelInternal).
+ if (!cancelRequested.get()) {
+ if (req.hasControl && req.getControl.hasFinish) {
+ // Set before onNext: directExecutor delivery may return FinishResponse
+ // reentrantly from inside this call.
+ finishSendStarted.set(true)
+ }
+ sendOnNext(req)
+ }
+ }
+
+ /**
+ * Sends a `Cancel` control message, returning `true` iff the request observer
+ * accepts the `Cancel`. Idempotent across ALL call sites (close()'s
+ * cancel, in-band cancels from `handleControl`'s INIT error / ERROR paths,
+ * and engine-internal cancels from the iterator) via [[cancelRequested]]:
+ * only the first caller attempts a `Cancel`, and the `cancel` thunk is evaluated
+ * only by that caller and only when the send is about to start --
+ * so a side-effecting thunk (e.g. one carrying a client cancel callback) runs
+ * at most once. 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(cancel: () => Cancel): Boolean = {
+ // Do not consume the one-shot cancellation intent until there is an observer
+ // that can carry it. doInit publishes the observer before sending Init while
+ // holding requestLock, so every response to Init sees a non-null observer.
+ if (requestObserver == null) return false
+ if (!cancelRequested.compareAndSet(false, true)) return false
+ if (currentState.isTerminal) return false
+ try {
+ requestLock.synchronized {
+ // A terminal that arrived first must win; bail without writing. Compute
+ // the block's value rather than `return` (see sendRequest for why non-local
+ // return is avoided -- here it would also escape the NonFatal catch below).
+ val cur = currentState
+ if (cur.isTerminal || requestSideTerminated.get()) {
+ false
+ } else {
+ val request = UdfRequest.newBuilder()
+ .setControl(UdfControlRequest.newBuilder().setCancel(cancel()).build())
+ .build()
+ // Record that a Cancel send has started by advancing to Cancelling.
+ Transitions.beginCancelFrom(cur)
+ sendOnNext(request)
+ true
+ }
+ }
+ } catch {
+ case NonFatal(e) =>
+ logger.debug(s"Cancel send failed (stream may already be torn down): ${e.getMessage}")
+ Transitions.transportFailed(e)
+ false
+ }
+ }
+
+ private def awaitTerminal(): Unit = {
+ if (currentState.isTerminal) return
+ try {
+ if (!terminalLatch.await(terminalTimeoutMs, TimeUnit.MILLISECONDS)) {
+ Transitions.transportFailed(new TimeoutException(
+ s"timed out waiting for stream terminator after ${terminalTimeoutMs}ms"))
+ }
+ } catch {
+ case _: InterruptedException => handleInterrupt("waiting for stream terminator")
+ }
+ }
+
+ /**
+ * Handles an [[InterruptedException]] observed while blocked on a worker event
+ * (init, terminal drain, close, or the result-iterator poll). An interrupt is
+ * an engine-side event -- typically a cancelled query / killed task -- not a
+ * worker fault, so we cancel cooperatively and keep the worker salvageable only
+ * when it acks:
+ *
+ * 1. Send a best-effort `Cancel` (idempotent across call sites).
+ * 2. Wait up to [[interruptCancelTimeoutMs]] -- short, so the interrupted
+ * thread unwinds promptly -- for the worker's `CancelResponse`. A healthy
+ * worker acks in that window and the callback settles a clean `Cancelled`
+ * terminal (worker salvageable, response drained).
+ * 3. If the ack does not arrive in time, settle `Interrupted`. Without a
+ * terminator or liveness proof the worker is not salvageable and must not
+ * be returned to a reuse pool.
+ *
+ * The bounded wait deliberately runs '''before''' the thread's interrupt flag
+ * is restored: [[InterruptedException]] clears the flag on throw, so re-setting
+ * it first would make the wait below throw immediately and defeat the drain.
+ * The flag is restored at the end so the caller's unwind still observes the
+ * interrupt.
+ *
+ * TODO [SPARK-57640]: add a worker liveness/heartbeat check so a future version
+ * may prove that an interrupted worker which missed the short ack window is
+ * nevertheless safe to recycle.
+ */
+ private def handleInterrupt(waitContext: String): Unit = {
+ // NB: flag is currently clear (InterruptedException cleared it); do not
+ // restore it until after the bounded drain below.
+ if (!currentState.isTerminal) {
+ sendCancelInternal(() => cancelWithReason(s"interrupted while $waitContext"))
+ val acked = try {
+ terminalLatch.await(interruptCancelTimeoutMs, TimeUnit.MILLISECONDS)
+ } catch {
+ case _: InterruptedException =>
+ // A second interrupt during the drain: give up the wait immediately.
+ false
+ }
+ if (!acked && !currentState.isTerminal) {
+ // No CancelResponse in time: settle Interrupted rather than
+ // TransportFailed, and stop a later close() from blocking on a terminator
+ // that will not be drained on this thread's behalf. Interrupted remains
+ // unsalvageable until a separate liveness check can prove otherwise.
+ Transitions.interrupted(
+ new InterruptedException(s"interrupted while $waitContext"))
+ }
+ }
+ Thread.currentThread().interrupt()
+ }
+
+ // ---- 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 settle the state) and with any thread that finalizes via
+ * [[close]]. It drives the `Streaming -> Finishing` transition (sending
+ * `Finish` once input is exhausted, built lazily from the `finish` thunk).
+ */
+ private class ProcessIterator(input: Iterator[DataRequest], finish: () => Finish)
+ extends Iterator[DataResponse] {
+
+ // 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 branches, and block branch 4 for the
+ // full terminalTimeoutMs before returning. Callers that probe hasNext
+ // an extra time (iterator.size, instrumentation wrappers) would hang.
+ // Iterator-local and only touched by the single engine thread.
+ private val exhausted = new AtomicBoolean(false)
+ @volatile private var prefetched: DataResponse = _
+
+ /**
+ * Runs a caller-supplied thunk (`input.hasNext`, `input.next()`, or the
+ * `finish` builder) that may throw, and on failure Cancels the in-flight
+ * stream before rethrowing. Without the Cancel a failure would leave the
+ * worker awaiting more input with no terminator owed on the wire; with it the
+ * worker tears down and the exception still propagates to the engine. Setting
+ * `cancelRequested` also suppresses any further Data/Finish this loop might
+ * otherwise attempt (see [[sendRequest]]).
+ */
+ private def cancelOnThrow[T](reason: String)(op: => T): T =
+ try op catch {
+ case NonFatal(e) =>
+ sendCancelInternal(() => cancelWithReason(reason))
Review Comment:
`cancelOnThrow` catches `NonFatal`, which does not match `InterruptedException` -- so this is the one blocking point that never reaches `handleInterrupt`.
When a task is killed while the engine thread is parked inside `input.next()` (a shuffle read, say), the raw `InterruptedException` propagates straight out of `advance()`. The engine's `finally` calls `close()`, `doClose` sees a non-terminal state, sends `Cancel`, and blocks on `terminalLatch.await(terminalTimeoutMs)`. The throw already cleared the interrupt flag, so the `InterruptedException` short-circuit at L604 does not fire either: against an unresponsive worker, teardown burns the full 30s `terminalTimeoutMs` and settles `TransportFailed`, instead of the 2s `interruptCancelTimeoutMs` `Interrupted` unwind this PR adds. The init wait (L453), the output poll (L970), and `close` (L604) all handle it; only the input pull is missed.
I'd add an interrupt branch here rather than widening `NonFatal`, so the cancel reason stays accurate and the bounded drain runs:
```suggestion
private def cancelOnThrow[T](reason: String)(op: => T): T =
try op catch {
case _: InterruptedException =>
handleInterrupt(reason)
throw new InterruptedException(s"interrupted while $reason")
case NonFatal(e) =>
sendCancelInternal(() => cancelWithReason(reason))
throw e
}
```
Worth a regression test alongside the two `interrupt during init` cases in `GrpcWorkerSessionConcurrencySuite.scala` -- one that interrupts a thread blocked in `input.next()` and asserts the terminal settles within `interruptCancelTimeoutMs`. That is the gap that let this through: the suite covers interrupts on the init wait in both directions but never on the data phase.
##########
udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/GrpcWorkerSession.scala:
##########
@@ -0,0 +1,1177 @@
+/*
+ * 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.grpc
+
+import java.util.Objects
+import java.util.concurrent.{CountDownLatch, LinkedBlockingQueue, TimeoutException, TimeUnit}
+import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference}
+
+import scala.util.control.NonFatal
+
+import io.grpc.{ConnectivityState, ManagedChannel}
+import io.grpc.stub.StreamObserver
+
+import org.apache.spark.annotation.Experimental
+import org.apache.spark.udf.worker.{Cancel, CancelResponse, DataRequest, DataResponse,
+ ExecutionError, Finish, FinishResponse, Init, InitResponse, UdfControlRequest,
+ UdfControlResponse, UdfPayload, UdfRequest, UdfResponse, UdfWorkerGrpc}
+import org.apache.spark.udf.worker.core.{Termination, WorkerHandle, WorkerLogger, WorkerSession}
+import org.apache.spark.udf.worker.core.WorkerSession.SessionState
+import org.apache.spark.udf.worker.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_message.proto` (`PayloadChunk*`
+ * between Init and InitResponse omitted here; chunking is not yet
+ * implemented -- see the TODO below):
+ * {{{
+ * Engine -> Worker: Init -> (DataRequest)* -> Finish (Cancel)?
+ * | Cancel
+ * Worker -> Engine: InitResponse -> (DataResponse)* ->
+ * (ErrorResponse)? -> (FinishResponse | CancelResponse)
+ * }}}
+ *
+ * Knows nothing about how the worker was provisioned -- the dispatcher
+ * constructs this with a [[WorkerHandle]] and channel; the base
+ * [[WorkerSession]] handles dispatcher-side cleanup on close.
+ *
+ * '''Driving model.''' Consumption-driven (Volcano / pull): the thread that
+ * consumes the [[doProcess]] result iterator is the one that pulls input and
+ * sends each `DataRequest`; the gRPC callback thread only receives output. It is
+ * pull-driven but not one-input-per-output -- `advance` sends the next input
+ * whenever the output queue is momentarily empty, so under async delivery it may
+ * push several input batches before any output is read. HTTP/2 flow control
+ * bounds wire traffic, but application-level buffering is intentionally left to
+ * a follow-up (see the TODO on [[outputQueue]]).
+ *
+ * '''State machine.''' This class does not keep its own state machine: it
+ * drives the single [[WorkerSession.SessionState]] owned by the base. The base
+ * advances `Created -> Initializing` (in `init`) and `Initialized -> Streaming`
+ * (in `process`); this class advances the protocol-event edges through
+ * [[compareAndSetState]] / [[completeTerminal]] as it exchanges messages:
+ * {{{
+ * Initializing --(InitResponse ok)--> Initialized [handleControl]
+ * Streaming ----(input exhausted)-------> Finishing [ProcessIterator]
+ * <any non-terminal> --(Cancel send starts)--> Cancelling [sendCancelInternal]
+ * <any non-terminal> --(terminator/error)--> terminal [completeTerminal]
+ * }}}
+ * The two clean terminals carry the worker's `FinishResponse` / `CancelResponse`
+ * (metrics + finish/cancel callback `data`/`error`) so [[close]] can return
+ * them. Cancellation intent is deliberately tracked outside the machine in
+ * [[cancelRequested]], which makes the wire write idempotent and suppresses any
+ * in-flight Data/Finish once cancellation begins.
+ *
+ * Threading:
+ * - [[doInit]] is synchronous: sends `Init` and blocks on `InitResponse`,
+ * returning it.
+ * - [[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.
+ * - [[doClose]] is thread-safe and idempotent: it settles + returns the
+ * terminator (cancelling in-flight work if the stream had not finished)
+ * and terminates the request side.
+ *
+ * TODO [SPARK-55278]: this class does not yet implement payload chunking;
+ * the entire [[UdfPayload]] 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 a gRPC channel built and owned by the caller (the
+ * dispatcher). 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`.
+ * @param interruptCancelTimeoutMs upper bound on the wait for a `CancelResponse`
+ * after an interrupt (cancelled query / killed
+ * task) sends `Cancel`. Short by design; on expiry
+ * the session settles `Interrupted`. See
+ * `handleInterrupt`.
+ */
+@Experimental
+class GrpcWorkerSession(
+ workerHandle: WorkerHandle,
+ channel: ManagedChannel,
+ logger: WorkerLogger = WorkerLogger.NoOp,
+ initResponseTimeoutMs: Long = DEFAULT_INIT_RESPONSE_TIMEOUT_MS,
+ terminalTimeoutMs: Long = DEFAULT_TERMINAL_TIMEOUT_MS,
+ interruptCancelTimeoutMs: Long = DEFAULT_INTERRUPT_CANCEL_TIMEOUT_MS)
+ extends WorkerSession(workerHandle, logger) {
+
+ require(channel != null, "channel is required")
+
+ private val asyncStub = UdfWorkerGrpc.newStub(channel)
+
+ // Output batches from the worker, drained by the process() iterator.
+ // Intentionally unbounded in this first implementation: a bounded queue would
+ // block the gRPC callback (Netty event-loop) thread when full, stalling
+ // terminator/control delivery on the whole channel. The consumer normally
+ // drains promptly, but HTTP/2 flow control alone does not bound this application
+ // queue or gRPC's asynchronous send buffers.
+ //
+ // TODO [SPARK-55278]: add application-level gRPC flow control in a follow-up,
+ // using ClientCallStreamObserver readiness for requests and manual inbound
+ // demand for responses, before wiring this transport into a production path.
+ // TODO [SPARK-57324]: expose queue depth as a metric (early warning for a
+ // stalled consumer).
+ private val outputQueue = new LinkedBlockingQueue[QueueItem]()
+
+ // This value couples the worker's `InitResponse` (success or error) with the
+ // latch on which init() blocks. The latch fires when the InitResponse arrives
+ // (`complete`), when a pre-init ErrorResponse arrives, or when a terminal
+ // settles first without an InitResponse (`signalWithoutValue`). Until it fires
+ // we have no proof the worker accepted the session.
+ //
+ // Settle-before-release rule (referenced from every callback that both settles
+ // a terminal and fires this latch): settle the terminal FIRST, then complete /
+ // signal. init() blocks on the latch, and only a latch await/release pair
+ // establishes a happens-before edge, so a woken init() is guaranteed to
+ // observe the terminal rather than a transient state. OneShotValue keeps that
+ // publish-then-release in one place instead of every caller remembering to
+ // count down after setting the reference.
+ private val initValue = new OneShotValue[InitResponse]
+
+ // Fired when the session reaches a terminal [[SessionState]]. doClose() and
+ // the init-error path block on this to drain the terminator.
+ private val terminalLatch = new CountDownLatch(1)
+
+ // Captures an ErrorResponse encountered during the data phase so that
+ // the CancelResponse terminator 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 immediately before Finish is handed to the request observer. It is set
+ // before onNext so a reentrant callback may deliver FinishResponse while the
+ // Finish request is still being written.
+ private val finishSendStarted = new AtomicBoolean(false)
+
+ // Cancellation INTENT -- distinct from the `Cancelling` state, which is reached
+ // once a Cancel send is attempted under [[requestLock]] ([[sendCancelInternal]]).
+ // Intent is set first and can outrun (or never reach) that attempt, so
+ // `cancelRequested` does NOT imply state `Cancelling`. It is kept outside the
+ // [[SessionState]] machine to (a) make cancellation idempotent across all call
+ // sites and (b) suppress any Data/Finish that would otherwise race a Cancel
+ // onto the wire (re-read inside [[requestLock]]).
+ private val cancelRequested = new AtomicBoolean(false)
+
+ // gRPC requires serialized writes to a request StreamObserver.
+ private val requestLock = new Object
+
+ // Initialized in init() -- before that, close() is a no-op on the request
+ // side, which is exactly the contract the wrapping WorkerSession expects.
+ @volatile private var requestObserver: StreamObserver[UdfRequest] = _
+
+ // True after the request side has been half-closed, aborted, or observed the
+ // transport closing. It prevents this session from issuing duplicate terminal
+ // calls, including following onError with onCompleted from close().
+ private val requestSideTerminated = new AtomicBoolean(false)
+
+ private val responseObserver: StreamObserver[UdfResponse] = new StreamObserver[UdfResponse] {
+ override def onNext(response: UdfResponse): Unit = {
+ response.getResponseCase match {
+ case UdfResponse.ResponseCase.DATA =>
+ // A DataResponse before InitResponse violates the protocol (InitResponse
+ // must precede any DataResponse): fast-fail rather than enqueue it and let
+ // init() block to its timeout. Settling the terminal wakes a blocked
+ // init() via onTerminalSettled.
+ if (!initResolved) {
+ Transitions.transportFailed(new IllegalStateException(
+ "worker sent a DataResponse before InitResponse"))
+ } else {
+ outputQueue.put(QueueItem.Batch(response.getData))
+ }
+
+ case UdfResponse.ResponseCase.CONTROL =>
+ handleControl(response.getControl)
+
+ case other =>
+ // A malformed response (empty / unknown oneof) is a terminal transport
+ // failure; fast-fail init the same way.
+ Transitions.transportFailed(new IllegalStateException(
+ s"unexpected response oneof: $other"))
+ }
+ }
+
+ override def onError(t: Throwable): Unit = {
+ // Transport-level failure: the stream is dead, no further writes possible.
+ // Settling the terminal wakes a blocked init() (via onTerminalSettled) so it
+ // surfaces the transport cause instead of the initResponseTimeoutMs error.
+ requestSideTerminated.set(true)
+ Transitions.transportFailed(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. Settling the terminal wakes a blocked
+ // init() via onTerminalSettled.
+ requestSideTerminated.set(true)
+ if (!currentState.isTerminal) {
+ Transitions.transportFailed(new IllegalStateException(
+ "worker response stream closed without a terminator"))
+ }
+ }
+ }
+
+ /**
+ * Wakes everything that can be blocked when the base settles a terminal:
+ * the result iterator (on [[outputQueue]]), a thread on [[terminalLatch]]
+ * (close), and a thread still blocked in [[doInit]] on [[initValue]]. Invoked
+ * once, by the caller that wins [[completeTerminal]], so this is the '''single'''
+ * place a settled terminal wakes a blocked [[doInit]] -- callers that settle a
+ * terminal (any response-callback failure/terminator branch, the timeout paths,
+ * or close()) do NOT signal [[initValue]] themselves; settling is enough.
+ * Settle-before-release (see [[initValue]]) holds because this runs after the
+ * terminal CAS in [[completeTerminal]]. The only direct [[initValue]] signals
+ * left are the non-terminal init paths ([[handleControl]]'s INIT-ok / INIT-error
+ * / pre-init-ERROR branches), which must wake [[doInit]] '''without''' settling a
+ * terminal. INIT-ok first advances the session to `Initialized`; the error
+ * branches remain `Initializing` only until [[doInit]] sends Cancel and throws.
+ */
+ override protected def onTerminalSettled(termination: Termination): Unit = {
+ outputQueue.put(QueueItem.EndOfStream)
+ initValue.signalWithoutValue()
+ terminalLatch.countDown()
+ }
+
+ private def handleControl(ctrl: UdfControlResponse): Unit = ctrl.getControlCase match {
+ case UdfControlResponse.ControlCase.INIT =>
+ val resp = ctrl.getInit
+ if (resp.hasError) {
+ // Record the error and publish the InitResponse; do NOT settle a terminal
+ // or send Cancel here. The proto requires the engine to Cancel after an
+ // init error (udf_message.proto); doInit owns that synchronous cleanup
+ // after the initial onNext returns, then drains the CancelResponse before
+ // throwing. Settle-before-release (see initValue): there is no terminal to
+ // settle first here, just publish.
+ executionError.compareAndSet(None, Some(resp.getError))
+ initValue.complete(resp)
+ } else {
+ // InitResponse OK. Only advance from Initializing so a terminal that raced
+ // in (e.g. a transport error) still wins; process() then opens the data
+ // phase. Settle-before-release (see initValue): publish after the CAS.
+ Transitions.initAccepted()
+ initValue.complete(resp)
+ }
+
+ case UdfControlResponse.ControlCase.ERROR =>
+ val err = ctrl.getError.getError
+ executionError.compareAndSet(None, Some(err))
+ if (!initResolved) {
+ // Pre-init ErrorResponse. Leave the session in Initializing and just wake
+ // init(): the recorded executionError tells doInit the worker failed, and
+ // doInit sends the Cancel and drains the CancelResponse after the initial
+ // onNext returns, just like the INIT-error branch.
+ initValue.signalWithoutValue()
+ } else {
+ // Data-phase ErrorResponse: requestObserver is published, so cancel here.
+ // Cancel -> CancelResponse settles the terminal; the iterator surfaces the
+ // recorded executionError. The init latch already fired in init(), so no
+ // signalWithoutValue is needed.
+ sendCancelInternal(() => cancelWithReason("aborting after ErrorResponse"))
+ }
+
+ case UdfControlResponse.ControlCase.FINISH =>
+ // The FinishResponse carries metrics + the finish-callback data/error.
+ // Keep it on the terminal so close() can return it; the iterator inspects
+ // its error field to decide whether to throw. A FinishResponse is valid only
+ // after the engine has started sending Finish; accepting one earlier could
+ // silently truncate input and report a clean result.
+ if (finishSendStarted.get()) {
+ Transitions.finished(ctrl.getFinish)
+ } else {
+ Transitions.transportFailed(new IllegalStateException(
+ "worker sent FinishResponse before the engine sent Finish"))
+ }
+
+ case UdfControlResponse.ControlCase.CANCEL =>
+ // The CancelResponse carries metrics + the cancel-callback error. Keep it
+ // on the terminal so close() can return it; any prior ErrorResponse is
+ // tracked in executionError and surfaced by the iterator. Settling the
+ // terminal wakes a blocked init() (a CANCEL before InitResponse) via
+ // onTerminalSettled.
+ Transitions.cancelled(ctrl.getCancel)
+
+ case UdfControlResponse.ControlCase.CONTROL_NOT_SET =>
+ // Settling the terminal wakes a blocked init() via onTerminalSettled.
+ Transitions.transportFailed(new IllegalStateException(
+ "empty UdfControlResponse oneof"))
+ }
+
+ /**
+ * True once init is no longer pending -- i.e. the stream is past `Initializing`.
+ * Not "init succeeded": a terminal (including a failure) also counts as resolved.
+ */
+ private def initResolved: Boolean = currentState match {
+ case SessionState.Created | SessionState.Initializing => false
+ case _ => true
+ }
+
+ private def cancelWithReason(reason: String): Cancel =
+ Cancel.newBuilder().setReason(reason).build()
+
+ /**
+ * The protocol transition graph in one place: names for the edges, not a
+ * second source of truth. Every edge acts on the single
+ * [[WorkerSession.SessionState]] owned by the base via [[compareAndSetState]]
+ * (non-terminal) or [[completeTerminal]] (terminal), so the base's CAS is the
+ * only synchronization and a terminal that arrived first always wins (the
+ * non-terminal CASes fail against it; [[completeTerminal]] is first-wins).
+ *
+ * Edges this class drives -- edge, method, then driver site(s) / thread (the
+ * base drives the API-call edges: `Created -> Initializing` in `init`,
+ * `Initialized -> Streaming` in `process`):
+ * {{{
+ * Initializing -> Initialized initAccepted handleControl INIT-ok [gRPC cb]
+ * Streaming -> Finishing beginFinish advance branch 3 [engine]
+ * non-terminal -> Cancelling beginCancelFrom sendCancelInternal, immediately
+ * before the Cancel send [gRPC cb | engine | init | close]
+ * non-terminal -> Terminal finished/cancelled/transportFailed/interrupted
+ * handleControl / doInit / doClose / advance / onError
+ * }}}
+ * `Cancelling` is reached only when a Cancel send starts; a pre-stream or
+ * raced cancel goes straight to a `Cancelled`/`TransportFailed` terminal (see
+ * [[sendCancelInternal]], [[doClose]], `ProcessIterator`) or nowhere -- so
+ * `cancelRequested` (intent) does not imply state `Cancelling`.
+ */
+ private object Transitions {
+ /** `InitResponse` OK: `Initializing -> Initialized`. */
+ def initAccepted(): Boolean =
+ compareAndSetState(SessionState.Initializing, SessionState.Initialized)
+
+ /** Input exhausted: `Streaming -> Finishing` (once). */
+ def beginFinish(): Boolean =
+ compareAndSetState(SessionState.Streaming, SessionState.Finishing)
+
+ /** `Cancel` send starts: `cur -> Cancelling`, from any non-terminal `cur`. */
+ def beginCancelFrom(cur: SessionState): Boolean =
+ !cur.isTerminal && compareAndSetState(cur, SessionState.Cancelling)
+
+ /** Clean terminal carrying the worker's `FinishResponse`. */
+ def finished(response: FinishResponse): Boolean =
+ completeTerminal(Termination.Finished(response))
+
+ /** Clean terminal carrying the worker's `CancelResponse`. */
+ def cancelled(response: CancelResponse): Boolean =
+ completeTerminal(Termination.Cancelled(response))
+
+ /** Failure terminal carrying a transport-level cause. */
+ def transportFailed(cause: Throwable): Boolean =
+ completeTerminal(Termination.TransportFailed(cause))
+
+ /**
+ * Terminal for an engine-thread interrupt (e.g. a Spark task kill). Distinct
+ * from [[transportFailed]] because the cause is an engine-side interrupt, but
+ * still unsalvageable without a worker acknowledgement. Settled by
+ * [[handleInterrupt]] only when the brief post-Cancel drain expires without a
+ * `CancelResponse`; if the worker acks in time the session settles a
+ * cooperative [[cancelled]] instead.
+ */
+ def interrupted(cause: Throwable): Boolean =
+ completeTerminal(Termination.Interrupted(cause))
+ }
+
+ // ---- WorkerSession hooks ------------------------------------------------
+
+ override protected def doInit(message: Init): InitResponse = {
+ // Construct the request before opening the RPC. A malformed Init must not
+ // leave an otherwise-unused Execute stream waiting for its first request.
+ val initRequest = UdfRequest.newBuilder()
+ .setControl(UdfControlRequest.newBuilder().setInit(message).build())
+ .build()
+
+ // 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")
+ Transitions.transportFailed(ex)
+ throw new GrpcWorkerSessionException("UDF worker channel is closed", ex)
+ }
+ try {
+ requestLock.synchronized {
+ // Serialize stream creation, publication, and Init with close/cancel. A
+ // close that wins the lock first prevents the RPC from opening; one that
+ // loses cannot release the worker until the initial write returns.
+ if (!currentState.isTerminal) {
+ val stream = asyncStub.execute(responseObserver)
+ requestObserver = stream
+ // Stream creation may invoke a response callback synchronously. Do not
+ // send Init if that callback has already settled a terminal.
+ if (!currentState.isTerminal) {
+ // With directExecutor, gRPC callbacks run synchronously on the caller's
+ // thread, so InitResponse can reach responseObserver/handleControl from
+ // *inside* this stream.onNext.
+ // requestObserver is already published, so an immediate ErrorResponse
+ // can write its required Cancel reentrantly without losing the request.
+ sendOnNext(initRequest)
+ }
+ }
+ }
+ } catch {
+ case NonFatal(e) =>
+ Transitions.transportFailed(e)
+ // Surface as GrpcWorkerSessionException so the engine integration layer
+ // (which catches that type and wraps it) sees a uniform init-failure
+ // exception rather than the raw transport error.
+ throw new GrpcWorkerSessionException("UDF worker stream failed during init", e)
+ }
+
+ try {
+ initValue.await(initResponseTimeoutMs)
+ } catch {
+ case _: InterruptedException =>
+ // Interrupt (cancelled query / killed task) while awaiting InitResponse:
+ // cooperatively Cancel with a bounded drain, settling Cancelled if the
+ // worker acks in time, else Interrupted. See handleInterrupt.
+ handleInterrupt("waiting for InitResponse")
+ throw new InterruptedException("interrupted while waiting for InitResponse")
+ case e: TimeoutException =>
+ sendCancelInternal(() => cancelWithReason("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.
+ Transitions.transportFailed(e)
+ // Surface as GrpcWorkerSessionException (carrying the timeout cause) so
+ // the engine integration layer that catches that type can wrap it.
+ throw new GrpcWorkerSessionException(
+ s"timed out waiting for InitResponse after ${initResponseTimeoutMs}ms", e)
+ }
+
+ initValue.get match {
+ case Some(resp) if resp.hasError =>
+ failInitWithError(resp.getError,
+ s"UDF worker init failed: ${describeError(resp.getError)}")
+ case Some(resp) =>
+ executionError.get() match {
+ case Some(err) =>
+ // A reentrant callback can report an ErrorResponse immediately after
+ // InitResponse and before the initial onNext returns. The callback
+ // has already sent Cancel; drain its response
+ // and surface the original error instead of returning init success
+ // for a session that is already cancelling or terminal.
+ failInitWithError(err,
+ s"UDF worker reported an error as init completed: ${describeError(err)}")
+ case None if currentState != SessionState.Initialized =>
+ // Likewise, an immediate protocol failure or concurrent close may
+ // have moved the session out of Initialized after publishing the
+ // InitResponse. A normal return would violate init()'s contract and
+ // leave process() to fail later with only an ordering error.
+ failInitFromCurrentState()
+ case None =>
+ resp
+ }
+ case None =>
+ // No InitResponse arrived but the latch fired.
+ //
+ // A pre-init ErrorResponse leaves the session in Initializing with the
+ // error recorded in executionError (see handleControl ERROR branch): the
+ // engine must now Cancel and drain the CancelResponse, which failInit does.
+ executionError.get() match {
+ case Some(err) if !currentState.isTerminal =>
+ failInitWithError(err,
+ s"UDF worker reported an error before init completed: ${describeError(err)}")
+ case _ =>
+ // Otherwise a terminal already settled -- the worker terminated the
+ // stream before sending InitResponse (transport error, half-close, or
+ // a premature FinishResponse/CancelResponse). Surface it as an init
+ // failure rather than letting the caller proceed as if init succeeded.
+ failInitFromCurrentState()
+ }
+ }
+ }
+
+ /** Surfaces the terminal that prevented init from completing normally. */
+ private def failInitFromCurrentState(): Nothing = {
+ executionError.get() match {
+ case Some(err) =>
+ failInitWithError(err,
+ s"UDF worker reported an error as init completed: ${describeError(err)}")
+ case None => ()
+ }
+ // Cancelling can be observed after a concurrent close or a reentrant error
+ // before its CancelResponse arrives. Drain it so the exception reflects the
+ // stable terminal rather than a transient state.
+ if (currentState == SessionState.Cancelling) {
+ awaitTerminal()
+ }
+ currentState match {
+ case SessionState.Terminal(Termination.TransportFailed(cause)) =>
+ throw new GrpcWorkerSessionException("UDF worker stream failed during init", cause)
+ case SessionState.Terminal(Termination.Failed(err)) =>
+ throw new GrpcWorkerSessionException(
+ s"UDF worker reported an error before init completed: ${describeError(err)}", err)
+ case SessionState.Terminal(Termination.Cancelled(_)) =>
+ throw new GrpcWorkerSessionException(
+ "UDF worker stream was cancelled before init completed")
+ case SessionState.Terminal(Termination.Interrupted(cause)) =>
+ throw new GrpcWorkerSessionException("UDF worker init was interrupted", cause)
+ case SessionState.Terminal(Termination.Finished(_)) =>
+ throw new GrpcWorkerSessionException("UDF worker finished before init completed")
+ case other =>
+ throw new IllegalStateException(
+ s"init completed without an accepted session or terminal: $other")
+ }
+ }
+
+ /**
+ * Fails init when the worker reports an error before init returns: an
+ * `InitResponse` carrying an error, a pre-init `ErrorResponse`, or an immediate
+ * post-init `ErrorResponse`. The protocol requires the engine to send `Cancel`
+ * and the worker to reply with `CancelResponse` (udf_message.proto); we send it,
+ * drain the terminator so no stream is left dangling, and throw the structured
+ * error. A responsive worker settles `Cancelled`; a failed drain settles
+ * `TransportFailed`.
+ */
+ private def failInitWithError(err: ExecutionError, message: String): Nothing = {
+ sendCancelInternal(() => cancelWithReason("init failed"))
+ awaitTerminal()
+ throw new GrpcWorkerSessionException(message, err)
+ }
+
+ override protected def doProcess(
+ input: Iterator[DataRequest],
+ finish: () => Finish): Iterator[DataResponse] = {
+ // Init success is guaranteed by the base [[WorkerSession]] lifecycle: if
+ // doInit had failed it would have thrown and process() would never run.
+ new ProcessIterator(input, finish)
+ }
+
+ override protected def doClose(cancel: () => Cancel): Termination = {
+ // Coordinate the no-stream decision with doInit's publication + Init write.
+ // If close wins this lock, doInit observes the terminal and never sends Init;
+ // if doInit wins, close observes the published observer and sends Cancel only
+ // after the Init send has returned.
+ val noPublishedStream = requestLock.synchronized {
+ if (requestObserver == null) {
+ // init() never put a stream on the wire (closed before/around init, or
+ // init threw before publishing). There is no protocol terminator; if no
+ // terminal has settled yet, treat the session as cancelled-before-start.
+ // A terminal may already be settled here (e.g. the channel-shutdown
+ // TransportFailed in doInit also leaves requestObserver null); return the
+ // settled terminal as-is rather than a bare Cancelled that disagrees with
+ // the state. The base WorkerSession still releases the worker handle.
+ if (!currentState.isTerminal) {
+ Transitions.cancelled(CancelResponse.getDefaultInstance)
+ }
+ true
+ } else {
+ false
+ }
+ }
+ if (noPublishedStream) {
+ return settledTermination
+ }
+ // If the stream has not finished on its own, cancel anything in flight so
+ // the worker can clean up, then wait for the terminator. sendCancelInternal
+ // evaluates the cancel thunk only when it attempts a Cancel, and at
+ // most once across all callers.
+ if (!currentState.isTerminal) {
+ sendCancelInternal(cancel)
+ try {
+ terminalLatch.await(terminalTimeoutMs, TimeUnit.MILLISECONDS)
+ } catch {
+ case _: InterruptedException =>
+ // Interrupted mid-close: settle Interrupted via the
+ // shared handler rather than falling through to the TransportFailed
+ // guard below. The Cancel was already attempted, so handleInterrupt's
+ // sendCancelInternal is a no-op and it just runs the bounded drain.
+ //
+ // TODO [SPARK-57640]: close() is a finalizer often already on the task
+ // kill path, yet this makes an interrupted close block up to another
+ // interruptCancelTimeoutMs (chasing a clean, salvageable Cancelled
+ // rather than immediately unwinding to Interrupted). Current choice:
+ // spend the bounded wait to obtain proof the worker is recyclable; it
+ // is small next to terminalTimeoutMs. Revisit if killed-task teardown
+ // latency (especially many sessions torn down at once) makes an
+ // immediate unwind preferable -- e.g. skip this second drain in the
+ // close() path, since close already gave the worker its terminalTimeoutMs window.
+ handleInterrupt("closing the session")
+ }
+ }
+ // If the worker still has not settled (timeout above, or an interrupt whose
+ // bounded drain did not settle a terminal), record a terminal so
+ // isWorkerSalvageable returns a definite answer and any other thread reading
+ // the state sees a stable value.
+ if (!currentState.isTerminal) {
+ Transitions.transportFailed(new TimeoutException(
+ s"timed out waiting for stream terminator after ${terminalTimeoutMs}ms"))
+ }
+ // Close the request side according to the settled outcome. Clean protocol
+ // terminators are half-closed; failures abort the RPC so close() never follows
+ // an onError with onCompleted.
+ currentState match {
+ case SessionState.Terminal(Termination.Finished(_)) |
+ SessionState.Terminal(Termination.Cancelled(_)) =>
+ completeRequestSide()
+ case SessionState.Terminal(Termination.Failed(error)) =>
+ abortRequestSide(new GrpcWorkerSessionException(
+ s"UDF execution failed: ${describeError(error)}", error))
+ case SessionState.Terminal(Termination.TransportFailed(cause)) =>
+ abortRequestSide(cause)
+ case SessionState.Terminal(Termination.Interrupted(cause)) =>
+ abortRequestSide(cause)
+ case other =>
+ logger.debug(s"UDF Execute stream closed without a terminal outcome: $other")
+ }
+ // Derive the Termination from the settled terminal. close() is the
+ // finalizer/cleanup path and must NOT re-throw: a UDF / data-phase error is
+ // already surfaced while the result iterator is consumed, and an init error
+ // is surfaced from init(). settledTermination returns the settled terminal
+ // as-is: the response proto for the clean terminators, and the failure
+ // terminals (Failed / TransportFailed / Interrupted) carrying their cause --
+ // not a bare Cancelled.
+ settledTermination
+ }
+
+ // ---- Internal request helpers ---------------------------------------------
+
+ /**
+ * Writes one request and aborts the request side if the observer rejects it.
+ * Callers hold [[requestLock]], so this also serializes the compensating
+ * `onError` with every other request-side operation.
+ */
+ private def sendOnNext(req: UdfRequest): Unit = {
+ try {
+ requestObserver.onNext(req)
+ } catch {
+ case NonFatal(e) =>
+ abortRequestSide(e)
+ throw e
+ }
+ }
+
+ /** Aborts the request side at most once. Safe to call while holding [[requestLock]]. */
+ private def abortRequestSide(cause: Throwable): Unit = requestLock.synchronized {
+ if (requestObserver != null && requestSideTerminated.compareAndSet(false, true)) {
+ try {
+ requestObserver.onError(cause)
+ } catch {
+ case NonFatal(e) => logger.debug("Error aborting UDF Execute stream", e)
+ }
+ }
+ }
+
+ /** Half-closes the request side at most once. */
+ private def completeRequestSide(): Unit = requestLock.synchronized {
+ if (requestObserver != null && requestSideTerminated.compareAndSet(false, true)) {
+ try {
+ requestObserver.onCompleted()
+ } catch {
+ case NonFatal(e) => logger.debug("Error half-closing UDF Execute stream", e)
+ }
+ }
+ }
+
+ /**
+ * Sends a Data or Finish request to the worker. Three invariants are
+ * checked inside [[requestLock]]:
+ * - terminal state: writes are unsafe (transport dead or terminal
+ * received). Throws, so the caller's terminal/exception path runs.
+ * - request-side termination: no write may follow `onError` / `onCompleted`.
+ * Throws so the caller settles or surfaces the failure.
+ * - [[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]] publishes
+ * [[requestObserver]] and writes Init together under [[requestLock]], so a
+ * concurrent cancel observes the observer but cannot acquire the lock and send
+ * Cancel until after Init.
+ */
+ private def sendRequest(req: UdfRequest): Unit =
+ requestLock.synchronized {
+ if (currentState.isTerminal) {
+ throw new IllegalStateException(
+ "cannot send request: UDF Execute stream is already closed")
+ }
+ if (requestSideTerminated.get()) {
+ throw new IllegalStateException(
+ "cannot send request: UDF Execute request stream is already closed")
+ }
+ // Suppress the write when a cancel has been (or is about to be) flushed
+ // through this lock; that preserves the proto ordering invariant (no
+ // Data/Finish after Cancel). Test the flag rather than `return`-ing from
+ // this by-name `synchronized` body: a non-local return compiles to a thrown
+ // NonLocalReturnControl, which is brittle (also relied on in sendCancelInternal).
+ if (!cancelRequested.get()) {
+ if (req.hasControl && req.getControl.hasFinish) {
+ // Set before onNext: directExecutor delivery may return FinishResponse
+ // reentrantly from inside this call.
+ finishSendStarted.set(true)
+ }
+ sendOnNext(req)
+ }
+ }
+
+ /**
+ * Sends a `Cancel` control message, returning `true` iff the request observer
+ * accepts the `Cancel`. Idempotent across ALL call sites (close()'s
+ * cancel, in-band cancels from `handleControl`'s INIT error / ERROR paths,
+ * and engine-internal cancels from the iterator) via [[cancelRequested]]:
+ * only the first caller attempts a `Cancel`, and the `cancel` thunk is evaluated
+ * only by that caller and only when the send is about to start --
+ * so a side-effecting thunk (e.g. one carrying a client cancel callback) runs
+ * at most once. 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(cancel: () => Cancel): Boolean = {
+ // Do not consume the one-shot cancellation intent until there is an observer
+ // that can carry it. doInit publishes the observer before sending Init while
+ // holding requestLock, so every response to Init sees a non-null observer.
+ if (requestObserver == null) return false
+ if (!cancelRequested.compareAndSet(false, true)) return false
+ if (currentState.isTerminal) return false
+ try {
+ requestLock.synchronized {
+ // A terminal that arrived first must win; bail without writing. Compute
+ // the block's value rather than `return` (see sendRequest for why non-local
+ // return is avoided -- here it would also escape the NonFatal catch below).
+ val cur = currentState
+ if (cur.isTerminal || requestSideTerminated.get()) {
+ false
+ } else {
+ val request = UdfRequest.newBuilder()
+ .setControl(UdfControlRequest.newBuilder().setCancel(cancel()).build())
+ .build()
+ // Record that a Cancel send has started by advancing to Cancelling.
+ Transitions.beginCancelFrom(cur)
+ sendOnNext(request)
+ true
+ }
+ }
+ } catch {
+ case NonFatal(e) =>
+ logger.debug(s"Cancel send failed (stream may already be torn down): ${e.getMessage}")
+ Transitions.transportFailed(e)
+ false
+ }
+ }
+
+ private def awaitTerminal(): Unit = {
+ if (currentState.isTerminal) return
+ try {
+ if (!terminalLatch.await(terminalTimeoutMs, TimeUnit.MILLISECONDS)) {
+ Transitions.transportFailed(new TimeoutException(
+ s"timed out waiting for stream terminator after ${terminalTimeoutMs}ms"))
+ }
+ } catch {
+ case _: InterruptedException => handleInterrupt("waiting for stream terminator")
+ }
+ }
+
+ /**
+ * Handles an [[InterruptedException]] observed while blocked on a worker event
+ * (init, terminal drain, close, or the result-iterator poll). An interrupt is
+ * an engine-side event -- typically a cancelled query / killed task -- not a
+ * worker fault, so we cancel cooperatively and keep the worker salvageable only
+ * when it acks:
+ *
+ * 1. Send a best-effort `Cancel` (idempotent across call sites).
+ * 2. Wait up to [[interruptCancelTimeoutMs]] -- short, so the interrupted
+ * thread unwinds promptly -- for the worker's `CancelResponse`. A healthy
+ * worker acks in that window and the callback settles a clean `Cancelled`
+ * terminal (worker salvageable, response drained).
+ * 3. If the ack does not arrive in time, settle `Interrupted`. Without a
+ * terminator or liveness proof the worker is not salvageable and must not
+ * be returned to a reuse pool.
+ *
+ * The bounded wait deliberately runs '''before''' the thread's interrupt flag
+ * is restored: [[InterruptedException]] clears the flag on throw, so re-setting
+ * it first would make the wait below throw immediately and defeat the drain.
+ * The flag is restored at the end so the caller's unwind still observes the
+ * interrupt.
+ *
+ * TODO [SPARK-57640]: add a worker liveness/heartbeat check so a future version
+ * may prove that an interrupted worker which missed the short ack window is
+ * nevertheless safe to recycle.
+ */
+ private def handleInterrupt(waitContext: String): Unit = {
+ // NB: flag is currently clear (InterruptedException cleared it); do not
+ // restore it until after the bounded drain below.
+ if (!currentState.isTerminal) {
+ sendCancelInternal(() => cancelWithReason(s"interrupted while $waitContext"))
+ val acked = try {
+ terminalLatch.await(interruptCancelTimeoutMs, TimeUnit.MILLISECONDS)
+ } catch {
+ case _: InterruptedException =>
+ // A second interrupt during the drain: give up the wait immediately.
+ false
+ }
+ if (!acked && !currentState.isTerminal) {
+ // No CancelResponse in time: settle Interrupted rather than
+ // TransportFailed, and stop a later close() from blocking on a terminator
+ // that will not be drained on this thread's behalf. Interrupted remains
+ // unsalvageable until a separate liveness check can prove otherwise.
+ Transitions.interrupted(
+ new InterruptedException(s"interrupted while $waitContext"))
+ }
+ }
+ Thread.currentThread().interrupt()
+ }
+
+ // ---- 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 settle the state) and with any thread that finalizes via
+ * [[close]]. It drives the `Streaming -> Finishing` transition (sending
+ * `Finish` once input is exhausted, built lazily from the `finish` thunk).
+ */
+ private class ProcessIterator(input: Iterator[DataRequest], finish: () => Finish)
+ extends Iterator[DataResponse] {
+
+ // 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 branches, and block branch 4 for the
+ // full terminalTimeoutMs before returning. Callers that probe hasNext
+ // an extra time (iterator.size, instrumentation wrappers) would hang.
+ // Iterator-local and only touched by the single engine thread.
+ private val exhausted = new AtomicBoolean(false)
+ @volatile private var prefetched: DataResponse = _
+
+ /**
+ * Runs a caller-supplied thunk (`input.hasNext`, `input.next()`, or the
+ * `finish` builder) that may throw, and on failure Cancels the in-flight
+ * stream before rethrowing. Without the Cancel a failure would leave the
+ * worker awaiting more input with no terminator owed on the wire; with it the
+ * worker tears down and the exception still propagates to the engine. Setting
+ * `cancelRequested` also suppresses any further Data/Finish this loop might
+ * otherwise attempt (see [[sendRequest]]).
+ */
+ private def cancelOnThrow[T](reason: String)(op: => T): T =
+ try op catch {
+ case NonFatal(e) =>
+ sendCancelInternal(() => cancelWithReason(reason))
+ throw e
+ }
+
+ override def hasNext: Boolean = {
+ if (prefetched ne null) return true
+ advance()
+ prefetched ne null
+ }
+
+ override def next(): DataResponse = {
+ if (prefetched eq null) advance()
+ val out = prefetched
+ if (out eq null) {
+ throw new NoSuchElementException("ProcessIterator exhausted")
+ }
+ prefetched = null
+ out
+ }
+
+ /**
+ * Fills [[prefetched]] with the next output batch, or leaves it null at the
+ * terminator (which also sets [[exhausted]], so `hasNext` reads null as "done").
+ * Each loop iteration tries in order: (1) drain queued output -- a batch (fill
+ * and return) or the terminator sentinel (return); (2) while `Streaming`, send
+ * the next input batch and loop; (3) once input is exhausted, send `Finish`
+ * (CAS `Streaming -> Finishing`, once) and loop; (4) otherwise block for late
+ * output or the terminator, then return. Branches 2-3 loop; 1 and 4 return.
+ *
+ * '''Per-poll timeout.''' Branch 4 waits up to [[terminalTimeoutMs]] per
+ * poll, reset by every worker event -- the contract is "emit at least one
+ * event every [[terminalTimeoutMs]] after Finish", not "finish the UDF within
+ * [[terminalTimeoutMs]]". A worker expecting a long post-Finish silence MAY
+ * emit an empty `DataResponse` as a heartbeat to reset the wait; it is
+ * surfaced to the caller, so it should be a batch the caller recognises as
+ * empty (e.g. a zero-row Arrow batch).
+ */
+ private def advance(): Unit = {
+ if (exhausted.get()) return // terminator already seen; never re-block
+ 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 while the stream is open for data.
+ // `input.hasNext` may itself fetch/compute the next element (many Spark
+ // iterators prefetch), so it can throw just like `input.next()`; both go
+ // through cancelOnThrow so a failing upstream Cancels the stream instead
+ // of stranding the worker waiting for input that will never arrive.
+ if (!cancelRequested.get() && currentState == SessionState.Streaming &&
+ cancelOnThrow("input iterator failed")(input.hasNext)) {
+ val request = cancelOnThrow("input iterator failed") {
+ val batch = Objects.requireNonNull(
+ input.next(), "input iterator returned null")
+ UdfRequest.newBuilder().setData(batch).build()
+ }
+ if (!sendOrEndOnRacedTerminal(request)) return
+ } else if (!cancelRequested.get() && Transitions.beginFinish()) {
+ // (3) No more input; send Finish exactly once (unless cancelled). The
+ // `finish` thunk is caller-supplied (it may run a finish callback), so
+ // it can throw; cancelOnThrow Cancels the stream before rethrowing. The
+ // Streaming -> Finishing CAS has already run, but we have not written
+ // Finish yet, so the Cancel is the only engine-to-worker message that
+ // reaches the wire -- the proto "nothing after Cancel" invariant holds.
+ val request = cancelOnThrow("finish callback failed") {
+ val finishMsg = Objects.requireNonNull(
+ finish(), "finish callback returned null")
+ UdfRequest.newBuilder()
+ .setControl(UdfControlRequest.newBuilder().setFinish(finishMsg).build())
+ .build()
+ }
+ if (!sendOrEndOnRacedTerminal(request)) {
+ return
+ }
+ } else {
+ // (4) Block for late output or the terminator. See class doc
+ // above for the per-poll-vs-total-session timeout semantics.
+ val item = try {
+ outputQueue.poll(terminalTimeoutMs, TimeUnit.MILLISECONDS)
+ } catch {
+ case _: InterruptedException =>
+ // Interrupt (cancelled query / killed task) while reading results:
+ // cooperatively Cancel with a bounded drain, settling Cancelled if
+ // the worker acks in time, else Interrupted. See handleInterrupt.
+ handleInterrupt("reading UDF result")
+ exhausted.set(true)
+ throw new InterruptedException("interrupted while reading UDF result")
+ }
+ item match {
+ case null =>
+ Transitions.transportFailed(new IllegalStateException(
+ s"timed out waiting for UDF output after ${terminalTimeoutMs}ms"))
+ exhausted.set(true)
+ throwIfTerminalError()
+ return
+ case QueueItem.EndOfStream =>
+ exhausted.set(true)
+ throwIfTerminalError()
+ return
+ case QueueItem.Batch(b) =>
+ prefetched = b
+ return
+ }
+ }
+ }
+ }
+
+ /**
+ * Sends one data-phase request (`DataRequest` / `Finish`), recovering from a
+ * terminator that raced the write. Returns `true` to keep looping (sent, or
+ * suppressed by a pending cancel), `false` if a terminator settled and the
+ * iterator should end (caller `return`s).
+ *
+ * On a write failure a terminator has usually already settled (the worker
+ * finished/failed early) and the write only failed because the stream is
+ * closed: record a transport terminal if none is set ([[completeTerminal]] is
+ * a no-op once settled), then [[throwIfTerminalError]] -- which throws for an
+ * error/transport terminal, or returns for a clean `Finished` that raced the
+ * send, in which case the benign "stream closed" error is dropped.
+ */
+ private def sendOrEndOnRacedTerminal(req: UdfRequest): Boolean = {
+ try {
+ sendRequest(req)
+ true
+ } catch {
+ case NonFatal(e) =>
+ if (!currentState.isTerminal) {
+ Transitions.transportFailed(e)
+ }
+ exhausted.set(true)
+ throwIfTerminalError()
+ false
+ }
+ }
+
+ /**
+ * Surfaces a failed terminator as an exception when the result iterator is
+ * drained. A data-phase [[ExecutionError]] (captured in [[executionError]])
+ * takes precedence over the terminator's own error, then the finish/cancel
+ * callback error carried on the terminator, then a bare cancellation.
+ */
+ private def throwIfTerminalError(): Unit = currentState match {
+ case SessionState.Terminal(Termination.Finished(response)) =>
+ responseError(response.hasError, response.getError, executionError.get())
+ .foreach(err => throw new GrpcWorkerSessionException(
+ s"UDF execution failed: ${describeError(err)}", err))
+ case SessionState.Terminal(Termination.Cancelled(response)) =>
+ responseError(response.hasError, response.getError, executionError.get()) match {
+ case Some(err) =>
+ throw new GrpcWorkerSessionException(
+ s"UDF execution failed: ${describeError(err)}", err)
+ case None =>
+ throw new GrpcWorkerSessionException("UDF execution was cancelled")
+ }
+ case SessionState.Terminal(Termination.Failed(err)) =>
+ throw new GrpcWorkerSessionException(
+ s"UDF execution failed: ${describeError(err)}", err)
+ case SessionState.Terminal(Termination.TransportFailed(t)) =>
+ throw new GrpcWorkerSessionException("UDF worker stream failed", t)
+ case SessionState.Terminal(Termination.Interrupted(t)) =>
+ throw new GrpcWorkerSessionException("UDF execution was interrupted", t)
+ case other =>
+ throw new IllegalStateException(s"terminator sentinel without terminal: $other")
+ }
+
+ /** Picks the error to surface: prior data-phase error, else the terminator's. */
+ private def responseError(
+ hasError: Boolean,
+ error: ExecutionError,
+ priorError: Option[ExecutionError]): Option[ExecutionError] =
+ priorError.orElse(if (hasError) Some(error) else None)
+ }
+}
+
+object GrpcWorkerSession {
+ /** Upper bound on the wait for `InitResponse`. */
+ val DEFAULT_INIT_RESPONSE_TIMEOUT_MS: Long = 30000L
+
+ /** Upper bound on the wait for `FinishResponse` / `CancelResponse`. */
+ val DEFAULT_TERMINAL_TIMEOUT_MS: Long = 30000L
+
+ /**
+ * Upper bound on the wait for a `CancelResponse` after an interrupt (e.g. a
+ * cancelled query / killed task) sends `Cancel`. Much shorter than
+ * [[DEFAULT_TERMINAL_TIMEOUT_MS]]: the interrupted thread must unwind
+ * promptly, so a healthy worker gets a brief window to ack the Cancel (clean
+ * `Cancelled` terminal, worker salvageable) before the session falls back to
+ * an unsalvageable `Interrupted` terminal. See `handleInterrupt`.
+ */
+ val DEFAULT_INTERRUPT_CANCEL_TIMEOUT_MS: Long = 2000L
+
+ // Distinguishes a (possibly empty) data batch from end-of-stream, and makes
+ // the iterator's match exhaustive.
+ private sealed trait QueueItem
+ private object QueueItem {
+ final case class Batch(response: DataResponse) extends QueueItem
+ case object EndOfStream extends QueueItem
+ }
+
+ /**
+ * A write-once value paired with the latch a waiter blocks on, so the "publish
+ * the value, then release the waiter" ordering lives in one place instead of
+ * every call site having to remember to count the latch down after setting the
+ * reference. The value is set at most once ([[complete]]); the latch can also
+ * be released without a value ([[signalWithoutValue]]) when init fails through
+ * a pre-init error or terminal. All reads/writes go through the latch, so a
+ * waiter released by [[await]] has a happens-before edge to the [[complete]]
+ * that set the value.
+ */
+ private final class OneShotValue[A] {
+ private val latch = new CountDownLatch(1)
+ private val ref = new AtomicReference[Option[A]](None)
+ private val completed = new AtomicBoolean(false)
+
+ /** Publishes the value (first writer wins) and releases any waiter. */
+ def complete(value: A): Unit = {
+ if (completed.compareAndSet(false, true)) {
+ ref.set(Some(value))
+ latch.countDown()
+ }
+ }
+
+ /** Completes without a value (init failed through another event) and releases the waiter. */
+ def signalWithoutValue(): Unit = {
+ if (completed.compareAndSet(false, true)) {
+ latch.countDown()
+ }
+ }
+
+ /**
+ * Blocks up to `timeoutMs` for a release, throwing [[TimeoutException]] if
+ * none arrives so the caller handles the timeout on the exception path rather
+ * than by inspecting a return value.
+ */
+ def await(timeoutMs: Long): Unit =
+ if (!latch.await(timeoutMs, TimeUnit.MILLISECONDS)) {
+ throw new TimeoutException(
+ s"timed out waiting for InitResponse after ${timeoutMs}ms")
Review Comment:
`OneShotValue[A]` is generic, but this message names `InitResponse`. `doInit` already wraps the `TimeoutException` with its own `InitResponse` wording (L467-468), so the string is both misplaced here and duplicated -- and a second use of the helper would report the wrong value.
```suggestion
throw new TimeoutException(
s"timed out waiting for value after ${timeoutMs}ms")
```
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57640] Implement gRPC WorkerSession and test it against an echo server. [spark]
Posted by "haiyangsun-db (via GitHub)" <gi...@apache.org>.
haiyangsun-db commented on code in PR #56702:
URL: https://github.com/apache/spark/pull/56702#discussion_r3583989410
##########
udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/GrpcWorkerSession.scala:
##########
@@ -0,0 +1,816 @@
+/*
+ * 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.grpc
+
+import java.util.concurrent.{CountDownLatch, LinkedBlockingQueue, TimeoutException, TimeUnit}
+import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference}
+
+import scala.util.control.NonFatal
+
+import io.grpc.{ConnectivityState, ManagedChannel}
+import io.grpc.stub.StreamObserver
+
+import org.apache.spark.annotation.Experimental
+import org.apache.spark.udf.worker.{Cancel, CancelResponse, DataRequest, DataResponse,
+ ExecutionError, Finish, Init, InitResponse, UdfControlRequest,
+ UdfControlResponse, UdfRequest, UdfResponse, UdfWorkerGrpc}
+import org.apache.spark.udf.worker.core.{Termination, WorkerHandle, WorkerLogger, WorkerSession}
+import org.apache.spark.udf.worker.core.WorkerSession.SessionState
+import org.apache.spark.udf.worker.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_message.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.
+ *
+ * '''Driving model.''' Consumption-driven (Volcano / pull): the thread that
+ * consumes the [[doProcess]] result iterator is the one that pulls the next
+ * input batch and sends its `DataRequest`. Nothing is sent until the iterator
+ * is consumed, and input flows one batch per output pull. The gRPC callback
+ * thread only receives output.
+ *
+ * '''State machine.''' This class does not keep its own state machine: it
+ * drives the single [[WorkerSession.SessionState]] owned by the base. The base
+ * advances `Created -> Initializing` (in `init`) and `Initialized -> Streaming`
+ * (in `process`); this class advances the protocol-event edges through
+ * [[compareAndSetState]] / [[completeTerminal]] as it exchanges messages:
+ * {{{
+ * Initializing --(InitResponse ok)--> Initialized [handleControl]
+ * Streaming ----(Finish written)----> Finishing [ProcessIterator]
+ * <any non-terminal> --(Cancel written)--> Cancelling [sendCancelInternal]
+ * <any non-terminal> --(terminator/error)--> terminal [completeTerminal]
+ * }}}
+ * The two clean terminals carry the worker's `FinishResponse` / `CancelResponse`
+ * (metrics + finish/cancel callback `data`/`error`) so [[close]] can return
+ * them. The only flag kept outside the machine is [[cancelRequested]]: a
+ * cancellation can be requested before the stream exists (so it cannot be a
+ * state transition yet), and it must both fast-fail the result iterator and
+ * suppress any in-flight Data/Finish.
+ *
+ * Threading:
+ * - [[doInit]] is synchronous: sends `Init` and blocks on `InitResponse`,
+ * returning it.
+ * - [[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.
+ * - [[doClose]] is thread-safe and idempotent: it settles + returns the
+ * terminator (cancelling in-flight work if the stream had not finished)
+ * and half-closes the request side.
+ *
+ * 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 a gRPC channel built and owned by the caller (the
+ * dispatcher). 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 from the worker, drained by the process() iterator.
+ // Intentionally unbounded: a bounded queue would block the gRPC callback
+ // (Netty event-loop) thread when full, stalling control-message and
+ // terminator delivery on the whole channel. HTTP/2 flow control bounds the
+ // wire and the consumer normally drains promptly, so it stays small. A
+ // stalled downstream can still grow it; the real fix is protocol-level
+ // back-pressure (out of scope here), not bounding the queue.
+ //
+ // TODO [SPARK-57324]: expose queue depth as a metric (early warning for a
+ // stalled consumer).
+ 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)
+
+ // The InitResponse the worker sent (success or error), captured so init()
+ // can return it. None until InitResponse arrives.
+ private val initResponse = new AtomicReference[Option[InitResponse]](None)
+
+ // Fired when the session reaches a terminal [[SessionState]]. doClose() and
+ // the init-error path block on this to drain the terminator.
+ private val terminalLatch = new CountDownLatch(1)
+
+ // Captures an ErrorResponse encountered during the data phase so that
+ // the CancelResponse terminator 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)
+
+ // Cancellation intent. Kept outside the [[SessionState]] machine because a
+ // cancel can be requested before the stream exists (pre-init), where there is
+ // no wire transition to make yet -- only an intent to record. Used to (a)
+ // make cancel idempotent across all call sites, (b) fast-fail ProcessIterator
+ // on a pre-init cancel, and (c) suppress any Data/Finish that would otherwise
+ // race a Cancel onto the wire (re-read inside [[requestLock]]).
+ private val cancelRequested = new AtomicBoolean(false)
+
+ // gRPC requires serialized writes to a request StreamObserver.
+ private val requestLock = new Object
+
+ // Initialised in init() -- before that, close() is a no-op 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))
+
+ case UdfResponse.ResponseCase.CONTROL =>
+ handleControl(response.getControl)
+
+ case other =>
+ // A malformed response (empty / unknown oneof) is a terminal
+ // transport failure. Count down initLatch like every other
+ // terminal-settling path here (onError / onCompleted / each
+ // handleControl branch): if this arrives before InitResponse, doInit
+ // must fail fast with this cause rather than block until
+ // initResponseTimeoutMs and report a misleading "timed out" error.
+ completeTerminal(Termination.TransportFailed(new IllegalStateException(
+ s"unexpected response oneof: $other")))
+ initLatch.countDown()
+ }
+ }
+
+ override def onError(t: Throwable): Unit = {
+ // Transport-level failure: the stream is dead. No further writes are
+ // possible (reaching a terminal state closes the write side). Settle the
+ // terminal BEFORE counting down initLatch: init() blocks on that latch,
+ // and only await/countDown establish a happens-before edge, so a thread
+ // woken by the countDown must already be able to observe the terminal.
+ // That lets init() surface the transport cause instead of timing out
+ // after initResponseTimeoutMs with a misleading "timed out" message --
+ // or, if it saw a transient non-terminal state, the defensive "init latch
+ // fired without an InitResponse or terminal" error.
+ completeTerminal(Termination.TransportFailed(t))
+ initLatch.countDown()
+ }
+
+ 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 (!currentState.isTerminal) {
+ completeTerminal(Termination.TransportFailed(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()
+ }
+ }
+
+ /**
+ * Wakes the result iterator (blocked on [[outputQueue]]) and any thread
+ * waiting on [[terminalLatch]] when the base settles a terminal. Invoked once,
+ * by the caller that wins [[completeTerminal]].
+ */
+ override protected def onTerminalSettled(termination: Termination): Unit = {
+ outputQueue.put(QueueItem.EndOfStream)
+ terminalLatch.countDown()
+ }
+
+ private def handleControl(ctrl: UdfControlResponse): Unit = ctrl.getControlCase match {
+ case UdfControlResponse.ControlCase.INIT =>
+ val resp = ctrl.getInit
+ // Capture the InitResponse so init() can return it (or throw on its error).
+ initResponse.set(Some(resp))
+ if (resp.hasError) {
+ initLatch.countDown()
+ // Per the protocol the engine follows an init error with Cancel; send it
+ // if the request stream is up. If the Cancel cannot be written -- e.g. a
+ // directExecutor worker delivered InitResponse reentrantly, before
+ // doInit published requestObserver -- settle the terminal here so doInit
+ // does not wait the full terminalTimeoutMs for a CancelResponse that can
+ // never arrive. (When the Cancel is sent, the CancelResponse settles the
+ // terminal instead.)
+ if (!sendCancelInternal(() => cancelWithReason("init failed")) &&
+ !currentState.isTerminal) {
+ completeTerminal(Termination.Failed(resp.getError))
+ }
+ } else {
+ // InitResponse OK: init succeeded. A terminal that arrived first (e.g. a
+ // racing transport error) must win, so only advance from the expected
+ // Initializing state; process() then opens the data phase (Streaming).
+ compareAndSetState(SessionState.Initializing, SessionState.Initialized)
+ initLatch.countDown()
+ }
+
+ case UdfControlResponse.ControlCase.ERROR =>
+ val err = ctrl.getError.getError
+ executionError.compareAndSet(None, Some(err))
+ // 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
+ // (as Failed, since there is no proto terminator to carry) so doInit
+ // doesn't return as if init succeeded. "Init not yet resolved" is exactly
+ // "the stream has not reached Initialized"; completeTerminal is idempotent,
+ // so the data-phase ERROR -> Cancel -> CancelResponse path is unaffected.
+ if (!initResolved) {
+ completeTerminal(Termination.Failed(err))
+ }
+ // Surface the error to doInit. Count down AFTER settling the terminal
+ // above: init() blocks on this latch and only await/countDown establish
+ // happens-before, so a woken doInit must be able to observe the terminal
+ // rather than a transient non-terminal state. Idempotent, so the normal
+ // data-phase path (latch already fired during init) is unaffected.
+ initLatch.countDown()
+ sendCancelInternal(() => cancelWithReason("aborting after ErrorResponse"))
+
+ case UdfControlResponse.ControlCase.FINISH =>
+ // The FinishResponse carries metrics + the finish-callback data/error.
+ // Keep it on the terminal so close() can return it; the iterator inspects
+ // its error field to decide whether to throw.
+ completeTerminal(Termination.Finished(ctrl.getFinish))
+ // 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 =>
+ // The CancelResponse carries metrics + the cancel-callback error. Keep it
+ // on the terminal so close() can return it; any prior ErrorResponse is
+ // tracked in executionError and surfaced by the iterator.
+ completeTerminal(Termination.Cancelled(ctrl.getCancel))
+ // 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(Termination.TransportFailed(new IllegalStateException(
+ "empty UdfControlResponse oneof")))
+ initLatch.countDown()
+ }
+
+ /** True once init has resolved (Initialized or beyond). */
+ private def initResolved: Boolean = currentState match {
+ case SessionState.Created | SessionState.Initializing => false
+ case _ => true
+ }
+
+ private def cancelWithReason(reason: String): Cancel =
+ Cancel.newBuilder().setReason(reason).build()
+
+ // ---- WorkerSession hooks ------------------------------------------------
+
+ override protected def doInit(message: Init): InitResponse = {
+ // 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(Termination.TransportFailed(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. close()'s 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 {
+ // The base advanced Created -> Initializing before calling doInit, so a
+ // directExecutor worker's InitResponse -- delivered reentrantly inside
+ // onNext -- finds Initializing and advances to Initialized.
+ 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
+ completeTerminal(Termination.TransportFailed(e))
+ // Surface as GrpcWorkerSessionException so the engine integration layer
+ // (which catches that type and wraps it) sees a uniform init-failure
+ // exception rather than the raw transport error.
+ throw new GrpcWorkerSessionException("UDF worker stream failed during init", e)
+ }
+
+ val responded = try {
+ initLatch.await(initResponseTimeoutMs, TimeUnit.MILLISECONDS)
+ } catch {
+ case _: InterruptedException =>
+ Thread.currentThread().interrupt()
+ sendCancelInternal(() => cancelWithReason("interrupted during init"))
+ // Make sure close() does not block waiting for a terminator that
+ // will never arrive on this thread's behalf.
+ completeTerminal(Termination.TransportFailed(
+ new InterruptedException("interrupted while waiting for InitResponse")))
+ throw new InterruptedException("interrupted while waiting for InitResponse")
+ }
+
+ if (!responded) {
+ sendCancelInternal(() => cancelWithReason("InitResponse timed out"))
+ val timeout = new TimeoutException(
+ s"timed out waiting for InitResponse after ${initResponseTimeoutMs}ms")
+ // 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(Termination.TransportFailed(timeout))
+ // Surface as GrpcWorkerSessionException (carrying the timeout cause) so
+ // the engine integration layer that catches that type can wrap it.
+ throw new GrpcWorkerSessionException(
+ s"timed out waiting for InitResponse after ${initResponseTimeoutMs}ms", timeout)
+ }
+
+ initResponse.get() match {
+ case Some(resp) if resp.hasError =>
+ // 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(resp.getError)}", resp.getError)
+ case Some(resp) =>
+ resp
+ case None =>
+ // No InitResponse arrived 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.
+ currentState match {
+ case SessionState.Terminal(Termination.TransportFailed(cause)) =>
+ throw new GrpcWorkerSessionException(
+ "UDF worker stream failed during init", cause)
+ case SessionState.Terminal(Termination.Failed(err)) =>
+ throw new GrpcWorkerSessionException(
+ s"UDF worker reported an error before init completed: " +
+ describeError(err), err)
+ case SessionState.Terminal(Termination.Cancelled(_)) =>
+ throw new GrpcWorkerSessionException(
+ "UDF worker stream was cancelled before init completed")
+ case SessionState.Terminal(Termination.Finished(_)) =>
+ throw new GrpcWorkerSessionException(
+ "UDF worker finished before init completed")
+ case other =>
+ throw new IllegalStateException(
+ s"init latch fired without an InitResponse or terminal: $other")
+ }
+ }
+ }
+
+ override protected def doProcess(
+ input: Iterator[DataRequest],
+ finish: () => Finish): Iterator[DataResponse] = {
+ // Init success is guaranteed by the base [[WorkerSession]] lifecycle: if
+ // doInit had failed it would have thrown and process() would never run.
+ new ProcessIterator(input, finish)
+ }
+
+ override protected def doClose(cancel: () => Cancel): Termination = {
+ if (requestObserver == null) {
+ // init() never put a stream on the wire (closed before/around init, or
+ // init threw before publishing). There is no protocol terminator; treat
+ // the session as cancelled-before-start. The base WorkerSession still
+ // releases the worker handle, so the worker is torn down.
+ completeTerminal(Termination.Cancelled(CancelResponse.getDefaultInstance))
+ return Termination.Cancelled(CancelResponse.getDefaultInstance)
+ }
+ // If the stream has not finished on its own, cancel anything in flight so
+ // the worker can clean up, then wait for the terminator. sendCancelInternal
+ // evaluates the cancel thunk only when it actually writes a Cancel, and at
+ // most once across all callers.
+ if (!currentState.isTerminal) {
+ sendCancelInternal(cancel)
+ 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 the state sees a stable value.
+ if (!currentState.isTerminal) {
+ completeTerminal(Termination.TransportFailed(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. onCompleted on an
+ // already-torn-down stream throws and is caught, so this stays idempotent
+ // across repeat close() calls.
+ try {
+ requestLock.synchronized { requestObserver.onCompleted() }
+ } catch {
+ case NonFatal(e) =>
+ logger.debug("Error half-closing UDF Execute stream", e)
+ }
+ // Derive the Termination from the settled terminal. close() is the
+ // finalizer/cleanup path and must NOT re-throw: a UDF / data-phase error is
+ // already surfaced while the result iterator is consumed, and an init error
+ // from init(). settledTermination returns the response proto for the clean
+ // terminators and a best-effort empty Cancelled for the failure terminals.
+ settledTermination
+ }
+
+ // ---- Internal request helpers ---------------------------------------------
+
+ /**
+ * Sends a Data or Finish request to the worker. Two invariants are
+ * checked inside [[requestLock]]:
+ * - terminal state: 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 (currentState.isTerminal) {
+ throw new IllegalStateException(
+ "cannot send request: UDF Execute stream is already closed")
+ }
+ // Suppress the write when a cancel has been (or is about to be) flushed
+ // through this lock; that preserves the proto ordering invariant (no
+ // Data/Finish after Cancel). We test the flag rather than `return`-ing
+ // from the block: a non-local return out of this by-name `synchronized`
+ // body compiles to a thrown NonLocalReturnControl, which is brittle.
+ if (!cancelRequested.get()) {
+ requestObserver.onNext(req)
+ }
+ }
+
+ /**
+ * Sends a `Cancel` control message, returning `true` iff a `Cancel` was
+ * actually written to the wire. Idempotent across ALL call sites (close()'s
+ * cancel, in-band cancels from `handleControl`'s INIT error / ERROR paths,
+ * and engine-internal cancels from the iterator) via [[cancelRequested]]:
+ * only the first caller puts a `Cancel` on the wire, and the `cancel` thunk is
+ * evaluated only by that caller and only when the message is actually sent --
+ * so a side-effecting thunk (e.g. one carrying a client cancel callback) runs
+ * at most once. 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(cancel: () => Cancel): Boolean = {
+ if (!cancelRequested.compareAndSet(false, true)) return false
+ if (requestObserver == null) return false
+ if (currentState.isTerminal) return false
+ try {
+ requestLock.synchronized {
+ // A terminal that arrived first must win; bail without writing. We
+ // compute the block's value rather than `return`-ing from this by-name
+ // `synchronized` body, whose non-local return would compile to a thrown
+ // NonLocalReturnControl that the surrounding NonFatal catch must not
+ // swallow -- brittle to rely on. Reads naturally as the block result.
+ val cur = currentState
+ if (cur.isTerminal) {
+ false
+ } else {
+ // Record that a Cancel is on the wire by advancing to Cancelling.
+ compareAndSetState(cur, SessionState.Cancelling)
+ requestObserver.onNext(UdfRequest.newBuilder()
+ .setControl(UdfControlRequest.newBuilder().setCancel(cancel()).build())
+ .build())
+ true
+ }
+ }
+ } catch {
+ case NonFatal(e) =>
+ logger.debug(s"Cancel send failed (stream may already be torn down): ${e.getMessage}")
+ completeTerminal(Termination.TransportFailed(e))
+ false
+ }
+ }
+
+ private def awaitTerminal(): Unit = {
+ if (currentState.isTerminal) return
+ try {
+ if (!terminalLatch.await(terminalTimeoutMs, TimeUnit.MILLISECONDS)) {
+ completeTerminal(Termination.TransportFailed(new TimeoutException(
+ s"timed out waiting for stream terminator after ${terminalTimeoutMs}ms")))
+ }
+ } catch {
+ case _: InterruptedException =>
+ // Attribute the terminal to the interrupt, not a timeout: the wait was
+ // cut short, it did not exhaust terminalTimeoutMs. Reporting "timed
+ // out" here would misdiagnose an interrupted close as a slow worker.
+ Thread.currentThread().interrupt()
+ completeTerminal(Termination.TransportFailed(
+ new InterruptedException("interrupted while waiting for stream terminator")))
+ }
+ }
+
+ // ---- 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 settle the state) and with any thread that finalizes via
+ * [[close]]. It drives the `Streaming -> Finishing` transition (sending
+ * `Finish` once input is exhausted, built lazily from the `finish` thunk).
+ */
Review Comment:
fixed
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57640] Implement gRPC WorkerSession and test it against an echo server. [spark]
Posted by "Yicong-Huang (via GitHub)" <gi...@apache.org>.
Yicong-Huang commented on code in PR #56702:
URL: https://github.com/apache/spark/pull/56702#discussion_r3549265974
##########
udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/GrpcWorkerSession.scala:
##########
@@ -0,0 +1,816 @@
+/*
+ * 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.grpc
+
+import java.util.concurrent.{CountDownLatch, LinkedBlockingQueue, TimeoutException, TimeUnit}
+import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference}
+
+import scala.util.control.NonFatal
+
+import io.grpc.{ConnectivityState, ManagedChannel}
+import io.grpc.stub.StreamObserver
+
+import org.apache.spark.annotation.Experimental
+import org.apache.spark.udf.worker.{Cancel, CancelResponse, DataRequest, DataResponse,
+ ExecutionError, Finish, Init, InitResponse, UdfControlRequest,
+ UdfControlResponse, UdfRequest, UdfResponse, UdfWorkerGrpc}
+import org.apache.spark.udf.worker.core.{Termination, WorkerHandle, WorkerLogger, WorkerSession}
+import org.apache.spark.udf.worker.core.WorkerSession.SessionState
+import org.apache.spark.udf.worker.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_message.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.
+ *
+ * '''Driving model.''' Consumption-driven (Volcano / pull): the thread that
+ * consumes the [[doProcess]] result iterator is the one that pulls the next
+ * input batch and sends its `DataRequest`. Nothing is sent until the iterator
+ * is consumed, and input flows one batch per output pull. The gRPC callback
+ * thread only receives output.
+ *
+ * '''State machine.''' This class does not keep its own state machine: it
+ * drives the single [[WorkerSession.SessionState]] owned by the base. The base
+ * advances `Created -> Initializing` (in `init`) and `Initialized -> Streaming`
+ * (in `process`); this class advances the protocol-event edges through
+ * [[compareAndSetState]] / [[completeTerminal]] as it exchanges messages:
+ * {{{
+ * Initializing --(InitResponse ok)--> Initialized [handleControl]
+ * Streaming ----(Finish written)----> Finishing [ProcessIterator]
+ * <any non-terminal> --(Cancel written)--> Cancelling [sendCancelInternal]
+ * <any non-terminal> --(terminator/error)--> terminal [completeTerminal]
+ * }}}
+ * The two clean terminals carry the worker's `FinishResponse` / `CancelResponse`
+ * (metrics + finish/cancel callback `data`/`error`) so [[close]] can return
+ * them. The only flag kept outside the machine is [[cancelRequested]]: a
+ * cancellation can be requested before the stream exists (so it cannot be a
+ * state transition yet), and it must both fast-fail the result iterator and
+ * suppress any in-flight Data/Finish.
+ *
+ * Threading:
+ * - [[doInit]] is synchronous: sends `Init` and blocks on `InitResponse`,
+ * returning it.
+ * - [[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.
+ * - [[doClose]] is thread-safe and idempotent: it settles + returns the
+ * terminator (cancelling in-flight work if the stream had not finished)
+ * and half-closes the request side.
+ *
+ * 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 a gRPC channel built and owned by the caller (the
+ * dispatcher). 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 from the worker, drained by the process() iterator.
+ // Intentionally unbounded: a bounded queue would block the gRPC callback
+ // (Netty event-loop) thread when full, stalling control-message and
+ // terminator delivery on the whole channel. HTTP/2 flow control bounds the
+ // wire and the consumer normally drains promptly, so it stays small. A
+ // stalled downstream can still grow it; the real fix is protocol-level
+ // back-pressure (out of scope here), not bounding the queue.
+ //
+ // TODO [SPARK-57324]: expose queue depth as a metric (early warning for a
+ // stalled consumer).
+ 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)
+
+ // The InitResponse the worker sent (success or error), captured so init()
+ // can return it. None until InitResponse arrives.
+ private val initResponse = new AtomicReference[Option[InitResponse]](None)
Review Comment:
I wonder if it is a good idea to use a latch + an atomic reference over a `Future`? I see you want to make sure it can only fire once, but the two values are decoupled and one needs to remember to count down the latch before access the reference. Does it make sense to wrap it with something like:
```
private final class OneShotValue[A] {
private val latch = new CountDownLatch(1)
private val value = new AtomicReference[Option[A]](None)
def complete(v: A): Unit = {
value.compareAndSet(None, Some(v))
latch.countDown()
}
def signalWithoutValue(): Unit = latch.countDown()
def await(timeoutMs: Long): Option[A] = {
if (!latch.await(timeoutMs, TimeUnit.MILLISECONDS)) None else value.get()
}
}
```
##########
udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/GrpcWorkerSession.scala:
##########
@@ -0,0 +1,816 @@
+/*
+ * 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.grpc
+
+import java.util.concurrent.{CountDownLatch, LinkedBlockingQueue, TimeoutException, TimeUnit}
+import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference}
+
+import scala.util.control.NonFatal
+
+import io.grpc.{ConnectivityState, ManagedChannel}
+import io.grpc.stub.StreamObserver
+
+import org.apache.spark.annotation.Experimental
+import org.apache.spark.udf.worker.{Cancel, CancelResponse, DataRequest, DataResponse,
+ ExecutionError, Finish, Init, InitResponse, UdfControlRequest,
+ UdfControlResponse, UdfRequest, UdfResponse, UdfWorkerGrpc}
+import org.apache.spark.udf.worker.core.{Termination, WorkerHandle, WorkerLogger, WorkerSession}
+import org.apache.spark.udf.worker.core.WorkerSession.SessionState
+import org.apache.spark.udf.worker.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_message.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.
+ *
+ * '''Driving model.''' Consumption-driven (Volcano / pull): the thread that
+ * consumes the [[doProcess]] result iterator is the one that pulls the next
+ * input batch and sends its `DataRequest`. Nothing is sent until the iterator
+ * is consumed, and input flows one batch per output pull. The gRPC callback
+ * thread only receives output.
+ *
+ * '''State machine.''' This class does not keep its own state machine: it
+ * drives the single [[WorkerSession.SessionState]] owned by the base. The base
+ * advances `Created -> Initializing` (in `init`) and `Initialized -> Streaming`
+ * (in `process`); this class advances the protocol-event edges through
+ * [[compareAndSetState]] / [[completeTerminal]] as it exchanges messages:
+ * {{{
+ * Initializing --(InitResponse ok)--> Initialized [handleControl]
+ * Streaming ----(Finish written)----> Finishing [ProcessIterator]
+ * <any non-terminal> --(Cancel written)--> Cancelling [sendCancelInternal]
+ * <any non-terminal> --(terminator/error)--> terminal [completeTerminal]
+ * }}}
+ * The two clean terminals carry the worker's `FinishResponse` / `CancelResponse`
+ * (metrics + finish/cancel callback `data`/`error`) so [[close]] can return
+ * them. The only flag kept outside the machine is [[cancelRequested]]: a
+ * cancellation can be requested before the stream exists (so it cannot be a
+ * state transition yet), and it must both fast-fail the result iterator and
+ * suppress any in-flight Data/Finish.
+ *
+ * Threading:
+ * - [[doInit]] is synchronous: sends `Init` and blocks on `InitResponse`,
+ * returning it.
+ * - [[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.
+ * - [[doClose]] is thread-safe and idempotent: it settles + returns the
+ * terminator (cancelling in-flight work if the stream had not finished)
+ * and half-closes the request side.
+ *
+ * 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 a gRPC channel built and owned by the caller (the
+ * dispatcher). 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 from the worker, drained by the process() iterator.
+ // Intentionally unbounded: a bounded queue would block the gRPC callback
+ // (Netty event-loop) thread when full, stalling control-message and
+ // terminator delivery on the whole channel. HTTP/2 flow control bounds the
+ // wire and the consumer normally drains promptly, so it stays small. A
+ // stalled downstream can still grow it; the real fix is protocol-level
+ // back-pressure (out of scope here), not bounding the queue.
+ //
+ // TODO [SPARK-57324]: expose queue depth as a metric (early warning for a
+ // stalled consumer).
+ 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)
+
+ // The InitResponse the worker sent (success or error), captured so init()
+ // can return it. None until InitResponse arrives.
+ private val initResponse = new AtomicReference[Option[InitResponse]](None)
+
+ // Fired when the session reaches a terminal [[SessionState]]. doClose() and
+ // the init-error path block on this to drain the terminator.
+ private val terminalLatch = new CountDownLatch(1)
+
+ // Captures an ErrorResponse encountered during the data phase so that
+ // the CancelResponse terminator 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)
+
+ // Cancellation intent. Kept outside the [[SessionState]] machine because a
+ // cancel can be requested before the stream exists (pre-init), where there is
+ // no wire transition to make yet -- only an intent to record. Used to (a)
+ // make cancel idempotent across all call sites, (b) fast-fail ProcessIterator
+ // on a pre-init cancel, and (c) suppress any Data/Finish that would otherwise
+ // race a Cancel onto the wire (re-read inside [[requestLock]]).
+ private val cancelRequested = new AtomicBoolean(false)
+
+ // gRPC requires serialized writes to a request StreamObserver.
+ private val requestLock = new Object
+
+ // Initialised in init() -- before that, close() is a no-op 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))
+
+ case UdfResponse.ResponseCase.CONTROL =>
+ handleControl(response.getControl)
+
+ case other =>
+ // A malformed response (empty / unknown oneof) is a terminal
+ // transport failure. Count down initLatch like every other
+ // terminal-settling path here (onError / onCompleted / each
+ // handleControl branch): if this arrives before InitResponse, doInit
+ // must fail fast with this cause rather than block until
+ // initResponseTimeoutMs and report a misleading "timed out" error.
+ completeTerminal(Termination.TransportFailed(new IllegalStateException(
+ s"unexpected response oneof: $other")))
+ initLatch.countDown()
+ }
+ }
+
+ override def onError(t: Throwable): Unit = {
+ // Transport-level failure: the stream is dead. No further writes are
+ // possible (reaching a terminal state closes the write side). Settle the
+ // terminal BEFORE counting down initLatch: init() blocks on that latch,
+ // and only await/countDown establish a happens-before edge, so a thread
+ // woken by the countDown must already be able to observe the terminal.
+ // That lets init() surface the transport cause instead of timing out
+ // after initResponseTimeoutMs with a misleading "timed out" message --
+ // or, if it saw a transient non-terminal state, the defensive "init latch
+ // fired without an InitResponse or terminal" error.
+ completeTerminal(Termination.TransportFailed(t))
+ initLatch.countDown()
+ }
+
+ 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 (!currentState.isTerminal) {
+ completeTerminal(Termination.TransportFailed(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()
+ }
+ }
+
+ /**
+ * Wakes the result iterator (blocked on [[outputQueue]]) and any thread
+ * waiting on [[terminalLatch]] when the base settles a terminal. Invoked once,
+ * by the caller that wins [[completeTerminal]].
+ */
+ override protected def onTerminalSettled(termination: Termination): Unit = {
+ outputQueue.put(QueueItem.EndOfStream)
+ terminalLatch.countDown()
+ }
+
+ private def handleControl(ctrl: UdfControlResponse): Unit = ctrl.getControlCase match {
+ case UdfControlResponse.ControlCase.INIT =>
+ val resp = ctrl.getInit
+ // Capture the InitResponse so init() can return it (or throw on its error).
+ initResponse.set(Some(resp))
+ if (resp.hasError) {
+ initLatch.countDown()
+ // Per the protocol the engine follows an init error with Cancel; send it
+ // if the request stream is up. If the Cancel cannot be written -- e.g. a
+ // directExecutor worker delivered InitResponse reentrantly, before
+ // doInit published requestObserver -- settle the terminal here so doInit
+ // does not wait the full terminalTimeoutMs for a CancelResponse that can
+ // never arrive. (When the Cancel is sent, the CancelResponse settles the
+ // terminal instead.)
+ if (!sendCancelInternal(() => cancelWithReason("init failed")) &&
Review Comment:
The INIT-error branch never records `executionError`, and if the Cancel is written the worker's CancelResponse settles the terminal as `Cancelled` (empty) -- so `close()` returns a bare `Termination.Cancelled` with the init error available only via the thrown exception. This is asymmetric with the ERROR branch's `Failed(err)`. The "close ... returns a Failed termination carrying the error" test only covers the ErrorResponse variant, so this gap is untested. Intended? If not, carry the error onto the terminal.
##########
udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/GrpcWorkerSession.scala:
##########
@@ -0,0 +1,816 @@
+/*
+ * 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.grpc
+
+import java.util.concurrent.{CountDownLatch, LinkedBlockingQueue, TimeoutException, TimeUnit}
+import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference}
+
+import scala.util.control.NonFatal
+
+import io.grpc.{ConnectivityState, ManagedChannel}
+import io.grpc.stub.StreamObserver
+
+import org.apache.spark.annotation.Experimental
+import org.apache.spark.udf.worker.{Cancel, CancelResponse, DataRequest, DataResponse,
+ ExecutionError, Finish, Init, InitResponse, UdfControlRequest,
+ UdfControlResponse, UdfRequest, UdfResponse, UdfWorkerGrpc}
+import org.apache.spark.udf.worker.core.{Termination, WorkerHandle, WorkerLogger, WorkerSession}
+import org.apache.spark.udf.worker.core.WorkerSession.SessionState
+import org.apache.spark.udf.worker.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_message.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.
+ *
+ * '''Driving model.''' Consumption-driven (Volcano / pull): the thread that
+ * consumes the [[doProcess]] result iterator is the one that pulls the next
+ * input batch and sends its `DataRequest`. Nothing is sent until the iterator
+ * is consumed, and input flows one batch per output pull. The gRPC callback
+ * thread only receives output.
+ *
+ * '''State machine.''' This class does not keep its own state machine: it
+ * drives the single [[WorkerSession.SessionState]] owned by the base. The base
+ * advances `Created -> Initializing` (in `init`) and `Initialized -> Streaming`
+ * (in `process`); this class advances the protocol-event edges through
+ * [[compareAndSetState]] / [[completeTerminal]] as it exchanges messages:
+ * {{{
+ * Initializing --(InitResponse ok)--> Initialized [handleControl]
+ * Streaming ----(Finish written)----> Finishing [ProcessIterator]
+ * <any non-terminal> --(Cancel written)--> Cancelling [sendCancelInternal]
+ * <any non-terminal> --(terminator/error)--> terminal [completeTerminal]
+ * }}}
+ * The two clean terminals carry the worker's `FinishResponse` / `CancelResponse`
+ * (metrics + finish/cancel callback `data`/`error`) so [[close]] can return
+ * them. The only flag kept outside the machine is [[cancelRequested]]: a
+ * cancellation can be requested before the stream exists (so it cannot be a
+ * state transition yet), and it must both fast-fail the result iterator and
+ * suppress any in-flight Data/Finish.
+ *
+ * Threading:
+ * - [[doInit]] is synchronous: sends `Init` and blocks on `InitResponse`,
+ * returning it.
+ * - [[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.
+ * - [[doClose]] is thread-safe and idempotent: it settles + returns the
+ * terminator (cancelling in-flight work if the stream had not finished)
+ * and half-closes the request side.
+ *
+ * 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 a gRPC channel built and owned by the caller (the
+ * dispatcher). 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 from the worker, drained by the process() iterator.
+ // Intentionally unbounded: a bounded queue would block the gRPC callback
+ // (Netty event-loop) thread when full, stalling control-message and
+ // terminator delivery on the whole channel. HTTP/2 flow control bounds the
+ // wire and the consumer normally drains promptly, so it stays small. A
+ // stalled downstream can still grow it; the real fix is protocol-level
+ // back-pressure (out of scope here), not bounding the queue.
+ //
+ // TODO [SPARK-57324]: expose queue depth as a metric (early warning for a
+ // stalled consumer).
+ 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)
+
+ // The InitResponse the worker sent (success or error), captured so init()
+ // can return it. None until InitResponse arrives.
+ private val initResponse = new AtomicReference[Option[InitResponse]](None)
+
+ // Fired when the session reaches a terminal [[SessionState]]. doClose() and
+ // the init-error path block on this to drain the terminator.
+ private val terminalLatch = new CountDownLatch(1)
+
+ // Captures an ErrorResponse encountered during the data phase so that
+ // the CancelResponse terminator 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)
+
+ // Cancellation intent. Kept outside the [[SessionState]] machine because a
+ // cancel can be requested before the stream exists (pre-init), where there is
+ // no wire transition to make yet -- only an intent to record. Used to (a)
+ // make cancel idempotent across all call sites, (b) fast-fail ProcessIterator
+ // on a pre-init cancel, and (c) suppress any Data/Finish that would otherwise
+ // race a Cancel onto the wire (re-read inside [[requestLock]]).
+ private val cancelRequested = new AtomicBoolean(false)
+
+ // gRPC requires serialized writes to a request StreamObserver.
+ private val requestLock = new Object
+
+ // Initialised in init() -- before that, close() is a no-op 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))
Review Comment:
I know this is hot path, but if we receive data batch before Init response, it could result in hang here. I think we need a way to prevent it from happening.
##########
udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/GrpcWorkerSession.scala:
##########
@@ -0,0 +1,816 @@
+/*
+ * 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.grpc
+
+import java.util.concurrent.{CountDownLatch, LinkedBlockingQueue, TimeoutException, TimeUnit}
+import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference}
+
+import scala.util.control.NonFatal
+
+import io.grpc.{ConnectivityState, ManagedChannel}
+import io.grpc.stub.StreamObserver
+
+import org.apache.spark.annotation.Experimental
+import org.apache.spark.udf.worker.{Cancel, CancelResponse, DataRequest, DataResponse,
+ ExecutionError, Finish, Init, InitResponse, UdfControlRequest,
+ UdfControlResponse, UdfRequest, UdfResponse, UdfWorkerGrpc}
+import org.apache.spark.udf.worker.core.{Termination, WorkerHandle, WorkerLogger, WorkerSession}
+import org.apache.spark.udf.worker.core.WorkerSession.SessionState
+import org.apache.spark.udf.worker.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_message.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.
+ *
+ * '''Driving model.''' Consumption-driven (Volcano / pull): the thread that
+ * consumes the [[doProcess]] result iterator is the one that pulls the next
+ * input batch and sends its `DataRequest`. Nothing is sent until the iterator
Review Comment:
"input flows one batch per output pull" seems to be overstating: `advance()` branch (2) sends the next input whenever the output queue is momentarily empty, so under async delivery it can push many input batches before any output is read (bounded only by HTTP/2 flow control).
is my understanding correct?
##########
udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/GrpcWorkerSessionConcurrencySuite.scala:
##########
@@ -0,0 +1,742 @@
+/*
+ * 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.grpc
+
+import java.util.Locale
+import java.util.concurrent.{ConcurrentLinkedQueue, CountDownLatch, TimeUnit}
+import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference}
+
+import scala.jdk.CollectionConverters._
+
+import com.google.protobuf.ByteString
+import io.grpc.{ManagedChannel, Server}
+import io.grpc.inprocess.{InProcessChannelBuilder, InProcessServerBuilder}
+import io.grpc.stub.StreamObserver
+import org.scalatest.BeforeAndAfterEach
+// scalastyle:off funsuite
+import org.scalatest.funsuite.AnyFunSuite
+
+import org.apache.spark.udf.worker.{Cancel, CancelResponse, DataRequest, DataResponse,
+ ErrorResponse, ExecutionError, Finish, FinishResponse, Init, InitResponse, UdfControlResponse,
+ UdfPayload, UdfRequest, UdfResponse, UDFWorkerDataFormat, UdfWorkerGrpc, UserError,
+ WorkerRequest, WorkerResponse}
+import org.apache.spark.udf.worker.core.{Termination, WorkerHandle, WorkerLogger}
+
+/**
+ * Concurrency tests for [[GrpcWorkerSession]] that pin the wire-ordering and
+ * fast-fail invariants under concurrent and worker-misbehavior scenarios:
+ * - Cancel must never appear on the wire before Init.
+ * - A worker terminator (ERROR / FINISH / CANCEL / onCompleted) arriving
+ * before InitResponse must fail [[GrpcWorkerSession#init]] fast, not hang
+ * for `initResponseTimeoutMs`.
+ * - Repeated [[Iterator#hasNext]] after natural iterator exhaustion must
+ * return immediately, not block for `terminalTimeoutMs`.
+ * - Cancel issued before [[GrpcWorkerSession#init]] makes a subsequent
+ * [[GrpcWorkerSession#doProcess]] surface the cancellation immediately,
+ * instead of blocking for `terminalTimeoutMs`.
+ * - Cancel / close concurrent with an in-progress data phase terminate
+ * cleanly without leaks or unbounded hangs.
+ *
+ * Runs entirely in-process: no subprocess, no UDS. Server services are
+ * custom-built per test so we can drive specific worker misbehavior.
+ */
+class GrpcWorkerSessionConcurrencySuite
+ extends AnyFunSuite with BeforeAndAfterEach {
+// scalastyle:on funsuite
+
+ /** Used by tests to keep stale in-flight infra reachable for teardown. */
+ private val openServers = new ConcurrentLinkedQueue[Server]()
+ private val openChannels = new ConcurrentLinkedQueue[ManagedChannel]()
+ private val openSessions = new ConcurrentLinkedQueue[GrpcWorkerSession]()
+
+ override def afterEach(): Unit = {
+ // Shut channels down first. This fires onError on any still-live stream,
+ // which settles the session terminal and counts down the init/terminal
+ // latches. That unblocks both the session.close() below and any worker
+ // thread a failing test left parked on a (deliberately large) timeout, so
+ // teardown never hangs even when a test asserts via assertFinishesWithin.
+ openChannels.asScala.foreach { c =>
+ try c.shutdownNow().awaitTermination(2, TimeUnit.SECONDS) catch { case _: Throwable => () }
+ }
+ openChannels.clear()
+ openSessions.asScala.foreach { s => try s.close(emptyCancel) catch { case _: Throwable => () } }
+ openSessions.clear()
+ openServers.asScala.foreach { s =>
+ try s.shutdownNow().awaitTermination(2, TimeUnit.SECONDS) catch { case _: Throwable => () }
+ }
+ openServers.clear()
+ super.afterEach()
+ }
+
+ // A session timeout large enough that a correct test never reaches it; a
+ // regression that fails to short-circuit blocks here for minutes and is caught
+ // by assertFinishesWithin (below) instead of a flaky `elapsed < timeout` bound.
+ private val NeverReachedTimeoutMs = TimeUnit.MINUTES.toMillis(10)
+
+ /**
+ * Runs `body` on a daemon thread and asserts it finishes within `withinMs`,
+ * rethrowing whatever `body` threw (so an `intercept` inside `body` still
+ * works). Pair with [[NeverReachedTimeoutMs]]: the correct fast path returns
+ * in milliseconds, so `withinMs` (seconds) has an enormous safety margin and
+ * does not flake, while a regression that parks on the timeout never returns
+ * within `withinMs` and fails the assertion.
+ */
+ private def assertFinishesWithin(withinMs: Long, name: String)(body: => Unit): Unit = {
+ val thrown = new AtomicReference[Throwable]()
+ val done = new CountDownLatch(1)
+ val worker = new Thread(() => {
+ try body catch { case t: Throwable => thrown.set(t) } finally done.countDown()
+ }, name)
+ worker.setDaemon(true)
+ worker.start()
+ assert(done.await(withinMs, TimeUnit.MILLISECONDS),
+ s"$name did not finish within ${withinMs}ms; it parked on a session timeout " +
+ "that the fast path should have short-circuited")
+ Option(thrown.get()).foreach(t => throw t)
+ }
+
+ // ---------------------------------------------------------------------------
+ // Helpers
+ // ---------------------------------------------------------------------------
+
+ private class TestWorkerHandle extends WorkerHandle {
+ val invalidated = new AtomicBoolean(false)
Review Comment:
`invalidated`/`released` are tracked but no test asserts them. The handle lifecycle (release exactly once; `markInvalid` on a non-salvageable `TransportFailed`, not on a clean finish) is a core `WorkerSession` contract that goes unverified here -- add assertions on these two flags.
##########
udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/GrpcWorkerSession.scala:
##########
@@ -0,0 +1,816 @@
+/*
+ * 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.grpc
+
+import java.util.concurrent.{CountDownLatch, LinkedBlockingQueue, TimeoutException, TimeUnit}
+import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference}
+
+import scala.util.control.NonFatal
+
+import io.grpc.{ConnectivityState, ManagedChannel}
+import io.grpc.stub.StreamObserver
+
+import org.apache.spark.annotation.Experimental
+import org.apache.spark.udf.worker.{Cancel, CancelResponse, DataRequest, DataResponse,
+ ExecutionError, Finish, Init, InitResponse, UdfControlRequest,
+ UdfControlResponse, UdfRequest, UdfResponse, UdfWorkerGrpc}
+import org.apache.spark.udf.worker.core.{Termination, WorkerHandle, WorkerLogger, WorkerSession}
+import org.apache.spark.udf.worker.core.WorkerSession.SessionState
+import org.apache.spark.udf.worker.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_message.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.
+ *
+ * '''Driving model.''' Consumption-driven (Volcano / pull): the thread that
+ * consumes the [[doProcess]] result iterator is the one that pulls the next
+ * input batch and sends its `DataRequest`. Nothing is sent until the iterator
+ * is consumed, and input flows one batch per output pull. The gRPC callback
+ * thread only receives output.
+ *
+ * '''State machine.''' This class does not keep its own state machine: it
+ * drives the single [[WorkerSession.SessionState]] owned by the base. The base
+ * advances `Created -> Initializing` (in `init`) and `Initialized -> Streaming`
+ * (in `process`); this class advances the protocol-event edges through
+ * [[compareAndSetState]] / [[completeTerminal]] as it exchanges messages:
+ * {{{
+ * Initializing --(InitResponse ok)--> Initialized [handleControl]
+ * Streaming ----(Finish written)----> Finishing [ProcessIterator]
+ * <any non-terminal> --(Cancel written)--> Cancelling [sendCancelInternal]
+ * <any non-terminal> --(terminator/error)--> terminal [completeTerminal]
+ * }}}
+ * The two clean terminals carry the worker's `FinishResponse` / `CancelResponse`
+ * (metrics + finish/cancel callback `data`/`error`) so [[close]] can return
+ * them. The only flag kept outside the machine is [[cancelRequested]]: a
+ * cancellation can be requested before the stream exists (so it cannot be a
+ * state transition yet), and it must both fast-fail the result iterator and
+ * suppress any in-flight Data/Finish.
+ *
+ * Threading:
+ * - [[doInit]] is synchronous: sends `Init` and blocks on `InitResponse`,
+ * returning it.
+ * - [[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.
+ * - [[doClose]] is thread-safe and idempotent: it settles + returns the
+ * terminator (cancelling in-flight work if the stream had not finished)
+ * and half-closes the request side.
+ *
+ * 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 a gRPC channel built and owned by the caller (the
+ * dispatcher). 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 from the worker, drained by the process() iterator.
+ // Intentionally unbounded: a bounded queue would block the gRPC callback
+ // (Netty event-loop) thread when full, stalling control-message and
+ // terminator delivery on the whole channel. HTTP/2 flow control bounds the
+ // wire and the consumer normally drains promptly, so it stays small. A
+ // stalled downstream can still grow it; the real fix is protocol-level
+ // back-pressure (out of scope here), not bounding the queue.
+ //
+ // TODO [SPARK-57324]: expose queue depth as a metric (early warning for a
+ // stalled consumer).
+ 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)
+
+ // The InitResponse the worker sent (success or error), captured so init()
+ // can return it. None until InitResponse arrives.
+ private val initResponse = new AtomicReference[Option[InitResponse]](None)
+
+ // Fired when the session reaches a terminal [[SessionState]]. doClose() and
+ // the init-error path block on this to drain the terminator.
+ private val terminalLatch = new CountDownLatch(1)
+
+ // Captures an ErrorResponse encountered during the data phase so that
+ // the CancelResponse terminator 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)
+
+ // Cancellation intent. Kept outside the [[SessionState]] machine because a
+ // cancel can be requested before the stream exists (pre-init), where there is
+ // no wire transition to make yet -- only an intent to record. Used to (a)
+ // make cancel idempotent across all call sites, (b) fast-fail ProcessIterator
+ // on a pre-init cancel, and (c) suppress any Data/Finish that would otherwise
+ // race a Cancel onto the wire (re-read inside [[requestLock]]).
+ private val cancelRequested = new AtomicBoolean(false)
+
+ // gRPC requires serialized writes to a request StreamObserver.
+ private val requestLock = new Object
+
+ // Initialised in init() -- before that, close() is a no-op 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))
+
+ case UdfResponse.ResponseCase.CONTROL =>
+ handleControl(response.getControl)
+
+ case other =>
+ // A malformed response (empty / unknown oneof) is a terminal
+ // transport failure. Count down initLatch like every other
+ // terminal-settling path here (onError / onCompleted / each
+ // handleControl branch): if this arrives before InitResponse, doInit
+ // must fail fast with this cause rather than block until
+ // initResponseTimeoutMs and report a misleading "timed out" error.
+ completeTerminal(Termination.TransportFailed(new IllegalStateException(
+ s"unexpected response oneof: $other")))
+ initLatch.countDown()
+ }
+ }
+
+ override def onError(t: Throwable): Unit = {
+ // Transport-level failure: the stream is dead. No further writes are
+ // possible (reaching a terminal state closes the write side). Settle the
+ // terminal BEFORE counting down initLatch: init() blocks on that latch,
+ // and only await/countDown establish a happens-before edge, so a thread
+ // woken by the countDown must already be able to observe the terminal.
+ // That lets init() surface the transport cause instead of timing out
+ // after initResponseTimeoutMs with a misleading "timed out" message --
+ // or, if it saw a transient non-terminal state, the defensive "init latch
+ // fired without an InitResponse or terminal" error.
+ completeTerminal(Termination.TransportFailed(t))
+ initLatch.countDown()
+ }
+
+ 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 (!currentState.isTerminal) {
+ completeTerminal(Termination.TransportFailed(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()
+ }
+ }
+
+ /**
+ * Wakes the result iterator (blocked on [[outputQueue]]) and any thread
+ * waiting on [[terminalLatch]] when the base settles a terminal. Invoked once,
+ * by the caller that wins [[completeTerminal]].
+ */
+ override protected def onTerminalSettled(termination: Termination): Unit = {
+ outputQueue.put(QueueItem.EndOfStream)
+ terminalLatch.countDown()
+ }
+
+ private def handleControl(ctrl: UdfControlResponse): Unit = ctrl.getControlCase match {
+ case UdfControlResponse.ControlCase.INIT =>
+ val resp = ctrl.getInit
+ // Capture the InitResponse so init() can return it (or throw on its error).
+ initResponse.set(Some(resp))
+ if (resp.hasError) {
+ initLatch.countDown()
+ // Per the protocol the engine follows an init error with Cancel; send it
+ // if the request stream is up. If the Cancel cannot be written -- e.g. a
+ // directExecutor worker delivered InitResponse reentrantly, before
+ // doInit published requestObserver -- settle the terminal here so doInit
+ // does not wait the full terminalTimeoutMs for a CancelResponse that can
+ // never arrive. (When the Cancel is sent, the CancelResponse settles the
+ // terminal instead.)
+ if (!sendCancelInternal(() => cancelWithReason("init failed")) &&
+ !currentState.isTerminal) {
+ completeTerminal(Termination.Failed(resp.getError))
+ }
+ } else {
+ // InitResponse OK: init succeeded. A terminal that arrived first (e.g. a
+ // racing transport error) must win, so only advance from the expected
+ // Initializing state; process() then opens the data phase (Streaming).
+ compareAndSetState(SessionState.Initializing, SessionState.Initialized)
+ initLatch.countDown()
+ }
+
+ case UdfControlResponse.ControlCase.ERROR =>
+ val err = ctrl.getError.getError
+ executionError.compareAndSet(None, Some(err))
+ // 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
+ // (as Failed, since there is no proto terminator to carry) so doInit
+ // doesn't return as if init succeeded. "Init not yet resolved" is exactly
+ // "the stream has not reached Initialized"; completeTerminal is idempotent,
+ // so the data-phase ERROR -> Cancel -> CancelResponse path is unaffected.
+ if (!initResolved) {
+ completeTerminal(Termination.Failed(err))
+ }
+ // Surface the error to doInit. Count down AFTER settling the terminal
+ // above: init() blocks on this latch and only await/countDown establish
+ // happens-before, so a woken doInit must be able to observe the terminal
+ // rather than a transient non-terminal state. Idempotent, so the normal
+ // data-phase path (latch already fired during init) is unaffected.
+ initLatch.countDown()
+ sendCancelInternal(() => cancelWithReason("aborting after ErrorResponse"))
+
+ case UdfControlResponse.ControlCase.FINISH =>
+ // The FinishResponse carries metrics + the finish-callback data/error.
+ // Keep it on the terminal so close() can return it; the iterator inspects
+ // its error field to decide whether to throw.
+ completeTerminal(Termination.Finished(ctrl.getFinish))
+ // 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 =>
+ // The CancelResponse carries metrics + the cancel-callback error. Keep it
+ // on the terminal so close() can return it; any prior ErrorResponse is
+ // tracked in executionError and surfaced by the iterator.
+ completeTerminal(Termination.Cancelled(ctrl.getCancel))
+ // 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(Termination.TransportFailed(new IllegalStateException(
+ "empty UdfControlResponse oneof")))
+ initLatch.countDown()
+ }
+
+ /** True once init has resolved (Initialized or beyond). */
+ private def initResolved: Boolean = currentState match {
+ case SessionState.Created | SessionState.Initializing => false
+ case _ => true
+ }
+
+ private def cancelWithReason(reason: String): Cancel =
+ Cancel.newBuilder().setReason(reason).build()
+
+ // ---- WorkerSession hooks ------------------------------------------------
+
+ override protected def doInit(message: Init): InitResponse = {
+ // 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(Termination.TransportFailed(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. close()'s 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 {
+ // The base advanced Created -> Initializing before calling doInit, so a
+ // directExecutor worker's InitResponse -- delivered reentrantly inside
+ // onNext -- finds Initializing and advances to Initialized.
+ 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
+ completeTerminal(Termination.TransportFailed(e))
+ // Surface as GrpcWorkerSessionException so the engine integration layer
+ // (which catches that type and wraps it) sees a uniform init-failure
+ // exception rather than the raw transport error.
+ throw new GrpcWorkerSessionException("UDF worker stream failed during init", e)
+ }
+
+ val responded = try {
+ initLatch.await(initResponseTimeoutMs, TimeUnit.MILLISECONDS)
+ } catch {
+ case _: InterruptedException =>
+ Thread.currentThread().interrupt()
+ sendCancelInternal(() => cancelWithReason("interrupted during init"))
+ // Make sure close() does not block waiting for a terminator that
+ // will never arrive on this thread's behalf.
+ completeTerminal(Termination.TransportFailed(
+ new InterruptedException("interrupted while waiting for InitResponse")))
+ throw new InterruptedException("interrupted while waiting for InitResponse")
+ }
+
+ if (!responded) {
+ sendCancelInternal(() => cancelWithReason("InitResponse timed out"))
+ val timeout = new TimeoutException(
+ s"timed out waiting for InitResponse after ${initResponseTimeoutMs}ms")
+ // 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(Termination.TransportFailed(timeout))
+ // Surface as GrpcWorkerSessionException (carrying the timeout cause) so
+ // the engine integration layer that catches that type can wrap it.
+ throw new GrpcWorkerSessionException(
+ s"timed out waiting for InitResponse after ${initResponseTimeoutMs}ms", timeout)
+ }
+
+ initResponse.get() match {
+ case Some(resp) if resp.hasError =>
+ // 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(resp.getError)}", resp.getError)
+ case Some(resp) =>
+ resp
+ case None =>
+ // No InitResponse arrived 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.
+ currentState match {
+ case SessionState.Terminal(Termination.TransportFailed(cause)) =>
+ throw new GrpcWorkerSessionException(
+ "UDF worker stream failed during init", cause)
+ case SessionState.Terminal(Termination.Failed(err)) =>
+ throw new GrpcWorkerSessionException(
+ s"UDF worker reported an error before init completed: " +
+ describeError(err), err)
+ case SessionState.Terminal(Termination.Cancelled(_)) =>
+ throw new GrpcWorkerSessionException(
+ "UDF worker stream was cancelled before init completed")
+ case SessionState.Terminal(Termination.Finished(_)) =>
+ throw new GrpcWorkerSessionException(
+ "UDF worker finished before init completed")
+ case other =>
+ throw new IllegalStateException(
+ s"init latch fired without an InitResponse or terminal: $other")
+ }
+ }
+ }
+
+ override protected def doProcess(
+ input: Iterator[DataRequest],
+ finish: () => Finish): Iterator[DataResponse] = {
+ // Init success is guaranteed by the base [[WorkerSession]] lifecycle: if
+ // doInit had failed it would have thrown and process() would never run.
+ new ProcessIterator(input, finish)
+ }
+
+ override protected def doClose(cancel: () => Cancel): Termination = {
+ if (requestObserver == null) {
+ // init() never put a stream on the wire (closed before/around init, or
+ // init threw before publishing). There is no protocol terminator; treat
+ // the session as cancelled-before-start. The base WorkerSession still
+ // releases the worker handle, so the worker is torn down.
+ completeTerminal(Termination.Cancelled(CancelResponse.getDefaultInstance))
+ return Termination.Cancelled(CancelResponse.getDefaultInstance)
+ }
+ // If the stream has not finished on its own, cancel anything in flight so
+ // the worker can clean up, then wait for the terminator. sendCancelInternal
+ // evaluates the cancel thunk only when it actually writes a Cancel, and at
+ // most once across all callers.
+ if (!currentState.isTerminal) {
+ sendCancelInternal(cancel)
+ 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 the state sees a stable value.
+ if (!currentState.isTerminal) {
+ completeTerminal(Termination.TransportFailed(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. onCompleted on an
+ // already-torn-down stream throws and is caught, so this stays idempotent
+ // across repeat close() calls.
+ try {
+ requestLock.synchronized { requestObserver.onCompleted() }
+ } catch {
+ case NonFatal(e) =>
+ logger.debug("Error half-closing UDF Execute stream", e)
+ }
+ // Derive the Termination from the settled terminal. close() is the
+ // finalizer/cleanup path and must NOT re-throw: a UDF / data-phase error is
+ // already surfaced while the result iterator is consumed, and an init error
+ // from init(). settledTermination returns the response proto for the clean
+ // terminators and a best-effort empty Cancelled for the failure terminals.
+ settledTermination
+ }
+
+ // ---- Internal request helpers ---------------------------------------------
+
+ /**
+ * Sends a Data or Finish request to the worker. Two invariants are
+ * checked inside [[requestLock]]:
+ * - terminal state: 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 (currentState.isTerminal) {
+ throw new IllegalStateException(
+ "cannot send request: UDF Execute stream is already closed")
+ }
+ // Suppress the write when a cancel has been (or is about to be) flushed
+ // through this lock; that preserves the proto ordering invariant (no
+ // Data/Finish after Cancel). We test the flag rather than `return`-ing
+ // from the block: a non-local return out of this by-name `synchronized`
+ // body compiles to a thrown NonLocalReturnControl, which is brittle.
+ if (!cancelRequested.get()) {
+ requestObserver.onNext(req)
+ }
+ }
+
+ /**
+ * Sends a `Cancel` control message, returning `true` iff a `Cancel` was
+ * actually written to the wire. Idempotent across ALL call sites (close()'s
+ * cancel, in-band cancels from `handleControl`'s INIT error / ERROR paths,
+ * and engine-internal cancels from the iterator) via [[cancelRequested]]:
+ * only the first caller puts a `Cancel` on the wire, and the `cancel` thunk is
+ * evaluated only by that caller and only when the message is actually sent --
+ * so a side-effecting thunk (e.g. one carrying a client cancel callback) runs
+ * at most once. 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(cancel: () => Cancel): Boolean = {
+ if (!cancelRequested.compareAndSet(false, true)) return false
+ if (requestObserver == null) return false
+ if (currentState.isTerminal) return false
+ try {
+ requestLock.synchronized {
+ // A terminal that arrived first must win; bail without writing. We
+ // compute the block's value rather than `return`-ing from this by-name
+ // `synchronized` body, whose non-local return would compile to a thrown
+ // NonLocalReturnControl that the surrounding NonFatal catch must not
+ // swallow -- brittle to rely on. Reads naturally as the block result.
+ val cur = currentState
+ if (cur.isTerminal) {
+ false
+ } else {
+ // Record that a Cancel is on the wire by advancing to Cancelling.
+ compareAndSetState(cur, SessionState.Cancelling)
+ requestObserver.onNext(UdfRequest.newBuilder()
+ .setControl(UdfControlRequest.newBuilder().setCancel(cancel()).build())
+ .build())
+ true
+ }
+ }
+ } catch {
+ case NonFatal(e) =>
+ logger.debug(s"Cancel send failed (stream may already be torn down): ${e.getMessage}")
+ completeTerminal(Termination.TransportFailed(e))
+ false
+ }
+ }
+
+ private def awaitTerminal(): Unit = {
+ if (currentState.isTerminal) return
+ try {
+ if (!terminalLatch.await(terminalTimeoutMs, TimeUnit.MILLISECONDS)) {
+ completeTerminal(Termination.TransportFailed(new TimeoutException(
+ s"timed out waiting for stream terminator after ${terminalTimeoutMs}ms")))
+ }
+ } catch {
+ case _: InterruptedException =>
+ // Attribute the terminal to the interrupt, not a timeout: the wait was
+ // cut short, it did not exhaust terminalTimeoutMs. Reporting "timed
+ // out" here would misdiagnose an interrupted close as a slow worker.
+ Thread.currentThread().interrupt()
+ completeTerminal(Termination.TransportFailed(
+ new InterruptedException("interrupted while waiting for stream terminator")))
+ }
+ }
+
+ // ---- 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 settle the state) and with any thread that finalizes via
+ * [[close]]. It drives the `Streaming -> Finishing` transition (sending
+ * `Finish` once input is exhausted, built lazily from the `finish` thunk).
+ */
+ private class ProcessIterator(input: Iterator[DataRequest], finish: () => Finish)
+ extends Iterator[DataResponse] {
+
+ // 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 branches, and block branch 4 for the
+ // full terminalTimeoutMs before returning. Callers that probe hasNext
+ // an extra time (iterator.size, instrumentation wrappers) would hang.
+ // Iterator-local and only touched by the single engine thread.
+ private val exhausted = new AtomicBoolean(false)
+ @volatile private var prefetched: DataResponse = _
+
+ // Pre-init cancel fast-fail: if a 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() && !currentState.isTerminal) {
+ completeTerminal(Termination.Cancelled(CancelResponse.getDefaultInstance))
+ }
+
+ override def hasNext: Boolean = {
+ if (prefetched ne null) return true
+ advance()
+ prefetched ne null
+ }
+
+ override def next(): DataResponse = {
+ if (prefetched eq null) advance()
+ val out = prefetched
+ if (out eq null) {
+ throw new NoSuchElementException("ProcessIterator exhausted")
+ }
+ prefetched = null
+ out
+ }
+
+ /**
+ * Single loop that, each iteration, tries in order: (1) drain queued output
+ * (a batch, or the terminator sentinel); (2) while `Streaming`, send the next
+ * input batch; (3) once input is exhausted, send `Finish` (CAS `Streaming ->
+ * Finishing`, once); (4) otherwise block for late output or the terminator.
+ * Each branch is commented inline below.
+ *
+ * '''Per-poll timeout.''' Branch 4 waits up to [[terminalTimeoutMs]] per
+ * poll, reset by every worker event -- the contract is "emit at least one
+ * event every [[terminalTimeoutMs]] after Finish", not "finish the UDF within
+ * [[terminalTimeoutMs]]". A worker expecting a long post-Finish silence MAY
+ * emit an empty `DataResponse` as a heartbeat to reset the wait; it is
+ * surfaced to the caller, so it should be a batch the caller recognises as
+ * empty (e.g. a zero-row Arrow batch).
+ */
+ private def advance(): Unit = {
+ if (exhausted.get()) return // terminator already seen; never re-block
+ 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 while the stream is open for data.
+ if (!cancelRequested.get() && currentState == SessionState.Streaming && input.hasNext) {
Review Comment:
can `input.hasNext` fail? it seems it fetches the next one and caches it, then `next()` can use it. if that's the case, `hasNext` could fail whiling fetch the next data right?
##########
udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/GrpcWorkerSession.scala:
##########
@@ -0,0 +1,816 @@
+/*
+ * 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.grpc
+
+import java.util.concurrent.{CountDownLatch, LinkedBlockingQueue, TimeoutException, TimeUnit}
+import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference}
+
+import scala.util.control.NonFatal
+
+import io.grpc.{ConnectivityState, ManagedChannel}
+import io.grpc.stub.StreamObserver
+
+import org.apache.spark.annotation.Experimental
+import org.apache.spark.udf.worker.{Cancel, CancelResponse, DataRequest, DataResponse,
+ ExecutionError, Finish, Init, InitResponse, UdfControlRequest,
+ UdfControlResponse, UdfRequest, UdfResponse, UdfWorkerGrpc}
+import org.apache.spark.udf.worker.core.{Termination, WorkerHandle, WorkerLogger, WorkerSession}
+import org.apache.spark.udf.worker.core.WorkerSession.SessionState
+import org.apache.spark.udf.worker.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_message.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.
+ *
+ * '''Driving model.''' Consumption-driven (Volcano / pull): the thread that
+ * consumes the [[doProcess]] result iterator is the one that pulls the next
+ * input batch and sends its `DataRequest`. Nothing is sent until the iterator
+ * is consumed, and input flows one batch per output pull. The gRPC callback
+ * thread only receives output.
+ *
+ * '''State machine.''' This class does not keep its own state machine: it
+ * drives the single [[WorkerSession.SessionState]] owned by the base. The base
+ * advances `Created -> Initializing` (in `init`) and `Initialized -> Streaming`
+ * (in `process`); this class advances the protocol-event edges through
+ * [[compareAndSetState]] / [[completeTerminal]] as it exchanges messages:
+ * {{{
+ * Initializing --(InitResponse ok)--> Initialized [handleControl]
+ * Streaming ----(Finish written)----> Finishing [ProcessIterator]
+ * <any non-terminal> --(Cancel written)--> Cancelling [sendCancelInternal]
+ * <any non-terminal> --(terminator/error)--> terminal [completeTerminal]
+ * }}}
+ * The two clean terminals carry the worker's `FinishResponse` / `CancelResponse`
+ * (metrics + finish/cancel callback `data`/`error`) so [[close]] can return
+ * them. The only flag kept outside the machine is [[cancelRequested]]: a
+ * cancellation can be requested before the stream exists (so it cannot be a
+ * state transition yet), and it must both fast-fail the result iterator and
+ * suppress any in-flight Data/Finish.
+ *
+ * Threading:
+ * - [[doInit]] is synchronous: sends `Init` and blocks on `InitResponse`,
+ * returning it.
+ * - [[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.
+ * - [[doClose]] is thread-safe and idempotent: it settles + returns the
+ * terminator (cancelling in-flight work if the stream had not finished)
+ * and half-closes the request side.
+ *
+ * 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 a gRPC channel built and owned by the caller (the
+ * dispatcher). 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 from the worker, drained by the process() iterator.
+ // Intentionally unbounded: a bounded queue would block the gRPC callback
+ // (Netty event-loop) thread when full, stalling control-message and
+ // terminator delivery on the whole channel. HTTP/2 flow control bounds the
+ // wire and the consumer normally drains promptly, so it stays small. A
+ // stalled downstream can still grow it; the real fix is protocol-level
+ // back-pressure (out of scope here), not bounding the queue.
+ //
+ // TODO [SPARK-57324]: expose queue depth as a metric (early warning for a
+ // stalled consumer).
+ 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)
+
+ // The InitResponse the worker sent (success or error), captured so init()
+ // can return it. None until InitResponse arrives.
+ private val initResponse = new AtomicReference[Option[InitResponse]](None)
+
+ // Fired when the session reaches a terminal [[SessionState]]. doClose() and
+ // the init-error path block on this to drain the terminator.
+ private val terminalLatch = new CountDownLatch(1)
+
+ // Captures an ErrorResponse encountered during the data phase so that
+ // the CancelResponse terminator 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)
+
+ // Cancellation intent. Kept outside the [[SessionState]] machine because a
+ // cancel can be requested before the stream exists (pre-init), where there is
+ // no wire transition to make yet -- only an intent to record. Used to (a)
+ // make cancel idempotent across all call sites, (b) fast-fail ProcessIterator
+ // on a pre-init cancel, and (c) suppress any Data/Finish that would otherwise
+ // race a Cancel onto the wire (re-read inside [[requestLock]]).
+ private val cancelRequested = new AtomicBoolean(false)
+
+ // gRPC requires serialized writes to a request StreamObserver.
+ private val requestLock = new Object
+
+ // Initialised in init() -- before that, close() is a no-op 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))
+
+ case UdfResponse.ResponseCase.CONTROL =>
+ handleControl(response.getControl)
+
+ case other =>
+ // A malformed response (empty / unknown oneof) is a terminal
+ // transport failure. Count down initLatch like every other
+ // terminal-settling path here (onError / onCompleted / each
+ // handleControl branch): if this arrives before InitResponse, doInit
+ // must fail fast with this cause rather than block until
+ // initResponseTimeoutMs and report a misleading "timed out" error.
+ completeTerminal(Termination.TransportFailed(new IllegalStateException(
+ s"unexpected response oneof: $other")))
+ initLatch.countDown()
+ }
+ }
+
+ override def onError(t: Throwable): Unit = {
+ // Transport-level failure: the stream is dead. No further writes are
+ // possible (reaching a terminal state closes the write side). Settle the
+ // terminal BEFORE counting down initLatch: init() blocks on that latch,
+ // and only await/countDown establish a happens-before edge, so a thread
+ // woken by the countDown must already be able to observe the terminal.
+ // That lets init() surface the transport cause instead of timing out
+ // after initResponseTimeoutMs with a misleading "timed out" message --
+ // or, if it saw a transient non-terminal state, the defensive "init latch
+ // fired without an InitResponse or terminal" error.
+ completeTerminal(Termination.TransportFailed(t))
+ initLatch.countDown()
+ }
+
+ 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 (!currentState.isTerminal) {
+ completeTerminal(Termination.TransportFailed(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()
+ }
+ }
+
+ /**
+ * Wakes the result iterator (blocked on [[outputQueue]]) and any thread
+ * waiting on [[terminalLatch]] when the base settles a terminal. Invoked once,
+ * by the caller that wins [[completeTerminal]].
+ */
+ override protected def onTerminalSettled(termination: Termination): Unit = {
+ outputQueue.put(QueueItem.EndOfStream)
+ terminalLatch.countDown()
+ }
+
+ private def handleControl(ctrl: UdfControlResponse): Unit = ctrl.getControlCase match {
+ case UdfControlResponse.ControlCase.INIT =>
+ val resp = ctrl.getInit
+ // Capture the InitResponse so init() can return it (or throw on its error).
+ initResponse.set(Some(resp))
+ if (resp.hasError) {
+ initLatch.countDown()
+ // Per the protocol the engine follows an init error with Cancel; send it
+ // if the request stream is up. If the Cancel cannot be written -- e.g. a
+ // directExecutor worker delivered InitResponse reentrantly, before
+ // doInit published requestObserver -- settle the terminal here so doInit
+ // does not wait the full terminalTimeoutMs for a CancelResponse that can
+ // never arrive. (When the Cancel is sent, the CancelResponse settles the
+ // terminal instead.)
+ if (!sendCancelInternal(() => cancelWithReason("init failed")) &&
+ !currentState.isTerminal) {
+ completeTerminal(Termination.Failed(resp.getError))
+ }
+ } else {
+ // InitResponse OK: init succeeded. A terminal that arrived first (e.g. a
+ // racing transport error) must win, so only advance from the expected
+ // Initializing state; process() then opens the data phase (Streaming).
+ compareAndSetState(SessionState.Initializing, SessionState.Initialized)
+ initLatch.countDown()
+ }
+
+ case UdfControlResponse.ControlCase.ERROR =>
+ val err = ctrl.getError.getError
+ executionError.compareAndSet(None, Some(err))
+ // 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
+ // (as Failed, since there is no proto terminator to carry) so doInit
+ // doesn't return as if init succeeded. "Init not yet resolved" is exactly
+ // "the stream has not reached Initialized"; completeTerminal is idempotent,
+ // so the data-phase ERROR -> Cancel -> CancelResponse path is unaffected.
+ if (!initResolved) {
+ completeTerminal(Termination.Failed(err))
+ }
+ // Surface the error to doInit. Count down AFTER settling the terminal
+ // above: init() blocks on this latch and only await/countDown establish
+ // happens-before, so a woken doInit must be able to observe the terminal
+ // rather than a transient non-terminal state. Idempotent, so the normal
+ // data-phase path (latch already fired during init) is unaffected.
+ initLatch.countDown()
+ sendCancelInternal(() => cancelWithReason("aborting after ErrorResponse"))
+
+ case UdfControlResponse.ControlCase.FINISH =>
+ // The FinishResponse carries metrics + the finish-callback data/error.
+ // Keep it on the terminal so close() can return it; the iterator inspects
+ // its error field to decide whether to throw.
+ completeTerminal(Termination.Finished(ctrl.getFinish))
+ // 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 =>
+ // The CancelResponse carries metrics + the cancel-callback error. Keep it
+ // on the terminal so close() can return it; any prior ErrorResponse is
+ // tracked in executionError and surfaced by the iterator.
+ completeTerminal(Termination.Cancelled(ctrl.getCancel))
+ // 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(Termination.TransportFailed(new IllegalStateException(
+ "empty UdfControlResponse oneof")))
+ initLatch.countDown()
+ }
+
+ /** True once init has resolved (Initialized or beyond). */
+ private def initResolved: Boolean = currentState match {
+ case SessionState.Created | SessionState.Initializing => false
+ case _ => true
+ }
+
+ private def cancelWithReason(reason: String): Cancel =
+ Cancel.newBuilder().setReason(reason).build()
+
+ // ---- WorkerSession hooks ------------------------------------------------
+
+ override protected def doInit(message: Init): InitResponse = {
+ // 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(Termination.TransportFailed(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. close()'s 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 {
+ // The base advanced Created -> Initializing before calling doInit, so a
+ // directExecutor worker's InitResponse -- delivered reentrantly inside
+ // onNext -- finds Initializing and advances to Initialized.
+ 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
+ completeTerminal(Termination.TransportFailed(e))
+ // Surface as GrpcWorkerSessionException so the engine integration layer
+ // (which catches that type and wraps it) sees a uniform init-failure
+ // exception rather than the raw transport error.
+ throw new GrpcWorkerSessionException("UDF worker stream failed during init", e)
+ }
+
+ val responded = try {
+ initLatch.await(initResponseTimeoutMs, TimeUnit.MILLISECONDS)
+ } catch {
+ case _: InterruptedException =>
+ Thread.currentThread().interrupt()
+ sendCancelInternal(() => cancelWithReason("interrupted during init"))
+ // Make sure close() does not block waiting for a terminator that
+ // will never arrive on this thread's behalf.
+ completeTerminal(Termination.TransportFailed(
+ new InterruptedException("interrupted while waiting for InitResponse")))
+ throw new InterruptedException("interrupted while waiting for InitResponse")
+ }
+
+ if (!responded) {
+ sendCancelInternal(() => cancelWithReason("InitResponse timed out"))
+ val timeout = new TimeoutException(
+ s"timed out waiting for InitResponse after ${initResponseTimeoutMs}ms")
+ // 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(Termination.TransportFailed(timeout))
+ // Surface as GrpcWorkerSessionException (carrying the timeout cause) so
+ // the engine integration layer that catches that type can wrap it.
+ throw new GrpcWorkerSessionException(
+ s"timed out waiting for InitResponse after ${initResponseTimeoutMs}ms", timeout)
+ }
+
+ initResponse.get() match {
+ case Some(resp) if resp.hasError =>
+ // 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(resp.getError)}", resp.getError)
+ case Some(resp) =>
+ resp
+ case None =>
+ // No InitResponse arrived 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.
+ currentState match {
+ case SessionState.Terminal(Termination.TransportFailed(cause)) =>
+ throw new GrpcWorkerSessionException(
+ "UDF worker stream failed during init", cause)
+ case SessionState.Terminal(Termination.Failed(err)) =>
+ throw new GrpcWorkerSessionException(
+ s"UDF worker reported an error before init completed: " +
+ describeError(err), err)
+ case SessionState.Terminal(Termination.Cancelled(_)) =>
+ throw new GrpcWorkerSessionException(
+ "UDF worker stream was cancelled before init completed")
+ case SessionState.Terminal(Termination.Finished(_)) =>
+ throw new GrpcWorkerSessionException(
+ "UDF worker finished before init completed")
+ case other =>
+ throw new IllegalStateException(
+ s"init latch fired without an InitResponse or terminal: $other")
+ }
+ }
+ }
+
+ override protected def doProcess(
+ input: Iterator[DataRequest],
+ finish: () => Finish): Iterator[DataResponse] = {
+ // Init success is guaranteed by the base [[WorkerSession]] lifecycle: if
+ // doInit had failed it would have thrown and process() would never run.
+ new ProcessIterator(input, finish)
+ }
+
+ override protected def doClose(cancel: () => Cancel): Termination = {
+ if (requestObserver == null) {
+ // init() never put a stream on the wire (closed before/around init, or
+ // init threw before publishing). There is no protocol terminator; treat
+ // the session as cancelled-before-start. The base WorkerSession still
+ // releases the worker handle, so the worker is torn down.
+ completeTerminal(Termination.Cancelled(CancelResponse.getDefaultInstance))
+ return Termination.Cancelled(CancelResponse.getDefaultInstance)
+ }
+ // If the stream has not finished on its own, cancel anything in flight so
+ // the worker can clean up, then wait for the terminator. sendCancelInternal
+ // evaluates the cancel thunk only when it actually writes a Cancel, and at
+ // most once across all callers.
+ if (!currentState.isTerminal) {
+ sendCancelInternal(cancel)
+ 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 the state sees a stable value.
+ if (!currentState.isTerminal) {
+ completeTerminal(Termination.TransportFailed(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. onCompleted on an
+ // already-torn-down stream throws and is caught, so this stays idempotent
+ // across repeat close() calls.
+ try {
+ requestLock.synchronized { requestObserver.onCompleted() }
+ } catch {
+ case NonFatal(e) =>
+ logger.debug("Error half-closing UDF Execute stream", e)
+ }
+ // Derive the Termination from the settled terminal. close() is the
+ // finalizer/cleanup path and must NOT re-throw: a UDF / data-phase error is
+ // already surfaced while the result iterator is consumed, and an init error
+ // from init(). settledTermination returns the response proto for the clean
+ // terminators and a best-effort empty Cancelled for the failure terminals.
+ settledTermination
+ }
+
+ // ---- Internal request helpers ---------------------------------------------
+
+ /**
+ * Sends a Data or Finish request to the worker. Two invariants are
+ * checked inside [[requestLock]]:
+ * - terminal state: 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 (currentState.isTerminal) {
+ throw new IllegalStateException(
+ "cannot send request: UDF Execute stream is already closed")
+ }
+ // Suppress the write when a cancel has been (or is about to be) flushed
+ // through this lock; that preserves the proto ordering invariant (no
+ // Data/Finish after Cancel). We test the flag rather than `return`-ing
+ // from the block: a non-local return out of this by-name `synchronized`
+ // body compiles to a thrown NonLocalReturnControl, which is brittle.
+ if (!cancelRequested.get()) {
+ requestObserver.onNext(req)
+ }
+ }
+
+ /**
+ * Sends a `Cancel` control message, returning `true` iff a `Cancel` was
+ * actually written to the wire. Idempotent across ALL call sites (close()'s
+ * cancel, in-band cancels from `handleControl`'s INIT error / ERROR paths,
+ * and engine-internal cancels from the iterator) via [[cancelRequested]]:
+ * only the first caller puts a `Cancel` on the wire, and the `cancel` thunk is
+ * evaluated only by that caller and only when the message is actually sent --
+ * so a side-effecting thunk (e.g. one carrying a client cancel callback) runs
+ * at most once. 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(cancel: () => Cancel): Boolean = {
+ if (!cancelRequested.compareAndSet(false, true)) return false
+ if (requestObserver == null) return false
+ if (currentState.isTerminal) return false
+ try {
+ requestLock.synchronized {
+ // A terminal that arrived first must win; bail without writing. We
+ // compute the block's value rather than `return`-ing from this by-name
+ // `synchronized` body, whose non-local return would compile to a thrown
+ // NonLocalReturnControl that the surrounding NonFatal catch must not
+ // swallow -- brittle to rely on. Reads naturally as the block result.
+ val cur = currentState
+ if (cur.isTerminal) {
+ false
+ } else {
+ // Record that a Cancel is on the wire by advancing to Cancelling.
+ compareAndSetState(cur, SessionState.Cancelling)
+ requestObserver.onNext(UdfRequest.newBuilder()
+ .setControl(UdfControlRequest.newBuilder().setCancel(cancel()).build())
+ .build())
+ true
+ }
+ }
+ } catch {
+ case NonFatal(e) =>
+ logger.debug(s"Cancel send failed (stream may already be torn down): ${e.getMessage}")
+ completeTerminal(Termination.TransportFailed(e))
+ false
+ }
+ }
+
+ private def awaitTerminal(): Unit = {
+ if (currentState.isTerminal) return
+ try {
+ if (!terminalLatch.await(terminalTimeoutMs, TimeUnit.MILLISECONDS)) {
+ completeTerminal(Termination.TransportFailed(new TimeoutException(
+ s"timed out waiting for stream terminator after ${terminalTimeoutMs}ms")))
+ }
+ } catch {
+ case _: InterruptedException =>
+ // Attribute the terminal to the interrupt, not a timeout: the wait was
+ // cut short, it did not exhaust terminalTimeoutMs. Reporting "timed
+ // out" here would misdiagnose an interrupted close as a slow worker.
+ Thread.currentThread().interrupt()
+ completeTerminal(Termination.TransportFailed(
+ new InterruptedException("interrupted while waiting for stream terminator")))
+ }
+ }
+
+ // ---- 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 settle the state) and with any thread that finalizes via
+ * [[close]]. It drives the `Streaming -> Finishing` transition (sending
+ * `Finish` once input is exhausted, built lazily from the `finish` thunk).
+ */
+ private class ProcessIterator(input: Iterator[DataRequest], finish: () => Finish)
+ extends Iterator[DataResponse] {
+
+ // 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 branches, and block branch 4 for the
+ // full terminalTimeoutMs before returning. Callers that probe hasNext
+ // an extra time (iterator.size, instrumentation wrappers) would hang.
+ // Iterator-local and only touched by the single engine thread.
+ private val exhausted = new AtomicBoolean(false)
+ @volatile private var prefetched: DataResponse = _
+
+ // Pre-init cancel fast-fail: if a 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() && !currentState.isTerminal) {
+ completeTerminal(Termination.Cancelled(CancelResponse.getDefaultInstance))
+ }
+
+ override def hasNext: Boolean = {
+ if (prefetched ne null) return true
+ advance()
+ prefetched ne null
+ }
+
+ override def next(): DataResponse = {
+ if (prefetched eq null) advance()
+ val out = prefetched
+ if (out eq null) {
+ throw new NoSuchElementException("ProcessIterator exhausted")
+ }
+ prefetched = null
+ out
+ }
+
+ /**
+ * Single loop that, each iteration, tries in order: (1) drain queued output
+ * (a batch, or the terminator sentinel); (2) while `Streaming`, send the next
+ * input batch; (3) once input is exhausted, send `Finish` (CAS `Streaming ->
+ * Finishing`, once); (4) otherwise block for late output or the terminator.
+ * Each branch is commented inline below.
+ *
+ * '''Per-poll timeout.''' Branch 4 waits up to [[terminalTimeoutMs]] per
+ * poll, reset by every worker event -- the contract is "emit at least one
+ * event every [[terminalTimeoutMs]] after Finish", not "finish the UDF within
+ * [[terminalTimeoutMs]]". A worker expecting a long post-Finish silence MAY
+ * emit an empty `DataResponse` as a heartbeat to reset the wait; it is
+ * surfaced to the caller, so it should be a batch the caller recognises as
+ * empty (e.g. a zero-row Arrow batch).
+ */
+ private def advance(): Unit = {
+ if (exhausted.get()) return // terminator already seen; never re-block
+ 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 while the stream is open for data.
+ if (!cancelRequested.get() && currentState == SessionState.Streaming && input.hasNext) {
+ val batch = try input.next() catch {
+ case NonFatal(e) =>
+ sendCancelInternal(() => cancelWithReason("input iterator failed"))
+ throw e
+ }
+ if (!sendOrEndOnRacedTerminal(UdfRequest.newBuilder().setData(batch).build())) return
+ } else if (!cancelRequested.get() &&
+ compareAndSetState(SessionState.Streaming, SessionState.Finishing)) {
+ // (3) No more input; send Finish exactly once (unless cancelled).
+ if (!sendOrEndOnRacedTerminal(UdfRequest.newBuilder()
+ .setControl(UdfControlRequest.newBuilder().setFinish(finish()).build())
Review Comment:
`finish()` can also potentially fail right? it is the callback code
##########
udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/GrpcWorkerSession.scala:
##########
@@ -0,0 +1,816 @@
+/*
+ * 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.grpc
+
+import java.util.concurrent.{CountDownLatch, LinkedBlockingQueue, TimeoutException, TimeUnit}
+import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference}
+
+import scala.util.control.NonFatal
+
+import io.grpc.{ConnectivityState, ManagedChannel}
+import io.grpc.stub.StreamObserver
+
+import org.apache.spark.annotation.Experimental
+import org.apache.spark.udf.worker.{Cancel, CancelResponse, DataRequest, DataResponse,
+ ExecutionError, Finish, Init, InitResponse, UdfControlRequest,
+ UdfControlResponse, UdfRequest, UdfResponse, UdfWorkerGrpc}
+import org.apache.spark.udf.worker.core.{Termination, WorkerHandle, WorkerLogger, WorkerSession}
+import org.apache.spark.udf.worker.core.WorkerSession.SessionState
+import org.apache.spark.udf.worker.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_message.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.
+ *
+ * '''Driving model.''' Consumption-driven (Volcano / pull): the thread that
+ * consumes the [[doProcess]] result iterator is the one that pulls the next
+ * input batch and sends its `DataRequest`. Nothing is sent until the iterator
+ * is consumed, and input flows one batch per output pull. The gRPC callback
+ * thread only receives output.
+ *
+ * '''State machine.''' This class does not keep its own state machine: it
+ * drives the single [[WorkerSession.SessionState]] owned by the base. The base
+ * advances `Created -> Initializing` (in `init`) and `Initialized -> Streaming`
+ * (in `process`); this class advances the protocol-event edges through
+ * [[compareAndSetState]] / [[completeTerminal]] as it exchanges messages:
+ * {{{
+ * Initializing --(InitResponse ok)--> Initialized [handleControl]
+ * Streaming ----(Finish written)----> Finishing [ProcessIterator]
+ * <any non-terminal> --(Cancel written)--> Cancelling [sendCancelInternal]
+ * <any non-terminal> --(terminator/error)--> terminal [completeTerminal]
+ * }}}
+ * The two clean terminals carry the worker's `FinishResponse` / `CancelResponse`
+ * (metrics + finish/cancel callback `data`/`error`) so [[close]] can return
+ * them. The only flag kept outside the machine is [[cancelRequested]]: a
+ * cancellation can be requested before the stream exists (so it cannot be a
+ * state transition yet), and it must both fast-fail the result iterator and
+ * suppress any in-flight Data/Finish.
+ *
+ * Threading:
+ * - [[doInit]] is synchronous: sends `Init` and blocks on `InitResponse`,
+ * returning it.
+ * - [[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.
+ * - [[doClose]] is thread-safe and idempotent: it settles + returns the
+ * terminator (cancelling in-flight work if the stream had not finished)
+ * and half-closes the request side.
+ *
+ * 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 a gRPC channel built and owned by the caller (the
+ * dispatcher). 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 from the worker, drained by the process() iterator.
+ // Intentionally unbounded: a bounded queue would block the gRPC callback
+ // (Netty event-loop) thread when full, stalling control-message and
+ // terminator delivery on the whole channel. HTTP/2 flow control bounds the
+ // wire and the consumer normally drains promptly, so it stays small. A
+ // stalled downstream can still grow it; the real fix is protocol-level
+ // back-pressure (out of scope here), not bounding the queue.
+ //
+ // TODO [SPARK-57324]: expose queue depth as a metric (early warning for a
+ // stalled consumer).
+ 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)
+
+ // The InitResponse the worker sent (success or error), captured so init()
+ // can return it. None until InitResponse arrives.
+ private val initResponse = new AtomicReference[Option[InitResponse]](None)
+
+ // Fired when the session reaches a terminal [[SessionState]]. doClose() and
+ // the init-error path block on this to drain the terminator.
+ private val terminalLatch = new CountDownLatch(1)
+
+ // Captures an ErrorResponse encountered during the data phase so that
+ // the CancelResponse terminator 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)
+
+ // Cancellation intent. Kept outside the [[SessionState]] machine because a
+ // cancel can be requested before the stream exists (pre-init), where there is
+ // no wire transition to make yet -- only an intent to record. Used to (a)
+ // make cancel idempotent across all call sites, (b) fast-fail ProcessIterator
+ // on a pre-init cancel, and (c) suppress any Data/Finish that would otherwise
+ // race a Cancel onto the wire (re-read inside [[requestLock]]).
+ private val cancelRequested = new AtomicBoolean(false)
+
+ // gRPC requires serialized writes to a request StreamObserver.
+ private val requestLock = new Object
+
+ // Initialised in init() -- before that, close() is a no-op 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))
+
+ case UdfResponse.ResponseCase.CONTROL =>
+ handleControl(response.getControl)
+
+ case other =>
+ // A malformed response (empty / unknown oneof) is a terminal
+ // transport failure. Count down initLatch like every other
+ // terminal-settling path here (onError / onCompleted / each
+ // handleControl branch): if this arrives before InitResponse, doInit
+ // must fail fast with this cause rather than block until
+ // initResponseTimeoutMs and report a misleading "timed out" error.
+ completeTerminal(Termination.TransportFailed(new IllegalStateException(
+ s"unexpected response oneof: $other")))
+ initLatch.countDown()
+ }
+ }
+
+ override def onError(t: Throwable): Unit = {
+ // Transport-level failure: the stream is dead. No further writes are
+ // possible (reaching a terminal state closes the write side). Settle the
+ // terminal BEFORE counting down initLatch: init() blocks on that latch,
+ // and only await/countDown establish a happens-before edge, so a thread
+ // woken by the countDown must already be able to observe the terminal.
+ // That lets init() surface the transport cause instead of timing out
+ // after initResponseTimeoutMs with a misleading "timed out" message --
+ // or, if it saw a transient non-terminal state, the defensive "init latch
+ // fired without an InitResponse or terminal" error.
+ completeTerminal(Termination.TransportFailed(t))
+ initLatch.countDown()
+ }
+
+ 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 (!currentState.isTerminal) {
+ completeTerminal(Termination.TransportFailed(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()
+ }
+ }
+
+ /**
+ * Wakes the result iterator (blocked on [[outputQueue]]) and any thread
+ * waiting on [[terminalLatch]] when the base settles a terminal. Invoked once,
+ * by the caller that wins [[completeTerminal]].
+ */
+ override protected def onTerminalSettled(termination: Termination): Unit = {
+ outputQueue.put(QueueItem.EndOfStream)
+ terminalLatch.countDown()
+ }
+
+ private def handleControl(ctrl: UdfControlResponse): Unit = ctrl.getControlCase match {
+ case UdfControlResponse.ControlCase.INIT =>
+ val resp = ctrl.getInit
+ // Capture the InitResponse so init() can return it (or throw on its error).
+ initResponse.set(Some(resp))
+ if (resp.hasError) {
+ initLatch.countDown()
+ // Per the protocol the engine follows an init error with Cancel; send it
+ // if the request stream is up. If the Cancel cannot be written -- e.g. a
+ // directExecutor worker delivered InitResponse reentrantly, before
+ // doInit published requestObserver -- settle the terminal here so doInit
+ // does not wait the full terminalTimeoutMs for a CancelResponse that can
+ // never arrive. (When the Cancel is sent, the CancelResponse settles the
+ // terminal instead.)
+ if (!sendCancelInternal(() => cancelWithReason("init failed")) &&
+ !currentState.isTerminal) {
+ completeTerminal(Termination.Failed(resp.getError))
+ }
+ } else {
+ // InitResponse OK: init succeeded. A terminal that arrived first (e.g. a
+ // racing transport error) must win, so only advance from the expected
+ // Initializing state; process() then opens the data phase (Streaming).
+ compareAndSetState(SessionState.Initializing, SessionState.Initialized)
+ initLatch.countDown()
+ }
+
+ case UdfControlResponse.ControlCase.ERROR =>
+ val err = ctrl.getError.getError
+ executionError.compareAndSet(None, Some(err))
+ // 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
+ // (as Failed, since there is no proto terminator to carry) so doInit
+ // doesn't return as if init succeeded. "Init not yet resolved" is exactly
+ // "the stream has not reached Initialized"; completeTerminal is idempotent,
+ // so the data-phase ERROR -> Cancel -> CancelResponse path is unaffected.
+ if (!initResolved) {
+ completeTerminal(Termination.Failed(err))
+ }
+ // Surface the error to doInit. Count down AFTER settling the terminal
+ // above: init() blocks on this latch and only await/countDown establish
+ // happens-before, so a woken doInit must be able to observe the terminal
+ // rather than a transient non-terminal state. Idempotent, so the normal
+ // data-phase path (latch already fired during init) is unaffected.
+ initLatch.countDown()
+ sendCancelInternal(() => cancelWithReason("aborting after ErrorResponse"))
+
+ case UdfControlResponse.ControlCase.FINISH =>
+ // The FinishResponse carries metrics + the finish-callback data/error.
+ // Keep it on the terminal so close() can return it; the iterator inspects
+ // its error field to decide whether to throw.
+ completeTerminal(Termination.Finished(ctrl.getFinish))
+ // 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 =>
+ // The CancelResponse carries metrics + the cancel-callback error. Keep it
+ // on the terminal so close() can return it; any prior ErrorResponse is
+ // tracked in executionError and surfaced by the iterator.
+ completeTerminal(Termination.Cancelled(ctrl.getCancel))
+ // 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(Termination.TransportFailed(new IllegalStateException(
+ "empty UdfControlResponse oneof")))
+ initLatch.countDown()
+ }
+
+ /** True once init has resolved (Initialized or beyond). */
+ private def initResolved: Boolean = currentState match {
+ case SessionState.Created | SessionState.Initializing => false
+ case _ => true
+ }
+
+ private def cancelWithReason(reason: String): Cancel =
+ Cancel.newBuilder().setReason(reason).build()
+
+ // ---- WorkerSession hooks ------------------------------------------------
+
+ override protected def doInit(message: Init): InitResponse = {
+ // 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(Termination.TransportFailed(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. close()'s 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 {
+ // The base advanced Created -> Initializing before calling doInit, so a
+ // directExecutor worker's InitResponse -- delivered reentrantly inside
+ // onNext -- finds Initializing and advances to Initialized.
+ 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
+ completeTerminal(Termination.TransportFailed(e))
+ // Surface as GrpcWorkerSessionException so the engine integration layer
+ // (which catches that type and wraps it) sees a uniform init-failure
+ // exception rather than the raw transport error.
+ throw new GrpcWorkerSessionException("UDF worker stream failed during init", e)
+ }
+
+ val responded = try {
+ initLatch.await(initResponseTimeoutMs, TimeUnit.MILLISECONDS)
+ } catch {
+ case _: InterruptedException =>
+ Thread.currentThread().interrupt()
+ sendCancelInternal(() => cancelWithReason("interrupted during init"))
+ // Make sure close() does not block waiting for a terminator that
+ // will never arrive on this thread's behalf.
+ completeTerminal(Termination.TransportFailed(
+ new InterruptedException("interrupted while waiting for InitResponse")))
+ throw new InterruptedException("interrupted while waiting for InitResponse")
+ }
+
+ if (!responded) {
+ sendCancelInternal(() => cancelWithReason("InitResponse timed out"))
+ val timeout = new TimeoutException(
+ s"timed out waiting for InitResponse after ${initResponseTimeoutMs}ms")
+ // 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(Termination.TransportFailed(timeout))
+ // Surface as GrpcWorkerSessionException (carrying the timeout cause) so
+ // the engine integration layer that catches that type can wrap it.
+ throw new GrpcWorkerSessionException(
+ s"timed out waiting for InitResponse after ${initResponseTimeoutMs}ms", timeout)
+ }
+
+ initResponse.get() match {
+ case Some(resp) if resp.hasError =>
+ // 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(resp.getError)}", resp.getError)
+ case Some(resp) =>
+ resp
+ case None =>
+ // No InitResponse arrived 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.
+ currentState match {
+ case SessionState.Terminal(Termination.TransportFailed(cause)) =>
+ throw new GrpcWorkerSessionException(
+ "UDF worker stream failed during init", cause)
+ case SessionState.Terminal(Termination.Failed(err)) =>
+ throw new GrpcWorkerSessionException(
+ s"UDF worker reported an error before init completed: " +
+ describeError(err), err)
+ case SessionState.Terminal(Termination.Cancelled(_)) =>
+ throw new GrpcWorkerSessionException(
+ "UDF worker stream was cancelled before init completed")
+ case SessionState.Terminal(Termination.Finished(_)) =>
+ throw new GrpcWorkerSessionException(
+ "UDF worker finished before init completed")
+ case other =>
+ throw new IllegalStateException(
+ s"init latch fired without an InitResponse or terminal: $other")
+ }
+ }
+ }
+
+ override protected def doProcess(
+ input: Iterator[DataRequest],
+ finish: () => Finish): Iterator[DataResponse] = {
+ // Init success is guaranteed by the base [[WorkerSession]] lifecycle: if
+ // doInit had failed it would have thrown and process() would never run.
+ new ProcessIterator(input, finish)
+ }
+
+ override protected def doClose(cancel: () => Cancel): Termination = {
+ if (requestObserver == null) {
+ // init() never put a stream on the wire (closed before/around init, or
+ // init threw before publishing). There is no protocol terminator; treat
+ // the session as cancelled-before-start. The base WorkerSession still
+ // releases the worker handle, so the worker is torn down.
+ completeTerminal(Termination.Cancelled(CancelResponse.getDefaultInstance))
+ return Termination.Cancelled(CancelResponse.getDefaultInstance)
+ }
+ // If the stream has not finished on its own, cancel anything in flight so
+ // the worker can clean up, then wait for the terminator. sendCancelInternal
+ // evaluates the cancel thunk only when it actually writes a Cancel, and at
+ // most once across all callers.
+ if (!currentState.isTerminal) {
+ sendCancelInternal(cancel)
+ 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 the state sees a stable value.
+ if (!currentState.isTerminal) {
+ completeTerminal(Termination.TransportFailed(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. onCompleted on an
+ // already-torn-down stream throws and is caught, so this stays idempotent
+ // across repeat close() calls.
+ try {
+ requestLock.synchronized { requestObserver.onCompleted() }
+ } catch {
+ case NonFatal(e) =>
+ logger.debug("Error half-closing UDF Execute stream", e)
+ }
+ // Derive the Termination from the settled terminal. close() is the
+ // finalizer/cleanup path and must NOT re-throw: a UDF / data-phase error is
+ // already surfaced while the result iterator is consumed, and an init error
+ // from init(). settledTermination returns the response proto for the clean
+ // terminators and a best-effort empty Cancelled for the failure terminals.
+ settledTermination
+ }
+
+ // ---- Internal request helpers ---------------------------------------------
+
+ /**
+ * Sends a Data or Finish request to the worker. Two invariants are
+ * checked inside [[requestLock]]:
+ * - terminal state: 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 (currentState.isTerminal) {
+ throw new IllegalStateException(
+ "cannot send request: UDF Execute stream is already closed")
+ }
+ // Suppress the write when a cancel has been (or is about to be) flushed
+ // through this lock; that preserves the proto ordering invariant (no
+ // Data/Finish after Cancel). We test the flag rather than `return`-ing
+ // from the block: a non-local return out of this by-name `synchronized`
+ // body compiles to a thrown NonLocalReturnControl, which is brittle.
+ if (!cancelRequested.get()) {
+ requestObserver.onNext(req)
+ }
+ }
+
+ /**
+ * Sends a `Cancel` control message, returning `true` iff a `Cancel` was
+ * actually written to the wire. Idempotent across ALL call sites (close()'s
+ * cancel, in-band cancels from `handleControl`'s INIT error / ERROR paths,
+ * and engine-internal cancels from the iterator) via [[cancelRequested]]:
+ * only the first caller puts a `Cancel` on the wire, and the `cancel` thunk is
+ * evaluated only by that caller and only when the message is actually sent --
+ * so a side-effecting thunk (e.g. one carrying a client cancel callback) runs
+ * at most once. 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(cancel: () => Cancel): Boolean = {
+ if (!cancelRequested.compareAndSet(false, true)) return false
+ if (requestObserver == null) return false
+ if (currentState.isTerminal) return false
+ try {
+ requestLock.synchronized {
+ // A terminal that arrived first must win; bail without writing. We
+ // compute the block's value rather than `return`-ing from this by-name
+ // `synchronized` body, whose non-local return would compile to a thrown
+ // NonLocalReturnControl that the surrounding NonFatal catch must not
+ // swallow -- brittle to rely on. Reads naturally as the block result.
+ val cur = currentState
+ if (cur.isTerminal) {
+ false
+ } else {
+ // Record that a Cancel is on the wire by advancing to Cancelling.
+ compareAndSetState(cur, SessionState.Cancelling)
+ requestObserver.onNext(UdfRequest.newBuilder()
+ .setControl(UdfControlRequest.newBuilder().setCancel(cancel()).build())
+ .build())
+ true
+ }
+ }
+ } catch {
+ case NonFatal(e) =>
+ logger.debug(s"Cancel send failed (stream may already be torn down): ${e.getMessage}")
+ completeTerminal(Termination.TransportFailed(e))
+ false
+ }
+ }
+
+ private def awaitTerminal(): Unit = {
+ if (currentState.isTerminal) return
+ try {
+ if (!terminalLatch.await(terminalTimeoutMs, TimeUnit.MILLISECONDS)) {
+ completeTerminal(Termination.TransportFailed(new TimeoutException(
+ s"timed out waiting for stream terminator after ${terminalTimeoutMs}ms")))
+ }
+ } catch {
+ case _: InterruptedException =>
+ // Attribute the terminal to the interrupt, not a timeout: the wait was
+ // cut short, it did not exhaust terminalTimeoutMs. Reporting "timed
+ // out" here would misdiagnose an interrupted close as a slow worker.
+ Thread.currentThread().interrupt()
+ completeTerminal(Termination.TransportFailed(
+ new InterruptedException("interrupted while waiting for stream terminator")))
+ }
+ }
+
+ // ---- 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 settle the state) and with any thread that finalizes via
+ * [[close]]. It drives the `Streaming -> Finishing` transition (sending
+ * `Finish` once input is exhausted, built lazily from the `finish` thunk).
+ */
Review Comment:
A readability suggestion: the state machine is well documented, but the actual transition edges are currently spread across `handleControl`, `doInit`, `doClose`, `sendCancelInternal`, and `ProcessIterator`. I found myself relying on the comments to reconstruct the transition graph.
Would it be worth centralizing the protocol transitions behind small semantic helpers, while still using the existing `WorkerSession.SessionState` as the single source of truth? For example:
```scala
private object Transitions {
def initAccepted(): Boolean =
compareAndSetState(SessionState.Initializing, SessionState.Initialized)
def finishSent(): Boolean =
compareAndSetState(SessionState.Streaming, SessionState.Finishing)
def cancelSentFrom(cur: SessionState): Boolean =
!cur.isTerminal && compareAndSetState(cur, SessionState.Cancelling)
def finished(response: FinishResponse): Boolean =
completeTerminal(Termination.Finished(response))
def cancelled(response: CancelResponse): Boolean =
completeTerminal(Termination.Cancelled(response))
def failed(error: ExecutionError): Boolean =
completeTerminal(Termination.Failed(error))
def transportFailed(cause: Throwable): Boolean =
completeTerminal(Termination.TransportFailed(cause))
}
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57640] Implement gRPC WorkerSession and test it against an echo server. [spark]
Posted by "haiyangsun-db (via GitHub)" <gi...@apache.org>.
haiyangsun-db commented on code in PR #56702:
URL: https://github.com/apache/spark/pull/56702#discussion_r3665894658
##########
udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/GrpcWorkerSession.scala:
##########
@@ -0,0 +1,933 @@
+/*
+ * 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.grpc
+
+import java.util.concurrent.{CountDownLatch, LinkedBlockingQueue, TimeoutException, TimeUnit}
+import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference}
+
+import scala.util.control.NonFatal
+
+import io.grpc.{ConnectivityState, ManagedChannel}
+import io.grpc.stub.StreamObserver
+
+import org.apache.spark.annotation.Experimental
+import org.apache.spark.udf.worker.{Cancel, CancelResponse, DataRequest, DataResponse,
+ ExecutionError, Finish, FinishResponse, Init, InitResponse, UdfControlRequest,
+ UdfControlResponse, UdfRequest, UdfResponse, UdfWorkerGrpc}
+import org.apache.spark.udf.worker.core.{Termination, WorkerHandle, WorkerLogger, WorkerSession}
+import org.apache.spark.udf.worker.core.WorkerSession.SessionState
+import org.apache.spark.udf.worker.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_message.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.
+ *
+ * '''Driving model.''' Consumption-driven (Volcano / pull): the thread that
+ * consumes the [[doProcess]] result iterator is the one that pulls input and
+ * sends each `DataRequest`; the gRPC callback thread only receives output. It is
+ * pull-driven but not one-input-per-output -- `advance` sends the next input
+ * whenever the output queue is momentarily empty, so under async delivery it may
+ * push several input batches before any output is read (bounded by HTTP/2 flow
+ * control).
+ *
+ * '''State machine.''' This class does not keep its own state machine: it
+ * drives the single [[WorkerSession.SessionState]] owned by the base. The base
+ * advances `Created -> Initializing` (in `init`) and `Initialized -> Streaming`
+ * (in `process`); this class advances the protocol-event edges through
+ * [[compareAndSetState]] / [[completeTerminal]] as it exchanges messages:
+ * {{{
+ * Initializing --(InitResponse ok)--> Initialized [handleControl]
+ * Streaming ----(Finish written)----> Finishing [ProcessIterator]
+ * <any non-terminal> --(Cancel written)--> Cancelling [sendCancelInternal]
+ * <any non-terminal> --(terminator/error)--> terminal [completeTerminal]
+ * }}}
+ * The two clean terminals carry the worker's `FinishResponse` / `CancelResponse`
+ * (metrics + finish/cancel callback `data`/`error`) so [[close]] can return
+ * them. The only flag kept outside the machine is [[cancelRequested]]: a
+ * cancellation can be requested before the stream exists (so it cannot be a
+ * state transition yet), and it must both fast-fail the result iterator and
+ * suppress any in-flight Data/Finish.
+ *
+ * Threading:
+ * - [[doInit]] is synchronous: sends `Init` and blocks on `InitResponse`,
+ * returning it.
+ * - [[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.
+ * - [[doClose]] is thread-safe and idempotent: it settles + returns the
+ * terminator (cancelling in-flight work if the stream had not finished)
+ * and half-closes the request side.
+ *
+ * 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 a gRPC channel built and owned by the caller (the
+ * dispatcher). 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 from the worker, drained by the process() iterator.
+ // Intentionally unbounded: a bounded queue would block the gRPC callback
+ // (Netty event-loop) thread when full, stalling terminator/control delivery on
+ // the whole channel. HTTP/2 flow control bounds the wire and the consumer
+ // normally drains promptly; a stalled downstream can still grow it, but the fix
+ // is protocol-level back-pressure (out of scope), not bounding the queue.
+ //
+ // TODO [SPARK-57324]: expose queue depth as a metric (early warning for a
+ // stalled consumer).
+ private val outputQueue = new LinkedBlockingQueue[QueueItem]()
+
+ // The worker's `InitResponse` (success or error) coupled with the latch that
+ // init() blocks on. Fires when the InitResponse arrives (`complete`) or when a
+ // terminal settles first without one -- transport error, half-close, or a
+ // premature terminator (`signalWithoutValue`). Until it fires we have no proof
+ // the worker accepted the session.
+ //
+ // Settle-before-release rule (referenced from every callback that both settles
+ // a terminal and fires this latch): settle the terminal FIRST, then complete /
+ // signal. init() blocks on the latch, and only latch await/release establish a
+ // happens-before edge, so a woken init() is guaranteed to observe the terminal
+ // rather than a transient state. OneShotValue keeps that publish-then-release
+ // in one place instead of every caller remembering to count down after setting
+ // the reference.
+ private val initValue = new OneShotValue[InitResponse]
+
+ // Fired when the session reaches a terminal [[SessionState]]. doClose() and
+ // the init-error path block on this to drain the terminator.
+ private val terminalLatch = new CountDownLatch(1)
+
+ // Captures an ErrorResponse encountered during the data phase so that
+ // the CancelResponse terminator 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)
+
+ // Cancellation INTENT -- distinct from the `Cancelling` state, which is reached
+ // only once a Cancel is actually written to the wire ([[sendCancelInternal]]).
+ // Intent is set first and can outrun (or never reach) that write, so
+ // `cancelRequested` does NOT imply state `Cancelling`. Kept outside the
+ // [[SessionState]] machine because a cancel can be requested before the stream
+ // exists (pre-init), where there is no wire transition to make yet. Used to (a)
+ // make cancel idempotent across all call sites, (b) fast-fail ProcessIterator
+ // on a pre-init cancel, and (c) suppress any Data/Finish that would otherwise
+ // race a Cancel onto the wire (re-read inside [[requestLock]]).
+ private val cancelRequested = new AtomicBoolean(false)
+
+ // gRPC requires serialized writes to a request StreamObserver.
+ private val requestLock = new Object
+
+ // Initialised in init() -- before that, close() is a no-op 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 =>
+ // A DataResponse before InitResponse violates the protocol (InitResponse
+ // must precede any DataResponse): fast-fail rather than enqueue it and let
+ // init() block to its timeout. Settle-before-release (see initValue).
+ if (!initResolved) {
+ Transitions.transportFailed(new IllegalStateException(
+ "worker sent a DataResponse before InitResponse"))
+ initValue.signalWithoutValue()
+ } else {
+ outputQueue.put(QueueItem.Batch(response.getData))
+ }
+
+ case UdfResponse.ResponseCase.CONTROL =>
+ handleControl(response.getControl)
+
+ case other =>
+ // A malformed response (empty / unknown oneof) is a terminal transport
+ // failure; fast-fail init the same way. Settle-before-release (see initValue).
+ Transitions.transportFailed(new IllegalStateException(
+ s"unexpected response oneof: $other"))
+ initValue.signalWithoutValue()
+ }
+ }
+
+ override def onError(t: Throwable): Unit = {
+ // Transport-level failure: the stream is dead, no further writes possible.
+ // Settle-before-release (see initValue) so init() surfaces the transport
+ // cause instead of the initResponseTimeoutMs "timed out" error.
+ Transitions.transportFailed(t)
+ initValue.signalWithoutValue()
+ }
+
+ 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 (!currentState.isTerminal) {
+ Transitions.transportFailed(new IllegalStateException(
+ "worker response stream closed without a terminator"))
+ }
+ // Defensive: if onCompleted reached us before InitResponse, doInit is
+ // still blocked on initValue and would otherwise time out.
+ initValue.signalWithoutValue()
+ }
+ }
+
+ /**
+ * Wakes the result iterator (blocked on [[outputQueue]]) and any thread
+ * waiting on [[terminalLatch]] when the base settles a terminal. Invoked once,
+ * by the caller that wins [[completeTerminal]].
+ */
+ override protected def onTerminalSettled(termination: Termination): Unit = {
+ outputQueue.put(QueueItem.EndOfStream)
+ terminalLatch.countDown()
+ }
+
+ private def handleControl(ctrl: UdfControlResponse): Unit = ctrl.getControlCase match {
+ case UdfControlResponse.ControlCase.INIT =>
+ val resp = ctrl.getInit
+ if (resp.hasError) {
+ // Attribute a subsequent close()'s terminator to the init error rather
+ // than a bare Cancelled, mirroring the data-phase ERROR branch.
+ executionError.compareAndSet(None, Some(resp.getError))
+ // Settle Failed(err) before the Cancel below (not after its CancelResponse):
+ // the sticky terminal makes close() report the init error, and the Cancel
+ // becomes best-effort worker cleanup -- so if it can't be written (a
+ // directExecutor worker delivered this reentrantly, before doInit published
+ // requestObserver), doInit still doesn't wait terminalTimeoutMs for a
+ // CancelResponse that can never arrive. Settle-before-release (see initValue).
+ Transitions.failed(resp.getError)
+ initValue.complete(resp)
+ sendCancelInternal(() => cancelWithReason("init failed"))
+ } else {
+ // InitResponse OK. Only advance from Initializing so a terminal that raced
+ // in (e.g. a transport error) still wins; process() then opens the data
+ // phase. Settle-before-release (see initValue): publish after the CAS.
+ Transitions.initAccepted()
+ initValue.complete(resp)
+ }
+
+ case UdfControlResponse.ControlCase.ERROR =>
+ val err = ctrl.getError.getError
+ executionError.compareAndSet(None, Some(err))
+ // Before init resolves, settle Failed here rather than rely on the Cancel
+ // below: sendCancelInternal may find requestObserver still null (reentrant
+ // delivery inside stream.onNext, before doInit published it), so no Cancel --
+ // hence no CancelResponse -- would settle the terminal, and doInit would
+ // return as if init succeeded. completeTerminal is idempotent, so the
+ // data-phase ERROR -> Cancel -> CancelResponse path is unaffected.
+ if (!initResolved) {
+ Transitions.failed(err)
+ }
+ // Settle-before-release (see initValue). No value to publish -- ERROR is not
+ // an InitResponse; on the data-phase path the latch already fired in init.
+ initValue.signalWithoutValue()
+ sendCancelInternal(() => cancelWithReason("aborting after ErrorResponse"))
+
+ case UdfControlResponse.ControlCase.FINISH =>
+ // The FinishResponse carries metrics + the finish-callback data/error.
+ // Keep it on the terminal so close() can return it; the iterator inspects
+ // its error field to decide whether to throw.
+ Transitions.finished(ctrl.getFinish)
+ // Defensive: FINISH before InitResponse is a worker protocol bug, but
+ // we should fail init fast rather than hang the 30s init timeout.
+ initValue.signalWithoutValue()
+
+ case UdfControlResponse.ControlCase.CANCEL =>
+ // The CancelResponse carries metrics + the cancel-callback error. Keep it
+ // on the terminal so close() can return it; any prior ErrorResponse is
+ // tracked in executionError and surfaced by the iterator.
+ Transitions.cancelled(ctrl.getCancel)
+ // Defensive: CANCEL before InitResponse unblocks doInit so it can
+ // surface the cancellation instead of timing out.
+ initValue.signalWithoutValue()
+
+ case UdfControlResponse.ControlCase.CONTROL_NOT_SET =>
+ Transitions.transportFailed(new IllegalStateException(
+ "empty UdfControlResponse oneof"))
+ initValue.signalWithoutValue()
+ }
+
+ /**
+ * True once init is no longer pending -- i.e. the stream is past `Initializing`.
+ * Not "init succeeded": a terminal (including a failure) also counts as resolved.
+ */
+ private def initResolved: Boolean = currentState match {
+ case SessionState.Created | SessionState.Initializing => false
+ case _ => true
+ }
+
+ private def cancelWithReason(reason: String): Cancel =
+ Cancel.newBuilder().setReason(reason).build()
+
+ /**
+ * The protocol transition graph in one place: names for the edges, not a
+ * second source of truth. Every edge acts on the single
+ * [[WorkerSession.SessionState]] owned by the base via [[compareAndSetState]]
+ * (non-terminal) or [[completeTerminal]] (terminal), so the base's CAS is the
+ * only synchronization and a terminal that arrived first always wins (the
+ * non-terminal CASes fail against it; [[completeTerminal]] is first-wins).
+ *
+ * Edges this class drives -- edge, method, then driver site(s) / thread (the
+ * base drives the API-call edges: `Created -> Initializing` in `init`,
+ * `Initialized -> Streaming` in `process`):
+ * {{{
+ * Initializing -> Initialized initAccepted handleControl INIT-ok [gRPC cb]
+ * Streaming -> Finishing finishSent advance branch 3 [engine]
+ * non-terminal -> Cancelling cancelSentFrom sendCancelInternal, iff a Cancel
+ * reaches the wire [gRPC cb | engine | init | close]
+ * non-terminal -> Terminal finished/cancelled/failed/transportFailed
+ * handleControl / doInit / doClose / advance / onError
+ * }}}
+ * `Cancelling` is reached only if a Cancel is actually written; a pre-stream or
+ * raced cancel goes straight to a `Cancelled`/`TransportFailed` terminal (see
+ * [[sendCancelInternal]], [[doClose]], `ProcessIterator`) or nowhere -- so
+ * `cancelRequested` (intent) does not imply state `Cancelling`.
+ */
+ private object Transitions {
+ /** `InitResponse` OK: `Initializing -> Initialized`. */
+ def initAccepted(): Boolean =
+ compareAndSetState(SessionState.Initializing, SessionState.Initialized)
+
+ /** Input exhausted and `Finish` written: `Streaming -> Finishing` (once). */
+ def finishSent(): Boolean =
+ compareAndSetState(SessionState.Streaming, SessionState.Finishing)
+
+ /** `Cancel` written: `cur -> Cancelling`, from any non-terminal `cur`. */
+ def cancelSentFrom(cur: SessionState): Boolean =
+ !cur.isTerminal && compareAndSetState(cur, SessionState.Cancelling)
+
+ /** Clean terminal carrying the worker's `FinishResponse`. */
+ def finished(response: FinishResponse): Boolean =
+ completeTerminal(Termination.Finished(response))
+
+ /** Clean terminal carrying the worker's `CancelResponse`. */
+ def cancelled(response: CancelResponse): Boolean =
+ completeTerminal(Termination.Cancelled(response))
+
+ /** Failure terminal carrying a structured [[ExecutionError]]. */
+ def failed(error: ExecutionError): Boolean =
+ completeTerminal(Termination.Failed(error))
+
+ /** Failure terminal carrying a transport-level cause. */
+ def transportFailed(cause: Throwable): Boolean =
+ completeTerminal(Termination.TransportFailed(cause))
+ }
+
+ // ---- WorkerSession hooks ------------------------------------------------
+
+ override protected def doInit(message: Init): InitResponse = {
+ // 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")
+ Transitions.transportFailed(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. close()'s 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 {
+ // Reentrant delivery: a directExecutor worker (one whose gRPC callbacks run
+ // synchronously on the caller's thread) can deliver InitResponse -- and thus
+ // run responseObserver/handleControl -- from *inside* this stream.onNext,
+ // before requestObserver below is published. The base advanced
+ // Created -> Initializing before calling doInit, so that reentrant
+ // InitResponse still finds Initializing and advances to Initialized; the
+ // handleControl error paths guard the still-null requestObserver. This is
+ // the scenario the "before doInit published requestObserver" comments mean.
+ 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
+ Transitions.transportFailed(e)
+ // Surface as GrpcWorkerSessionException so the engine integration layer
+ // (which catches that type and wraps it) sees a uniform init-failure
+ // exception rather than the raw transport error.
+ throw new GrpcWorkerSessionException("UDF worker stream failed during init", e)
+ }
+
+ val responded = try {
+ initValue.await(initResponseTimeoutMs)
+ } catch {
+ case _: InterruptedException =>
+ Thread.currentThread().interrupt()
+ sendCancelInternal(() => cancelWithReason("interrupted during init"))
+ // Make sure close() does not block waiting for a terminator that
+ // will never arrive on this thread's behalf.
+ Transitions.transportFailed(
+ new InterruptedException("interrupted while waiting for InitResponse"))
+ throw new InterruptedException("interrupted while waiting for InitResponse")
+ }
+
+ if (!responded) {
+ sendCancelInternal(() => cancelWithReason("InitResponse timed out"))
+ val timeout = new TimeoutException(
+ s"timed out waiting for InitResponse after ${initResponseTimeoutMs}ms")
+ // Settle the terminal so a subsequent close() does not stall for a
+ // second `terminalTimeoutMs` waiting for a worker that already missed
+ // its init deadline.
+ Transitions.transportFailed(timeout)
+ // Surface as GrpcWorkerSessionException (carrying the timeout cause) so
+ // the engine integration layer that catches that type can wrap it.
+ throw new GrpcWorkerSessionException(
+ s"timed out waiting for InitResponse after ${initResponseTimeoutMs}ms", timeout)
+ }
+
+ initValue.get match {
+ case Some(resp) if resp.hasError =>
+ // 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(resp.getError)}", resp.getError)
+ case Some(resp) =>
+ resp
+ case None =>
+ // No InitResponse arrived but the latch fired. The defensive
+ // initValue.signalWithoutValue() 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.
+ currentState match {
+ case SessionState.Terminal(Termination.TransportFailed(cause)) =>
+ throw new GrpcWorkerSessionException(
+ "UDF worker stream failed during init", cause)
+ case SessionState.Terminal(Termination.Failed(err)) =>
+ throw new GrpcWorkerSessionException(
+ s"UDF worker reported an error before init completed: " +
+ describeError(err), err)
+ case SessionState.Terminal(Termination.Cancelled(_)) =>
+ throw new GrpcWorkerSessionException(
+ "UDF worker stream was cancelled before init completed")
+ case SessionState.Terminal(Termination.Finished(_)) =>
+ throw new GrpcWorkerSessionException(
+ "UDF worker finished before init completed")
+ case other =>
+ throw new IllegalStateException(
+ s"init latch fired without an InitResponse or terminal: $other")
+ }
+ }
+ }
+
+ override protected def doProcess(
+ input: Iterator[DataRequest],
+ finish: () => Finish): Iterator[DataResponse] = {
+ // Init success is guaranteed by the base [[WorkerSession]] lifecycle: if
+ // doInit had failed it would have thrown and process() would never run.
+ new ProcessIterator(input, finish)
+ }
+
+ override protected def doClose(cancel: () => Cancel): Termination = {
+ if (requestObserver == null) {
+ // init() never put a stream on the wire (closed before/around init, or
+ // init threw before publishing). There is no protocol terminator; treat
+ // the session as cancelled-before-start. The base WorkerSession still
+ // releases the worker handle, so the worker is torn down.
+ Transitions.cancelled(CancelResponse.getDefaultInstance)
+ return Termination.Cancelled(CancelResponse.getDefaultInstance)
Review Comment:
L529 only settles cancelled if (!currentState.isTerminal), then return settledTermination (L531).
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57640][UDF] Implement gRPC WorkerSession and test it against an echo server [spark]
Posted by "Yicong-Huang (via GitHub)" <gi...@apache.org>.
Yicong-Huang commented on PR #56702:
URL: https://github.com/apache/spark/pull/56702#issuecomment-5175199991
**Merge Summary:**
- merged into master https://github.com/apache/spark/commit/fc6b3cd23e576520f67f046f7bf88509181ad4b0
- merged into branch-4.x https://github.com/apache/spark/commit/3dd9265e070891b95764c9f0dc7f0ccb02369fb6
*Posted by `merge_spark_pr.py`*
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57640] Implement gRPC WorkerSession and test it against an echo server. [spark]
Posted by "haiyangsun-db (via GitHub)" <gi...@apache.org>.
haiyangsun-db commented on code in PR #56702:
URL: https://github.com/apache/spark/pull/56702#discussion_r3665874818
##########
udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/GrpcWorkerSession.scala:
##########
@@ -0,0 +1,933 @@
+/*
+ * 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.grpc
+
+import java.util.concurrent.{CountDownLatch, LinkedBlockingQueue, TimeoutException, TimeUnit}
+import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference}
+
+import scala.util.control.NonFatal
+
+import io.grpc.{ConnectivityState, ManagedChannel}
+import io.grpc.stub.StreamObserver
+
+import org.apache.spark.annotation.Experimental
+import org.apache.spark.udf.worker.{Cancel, CancelResponse, DataRequest, DataResponse,
+ ExecutionError, Finish, FinishResponse, Init, InitResponse, UdfControlRequest,
+ UdfControlResponse, UdfRequest, UdfResponse, UdfWorkerGrpc}
+import org.apache.spark.udf.worker.core.{Termination, WorkerHandle, WorkerLogger, WorkerSession}
+import org.apache.spark.udf.worker.core.WorkerSession.SessionState
+import org.apache.spark.udf.worker.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_message.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.
+ *
+ * '''Driving model.''' Consumption-driven (Volcano / pull): the thread that
+ * consumes the [[doProcess]] result iterator is the one that pulls input and
+ * sends each `DataRequest`; the gRPC callback thread only receives output. It is
+ * pull-driven but not one-input-per-output -- `advance` sends the next input
+ * whenever the output queue is momentarily empty, so under async delivery it may
+ * push several input batches before any output is read (bounded by HTTP/2 flow
+ * control).
+ *
+ * '''State machine.''' This class does not keep its own state machine: it
+ * drives the single [[WorkerSession.SessionState]] owned by the base. The base
+ * advances `Created -> Initializing` (in `init`) and `Initialized -> Streaming`
+ * (in `process`); this class advances the protocol-event edges through
+ * [[compareAndSetState]] / [[completeTerminal]] as it exchanges messages:
+ * {{{
+ * Initializing --(InitResponse ok)--> Initialized [handleControl]
+ * Streaming ----(Finish written)----> Finishing [ProcessIterator]
+ * <any non-terminal> --(Cancel written)--> Cancelling [sendCancelInternal]
+ * <any non-terminal> --(terminator/error)--> terminal [completeTerminal]
+ * }}}
+ * The two clean terminals carry the worker's `FinishResponse` / `CancelResponse`
+ * (metrics + finish/cancel callback `data`/`error`) so [[close]] can return
+ * them. The only flag kept outside the machine is [[cancelRequested]]: a
+ * cancellation can be requested before the stream exists (so it cannot be a
+ * state transition yet), and it must both fast-fail the result iterator and
+ * suppress any in-flight Data/Finish.
+ *
+ * Threading:
+ * - [[doInit]] is synchronous: sends `Init` and blocks on `InitResponse`,
+ * returning it.
+ * - [[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.
+ * - [[doClose]] is thread-safe and idempotent: it settles + returns the
+ * terminator (cancelling in-flight work if the stream had not finished)
+ * and half-closes the request side.
+ *
+ * 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 a gRPC channel built and owned by the caller (the
+ * dispatcher). 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 from the worker, drained by the process() iterator.
+ // Intentionally unbounded: a bounded queue would block the gRPC callback
+ // (Netty event-loop) thread when full, stalling terminator/control delivery on
+ // the whole channel. HTTP/2 flow control bounds the wire and the consumer
+ // normally drains promptly; a stalled downstream can still grow it, but the fix
+ // is protocol-level back-pressure (out of scope), not bounding the queue.
+ //
+ // TODO [SPARK-57324]: expose queue depth as a metric (early warning for a
+ // stalled consumer).
+ private val outputQueue = new LinkedBlockingQueue[QueueItem]()
+
+ // The worker's `InitResponse` (success or error) coupled with the latch that
+ // init() blocks on. Fires when the InitResponse arrives (`complete`) or when a
+ // terminal settles first without one -- transport error, half-close, or a
+ // premature terminator (`signalWithoutValue`). Until it fires we have no proof
+ // the worker accepted the session.
+ //
+ // Settle-before-release rule (referenced from every callback that both settles
+ // a terminal and fires this latch): settle the terminal FIRST, then complete /
+ // signal. init() blocks on the latch, and only latch await/release establish a
+ // happens-before edge, so a woken init() is guaranteed to observe the terminal
+ // rather than a transient state. OneShotValue keeps that publish-then-release
+ // in one place instead of every caller remembering to count down after setting
+ // the reference.
+ private val initValue = new OneShotValue[InitResponse]
+
+ // Fired when the session reaches a terminal [[SessionState]]. doClose() and
+ // the init-error path block on this to drain the terminator.
+ private val terminalLatch = new CountDownLatch(1)
+
+ // Captures an ErrorResponse encountered during the data phase so that
+ // the CancelResponse terminator 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)
+
+ // Cancellation INTENT -- distinct from the `Cancelling` state, which is reached
+ // only once a Cancel is actually written to the wire ([[sendCancelInternal]]).
+ // Intent is set first and can outrun (or never reach) that write, so
+ // `cancelRequested` does NOT imply state `Cancelling`. Kept outside the
+ // [[SessionState]] machine because a cancel can be requested before the stream
+ // exists (pre-init), where there is no wire transition to make yet. Used to (a)
+ // make cancel idempotent across all call sites, (b) fast-fail ProcessIterator
+ // on a pre-init cancel, and (c) suppress any Data/Finish that would otherwise
+ // race a Cancel onto the wire (re-read inside [[requestLock]]).
+ private val cancelRequested = new AtomicBoolean(false)
+
+ // gRPC requires serialized writes to a request StreamObserver.
+ private val requestLock = new Object
+
+ // Initialised in init() -- before that, close() is a no-op 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 =>
+ // A DataResponse before InitResponse violates the protocol (InitResponse
+ // must precede any DataResponse): fast-fail rather than enqueue it and let
+ // init() block to its timeout. Settle-before-release (see initValue).
+ if (!initResolved) {
+ Transitions.transportFailed(new IllegalStateException(
+ "worker sent a DataResponse before InitResponse"))
+ initValue.signalWithoutValue()
+ } else {
+ outputQueue.put(QueueItem.Batch(response.getData))
+ }
+
+ case UdfResponse.ResponseCase.CONTROL =>
+ handleControl(response.getControl)
+
+ case other =>
+ // A malformed response (empty / unknown oneof) is a terminal transport
+ // failure; fast-fail init the same way. Settle-before-release (see initValue).
+ Transitions.transportFailed(new IllegalStateException(
+ s"unexpected response oneof: $other"))
+ initValue.signalWithoutValue()
+ }
+ }
+
+ override def onError(t: Throwable): Unit = {
+ // Transport-level failure: the stream is dead, no further writes possible.
+ // Settle-before-release (see initValue) so init() surfaces the transport
+ // cause instead of the initResponseTimeoutMs "timed out" error.
+ Transitions.transportFailed(t)
+ initValue.signalWithoutValue()
+ }
+
+ 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 (!currentState.isTerminal) {
+ Transitions.transportFailed(new IllegalStateException(
+ "worker response stream closed without a terminator"))
+ }
+ // Defensive: if onCompleted reached us before InitResponse, doInit is
+ // still blocked on initValue and would otherwise time out.
+ initValue.signalWithoutValue()
+ }
+ }
+
+ /**
+ * Wakes the result iterator (blocked on [[outputQueue]]) and any thread
+ * waiting on [[terminalLatch]] when the base settles a terminal. Invoked once,
+ * by the caller that wins [[completeTerminal]].
+ */
+ override protected def onTerminalSettled(termination: Termination): Unit = {
+ outputQueue.put(QueueItem.EndOfStream)
+ terminalLatch.countDown()
+ }
+
+ private def handleControl(ctrl: UdfControlResponse): Unit = ctrl.getControlCase match {
+ case UdfControlResponse.ControlCase.INIT =>
+ val resp = ctrl.getInit
+ if (resp.hasError) {
+ // Attribute a subsequent close()'s terminator to the init error rather
+ // than a bare Cancelled, mirroring the data-phase ERROR branch.
+ executionError.compareAndSet(None, Some(resp.getError))
+ // Settle Failed(err) before the Cancel below (not after its CancelResponse):
+ // the sticky terminal makes close() report the init error, and the Cancel
+ // becomes best-effort worker cleanup -- so if it can't be written (a
+ // directExecutor worker delivered this reentrantly, before doInit published
+ // requestObserver), doInit still doesn't wait terminalTimeoutMs for a
+ // CancelResponse that can never arrive. Settle-before-release (see initValue).
+ Transitions.failed(resp.getError)
+ initValue.complete(resp)
+ sendCancelInternal(() => cancelWithReason("init failed"))
+ } else {
+ // InitResponse OK. Only advance from Initializing so a terminal that raced
+ // in (e.g. a transport error) still wins; process() then opens the data
+ // phase. Settle-before-release (see initValue): publish after the CAS.
+ Transitions.initAccepted()
+ initValue.complete(resp)
Review Comment:
Both INIT-ok (initAccepted() then complete) and the error branches now order settle/record before signal
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57640] Implement gRPC WorkerSession and test it against an echo server. [spark]
Posted by "haiyangsun-db (via GitHub)" <gi...@apache.org>.
haiyangsun-db commented on PR #56702:
URL: https://github.com/apache/spark/pull/56702#issuecomment-5106945232
@Yicong-Huang , Thanks for the thorough pass — this was really helpful. I've pushed a series of commits addressing everything; summary below, with per-thread replies inline. Would you please help take another look?
Some summary of the changes:
▎ HL-2 (Cancel-after-init-error via Cancelling, not early Failed): done — the init-error path no longer settles Failed in the callback; it records the error, and doInit (where requestObserver is published) sends the Cancel and drains the CancelResponse, so the terminal is Cancelled. Transitions.failed was removed entirely.
â–Ž
▎ HL-1 (init-handshake complexity): addressed by cleanup rather than the full handler-split — details in the inline reply on that thread.
â–Ž
▎ I also took the interrupt discussion (#5) further into a dedicated Interrupted termination — see that thread.
---
HL-1 — "init logic still feels too complex / isolate init as a one-shot"
▎ You're right that the signaling was the main source of noise, and it turned out to be mostly accidental. onTerminalSettled is now the single place a settled terminal wakes a blocked init(), so every branch that settles a terminal (the responseObserver failure/terminator cases, the timeout paths, close()) no longer signals the init waiter itself — it just settles. That removed 7 redundant signalWithoutValue() calls; the only direct init signals left are the three genuinely non-terminal paths (INIT-ok, INIT-error, pre-init ERROR), which must wake init() without settling a terminal so the session stays Initializing and process() can still run. That's the irreducible core of the handshake.
▎ On fully isolating init behind a one-shot handler and swapping to a data-phase handler once it resolves: I looked at this and decided against it. gRPC gives us a single StreamObserver for the stream's whole life, so "switching" means either a swappable @volatile var handler or an if (initResolved) … at the top of onNext. The former adds a mutable concurrency surface and a swap-vs-reentrant-delivery race — exactly the trickiest part of this class — and the latter just relocates the two initResolved checks that already exist. Neither removes the essential constraint (reentrant directExecutor delivery forcing the deferred Cancel), so I don't think the extra indirection pays for itself. Happy to revisit if you feel strongly.
---
#1 (id 3593697357) — "settling failed before sendCancelInternal means Cancel is never written; contradicts proto"
▎ Fixed. The INIT-error branch no longer settles a terminal in the callback — it records the error (executionError) and completes the init waiter, leaving the session in Initializing. doInit then runs failInitWithError, which sends the Cancel (where requestObserver is guaranteed published) and drains the CancelResponse, so the terminal settles Cancelled and the proto's "Cancel after init error" holds on both the directExecutor and cross-thread paths.
---
#2 (id 3593697365) — "awaitTerminal() is a no-op; drain comment is stale"
▎ Fixed by the same change as #1. Since the terminal is no longer settled early, awaitTerminal() in failInitWithError now genuinely blocks for the CancelResponse before throwing — the drain is real, and the comment is accurate again.
---
#3 (id 3593697374) — "doClose null-path returns a fresh Cancelled even when a terminal already settled"
▎ Fixed. The null-path now only settles Cancelled if (!currentState.isTerminal), then returns settledTermination — so an already-settled terminal (e.g. the channel-shutdown TransportFailed) is returned as-is rather than a bare Cancelled that disagrees with the state.
---
#4 (id 3593856689) — "can TransportFailed be merged to failed?"
▎ This one's moot now: Transitions.failed was removed (the init-error path settles Cancelled via the drain, not Failed). The failure terminals are deliberately distinct — TransportFailed (unsalvageable, unknown worker state) vs. the new Interrupted (salvageable, engine-side event) — so I'd keep them separate rather than merge. See #5.
---
#5 (id 3593874187) — "who can interrupt this thread? cancel vs failed?"
▎ Good question — it prompted a dedicated outcome. An interrupt here is an engine-side event (cancelled query / killed task), not a worker fault, so it's now neither Cancelled nor Failed: it settles a new Termination.Interrupted. handleInterrupt sends a best-effort Cancel and waits a short, configurable interruptCancelTimeoutMs (default 2s) for the ack — if the worker responds in time the session settles a clean, salvageable Cancelled; otherwise it falls back to Interrupted (also salvageable, since the worker isn't at fault). It's applied at every blocking point (init, terminal drain, close, result poll), not just init. There's a TODO [SPARK-57640] to back the optimistic "salvageable" assumption with a worker liveness check.
---
#6 (id 3593884935) — "make await raise a timeout error instead of returning a Boolean"
â–Ž Done. OneShotValue.await now returns Unit and throws TimeoutException on expiry, so doInit handles the timeout on the exception path instead of inspecting a return value.
---
#7 (id 3593892046) — "complete the value before transit the state?"
â–Ž I kept the order as transition-then-complete here, deliberately. Transitions.initAccepted() is a CAS from Initializing that a raced terminal (e.g. a transport error) must be able to lose against; completing the value first would publish an InitResponse that a subsequently-winning terminal contradicts. Completing after the CAS preserves settle/transition-before-release: init() wakes to a state consistent with the value it reads. Added a comment noting this.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57640] Implement gRPC WorkerSession and test it against an echo server. [spark]
Posted by "Yicong-Huang (via GitHub)" <gi...@apache.org>.
Yicong-Huang commented on code in PR #56702:
URL: https://github.com/apache/spark/pull/56702#discussion_r3593697357
##########
udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/GrpcWorkerSession.scala:
##########
@@ -0,0 +1,933 @@
+/*
+ * 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.grpc
+
+import java.util.concurrent.{CountDownLatch, LinkedBlockingQueue, TimeoutException, TimeUnit}
+import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference}
+
+import scala.util.control.NonFatal
+
+import io.grpc.{ConnectivityState, ManagedChannel}
+import io.grpc.stub.StreamObserver
+
+import org.apache.spark.annotation.Experimental
+import org.apache.spark.udf.worker.{Cancel, CancelResponse, DataRequest, DataResponse,
+ ExecutionError, Finish, FinishResponse, Init, InitResponse, UdfControlRequest,
+ UdfControlResponse, UdfRequest, UdfResponse, UdfWorkerGrpc}
+import org.apache.spark.udf.worker.core.{Termination, WorkerHandle, WorkerLogger, WorkerSession}
+import org.apache.spark.udf.worker.core.WorkerSession.SessionState
+import org.apache.spark.udf.worker.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_message.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.
+ *
+ * '''Driving model.''' Consumption-driven (Volcano / pull): the thread that
+ * consumes the [[doProcess]] result iterator is the one that pulls input and
+ * sends each `DataRequest`; the gRPC callback thread only receives output. It is
+ * pull-driven but not one-input-per-output -- `advance` sends the next input
+ * whenever the output queue is momentarily empty, so under async delivery it may
+ * push several input batches before any output is read (bounded by HTTP/2 flow
+ * control).
+ *
+ * '''State machine.''' This class does not keep its own state machine: it
+ * drives the single [[WorkerSession.SessionState]] owned by the base. The base
+ * advances `Created -> Initializing` (in `init`) and `Initialized -> Streaming`
+ * (in `process`); this class advances the protocol-event edges through
+ * [[compareAndSetState]] / [[completeTerminal]] as it exchanges messages:
+ * {{{
+ * Initializing --(InitResponse ok)--> Initialized [handleControl]
+ * Streaming ----(Finish written)----> Finishing [ProcessIterator]
+ * <any non-terminal> --(Cancel written)--> Cancelling [sendCancelInternal]
+ * <any non-terminal> --(terminator/error)--> terminal [completeTerminal]
+ * }}}
+ * The two clean terminals carry the worker's `FinishResponse` / `CancelResponse`
+ * (metrics + finish/cancel callback `data`/`error`) so [[close]] can return
+ * them. The only flag kept outside the machine is [[cancelRequested]]: a
+ * cancellation can be requested before the stream exists (so it cannot be a
+ * state transition yet), and it must both fast-fail the result iterator and
+ * suppress any in-flight Data/Finish.
+ *
+ * Threading:
+ * - [[doInit]] is synchronous: sends `Init` and blocks on `InitResponse`,
+ * returning it.
+ * - [[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.
+ * - [[doClose]] is thread-safe and idempotent: it settles + returns the
+ * terminator (cancelling in-flight work if the stream had not finished)
+ * and half-closes the request side.
+ *
+ * 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 a gRPC channel built and owned by the caller (the
+ * dispatcher). 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 from the worker, drained by the process() iterator.
+ // Intentionally unbounded: a bounded queue would block the gRPC callback
+ // (Netty event-loop) thread when full, stalling terminator/control delivery on
+ // the whole channel. HTTP/2 flow control bounds the wire and the consumer
+ // normally drains promptly; a stalled downstream can still grow it, but the fix
+ // is protocol-level back-pressure (out of scope), not bounding the queue.
+ //
+ // TODO [SPARK-57324]: expose queue depth as a metric (early warning for a
+ // stalled consumer).
+ private val outputQueue = new LinkedBlockingQueue[QueueItem]()
+
+ // The worker's `InitResponse` (success or error) coupled with the latch that
+ // init() blocks on. Fires when the InitResponse arrives (`complete`) or when a
+ // terminal settles first without one -- transport error, half-close, or a
+ // premature terminator (`signalWithoutValue`). Until it fires we have no proof
+ // the worker accepted the session.
+ //
+ // Settle-before-release rule (referenced from every callback that both settles
+ // a terminal and fires this latch): settle the terminal FIRST, then complete /
+ // signal. init() blocks on the latch, and only latch await/release establish a
+ // happens-before edge, so a woken init() is guaranteed to observe the terminal
+ // rather than a transient state. OneShotValue keeps that publish-then-release
+ // in one place instead of every caller remembering to count down after setting
+ // the reference.
+ private val initValue = new OneShotValue[InitResponse]
+
+ // Fired when the session reaches a terminal [[SessionState]]. doClose() and
+ // the init-error path block on this to drain the terminator.
+ private val terminalLatch = new CountDownLatch(1)
+
+ // Captures an ErrorResponse encountered during the data phase so that
+ // the CancelResponse terminator 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)
+
+ // Cancellation INTENT -- distinct from the `Cancelling` state, which is reached
+ // only once a Cancel is actually written to the wire ([[sendCancelInternal]]).
+ // Intent is set first and can outrun (or never reach) that write, so
+ // `cancelRequested` does NOT imply state `Cancelling`. Kept outside the
+ // [[SessionState]] machine because a cancel can be requested before the stream
+ // exists (pre-init), where there is no wire transition to make yet. Used to (a)
+ // make cancel idempotent across all call sites, (b) fast-fail ProcessIterator
+ // on a pre-init cancel, and (c) suppress any Data/Finish that would otherwise
+ // race a Cancel onto the wire (re-read inside [[requestLock]]).
+ private val cancelRequested = new AtomicBoolean(false)
+
+ // gRPC requires serialized writes to a request StreamObserver.
+ private val requestLock = new Object
+
+ // Initialised in init() -- before that, close() is a no-op 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 =>
+ // A DataResponse before InitResponse violates the protocol (InitResponse
+ // must precede any DataResponse): fast-fail rather than enqueue it and let
+ // init() block to its timeout. Settle-before-release (see initValue).
+ if (!initResolved) {
+ Transitions.transportFailed(new IllegalStateException(
+ "worker sent a DataResponse before InitResponse"))
+ initValue.signalWithoutValue()
+ } else {
+ outputQueue.put(QueueItem.Batch(response.getData))
+ }
+
+ case UdfResponse.ResponseCase.CONTROL =>
+ handleControl(response.getControl)
+
+ case other =>
+ // A malformed response (empty / unknown oneof) is a terminal transport
+ // failure; fast-fail init the same way. Settle-before-release (see initValue).
+ Transitions.transportFailed(new IllegalStateException(
+ s"unexpected response oneof: $other"))
+ initValue.signalWithoutValue()
+ }
+ }
+
+ override def onError(t: Throwable): Unit = {
+ // Transport-level failure: the stream is dead, no further writes possible.
+ // Settle-before-release (see initValue) so init() surfaces the transport
+ // cause instead of the initResponseTimeoutMs "timed out" error.
+ Transitions.transportFailed(t)
+ initValue.signalWithoutValue()
+ }
+
+ 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 (!currentState.isTerminal) {
+ Transitions.transportFailed(new IllegalStateException(
+ "worker response stream closed without a terminator"))
+ }
+ // Defensive: if onCompleted reached us before InitResponse, doInit is
+ // still blocked on initValue and would otherwise time out.
+ initValue.signalWithoutValue()
+ }
+ }
+
+ /**
+ * Wakes the result iterator (blocked on [[outputQueue]]) and any thread
+ * waiting on [[terminalLatch]] when the base settles a terminal. Invoked once,
+ * by the caller that wins [[completeTerminal]].
+ */
+ override protected def onTerminalSettled(termination: Termination): Unit = {
+ outputQueue.put(QueueItem.EndOfStream)
+ terminalLatch.countDown()
+ }
+
+ private def handleControl(ctrl: UdfControlResponse): Unit = ctrl.getControlCase match {
+ case UdfControlResponse.ControlCase.INIT =>
+ val resp = ctrl.getInit
+ if (resp.hasError) {
+ // Attribute a subsequent close()'s terminator to the init error rather
+ // than a bare Cancelled, mirroring the data-phase ERROR branch.
+ executionError.compareAndSet(None, Some(resp.getError))
+ // Settle Failed(err) before the Cancel below (not after its CancelResponse):
+ // the sticky terminal makes close() report the init error, and the Cancel
+ // becomes best-effort worker cleanup -- so if it can't be written (a
+ // directExecutor worker delivered this reentrantly, before doInit published
+ // requestObserver), doInit still doesn't wait terminalTimeoutMs for a
+ // CancelResponse that can never arrive. Settle-before-release (see initValue).
+ Transitions.failed(resp.getError)
Review Comment:
Settling `failed` before `sendCancelInternal` means the Cancel is now never written on an init error: `sendCancelInternal` bails on the directExecutor path (requestObserver null) AND the cross-thread path (state already terminal). But udf_message.proto (~L91) says the engine MUST send Cancel after an init error. These two are contradictory?
##########
udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/GrpcWorkerSession.scala:
##########
@@ -0,0 +1,933 @@
+/*
+ * 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.grpc
+
+import java.util.concurrent.{CountDownLatch, LinkedBlockingQueue, TimeoutException, TimeUnit}
+import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference}
+
+import scala.util.control.NonFatal
+
+import io.grpc.{ConnectivityState, ManagedChannel}
+import io.grpc.stub.StreamObserver
+
+import org.apache.spark.annotation.Experimental
+import org.apache.spark.udf.worker.{Cancel, CancelResponse, DataRequest, DataResponse,
+ ExecutionError, Finish, FinishResponse, Init, InitResponse, UdfControlRequest,
+ UdfControlResponse, UdfRequest, UdfResponse, UdfWorkerGrpc}
+import org.apache.spark.udf.worker.core.{Termination, WorkerHandle, WorkerLogger, WorkerSession}
+import org.apache.spark.udf.worker.core.WorkerSession.SessionState
+import org.apache.spark.udf.worker.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_message.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.
+ *
+ * '''Driving model.''' Consumption-driven (Volcano / pull): the thread that
+ * consumes the [[doProcess]] result iterator is the one that pulls input and
+ * sends each `DataRequest`; the gRPC callback thread only receives output. It is
+ * pull-driven but not one-input-per-output -- `advance` sends the next input
+ * whenever the output queue is momentarily empty, so under async delivery it may
+ * push several input batches before any output is read (bounded by HTTP/2 flow
+ * control).
+ *
+ * '''State machine.''' This class does not keep its own state machine: it
+ * drives the single [[WorkerSession.SessionState]] owned by the base. The base
+ * advances `Created -> Initializing` (in `init`) and `Initialized -> Streaming`
+ * (in `process`); this class advances the protocol-event edges through
+ * [[compareAndSetState]] / [[completeTerminal]] as it exchanges messages:
+ * {{{
+ * Initializing --(InitResponse ok)--> Initialized [handleControl]
+ * Streaming ----(Finish written)----> Finishing [ProcessIterator]
+ * <any non-terminal> --(Cancel written)--> Cancelling [sendCancelInternal]
+ * <any non-terminal> --(terminator/error)--> terminal [completeTerminal]
+ * }}}
+ * The two clean terminals carry the worker's `FinishResponse` / `CancelResponse`
+ * (metrics + finish/cancel callback `data`/`error`) so [[close]] can return
+ * them. The only flag kept outside the machine is [[cancelRequested]]: a
+ * cancellation can be requested before the stream exists (so it cannot be a
+ * state transition yet), and it must both fast-fail the result iterator and
+ * suppress any in-flight Data/Finish.
+ *
+ * Threading:
+ * - [[doInit]] is synchronous: sends `Init` and blocks on `InitResponse`,
+ * returning it.
+ * - [[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.
+ * - [[doClose]] is thread-safe and idempotent: it settles + returns the
+ * terminator (cancelling in-flight work if the stream had not finished)
+ * and half-closes the request side.
+ *
+ * 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 a gRPC channel built and owned by the caller (the
+ * dispatcher). 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 from the worker, drained by the process() iterator.
+ // Intentionally unbounded: a bounded queue would block the gRPC callback
+ // (Netty event-loop) thread when full, stalling terminator/control delivery on
+ // the whole channel. HTTP/2 flow control bounds the wire and the consumer
+ // normally drains promptly; a stalled downstream can still grow it, but the fix
+ // is protocol-level back-pressure (out of scope), not bounding the queue.
+ //
+ // TODO [SPARK-57324]: expose queue depth as a metric (early warning for a
+ // stalled consumer).
+ private val outputQueue = new LinkedBlockingQueue[QueueItem]()
+
+ // The worker's `InitResponse` (success or error) coupled with the latch that
+ // init() blocks on. Fires when the InitResponse arrives (`complete`) or when a
+ // terminal settles first without one -- transport error, half-close, or a
+ // premature terminator (`signalWithoutValue`). Until it fires we have no proof
+ // the worker accepted the session.
+ //
+ // Settle-before-release rule (referenced from every callback that both settles
+ // a terminal and fires this latch): settle the terminal FIRST, then complete /
+ // signal. init() blocks on the latch, and only latch await/release establish a
+ // happens-before edge, so a woken init() is guaranteed to observe the terminal
+ // rather than a transient state. OneShotValue keeps that publish-then-release
+ // in one place instead of every caller remembering to count down after setting
+ // the reference.
+ private val initValue = new OneShotValue[InitResponse]
+
+ // Fired when the session reaches a terminal [[SessionState]]. doClose() and
+ // the init-error path block on this to drain the terminator.
+ private val terminalLatch = new CountDownLatch(1)
+
+ // Captures an ErrorResponse encountered during the data phase so that
+ // the CancelResponse terminator 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)
+
+ // Cancellation INTENT -- distinct from the `Cancelling` state, which is reached
+ // only once a Cancel is actually written to the wire ([[sendCancelInternal]]).
+ // Intent is set first and can outrun (or never reach) that write, so
+ // `cancelRequested` does NOT imply state `Cancelling`. Kept outside the
+ // [[SessionState]] machine because a cancel can be requested before the stream
+ // exists (pre-init), where there is no wire transition to make yet. Used to (a)
+ // make cancel idempotent across all call sites, (b) fast-fail ProcessIterator
+ // on a pre-init cancel, and (c) suppress any Data/Finish that would otherwise
+ // race a Cancel onto the wire (re-read inside [[requestLock]]).
+ private val cancelRequested = new AtomicBoolean(false)
+
+ // gRPC requires serialized writes to a request StreamObserver.
+ private val requestLock = new Object
+
+ // Initialised in init() -- before that, close() is a no-op 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 =>
+ // A DataResponse before InitResponse violates the protocol (InitResponse
+ // must precede any DataResponse): fast-fail rather than enqueue it and let
+ // init() block to its timeout. Settle-before-release (see initValue).
+ if (!initResolved) {
+ Transitions.transportFailed(new IllegalStateException(
+ "worker sent a DataResponse before InitResponse"))
+ initValue.signalWithoutValue()
+ } else {
+ outputQueue.put(QueueItem.Batch(response.getData))
+ }
+
+ case UdfResponse.ResponseCase.CONTROL =>
+ handleControl(response.getControl)
+
+ case other =>
+ // A malformed response (empty / unknown oneof) is a terminal transport
+ // failure; fast-fail init the same way. Settle-before-release (see initValue).
+ Transitions.transportFailed(new IllegalStateException(
+ s"unexpected response oneof: $other"))
+ initValue.signalWithoutValue()
+ }
+ }
+
+ override def onError(t: Throwable): Unit = {
+ // Transport-level failure: the stream is dead, no further writes possible.
+ // Settle-before-release (see initValue) so init() surfaces the transport
+ // cause instead of the initResponseTimeoutMs "timed out" error.
+ Transitions.transportFailed(t)
+ initValue.signalWithoutValue()
+ }
+
+ 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 (!currentState.isTerminal) {
+ Transitions.transportFailed(new IllegalStateException(
+ "worker response stream closed without a terminator"))
+ }
+ // Defensive: if onCompleted reached us before InitResponse, doInit is
+ // still blocked on initValue and would otherwise time out.
+ initValue.signalWithoutValue()
+ }
+ }
+
+ /**
+ * Wakes the result iterator (blocked on [[outputQueue]]) and any thread
+ * waiting on [[terminalLatch]] when the base settles a terminal. Invoked once,
+ * by the caller that wins [[completeTerminal]].
+ */
+ override protected def onTerminalSettled(termination: Termination): Unit = {
+ outputQueue.put(QueueItem.EndOfStream)
+ terminalLatch.countDown()
+ }
+
+ private def handleControl(ctrl: UdfControlResponse): Unit = ctrl.getControlCase match {
+ case UdfControlResponse.ControlCase.INIT =>
+ val resp = ctrl.getInit
+ if (resp.hasError) {
+ // Attribute a subsequent close()'s terminator to the init error rather
+ // than a bare Cancelled, mirroring the data-phase ERROR branch.
+ executionError.compareAndSet(None, Some(resp.getError))
+ // Settle Failed(err) before the Cancel below (not after its CancelResponse):
+ // the sticky terminal makes close() report the init error, and the Cancel
+ // becomes best-effort worker cleanup -- so if it can't be written (a
+ // directExecutor worker delivered this reentrantly, before doInit published
+ // requestObserver), doInit still doesn't wait terminalTimeoutMs for a
+ // CancelResponse that can never arrive. Settle-before-release (see initValue).
+ Transitions.failed(resp.getError)
+ initValue.complete(resp)
+ sendCancelInternal(() => cancelWithReason("init failed"))
+ } else {
+ // InitResponse OK. Only advance from Initializing so a terminal that raced
+ // in (e.g. a transport error) still wins; process() then opens the data
+ // phase. Settle-before-release (see initValue): publish after the CAS.
+ Transitions.initAccepted()
+ initValue.complete(resp)
+ }
+
+ case UdfControlResponse.ControlCase.ERROR =>
+ val err = ctrl.getError.getError
+ executionError.compareAndSet(None, Some(err))
+ // Before init resolves, settle Failed here rather than rely on the Cancel
+ // below: sendCancelInternal may find requestObserver still null (reentrant
+ // delivery inside stream.onNext, before doInit published it), so no Cancel --
+ // hence no CancelResponse -- would settle the terminal, and doInit would
+ // return as if init succeeded. completeTerminal is idempotent, so the
+ // data-phase ERROR -> Cancel -> CancelResponse path is unaffected.
+ if (!initResolved) {
+ Transitions.failed(err)
+ }
+ // Settle-before-release (see initValue). No value to publish -- ERROR is not
+ // an InitResponse; on the data-phase path the latch already fired in init.
+ initValue.signalWithoutValue()
+ sendCancelInternal(() => cancelWithReason("aborting after ErrorResponse"))
+
+ case UdfControlResponse.ControlCase.FINISH =>
+ // The FinishResponse carries metrics + the finish-callback data/error.
+ // Keep it on the terminal so close() can return it; the iterator inspects
+ // its error field to decide whether to throw.
+ Transitions.finished(ctrl.getFinish)
+ // Defensive: FINISH before InitResponse is a worker protocol bug, but
+ // we should fail init fast rather than hang the 30s init timeout.
+ initValue.signalWithoutValue()
+
+ case UdfControlResponse.ControlCase.CANCEL =>
+ // The CancelResponse carries metrics + the cancel-callback error. Keep it
+ // on the terminal so close() can return it; any prior ErrorResponse is
+ // tracked in executionError and surfaced by the iterator.
+ Transitions.cancelled(ctrl.getCancel)
+ // Defensive: CANCEL before InitResponse unblocks doInit so it can
+ // surface the cancellation instead of timing out.
+ initValue.signalWithoutValue()
+
+ case UdfControlResponse.ControlCase.CONTROL_NOT_SET =>
+ Transitions.transportFailed(new IllegalStateException(
+ "empty UdfControlResponse oneof"))
+ initValue.signalWithoutValue()
+ }
+
+ /**
+ * True once init is no longer pending -- i.e. the stream is past `Initializing`.
+ * Not "init succeeded": a terminal (including a failure) also counts as resolved.
+ */
+ private def initResolved: Boolean = currentState match {
+ case SessionState.Created | SessionState.Initializing => false
+ case _ => true
+ }
+
+ private def cancelWithReason(reason: String): Cancel =
+ Cancel.newBuilder().setReason(reason).build()
+
+ /**
+ * The protocol transition graph in one place: names for the edges, not a
+ * second source of truth. Every edge acts on the single
+ * [[WorkerSession.SessionState]] owned by the base via [[compareAndSetState]]
+ * (non-terminal) or [[completeTerminal]] (terminal), so the base's CAS is the
+ * only synchronization and a terminal that arrived first always wins (the
+ * non-terminal CASes fail against it; [[completeTerminal]] is first-wins).
+ *
+ * Edges this class drives -- edge, method, then driver site(s) / thread (the
+ * base drives the API-call edges: `Created -> Initializing` in `init`,
+ * `Initialized -> Streaming` in `process`):
+ * {{{
+ * Initializing -> Initialized initAccepted handleControl INIT-ok [gRPC cb]
+ * Streaming -> Finishing finishSent advance branch 3 [engine]
+ * non-terminal -> Cancelling cancelSentFrom sendCancelInternal, iff a Cancel
+ * reaches the wire [gRPC cb | engine | init | close]
+ * non-terminal -> Terminal finished/cancelled/failed/transportFailed
+ * handleControl / doInit / doClose / advance / onError
+ * }}}
+ * `Cancelling` is reached only if a Cancel is actually written; a pre-stream or
+ * raced cancel goes straight to a `Cancelled`/`TransportFailed` terminal (see
+ * [[sendCancelInternal]], [[doClose]], `ProcessIterator`) or nowhere -- so
+ * `cancelRequested` (intent) does not imply state `Cancelling`.
+ */
+ private object Transitions {
+ /** `InitResponse` OK: `Initializing -> Initialized`. */
+ def initAccepted(): Boolean =
+ compareAndSetState(SessionState.Initializing, SessionState.Initialized)
+
+ /** Input exhausted and `Finish` written: `Streaming -> Finishing` (once). */
+ def finishSent(): Boolean =
+ compareAndSetState(SessionState.Streaming, SessionState.Finishing)
+
+ /** `Cancel` written: `cur -> Cancelling`, from any non-terminal `cur`. */
+ def cancelSentFrom(cur: SessionState): Boolean =
+ !cur.isTerminal && compareAndSetState(cur, SessionState.Cancelling)
+
+ /** Clean terminal carrying the worker's `FinishResponse`. */
+ def finished(response: FinishResponse): Boolean =
+ completeTerminal(Termination.Finished(response))
+
+ /** Clean terminal carrying the worker's `CancelResponse`. */
+ def cancelled(response: CancelResponse): Boolean =
+ completeTerminal(Termination.Cancelled(response))
+
+ /** Failure terminal carrying a structured [[ExecutionError]]. */
+ def failed(error: ExecutionError): Boolean =
+ completeTerminal(Termination.Failed(error))
+
+ /** Failure terminal carrying a transport-level cause. */
+ def transportFailed(cause: Throwable): Boolean =
+ completeTerminal(Termination.TransportFailed(cause))
+ }
+
+ // ---- WorkerSession hooks ------------------------------------------------
+
+ override protected def doInit(message: Init): InitResponse = {
+ // 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")
+ Transitions.transportFailed(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. close()'s 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 {
+ // Reentrant delivery: a directExecutor worker (one whose gRPC callbacks run
+ // synchronously on the caller's thread) can deliver InitResponse -- and thus
+ // run responseObserver/handleControl -- from *inside* this stream.onNext,
+ // before requestObserver below is published. The base advanced
+ // Created -> Initializing before calling doInit, so that reentrant
+ // InitResponse still finds Initializing and advances to Initialized; the
+ // handleControl error paths guard the still-null requestObserver. This is
+ // the scenario the "before doInit published requestObserver" comments mean.
+ 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
+ Transitions.transportFailed(e)
+ // Surface as GrpcWorkerSessionException so the engine integration layer
+ // (which catches that type and wraps it) sees a uniform init-failure
+ // exception rather than the raw transport error.
+ throw new GrpcWorkerSessionException("UDF worker stream failed during init", e)
+ }
+
+ val responded = try {
+ initValue.await(initResponseTimeoutMs)
+ } catch {
+ case _: InterruptedException =>
+ Thread.currentThread().interrupt()
+ sendCancelInternal(() => cancelWithReason("interrupted during init"))
+ // Make sure close() does not block waiting for a terminator that
+ // will never arrive on this thread's behalf.
+ Transitions.transportFailed(
+ new InterruptedException("interrupted while waiting for InitResponse"))
+ throw new InterruptedException("interrupted while waiting for InitResponse")
+ }
+
+ if (!responded) {
+ sendCancelInternal(() => cancelWithReason("InitResponse timed out"))
+ val timeout = new TimeoutException(
+ s"timed out waiting for InitResponse after ${initResponseTimeoutMs}ms")
+ // Settle the terminal so a subsequent close() does not stall for a
+ // second `terminalTimeoutMs` waiting for a worker that already missed
+ // its init deadline.
+ Transitions.transportFailed(timeout)
+ // Surface as GrpcWorkerSessionException (carrying the timeout cause) so
+ // the engine integration layer that catches that type can wrap it.
+ throw new GrpcWorkerSessionException(
+ s"timed out waiting for InitResponse after ${initResponseTimeoutMs}ms", timeout)
+ }
+
+ initValue.get match {
+ case Some(resp) if resp.hasError =>
+ // 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()
Review Comment:
`awaitTerminal()` is a no-op here: the terminal was already settled Failed at L250 before `initValue.complete`, so this returns immediately. The "drain it before throwing so we don't leave a dangling stream" comment (L443-445) is now stale -- nothing is drained. Drop the dead call + comment, or restore an actual drain if one is intended.
##########
udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/GrpcWorkerSession.scala:
##########
@@ -0,0 +1,933 @@
+/*
+ * 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.grpc
+
+import java.util.concurrent.{CountDownLatch, LinkedBlockingQueue, TimeoutException, TimeUnit}
+import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference}
+
+import scala.util.control.NonFatal
+
+import io.grpc.{ConnectivityState, ManagedChannel}
+import io.grpc.stub.StreamObserver
+
+import org.apache.spark.annotation.Experimental
+import org.apache.spark.udf.worker.{Cancel, CancelResponse, DataRequest, DataResponse,
+ ExecutionError, Finish, FinishResponse, Init, InitResponse, UdfControlRequest,
+ UdfControlResponse, UdfRequest, UdfResponse, UdfWorkerGrpc}
+import org.apache.spark.udf.worker.core.{Termination, WorkerHandle, WorkerLogger, WorkerSession}
+import org.apache.spark.udf.worker.core.WorkerSession.SessionState
+import org.apache.spark.udf.worker.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_message.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.
+ *
+ * '''Driving model.''' Consumption-driven (Volcano / pull): the thread that
+ * consumes the [[doProcess]] result iterator is the one that pulls input and
+ * sends each `DataRequest`; the gRPC callback thread only receives output. It is
+ * pull-driven but not one-input-per-output -- `advance` sends the next input
+ * whenever the output queue is momentarily empty, so under async delivery it may
+ * push several input batches before any output is read (bounded by HTTP/2 flow
+ * control).
+ *
+ * '''State machine.''' This class does not keep its own state machine: it
+ * drives the single [[WorkerSession.SessionState]] owned by the base. The base
+ * advances `Created -> Initializing` (in `init`) and `Initialized -> Streaming`
+ * (in `process`); this class advances the protocol-event edges through
+ * [[compareAndSetState]] / [[completeTerminal]] as it exchanges messages:
+ * {{{
+ * Initializing --(InitResponse ok)--> Initialized [handleControl]
+ * Streaming ----(Finish written)----> Finishing [ProcessIterator]
+ * <any non-terminal> --(Cancel written)--> Cancelling [sendCancelInternal]
+ * <any non-terminal> --(terminator/error)--> terminal [completeTerminal]
+ * }}}
+ * The two clean terminals carry the worker's `FinishResponse` / `CancelResponse`
+ * (metrics + finish/cancel callback `data`/`error`) so [[close]] can return
+ * them. The only flag kept outside the machine is [[cancelRequested]]: a
+ * cancellation can be requested before the stream exists (so it cannot be a
+ * state transition yet), and it must both fast-fail the result iterator and
+ * suppress any in-flight Data/Finish.
+ *
+ * Threading:
+ * - [[doInit]] is synchronous: sends `Init` and blocks on `InitResponse`,
+ * returning it.
+ * - [[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.
+ * - [[doClose]] is thread-safe and idempotent: it settles + returns the
+ * terminator (cancelling in-flight work if the stream had not finished)
+ * and half-closes the request side.
+ *
+ * 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 a gRPC channel built and owned by the caller (the
+ * dispatcher). 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 from the worker, drained by the process() iterator.
+ // Intentionally unbounded: a bounded queue would block the gRPC callback
+ // (Netty event-loop) thread when full, stalling terminator/control delivery on
+ // the whole channel. HTTP/2 flow control bounds the wire and the consumer
+ // normally drains promptly; a stalled downstream can still grow it, but the fix
+ // is protocol-level back-pressure (out of scope), not bounding the queue.
+ //
+ // TODO [SPARK-57324]: expose queue depth as a metric (early warning for a
+ // stalled consumer).
+ private val outputQueue = new LinkedBlockingQueue[QueueItem]()
+
+ // The worker's `InitResponse` (success or error) coupled with the latch that
+ // init() blocks on. Fires when the InitResponse arrives (`complete`) or when a
+ // terminal settles first without one -- transport error, half-close, or a
+ // premature terminator (`signalWithoutValue`). Until it fires we have no proof
+ // the worker accepted the session.
+ //
+ // Settle-before-release rule (referenced from every callback that both settles
+ // a terminal and fires this latch): settle the terminal FIRST, then complete /
+ // signal. init() blocks on the latch, and only latch await/release establish a
+ // happens-before edge, so a woken init() is guaranteed to observe the terminal
+ // rather than a transient state. OneShotValue keeps that publish-then-release
+ // in one place instead of every caller remembering to count down after setting
+ // the reference.
+ private val initValue = new OneShotValue[InitResponse]
+
+ // Fired when the session reaches a terminal [[SessionState]]. doClose() and
+ // the init-error path block on this to drain the terminator.
+ private val terminalLatch = new CountDownLatch(1)
+
+ // Captures an ErrorResponse encountered during the data phase so that
+ // the CancelResponse terminator 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)
+
+ // Cancellation INTENT -- distinct from the `Cancelling` state, which is reached
+ // only once a Cancel is actually written to the wire ([[sendCancelInternal]]).
+ // Intent is set first and can outrun (or never reach) that write, so
+ // `cancelRequested` does NOT imply state `Cancelling`. Kept outside the
+ // [[SessionState]] machine because a cancel can be requested before the stream
+ // exists (pre-init), where there is no wire transition to make yet. Used to (a)
+ // make cancel idempotent across all call sites, (b) fast-fail ProcessIterator
+ // on a pre-init cancel, and (c) suppress any Data/Finish that would otherwise
+ // race a Cancel onto the wire (re-read inside [[requestLock]]).
+ private val cancelRequested = new AtomicBoolean(false)
+
+ // gRPC requires serialized writes to a request StreamObserver.
+ private val requestLock = new Object
+
+ // Initialised in init() -- before that, close() is a no-op 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 =>
+ // A DataResponse before InitResponse violates the protocol (InitResponse
+ // must precede any DataResponse): fast-fail rather than enqueue it and let
+ // init() block to its timeout. Settle-before-release (see initValue).
+ if (!initResolved) {
+ Transitions.transportFailed(new IllegalStateException(
+ "worker sent a DataResponse before InitResponse"))
+ initValue.signalWithoutValue()
+ } else {
+ outputQueue.put(QueueItem.Batch(response.getData))
+ }
+
+ case UdfResponse.ResponseCase.CONTROL =>
+ handleControl(response.getControl)
+
+ case other =>
+ // A malformed response (empty / unknown oneof) is a terminal transport
+ // failure; fast-fail init the same way. Settle-before-release (see initValue).
+ Transitions.transportFailed(new IllegalStateException(
+ s"unexpected response oneof: $other"))
+ initValue.signalWithoutValue()
+ }
+ }
+
+ override def onError(t: Throwable): Unit = {
+ // Transport-level failure: the stream is dead, no further writes possible.
+ // Settle-before-release (see initValue) so init() surfaces the transport
+ // cause instead of the initResponseTimeoutMs "timed out" error.
+ Transitions.transportFailed(t)
+ initValue.signalWithoutValue()
+ }
+
+ 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 (!currentState.isTerminal) {
+ Transitions.transportFailed(new IllegalStateException(
+ "worker response stream closed without a terminator"))
+ }
+ // Defensive: if onCompleted reached us before InitResponse, doInit is
+ // still blocked on initValue and would otherwise time out.
+ initValue.signalWithoutValue()
+ }
+ }
+
+ /**
+ * Wakes the result iterator (blocked on [[outputQueue]]) and any thread
+ * waiting on [[terminalLatch]] when the base settles a terminal. Invoked once,
+ * by the caller that wins [[completeTerminal]].
+ */
+ override protected def onTerminalSettled(termination: Termination): Unit = {
+ outputQueue.put(QueueItem.EndOfStream)
+ terminalLatch.countDown()
+ }
+
+ private def handleControl(ctrl: UdfControlResponse): Unit = ctrl.getControlCase match {
+ case UdfControlResponse.ControlCase.INIT =>
+ val resp = ctrl.getInit
+ if (resp.hasError) {
+ // Attribute a subsequent close()'s terminator to the init error rather
+ // than a bare Cancelled, mirroring the data-phase ERROR branch.
+ executionError.compareAndSet(None, Some(resp.getError))
+ // Settle Failed(err) before the Cancel below (not after its CancelResponse):
+ // the sticky terminal makes close() report the init error, and the Cancel
+ // becomes best-effort worker cleanup -- so if it can't be written (a
+ // directExecutor worker delivered this reentrantly, before doInit published
+ // requestObserver), doInit still doesn't wait terminalTimeoutMs for a
+ // CancelResponse that can never arrive. Settle-before-release (see initValue).
+ Transitions.failed(resp.getError)
+ initValue.complete(resp)
+ sendCancelInternal(() => cancelWithReason("init failed"))
+ } else {
+ // InitResponse OK. Only advance from Initializing so a terminal that raced
+ // in (e.g. a transport error) still wins; process() then opens the data
+ // phase. Settle-before-release (see initValue): publish after the CAS.
+ Transitions.initAccepted()
+ initValue.complete(resp)
+ }
+
+ case UdfControlResponse.ControlCase.ERROR =>
+ val err = ctrl.getError.getError
+ executionError.compareAndSet(None, Some(err))
+ // Before init resolves, settle Failed here rather than rely on the Cancel
+ // below: sendCancelInternal may find requestObserver still null (reentrant
+ // delivery inside stream.onNext, before doInit published it), so no Cancel --
+ // hence no CancelResponse -- would settle the terminal, and doInit would
+ // return as if init succeeded. completeTerminal is idempotent, so the
+ // data-phase ERROR -> Cancel -> CancelResponse path is unaffected.
+ if (!initResolved) {
+ Transitions.failed(err)
+ }
+ // Settle-before-release (see initValue). No value to publish -- ERROR is not
+ // an InitResponse; on the data-phase path the latch already fired in init.
+ initValue.signalWithoutValue()
+ sendCancelInternal(() => cancelWithReason("aborting after ErrorResponse"))
+
+ case UdfControlResponse.ControlCase.FINISH =>
+ // The FinishResponse carries metrics + the finish-callback data/error.
+ // Keep it on the terminal so close() can return it; the iterator inspects
+ // its error field to decide whether to throw.
+ Transitions.finished(ctrl.getFinish)
+ // Defensive: FINISH before InitResponse is a worker protocol bug, but
+ // we should fail init fast rather than hang the 30s init timeout.
+ initValue.signalWithoutValue()
+
+ case UdfControlResponse.ControlCase.CANCEL =>
+ // The CancelResponse carries metrics + the cancel-callback error. Keep it
+ // on the terminal so close() can return it; any prior ErrorResponse is
+ // tracked in executionError and surfaced by the iterator.
+ Transitions.cancelled(ctrl.getCancel)
+ // Defensive: CANCEL before InitResponse unblocks doInit so it can
+ // surface the cancellation instead of timing out.
+ initValue.signalWithoutValue()
+
+ case UdfControlResponse.ControlCase.CONTROL_NOT_SET =>
+ Transitions.transportFailed(new IllegalStateException(
+ "empty UdfControlResponse oneof"))
+ initValue.signalWithoutValue()
+ }
+
+ /**
+ * True once init is no longer pending -- i.e. the stream is past `Initializing`.
+ * Not "init succeeded": a terminal (including a failure) also counts as resolved.
+ */
+ private def initResolved: Boolean = currentState match {
+ case SessionState.Created | SessionState.Initializing => false
+ case _ => true
+ }
+
+ private def cancelWithReason(reason: String): Cancel =
+ Cancel.newBuilder().setReason(reason).build()
+
+ /**
+ * The protocol transition graph in one place: names for the edges, not a
+ * second source of truth. Every edge acts on the single
+ * [[WorkerSession.SessionState]] owned by the base via [[compareAndSetState]]
+ * (non-terminal) or [[completeTerminal]] (terminal), so the base's CAS is the
+ * only synchronization and a terminal that arrived first always wins (the
+ * non-terminal CASes fail against it; [[completeTerminal]] is first-wins).
+ *
+ * Edges this class drives -- edge, method, then driver site(s) / thread (the
+ * base drives the API-call edges: `Created -> Initializing` in `init`,
+ * `Initialized -> Streaming` in `process`):
+ * {{{
+ * Initializing -> Initialized initAccepted handleControl INIT-ok [gRPC cb]
+ * Streaming -> Finishing finishSent advance branch 3 [engine]
+ * non-terminal -> Cancelling cancelSentFrom sendCancelInternal, iff a Cancel
+ * reaches the wire [gRPC cb | engine | init | close]
+ * non-terminal -> Terminal finished/cancelled/failed/transportFailed
+ * handleControl / doInit / doClose / advance / onError
+ * }}}
+ * `Cancelling` is reached only if a Cancel is actually written; a pre-stream or
+ * raced cancel goes straight to a `Cancelled`/`TransportFailed` terminal (see
+ * [[sendCancelInternal]], [[doClose]], `ProcessIterator`) or nowhere -- so
+ * `cancelRequested` (intent) does not imply state `Cancelling`.
+ */
+ private object Transitions {
+ /** `InitResponse` OK: `Initializing -> Initialized`. */
+ def initAccepted(): Boolean =
+ compareAndSetState(SessionState.Initializing, SessionState.Initialized)
+
+ /** Input exhausted and `Finish` written: `Streaming -> Finishing` (once). */
+ def finishSent(): Boolean =
+ compareAndSetState(SessionState.Streaming, SessionState.Finishing)
+
+ /** `Cancel` written: `cur -> Cancelling`, from any non-terminal `cur`. */
+ def cancelSentFrom(cur: SessionState): Boolean =
+ !cur.isTerminal && compareAndSetState(cur, SessionState.Cancelling)
+
+ /** Clean terminal carrying the worker's `FinishResponse`. */
+ def finished(response: FinishResponse): Boolean =
+ completeTerminal(Termination.Finished(response))
+
+ /** Clean terminal carrying the worker's `CancelResponse`. */
+ def cancelled(response: CancelResponse): Boolean =
+ completeTerminal(Termination.Cancelled(response))
+
+ /** Failure terminal carrying a structured [[ExecutionError]]. */
+ def failed(error: ExecutionError): Boolean =
+ completeTerminal(Termination.Failed(error))
+
+ /** Failure terminal carrying a transport-level cause. */
+ def transportFailed(cause: Throwable): Boolean =
+ completeTerminal(Termination.TransportFailed(cause))
+ }
+
+ // ---- WorkerSession hooks ------------------------------------------------
+
+ override protected def doInit(message: Init): InitResponse = {
+ // 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")
+ Transitions.transportFailed(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. close()'s 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 {
+ // Reentrant delivery: a directExecutor worker (one whose gRPC callbacks run
+ // synchronously on the caller's thread) can deliver InitResponse -- and thus
+ // run responseObserver/handleControl -- from *inside* this stream.onNext,
+ // before requestObserver below is published. The base advanced
+ // Created -> Initializing before calling doInit, so that reentrant
+ // InitResponse still finds Initializing and advances to Initialized; the
+ // handleControl error paths guard the still-null requestObserver. This is
+ // the scenario the "before doInit published requestObserver" comments mean.
+ 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
+ Transitions.transportFailed(e)
+ // Surface as GrpcWorkerSessionException so the engine integration layer
+ // (which catches that type and wraps it) sees a uniform init-failure
+ // exception rather than the raw transport error.
+ throw new GrpcWorkerSessionException("UDF worker stream failed during init", e)
+ }
+
+ val responded = try {
+ initValue.await(initResponseTimeoutMs)
+ } catch {
+ case _: InterruptedException =>
+ Thread.currentThread().interrupt()
+ sendCancelInternal(() => cancelWithReason("interrupted during init"))
+ // Make sure close() does not block waiting for a terminator that
+ // will never arrive on this thread's behalf.
+ Transitions.transportFailed(
+ new InterruptedException("interrupted while waiting for InitResponse"))
+ throw new InterruptedException("interrupted while waiting for InitResponse")
+ }
+
+ if (!responded) {
+ sendCancelInternal(() => cancelWithReason("InitResponse timed out"))
+ val timeout = new TimeoutException(
+ s"timed out waiting for InitResponse after ${initResponseTimeoutMs}ms")
+ // Settle the terminal so a subsequent close() does not stall for a
+ // second `terminalTimeoutMs` waiting for a worker that already missed
+ // its init deadline.
+ Transitions.transportFailed(timeout)
+ // Surface as GrpcWorkerSessionException (carrying the timeout cause) so
+ // the engine integration layer that catches that type can wrap it.
+ throw new GrpcWorkerSessionException(
+ s"timed out waiting for InitResponse after ${initResponseTimeoutMs}ms", timeout)
+ }
+
+ initValue.get match {
+ case Some(resp) if resp.hasError =>
+ // 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(resp.getError)}", resp.getError)
+ case Some(resp) =>
+ resp
+ case None =>
+ // No InitResponse arrived but the latch fired. The defensive
+ // initValue.signalWithoutValue() 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.
+ currentState match {
+ case SessionState.Terminal(Termination.TransportFailed(cause)) =>
+ throw new GrpcWorkerSessionException(
+ "UDF worker stream failed during init", cause)
+ case SessionState.Terminal(Termination.Failed(err)) =>
+ throw new GrpcWorkerSessionException(
+ s"UDF worker reported an error before init completed: " +
+ describeError(err), err)
+ case SessionState.Terminal(Termination.Cancelled(_)) =>
+ throw new GrpcWorkerSessionException(
+ "UDF worker stream was cancelled before init completed")
+ case SessionState.Terminal(Termination.Finished(_)) =>
+ throw new GrpcWorkerSessionException(
+ "UDF worker finished before init completed")
+ case other =>
+ throw new IllegalStateException(
+ s"init latch fired without an InitResponse or terminal: $other")
+ }
+ }
+ }
+
+ override protected def doProcess(
+ input: Iterator[DataRequest],
+ finish: () => Finish): Iterator[DataResponse] = {
+ // Init success is guaranteed by the base [[WorkerSession]] lifecycle: if
+ // doInit had failed it would have thrown and process() would never run.
+ new ProcessIterator(input, finish)
+ }
+
+ override protected def doClose(cancel: () => Cancel): Termination = {
+ if (requestObserver == null) {
+ // init() never put a stream on the wire (closed before/around init, or
+ // init threw before publishing). There is no protocol terminator; treat
+ // the session as cancelled-before-start. The base WorkerSession still
+ // releases the worker handle, so the worker is torn down.
+ Transitions.cancelled(CancelResponse.getDefaultInstance)
+ return Termination.Cancelled(CancelResponse.getDefaultInstance)
Review Comment:
Returns a fresh `Cancelled` even when a terminal already settled (e.g. the channel-shutdown TransportFailed in doInit also leaves requestObserver null), so the return disagrees with the state. Return `settledTermination` instead. Harmless today since base close() discards this value, but misleading.
##########
udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/GrpcWorkerSession.scala:
##########
@@ -0,0 +1,933 @@
+/*
+ * 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.grpc
+
+import java.util.concurrent.{CountDownLatch, LinkedBlockingQueue, TimeoutException, TimeUnit}
+import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference}
+
+import scala.util.control.NonFatal
+
+import io.grpc.{ConnectivityState, ManagedChannel}
+import io.grpc.stub.StreamObserver
+
+import org.apache.spark.annotation.Experimental
+import org.apache.spark.udf.worker.{Cancel, CancelResponse, DataRequest, DataResponse,
+ ExecutionError, Finish, FinishResponse, Init, InitResponse, UdfControlRequest,
+ UdfControlResponse, UdfRequest, UdfResponse, UdfWorkerGrpc}
+import org.apache.spark.udf.worker.core.{Termination, WorkerHandle, WorkerLogger, WorkerSession}
+import org.apache.spark.udf.worker.core.WorkerSession.SessionState
+import org.apache.spark.udf.worker.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_message.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.
+ *
+ * '''Driving model.''' Consumption-driven (Volcano / pull): the thread that
+ * consumes the [[doProcess]] result iterator is the one that pulls input and
+ * sends each `DataRequest`; the gRPC callback thread only receives output. It is
+ * pull-driven but not one-input-per-output -- `advance` sends the next input
+ * whenever the output queue is momentarily empty, so under async delivery it may
+ * push several input batches before any output is read (bounded by HTTP/2 flow
+ * control).
+ *
+ * '''State machine.''' This class does not keep its own state machine: it
+ * drives the single [[WorkerSession.SessionState]] owned by the base. The base
+ * advances `Created -> Initializing` (in `init`) and `Initialized -> Streaming`
+ * (in `process`); this class advances the protocol-event edges through
+ * [[compareAndSetState]] / [[completeTerminal]] as it exchanges messages:
+ * {{{
+ * Initializing --(InitResponse ok)--> Initialized [handleControl]
+ * Streaming ----(Finish written)----> Finishing [ProcessIterator]
+ * <any non-terminal> --(Cancel written)--> Cancelling [sendCancelInternal]
+ * <any non-terminal> --(terminator/error)--> terminal [completeTerminal]
+ * }}}
+ * The two clean terminals carry the worker's `FinishResponse` / `CancelResponse`
+ * (metrics + finish/cancel callback `data`/`error`) so [[close]] can return
+ * them. The only flag kept outside the machine is [[cancelRequested]]: a
+ * cancellation can be requested before the stream exists (so it cannot be a
+ * state transition yet), and it must both fast-fail the result iterator and
+ * suppress any in-flight Data/Finish.
+ *
+ * Threading:
+ * - [[doInit]] is synchronous: sends `Init` and blocks on `InitResponse`,
+ * returning it.
+ * - [[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.
+ * - [[doClose]] is thread-safe and idempotent: it settles + returns the
+ * terminator (cancelling in-flight work if the stream had not finished)
+ * and half-closes the request side.
+ *
+ * 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 a gRPC channel built and owned by the caller (the
+ * dispatcher). 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 from the worker, drained by the process() iterator.
+ // Intentionally unbounded: a bounded queue would block the gRPC callback
+ // (Netty event-loop) thread when full, stalling terminator/control delivery on
+ // the whole channel. HTTP/2 flow control bounds the wire and the consumer
+ // normally drains promptly; a stalled downstream can still grow it, but the fix
+ // is protocol-level back-pressure (out of scope), not bounding the queue.
+ //
+ // TODO [SPARK-57324]: expose queue depth as a metric (early warning for a
+ // stalled consumer).
+ private val outputQueue = new LinkedBlockingQueue[QueueItem]()
+
+ // The worker's `InitResponse` (success or error) coupled with the latch that
+ // init() blocks on. Fires when the InitResponse arrives (`complete`) or when a
+ // terminal settles first without one -- transport error, half-close, or a
+ // premature terminator (`signalWithoutValue`). Until it fires we have no proof
+ // the worker accepted the session.
+ //
+ // Settle-before-release rule (referenced from every callback that both settles
+ // a terminal and fires this latch): settle the terminal FIRST, then complete /
+ // signal. init() blocks on the latch, and only latch await/release establish a
+ // happens-before edge, so a woken init() is guaranteed to observe the terminal
+ // rather than a transient state. OneShotValue keeps that publish-then-release
+ // in one place instead of every caller remembering to count down after setting
+ // the reference.
+ private val initValue = new OneShotValue[InitResponse]
+
+ // Fired when the session reaches a terminal [[SessionState]]. doClose() and
+ // the init-error path block on this to drain the terminator.
+ private val terminalLatch = new CountDownLatch(1)
+
+ // Captures an ErrorResponse encountered during the data phase so that
+ // the CancelResponse terminator 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)
+
+ // Cancellation INTENT -- distinct from the `Cancelling` state, which is reached
+ // only once a Cancel is actually written to the wire ([[sendCancelInternal]]).
+ // Intent is set first and can outrun (or never reach) that write, so
+ // `cancelRequested` does NOT imply state `Cancelling`. Kept outside the
+ // [[SessionState]] machine because a cancel can be requested before the stream
+ // exists (pre-init), where there is no wire transition to make yet. Used to (a)
+ // make cancel idempotent across all call sites, (b) fast-fail ProcessIterator
+ // on a pre-init cancel, and (c) suppress any Data/Finish that would otherwise
+ // race a Cancel onto the wire (re-read inside [[requestLock]]).
+ private val cancelRequested = new AtomicBoolean(false)
+
+ // gRPC requires serialized writes to a request StreamObserver.
+ private val requestLock = new Object
+
+ // Initialised in init() -- before that, close() is a no-op 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 =>
+ // A DataResponse before InitResponse violates the protocol (InitResponse
+ // must precede any DataResponse): fast-fail rather than enqueue it and let
+ // init() block to its timeout. Settle-before-release (see initValue).
+ if (!initResolved) {
+ Transitions.transportFailed(new IllegalStateException(
+ "worker sent a DataResponse before InitResponse"))
+ initValue.signalWithoutValue()
+ } else {
+ outputQueue.put(QueueItem.Batch(response.getData))
+ }
+
+ case UdfResponse.ResponseCase.CONTROL =>
+ handleControl(response.getControl)
+
+ case other =>
+ // A malformed response (empty / unknown oneof) is a terminal transport
+ // failure; fast-fail init the same way. Settle-before-release (see initValue).
+ Transitions.transportFailed(new IllegalStateException(
+ s"unexpected response oneof: $other"))
+ initValue.signalWithoutValue()
+ }
+ }
+
+ override def onError(t: Throwable): Unit = {
+ // Transport-level failure: the stream is dead, no further writes possible.
+ // Settle-before-release (see initValue) so init() surfaces the transport
+ // cause instead of the initResponseTimeoutMs "timed out" error.
+ Transitions.transportFailed(t)
+ initValue.signalWithoutValue()
+ }
+
+ 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 (!currentState.isTerminal) {
+ Transitions.transportFailed(new IllegalStateException(
+ "worker response stream closed without a terminator"))
+ }
+ // Defensive: if onCompleted reached us before InitResponse, doInit is
+ // still blocked on initValue and would otherwise time out.
+ initValue.signalWithoutValue()
+ }
+ }
+
+ /**
+ * Wakes the result iterator (blocked on [[outputQueue]]) and any thread
+ * waiting on [[terminalLatch]] when the base settles a terminal. Invoked once,
+ * by the caller that wins [[completeTerminal]].
+ */
+ override protected def onTerminalSettled(termination: Termination): Unit = {
+ outputQueue.put(QueueItem.EndOfStream)
+ terminalLatch.countDown()
+ }
+
+ private def handleControl(ctrl: UdfControlResponse): Unit = ctrl.getControlCase match {
+ case UdfControlResponse.ControlCase.INIT =>
+ val resp = ctrl.getInit
+ if (resp.hasError) {
+ // Attribute a subsequent close()'s terminator to the init error rather
+ // than a bare Cancelled, mirroring the data-phase ERROR branch.
+ executionError.compareAndSet(None, Some(resp.getError))
+ // Settle Failed(err) before the Cancel below (not after its CancelResponse):
+ // the sticky terminal makes close() report the init error, and the Cancel
+ // becomes best-effort worker cleanup -- so if it can't be written (a
+ // directExecutor worker delivered this reentrantly, before doInit published
+ // requestObserver), doInit still doesn't wait terminalTimeoutMs for a
+ // CancelResponse that can never arrive. Settle-before-release (see initValue).
+ Transitions.failed(resp.getError)
+ initValue.complete(resp)
+ sendCancelInternal(() => cancelWithReason("init failed"))
+ } else {
+ // InitResponse OK. Only advance from Initializing so a terminal that raced
+ // in (e.g. a transport error) still wins; process() then opens the data
+ // phase. Settle-before-release (see initValue): publish after the CAS.
+ Transitions.initAccepted()
+ initValue.complete(resp)
+ }
+
+ case UdfControlResponse.ControlCase.ERROR =>
+ val err = ctrl.getError.getError
+ executionError.compareAndSet(None, Some(err))
+ // Before init resolves, settle Failed here rather than rely on the Cancel
+ // below: sendCancelInternal may find requestObserver still null (reentrant
+ // delivery inside stream.onNext, before doInit published it), so no Cancel --
+ // hence no CancelResponse -- would settle the terminal, and doInit would
+ // return as if init succeeded. completeTerminal is idempotent, so the
+ // data-phase ERROR -> Cancel -> CancelResponse path is unaffected.
+ if (!initResolved) {
+ Transitions.failed(err)
+ }
+ // Settle-before-release (see initValue). No value to publish -- ERROR is not
+ // an InitResponse; on the data-phase path the latch already fired in init.
+ initValue.signalWithoutValue()
+ sendCancelInternal(() => cancelWithReason("aborting after ErrorResponse"))
+
+ case UdfControlResponse.ControlCase.FINISH =>
+ // The FinishResponse carries metrics + the finish-callback data/error.
+ // Keep it on the terminal so close() can return it; the iterator inspects
+ // its error field to decide whether to throw.
+ Transitions.finished(ctrl.getFinish)
+ // Defensive: FINISH before InitResponse is a worker protocol bug, but
+ // we should fail init fast rather than hang the 30s init timeout.
+ initValue.signalWithoutValue()
+
+ case UdfControlResponse.ControlCase.CANCEL =>
+ // The CancelResponse carries metrics + the cancel-callback error. Keep it
+ // on the terminal so close() can return it; any prior ErrorResponse is
+ // tracked in executionError and surfaced by the iterator.
+ Transitions.cancelled(ctrl.getCancel)
+ // Defensive: CANCEL before InitResponse unblocks doInit so it can
+ // surface the cancellation instead of timing out.
+ initValue.signalWithoutValue()
+
+ case UdfControlResponse.ControlCase.CONTROL_NOT_SET =>
+ Transitions.transportFailed(new IllegalStateException(
+ "empty UdfControlResponse oneof"))
+ initValue.signalWithoutValue()
+ }
+
+ /**
+ * True once init is no longer pending -- i.e. the stream is past `Initializing`.
+ * Not "init succeeded": a terminal (including a failure) also counts as resolved.
+ */
+ private def initResolved: Boolean = currentState match {
+ case SessionState.Created | SessionState.Initializing => false
+ case _ => true
+ }
+
+ private def cancelWithReason(reason: String): Cancel =
+ Cancel.newBuilder().setReason(reason).build()
+
+ /**
+ * The protocol transition graph in one place: names for the edges, not a
+ * second source of truth. Every edge acts on the single
+ * [[WorkerSession.SessionState]] owned by the base via [[compareAndSetState]]
+ * (non-terminal) or [[completeTerminal]] (terminal), so the base's CAS is the
+ * only synchronization and a terminal that arrived first always wins (the
+ * non-terminal CASes fail against it; [[completeTerminal]] is first-wins).
+ *
+ * Edges this class drives -- edge, method, then driver site(s) / thread (the
+ * base drives the API-call edges: `Created -> Initializing` in `init`,
+ * `Initialized -> Streaming` in `process`):
+ * {{{
+ * Initializing -> Initialized initAccepted handleControl INIT-ok [gRPC cb]
+ * Streaming -> Finishing finishSent advance branch 3 [engine]
+ * non-terminal -> Cancelling cancelSentFrom sendCancelInternal, iff a Cancel
+ * reaches the wire [gRPC cb | engine | init | close]
+ * non-terminal -> Terminal finished/cancelled/failed/transportFailed
+ * handleControl / doInit / doClose / advance / onError
+ * }}}
+ * `Cancelling` is reached only if a Cancel is actually written; a pre-stream or
+ * raced cancel goes straight to a `Cancelled`/`TransportFailed` terminal (see
+ * [[sendCancelInternal]], [[doClose]], `ProcessIterator`) or nowhere -- so
+ * `cancelRequested` (intent) does not imply state `Cancelling`.
+ */
+ private object Transitions {
+ /** `InitResponse` OK: `Initializing -> Initialized`. */
+ def initAccepted(): Boolean =
+ compareAndSetState(SessionState.Initializing, SessionState.Initialized)
+
+ /** Input exhausted and `Finish` written: `Streaming -> Finishing` (once). */
+ def finishSent(): Boolean =
+ compareAndSetState(SessionState.Streaming, SessionState.Finishing)
+
+ /** `Cancel` written: `cur -> Cancelling`, from any non-terminal `cur`. */
+ def cancelSentFrom(cur: SessionState): Boolean =
+ !cur.isTerminal && compareAndSetState(cur, SessionState.Cancelling)
+
+ /** Clean terminal carrying the worker's `FinishResponse`. */
+ def finished(response: FinishResponse): Boolean =
+ completeTerminal(Termination.Finished(response))
+
+ /** Clean terminal carrying the worker's `CancelResponse`. */
+ def cancelled(response: CancelResponse): Boolean =
+ completeTerminal(Termination.Cancelled(response))
+
+ /** Failure terminal carrying a structured [[ExecutionError]]. */
+ def failed(error: ExecutionError): Boolean =
+ completeTerminal(Termination.Failed(error))
+
+ /** Failure terminal carrying a transport-level cause. */
+ def transportFailed(cause: Throwable): Boolean =
+ completeTerminal(Termination.TransportFailed(cause))
Review Comment:
can this be merged to `failed`?
##########
udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/GrpcWorkerSession.scala:
##########
@@ -0,0 +1,933 @@
+/*
+ * 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.grpc
+
+import java.util.concurrent.{CountDownLatch, LinkedBlockingQueue, TimeoutException, TimeUnit}
+import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference}
+
+import scala.util.control.NonFatal
+
+import io.grpc.{ConnectivityState, ManagedChannel}
+import io.grpc.stub.StreamObserver
+
+import org.apache.spark.annotation.Experimental
+import org.apache.spark.udf.worker.{Cancel, CancelResponse, DataRequest, DataResponse,
+ ExecutionError, Finish, FinishResponse, Init, InitResponse, UdfControlRequest,
+ UdfControlResponse, UdfRequest, UdfResponse, UdfWorkerGrpc}
+import org.apache.spark.udf.worker.core.{Termination, WorkerHandle, WorkerLogger, WorkerSession}
+import org.apache.spark.udf.worker.core.WorkerSession.SessionState
+import org.apache.spark.udf.worker.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_message.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.
+ *
+ * '''Driving model.''' Consumption-driven (Volcano / pull): the thread that
+ * consumes the [[doProcess]] result iterator is the one that pulls input and
+ * sends each `DataRequest`; the gRPC callback thread only receives output. It is
+ * pull-driven but not one-input-per-output -- `advance` sends the next input
+ * whenever the output queue is momentarily empty, so under async delivery it may
+ * push several input batches before any output is read (bounded by HTTP/2 flow
+ * control).
+ *
+ * '''State machine.''' This class does not keep its own state machine: it
+ * drives the single [[WorkerSession.SessionState]] owned by the base. The base
+ * advances `Created -> Initializing` (in `init`) and `Initialized -> Streaming`
+ * (in `process`); this class advances the protocol-event edges through
+ * [[compareAndSetState]] / [[completeTerminal]] as it exchanges messages:
+ * {{{
+ * Initializing --(InitResponse ok)--> Initialized [handleControl]
+ * Streaming ----(Finish written)----> Finishing [ProcessIterator]
+ * <any non-terminal> --(Cancel written)--> Cancelling [sendCancelInternal]
+ * <any non-terminal> --(terminator/error)--> terminal [completeTerminal]
+ * }}}
+ * The two clean terminals carry the worker's `FinishResponse` / `CancelResponse`
+ * (metrics + finish/cancel callback `data`/`error`) so [[close]] can return
+ * them. The only flag kept outside the machine is [[cancelRequested]]: a
+ * cancellation can be requested before the stream exists (so it cannot be a
+ * state transition yet), and it must both fast-fail the result iterator and
+ * suppress any in-flight Data/Finish.
+ *
+ * Threading:
+ * - [[doInit]] is synchronous: sends `Init` and blocks on `InitResponse`,
+ * returning it.
+ * - [[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.
+ * - [[doClose]] is thread-safe and idempotent: it settles + returns the
+ * terminator (cancelling in-flight work if the stream had not finished)
+ * and half-closes the request side.
+ *
+ * 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 a gRPC channel built and owned by the caller (the
+ * dispatcher). 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 from the worker, drained by the process() iterator.
+ // Intentionally unbounded: a bounded queue would block the gRPC callback
+ // (Netty event-loop) thread when full, stalling terminator/control delivery on
+ // the whole channel. HTTP/2 flow control bounds the wire and the consumer
+ // normally drains promptly; a stalled downstream can still grow it, but the fix
+ // is protocol-level back-pressure (out of scope), not bounding the queue.
+ //
+ // TODO [SPARK-57324]: expose queue depth as a metric (early warning for a
+ // stalled consumer).
+ private val outputQueue = new LinkedBlockingQueue[QueueItem]()
+
+ // The worker's `InitResponse` (success or error) coupled with the latch that
+ // init() blocks on. Fires when the InitResponse arrives (`complete`) or when a
+ // terminal settles first without one -- transport error, half-close, or a
+ // premature terminator (`signalWithoutValue`). Until it fires we have no proof
+ // the worker accepted the session.
+ //
+ // Settle-before-release rule (referenced from every callback that both settles
+ // a terminal and fires this latch): settle the terminal FIRST, then complete /
+ // signal. init() blocks on the latch, and only latch await/release establish a
+ // happens-before edge, so a woken init() is guaranteed to observe the terminal
+ // rather than a transient state. OneShotValue keeps that publish-then-release
+ // in one place instead of every caller remembering to count down after setting
+ // the reference.
+ private val initValue = new OneShotValue[InitResponse]
+
+ // Fired when the session reaches a terminal [[SessionState]]. doClose() and
+ // the init-error path block on this to drain the terminator.
+ private val terminalLatch = new CountDownLatch(1)
+
+ // Captures an ErrorResponse encountered during the data phase so that
+ // the CancelResponse terminator 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)
+
+ // Cancellation INTENT -- distinct from the `Cancelling` state, which is reached
+ // only once a Cancel is actually written to the wire ([[sendCancelInternal]]).
+ // Intent is set first and can outrun (or never reach) that write, so
+ // `cancelRequested` does NOT imply state `Cancelling`. Kept outside the
+ // [[SessionState]] machine because a cancel can be requested before the stream
+ // exists (pre-init), where there is no wire transition to make yet. Used to (a)
+ // make cancel idempotent across all call sites, (b) fast-fail ProcessIterator
+ // on a pre-init cancel, and (c) suppress any Data/Finish that would otherwise
+ // race a Cancel onto the wire (re-read inside [[requestLock]]).
+ private val cancelRequested = new AtomicBoolean(false)
+
+ // gRPC requires serialized writes to a request StreamObserver.
+ private val requestLock = new Object
+
+ // Initialised in init() -- before that, close() is a no-op 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 =>
+ // A DataResponse before InitResponse violates the protocol (InitResponse
+ // must precede any DataResponse): fast-fail rather than enqueue it and let
+ // init() block to its timeout. Settle-before-release (see initValue).
+ if (!initResolved) {
+ Transitions.transportFailed(new IllegalStateException(
+ "worker sent a DataResponse before InitResponse"))
+ initValue.signalWithoutValue()
+ } else {
+ outputQueue.put(QueueItem.Batch(response.getData))
+ }
+
+ case UdfResponse.ResponseCase.CONTROL =>
+ handleControl(response.getControl)
+
+ case other =>
+ // A malformed response (empty / unknown oneof) is a terminal transport
+ // failure; fast-fail init the same way. Settle-before-release (see initValue).
+ Transitions.transportFailed(new IllegalStateException(
+ s"unexpected response oneof: $other"))
+ initValue.signalWithoutValue()
+ }
+ }
+
+ override def onError(t: Throwable): Unit = {
+ // Transport-level failure: the stream is dead, no further writes possible.
+ // Settle-before-release (see initValue) so init() surfaces the transport
+ // cause instead of the initResponseTimeoutMs "timed out" error.
+ Transitions.transportFailed(t)
+ initValue.signalWithoutValue()
+ }
+
+ 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 (!currentState.isTerminal) {
+ Transitions.transportFailed(new IllegalStateException(
+ "worker response stream closed without a terminator"))
+ }
+ // Defensive: if onCompleted reached us before InitResponse, doInit is
+ // still blocked on initValue and would otherwise time out.
+ initValue.signalWithoutValue()
+ }
+ }
+
+ /**
+ * Wakes the result iterator (blocked on [[outputQueue]]) and any thread
+ * waiting on [[terminalLatch]] when the base settles a terminal. Invoked once,
+ * by the caller that wins [[completeTerminal]].
+ */
+ override protected def onTerminalSettled(termination: Termination): Unit = {
+ outputQueue.put(QueueItem.EndOfStream)
+ terminalLatch.countDown()
+ }
+
+ private def handleControl(ctrl: UdfControlResponse): Unit = ctrl.getControlCase match {
+ case UdfControlResponse.ControlCase.INIT =>
+ val resp = ctrl.getInit
+ if (resp.hasError) {
+ // Attribute a subsequent close()'s terminator to the init error rather
+ // than a bare Cancelled, mirroring the data-phase ERROR branch.
+ executionError.compareAndSet(None, Some(resp.getError))
+ // Settle Failed(err) before the Cancel below (not after its CancelResponse):
+ // the sticky terminal makes close() report the init error, and the Cancel
+ // becomes best-effort worker cleanup -- so if it can't be written (a
+ // directExecutor worker delivered this reentrantly, before doInit published
+ // requestObserver), doInit still doesn't wait terminalTimeoutMs for a
+ // CancelResponse that can never arrive. Settle-before-release (see initValue).
+ Transitions.failed(resp.getError)
+ initValue.complete(resp)
+ sendCancelInternal(() => cancelWithReason("init failed"))
+ } else {
+ // InitResponse OK. Only advance from Initializing so a terminal that raced
+ // in (e.g. a transport error) still wins; process() then opens the data
+ // phase. Settle-before-release (see initValue): publish after the CAS.
+ Transitions.initAccepted()
+ initValue.complete(resp)
+ }
+
+ case UdfControlResponse.ControlCase.ERROR =>
+ val err = ctrl.getError.getError
+ executionError.compareAndSet(None, Some(err))
+ // Before init resolves, settle Failed here rather than rely on the Cancel
+ // below: sendCancelInternal may find requestObserver still null (reentrant
+ // delivery inside stream.onNext, before doInit published it), so no Cancel --
+ // hence no CancelResponse -- would settle the terminal, and doInit would
+ // return as if init succeeded. completeTerminal is idempotent, so the
+ // data-phase ERROR -> Cancel -> CancelResponse path is unaffected.
+ if (!initResolved) {
+ Transitions.failed(err)
+ }
+ // Settle-before-release (see initValue). No value to publish -- ERROR is not
+ // an InitResponse; on the data-phase path the latch already fired in init.
+ initValue.signalWithoutValue()
+ sendCancelInternal(() => cancelWithReason("aborting after ErrorResponse"))
+
+ case UdfControlResponse.ControlCase.FINISH =>
+ // The FinishResponse carries metrics + the finish-callback data/error.
+ // Keep it on the terminal so close() can return it; the iterator inspects
+ // its error field to decide whether to throw.
+ Transitions.finished(ctrl.getFinish)
+ // Defensive: FINISH before InitResponse is a worker protocol bug, but
+ // we should fail init fast rather than hang the 30s init timeout.
+ initValue.signalWithoutValue()
+
+ case UdfControlResponse.ControlCase.CANCEL =>
+ // The CancelResponse carries metrics + the cancel-callback error. Keep it
+ // on the terminal so close() can return it; any prior ErrorResponse is
+ // tracked in executionError and surfaced by the iterator.
+ Transitions.cancelled(ctrl.getCancel)
+ // Defensive: CANCEL before InitResponse unblocks doInit so it can
+ // surface the cancellation instead of timing out.
+ initValue.signalWithoutValue()
+
+ case UdfControlResponse.ControlCase.CONTROL_NOT_SET =>
+ Transitions.transportFailed(new IllegalStateException(
+ "empty UdfControlResponse oneof"))
+ initValue.signalWithoutValue()
+ }
+
+ /**
+ * True once init is no longer pending -- i.e. the stream is past `Initializing`.
+ * Not "init succeeded": a terminal (including a failure) also counts as resolved.
+ */
+ private def initResolved: Boolean = currentState match {
+ case SessionState.Created | SessionState.Initializing => false
+ case _ => true
+ }
+
+ private def cancelWithReason(reason: String): Cancel =
+ Cancel.newBuilder().setReason(reason).build()
+
+ /**
+ * The protocol transition graph in one place: names for the edges, not a
+ * second source of truth. Every edge acts on the single
+ * [[WorkerSession.SessionState]] owned by the base via [[compareAndSetState]]
+ * (non-terminal) or [[completeTerminal]] (terminal), so the base's CAS is the
+ * only synchronization and a terminal that arrived first always wins (the
+ * non-terminal CASes fail against it; [[completeTerminal]] is first-wins).
+ *
+ * Edges this class drives -- edge, method, then driver site(s) / thread (the
+ * base drives the API-call edges: `Created -> Initializing` in `init`,
+ * `Initialized -> Streaming` in `process`):
+ * {{{
+ * Initializing -> Initialized initAccepted handleControl INIT-ok [gRPC cb]
+ * Streaming -> Finishing finishSent advance branch 3 [engine]
+ * non-terminal -> Cancelling cancelSentFrom sendCancelInternal, iff a Cancel
+ * reaches the wire [gRPC cb | engine | init | close]
+ * non-terminal -> Terminal finished/cancelled/failed/transportFailed
+ * handleControl / doInit / doClose / advance / onError
+ * }}}
+ * `Cancelling` is reached only if a Cancel is actually written; a pre-stream or
+ * raced cancel goes straight to a `Cancelled`/`TransportFailed` terminal (see
+ * [[sendCancelInternal]], [[doClose]], `ProcessIterator`) or nowhere -- so
+ * `cancelRequested` (intent) does not imply state `Cancelling`.
+ */
+ private object Transitions {
+ /** `InitResponse` OK: `Initializing -> Initialized`. */
+ def initAccepted(): Boolean =
+ compareAndSetState(SessionState.Initializing, SessionState.Initialized)
+
+ /** Input exhausted and `Finish` written: `Streaming -> Finishing` (once). */
+ def finishSent(): Boolean =
+ compareAndSetState(SessionState.Streaming, SessionState.Finishing)
+
+ /** `Cancel` written: `cur -> Cancelling`, from any non-terminal `cur`. */
+ def cancelSentFrom(cur: SessionState): Boolean =
+ !cur.isTerminal && compareAndSetState(cur, SessionState.Cancelling)
+
+ /** Clean terminal carrying the worker's `FinishResponse`. */
+ def finished(response: FinishResponse): Boolean =
+ completeTerminal(Termination.Finished(response))
+
+ /** Clean terminal carrying the worker's `CancelResponse`. */
+ def cancelled(response: CancelResponse): Boolean =
+ completeTerminal(Termination.Cancelled(response))
+
+ /** Failure terminal carrying a structured [[ExecutionError]]. */
+ def failed(error: ExecutionError): Boolean =
+ completeTerminal(Termination.Failed(error))
+
+ /** Failure terminal carrying a transport-level cause. */
+ def transportFailed(cause: Throwable): Boolean =
+ completeTerminal(Termination.TransportFailed(cause))
+ }
+
+ // ---- WorkerSession hooks ------------------------------------------------
+
+ override protected def doInit(message: Init): InitResponse = {
+ // 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")
+ Transitions.transportFailed(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. close()'s 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 {
+ // Reentrant delivery: a directExecutor worker (one whose gRPC callbacks run
+ // synchronously on the caller's thread) can deliver InitResponse -- and thus
+ // run responseObserver/handleControl -- from *inside* this stream.onNext,
+ // before requestObserver below is published. The base advanced
+ // Created -> Initializing before calling doInit, so that reentrant
+ // InitResponse still finds Initializing and advances to Initialized; the
+ // handleControl error paths guard the still-null requestObserver. This is
+ // the scenario the "before doInit published requestObserver" comments mean.
+ 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
+ Transitions.transportFailed(e)
+ // Surface as GrpcWorkerSessionException so the engine integration layer
+ // (which catches that type and wraps it) sees a uniform init-failure
+ // exception rather than the raw transport error.
+ throw new GrpcWorkerSessionException("UDF worker stream failed during init", e)
+ }
+
+ val responded = try {
+ initValue.await(initResponseTimeoutMs)
+ } catch {
+ case _: InterruptedException =>
+ Thread.currentThread().interrupt()
+ sendCancelInternal(() => cancelWithReason("interrupted during init"))
+ // Make sure close() does not block waiting for a terminator that
+ // will never arrive on this thread's behalf.
+ Transitions.transportFailed(
+ new InterruptedException("interrupted while waiting for InitResponse"))
+ throw new InterruptedException("interrupted while waiting for InitResponse")
+ }
+
+ if (!responded) {
+ sendCancelInternal(() => cancelWithReason("InitResponse timed out"))
+ val timeout = new TimeoutException(
+ s"timed out waiting for InitResponse after ${initResponseTimeoutMs}ms")
+ // Settle the terminal so a subsequent close() does not stall for a
+ // second `terminalTimeoutMs` waiting for a worker that already missed
+ // its init deadline.
+ Transitions.transportFailed(timeout)
+ // Surface as GrpcWorkerSessionException (carrying the timeout cause) so
+ // the engine integration layer that catches that type can wrap it.
+ throw new GrpcWorkerSessionException(
+ s"timed out waiting for InitResponse after ${initResponseTimeoutMs}ms", timeout)
+ }
+
+ initValue.get match {
+ case Some(resp) if resp.hasError =>
+ // 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(resp.getError)}", resp.getError)
+ case Some(resp) =>
+ resp
+ case None =>
+ // No InitResponse arrived but the latch fired. The defensive
+ // initValue.signalWithoutValue() 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.
+ currentState match {
+ case SessionState.Terminal(Termination.TransportFailed(cause)) =>
+ throw new GrpcWorkerSessionException(
+ "UDF worker stream failed during init", cause)
+ case SessionState.Terminal(Termination.Failed(err)) =>
+ throw new GrpcWorkerSessionException(
+ s"UDF worker reported an error before init completed: " +
+ describeError(err), err)
+ case SessionState.Terminal(Termination.Cancelled(_)) =>
+ throw new GrpcWorkerSessionException(
+ "UDF worker stream was cancelled before init completed")
+ case SessionState.Terminal(Termination.Finished(_)) =>
+ throw new GrpcWorkerSessionException(
+ "UDF worker finished before init completed")
+ case other =>
+ throw new IllegalStateException(
+ s"init latch fired without an InitResponse or terminal: $other")
+ }
+ }
+ }
+
+ override protected def doProcess(
+ input: Iterator[DataRequest],
+ finish: () => Finish): Iterator[DataResponse] = {
+ // Init success is guaranteed by the base [[WorkerSession]] lifecycle: if
+ // doInit had failed it would have thrown and process() would never run.
+ new ProcessIterator(input, finish)
+ }
+
+ override protected def doClose(cancel: () => Cancel): Termination = {
+ if (requestObserver == null) {
+ // init() never put a stream on the wire (closed before/around init, or
+ // init threw before publishing). There is no protocol terminator; treat
+ // the session as cancelled-before-start. The base WorkerSession still
+ // releases the worker handle, so the worker is torn down.
+ Transitions.cancelled(CancelResponse.getDefaultInstance)
+ return Termination.Cancelled(CancelResponse.getDefaultInstance)
+ }
+ // If the stream has not finished on its own, cancel anything in flight so
+ // the worker can clean up, then wait for the terminator. sendCancelInternal
+ // evaluates the cancel thunk only when it actually writes a Cancel, and at
+ // most once across all callers.
+ if (!currentState.isTerminal) {
+ sendCancelInternal(cancel)
+ 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 the state sees a stable value.
+ if (!currentState.isTerminal) {
+ Transitions.transportFailed(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. onCompleted on an
+ // already-torn-down stream throws and is caught, so this stays idempotent
+ // across repeat close() calls.
+ try {
+ requestLock.synchronized { requestObserver.onCompleted() }
+ } catch {
+ case NonFatal(e) =>
+ logger.debug("Error half-closing UDF Execute stream", e)
+ }
+ // Derive the Termination from the settled terminal. close() is the
+ // finalizer/cleanup path and must NOT re-throw: a UDF / data-phase error is
+ // already surfaced while the result iterator is consumed, and an init error
+ // is surfaced from init(). settledTermination returns the settled terminal
+ // as-is: the response proto for the clean terminators, and the failure
+ // terminals (Failed / TransportFailed) carrying their cause -- not a bare
+ // Cancelled.
+ settledTermination
+ }
+
+ // ---- Internal request helpers ---------------------------------------------
+
+ /**
+ * Sends a Data or Finish request to the worker. Two invariants are
+ * checked inside [[requestLock]]:
+ * - terminal state: 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 (currentState.isTerminal) {
+ throw new IllegalStateException(
+ "cannot send request: UDF Execute stream is already closed")
+ }
+ // Suppress the write when a cancel has been (or is about to be) flushed
+ // through this lock; that preserves the proto ordering invariant (no
+ // Data/Finish after Cancel). Test the flag rather than `return`-ing from
+ // this by-name `synchronized` body: a non-local return compiles to a thrown
+ // NonLocalReturnControl, which is brittle (also relied on in sendCancelInternal).
+ if (!cancelRequested.get()) {
+ requestObserver.onNext(req)
+ }
+ }
+
+ /**
+ * Sends a `Cancel` control message, returning `true` iff a `Cancel` was
+ * actually written to the wire. Idempotent across ALL call sites (close()'s
+ * cancel, in-band cancels from `handleControl`'s INIT error / ERROR paths,
+ * and engine-internal cancels from the iterator) via [[cancelRequested]]:
+ * only the first caller puts a `Cancel` on the wire, and the `cancel` thunk is
+ * evaluated only by that caller and only when the message is actually sent --
+ * so a side-effecting thunk (e.g. one carrying a client cancel callback) runs
+ * at most once. 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(cancel: () => Cancel): Boolean = {
+ if (!cancelRequested.compareAndSet(false, true)) return false
+ if (requestObserver == null) return false
+ if (currentState.isTerminal) return false
+ try {
+ requestLock.synchronized {
+ // A terminal that arrived first must win; bail without writing. Compute
+ // the block's value rather than `return` (see sendRequest on why non-local
+ // return is avoided -- here it would also escape the NonFatal catch below).
+ val cur = currentState
+ if (cur.isTerminal) {
+ false
+ } else {
+ // Record that a Cancel is on the wire by advancing to Cancelling.
+ Transitions.cancelSentFrom(cur)
+ requestObserver.onNext(UdfRequest.newBuilder()
+ .setControl(UdfControlRequest.newBuilder().setCancel(cancel()).build())
+ .build())
+ true
+ }
+ }
+ } catch {
+ case NonFatal(e) =>
+ logger.debug(s"Cancel send failed (stream may already be torn down): ${e.getMessage}")
+ Transitions.transportFailed(e)
+ false
+ }
+ }
+
+ private def awaitTerminal(): Unit = {
+ if (currentState.isTerminal) return
+ try {
+ if (!terminalLatch.await(terminalTimeoutMs, TimeUnit.MILLISECONDS)) {
+ Transitions.transportFailed(new TimeoutException(
+ s"timed out waiting for stream terminator after ${terminalTimeoutMs}ms"))
+ }
+ } catch {
+ case _: InterruptedException =>
+ // Attribute the terminal to the interrupt, not a timeout: the wait was
+ // cut short, it did not exhaust terminalTimeoutMs. Reporting "timed
+ // out" here would misdiagnose an interrupted close as a slow worker.
+ Thread.currentThread().interrupt()
+ Transitions.transportFailed(
+ new InterruptedException("interrupted while waiting for stream terminator"))
+ }
+ }
+
+ // ---- 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 settle the state) and with any thread that finalizes via
+ * [[close]]. It drives the `Streaming -> Finishing` transition (sending
+ * `Finish` once input is exhausted, built lazily from the `finish` thunk).
+ */
+ private class ProcessIterator(input: Iterator[DataRequest], finish: () => Finish)
+ extends Iterator[DataResponse] {
+
+ // 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 branches, and block branch 4 for the
+ // full terminalTimeoutMs before returning. Callers that probe hasNext
+ // an extra time (iterator.size, instrumentation wrappers) would hang.
+ // Iterator-local and only touched by the single engine thread.
+ private val exhausted = new AtomicBoolean(false)
+ @volatile private var prefetched: DataResponse = _
+
+ // Pre-init cancel fast-fail: if a 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() && !currentState.isTerminal) {
+ Transitions.cancelled(CancelResponse.getDefaultInstance)
+ }
+
+ /**
+ * Runs a caller-supplied thunk (`input.hasNext`, `input.next()`, or the
+ * `finish` builder) that may throw, and on failure Cancels the in-flight
+ * stream before rethrowing. Without the Cancel a failure would leave the
+ * worker awaiting more input with no terminator owed on the wire; with it the
+ * worker tears down and the exception still propagates to the engine. Setting
+ * `cancelRequested` also suppresses any further Data/Finish this loop might
+ * otherwise attempt (see [[sendRequest]]).
+ */
+ private def cancelOnThrow[T](reason: String)(op: => T): T =
+ try op catch {
+ case NonFatal(e) =>
+ sendCancelInternal(() => cancelWithReason(reason))
+ throw e
+ }
+
+ override def hasNext: Boolean = {
+ if (prefetched ne null) return true
+ advance()
+ prefetched ne null
+ }
+
+ override def next(): DataResponse = {
+ if (prefetched eq null) advance()
+ val out = prefetched
+ if (out eq null) {
+ throw new NoSuchElementException("ProcessIterator exhausted")
+ }
+ prefetched = null
+ out
+ }
+
+ /**
+ * Fills [[prefetched]] with the next output batch, or leaves it null at the
+ * terminator (which also sets [[exhausted]], so `hasNext` reads null as "done").
+ * Each loop iteration tries in order: (1) drain queued output -- a batch (fill
+ * and return) or the terminator sentinel (return); (2) while `Streaming`, send
+ * the next input batch and loop; (3) once input is exhausted, send `Finish`
+ * (CAS `Streaming -> Finishing`, once) and loop; (4) otherwise block for late
+ * output or the terminator, then return. Branches 2-3 loop; 1 and 4 return.
+ *
+ * '''Per-poll timeout.''' Branch 4 waits up to [[terminalTimeoutMs]] per
+ * poll, reset by every worker event -- the contract is "emit at least one
+ * event every [[terminalTimeoutMs]] after Finish", not "finish the UDF within
+ * [[terminalTimeoutMs]]". A worker expecting a long post-Finish silence MAY
+ * emit an empty `DataResponse` as a heartbeat to reset the wait; it is
+ * surfaced to the caller, so it should be a batch the caller recognises as
+ * empty (e.g. a zero-row Arrow batch).
+ */
+ private def advance(): Unit = {
+ if (exhausted.get()) return // terminator already seen; never re-block
+ 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 while the stream is open for data.
+ // `input.hasNext` may itself fetch/compute the next element (many Spark
+ // iterators prefetch), so it can throw just like `input.next()`; both go
+ // through cancelOnThrow so a failing upstream Cancels the stream instead
+ // of stranding the worker waiting for input that will never arrive.
+ if (!cancelRequested.get() && currentState == SessionState.Streaming &&
+ cancelOnThrow("input iterator failed")(input.hasNext)) {
+ val batch = cancelOnThrow("input iterator failed")(input.next())
+ if (!sendOrEndOnRacedTerminal(UdfRequest.newBuilder().setData(batch).build())) return
+ } else if (!cancelRequested.get() && Transitions.finishSent()) {
+ // (3) No more input; send Finish exactly once (unless cancelled). The
+ // `finish` thunk is caller-supplied (it may run a finish callback), so
+ // it can throw; cancelOnThrow Cancels the stream before rethrowing. The
+ // Streaming -> Finishing CAS has already run, but we have not written
+ // Finish yet, so the Cancel is the only engine-to-worker message that
+ // reaches the wire -- the proto "nothing after Cancel" invariant holds.
+ val finishMsg = cancelOnThrow("finish callback failed")(finish())
+ if (!sendOrEndOnRacedTerminal(UdfRequest.newBuilder()
+ .setControl(UdfControlRequest.newBuilder().setFinish(finishMsg).build())
+ .build())) {
+ return
+ }
+ } else {
+ // (4) Block for late output or the terminator. See class doc
+ // above for the per-poll-vs-total-session timeout semantics.
+ val item = try {
+ outputQueue.poll(terminalTimeoutMs, TimeUnit.MILLISECONDS)
+ } catch {
+ case _: InterruptedException =>
+ Thread.currentThread().interrupt()
+ sendCancelInternal(() => cancelWithReason("iterator interrupted"))
+ exhausted.set(true)
+ throw new InterruptedException("interrupted while reading UDF result")
+ }
+ item match {
+ case null =>
+ Transitions.transportFailed(new IllegalStateException(
+ s"timed out waiting for UDF output after ${terminalTimeoutMs}ms"))
+ exhausted.set(true)
+ throwIfTerminalError()
+ return
+ case QueueItem.EndOfStream =>
+ exhausted.set(true)
+ throwIfTerminalError()
+ return
+ case QueueItem.Batch(b) =>
+ prefetched = b
+ return
+ }
+ }
+ }
+ }
+
+ /**
+ * Sends one data-phase request (`DataRequest` / `Finish`), recovering from a
+ * terminator that raced the write. Returns `true` to keep looping (sent, or
+ * suppressed by a pending cancel), `false` if a terminator settled and the
+ * iterator should end (caller `return`s).
+ *
+ * On a write failure a terminator usually already settled (the worker
+ * finished/failed early) and the write only failed because the stream is
+ * closed: record a transport terminal if none is set ([[completeTerminal]] is
+ * a no-op once settled), then [[throwIfTerminalError]] -- which throws for an
+ * error/transport terminal, or returns for a clean `Finished` that raced the
+ * send, in which case the benign "stream closed" error is dropped.
+ */
+ private def sendOrEndOnRacedTerminal(req: UdfRequest): Boolean = {
+ try {
+ sendRequest(req)
+ true
+ } catch {
+ case NonFatal(e) =>
+ if (!currentState.isTerminal) {
+ Transitions.transportFailed(e)
+ }
+ exhausted.set(true)
+ throwIfTerminalError()
+ false
+ }
+ }
+
+ /**
+ * Surfaces a failed terminator as an exception when the result iterator is
+ * drained. A data-phase [[ExecutionError]] (captured in [[executionError]])
+ * takes precedence over the terminator's own error, then the finish/cancel
+ * callback error carried on the terminator, then a bare cancellation.
+ */
+ private def throwIfTerminalError(): Unit = currentState match {
+ case SessionState.Terminal(Termination.Finished(response)) =>
+ responseError(response.hasError, response.getError, executionError.get())
+ .foreach(err => throw new GrpcWorkerSessionException(
+ s"UDF execution failed: ${describeError(err)}", err))
+ case SessionState.Terminal(Termination.Cancelled(response)) =>
+ responseError(response.hasError, response.getError, executionError.get()) match {
+ case Some(err) =>
+ throw new GrpcWorkerSessionException(
+ s"UDF execution failed: ${describeError(err)}", err)
+ case None =>
+ throw new GrpcWorkerSessionException("UDF execution was cancelled")
+ }
+ case SessionState.Terminal(Termination.Failed(err)) =>
+ throw new GrpcWorkerSessionException(
+ s"UDF execution failed: ${describeError(err)}", err)
+ case SessionState.Terminal(Termination.TransportFailed(t)) =>
+ throw new GrpcWorkerSessionException("UDF worker stream failed", t)
+ case other =>
+ throw new IllegalStateException(s"terminator sentinel without terminal: $other")
+ }
+
+ /** Picks the error to surface: prior data-phase error, else the terminator's. */
+ private def responseError(
+ hasError: Boolean,
+ error: ExecutionError,
+ priorError: Option[ExecutionError]): Option[ExecutionError] =
+ priorError.orElse(if (hasError) Some(error) else None)
+ }
+}
+
+object GrpcWorkerSession {
+ /** Upper bound on the wait for `InitResponse`. */
+ val DEFAULT_INIT_RESPONSE_TIMEOUT_MS: Long = 30000L
+
+ /** Upper bound on the wait for `FinishResponse` / `CancelResponse`. */
+ val DEFAULT_TERMINAL_TIMEOUT_MS: Long = 30000L
+
+ // Distinguishes a (possibly empty) data batch from end-of-stream, and makes
+ // the iterator's match exhaustive.
+ private sealed trait QueueItem
+ private object QueueItem {
+ final case class Batch(response: DataResponse) extends QueueItem
+ case object EndOfStream extends QueueItem
+ }
+
+ /**
+ * A write-once value paired with the latch a waiter blocks on, so the "publish
+ * the value, then release the waiter" ordering lives in one place instead of
+ * every call site having to remember to count the latch down after setting the
+ * reference. The value is set at most once ([[complete]]); the latch can also
+ * be released without a value ([[signalWithoutValue]]) when the waiter must be
+ * woken by a terminal that arrived instead of the value. All reads/writes go
+ * through the latch, so a waiter released by [[await]] has a happens-before
+ * edge to the [[complete]] that set the value.
+ */
+ private final class OneShotValue[A] {
+ private val latch = new CountDownLatch(1)
+ private val ref = new AtomicReference[Option[A]](None)
+
+ /** Publishes the value (first writer wins) and releases any waiter. */
+ def complete(value: A): Unit = {
+ ref.compareAndSet(None, Some(value))
+ latch.countDown()
+ }
+
+ /** Releases the waiter without a value (a terminal arrived, not the value). */
+ def signalWithoutValue(): Unit = latch.countDown()
+
+ /** Blocks up to `timeoutMs` for a release; true iff released in time. */
+ def await(timeoutMs: Long): Boolean = latch.await(timeoutMs, TimeUnit.MILLISECONDS)
Review Comment:
shall we make it raise timeout error so that it can be handled easily from the caller? (instead of checking a return value)
##########
udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/GrpcWorkerSession.scala:
##########
@@ -0,0 +1,933 @@
+/*
+ * 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.grpc
+
+import java.util.concurrent.{CountDownLatch, LinkedBlockingQueue, TimeoutException, TimeUnit}
+import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference}
+
+import scala.util.control.NonFatal
+
+import io.grpc.{ConnectivityState, ManagedChannel}
+import io.grpc.stub.StreamObserver
+
+import org.apache.spark.annotation.Experimental
+import org.apache.spark.udf.worker.{Cancel, CancelResponse, DataRequest, DataResponse,
+ ExecutionError, Finish, FinishResponse, Init, InitResponse, UdfControlRequest,
+ UdfControlResponse, UdfRequest, UdfResponse, UdfWorkerGrpc}
+import org.apache.spark.udf.worker.core.{Termination, WorkerHandle, WorkerLogger, WorkerSession}
+import org.apache.spark.udf.worker.core.WorkerSession.SessionState
+import org.apache.spark.udf.worker.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_message.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.
+ *
+ * '''Driving model.''' Consumption-driven (Volcano / pull): the thread that
+ * consumes the [[doProcess]] result iterator is the one that pulls input and
+ * sends each `DataRequest`; the gRPC callback thread only receives output. It is
+ * pull-driven but not one-input-per-output -- `advance` sends the next input
+ * whenever the output queue is momentarily empty, so under async delivery it may
+ * push several input batches before any output is read (bounded by HTTP/2 flow
+ * control).
+ *
+ * '''State machine.''' This class does not keep its own state machine: it
+ * drives the single [[WorkerSession.SessionState]] owned by the base. The base
+ * advances `Created -> Initializing` (in `init`) and `Initialized -> Streaming`
+ * (in `process`); this class advances the protocol-event edges through
+ * [[compareAndSetState]] / [[completeTerminal]] as it exchanges messages:
+ * {{{
+ * Initializing --(InitResponse ok)--> Initialized [handleControl]
+ * Streaming ----(Finish written)----> Finishing [ProcessIterator]
+ * <any non-terminal> --(Cancel written)--> Cancelling [sendCancelInternal]
+ * <any non-terminal> --(terminator/error)--> terminal [completeTerminal]
+ * }}}
+ * The two clean terminals carry the worker's `FinishResponse` / `CancelResponse`
+ * (metrics + finish/cancel callback `data`/`error`) so [[close]] can return
+ * them. The only flag kept outside the machine is [[cancelRequested]]: a
+ * cancellation can be requested before the stream exists (so it cannot be a
+ * state transition yet), and it must both fast-fail the result iterator and
+ * suppress any in-flight Data/Finish.
+ *
+ * Threading:
+ * - [[doInit]] is synchronous: sends `Init` and blocks on `InitResponse`,
+ * returning it.
+ * - [[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.
+ * - [[doClose]] is thread-safe and idempotent: it settles + returns the
+ * terminator (cancelling in-flight work if the stream had not finished)
+ * and half-closes the request side.
+ *
+ * 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 a gRPC channel built and owned by the caller (the
+ * dispatcher). 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 from the worker, drained by the process() iterator.
+ // Intentionally unbounded: a bounded queue would block the gRPC callback
+ // (Netty event-loop) thread when full, stalling terminator/control delivery on
+ // the whole channel. HTTP/2 flow control bounds the wire and the consumer
+ // normally drains promptly; a stalled downstream can still grow it, but the fix
+ // is protocol-level back-pressure (out of scope), not bounding the queue.
+ //
+ // TODO [SPARK-57324]: expose queue depth as a metric (early warning for a
+ // stalled consumer).
+ private val outputQueue = new LinkedBlockingQueue[QueueItem]()
+
+ // The worker's `InitResponse` (success or error) coupled with the latch that
+ // init() blocks on. Fires when the InitResponse arrives (`complete`) or when a
+ // terminal settles first without one -- transport error, half-close, or a
+ // premature terminator (`signalWithoutValue`). Until it fires we have no proof
+ // the worker accepted the session.
+ //
+ // Settle-before-release rule (referenced from every callback that both settles
+ // a terminal and fires this latch): settle the terminal FIRST, then complete /
+ // signal. init() blocks on the latch, and only latch await/release establish a
+ // happens-before edge, so a woken init() is guaranteed to observe the terminal
+ // rather than a transient state. OneShotValue keeps that publish-then-release
+ // in one place instead of every caller remembering to count down after setting
+ // the reference.
+ private val initValue = new OneShotValue[InitResponse]
+
+ // Fired when the session reaches a terminal [[SessionState]]. doClose() and
+ // the init-error path block on this to drain the terminator.
+ private val terminalLatch = new CountDownLatch(1)
+
+ // Captures an ErrorResponse encountered during the data phase so that
+ // the CancelResponse terminator 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)
+
+ // Cancellation INTENT -- distinct from the `Cancelling` state, which is reached
+ // only once a Cancel is actually written to the wire ([[sendCancelInternal]]).
+ // Intent is set first and can outrun (or never reach) that write, so
+ // `cancelRequested` does NOT imply state `Cancelling`. Kept outside the
+ // [[SessionState]] machine because a cancel can be requested before the stream
+ // exists (pre-init), where there is no wire transition to make yet. Used to (a)
+ // make cancel idempotent across all call sites, (b) fast-fail ProcessIterator
+ // on a pre-init cancel, and (c) suppress any Data/Finish that would otherwise
+ // race a Cancel onto the wire (re-read inside [[requestLock]]).
+ private val cancelRequested = new AtomicBoolean(false)
+
+ // gRPC requires serialized writes to a request StreamObserver.
+ private val requestLock = new Object
+
+ // Initialised in init() -- before that, close() is a no-op 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 =>
+ // A DataResponse before InitResponse violates the protocol (InitResponse
+ // must precede any DataResponse): fast-fail rather than enqueue it and let
+ // init() block to its timeout. Settle-before-release (see initValue).
+ if (!initResolved) {
+ Transitions.transportFailed(new IllegalStateException(
+ "worker sent a DataResponse before InitResponse"))
+ initValue.signalWithoutValue()
+ } else {
+ outputQueue.put(QueueItem.Batch(response.getData))
+ }
+
+ case UdfResponse.ResponseCase.CONTROL =>
+ handleControl(response.getControl)
+
+ case other =>
+ // A malformed response (empty / unknown oneof) is a terminal transport
+ // failure; fast-fail init the same way. Settle-before-release (see initValue).
+ Transitions.transportFailed(new IllegalStateException(
+ s"unexpected response oneof: $other"))
+ initValue.signalWithoutValue()
+ }
+ }
+
+ override def onError(t: Throwable): Unit = {
+ // Transport-level failure: the stream is dead, no further writes possible.
+ // Settle-before-release (see initValue) so init() surfaces the transport
+ // cause instead of the initResponseTimeoutMs "timed out" error.
+ Transitions.transportFailed(t)
+ initValue.signalWithoutValue()
+ }
+
+ 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 (!currentState.isTerminal) {
+ Transitions.transportFailed(new IllegalStateException(
+ "worker response stream closed without a terminator"))
+ }
+ // Defensive: if onCompleted reached us before InitResponse, doInit is
+ // still blocked on initValue and would otherwise time out.
+ initValue.signalWithoutValue()
+ }
+ }
+
+ /**
+ * Wakes the result iterator (blocked on [[outputQueue]]) and any thread
+ * waiting on [[terminalLatch]] when the base settles a terminal. Invoked once,
+ * by the caller that wins [[completeTerminal]].
+ */
+ override protected def onTerminalSettled(termination: Termination): Unit = {
+ outputQueue.put(QueueItem.EndOfStream)
+ terminalLatch.countDown()
+ }
+
+ private def handleControl(ctrl: UdfControlResponse): Unit = ctrl.getControlCase match {
+ case UdfControlResponse.ControlCase.INIT =>
+ val resp = ctrl.getInit
+ if (resp.hasError) {
+ // Attribute a subsequent close()'s terminator to the init error rather
+ // than a bare Cancelled, mirroring the data-phase ERROR branch.
+ executionError.compareAndSet(None, Some(resp.getError))
+ // Settle Failed(err) before the Cancel below (not after its CancelResponse):
+ // the sticky terminal makes close() report the init error, and the Cancel
+ // becomes best-effort worker cleanup -- so if it can't be written (a
+ // directExecutor worker delivered this reentrantly, before doInit published
+ // requestObserver), doInit still doesn't wait terminalTimeoutMs for a
+ // CancelResponse that can never arrive. Settle-before-release (see initValue).
+ Transitions.failed(resp.getError)
+ initValue.complete(resp)
+ sendCancelInternal(() => cancelWithReason("init failed"))
+ } else {
+ // InitResponse OK. Only advance from Initializing so a terminal that raced
+ // in (e.g. a transport error) still wins; process() then opens the data
+ // phase. Settle-before-release (see initValue): publish after the CAS.
+ Transitions.initAccepted()
+ initValue.complete(resp)
+ }
+
+ case UdfControlResponse.ControlCase.ERROR =>
+ val err = ctrl.getError.getError
+ executionError.compareAndSet(None, Some(err))
+ // Before init resolves, settle Failed here rather than rely on the Cancel
+ // below: sendCancelInternal may find requestObserver still null (reentrant
+ // delivery inside stream.onNext, before doInit published it), so no Cancel --
+ // hence no CancelResponse -- would settle the terminal, and doInit would
+ // return as if init succeeded. completeTerminal is idempotent, so the
+ // data-phase ERROR -> Cancel -> CancelResponse path is unaffected.
+ if (!initResolved) {
+ Transitions.failed(err)
+ }
+ // Settle-before-release (see initValue). No value to publish -- ERROR is not
+ // an InitResponse; on the data-phase path the latch already fired in init.
+ initValue.signalWithoutValue()
+ sendCancelInternal(() => cancelWithReason("aborting after ErrorResponse"))
+
+ case UdfControlResponse.ControlCase.FINISH =>
+ // The FinishResponse carries metrics + the finish-callback data/error.
+ // Keep it on the terminal so close() can return it; the iterator inspects
+ // its error field to decide whether to throw.
+ Transitions.finished(ctrl.getFinish)
+ // Defensive: FINISH before InitResponse is a worker protocol bug, but
+ // we should fail init fast rather than hang the 30s init timeout.
+ initValue.signalWithoutValue()
+
+ case UdfControlResponse.ControlCase.CANCEL =>
+ // The CancelResponse carries metrics + the cancel-callback error. Keep it
+ // on the terminal so close() can return it; any prior ErrorResponse is
+ // tracked in executionError and surfaced by the iterator.
+ Transitions.cancelled(ctrl.getCancel)
+ // Defensive: CANCEL before InitResponse unblocks doInit so it can
+ // surface the cancellation instead of timing out.
+ initValue.signalWithoutValue()
+
+ case UdfControlResponse.ControlCase.CONTROL_NOT_SET =>
+ Transitions.transportFailed(new IllegalStateException(
+ "empty UdfControlResponse oneof"))
+ initValue.signalWithoutValue()
+ }
+
+ /**
+ * True once init is no longer pending -- i.e. the stream is past `Initializing`.
+ * Not "init succeeded": a terminal (including a failure) also counts as resolved.
+ */
+ private def initResolved: Boolean = currentState match {
+ case SessionState.Created | SessionState.Initializing => false
+ case _ => true
+ }
+
+ private def cancelWithReason(reason: String): Cancel =
+ Cancel.newBuilder().setReason(reason).build()
+
+ /**
+ * The protocol transition graph in one place: names for the edges, not a
+ * second source of truth. Every edge acts on the single
+ * [[WorkerSession.SessionState]] owned by the base via [[compareAndSetState]]
+ * (non-terminal) or [[completeTerminal]] (terminal), so the base's CAS is the
+ * only synchronization and a terminal that arrived first always wins (the
+ * non-terminal CASes fail against it; [[completeTerminal]] is first-wins).
+ *
+ * Edges this class drives -- edge, method, then driver site(s) / thread (the
+ * base drives the API-call edges: `Created -> Initializing` in `init`,
+ * `Initialized -> Streaming` in `process`):
+ * {{{
+ * Initializing -> Initialized initAccepted handleControl INIT-ok [gRPC cb]
+ * Streaming -> Finishing finishSent advance branch 3 [engine]
+ * non-terminal -> Cancelling cancelSentFrom sendCancelInternal, iff a Cancel
+ * reaches the wire [gRPC cb | engine | init | close]
+ * non-terminal -> Terminal finished/cancelled/failed/transportFailed
+ * handleControl / doInit / doClose / advance / onError
+ * }}}
+ * `Cancelling` is reached only if a Cancel is actually written; a pre-stream or
+ * raced cancel goes straight to a `Cancelled`/`TransportFailed` terminal (see
+ * [[sendCancelInternal]], [[doClose]], `ProcessIterator`) or nowhere -- so
+ * `cancelRequested` (intent) does not imply state `Cancelling`.
+ */
+ private object Transitions {
+ /** `InitResponse` OK: `Initializing -> Initialized`. */
+ def initAccepted(): Boolean =
+ compareAndSetState(SessionState.Initializing, SessionState.Initialized)
+
+ /** Input exhausted and `Finish` written: `Streaming -> Finishing` (once). */
+ def finishSent(): Boolean =
+ compareAndSetState(SessionState.Streaming, SessionState.Finishing)
+
+ /** `Cancel` written: `cur -> Cancelling`, from any non-terminal `cur`. */
+ def cancelSentFrom(cur: SessionState): Boolean =
+ !cur.isTerminal && compareAndSetState(cur, SessionState.Cancelling)
+
+ /** Clean terminal carrying the worker's `FinishResponse`. */
+ def finished(response: FinishResponse): Boolean =
+ completeTerminal(Termination.Finished(response))
+
+ /** Clean terminal carrying the worker's `CancelResponse`. */
+ def cancelled(response: CancelResponse): Boolean =
+ completeTerminal(Termination.Cancelled(response))
+
+ /** Failure terminal carrying a structured [[ExecutionError]]. */
+ def failed(error: ExecutionError): Boolean =
+ completeTerminal(Termination.Failed(error))
+
+ /** Failure terminal carrying a transport-level cause. */
+ def transportFailed(cause: Throwable): Boolean =
+ completeTerminal(Termination.TransportFailed(cause))
+ }
+
+ // ---- WorkerSession hooks ------------------------------------------------
+
+ override protected def doInit(message: Init): InitResponse = {
+ // 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")
+ Transitions.transportFailed(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. close()'s 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 {
+ // Reentrant delivery: a directExecutor worker (one whose gRPC callbacks run
+ // synchronously on the caller's thread) can deliver InitResponse -- and thus
+ // run responseObserver/handleControl -- from *inside* this stream.onNext,
+ // before requestObserver below is published. The base advanced
+ // Created -> Initializing before calling doInit, so that reentrant
+ // InitResponse still finds Initializing and advances to Initialized; the
+ // handleControl error paths guard the still-null requestObserver. This is
+ // the scenario the "before doInit published requestObserver" comments mean.
+ 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
+ Transitions.transportFailed(e)
+ // Surface as GrpcWorkerSessionException so the engine integration layer
+ // (which catches that type and wraps it) sees a uniform init-failure
+ // exception rather than the raw transport error.
+ throw new GrpcWorkerSessionException("UDF worker stream failed during init", e)
+ }
+
+ val responded = try {
+ initValue.await(initResponseTimeoutMs)
+ } catch {
+ case _: InterruptedException =>
+ Thread.currentThread().interrupt()
+ sendCancelInternal(() => cancelWithReason("interrupted during init"))
+ // Make sure close() does not block waiting for a terminator that
+ // will never arrive on this thread's behalf.
+ Transitions.transportFailed(
+ new InterruptedException("interrupted while waiting for InitResponse"))
+ throw new InterruptedException("interrupted while waiting for InitResponse")
Review Comment:
who can interrupt this thread? if it is user, then this should be `cancel`? if this is the engine/system, then this should be `failed`?
##########
udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/GrpcWorkerSession.scala:
##########
@@ -0,0 +1,933 @@
+/*
+ * 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.grpc
+
+import java.util.concurrent.{CountDownLatch, LinkedBlockingQueue, TimeoutException, TimeUnit}
+import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference}
+
+import scala.util.control.NonFatal
+
+import io.grpc.{ConnectivityState, ManagedChannel}
+import io.grpc.stub.StreamObserver
+
+import org.apache.spark.annotation.Experimental
+import org.apache.spark.udf.worker.{Cancel, CancelResponse, DataRequest, DataResponse,
+ ExecutionError, Finish, FinishResponse, Init, InitResponse, UdfControlRequest,
+ UdfControlResponse, UdfRequest, UdfResponse, UdfWorkerGrpc}
+import org.apache.spark.udf.worker.core.{Termination, WorkerHandle, WorkerLogger, WorkerSession}
+import org.apache.spark.udf.worker.core.WorkerSession.SessionState
+import org.apache.spark.udf.worker.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_message.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.
+ *
+ * '''Driving model.''' Consumption-driven (Volcano / pull): the thread that
+ * consumes the [[doProcess]] result iterator is the one that pulls input and
+ * sends each `DataRequest`; the gRPC callback thread only receives output. It is
+ * pull-driven but not one-input-per-output -- `advance` sends the next input
+ * whenever the output queue is momentarily empty, so under async delivery it may
+ * push several input batches before any output is read (bounded by HTTP/2 flow
+ * control).
+ *
+ * '''State machine.''' This class does not keep its own state machine: it
+ * drives the single [[WorkerSession.SessionState]] owned by the base. The base
+ * advances `Created -> Initializing` (in `init`) and `Initialized -> Streaming`
+ * (in `process`); this class advances the protocol-event edges through
+ * [[compareAndSetState]] / [[completeTerminal]] as it exchanges messages:
+ * {{{
+ * Initializing --(InitResponse ok)--> Initialized [handleControl]
+ * Streaming ----(Finish written)----> Finishing [ProcessIterator]
+ * <any non-terminal> --(Cancel written)--> Cancelling [sendCancelInternal]
+ * <any non-terminal> --(terminator/error)--> terminal [completeTerminal]
+ * }}}
+ * The two clean terminals carry the worker's `FinishResponse` / `CancelResponse`
+ * (metrics + finish/cancel callback `data`/`error`) so [[close]] can return
+ * them. The only flag kept outside the machine is [[cancelRequested]]: a
+ * cancellation can be requested before the stream exists (so it cannot be a
+ * state transition yet), and it must both fast-fail the result iterator and
+ * suppress any in-flight Data/Finish.
+ *
+ * Threading:
+ * - [[doInit]] is synchronous: sends `Init` and blocks on `InitResponse`,
+ * returning it.
+ * - [[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.
+ * - [[doClose]] is thread-safe and idempotent: it settles + returns the
+ * terminator (cancelling in-flight work if the stream had not finished)
+ * and half-closes the request side.
+ *
+ * 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 a gRPC channel built and owned by the caller (the
+ * dispatcher). 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 from the worker, drained by the process() iterator.
+ // Intentionally unbounded: a bounded queue would block the gRPC callback
+ // (Netty event-loop) thread when full, stalling terminator/control delivery on
+ // the whole channel. HTTP/2 flow control bounds the wire and the consumer
+ // normally drains promptly; a stalled downstream can still grow it, but the fix
+ // is protocol-level back-pressure (out of scope), not bounding the queue.
+ //
+ // TODO [SPARK-57324]: expose queue depth as a metric (early warning for a
+ // stalled consumer).
+ private val outputQueue = new LinkedBlockingQueue[QueueItem]()
+
+ // The worker's `InitResponse` (success or error) coupled with the latch that
+ // init() blocks on. Fires when the InitResponse arrives (`complete`) or when a
+ // terminal settles first without one -- transport error, half-close, or a
+ // premature terminator (`signalWithoutValue`). Until it fires we have no proof
+ // the worker accepted the session.
+ //
+ // Settle-before-release rule (referenced from every callback that both settles
+ // a terminal and fires this latch): settle the terminal FIRST, then complete /
+ // signal. init() blocks on the latch, and only latch await/release establish a
+ // happens-before edge, so a woken init() is guaranteed to observe the terminal
+ // rather than a transient state. OneShotValue keeps that publish-then-release
+ // in one place instead of every caller remembering to count down after setting
+ // the reference.
+ private val initValue = new OneShotValue[InitResponse]
+
+ // Fired when the session reaches a terminal [[SessionState]]. doClose() and
+ // the init-error path block on this to drain the terminator.
+ private val terminalLatch = new CountDownLatch(1)
+
+ // Captures an ErrorResponse encountered during the data phase so that
+ // the CancelResponse terminator 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)
+
+ // Cancellation INTENT -- distinct from the `Cancelling` state, which is reached
+ // only once a Cancel is actually written to the wire ([[sendCancelInternal]]).
+ // Intent is set first and can outrun (or never reach) that write, so
+ // `cancelRequested` does NOT imply state `Cancelling`. Kept outside the
+ // [[SessionState]] machine because a cancel can be requested before the stream
+ // exists (pre-init), where there is no wire transition to make yet. Used to (a)
+ // make cancel idempotent across all call sites, (b) fast-fail ProcessIterator
+ // on a pre-init cancel, and (c) suppress any Data/Finish that would otherwise
+ // race a Cancel onto the wire (re-read inside [[requestLock]]).
+ private val cancelRequested = new AtomicBoolean(false)
+
+ // gRPC requires serialized writes to a request StreamObserver.
+ private val requestLock = new Object
+
+ // Initialised in init() -- before that, close() is a no-op 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 =>
+ // A DataResponse before InitResponse violates the protocol (InitResponse
+ // must precede any DataResponse): fast-fail rather than enqueue it and let
+ // init() block to its timeout. Settle-before-release (see initValue).
+ if (!initResolved) {
+ Transitions.transportFailed(new IllegalStateException(
+ "worker sent a DataResponse before InitResponse"))
+ initValue.signalWithoutValue()
+ } else {
+ outputQueue.put(QueueItem.Batch(response.getData))
+ }
+
+ case UdfResponse.ResponseCase.CONTROL =>
+ handleControl(response.getControl)
+
+ case other =>
+ // A malformed response (empty / unknown oneof) is a terminal transport
+ // failure; fast-fail init the same way. Settle-before-release (see initValue).
+ Transitions.transportFailed(new IllegalStateException(
+ s"unexpected response oneof: $other"))
+ initValue.signalWithoutValue()
+ }
+ }
+
+ override def onError(t: Throwable): Unit = {
+ // Transport-level failure: the stream is dead, no further writes possible.
+ // Settle-before-release (see initValue) so init() surfaces the transport
+ // cause instead of the initResponseTimeoutMs "timed out" error.
+ Transitions.transportFailed(t)
+ initValue.signalWithoutValue()
+ }
+
+ 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 (!currentState.isTerminal) {
+ Transitions.transportFailed(new IllegalStateException(
+ "worker response stream closed without a terminator"))
+ }
+ // Defensive: if onCompleted reached us before InitResponse, doInit is
+ // still blocked on initValue and would otherwise time out.
+ initValue.signalWithoutValue()
+ }
+ }
+
+ /**
+ * Wakes the result iterator (blocked on [[outputQueue]]) and any thread
+ * waiting on [[terminalLatch]] when the base settles a terminal. Invoked once,
+ * by the caller that wins [[completeTerminal]].
+ */
+ override protected def onTerminalSettled(termination: Termination): Unit = {
+ outputQueue.put(QueueItem.EndOfStream)
+ terminalLatch.countDown()
+ }
+
+ private def handleControl(ctrl: UdfControlResponse): Unit = ctrl.getControlCase match {
+ case UdfControlResponse.ControlCase.INIT =>
+ val resp = ctrl.getInit
+ if (resp.hasError) {
+ // Attribute a subsequent close()'s terminator to the init error rather
+ // than a bare Cancelled, mirroring the data-phase ERROR branch.
+ executionError.compareAndSet(None, Some(resp.getError))
+ // Settle Failed(err) before the Cancel below (not after its CancelResponse):
+ // the sticky terminal makes close() report the init error, and the Cancel
+ // becomes best-effort worker cleanup -- so if it can't be written (a
+ // directExecutor worker delivered this reentrantly, before doInit published
+ // requestObserver), doInit still doesn't wait terminalTimeoutMs for a
+ // CancelResponse that can never arrive. Settle-before-release (see initValue).
+ Transitions.failed(resp.getError)
+ initValue.complete(resp)
+ sendCancelInternal(() => cancelWithReason("init failed"))
+ } else {
+ // InitResponse OK. Only advance from Initializing so a terminal that raced
+ // in (e.g. a transport error) still wins; process() then opens the data
+ // phase. Settle-before-release (see initValue): publish after the CAS.
+ Transitions.initAccepted()
+ initValue.complete(resp)
Review Comment:
complete the value before transit the state?
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57640] Implement gRPC WorkerSession and test it against an echo server. [spark]
Posted by "haiyangsun-db (via GitHub)" <gi...@apache.org>.
haiyangsun-db commented on code in PR #56702:
URL: https://github.com/apache/spark/pull/56702#discussion_r3707814918
##########
udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/GrpcWorkerSession.scala:
##########
@@ -0,0 +1,1177 @@
+/*
+ * 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.grpc
+
+import java.util.Objects
+import java.util.concurrent.{CountDownLatch, LinkedBlockingQueue, TimeoutException, TimeUnit}
+import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference}
+
+import scala.util.control.NonFatal
+
+import io.grpc.{ConnectivityState, ManagedChannel}
+import io.grpc.stub.StreamObserver
+
+import org.apache.spark.annotation.Experimental
+import org.apache.spark.udf.worker.{Cancel, CancelResponse, DataRequest, DataResponse,
+ ExecutionError, Finish, FinishResponse, Init, InitResponse, UdfControlRequest,
+ UdfControlResponse, UdfPayload, UdfRequest, UdfResponse, UdfWorkerGrpc}
+import org.apache.spark.udf.worker.core.{Termination, WorkerHandle, WorkerLogger, WorkerSession}
+import org.apache.spark.udf.worker.core.WorkerSession.SessionState
+import org.apache.spark.udf.worker.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_message.proto` (`PayloadChunk*`
+ * between Init and InitResponse omitted here; chunking is not yet
+ * implemented -- see the TODO below):
+ * {{{
+ * Engine -> Worker: Init -> (DataRequest)* -> Finish (Cancel)?
+ * | Cancel
+ * Worker -> Engine: InitResponse -> (DataResponse)* ->
+ * (ErrorResponse)? -> (FinishResponse | CancelResponse)
+ * }}}
+ *
+ * Knows nothing about how the worker was provisioned -- the dispatcher
+ * constructs this with a [[WorkerHandle]] and channel; the base
+ * [[WorkerSession]] handles dispatcher-side cleanup on close.
+ *
+ * '''Driving model.''' Consumption-driven (Volcano / pull): the thread that
+ * consumes the [[doProcess]] result iterator is the one that pulls input and
+ * sends each `DataRequest`; the gRPC callback thread only receives output. It is
+ * pull-driven but not one-input-per-output -- `advance` sends the next input
+ * whenever the output queue is momentarily empty, so under async delivery it may
+ * push several input batches before any output is read. HTTP/2 flow control
+ * bounds wire traffic, but application-level buffering is intentionally left to
+ * a follow-up (see the TODO on [[outputQueue]]).
+ *
+ * '''State machine.''' This class does not keep its own state machine: it
+ * drives the single [[WorkerSession.SessionState]] owned by the base. The base
+ * advances `Created -> Initializing` (in `init`) and `Initialized -> Streaming`
+ * (in `process`); this class advances the protocol-event edges through
+ * [[compareAndSetState]] / [[completeTerminal]] as it exchanges messages:
+ * {{{
+ * Initializing --(InitResponse ok)--> Initialized [handleControl]
+ * Streaming ----(input exhausted)-------> Finishing [ProcessIterator]
+ * <any non-terminal> --(Cancel send starts)--> Cancelling [sendCancelInternal]
+ * <any non-terminal> --(terminator/error)--> terminal [completeTerminal]
+ * }}}
+ * The two clean terminals carry the worker's `FinishResponse` / `CancelResponse`
+ * (metrics + finish/cancel callback `data`/`error`) so [[close]] can return
+ * them. Cancellation intent is deliberately tracked outside the machine in
+ * [[cancelRequested]], which makes the wire write idempotent and suppresses any
+ * in-flight Data/Finish once cancellation begins.
+ *
+ * Threading:
+ * - [[doInit]] is synchronous: sends `Init` and blocks on `InitResponse`,
+ * returning it.
+ * - [[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.
+ * - [[doClose]] is thread-safe and idempotent: it settles + returns the
+ * terminator (cancelling in-flight work if the stream had not finished)
+ * and terminates the request side.
+ *
+ * TODO [SPARK-55278]: this class does not yet implement payload chunking;
+ * the entire [[UdfPayload]] 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 a gRPC channel built and owned by the caller (the
+ * dispatcher). 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`.
+ * @param interruptCancelTimeoutMs upper bound on the wait for a `CancelResponse`
+ * after an interrupt (cancelled query / killed
+ * task) sends `Cancel`. Short by design; on expiry
+ * the session settles `Interrupted`. See
+ * `handleInterrupt`.
+ */
+@Experimental
+class GrpcWorkerSession(
+ workerHandle: WorkerHandle,
+ channel: ManagedChannel,
+ logger: WorkerLogger = WorkerLogger.NoOp,
+ initResponseTimeoutMs: Long = DEFAULT_INIT_RESPONSE_TIMEOUT_MS,
+ terminalTimeoutMs: Long = DEFAULT_TERMINAL_TIMEOUT_MS,
+ interruptCancelTimeoutMs: Long = DEFAULT_INTERRUPT_CANCEL_TIMEOUT_MS)
+ extends WorkerSession(workerHandle, logger) {
+
+ require(channel != null, "channel is required")
+
+ private val asyncStub = UdfWorkerGrpc.newStub(channel)
+
+ // Output batches from the worker, drained by the process() iterator.
+ // Intentionally unbounded in this first implementation: a bounded queue would
+ // block the gRPC callback (Netty event-loop) thread when full, stalling
+ // terminator/control delivery on the whole channel. The consumer normally
+ // drains promptly, but HTTP/2 flow control alone does not bound this application
+ // queue or gRPC's asynchronous send buffers.
+ //
+ // TODO [SPARK-55278]: add application-level gRPC flow control in a follow-up,
+ // using ClientCallStreamObserver readiness for requests and manual inbound
+ // demand for responses, before wiring this transport into a production path.
+ // TODO [SPARK-57324]: expose queue depth as a metric (early warning for a
+ // stalled consumer).
+ private val outputQueue = new LinkedBlockingQueue[QueueItem]()
+
+ // This value couples the worker's `InitResponse` (success or error) with the
+ // latch on which init() blocks. The latch fires when the InitResponse arrives
+ // (`complete`), when a pre-init ErrorResponse arrives, or when a terminal
+ // settles first without an InitResponse (`signalWithoutValue`). Until it fires
+ // we have no proof the worker accepted the session.
+ //
+ // Settle-before-release rule (referenced from every callback that both settles
+ // a terminal and fires this latch): settle the terminal FIRST, then complete /
+ // signal. init() blocks on the latch, and only a latch await/release pair
+ // establishes a happens-before edge, so a woken init() is guaranteed to
+ // observe the terminal rather than a transient state. OneShotValue keeps that
+ // publish-then-release in one place instead of every caller remembering to
+ // count down after setting the reference.
+ private val initValue = new OneShotValue[InitResponse]
+
+ // Fired when the session reaches a terminal [[SessionState]]. doClose() and
+ // the init-error path block on this to drain the terminator.
+ private val terminalLatch = new CountDownLatch(1)
+
+ // Captures an ErrorResponse encountered during the data phase so that
+ // the CancelResponse terminator 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 immediately before Finish is handed to the request observer. It is set
+ // before onNext so a reentrant callback may deliver FinishResponse while the
+ // Finish request is still being written.
+ private val finishSendStarted = new AtomicBoolean(false)
+
+ // Cancellation INTENT -- distinct from the `Cancelling` state, which is reached
+ // once a Cancel send is attempted under [[requestLock]] ([[sendCancelInternal]]).
+ // Intent is set first and can outrun (or never reach) that attempt, so
+ // `cancelRequested` does NOT imply state `Cancelling`. It is kept outside the
+ // [[SessionState]] machine to (a) make cancellation idempotent across all call
+ // sites and (b) suppress any Data/Finish that would otherwise race a Cancel
+ // onto the wire (re-read inside [[requestLock]]).
+ private val cancelRequested = new AtomicBoolean(false)
+
+ // gRPC requires serialized writes to a request StreamObserver.
+ private val requestLock = new Object
+
+ // Initialized in init() -- before that, close() is a no-op on the request
+ // side, which is exactly the contract the wrapping WorkerSession expects.
+ @volatile private var requestObserver: StreamObserver[UdfRequest] = _
+
+ // True after the request side has been half-closed, aborted, or observed the
+ // transport closing. It prevents this session from issuing duplicate terminal
+ // calls, including following onError with onCompleted from close().
+ private val requestSideTerminated = new AtomicBoolean(false)
+
+ private val responseObserver: StreamObserver[UdfResponse] = new StreamObserver[UdfResponse] {
+ override def onNext(response: UdfResponse): Unit = {
+ response.getResponseCase match {
+ case UdfResponse.ResponseCase.DATA =>
+ // A DataResponse before InitResponse violates the protocol (InitResponse
+ // must precede any DataResponse): fast-fail rather than enqueue it and let
+ // init() block to its timeout. Settling the terminal wakes a blocked
+ // init() via onTerminalSettled.
+ if (!initResolved) {
+ Transitions.transportFailed(new IllegalStateException(
+ "worker sent a DataResponse before InitResponse"))
+ } else {
+ outputQueue.put(QueueItem.Batch(response.getData))
+ }
+
+ case UdfResponse.ResponseCase.CONTROL =>
+ handleControl(response.getControl)
+
+ case other =>
+ // A malformed response (empty / unknown oneof) is a terminal transport
+ // failure; fast-fail init the same way.
+ Transitions.transportFailed(new IllegalStateException(
+ s"unexpected response oneof: $other"))
+ }
+ }
+
+ override def onError(t: Throwable): Unit = {
+ // Transport-level failure: the stream is dead, no further writes possible.
+ // Settling the terminal wakes a blocked init() (via onTerminalSettled) so it
+ // surfaces the transport cause instead of the initResponseTimeoutMs error.
+ requestSideTerminated.set(true)
+ Transitions.transportFailed(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. Settling the terminal wakes a blocked
+ // init() via onTerminalSettled.
+ requestSideTerminated.set(true)
+ if (!currentState.isTerminal) {
+ Transitions.transportFailed(new IllegalStateException(
+ "worker response stream closed without a terminator"))
+ }
+ }
+ }
+
+ /**
+ * Wakes everything that can be blocked when the base settles a terminal:
+ * the result iterator (on [[outputQueue]]), a thread on [[terminalLatch]]
+ * (close), and a thread still blocked in [[doInit]] on [[initValue]]. Invoked
+ * once, by the caller that wins [[completeTerminal]], so this is the '''single'''
+ * place a settled terminal wakes a blocked [[doInit]] -- callers that settle a
+ * terminal (any response-callback failure/terminator branch, the timeout paths,
+ * or close()) do NOT signal [[initValue]] themselves; settling is enough.
+ * Settle-before-release (see [[initValue]]) holds because this runs after the
+ * terminal CAS in [[completeTerminal]]. The only direct [[initValue]] signals
+ * left are the non-terminal init paths ([[handleControl]]'s INIT-ok / INIT-error
+ * / pre-init-ERROR branches), which must wake [[doInit]] '''without''' settling a
+ * terminal. INIT-ok first advances the session to `Initialized`; the error
+ * branches remain `Initializing` only until [[doInit]] sends Cancel and throws.
+ */
+ override protected def onTerminalSettled(termination: Termination): Unit = {
+ outputQueue.put(QueueItem.EndOfStream)
+ initValue.signalWithoutValue()
+ terminalLatch.countDown()
+ }
+
+ private def handleControl(ctrl: UdfControlResponse): Unit = ctrl.getControlCase match {
+ case UdfControlResponse.ControlCase.INIT =>
+ val resp = ctrl.getInit
+ if (resp.hasError) {
+ // Record the error and publish the InitResponse; do NOT settle a terminal
+ // or send Cancel here. The proto requires the engine to Cancel after an
+ // init error (udf_message.proto); doInit owns that synchronous cleanup
+ // after the initial onNext returns, then drains the CancelResponse before
+ // throwing. Settle-before-release (see initValue): there is no terminal to
+ // settle first here, just publish.
+ executionError.compareAndSet(None, Some(resp.getError))
+ initValue.complete(resp)
+ } else {
+ // InitResponse OK. Only advance from Initializing so a terminal that raced
+ // in (e.g. a transport error) still wins; process() then opens the data
+ // phase. Settle-before-release (see initValue): publish after the CAS.
+ Transitions.initAccepted()
+ initValue.complete(resp)
+ }
+
+ case UdfControlResponse.ControlCase.ERROR =>
+ val err = ctrl.getError.getError
+ executionError.compareAndSet(None, Some(err))
+ if (!initResolved) {
+ // Pre-init ErrorResponse. Leave the session in Initializing and just wake
+ // init(): the recorded executionError tells doInit the worker failed, and
+ // doInit sends the Cancel and drains the CancelResponse after the initial
+ // onNext returns, just like the INIT-error branch.
+ initValue.signalWithoutValue()
+ } else {
+ // Data-phase ErrorResponse: requestObserver is published, so cancel here.
+ // Cancel -> CancelResponse settles the terminal; the iterator surfaces the
+ // recorded executionError. The init latch already fired in init(), so no
+ // signalWithoutValue is needed.
+ sendCancelInternal(() => cancelWithReason("aborting after ErrorResponse"))
+ }
+
+ case UdfControlResponse.ControlCase.FINISH =>
+ // The FinishResponse carries metrics + the finish-callback data/error.
+ // Keep it on the terminal so close() can return it; the iterator inspects
+ // its error field to decide whether to throw. A FinishResponse is valid only
+ // after the engine has started sending Finish; accepting one earlier could
+ // silently truncate input and report a clean result.
+ if (finishSendStarted.get()) {
+ Transitions.finished(ctrl.getFinish)
+ } else {
+ Transitions.transportFailed(new IllegalStateException(
+ "worker sent FinishResponse before the engine sent Finish"))
+ }
+
+ case UdfControlResponse.ControlCase.CANCEL =>
+ // The CancelResponse carries metrics + the cancel-callback error. Keep it
+ // on the terminal so close() can return it; any prior ErrorResponse is
+ // tracked in executionError and surfaced by the iterator. Settling the
+ // terminal wakes a blocked init() (a CANCEL before InitResponse) via
+ // onTerminalSettled.
+ Transitions.cancelled(ctrl.getCancel)
+
+ case UdfControlResponse.ControlCase.CONTROL_NOT_SET =>
+ // Settling the terminal wakes a blocked init() via onTerminalSettled.
+ Transitions.transportFailed(new IllegalStateException(
+ "empty UdfControlResponse oneof"))
+ }
+
+ /**
+ * True once init is no longer pending -- i.e. the stream is past `Initializing`.
+ * Not "init succeeded": a terminal (including a failure) also counts as resolved.
+ */
+ private def initResolved: Boolean = currentState match {
+ case SessionState.Created | SessionState.Initializing => false
+ case _ => true
+ }
+
+ private def cancelWithReason(reason: String): Cancel =
+ Cancel.newBuilder().setReason(reason).build()
+
+ /**
+ * The protocol transition graph in one place: names for the edges, not a
+ * second source of truth. Every edge acts on the single
+ * [[WorkerSession.SessionState]] owned by the base via [[compareAndSetState]]
+ * (non-terminal) or [[completeTerminal]] (terminal), so the base's CAS is the
+ * only synchronization and a terminal that arrived first always wins (the
+ * non-terminal CASes fail against it; [[completeTerminal]] is first-wins).
+ *
+ * Edges this class drives -- edge, method, then driver site(s) / thread (the
+ * base drives the API-call edges: `Created -> Initializing` in `init`,
+ * `Initialized -> Streaming` in `process`):
+ * {{{
+ * Initializing -> Initialized initAccepted handleControl INIT-ok [gRPC cb]
+ * Streaming -> Finishing beginFinish advance branch 3 [engine]
+ * non-terminal -> Cancelling beginCancelFrom sendCancelInternal, immediately
+ * before the Cancel send [gRPC cb | engine | init | close]
+ * non-terminal -> Terminal finished/cancelled/transportFailed/interrupted
+ * handleControl / doInit / doClose / advance / onError
+ * }}}
+ * `Cancelling` is reached only when a Cancel send starts; a pre-stream or
+ * raced cancel goes straight to a `Cancelled`/`TransportFailed` terminal (see
+ * [[sendCancelInternal]], [[doClose]], `ProcessIterator`) or nowhere -- so
+ * `cancelRequested` (intent) does not imply state `Cancelling`.
+ */
+ private object Transitions {
+ /** `InitResponse` OK: `Initializing -> Initialized`. */
+ def initAccepted(): Boolean =
+ compareAndSetState(SessionState.Initializing, SessionState.Initialized)
+
+ /** Input exhausted: `Streaming -> Finishing` (once). */
+ def beginFinish(): Boolean =
+ compareAndSetState(SessionState.Streaming, SessionState.Finishing)
+
+ /** `Cancel` send starts: `cur -> Cancelling`, from any non-terminal `cur`. */
+ def beginCancelFrom(cur: SessionState): Boolean =
+ !cur.isTerminal && compareAndSetState(cur, SessionState.Cancelling)
+
+ /** Clean terminal carrying the worker's `FinishResponse`. */
+ def finished(response: FinishResponse): Boolean =
+ completeTerminal(Termination.Finished(response))
+
+ /** Clean terminal carrying the worker's `CancelResponse`. */
+ def cancelled(response: CancelResponse): Boolean =
+ completeTerminal(Termination.Cancelled(response))
+
+ /** Failure terminal carrying a transport-level cause. */
+ def transportFailed(cause: Throwable): Boolean =
+ completeTerminal(Termination.TransportFailed(cause))
+
+ /**
+ * Terminal for an engine-thread interrupt (e.g. a Spark task kill). Distinct
+ * from [[transportFailed]] because the cause is an engine-side interrupt, but
+ * still unsalvageable without a worker acknowledgement. Settled by
+ * [[handleInterrupt]] only when the brief post-Cancel drain expires without a
+ * `CancelResponse`; if the worker acks in time the session settles a
+ * cooperative [[cancelled]] instead.
+ */
+ def interrupted(cause: Throwable): Boolean =
+ completeTerminal(Termination.Interrupted(cause))
+ }
+
+ // ---- WorkerSession hooks ------------------------------------------------
+
+ override protected def doInit(message: Init): InitResponse = {
+ // Construct the request before opening the RPC. A malformed Init must not
+ // leave an otherwise-unused Execute stream waiting for its first request.
+ val initRequest = UdfRequest.newBuilder()
+ .setControl(UdfControlRequest.newBuilder().setInit(message).build())
+ .build()
+
+ // 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")
+ Transitions.transportFailed(ex)
+ throw new GrpcWorkerSessionException("UDF worker channel is closed", ex)
+ }
+ try {
+ requestLock.synchronized {
+ // Serialize stream creation, publication, and Init with close/cancel. A
+ // close that wins the lock first prevents the RPC from opening; one that
+ // loses cannot release the worker until the initial write returns.
+ if (!currentState.isTerminal) {
+ val stream = asyncStub.execute(responseObserver)
+ requestObserver = stream
+ // Stream creation may invoke a response callback synchronously. Do not
+ // send Init if that callback has already settled a terminal.
+ if (!currentState.isTerminal) {
+ // With directExecutor, gRPC callbacks run synchronously on the caller's
+ // thread, so InitResponse can reach responseObserver/handleControl from
+ // *inside* this stream.onNext.
+ // requestObserver is already published, so an immediate ErrorResponse
+ // can write its required Cancel reentrantly without losing the request.
+ sendOnNext(initRequest)
+ }
+ }
+ }
+ } catch {
+ case NonFatal(e) =>
+ Transitions.transportFailed(e)
+ // Surface as GrpcWorkerSessionException so the engine integration layer
+ // (which catches that type and wraps it) sees a uniform init-failure
+ // exception rather than the raw transport error.
+ throw new GrpcWorkerSessionException("UDF worker stream failed during init", e)
+ }
+
+ try {
+ initValue.await(initResponseTimeoutMs)
+ } catch {
+ case _: InterruptedException =>
+ // Interrupt (cancelled query / killed task) while awaiting InitResponse:
+ // cooperatively Cancel with a bounded drain, settling Cancelled if the
+ // worker acks in time, else Interrupted. See handleInterrupt.
+ handleInterrupt("waiting for InitResponse")
+ throw new InterruptedException("interrupted while waiting for InitResponse")
+ case e: TimeoutException =>
+ sendCancelInternal(() => cancelWithReason("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.
+ Transitions.transportFailed(e)
+ // Surface as GrpcWorkerSessionException (carrying the timeout cause) so
+ // the engine integration layer that catches that type can wrap it.
+ throw new GrpcWorkerSessionException(
+ s"timed out waiting for InitResponse after ${initResponseTimeoutMs}ms", e)
+ }
+
+ initValue.get match {
+ case Some(resp) if resp.hasError =>
+ failInitWithError(resp.getError,
+ s"UDF worker init failed: ${describeError(resp.getError)}")
+ case Some(resp) =>
+ executionError.get() match {
+ case Some(err) =>
+ // A reentrant callback can report an ErrorResponse immediately after
+ // InitResponse and before the initial onNext returns. The callback
+ // has already sent Cancel; drain its response
+ // and surface the original error instead of returning init success
+ // for a session that is already cancelling or terminal.
+ failInitWithError(err,
+ s"UDF worker reported an error as init completed: ${describeError(err)}")
+ case None if currentState != SessionState.Initialized =>
+ // Likewise, an immediate protocol failure or concurrent close may
+ // have moved the session out of Initialized after publishing the
+ // InitResponse. A normal return would violate init()'s contract and
+ // leave process() to fail later with only an ordering error.
+ failInitFromCurrentState()
+ case None =>
+ resp
+ }
+ case None =>
+ // No InitResponse arrived but the latch fired.
+ //
+ // A pre-init ErrorResponse leaves the session in Initializing with the
+ // error recorded in executionError (see handleControl ERROR branch): the
+ // engine must now Cancel and drain the CancelResponse, which failInit does.
+ executionError.get() match {
+ case Some(err) if !currentState.isTerminal =>
+ failInitWithError(err,
+ s"UDF worker reported an error before init completed: ${describeError(err)}")
+ case _ =>
+ // Otherwise a terminal already settled -- the worker terminated the
+ // stream before sending InitResponse (transport error, half-close, or
+ // a premature FinishResponse/CancelResponse). Surface it as an init
+ // failure rather than letting the caller proceed as if init succeeded.
+ failInitFromCurrentState()
+ }
+ }
+ }
+
+ /** Surfaces the terminal that prevented init from completing normally. */
+ private def failInitFromCurrentState(): Nothing = {
+ executionError.get() match {
+ case Some(err) =>
+ failInitWithError(err,
+ s"UDF worker reported an error as init completed: ${describeError(err)}")
+ case None => ()
+ }
+ // Cancelling can be observed after a concurrent close or a reentrant error
+ // before its CancelResponse arrives. Drain it so the exception reflects the
+ // stable terminal rather than a transient state.
+ if (currentState == SessionState.Cancelling) {
+ awaitTerminal()
+ }
+ currentState match {
+ case SessionState.Terminal(Termination.TransportFailed(cause)) =>
+ throw new GrpcWorkerSessionException("UDF worker stream failed during init", cause)
+ case SessionState.Terminal(Termination.Failed(err)) =>
+ throw new GrpcWorkerSessionException(
+ s"UDF worker reported an error before init completed: ${describeError(err)}", err)
+ case SessionState.Terminal(Termination.Cancelled(_)) =>
+ throw new GrpcWorkerSessionException(
+ "UDF worker stream was cancelled before init completed")
+ case SessionState.Terminal(Termination.Interrupted(cause)) =>
+ throw new GrpcWorkerSessionException("UDF worker init was interrupted", cause)
+ case SessionState.Terminal(Termination.Finished(_)) =>
+ throw new GrpcWorkerSessionException("UDF worker finished before init completed")
+ case other =>
+ throw new IllegalStateException(
+ s"init completed without an accepted session or terminal: $other")
+ }
+ }
+
+ /**
+ * Fails init when the worker reports an error before init returns: an
+ * `InitResponse` carrying an error, a pre-init `ErrorResponse`, or an immediate
+ * post-init `ErrorResponse`. The protocol requires the engine to send `Cancel`
+ * and the worker to reply with `CancelResponse` (udf_message.proto); we send it,
+ * drain the terminator so no stream is left dangling, and throw the structured
+ * error. A responsive worker settles `Cancelled`; a failed drain settles
+ * `TransportFailed`.
+ */
+ private def failInitWithError(err: ExecutionError, message: String): Nothing = {
+ sendCancelInternal(() => cancelWithReason("init failed"))
+ awaitTerminal()
+ throw new GrpcWorkerSessionException(message, err)
+ }
+
+ override protected def doProcess(
+ input: Iterator[DataRequest],
+ finish: () => Finish): Iterator[DataResponse] = {
+ // Init success is guaranteed by the base [[WorkerSession]] lifecycle: if
+ // doInit had failed it would have thrown and process() would never run.
+ new ProcessIterator(input, finish)
+ }
+
+ override protected def doClose(cancel: () => Cancel): Termination = {
+ // Coordinate the no-stream decision with doInit's publication + Init write.
+ // If close wins this lock, doInit observes the terminal and never sends Init;
+ // if doInit wins, close observes the published observer and sends Cancel only
+ // after the Init send has returned.
+ val noPublishedStream = requestLock.synchronized {
+ if (requestObserver == null) {
+ // init() never put a stream on the wire (closed before/around init, or
+ // init threw before publishing). There is no protocol terminator; if no
+ // terminal has settled yet, treat the session as cancelled-before-start.
+ // A terminal may already be settled here (e.g. the channel-shutdown
+ // TransportFailed in doInit also leaves requestObserver null); return the
+ // settled terminal as-is rather than a bare Cancelled that disagrees with
+ // the state. The base WorkerSession still releases the worker handle.
+ if (!currentState.isTerminal) {
+ Transitions.cancelled(CancelResponse.getDefaultInstance)
+ }
+ true
+ } else {
+ false
+ }
+ }
+ if (noPublishedStream) {
+ return settledTermination
+ }
+ // If the stream has not finished on its own, cancel anything in flight so
+ // the worker can clean up, then wait for the terminator. sendCancelInternal
+ // evaluates the cancel thunk only when it attempts a Cancel, and at
+ // most once across all callers.
+ if (!currentState.isTerminal) {
+ sendCancelInternal(cancel)
+ try {
+ terminalLatch.await(terminalTimeoutMs, TimeUnit.MILLISECONDS)
+ } catch {
+ case _: InterruptedException =>
+ // Interrupted mid-close: settle Interrupted via the
+ // shared handler rather than falling through to the TransportFailed
+ // guard below. The Cancel was already attempted, so handleInterrupt's
+ // sendCancelInternal is a no-op and it just runs the bounded drain.
+ //
+ // TODO [SPARK-57640]: close() is a finalizer often already on the task
+ // kill path, yet this makes an interrupted close block up to another
+ // interruptCancelTimeoutMs (chasing a clean, salvageable Cancelled
+ // rather than immediately unwinding to Interrupted). Current choice:
+ // spend the bounded wait to obtain proof the worker is recyclable; it
+ // is small next to terminalTimeoutMs. Revisit if killed-task teardown
+ // latency (especially many sessions torn down at once) makes an
+ // immediate unwind preferable -- e.g. skip this second drain in the
+ // close() path, since close already gave the worker its terminalTimeoutMs window.
+ handleInterrupt("closing the session")
+ }
+ }
+ // If the worker still has not settled (timeout above, or an interrupt whose
+ // bounded drain did not settle a terminal), record a terminal so
+ // isWorkerSalvageable returns a definite answer and any other thread reading
+ // the state sees a stable value.
+ if (!currentState.isTerminal) {
+ Transitions.transportFailed(new TimeoutException(
+ s"timed out waiting for stream terminator after ${terminalTimeoutMs}ms"))
+ }
+ // Close the request side according to the settled outcome. Clean protocol
+ // terminators are half-closed; failures abort the RPC so close() never follows
+ // an onError with onCompleted.
+ currentState match {
+ case SessionState.Terminal(Termination.Finished(_)) |
+ SessionState.Terminal(Termination.Cancelled(_)) =>
+ completeRequestSide()
+ case SessionState.Terminal(Termination.Failed(error)) =>
+ abortRequestSide(new GrpcWorkerSessionException(
+ s"UDF execution failed: ${describeError(error)}", error))
+ case SessionState.Terminal(Termination.TransportFailed(cause)) =>
+ abortRequestSide(cause)
+ case SessionState.Terminal(Termination.Interrupted(cause)) =>
+ abortRequestSide(cause)
+ case other =>
+ logger.debug(s"UDF Execute stream closed without a terminal outcome: $other")
+ }
+ // Derive the Termination from the settled terminal. close() is the
+ // finalizer/cleanup path and must NOT re-throw: a UDF / data-phase error is
+ // already surfaced while the result iterator is consumed, and an init error
+ // is surfaced from init(). settledTermination returns the settled terminal
+ // as-is: the response proto for the clean terminators, and the failure
+ // terminals (Failed / TransportFailed / Interrupted) carrying their cause --
+ // not a bare Cancelled.
+ settledTermination
+ }
+
+ // ---- Internal request helpers ---------------------------------------------
+
+ /**
+ * Writes one request and aborts the request side if the observer rejects it.
+ * Callers hold [[requestLock]], so this also serializes the compensating
+ * `onError` with every other request-side operation.
+ */
+ private def sendOnNext(req: UdfRequest): Unit = {
+ try {
+ requestObserver.onNext(req)
+ } catch {
+ case NonFatal(e) =>
+ abortRequestSide(e)
+ throw e
+ }
+ }
+
+ /** Aborts the request side at most once. Safe to call while holding [[requestLock]]. */
+ private def abortRequestSide(cause: Throwable): Unit = requestLock.synchronized {
+ if (requestObserver != null && requestSideTerminated.compareAndSet(false, true)) {
+ try {
+ requestObserver.onError(cause)
+ } catch {
+ case NonFatal(e) => logger.debug("Error aborting UDF Execute stream", e)
+ }
+ }
+ }
+
+ /** Half-closes the request side at most once. */
+ private def completeRequestSide(): Unit = requestLock.synchronized {
+ if (requestObserver != null && requestSideTerminated.compareAndSet(false, true)) {
+ try {
+ requestObserver.onCompleted()
+ } catch {
+ case NonFatal(e) => logger.debug("Error half-closing UDF Execute stream", e)
+ }
+ }
+ }
+
+ /**
+ * Sends a Data or Finish request to the worker. Three invariants are
+ * checked inside [[requestLock]]:
+ * - terminal state: writes are unsafe (transport dead or terminal
+ * received). Throws, so the caller's terminal/exception path runs.
+ * - request-side termination: no write may follow `onError` / `onCompleted`.
+ * Throws so the caller settles or surfaces the failure.
+ * - [[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]] publishes
+ * [[requestObserver]] and writes Init together under [[requestLock]], so a
+ * concurrent cancel observes the observer but cannot acquire the lock and send
+ * Cancel until after Init.
+ */
+ private def sendRequest(req: UdfRequest): Unit =
+ requestLock.synchronized {
+ if (currentState.isTerminal) {
+ throw new IllegalStateException(
+ "cannot send request: UDF Execute stream is already closed")
+ }
+ if (requestSideTerminated.get()) {
+ throw new IllegalStateException(
+ "cannot send request: UDF Execute request stream is already closed")
+ }
+ // Suppress the write when a cancel has been (or is about to be) flushed
+ // through this lock; that preserves the proto ordering invariant (no
+ // Data/Finish after Cancel). Test the flag rather than `return`-ing from
+ // this by-name `synchronized` body: a non-local return compiles to a thrown
+ // NonLocalReturnControl, which is brittle (also relied on in sendCancelInternal).
+ if (!cancelRequested.get()) {
+ if (req.hasControl && req.getControl.hasFinish) {
+ // Set before onNext: directExecutor delivery may return FinishResponse
+ // reentrantly from inside this call.
+ finishSendStarted.set(true)
+ }
+ sendOnNext(req)
+ }
+ }
+
+ /**
+ * Sends a `Cancel` control message, returning `true` iff the request observer
+ * accepts the `Cancel`. Idempotent across ALL call sites (close()'s
+ * cancel, in-band cancels from `handleControl`'s INIT error / ERROR paths,
+ * and engine-internal cancels from the iterator) via [[cancelRequested]]:
+ * only the first caller attempts a `Cancel`, and the `cancel` thunk is evaluated
+ * only by that caller and only when the send is about to start --
+ * so a side-effecting thunk (e.g. one carrying a client cancel callback) runs
+ * at most once. 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(cancel: () => Cancel): Boolean = {
+ // Do not consume the one-shot cancellation intent until there is an observer
+ // that can carry it. doInit publishes the observer before sending Init while
+ // holding requestLock, so every response to Init sees a non-null observer.
+ if (requestObserver == null) return false
+ if (!cancelRequested.compareAndSet(false, true)) return false
+ if (currentState.isTerminal) return false
+ try {
+ requestLock.synchronized {
+ // A terminal that arrived first must win; bail without writing. Compute
+ // the block's value rather than `return` (see sendRequest for why non-local
+ // return is avoided -- here it would also escape the NonFatal catch below).
+ val cur = currentState
+ if (cur.isTerminal || requestSideTerminated.get()) {
+ false
+ } else {
+ val request = UdfRequest.newBuilder()
+ .setControl(UdfControlRequest.newBuilder().setCancel(cancel()).build())
+ .build()
+ // Record that a Cancel send has started by advancing to Cancelling.
+ Transitions.beginCancelFrom(cur)
+ sendOnNext(request)
+ true
+ }
+ }
+ } catch {
+ case NonFatal(e) =>
+ logger.debug(s"Cancel send failed (stream may already be torn down): ${e.getMessage}")
+ Transitions.transportFailed(e)
+ false
+ }
+ }
+
+ private def awaitTerminal(): Unit = {
+ if (currentState.isTerminal) return
+ try {
+ if (!terminalLatch.await(terminalTimeoutMs, TimeUnit.MILLISECONDS)) {
+ Transitions.transportFailed(new TimeoutException(
+ s"timed out waiting for stream terminator after ${terminalTimeoutMs}ms"))
+ }
+ } catch {
+ case _: InterruptedException => handleInterrupt("waiting for stream terminator")
+ }
+ }
+
+ /**
+ * Handles an [[InterruptedException]] observed while blocked on a worker event
+ * (init, terminal drain, close, or the result-iterator poll). An interrupt is
+ * an engine-side event -- typically a cancelled query / killed task -- not a
+ * worker fault, so we cancel cooperatively and keep the worker salvageable only
+ * when it acks:
+ *
+ * 1. Send a best-effort `Cancel` (idempotent across call sites).
+ * 2. Wait up to [[interruptCancelTimeoutMs]] -- short, so the interrupted
+ * thread unwinds promptly -- for the worker's `CancelResponse`. A healthy
+ * worker acks in that window and the callback settles a clean `Cancelled`
+ * terminal (worker salvageable, response drained).
+ * 3. If the ack does not arrive in time, settle `Interrupted`. Without a
+ * terminator or liveness proof the worker is not salvageable and must not
+ * be returned to a reuse pool.
+ *
+ * The bounded wait deliberately runs '''before''' the thread's interrupt flag
+ * is restored: [[InterruptedException]] clears the flag on throw, so re-setting
+ * it first would make the wait below throw immediately and defeat the drain.
+ * The flag is restored at the end so the caller's unwind still observes the
+ * interrupt.
+ *
+ * TODO [SPARK-57640]: add a worker liveness/heartbeat check so a future version
+ * may prove that an interrupted worker which missed the short ack window is
+ * nevertheless safe to recycle.
+ */
+ private def handleInterrupt(waitContext: String): Unit = {
+ // NB: flag is currently clear (InterruptedException cleared it); do not
+ // restore it until after the bounded drain below.
+ if (!currentState.isTerminal) {
+ sendCancelInternal(() => cancelWithReason(s"interrupted while $waitContext"))
+ val acked = try {
+ terminalLatch.await(interruptCancelTimeoutMs, TimeUnit.MILLISECONDS)
+ } catch {
+ case _: InterruptedException =>
+ // A second interrupt during the drain: give up the wait immediately.
+ false
+ }
+ if (!acked && !currentState.isTerminal) {
+ // No CancelResponse in time: settle Interrupted rather than
+ // TransportFailed, and stop a later close() from blocking on a terminator
+ // that will not be drained on this thread's behalf. Interrupted remains
+ // unsalvageable until a separate liveness check can prove otherwise.
+ Transitions.interrupted(
+ new InterruptedException(s"interrupted while $waitContext"))
+ }
+ }
+ Thread.currentThread().interrupt()
+ }
+
+ // ---- 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 settle the state) and with any thread that finalizes via
+ * [[close]]. It drives the `Streaming -> Finishing` transition (sending
+ * `Finish` once input is exhausted, built lazily from the `finish` thunk).
+ */
+ private class ProcessIterator(input: Iterator[DataRequest], finish: () => Finish)
+ extends Iterator[DataResponse] {
+
+ // 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 branches, and block branch 4 for the
+ // full terminalTimeoutMs before returning. Callers that probe hasNext
+ // an extra time (iterator.size, instrumentation wrappers) would hang.
+ // Iterator-local and only touched by the single engine thread.
+ private val exhausted = new AtomicBoolean(false)
+ @volatile private var prefetched: DataResponse = _
+
+ /**
+ * Runs a caller-supplied thunk (`input.hasNext`, `input.next()`, or the
+ * `finish` builder) that may throw, and on failure Cancels the in-flight
+ * stream before rethrowing. Without the Cancel a failure would leave the
+ * worker awaiting more input with no terminator owed on the wire; with it the
+ * worker tears down and the exception still propagates to the engine. Setting
+ * `cancelRequested` also suppresses any further Data/Finish this loop might
+ * otherwise attempt (see [[sendRequest]]).
+ */
+ private def cancelOnThrow[T](reason: String)(op: => T): T =
+ try op catch {
+ case NonFatal(e) =>
+ sendCancelInternal(() => cancelWithReason(reason))
Review Comment:
THank you, fixed
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57640] Implement gRPC WorkerSession and test it against an echo server. [spark]
Posted by "haiyangsun-db (via GitHub)" <gi...@apache.org>.
haiyangsun-db commented on code in PR #56702:
URL: https://github.com/apache/spark/pull/56702#discussion_r3665892025
##########
udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/GrpcWorkerSession.scala:
##########
@@ -0,0 +1,933 @@
+/*
+ * 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.grpc
+
+import java.util.concurrent.{CountDownLatch, LinkedBlockingQueue, TimeoutException, TimeUnit}
+import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference}
+
+import scala.util.control.NonFatal
+
+import io.grpc.{ConnectivityState, ManagedChannel}
+import io.grpc.stub.StreamObserver
+
+import org.apache.spark.annotation.Experimental
+import org.apache.spark.udf.worker.{Cancel, CancelResponse, DataRequest, DataResponse,
+ ExecutionError, Finish, FinishResponse, Init, InitResponse, UdfControlRequest,
+ UdfControlResponse, UdfRequest, UdfResponse, UdfWorkerGrpc}
+import org.apache.spark.udf.worker.core.{Termination, WorkerHandle, WorkerLogger, WorkerSession}
+import org.apache.spark.udf.worker.core.WorkerSession.SessionState
+import org.apache.spark.udf.worker.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_message.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.
+ *
+ * '''Driving model.''' Consumption-driven (Volcano / pull): the thread that
+ * consumes the [[doProcess]] result iterator is the one that pulls input and
+ * sends each `DataRequest`; the gRPC callback thread only receives output. It is
+ * pull-driven but not one-input-per-output -- `advance` sends the next input
+ * whenever the output queue is momentarily empty, so under async delivery it may
+ * push several input batches before any output is read (bounded by HTTP/2 flow
+ * control).
+ *
+ * '''State machine.''' This class does not keep its own state machine: it
+ * drives the single [[WorkerSession.SessionState]] owned by the base. The base
+ * advances `Created -> Initializing` (in `init`) and `Initialized -> Streaming`
+ * (in `process`); this class advances the protocol-event edges through
+ * [[compareAndSetState]] / [[completeTerminal]] as it exchanges messages:
+ * {{{
+ * Initializing --(InitResponse ok)--> Initialized [handleControl]
+ * Streaming ----(Finish written)----> Finishing [ProcessIterator]
+ * <any non-terminal> --(Cancel written)--> Cancelling [sendCancelInternal]
+ * <any non-terminal> --(terminator/error)--> terminal [completeTerminal]
+ * }}}
+ * The two clean terminals carry the worker's `FinishResponse` / `CancelResponse`
+ * (metrics + finish/cancel callback `data`/`error`) so [[close]] can return
+ * them. The only flag kept outside the machine is [[cancelRequested]]: a
+ * cancellation can be requested before the stream exists (so it cannot be a
+ * state transition yet), and it must both fast-fail the result iterator and
+ * suppress any in-flight Data/Finish.
+ *
+ * Threading:
+ * - [[doInit]] is synchronous: sends `Init` and blocks on `InitResponse`,
+ * returning it.
+ * - [[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.
+ * - [[doClose]] is thread-safe and idempotent: it settles + returns the
+ * terminator (cancelling in-flight work if the stream had not finished)
+ * and half-closes the request side.
+ *
+ * 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 a gRPC channel built and owned by the caller (the
+ * dispatcher). 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 from the worker, drained by the process() iterator.
+ // Intentionally unbounded: a bounded queue would block the gRPC callback
+ // (Netty event-loop) thread when full, stalling terminator/control delivery on
+ // the whole channel. HTTP/2 flow control bounds the wire and the consumer
+ // normally drains promptly; a stalled downstream can still grow it, but the fix
+ // is protocol-level back-pressure (out of scope), not bounding the queue.
+ //
+ // TODO [SPARK-57324]: expose queue depth as a metric (early warning for a
+ // stalled consumer).
+ private val outputQueue = new LinkedBlockingQueue[QueueItem]()
+
+ // The worker's `InitResponse` (success or error) coupled with the latch that
+ // init() blocks on. Fires when the InitResponse arrives (`complete`) or when a
+ // terminal settles first without one -- transport error, half-close, or a
+ // premature terminator (`signalWithoutValue`). Until it fires we have no proof
+ // the worker accepted the session.
+ //
+ // Settle-before-release rule (referenced from every callback that both settles
+ // a terminal and fires this latch): settle the terminal FIRST, then complete /
+ // signal. init() blocks on the latch, and only latch await/release establish a
+ // happens-before edge, so a woken init() is guaranteed to observe the terminal
+ // rather than a transient state. OneShotValue keeps that publish-then-release
+ // in one place instead of every caller remembering to count down after setting
+ // the reference.
+ private val initValue = new OneShotValue[InitResponse]
+
+ // Fired when the session reaches a terminal [[SessionState]]. doClose() and
+ // the init-error path block on this to drain the terminator.
+ private val terminalLatch = new CountDownLatch(1)
+
+ // Captures an ErrorResponse encountered during the data phase so that
+ // the CancelResponse terminator 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)
+
+ // Cancellation INTENT -- distinct from the `Cancelling` state, which is reached
+ // only once a Cancel is actually written to the wire ([[sendCancelInternal]]).
+ // Intent is set first and can outrun (or never reach) that write, so
+ // `cancelRequested` does NOT imply state `Cancelling`. Kept outside the
+ // [[SessionState]] machine because a cancel can be requested before the stream
+ // exists (pre-init), where there is no wire transition to make yet. Used to (a)
+ // make cancel idempotent across all call sites, (b) fast-fail ProcessIterator
+ // on a pre-init cancel, and (c) suppress any Data/Finish that would otherwise
+ // race a Cancel onto the wire (re-read inside [[requestLock]]).
+ private val cancelRequested = new AtomicBoolean(false)
+
+ // gRPC requires serialized writes to a request StreamObserver.
+ private val requestLock = new Object
+
+ // Initialised in init() -- before that, close() is a no-op 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 =>
+ // A DataResponse before InitResponse violates the protocol (InitResponse
+ // must precede any DataResponse): fast-fail rather than enqueue it and let
+ // init() block to its timeout. Settle-before-release (see initValue).
+ if (!initResolved) {
+ Transitions.transportFailed(new IllegalStateException(
+ "worker sent a DataResponse before InitResponse"))
+ initValue.signalWithoutValue()
+ } else {
+ outputQueue.put(QueueItem.Batch(response.getData))
+ }
+
+ case UdfResponse.ResponseCase.CONTROL =>
+ handleControl(response.getControl)
+
+ case other =>
+ // A malformed response (empty / unknown oneof) is a terminal transport
+ // failure; fast-fail init the same way. Settle-before-release (see initValue).
+ Transitions.transportFailed(new IllegalStateException(
+ s"unexpected response oneof: $other"))
+ initValue.signalWithoutValue()
+ }
+ }
+
+ override def onError(t: Throwable): Unit = {
+ // Transport-level failure: the stream is dead, no further writes possible.
+ // Settle-before-release (see initValue) so init() surfaces the transport
+ // cause instead of the initResponseTimeoutMs "timed out" error.
+ Transitions.transportFailed(t)
+ initValue.signalWithoutValue()
+ }
+
+ 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 (!currentState.isTerminal) {
+ Transitions.transportFailed(new IllegalStateException(
+ "worker response stream closed without a terminator"))
+ }
+ // Defensive: if onCompleted reached us before InitResponse, doInit is
+ // still blocked on initValue and would otherwise time out.
+ initValue.signalWithoutValue()
+ }
+ }
+
+ /**
+ * Wakes the result iterator (blocked on [[outputQueue]]) and any thread
+ * waiting on [[terminalLatch]] when the base settles a terminal. Invoked once,
+ * by the caller that wins [[completeTerminal]].
+ */
+ override protected def onTerminalSettled(termination: Termination): Unit = {
+ outputQueue.put(QueueItem.EndOfStream)
+ terminalLatch.countDown()
+ }
+
+ private def handleControl(ctrl: UdfControlResponse): Unit = ctrl.getControlCase match {
+ case UdfControlResponse.ControlCase.INIT =>
+ val resp = ctrl.getInit
+ if (resp.hasError) {
+ // Attribute a subsequent close()'s terminator to the init error rather
+ // than a bare Cancelled, mirroring the data-phase ERROR branch.
+ executionError.compareAndSet(None, Some(resp.getError))
+ // Settle Failed(err) before the Cancel below (not after its CancelResponse):
+ // the sticky terminal makes close() report the init error, and the Cancel
+ // becomes best-effort worker cleanup -- so if it can't be written (a
+ // directExecutor worker delivered this reentrantly, before doInit published
+ // requestObserver), doInit still doesn't wait terminalTimeoutMs for a
+ // CancelResponse that can never arrive. Settle-before-release (see initValue).
+ Transitions.failed(resp.getError)
+ initValue.complete(resp)
+ sendCancelInternal(() => cancelWithReason("init failed"))
+ } else {
+ // InitResponse OK. Only advance from Initializing so a terminal that raced
+ // in (e.g. a transport error) still wins; process() then opens the data
+ // phase. Settle-before-release (see initValue): publish after the CAS.
+ Transitions.initAccepted()
+ initValue.complete(resp)
+ }
+
+ case UdfControlResponse.ControlCase.ERROR =>
+ val err = ctrl.getError.getError
+ executionError.compareAndSet(None, Some(err))
+ // Before init resolves, settle Failed here rather than rely on the Cancel
+ // below: sendCancelInternal may find requestObserver still null (reentrant
+ // delivery inside stream.onNext, before doInit published it), so no Cancel --
+ // hence no CancelResponse -- would settle the terminal, and doInit would
+ // return as if init succeeded. completeTerminal is idempotent, so the
+ // data-phase ERROR -> Cancel -> CancelResponse path is unaffected.
+ if (!initResolved) {
+ Transitions.failed(err)
+ }
+ // Settle-before-release (see initValue). No value to publish -- ERROR is not
+ // an InitResponse; on the data-phase path the latch already fired in init.
+ initValue.signalWithoutValue()
+ sendCancelInternal(() => cancelWithReason("aborting after ErrorResponse"))
+
+ case UdfControlResponse.ControlCase.FINISH =>
+ // The FinishResponse carries metrics + the finish-callback data/error.
+ // Keep it on the terminal so close() can return it; the iterator inspects
+ // its error field to decide whether to throw.
+ Transitions.finished(ctrl.getFinish)
+ // Defensive: FINISH before InitResponse is a worker protocol bug, but
+ // we should fail init fast rather than hang the 30s init timeout.
+ initValue.signalWithoutValue()
+
+ case UdfControlResponse.ControlCase.CANCEL =>
+ // The CancelResponse carries metrics + the cancel-callback error. Keep it
+ // on the terminal so close() can return it; any prior ErrorResponse is
+ // tracked in executionError and surfaced by the iterator.
+ Transitions.cancelled(ctrl.getCancel)
+ // Defensive: CANCEL before InitResponse unblocks doInit so it can
+ // surface the cancellation instead of timing out.
+ initValue.signalWithoutValue()
+
+ case UdfControlResponse.ControlCase.CONTROL_NOT_SET =>
+ Transitions.transportFailed(new IllegalStateException(
+ "empty UdfControlResponse oneof"))
+ initValue.signalWithoutValue()
+ }
+
+ /**
+ * True once init is no longer pending -- i.e. the stream is past `Initializing`.
+ * Not "init succeeded": a terminal (including a failure) also counts as resolved.
+ */
+ private def initResolved: Boolean = currentState match {
+ case SessionState.Created | SessionState.Initializing => false
+ case _ => true
+ }
+
+ private def cancelWithReason(reason: String): Cancel =
+ Cancel.newBuilder().setReason(reason).build()
+
+ /**
+ * The protocol transition graph in one place: names for the edges, not a
+ * second source of truth. Every edge acts on the single
+ * [[WorkerSession.SessionState]] owned by the base via [[compareAndSetState]]
+ * (non-terminal) or [[completeTerminal]] (terminal), so the base's CAS is the
+ * only synchronization and a terminal that arrived first always wins (the
+ * non-terminal CASes fail against it; [[completeTerminal]] is first-wins).
+ *
+ * Edges this class drives -- edge, method, then driver site(s) / thread (the
+ * base drives the API-call edges: `Created -> Initializing` in `init`,
+ * `Initialized -> Streaming` in `process`):
+ * {{{
+ * Initializing -> Initialized initAccepted handleControl INIT-ok [gRPC cb]
+ * Streaming -> Finishing finishSent advance branch 3 [engine]
+ * non-terminal -> Cancelling cancelSentFrom sendCancelInternal, iff a Cancel
+ * reaches the wire [gRPC cb | engine | init | close]
+ * non-terminal -> Terminal finished/cancelled/failed/transportFailed
+ * handleControl / doInit / doClose / advance / onError
+ * }}}
+ * `Cancelling` is reached only if a Cancel is actually written; a pre-stream or
+ * raced cancel goes straight to a `Cancelled`/`TransportFailed` terminal (see
+ * [[sendCancelInternal]], [[doClose]], `ProcessIterator`) or nowhere -- so
+ * `cancelRequested` (intent) does not imply state `Cancelling`.
+ */
+ private object Transitions {
+ /** `InitResponse` OK: `Initializing -> Initialized`. */
+ def initAccepted(): Boolean =
+ compareAndSetState(SessionState.Initializing, SessionState.Initialized)
+
+ /** Input exhausted and `Finish` written: `Streaming -> Finishing` (once). */
+ def finishSent(): Boolean =
+ compareAndSetState(SessionState.Streaming, SessionState.Finishing)
+
+ /** `Cancel` written: `cur -> Cancelling`, from any non-terminal `cur`. */
+ def cancelSentFrom(cur: SessionState): Boolean =
+ !cur.isTerminal && compareAndSetState(cur, SessionState.Cancelling)
+
+ /** Clean terminal carrying the worker's `FinishResponse`. */
+ def finished(response: FinishResponse): Boolean =
+ completeTerminal(Termination.Finished(response))
+
+ /** Clean terminal carrying the worker's `CancelResponse`. */
+ def cancelled(response: CancelResponse): Boolean =
+ completeTerminal(Termination.Cancelled(response))
+
+ /** Failure terminal carrying a structured [[ExecutionError]]. */
+ def failed(error: ExecutionError): Boolean =
+ completeTerminal(Termination.Failed(error))
+
+ /** Failure terminal carrying a transport-level cause. */
+ def transportFailed(cause: Throwable): Boolean =
+ completeTerminal(Termination.TransportFailed(cause))
Review Comment:
There's no Transitions.failed anymore - transportFailed and the terminal outcomes are now distinct by design
##########
udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/GrpcWorkerSession.scala:
##########
@@ -0,0 +1,933 @@
+/*
+ * 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.grpc
+
+import java.util.concurrent.{CountDownLatch, LinkedBlockingQueue, TimeoutException, TimeUnit}
+import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference}
+
+import scala.util.control.NonFatal
+
+import io.grpc.{ConnectivityState, ManagedChannel}
+import io.grpc.stub.StreamObserver
+
+import org.apache.spark.annotation.Experimental
+import org.apache.spark.udf.worker.{Cancel, CancelResponse, DataRequest, DataResponse,
+ ExecutionError, Finish, FinishResponse, Init, InitResponse, UdfControlRequest,
+ UdfControlResponse, UdfRequest, UdfResponse, UdfWorkerGrpc}
+import org.apache.spark.udf.worker.core.{Termination, WorkerHandle, WorkerLogger, WorkerSession}
+import org.apache.spark.udf.worker.core.WorkerSession.SessionState
+import org.apache.spark.udf.worker.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_message.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.
+ *
+ * '''Driving model.''' Consumption-driven (Volcano / pull): the thread that
+ * consumes the [[doProcess]] result iterator is the one that pulls input and
+ * sends each `DataRequest`; the gRPC callback thread only receives output. It is
+ * pull-driven but not one-input-per-output -- `advance` sends the next input
+ * whenever the output queue is momentarily empty, so under async delivery it may
+ * push several input batches before any output is read (bounded by HTTP/2 flow
+ * control).
+ *
+ * '''State machine.''' This class does not keep its own state machine: it
+ * drives the single [[WorkerSession.SessionState]] owned by the base. The base
+ * advances `Created -> Initializing` (in `init`) and `Initialized -> Streaming`
+ * (in `process`); this class advances the protocol-event edges through
+ * [[compareAndSetState]] / [[completeTerminal]] as it exchanges messages:
+ * {{{
+ * Initializing --(InitResponse ok)--> Initialized [handleControl]
+ * Streaming ----(Finish written)----> Finishing [ProcessIterator]
+ * <any non-terminal> --(Cancel written)--> Cancelling [sendCancelInternal]
+ * <any non-terminal> --(terminator/error)--> terminal [completeTerminal]
+ * }}}
+ * The two clean terminals carry the worker's `FinishResponse` / `CancelResponse`
+ * (metrics + finish/cancel callback `data`/`error`) so [[close]] can return
+ * them. The only flag kept outside the machine is [[cancelRequested]]: a
+ * cancellation can be requested before the stream exists (so it cannot be a
+ * state transition yet), and it must both fast-fail the result iterator and
+ * suppress any in-flight Data/Finish.
+ *
+ * Threading:
+ * - [[doInit]] is synchronous: sends `Init` and blocks on `InitResponse`,
+ * returning it.
+ * - [[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.
+ * - [[doClose]] is thread-safe and idempotent: it settles + returns the
+ * terminator (cancelling in-flight work if the stream had not finished)
+ * and half-closes the request side.
+ *
+ * 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 a gRPC channel built and owned by the caller (the
+ * dispatcher). 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 from the worker, drained by the process() iterator.
+ // Intentionally unbounded: a bounded queue would block the gRPC callback
+ // (Netty event-loop) thread when full, stalling terminator/control delivery on
+ // the whole channel. HTTP/2 flow control bounds the wire and the consumer
+ // normally drains promptly; a stalled downstream can still grow it, but the fix
+ // is protocol-level back-pressure (out of scope), not bounding the queue.
+ //
+ // TODO [SPARK-57324]: expose queue depth as a metric (early warning for a
+ // stalled consumer).
+ private val outputQueue = new LinkedBlockingQueue[QueueItem]()
+
+ // The worker's `InitResponse` (success or error) coupled with the latch that
+ // init() blocks on. Fires when the InitResponse arrives (`complete`) or when a
+ // terminal settles first without one -- transport error, half-close, or a
+ // premature terminator (`signalWithoutValue`). Until it fires we have no proof
+ // the worker accepted the session.
+ //
+ // Settle-before-release rule (referenced from every callback that both settles
+ // a terminal and fires this latch): settle the terminal FIRST, then complete /
+ // signal. init() blocks on the latch, and only latch await/release establish a
+ // happens-before edge, so a woken init() is guaranteed to observe the terminal
+ // rather than a transient state. OneShotValue keeps that publish-then-release
+ // in one place instead of every caller remembering to count down after setting
+ // the reference.
+ private val initValue = new OneShotValue[InitResponse]
+
+ // Fired when the session reaches a terminal [[SessionState]]. doClose() and
+ // the init-error path block on this to drain the terminator.
+ private val terminalLatch = new CountDownLatch(1)
+
+ // Captures an ErrorResponse encountered during the data phase so that
+ // the CancelResponse terminator 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)
+
+ // Cancellation INTENT -- distinct from the `Cancelling` state, which is reached
+ // only once a Cancel is actually written to the wire ([[sendCancelInternal]]).
+ // Intent is set first and can outrun (or never reach) that write, so
+ // `cancelRequested` does NOT imply state `Cancelling`. Kept outside the
+ // [[SessionState]] machine because a cancel can be requested before the stream
+ // exists (pre-init), where there is no wire transition to make yet. Used to (a)
+ // make cancel idempotent across all call sites, (b) fast-fail ProcessIterator
+ // on a pre-init cancel, and (c) suppress any Data/Finish that would otherwise
+ // race a Cancel onto the wire (re-read inside [[requestLock]]).
+ private val cancelRequested = new AtomicBoolean(false)
+
+ // gRPC requires serialized writes to a request StreamObserver.
+ private val requestLock = new Object
+
+ // Initialised in init() -- before that, close() is a no-op 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 =>
+ // A DataResponse before InitResponse violates the protocol (InitResponse
+ // must precede any DataResponse): fast-fail rather than enqueue it and let
+ // init() block to its timeout. Settle-before-release (see initValue).
+ if (!initResolved) {
+ Transitions.transportFailed(new IllegalStateException(
+ "worker sent a DataResponse before InitResponse"))
+ initValue.signalWithoutValue()
+ } else {
+ outputQueue.put(QueueItem.Batch(response.getData))
+ }
+
+ case UdfResponse.ResponseCase.CONTROL =>
+ handleControl(response.getControl)
+
+ case other =>
+ // A malformed response (empty / unknown oneof) is a terminal transport
+ // failure; fast-fail init the same way. Settle-before-release (see initValue).
+ Transitions.transportFailed(new IllegalStateException(
+ s"unexpected response oneof: $other"))
+ initValue.signalWithoutValue()
+ }
+ }
+
+ override def onError(t: Throwable): Unit = {
+ // Transport-level failure: the stream is dead, no further writes possible.
+ // Settle-before-release (see initValue) so init() surfaces the transport
+ // cause instead of the initResponseTimeoutMs "timed out" error.
+ Transitions.transportFailed(t)
+ initValue.signalWithoutValue()
+ }
+
+ 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 (!currentState.isTerminal) {
+ Transitions.transportFailed(new IllegalStateException(
+ "worker response stream closed without a terminator"))
+ }
+ // Defensive: if onCompleted reached us before InitResponse, doInit is
+ // still blocked on initValue and would otherwise time out.
+ initValue.signalWithoutValue()
+ }
+ }
+
+ /**
+ * Wakes the result iterator (blocked on [[outputQueue]]) and any thread
+ * waiting on [[terminalLatch]] when the base settles a terminal. Invoked once,
+ * by the caller that wins [[completeTerminal]].
+ */
+ override protected def onTerminalSettled(termination: Termination): Unit = {
+ outputQueue.put(QueueItem.EndOfStream)
+ terminalLatch.countDown()
+ }
+
+ private def handleControl(ctrl: UdfControlResponse): Unit = ctrl.getControlCase match {
+ case UdfControlResponse.ControlCase.INIT =>
+ val resp = ctrl.getInit
+ if (resp.hasError) {
+ // Attribute a subsequent close()'s terminator to the init error rather
+ // than a bare Cancelled, mirroring the data-phase ERROR branch.
+ executionError.compareAndSet(None, Some(resp.getError))
+ // Settle Failed(err) before the Cancel below (not after its CancelResponse):
+ // the sticky terminal makes close() report the init error, and the Cancel
+ // becomes best-effort worker cleanup -- so if it can't be written (a
+ // directExecutor worker delivered this reentrantly, before doInit published
+ // requestObserver), doInit still doesn't wait terminalTimeoutMs for a
+ // CancelResponse that can never arrive. Settle-before-release (see initValue).
+ Transitions.failed(resp.getError)
+ initValue.complete(resp)
+ sendCancelInternal(() => cancelWithReason("init failed"))
+ } else {
+ // InitResponse OK. Only advance from Initializing so a terminal that raced
+ // in (e.g. a transport error) still wins; process() then opens the data
+ // phase. Settle-before-release (see initValue): publish after the CAS.
+ Transitions.initAccepted()
+ initValue.complete(resp)
+ }
+
+ case UdfControlResponse.ControlCase.ERROR =>
+ val err = ctrl.getError.getError
+ executionError.compareAndSet(None, Some(err))
+ // Before init resolves, settle Failed here rather than rely on the Cancel
+ // below: sendCancelInternal may find requestObserver still null (reentrant
+ // delivery inside stream.onNext, before doInit published it), so no Cancel --
+ // hence no CancelResponse -- would settle the terminal, and doInit would
+ // return as if init succeeded. completeTerminal is idempotent, so the
+ // data-phase ERROR -> Cancel -> CancelResponse path is unaffected.
+ if (!initResolved) {
+ Transitions.failed(err)
+ }
+ // Settle-before-release (see initValue). No value to publish -- ERROR is not
+ // an InitResponse; on the data-phase path the latch already fired in init.
+ initValue.signalWithoutValue()
+ sendCancelInternal(() => cancelWithReason("aborting after ErrorResponse"))
+
+ case UdfControlResponse.ControlCase.FINISH =>
+ // The FinishResponse carries metrics + the finish-callback data/error.
+ // Keep it on the terminal so close() can return it; the iterator inspects
+ // its error field to decide whether to throw.
+ Transitions.finished(ctrl.getFinish)
+ // Defensive: FINISH before InitResponse is a worker protocol bug, but
+ // we should fail init fast rather than hang the 30s init timeout.
+ initValue.signalWithoutValue()
+
+ case UdfControlResponse.ControlCase.CANCEL =>
+ // The CancelResponse carries metrics + the cancel-callback error. Keep it
+ // on the terminal so close() can return it; any prior ErrorResponse is
+ // tracked in executionError and surfaced by the iterator.
+ Transitions.cancelled(ctrl.getCancel)
+ // Defensive: CANCEL before InitResponse unblocks doInit so it can
+ // surface the cancellation instead of timing out.
+ initValue.signalWithoutValue()
+
+ case UdfControlResponse.ControlCase.CONTROL_NOT_SET =>
+ Transitions.transportFailed(new IllegalStateException(
+ "empty UdfControlResponse oneof"))
+ initValue.signalWithoutValue()
+ }
+
+ /**
+ * True once init is no longer pending -- i.e. the stream is past `Initializing`.
+ * Not "init succeeded": a terminal (including a failure) also counts as resolved.
+ */
+ private def initResolved: Boolean = currentState match {
+ case SessionState.Created | SessionState.Initializing => false
+ case _ => true
+ }
+
+ private def cancelWithReason(reason: String): Cancel =
+ Cancel.newBuilder().setReason(reason).build()
+
+ /**
+ * The protocol transition graph in one place: names for the edges, not a
+ * second source of truth. Every edge acts on the single
+ * [[WorkerSession.SessionState]] owned by the base via [[compareAndSetState]]
+ * (non-terminal) or [[completeTerminal]] (terminal), so the base's CAS is the
+ * only synchronization and a terminal that arrived first always wins (the
+ * non-terminal CASes fail against it; [[completeTerminal]] is first-wins).
+ *
+ * Edges this class drives -- edge, method, then driver site(s) / thread (the
+ * base drives the API-call edges: `Created -> Initializing` in `init`,
+ * `Initialized -> Streaming` in `process`):
+ * {{{
+ * Initializing -> Initialized initAccepted handleControl INIT-ok [gRPC cb]
+ * Streaming -> Finishing finishSent advance branch 3 [engine]
+ * non-terminal -> Cancelling cancelSentFrom sendCancelInternal, iff a Cancel
+ * reaches the wire [gRPC cb | engine | init | close]
+ * non-terminal -> Terminal finished/cancelled/failed/transportFailed
+ * handleControl / doInit / doClose / advance / onError
+ * }}}
+ * `Cancelling` is reached only if a Cancel is actually written; a pre-stream or
+ * raced cancel goes straight to a `Cancelled`/`TransportFailed` terminal (see
+ * [[sendCancelInternal]], [[doClose]], `ProcessIterator`) or nowhere -- so
+ * `cancelRequested` (intent) does not imply state `Cancelling`.
+ */
+ private object Transitions {
+ /** `InitResponse` OK: `Initializing -> Initialized`. */
+ def initAccepted(): Boolean =
+ compareAndSetState(SessionState.Initializing, SessionState.Initialized)
+
+ /** Input exhausted and `Finish` written: `Streaming -> Finishing` (once). */
+ def finishSent(): Boolean =
+ compareAndSetState(SessionState.Streaming, SessionState.Finishing)
+
+ /** `Cancel` written: `cur -> Cancelling`, from any non-terminal `cur`. */
+ def cancelSentFrom(cur: SessionState): Boolean =
+ !cur.isTerminal && compareAndSetState(cur, SessionState.Cancelling)
+
+ /** Clean terminal carrying the worker's `FinishResponse`. */
+ def finished(response: FinishResponse): Boolean =
+ completeTerminal(Termination.Finished(response))
+
+ /** Clean terminal carrying the worker's `CancelResponse`. */
+ def cancelled(response: CancelResponse): Boolean =
+ completeTerminal(Termination.Cancelled(response))
+
+ /** Failure terminal carrying a structured [[ExecutionError]]. */
+ def failed(error: ExecutionError): Boolean =
+ completeTerminal(Termination.Failed(error))
+
+ /** Failure terminal carrying a transport-level cause. */
+ def transportFailed(cause: Throwable): Boolean =
+ completeTerminal(Termination.TransportFailed(cause))
+ }
+
+ // ---- WorkerSession hooks ------------------------------------------------
+
+ override protected def doInit(message: Init): InitResponse = {
+ // 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")
+ Transitions.transportFailed(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. close()'s 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 {
+ // Reentrant delivery: a directExecutor worker (one whose gRPC callbacks run
+ // synchronously on the caller's thread) can deliver InitResponse -- and thus
+ // run responseObserver/handleControl -- from *inside* this stream.onNext,
+ // before requestObserver below is published. The base advanced
+ // Created -> Initializing before calling doInit, so that reentrant
+ // InitResponse still finds Initializing and advances to Initialized; the
+ // handleControl error paths guard the still-null requestObserver. This is
+ // the scenario the "before doInit published requestObserver" comments mean.
+ 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
+ Transitions.transportFailed(e)
+ // Surface as GrpcWorkerSessionException so the engine integration layer
+ // (which catches that type and wraps it) sees a uniform init-failure
+ // exception rather than the raw transport error.
+ throw new GrpcWorkerSessionException("UDF worker stream failed during init", e)
+ }
+
+ val responded = try {
+ initValue.await(initResponseTimeoutMs)
+ } catch {
+ case _: InterruptedException =>
+ Thread.currentThread().interrupt()
+ sendCancelInternal(() => cancelWithReason("interrupted during init"))
+ // Make sure close() does not block waiting for a terminator that
+ // will never arrive on this thread's behalf.
+ Transitions.transportFailed(
+ new InterruptedException("interrupted while waiting for InitResponse"))
+ throw new InterruptedException("interrupted while waiting for InitResponse")
+ }
+
+ if (!responded) {
+ sendCancelInternal(() => cancelWithReason("InitResponse timed out"))
+ val timeout = new TimeoutException(
+ s"timed out waiting for InitResponse after ${initResponseTimeoutMs}ms")
+ // Settle the terminal so a subsequent close() does not stall for a
+ // second `terminalTimeoutMs` waiting for a worker that already missed
+ // its init deadline.
+ Transitions.transportFailed(timeout)
+ // Surface as GrpcWorkerSessionException (carrying the timeout cause) so
+ // the engine integration layer that catches that type can wrap it.
+ throw new GrpcWorkerSessionException(
+ s"timed out waiting for InitResponse after ${initResponseTimeoutMs}ms", timeout)
+ }
+
+ initValue.get match {
+ case Some(resp) if resp.hasError =>
+ // 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(resp.getError)}", resp.getError)
+ case Some(resp) =>
+ resp
+ case None =>
+ // No InitResponse arrived but the latch fired. The defensive
+ // initValue.signalWithoutValue() 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.
+ currentState match {
+ case SessionState.Terminal(Termination.TransportFailed(cause)) =>
+ throw new GrpcWorkerSessionException(
+ "UDF worker stream failed during init", cause)
+ case SessionState.Terminal(Termination.Failed(err)) =>
+ throw new GrpcWorkerSessionException(
+ s"UDF worker reported an error before init completed: " +
+ describeError(err), err)
+ case SessionState.Terminal(Termination.Cancelled(_)) =>
+ throw new GrpcWorkerSessionException(
+ "UDF worker stream was cancelled before init completed")
+ case SessionState.Terminal(Termination.Finished(_)) =>
+ throw new GrpcWorkerSessionException(
+ "UDF worker finished before init completed")
+ case other =>
+ throw new IllegalStateException(
+ s"init latch fired without an InitResponse or terminal: $other")
+ }
+ }
+ }
+
+ override protected def doProcess(
+ input: Iterator[DataRequest],
+ finish: () => Finish): Iterator[DataResponse] = {
+ // Init success is guaranteed by the base [[WorkerSession]] lifecycle: if
+ // doInit had failed it would have thrown and process() would never run.
+ new ProcessIterator(input, finish)
+ }
+
+ override protected def doClose(cancel: () => Cancel): Termination = {
+ if (requestObserver == null) {
+ // init() never put a stream on the wire (closed before/around init, or
+ // init threw before publishing). There is no protocol terminator; treat
+ // the session as cancelled-before-start. The base WorkerSession still
+ // releases the worker handle, so the worker is torn down.
+ Transitions.cancelled(CancelResponse.getDefaultInstance)
+ return Termination.Cancelled(CancelResponse.getDefaultInstance)
Review Comment:
There's no Transitions.failed anymore - transportFailed and the terminal outcomes are now distinct by design
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57640] Implement gRPC WorkerSession and test it against an echo server. [spark]
Posted by "haiyangsun-db (via GitHub)" <gi...@apache.org>.
haiyangsun-db commented on code in PR #56702:
URL: https://github.com/apache/spark/pull/56702#discussion_r3665872582
##########
udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/GrpcWorkerSession.scala:
##########
@@ -0,0 +1,933 @@
+/*
+ * 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.grpc
+
+import java.util.concurrent.{CountDownLatch, LinkedBlockingQueue, TimeoutException, TimeUnit}
+import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference}
+
+import scala.util.control.NonFatal
+
+import io.grpc.{ConnectivityState, ManagedChannel}
+import io.grpc.stub.StreamObserver
+
+import org.apache.spark.annotation.Experimental
+import org.apache.spark.udf.worker.{Cancel, CancelResponse, DataRequest, DataResponse,
+ ExecutionError, Finish, FinishResponse, Init, InitResponse, UdfControlRequest,
+ UdfControlResponse, UdfRequest, UdfResponse, UdfWorkerGrpc}
+import org.apache.spark.udf.worker.core.{Termination, WorkerHandle, WorkerLogger, WorkerSession}
+import org.apache.spark.udf.worker.core.WorkerSession.SessionState
+import org.apache.spark.udf.worker.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_message.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.
+ *
+ * '''Driving model.''' Consumption-driven (Volcano / pull): the thread that
+ * consumes the [[doProcess]] result iterator is the one that pulls input and
+ * sends each `DataRequest`; the gRPC callback thread only receives output. It is
+ * pull-driven but not one-input-per-output -- `advance` sends the next input
+ * whenever the output queue is momentarily empty, so under async delivery it may
+ * push several input batches before any output is read (bounded by HTTP/2 flow
+ * control).
+ *
+ * '''State machine.''' This class does not keep its own state machine: it
+ * drives the single [[WorkerSession.SessionState]] owned by the base. The base
+ * advances `Created -> Initializing` (in `init`) and `Initialized -> Streaming`
+ * (in `process`); this class advances the protocol-event edges through
+ * [[compareAndSetState]] / [[completeTerminal]] as it exchanges messages:
+ * {{{
+ * Initializing --(InitResponse ok)--> Initialized [handleControl]
+ * Streaming ----(Finish written)----> Finishing [ProcessIterator]
+ * <any non-terminal> --(Cancel written)--> Cancelling [sendCancelInternal]
+ * <any non-terminal> --(terminator/error)--> terminal [completeTerminal]
+ * }}}
+ * The two clean terminals carry the worker's `FinishResponse` / `CancelResponse`
+ * (metrics + finish/cancel callback `data`/`error`) so [[close]] can return
+ * them. The only flag kept outside the machine is [[cancelRequested]]: a
+ * cancellation can be requested before the stream exists (so it cannot be a
+ * state transition yet), and it must both fast-fail the result iterator and
+ * suppress any in-flight Data/Finish.
+ *
+ * Threading:
+ * - [[doInit]] is synchronous: sends `Init` and blocks on `InitResponse`,
+ * returning it.
+ * - [[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.
+ * - [[doClose]] is thread-safe and idempotent: it settles + returns the
+ * terminator (cancelling in-flight work if the stream had not finished)
+ * and half-closes the request side.
+ *
+ * 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 a gRPC channel built and owned by the caller (the
+ * dispatcher). 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 from the worker, drained by the process() iterator.
+ // Intentionally unbounded: a bounded queue would block the gRPC callback
+ // (Netty event-loop) thread when full, stalling terminator/control delivery on
+ // the whole channel. HTTP/2 flow control bounds the wire and the consumer
+ // normally drains promptly; a stalled downstream can still grow it, but the fix
+ // is protocol-level back-pressure (out of scope), not bounding the queue.
+ //
+ // TODO [SPARK-57324]: expose queue depth as a metric (early warning for a
+ // stalled consumer).
+ private val outputQueue = new LinkedBlockingQueue[QueueItem]()
+
+ // The worker's `InitResponse` (success or error) coupled with the latch that
+ // init() blocks on. Fires when the InitResponse arrives (`complete`) or when a
+ // terminal settles first without one -- transport error, half-close, or a
+ // premature terminator (`signalWithoutValue`). Until it fires we have no proof
+ // the worker accepted the session.
+ //
+ // Settle-before-release rule (referenced from every callback that both settles
+ // a terminal and fires this latch): settle the terminal FIRST, then complete /
+ // signal. init() blocks on the latch, and only latch await/release establish a
+ // happens-before edge, so a woken init() is guaranteed to observe the terminal
+ // rather than a transient state. OneShotValue keeps that publish-then-release
+ // in one place instead of every caller remembering to count down after setting
+ // the reference.
+ private val initValue = new OneShotValue[InitResponse]
+
+ // Fired when the session reaches a terminal [[SessionState]]. doClose() and
+ // the init-error path block on this to drain the terminator.
+ private val terminalLatch = new CountDownLatch(1)
+
+ // Captures an ErrorResponse encountered during the data phase so that
+ // the CancelResponse terminator 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)
+
+ // Cancellation INTENT -- distinct from the `Cancelling` state, which is reached
+ // only once a Cancel is actually written to the wire ([[sendCancelInternal]]).
+ // Intent is set first and can outrun (or never reach) that write, so
+ // `cancelRequested` does NOT imply state `Cancelling`. Kept outside the
+ // [[SessionState]] machine because a cancel can be requested before the stream
+ // exists (pre-init), where there is no wire transition to make yet. Used to (a)
+ // make cancel idempotent across all call sites, (b) fast-fail ProcessIterator
+ // on a pre-init cancel, and (c) suppress any Data/Finish that would otherwise
+ // race a Cancel onto the wire (re-read inside [[requestLock]]).
+ private val cancelRequested = new AtomicBoolean(false)
+
+ // gRPC requires serialized writes to a request StreamObserver.
+ private val requestLock = new Object
+
+ // Initialised in init() -- before that, close() is a no-op 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 =>
+ // A DataResponse before InitResponse violates the protocol (InitResponse
+ // must precede any DataResponse): fast-fail rather than enqueue it and let
+ // init() block to its timeout. Settle-before-release (see initValue).
+ if (!initResolved) {
+ Transitions.transportFailed(new IllegalStateException(
+ "worker sent a DataResponse before InitResponse"))
+ initValue.signalWithoutValue()
+ } else {
+ outputQueue.put(QueueItem.Batch(response.getData))
+ }
+
+ case UdfResponse.ResponseCase.CONTROL =>
+ handleControl(response.getControl)
+
+ case other =>
+ // A malformed response (empty / unknown oneof) is a terminal transport
+ // failure; fast-fail init the same way. Settle-before-release (see initValue).
+ Transitions.transportFailed(new IllegalStateException(
+ s"unexpected response oneof: $other"))
+ initValue.signalWithoutValue()
+ }
+ }
+
+ override def onError(t: Throwable): Unit = {
+ // Transport-level failure: the stream is dead, no further writes possible.
+ // Settle-before-release (see initValue) so init() surfaces the transport
+ // cause instead of the initResponseTimeoutMs "timed out" error.
+ Transitions.transportFailed(t)
+ initValue.signalWithoutValue()
+ }
+
+ 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 (!currentState.isTerminal) {
+ Transitions.transportFailed(new IllegalStateException(
+ "worker response stream closed without a terminator"))
+ }
+ // Defensive: if onCompleted reached us before InitResponse, doInit is
+ // still blocked on initValue and would otherwise time out.
+ initValue.signalWithoutValue()
+ }
+ }
+
+ /**
+ * Wakes the result iterator (blocked on [[outputQueue]]) and any thread
+ * waiting on [[terminalLatch]] when the base settles a terminal. Invoked once,
+ * by the caller that wins [[completeTerminal]].
+ */
+ override protected def onTerminalSettled(termination: Termination): Unit = {
+ outputQueue.put(QueueItem.EndOfStream)
+ terminalLatch.countDown()
+ }
+
+ private def handleControl(ctrl: UdfControlResponse): Unit = ctrl.getControlCase match {
+ case UdfControlResponse.ControlCase.INIT =>
+ val resp = ctrl.getInit
+ if (resp.hasError) {
+ // Attribute a subsequent close()'s terminator to the init error rather
+ // than a bare Cancelled, mirroring the data-phase ERROR branch.
+ executionError.compareAndSet(None, Some(resp.getError))
+ // Settle Failed(err) before the Cancel below (not after its CancelResponse):
+ // the sticky terminal makes close() report the init error, and the Cancel
+ // becomes best-effort worker cleanup -- so if it can't be written (a
+ // directExecutor worker delivered this reentrantly, before doInit published
+ // requestObserver), doInit still doesn't wait terminalTimeoutMs for a
+ // CancelResponse that can never arrive. Settle-before-release (see initValue).
+ Transitions.failed(resp.getError)
+ initValue.complete(resp)
+ sendCancelInternal(() => cancelWithReason("init failed"))
+ } else {
+ // InitResponse OK. Only advance from Initializing so a terminal that raced
+ // in (e.g. a transport error) still wins; process() then opens the data
+ // phase. Settle-before-release (see initValue): publish after the CAS.
+ Transitions.initAccepted()
+ initValue.complete(resp)
+ }
+
+ case UdfControlResponse.ControlCase.ERROR =>
+ val err = ctrl.getError.getError
+ executionError.compareAndSet(None, Some(err))
+ // Before init resolves, settle Failed here rather than rely on the Cancel
+ // below: sendCancelInternal may find requestObserver still null (reentrant
+ // delivery inside stream.onNext, before doInit published it), so no Cancel --
+ // hence no CancelResponse -- would settle the terminal, and doInit would
+ // return as if init succeeded. completeTerminal is idempotent, so the
+ // data-phase ERROR -> Cancel -> CancelResponse path is unaffected.
+ if (!initResolved) {
+ Transitions.failed(err)
+ }
+ // Settle-before-release (see initValue). No value to publish -- ERROR is not
+ // an InitResponse; on the data-phase path the latch already fired in init.
+ initValue.signalWithoutValue()
+ sendCancelInternal(() => cancelWithReason("aborting after ErrorResponse"))
+
+ case UdfControlResponse.ControlCase.FINISH =>
+ // The FinishResponse carries metrics + the finish-callback data/error.
+ // Keep it on the terminal so close() can return it; the iterator inspects
+ // its error field to decide whether to throw.
+ Transitions.finished(ctrl.getFinish)
+ // Defensive: FINISH before InitResponse is a worker protocol bug, but
+ // we should fail init fast rather than hang the 30s init timeout.
+ initValue.signalWithoutValue()
+
+ case UdfControlResponse.ControlCase.CANCEL =>
+ // The CancelResponse carries metrics + the cancel-callback error. Keep it
+ // on the terminal so close() can return it; any prior ErrorResponse is
+ // tracked in executionError and surfaced by the iterator.
+ Transitions.cancelled(ctrl.getCancel)
+ // Defensive: CANCEL before InitResponse unblocks doInit so it can
+ // surface the cancellation instead of timing out.
+ initValue.signalWithoutValue()
+
+ case UdfControlResponse.ControlCase.CONTROL_NOT_SET =>
+ Transitions.transportFailed(new IllegalStateException(
+ "empty UdfControlResponse oneof"))
+ initValue.signalWithoutValue()
+ }
+
+ /**
+ * True once init is no longer pending -- i.e. the stream is past `Initializing`.
+ * Not "init succeeded": a terminal (including a failure) also counts as resolved.
+ */
+ private def initResolved: Boolean = currentState match {
+ case SessionState.Created | SessionState.Initializing => false
+ case _ => true
+ }
+
+ private def cancelWithReason(reason: String): Cancel =
+ Cancel.newBuilder().setReason(reason).build()
+
+ /**
+ * The protocol transition graph in one place: names for the edges, not a
+ * second source of truth. Every edge acts on the single
+ * [[WorkerSession.SessionState]] owned by the base via [[compareAndSetState]]
+ * (non-terminal) or [[completeTerminal]] (terminal), so the base's CAS is the
+ * only synchronization and a terminal that arrived first always wins (the
+ * non-terminal CASes fail against it; [[completeTerminal]] is first-wins).
+ *
+ * Edges this class drives -- edge, method, then driver site(s) / thread (the
+ * base drives the API-call edges: `Created -> Initializing` in `init`,
+ * `Initialized -> Streaming` in `process`):
+ * {{{
+ * Initializing -> Initialized initAccepted handleControl INIT-ok [gRPC cb]
+ * Streaming -> Finishing finishSent advance branch 3 [engine]
+ * non-terminal -> Cancelling cancelSentFrom sendCancelInternal, iff a Cancel
+ * reaches the wire [gRPC cb | engine | init | close]
+ * non-terminal -> Terminal finished/cancelled/failed/transportFailed
+ * handleControl / doInit / doClose / advance / onError
+ * }}}
+ * `Cancelling` is reached only if a Cancel is actually written; a pre-stream or
+ * raced cancel goes straight to a `Cancelled`/`TransportFailed` terminal (see
+ * [[sendCancelInternal]], [[doClose]], `ProcessIterator`) or nowhere -- so
+ * `cancelRequested` (intent) does not imply state `Cancelling`.
+ */
+ private object Transitions {
+ /** `InitResponse` OK: `Initializing -> Initialized`. */
+ def initAccepted(): Boolean =
+ compareAndSetState(SessionState.Initializing, SessionState.Initialized)
+
+ /** Input exhausted and `Finish` written: `Streaming -> Finishing` (once). */
+ def finishSent(): Boolean =
+ compareAndSetState(SessionState.Streaming, SessionState.Finishing)
+
+ /** `Cancel` written: `cur -> Cancelling`, from any non-terminal `cur`. */
+ def cancelSentFrom(cur: SessionState): Boolean =
+ !cur.isTerminal && compareAndSetState(cur, SessionState.Cancelling)
+
+ /** Clean terminal carrying the worker's `FinishResponse`. */
+ def finished(response: FinishResponse): Boolean =
+ completeTerminal(Termination.Finished(response))
+
+ /** Clean terminal carrying the worker's `CancelResponse`. */
+ def cancelled(response: CancelResponse): Boolean =
+ completeTerminal(Termination.Cancelled(response))
+
+ /** Failure terminal carrying a structured [[ExecutionError]]. */
+ def failed(error: ExecutionError): Boolean =
+ completeTerminal(Termination.Failed(error))
+
+ /** Failure terminal carrying a transport-level cause. */
+ def transportFailed(cause: Throwable): Boolean =
+ completeTerminal(Termination.TransportFailed(cause))
+ }
+
+ // ---- WorkerSession hooks ------------------------------------------------
+
+ override protected def doInit(message: Init): InitResponse = {
+ // 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")
+ Transitions.transportFailed(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. close()'s 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 {
+ // Reentrant delivery: a directExecutor worker (one whose gRPC callbacks run
+ // synchronously on the caller's thread) can deliver InitResponse -- and thus
+ // run responseObserver/handleControl -- from *inside* this stream.onNext,
+ // before requestObserver below is published. The base advanced
+ // Created -> Initializing before calling doInit, so that reentrant
+ // InitResponse still finds Initializing and advances to Initialized; the
+ // handleControl error paths guard the still-null requestObserver. This is
+ // the scenario the "before doInit published requestObserver" comments mean.
+ 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
+ Transitions.transportFailed(e)
+ // Surface as GrpcWorkerSessionException so the engine integration layer
+ // (which catches that type and wraps it) sees a uniform init-failure
+ // exception rather than the raw transport error.
+ throw new GrpcWorkerSessionException("UDF worker stream failed during init", e)
+ }
+
+ val responded = try {
+ initValue.await(initResponseTimeoutMs)
+ } catch {
+ case _: InterruptedException =>
+ Thread.currentThread().interrupt()
+ sendCancelInternal(() => cancelWithReason("interrupted during init"))
+ // Make sure close() does not block waiting for a terminator that
+ // will never arrive on this thread's behalf.
+ Transitions.transportFailed(
+ new InterruptedException("interrupted while waiting for InitResponse"))
+ throw new InterruptedException("interrupted while waiting for InitResponse")
+ }
+
+ if (!responded) {
+ sendCancelInternal(() => cancelWithReason("InitResponse timed out"))
+ val timeout = new TimeoutException(
+ s"timed out waiting for InitResponse after ${initResponseTimeoutMs}ms")
+ // Settle the terminal so a subsequent close() does not stall for a
+ // second `terminalTimeoutMs` waiting for a worker that already missed
+ // its init deadline.
+ Transitions.transportFailed(timeout)
+ // Surface as GrpcWorkerSessionException (carrying the timeout cause) so
+ // the engine integration layer that catches that type can wrap it.
+ throw new GrpcWorkerSessionException(
+ s"timed out waiting for InitResponse after ${initResponseTimeoutMs}ms", timeout)
+ }
+
+ initValue.get match {
+ case Some(resp) if resp.hasError =>
+ // 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(resp.getError)}", resp.getError)
+ case Some(resp) =>
+ resp
+ case None =>
+ // No InitResponse arrived but the latch fired. The defensive
+ // initValue.signalWithoutValue() 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.
+ currentState match {
+ case SessionState.Terminal(Termination.TransportFailed(cause)) =>
+ throw new GrpcWorkerSessionException(
+ "UDF worker stream failed during init", cause)
+ case SessionState.Terminal(Termination.Failed(err)) =>
+ throw new GrpcWorkerSessionException(
+ s"UDF worker reported an error before init completed: " +
+ describeError(err), err)
+ case SessionState.Terminal(Termination.Cancelled(_)) =>
+ throw new GrpcWorkerSessionException(
+ "UDF worker stream was cancelled before init completed")
+ case SessionState.Terminal(Termination.Finished(_)) =>
+ throw new GrpcWorkerSessionException(
+ "UDF worker finished before init completed")
+ case other =>
+ throw new IllegalStateException(
+ s"init latch fired without an InitResponse or terminal: $other")
+ }
+ }
+ }
+
+ override protected def doProcess(
+ input: Iterator[DataRequest],
+ finish: () => Finish): Iterator[DataResponse] = {
+ // Init success is guaranteed by the base [[WorkerSession]] lifecycle: if
+ // doInit had failed it would have thrown and process() would never run.
+ new ProcessIterator(input, finish)
+ }
+
+ override protected def doClose(cancel: () => Cancel): Termination = {
+ if (requestObserver == null) {
+ // init() never put a stream on the wire (closed before/around init, or
+ // init threw before publishing). There is no protocol terminator; treat
+ // the session as cancelled-before-start. The base WorkerSession still
+ // releases the worker handle, so the worker is torn down.
+ Transitions.cancelled(CancelResponse.getDefaultInstance)
+ return Termination.Cancelled(CancelResponse.getDefaultInstance)
+ }
+ // If the stream has not finished on its own, cancel anything in flight so
+ // the worker can clean up, then wait for the terminator. sendCancelInternal
+ // evaluates the cancel thunk only when it actually writes a Cancel, and at
+ // most once across all callers.
+ if (!currentState.isTerminal) {
+ sendCancelInternal(cancel)
+ 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 the state sees a stable value.
+ if (!currentState.isTerminal) {
+ Transitions.transportFailed(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. onCompleted on an
+ // already-torn-down stream throws and is caught, so this stays idempotent
+ // across repeat close() calls.
+ try {
+ requestLock.synchronized { requestObserver.onCompleted() }
+ } catch {
+ case NonFatal(e) =>
+ logger.debug("Error half-closing UDF Execute stream", e)
+ }
+ // Derive the Termination from the settled terminal. close() is the
+ // finalizer/cleanup path and must NOT re-throw: a UDF / data-phase error is
+ // already surfaced while the result iterator is consumed, and an init error
+ // is surfaced from init(). settledTermination returns the settled terminal
+ // as-is: the response proto for the clean terminators, and the failure
+ // terminals (Failed / TransportFailed) carrying their cause -- not a bare
+ // Cancelled.
+ settledTermination
+ }
+
+ // ---- Internal request helpers ---------------------------------------------
+
+ /**
+ * Sends a Data or Finish request to the worker. Two invariants are
+ * checked inside [[requestLock]]:
+ * - terminal state: 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 (currentState.isTerminal) {
+ throw new IllegalStateException(
+ "cannot send request: UDF Execute stream is already closed")
+ }
+ // Suppress the write when a cancel has been (or is about to be) flushed
+ // through this lock; that preserves the proto ordering invariant (no
+ // Data/Finish after Cancel). Test the flag rather than `return`-ing from
+ // this by-name `synchronized` body: a non-local return compiles to a thrown
+ // NonLocalReturnControl, which is brittle (also relied on in sendCancelInternal).
+ if (!cancelRequested.get()) {
+ requestObserver.onNext(req)
+ }
+ }
+
+ /**
+ * Sends a `Cancel` control message, returning `true` iff a `Cancel` was
+ * actually written to the wire. Idempotent across ALL call sites (close()'s
+ * cancel, in-band cancels from `handleControl`'s INIT error / ERROR paths,
+ * and engine-internal cancels from the iterator) via [[cancelRequested]]:
+ * only the first caller puts a `Cancel` on the wire, and the `cancel` thunk is
+ * evaluated only by that caller and only when the message is actually sent --
+ * so a side-effecting thunk (e.g. one carrying a client cancel callback) runs
+ * at most once. 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(cancel: () => Cancel): Boolean = {
+ if (!cancelRequested.compareAndSet(false, true)) return false
+ if (requestObserver == null) return false
+ if (currentState.isTerminal) return false
+ try {
+ requestLock.synchronized {
+ // A terminal that arrived first must win; bail without writing. Compute
+ // the block's value rather than `return` (see sendRequest on why non-local
+ // return is avoided -- here it would also escape the NonFatal catch below).
+ val cur = currentState
+ if (cur.isTerminal) {
+ false
+ } else {
+ // Record that a Cancel is on the wire by advancing to Cancelling.
+ Transitions.cancelSentFrom(cur)
+ requestObserver.onNext(UdfRequest.newBuilder()
+ .setControl(UdfControlRequest.newBuilder().setCancel(cancel()).build())
+ .build())
+ true
+ }
+ }
+ } catch {
+ case NonFatal(e) =>
+ logger.debug(s"Cancel send failed (stream may already be torn down): ${e.getMessage}")
+ Transitions.transportFailed(e)
+ false
+ }
+ }
+
+ private def awaitTerminal(): Unit = {
+ if (currentState.isTerminal) return
+ try {
+ if (!terminalLatch.await(terminalTimeoutMs, TimeUnit.MILLISECONDS)) {
+ Transitions.transportFailed(new TimeoutException(
+ s"timed out waiting for stream terminator after ${terminalTimeoutMs}ms"))
+ }
+ } catch {
+ case _: InterruptedException =>
+ // Attribute the terminal to the interrupt, not a timeout: the wait was
+ // cut short, it did not exhaust terminalTimeoutMs. Reporting "timed
+ // out" here would misdiagnose an interrupted close as a slow worker.
+ Thread.currentThread().interrupt()
+ Transitions.transportFailed(
+ new InterruptedException("interrupted while waiting for stream terminator"))
+ }
+ }
+
+ // ---- 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 settle the state) and with any thread that finalizes via
+ * [[close]]. It drives the `Streaming -> Finishing` transition (sending
+ * `Finish` once input is exhausted, built lazily from the `finish` thunk).
+ */
+ private class ProcessIterator(input: Iterator[DataRequest], finish: () => Finish)
+ extends Iterator[DataResponse] {
+
+ // 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 branches, and block branch 4 for the
+ // full terminalTimeoutMs before returning. Callers that probe hasNext
+ // an extra time (iterator.size, instrumentation wrappers) would hang.
+ // Iterator-local and only touched by the single engine thread.
+ private val exhausted = new AtomicBoolean(false)
+ @volatile private var prefetched: DataResponse = _
+
+ // Pre-init cancel fast-fail: if a 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() && !currentState.isTerminal) {
+ Transitions.cancelled(CancelResponse.getDefaultInstance)
+ }
+
+ /**
+ * Runs a caller-supplied thunk (`input.hasNext`, `input.next()`, or the
+ * `finish` builder) that may throw, and on failure Cancels the in-flight
+ * stream before rethrowing. Without the Cancel a failure would leave the
+ * worker awaiting more input with no terminator owed on the wire; with it the
+ * worker tears down and the exception still propagates to the engine. Setting
+ * `cancelRequested` also suppresses any further Data/Finish this loop might
+ * otherwise attempt (see [[sendRequest]]).
+ */
+ private def cancelOnThrow[T](reason: String)(op: => T): T =
+ try op catch {
+ case NonFatal(e) =>
+ sendCancelInternal(() => cancelWithReason(reason))
+ throw e
+ }
+
+ override def hasNext: Boolean = {
+ if (prefetched ne null) return true
+ advance()
+ prefetched ne null
+ }
+
+ override def next(): DataResponse = {
+ if (prefetched eq null) advance()
+ val out = prefetched
+ if (out eq null) {
+ throw new NoSuchElementException("ProcessIterator exhausted")
+ }
+ prefetched = null
+ out
+ }
+
+ /**
+ * Fills [[prefetched]] with the next output batch, or leaves it null at the
+ * terminator (which also sets [[exhausted]], so `hasNext` reads null as "done").
+ * Each loop iteration tries in order: (1) drain queued output -- a batch (fill
+ * and return) or the terminator sentinel (return); (2) while `Streaming`, send
+ * the next input batch and loop; (3) once input is exhausted, send `Finish`
+ * (CAS `Streaming -> Finishing`, once) and loop; (4) otherwise block for late
+ * output or the terminator, then return. Branches 2-3 loop; 1 and 4 return.
+ *
+ * '''Per-poll timeout.''' Branch 4 waits up to [[terminalTimeoutMs]] per
+ * poll, reset by every worker event -- the contract is "emit at least one
+ * event every [[terminalTimeoutMs]] after Finish", not "finish the UDF within
+ * [[terminalTimeoutMs]]". A worker expecting a long post-Finish silence MAY
+ * emit an empty `DataResponse` as a heartbeat to reset the wait; it is
+ * surfaced to the caller, so it should be a batch the caller recognises as
+ * empty (e.g. a zero-row Arrow batch).
+ */
+ private def advance(): Unit = {
+ if (exhausted.get()) return // terminator already seen; never re-block
+ 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 while the stream is open for data.
+ // `input.hasNext` may itself fetch/compute the next element (many Spark
+ // iterators prefetch), so it can throw just like `input.next()`; both go
+ // through cancelOnThrow so a failing upstream Cancels the stream instead
+ // of stranding the worker waiting for input that will never arrive.
+ if (!cancelRequested.get() && currentState == SessionState.Streaming &&
+ cancelOnThrow("input iterator failed")(input.hasNext)) {
+ val batch = cancelOnThrow("input iterator failed")(input.next())
+ if (!sendOrEndOnRacedTerminal(UdfRequest.newBuilder().setData(batch).build())) return
+ } else if (!cancelRequested.get() && Transitions.finishSent()) {
+ // (3) No more input; send Finish exactly once (unless cancelled). The
+ // `finish` thunk is caller-supplied (it may run a finish callback), so
+ // it can throw; cancelOnThrow Cancels the stream before rethrowing. The
+ // Streaming -> Finishing CAS has already run, but we have not written
+ // Finish yet, so the Cancel is the only engine-to-worker message that
+ // reaches the wire -- the proto "nothing after Cancel" invariant holds.
+ val finishMsg = cancelOnThrow("finish callback failed")(finish())
+ if (!sendOrEndOnRacedTerminal(UdfRequest.newBuilder()
+ .setControl(UdfControlRequest.newBuilder().setFinish(finishMsg).build())
+ .build())) {
+ return
+ }
+ } else {
+ // (4) Block for late output or the terminator. See class doc
+ // above for the per-poll-vs-total-session timeout semantics.
+ val item = try {
+ outputQueue.poll(terminalTimeoutMs, TimeUnit.MILLISECONDS)
+ } catch {
+ case _: InterruptedException =>
+ Thread.currentThread().interrupt()
+ sendCancelInternal(() => cancelWithReason("iterator interrupted"))
+ exhausted.set(true)
+ throw new InterruptedException("interrupted while reading UDF result")
+ }
+ item match {
+ case null =>
+ Transitions.transportFailed(new IllegalStateException(
+ s"timed out waiting for UDF output after ${terminalTimeoutMs}ms"))
+ exhausted.set(true)
+ throwIfTerminalError()
+ return
+ case QueueItem.EndOfStream =>
+ exhausted.set(true)
+ throwIfTerminalError()
+ return
+ case QueueItem.Batch(b) =>
+ prefetched = b
+ return
+ }
+ }
+ }
+ }
+
+ /**
+ * Sends one data-phase request (`DataRequest` / `Finish`), recovering from a
+ * terminator that raced the write. Returns `true` to keep looping (sent, or
+ * suppressed by a pending cancel), `false` if a terminator settled and the
+ * iterator should end (caller `return`s).
+ *
+ * On a write failure a terminator usually already settled (the worker
+ * finished/failed early) and the write only failed because the stream is
+ * closed: record a transport terminal if none is set ([[completeTerminal]] is
+ * a no-op once settled), then [[throwIfTerminalError]] -- which throws for an
+ * error/transport terminal, or returns for a clean `Finished` that raced the
+ * send, in which case the benign "stream closed" error is dropped.
+ */
+ private def sendOrEndOnRacedTerminal(req: UdfRequest): Boolean = {
+ try {
+ sendRequest(req)
+ true
+ } catch {
+ case NonFatal(e) =>
+ if (!currentState.isTerminal) {
+ Transitions.transportFailed(e)
+ }
+ exhausted.set(true)
+ throwIfTerminalError()
+ false
+ }
+ }
+
+ /**
+ * Surfaces a failed terminator as an exception when the result iterator is
+ * drained. A data-phase [[ExecutionError]] (captured in [[executionError]])
+ * takes precedence over the terminator's own error, then the finish/cancel
+ * callback error carried on the terminator, then a bare cancellation.
+ */
+ private def throwIfTerminalError(): Unit = currentState match {
+ case SessionState.Terminal(Termination.Finished(response)) =>
+ responseError(response.hasError, response.getError, executionError.get())
+ .foreach(err => throw new GrpcWorkerSessionException(
+ s"UDF execution failed: ${describeError(err)}", err))
+ case SessionState.Terminal(Termination.Cancelled(response)) =>
+ responseError(response.hasError, response.getError, executionError.get()) match {
+ case Some(err) =>
+ throw new GrpcWorkerSessionException(
+ s"UDF execution failed: ${describeError(err)}", err)
+ case None =>
+ throw new GrpcWorkerSessionException("UDF execution was cancelled")
+ }
+ case SessionState.Terminal(Termination.Failed(err)) =>
+ throw new GrpcWorkerSessionException(
+ s"UDF execution failed: ${describeError(err)}", err)
+ case SessionState.Terminal(Termination.TransportFailed(t)) =>
+ throw new GrpcWorkerSessionException("UDF worker stream failed", t)
+ case other =>
+ throw new IllegalStateException(s"terminator sentinel without terminal: $other")
+ }
+
+ /** Picks the error to surface: prior data-phase error, else the terminator's. */
+ private def responseError(
+ hasError: Boolean,
+ error: ExecutionError,
+ priorError: Option[ExecutionError]): Option[ExecutionError] =
+ priorError.orElse(if (hasError) Some(error) else None)
+ }
+}
+
+object GrpcWorkerSession {
+ /** Upper bound on the wait for `InitResponse`. */
+ val DEFAULT_INIT_RESPONSE_TIMEOUT_MS: Long = 30000L
+
+ /** Upper bound on the wait for `FinishResponse` / `CancelResponse`. */
+ val DEFAULT_TERMINAL_TIMEOUT_MS: Long = 30000L
+
+ // Distinguishes a (possibly empty) data batch from end-of-stream, and makes
+ // the iterator's match exhaustive.
+ private sealed trait QueueItem
+ private object QueueItem {
+ final case class Batch(response: DataResponse) extends QueueItem
+ case object EndOfStream extends QueueItem
+ }
+
+ /**
+ * A write-once value paired with the latch a waiter blocks on, so the "publish
+ * the value, then release the waiter" ordering lives in one place instead of
+ * every call site having to remember to count the latch down after setting the
+ * reference. The value is set at most once ([[complete]]); the latch can also
+ * be released without a value ([[signalWithoutValue]]) when the waiter must be
+ * woken by a terminal that arrived instead of the value. All reads/writes go
+ * through the latch, so a waiter released by [[await]] has a happens-before
+ * edge to the [[complete]] that set the value.
+ */
+ private final class OneShotValue[A] {
+ private val latch = new CountDownLatch(1)
+ private val ref = new AtomicReference[Option[A]](None)
+
+ /** Publishes the value (first writer wins) and releases any waiter. */
+ def complete(value: A): Unit = {
+ ref.compareAndSet(None, Some(value))
+ latch.countDown()
+ }
+
+ /** Releases the waiter without a value (a terminal arrived, not the value). */
+ def signalWithoutValue(): Unit = latch.countDown()
+
+ /** Blocks up to `timeoutMs` for a release; true iff released in time. */
+ def await(timeoutMs: Long): Boolean = latch.await(timeoutMs, TimeUnit.MILLISECONDS)
Review Comment:
OneShotValue.await now returns Unit and throws TimeoutException (L924-928). Exactly the change requested.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org