Yicong-Huang commented on code in PR #56702: URL: https://github.com/apache/spark/pull/56702#discussion_r3707554551
########## udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/GrpcWorkerSession.scala: ########## @@ -0,0 +1,816 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.udf.worker.grpc + +import java.util.concurrent.{CountDownLatch, LinkedBlockingQueue, TimeoutException, TimeUnit} +import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference} + +import scala.util.control.NonFatal + +import io.grpc.{ConnectivityState, ManagedChannel} +import io.grpc.stub.StreamObserver + +import org.apache.spark.annotation.Experimental +import org.apache.spark.udf.worker.{Cancel, CancelResponse, DataRequest, DataResponse, + ExecutionError, Finish, Init, InitResponse, UdfControlRequest, + UdfControlResponse, UdfRequest, UdfResponse, UdfWorkerGrpc} +import org.apache.spark.udf.worker.core.{Termination, WorkerHandle, WorkerLogger, WorkerSession} +import org.apache.spark.udf.worker.core.WorkerSession.SessionState +import org.apache.spark.udf.worker.grpc.GrpcWorkerSession._ + +/** + * :: Experimental :: + * gRPC implementation of [[WorkerSession]] for the `UdfWorker.Execute` + * bidirectional RPC. + * + * Drives one bidirectional `Execute` stream against the worker per the + * ordering invariants documented in `udf_message.proto`: + * {{{ + * Engine -> Worker: Init -> (DataRequest)* -> Finish (Cancel)? + * | Cancel + * Worker -> Engine: InitResponse -> (DataResponse)* -> + * (ErrorResponse)? -> (FinishResponse | CancelResponse) + * }}} + * + * Knows nothing about how the worker was provisioned (locally spawned, + * indirectly looked up, ...) -- the dispatcher constructs this with a + * [[WorkerHandle]] and channel; the base [[WorkerSession]] handles + * dispatcher-side cleanup on close. + * + * '''Driving model.''' Consumption-driven (Volcano / pull): the thread that + * consumes the [[doProcess]] result iterator is the one that pulls the next + * input batch and sends its `DataRequest`. Nothing is sent until the iterator + * is consumed, and input flows one batch per output pull. The gRPC callback + * thread only receives output. + * + * '''State machine.''' This class does not keep its own state machine: it + * drives the single [[WorkerSession.SessionState]] owned by the base. The base + * advances `Created -> Initializing` (in `init`) and `Initialized -> Streaming` + * (in `process`); this class advances the protocol-event edges through + * [[compareAndSetState]] / [[completeTerminal]] as it exchanges messages: + * {{{ + * Initializing --(InitResponse ok)--> Initialized [handleControl] + * Streaming ----(Finish written)----> Finishing [ProcessIterator] + * <any non-terminal> --(Cancel written)--> Cancelling [sendCancelInternal] + * <any non-terminal> --(terminator/error)--> terminal [completeTerminal] + * }}} + * The two clean terminals carry the worker's `FinishResponse` / `CancelResponse` + * (metrics + finish/cancel callback `data`/`error`) so [[close]] can return + * them. The only flag kept outside the machine is [[cancelRequested]]: a + * cancellation can be requested before the stream exists (so it cannot be a + * state transition yet), and it must both fast-fail the result iterator and + * suppress any in-flight Data/Finish. + * + * Threading: + * - [[doInit]] is synchronous: sends `Init` and blocks on `InitResponse`, + * returning it. + * - [[doProcess]] returns an iterator. Input batches are forwarded inline + * (the iterator's `next()` thread also sends `DataRequest`). Output + * batches arrive via the response observer (gRPC callback thread) and + * are consumed by the same iterator. A terminator (`FinishResponse`, + * `CancelResponse`, `ErrorResponse`, gRPC stream error) is published + * once. + * - [[doClose]] is thread-safe and idempotent: it settles + returns the + * terminator (cancelling in-flight work if the stream had not finished) + * and half-closes the request side. + * + * TODO [SPARK-55278]: this class does not yet implement payload chunking; + * the entire [[Init.udf]] payload is sent inline. Chunking will be added + * when a UDF payload large enough to exceed gRPC's default message size + * limit is introduced. + * + * @param workerHandle dispatcher-side handle for releasing the worker on + * [[close]] (see [[WorkerSession]]). + * @param channel a gRPC channel built and owned by the caller (the + * dispatcher). Not closed here -- the dispatcher tears it + * down via [[WorkerHandle]]. + * @param logger diagnostics. Defaults to [[WorkerLogger.NoOp]]. + * @param initResponseTimeoutMs upper bound on the wait for `InitResponse` + * after [[doInit]] sends `Init`. + * @param terminalTimeoutMs upper bound on the wait for a stream + * terminator (`FinishResponse`, + * `CancelResponse`, or `ErrorResponse`). + * Each output-queue poll resets this wait; + * see [[doProcess]] / `ProcessIterator`. + */ +@Experimental +class GrpcWorkerSession( + workerHandle: WorkerHandle, + channel: ManagedChannel, + logger: WorkerLogger = WorkerLogger.NoOp, + initResponseTimeoutMs: Long = DEFAULT_INIT_RESPONSE_TIMEOUT_MS, + terminalTimeoutMs: Long = DEFAULT_TERMINAL_TIMEOUT_MS) + extends WorkerSession(workerHandle, logger) { + + require(channel != null, "channel is required") + + private val asyncStub = UdfWorkerGrpc.newStub(channel) + + // Output batches from the worker, drained by the process() iterator. + // Intentionally unbounded: a bounded queue would block the gRPC callback + // (Netty event-loop) thread when full, stalling control-message and + // terminator delivery on the whole channel. HTTP/2 flow control bounds the + // wire and the consumer normally drains promptly, so it stays small. A + // stalled downstream can still grow it; the real fix is protocol-level + // back-pressure (out of scope here), not bounding the queue. + // + // TODO [SPARK-57324]: expose queue depth as a metric (early warning for a + // stalled consumer). + private val outputQueue = new LinkedBlockingQueue[QueueItem]() + + // Latch fired when `InitResponse` (success or error) or a transport error + // arrives. init() blocks on this; until it fires we have no proof the + // worker actually accepted the session. + private val initLatch = new CountDownLatch(1) + + // The InitResponse the worker sent (success or error), captured so init() + // can return it. None until InitResponse arrives. + private val initResponse = new AtomicReference[Option[InitResponse]](None) + + // Fired when the session reaches a terminal [[SessionState]]. doClose() and + // the init-error path block on this to drain the terminator. + private val terminalLatch = new CountDownLatch(1) + + // Captures an ErrorResponse encountered during the data phase so that + // the CancelResponse terminator can attribute the failure to the original + // user / worker / protocol error rather than reporting a bare "Cancelled". + private val executionError = new AtomicReference[Option[ExecutionError]](None) + + // Cancellation intent. Kept outside the [[SessionState]] machine because a + // cancel can be requested before the stream exists (pre-init), where there is + // no wire transition to make yet -- only an intent to record. Used to (a) + // make cancel idempotent across all call sites, (b) fast-fail ProcessIterator + // on a pre-init cancel, and (c) suppress any Data/Finish that would otherwise + // race a Cancel onto the wire (re-read inside [[requestLock]]). + private val cancelRequested = new AtomicBoolean(false) + + // gRPC requires serialized writes to a request StreamObserver. + private val requestLock = new Object + + // Initialised in init() -- before that, close() is a no-op on the request + // side, which is exactly the contract the wrapping WorkerSession expects. + @volatile private var requestObserver: StreamObserver[UdfRequest] = _ + + private val responseObserver: StreamObserver[UdfResponse] = new StreamObserver[UdfResponse] { + override def onNext(response: UdfResponse): Unit = { + response.getResponseCase match { + case UdfResponse.ResponseCase.DATA => + outputQueue.put(QueueItem.Batch(response.getData)) + + case UdfResponse.ResponseCase.CONTROL => + handleControl(response.getControl) + + case other => + // A malformed response (empty / unknown oneof) is a terminal + // transport failure. Count down initLatch like every other + // terminal-settling path here (onError / onCompleted / each + // handleControl branch): if this arrives before InitResponse, doInit + // must fail fast with this cause rather than block until + // initResponseTimeoutMs and report a misleading "timed out" error. + completeTerminal(Termination.TransportFailed(new IllegalStateException( + s"unexpected response oneof: $other"))) + initLatch.countDown() + } + } + + override def onError(t: Throwable): Unit = { + // Transport-level failure: the stream is dead. No further writes are + // possible (reaching a terminal state closes the write side). Settle the + // terminal BEFORE counting down initLatch: init() blocks on that latch, + // and only await/countDown establish a happens-before edge, so a thread + // woken by the countDown must already be able to observe the terminal. + // That lets init() surface the transport cause instead of timing out + // after initResponseTimeoutMs with a misleading "timed out" message -- + // or, if it saw a transient non-terminal state, the defensive "init latch + // fired without an InitResponse or terminal" error. + completeTerminal(Termination.TransportFailed(t)) + initLatch.countDown() + } + + override def onCompleted(): Unit = { + // Worker half-closed its side without sending a terminator (FinishResponse + // / CancelResponse). Treat as transport error so the engine sees a + // failure, not a silent end-of-stream. + if (!currentState.isTerminal) { + completeTerminal(Termination.TransportFailed(new IllegalStateException( + "worker response stream closed without a terminator"))) + } + // Defensive: if onCompleted reached us before InitResponse, doInit is + // still blocked on initLatch and would otherwise time out. + initLatch.countDown() + } + } + + /** + * Wakes the result iterator (blocked on [[outputQueue]]) and any thread + * waiting on [[terminalLatch]] when the base settles a terminal. Invoked once, + * by the caller that wins [[completeTerminal]]. + */ + override protected def onTerminalSettled(termination: Termination): Unit = { + outputQueue.put(QueueItem.EndOfStream) + terminalLatch.countDown() + } + + private def handleControl(ctrl: UdfControlResponse): Unit = ctrl.getControlCase match { + case UdfControlResponse.ControlCase.INIT => + val resp = ctrl.getInit + // Capture the InitResponse so init() can return it (or throw on its error). + initResponse.set(Some(resp)) + if (resp.hasError) { + initLatch.countDown() + // Per the protocol the engine follows an init error with Cancel; send it + // if the request stream is up. If the Cancel cannot be written -- e.g. a + // directExecutor worker delivered InitResponse reentrantly, before + // doInit published requestObserver -- settle the terminal here so doInit + // does not wait the full terminalTimeoutMs for a CancelResponse that can + // never arrive. (When the Cancel is sent, the CancelResponse settles the + // terminal instead.) + if (!sendCancelInternal(() => cancelWithReason("init failed")) && Review Comment: Confirmed fixed: the INIT-error branch now records `executionError` before publishing the `InitResponse`, and `doInit` owns the Cancel plus `CancelResponse` drain, so the structured error reaches the caller and the terminal is the proto `Cancelled`. ########## udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/GrpcWorkerSessionConcurrencySuite.scala: ########## @@ -0,0 +1,742 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.udf.worker.grpc + +import java.util.Locale +import java.util.concurrent.{ConcurrentLinkedQueue, CountDownLatch, TimeUnit} +import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference} + +import scala.jdk.CollectionConverters._ + +import com.google.protobuf.ByteString +import io.grpc.{ManagedChannel, Server} +import io.grpc.inprocess.{InProcessChannelBuilder, InProcessServerBuilder} +import io.grpc.stub.StreamObserver +import org.scalatest.BeforeAndAfterEach +// scalastyle:off funsuite +import org.scalatest.funsuite.AnyFunSuite + +import org.apache.spark.udf.worker.{Cancel, CancelResponse, DataRequest, DataResponse, + ErrorResponse, ExecutionError, Finish, FinishResponse, Init, InitResponse, UdfControlResponse, + UdfPayload, UdfRequest, UdfResponse, UDFWorkerDataFormat, UdfWorkerGrpc, UserError, + WorkerRequest, WorkerResponse} +import org.apache.spark.udf.worker.core.{Termination, WorkerHandle, WorkerLogger} + +/** + * Concurrency tests for [[GrpcWorkerSession]] that pin the wire-ordering and + * fast-fail invariants under concurrent and worker-misbehavior scenarios: + * - Cancel must never appear on the wire before Init. + * - A worker terminator (ERROR / FINISH / CANCEL / onCompleted) arriving + * before InitResponse must fail [[GrpcWorkerSession#init]] fast, not hang + * for `initResponseTimeoutMs`. + * - Repeated [[Iterator#hasNext]] after natural iterator exhaustion must + * return immediately, not block for `terminalTimeoutMs`. + * - Cancel issued before [[GrpcWorkerSession#init]] makes a subsequent + * [[GrpcWorkerSession#doProcess]] surface the cancellation immediately, + * instead of blocking for `terminalTimeoutMs`. + * - Cancel / close concurrent with an in-progress data phase terminate + * cleanly without leaks or unbounded hangs. + * + * Runs entirely in-process: no subprocess, no UDS. Server services are + * custom-built per test so we can drive specific worker misbehavior. + */ +class GrpcWorkerSessionConcurrencySuite + extends AnyFunSuite with BeforeAndAfterEach { +// scalastyle:on funsuite + + /** Used by tests to keep stale in-flight infra reachable for teardown. */ + private val openServers = new ConcurrentLinkedQueue[Server]() + private val openChannels = new ConcurrentLinkedQueue[ManagedChannel]() + private val openSessions = new ConcurrentLinkedQueue[GrpcWorkerSession]() + + override def afterEach(): Unit = { + // Shut channels down first. This fires onError on any still-live stream, + // which settles the session terminal and counts down the init/terminal + // latches. That unblocks both the session.close() below and any worker + // thread a failing test left parked on a (deliberately large) timeout, so + // teardown never hangs even when a test asserts via assertFinishesWithin. + openChannels.asScala.foreach { c => + try c.shutdownNow().awaitTermination(2, TimeUnit.SECONDS) catch { case _: Throwable => () } + } + openChannels.clear() + openSessions.asScala.foreach { s => try s.close(emptyCancel) catch { case _: Throwable => () } } + openSessions.clear() + openServers.asScala.foreach { s => + try s.shutdownNow().awaitTermination(2, TimeUnit.SECONDS) catch { case _: Throwable => () } + } + openServers.clear() + super.afterEach() + } + + // A session timeout large enough that a correct test never reaches it; a + // regression that fails to short-circuit blocks here for minutes and is caught + // by assertFinishesWithin (below) instead of a flaky `elapsed < timeout` bound. + private val NeverReachedTimeoutMs = TimeUnit.MINUTES.toMillis(10) + + /** + * Runs `body` on a daemon thread and asserts it finishes within `withinMs`, + * rethrowing whatever `body` threw (so an `intercept` inside `body` still + * works). Pair with [[NeverReachedTimeoutMs]]: the correct fast path returns + * in milliseconds, so `withinMs` (seconds) has an enormous safety margin and + * does not flake, while a regression that parks on the timeout never returns + * within `withinMs` and fails the assertion. + */ + private def assertFinishesWithin(withinMs: Long, name: String)(body: => Unit): Unit = { + val thrown = new AtomicReference[Throwable]() + val done = new CountDownLatch(1) + val worker = new Thread(() => { + try body catch { case t: Throwable => thrown.set(t) } finally done.countDown() + }, name) + worker.setDaemon(true) + worker.start() + assert(done.await(withinMs, TimeUnit.MILLISECONDS), + s"$name did not finish within ${withinMs}ms; it parked on a session timeout " + + "that the fast path should have short-circuited") + Option(thrown.get()).foreach(t => throw t) + } + + // --------------------------------------------------------------------------- + // Helpers + // --------------------------------------------------------------------------- + + private class TestWorkerHandle extends WorkerHandle { + val invalidated = new AtomicBoolean(false) Review Comment: Confirmed: `released`/`invalidated` are now asserted across the clean-finish, TransportFailed, acknowledged-Cancelled and unacknowledged-Interrupted paths, including the idempotent repeat `close()`. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [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]
