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: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
