diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/externalUDF/ExternalUDFExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/externalUDF/ExternalUDFExec.scala index 883e7dce88b10..57630ce1cba92 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/externalUDF/ExternalUDFExec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/externalUDF/ExternalUDFExec.scala @@ -49,7 +49,7 @@ trait ExternalUDFExec extends UnaryExecNode { // --------------------------------------------------------------------------- protected def externalUdfMetrics: Map[String, SQLMetric] = Map( - // TODO [SPARK-55278]: Emit the correct metrics here + // TODO [SPARK-57324]: Emit the correct metrics here ) override lazy val metrics: Map[String, SQLMetric] = externalUdfMetrics @@ -60,10 +60,12 @@ trait ExternalUDFExec extends UnaryExecNode { /** * Creates a [[WorkerSession]] via [[SparkEnv#getExternalUDFDispatcher]]. - * Registers session cancellation on task failure and session termination - * on task completion. The provided function receives the session - * and must return the result iterator. The function may use the - * session but MUST NOT cancel or close it. + * Finalizes the session on task completion (which fires on both success and + * failure). [[WorkerSession#close]] is the single finalizer: it fetches the + * `FinishResponse` if processing completed, or cancels anything still in + * flight and waits for the `CancelResponse`. The provided function receives + * the session and must return the result iterator. It may use the session + * but MUST NOT close it. */ protected def withUDFWorkerSession( taskContext: TaskContext, @@ -74,12 +76,29 @@ trait ExternalUDFExec extends UnaryExecNode { workerSpec) val session = dispatcher.createSession(securityScope) - // Make sure to cancel the session, if the task fails - taskContext.addTaskFailureListener { (_, _) => - session.cancel() - } - - // Make sure to close the session once we are done + // Finalize the session when the task ends. The completion listener fires on + // both success and failure, and close() is the single finalizer that + // resolves to whichever terminator the stream reached: + // + // - Task completed and the result iterator was fully consumed: process() + // sent Finish (input exhausted) and the data drained. close() then + // typically returns the Finished termination -- but a failure during the + // finish/cleanup phase (raised after the data drained, so it reached no one + // through the iterator) is surfaced here too, as Failed / TransportFailed. + // - Task failed, was killed, or stopped before draining (e.g. a downstream + // LIMIT or exception): the stream has not finished, so close() sends a + // Cancel, the worker runs its cleanup, and its CancelResponse is returned + // as a Cancelled termination. An empty Cancel is enough here -- there is + // no extra information to convey to the worker on cancellation -- so we + // rely on close()'s default and pass no cancel thunk. + // - The stream died without a terminator (transport failure / timeout): + // close() returns a best-effort TransportFailed termination rather than + // raising it (a thread interrupt may still propagate); the underlying + // failure has already surfaced through the result iterator. + // + // The returned Termination (per-execution metrics, finish/cancel callback + // result) is not consumed yet -- TODO [SPARK-57324] surface it once metrics + // wiring lands. taskContext.addTaskCompletionListener[Unit] { _ => session.close() } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/externalUDF/PythonUDFWorkerSpecificationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/externalUDF/PythonUDFWorkerSpecificationSuite.scala index f44d3adf7b7d7..bbd47741aba1c 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/externalUDF/PythonUDFWorkerSpecificationSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/externalUDF/PythonUDFWorkerSpecificationSuite.scala @@ -17,6 +17,8 @@ package org.apache.spark.sql.execution.externalUDF import java.io.File +import java.net.{StandardProtocolFamily, UnixDomainSocketAddress} +import java.nio.channels.SocketChannel import java.nio.charset.StandardCharsets import java.nio.file.Files @@ -26,72 +28,57 @@ import scala.jdk.CollectionConverters._ import org.apache.spark.api.python.SimplePythonFunction import org.apache.spark.sql.IntegratedUDFTestUtils import org.apache.spark.sql.test.SharedSparkSession -import org.apache.spark.udf.worker.UDFWorkerSpecification -import org.apache.spark.udf.worker.core.{TestDirectWorkerDispatcher, - UnixSocketWorkerConnection} +import org.apache.spark.udf.worker.{Cancel, UDFWorkerSpecification} +import org.apache.spark.udf.worker.core.{TestDirectWorkerDispatcher, WorkerConnection} /** - * A test [[UnixSocketWorkerConnection]] that opens a real Unix - * domain socket channel to the worker. + * A [[WorkerConnection]] that opens a real Unix-domain-socket channel to the + * worker at construction and considers itself active while the channel is + * open. Verifies the worker's socket is actually connectable, not merely + * that the file exists. */ -private class RealSocketConnection( - socketPath: String, - private val channel: java.nio.channels.SocketChannel) - extends UnixSocketWorkerConnection(socketPath) { - +private class ConnectingSocketConnection( + private val channel: SocketChannel) extends WorkerConnection { override def isActive: Boolean = channel.isOpen - - override def close(): Unit = { - channel.close() - super.close() - } + override def close(): Unit = channel.close() } -private object RealSocketConnection { - def connect(socketPath: String): RealSocketConnection = { - val address = - java.net.UnixDomainSocketAddress.of(socketPath) - val channel = java.nio.channels.SocketChannel.open( - java.net.StandardProtocolFamily.UNIX) - channel.connect(address) - new RealSocketConnection(socketPath, channel) +private object ConnectingSocketConnection { + def connect(socketPath: String): ConnectingSocketConnection = { + val channel = SocketChannel.open(StandardProtocolFamily.UNIX) + channel.connect(UnixDomainSocketAddress.of(socketPath)) + new ConnectingSocketConnection(channel) } } /** - * Extends [[TestDirectWorkerDispatcher]] to use a real Unix - * domain socket connection instead of just checking file - * existence. + * A [[TestDirectWorkerDispatcher]] whose connection actually connects to the + * worker socket (rather than just checking file existence), so the test + * verifies the spawned Python worker is reachable. */ -private class RealSocketTestDispatcher( - spec: UDFWorkerSpecification) +private class ConnectingTestDispatcher(spec: UDFWorkerSpecification) extends TestDirectWorkerDispatcher(spec) { - - // Use /tmp to avoid UDS path length limits in deep - // worktree paths. - override protected def newEndpointAddress( - workerId: String): String = { - val socketDir = java.nio.file.Files.createTempDirectory( - java.nio.file.Paths.get("/tmp"), "udf-test-") - socketDir.toFile.deleteOnExit() - socketDir.resolve(s"w-$workerId.sock").toString - } - - override protected def createConnection( - socketPath: String): UnixSocketWorkerConnection = - RealSocketConnection.connect(socketPath) + override protected def newConnection(address: String): WorkerConnection = + ConnectingSocketConnection.connect(address) } /** * Tests that [[PythonUDFWorkerSpecification#fromPythonFunction]] - * produces a valid [[UDFWorkerSpecification]] that can be used by - * a [[org.apache.spark.udf.worker.core.WorkerDispatcher]] + * produces a valid [[org.apache.spark.udf.worker.UDFWorkerSpecification]] + * that can be used by a + * [[org.apache.spark.udf.worker.core.WorkerDispatcher]] * to spawn a real Python worker process. * * The test overrides `spark.python.worker.module` to point to a * test-only Python module that creates the UDS socket and waits * for SIGTERM, verifying that pythonExec, PYTHONPATH, env vars, * and command construction all work end-to-end. + * + * Reuses [[TestDirectWorkerDispatcher]]'s spawn / wait-for-ready machinery + * (shared from `udf-worker-core` test sources) but overrides the connection + * to open a real UDS channel -- the Python test worker only binds a raw + * socket, it does not host a gRPC server, so a full gRPC session is not + * exercised here. */ class PythonUDFWorkerSpecificationSuite extends SharedSparkSession { @@ -185,15 +172,16 @@ class PythonUDFWorkerSpecificationSuite val workerSpec = PythonUDFWorkerSpecification.fromPythonFunction(func, conf) - // Verify the spec works end-to-end with a real - // socket connection - val dispatcher = new RealSocketTestDispatcher(workerSpec) + // Verify the spec works end-to-end: the dispatcher spawns the Python + // worker, waits for the socket, and opens a real UDS connection to it + // (proving the worker is reachable, not just that the file exists). + val dispatcher = new ConnectingTestDispatcher(workerSpec) try { val session = dispatcher.createSession( securityScope = None) assert(session != null, "Expected a non-null session from the dispatcher") - session.close() + session.close(() => Cancel.getDefaultInstance) } finally { dispatcher.close() } diff --git a/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/Termination.scala b/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/Termination.scala new file mode 100644 index 0000000000000..4278313fdf394 --- /dev/null +++ b/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/Termination.scala @@ -0,0 +1,66 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.udf.worker.core + +import org.apache.spark.annotation.Experimental +import org.apache.spark.udf.worker.{CancelResponse, ExecutionError, FinishResponse} + +/** + * :: Experimental :: + * The terminal outcome a [[WorkerSession]] settles on, returned by + * [[WorkerSession#close]]. Mirrors the four terminal `WorkerSession.SessionState`s, + * so close() reports the outcome faithfully rather than collapsing failures into + * a clean cancel. + * + * '''Clean outcomes''' ([[Finished]] / [[Cancelled]]) wrap the worker's + * `FinishResponse` / `CancelResponse` -- per-execution metrics, an optional + * inline `data` payload, and an optional finish/cancel '''callback error''' -- + * which the engine surfaces to the UDF's client-side finish/cancel callbacks. + * The proto allows `Cancel` to follow `Finish` on the same stream, so close() + * may observe either depending on arrival timing (see `udf_message.proto`): if + * `FinishResponse` was already produced when a `Cancel` arrives the engine still + * receives [[Finished]], otherwise [[Cancelled]]. + * + * '''Failure outcomes''' ([[Failed]] / [[TransportFailed]]) carry the cause + * instead of a proto terminator (none arrived, so they have no metrics). They + * exist so a failure is not reported as a benign cancel -- in particular an + * error raised during finish/close, '''after all data has been drained''', + * reaches the caller only through this value, never through the result iterator. + */ +@Experimental +sealed trait Termination + +object Termination { + /** The stream ended with a `FinishResponse` (normal completion). */ + final case class Finished(response: FinishResponse) extends Termination + + /** The stream ended with a `CancelResponse` (cooperative cancellation). */ + final case class Cancelled(response: CancelResponse) extends Termination + + /** + * An execution error settled the session without a proto terminator (e.g. an + * `ErrorResponse`, or an error before init completed). Carries the structured + * [[ExecutionError]]. + */ + final case class Failed(error: ExecutionError) extends Termination + + /** + * A transport failure, timeout, or interrupt tore the stream down before any + * terminator arrived. Carries the underlying cause. + */ + final case class TransportFailed(cause: Throwable) extends Termination +} diff --git a/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/WorkerConnection.scala b/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/WorkerConnection.scala index 82b2fff8df585..921dc910fd34e 100644 --- a/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/WorkerConnection.scala +++ b/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/WorkerConnection.scala @@ -33,7 +33,7 @@ import org.apache.spark.annotation.Experimental * multiplexed streams over this channel. * * Implementations expose only lifecycle. Data transmission happens at - * the [[WorkerSession]] level -- this class is solely about whether the + * the [[WorkerSession]] level -- this trait is solely about whether the * channel is open. * * '''Relationship to other classes (direct creation mode):''' @@ -43,7 +43,7 @@ import org.apache.spark.annotation.Experimental * }}} */ @Experimental -abstract class WorkerConnection extends AutoCloseable { +trait WorkerConnection extends AutoCloseable { /** Returns true if the underlying transport channel is still usable. */ def isActive: Boolean } diff --git a/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/WorkerHandle.scala b/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/WorkerHandle.scala new file mode 100644 index 0000000000000..c66cb5689b875 --- /dev/null +++ b/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/WorkerHandle.scala @@ -0,0 +1,66 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.udf.worker.core + +import org.apache.spark.annotation.Experimental + +/** + * :: Experimental :: + * The lifecycle-side handle the engine uses to release a worker back to + * its dispatcher. + * + * Sessions hold a [[WorkerHandle]] -- not the concrete worker type -- + * so the engine-facing layer (the `core` package) is decoupled from any + * specific worker provisioning model (direct local spawn, indirect + * daemon-provided, etc.). Each [[WorkerDispatcher]] supplies its own + * implementation; the session calls back into it on close. + * + * Contract: + * - [[releaseSession]] decrements the dispatcher-side reference count. + * The caller ([[WorkerSession.close]]) invokes it exactly once per + * session, so implementations need not be idempotent and MAY fail fast + * on an unbalanced (negative) count. Sessions sharing the same worker + * call it independently (each exactly once). + * - [[markInvalid]] is sticky: once set, the dispatcher MUST NOT return + * this worker to any reuse pool. Sessions set this when the worker + * is observed in an unsafe state (transport error, hung worker, + * protocol violation). + * - [[id]] is a stable string for diagnostics. The session uses it in + * error messages; the dispatcher uses it to key the worker in its + * own tracking map. Implementations must keep it constant for the + * lifetime of the handle. + */ +@Experimental +trait WorkerHandle { + + /** Stable identifier for logs and diagnostic messages. */ + def id: String + + /** + * Marks the worker as unsafe to recycle. Idempotent and sticky. + * Sessions call this in their close path when the protocol layer + * reports an unsalvageable termination. + */ + def markInvalid(): Unit + + /** + * Decrements the dispatcher-side session ref count. The caller + * ([[WorkerSession.close]]) invokes this exactly once per session, so + * implementations need not be idempotent. + */ + def releaseSession(): Unit +} diff --git a/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/WorkerSession.scala b/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/WorkerSession.scala index ef18ce8d70853..6ee9b129ac945 100644 --- a/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/WorkerSession.scala +++ b/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/WorkerSession.scala @@ -16,10 +16,13 @@ */ package org.apache.spark.udf.worker.core -import java.util.concurrent.atomic.AtomicBoolean +import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference} + +import scala.util.control.NonFatal import org.apache.spark.annotation.Experimental -import org.apache.spark.udf.worker.Init +import org.apache.spark.udf.worker.{Cancel, DataRequest, DataResponse, Finish, Init, + InitResponse} /** * :: Experimental :: @@ -28,145 +31,441 @@ import org.apache.spark.udf.worker.Init * A [[WorkerSession]] is the '''per-UDF-invocation''' handle that Spark * obtains from [[WorkerDispatcher#createSession]]. It carries the full * init / data-stream / finish lifecycle for a single UDF evaluation. - * - * A [[WorkerSession]] does ''not'' own the underlying worker or its - * transport channel -- those are managed by the [[WorkerDispatcher]]. - * Multiple sessions may share the same worker when the worker supports - * concurrency. + * Concrete protocol implementations subclass this and provide the + * protocol-specific request/response I/O; the base class owns the + * [[SessionState]] machine and handles dispatcher-side cleanup via + * [[WorkerHandle]]. * * '''Usage:''' * {{{ * val session = dispatcher.createSession(securityScope = None) + * val initResponse = session.init(Init.newBuilder()...build()) + * val results = session.process(inputBatches, () => Finish.newBuilder()...build()) * try { - * session.init(Init.newBuilder() - * .setProtocolVersion(1) - * .setUdf(UdfPayload.newBuilder().setPayload(callable).setFormat(fmt).build()) - * .setDataFormat(UDFWorkerDataFormat.ARROW) - * .build()) - * val results = session.process(inputBatches) - * results.foreach(handleBatch) + * results.foreach(handleBatch) // may throw on a UDF / data-phase error * } finally { - * session.close() + * val termination = session.close(() => Cancel.newBuilder()...build()) * } * }}} * * '''Lifecycle:''' - * - [[init]] must be called exactly once before [[process]]. - * - [[process]] must be called at most once per session. - * - [[close]] must always be called (use try-finally). - * - [[cancel]] may be called at any time from any execution context. - * See [[cancel]] for the full contract. + * - [[init]] must be called exactly once before [[process]]; it returns the + * worker's [[InitResponse]] (init-callback result + metrics). + * - [[process]] must be called at most once per session. It returns the + * result iterator; consuming that iterator drives the data exchange and + * surfaces any UDF / data-phase error (as an exception from `next()` / + * `hasNext`). `Finish` is sent (via the supplied `finish` thunk) when the + * input iterator is exhausted, so the result iterator yields the trailing + * output before ending. + * - [[close]] must always be called (use try-finally). It finalizes the + * session: if the stream already finished it returns the [[Finish]] + * terminator, otherwise it cancels anything still in flight (sending the + * [[Cancel]] from the supplied thunk so the cancel callback is carried) and + * waits for the terminator. It returns the [[Termination]] either way. + * + * '''State machine.''' A single [[SessionState]] reference tracks the whole + * session -- the call-ordering contract, the protocol phase, and the terminal + * outcome -- so there is one source of truth rather than separate lifecycle and + * wire machines: + * {{{ + * Created + * | init() + * v + * Initializing --(InitResponse ok)--> Initialized + * | | process() + * | init failure / transport v + * | Streaming + * | | input exhausted (Finish sent) + * | v + * | Finishing + * | | terminator + * | v + * +---------------------------------> (terminal) + * ^ + * | terminator + * Cancelling + * ^ + * | close() / cancel / ErrorResponse + * | (from any non-terminal state) + * }}} + * The clean path (via Finishing) and the cancel path (via Cancelling) settle the + * same `(terminal)`. The four terminals are: + * {{{ + * Finished(FinishResponse) | Cancelled(CancelResponse) + * Failed(ExecutionError) | TransportFailed(Throwable) + * }}} + * The two clean terminals carry the worker's `FinishResponse` / `CancelResponse` + * (metrics + finish/cancel callback `data`/`error`); the failure terminals carry + * the cause. Reaching any terminal means the session is over -- no further + * protocol writes are valid and [[close]] derives its [[Termination]] from it. + * The terminal states subsume what would otherwise be a separate "closed" state: + * [[close]] settles a terminal (cancelling first if needed) rather than moving to + * a distinct closed marker. + * + * '''Who drives transitions.''' Transitions are split between this base and the + * subclass, but all act on the one [[SessionState]]: + * - the base drives the API-call edges -- `Created -> Initializing` in [[init]] + * and `Initialized -> Streaming` in [[process]] -- with a single atomic + * transition before delegating to [[doInit]] / [[doProcess]]; an out-of-order + * or repeat call fails the transition and throws `IllegalStateException`; + * - the subclass drives the protocol-event edges as it exchanges messages + * (`Initializing -> Initialized` on `InitResponse`, `Streaming -> Finishing` + * on `Finish`, any non-terminal `-> Cancelling` on `Cancel`, and every + * terminal) through [[compareAndSetState]] and [[completeTerminal]]. + * + * [[close]] is idempotent (it finalizes once and releases the worker exactly + * once); a failed init leaves a terminal state, so a subsequent [[process]] is + * rejected just like a call after [[close]]. * - * The lifecycle is enforced here: [[init]] and [[process]] are `final` - * and delegate to [[doInit]] / [[doProcess]] after AtomicBoolean guards. - * Subclasses implement the protocol-specific work and do not re-check - * the contract. + * @param workerHandle dispatcher-side handle used to release the worker + * on [[close]]. Package-private to the udf-worker + * modules (`private[worker]`): the handle is a dispatcher + * implementation detail. + * @param logger logger for lifecycle diagnostics. */ @Experimental -abstract class WorkerSession extends AutoCloseable { +abstract class WorkerSession( + private[worker] val workerHandle: WorkerHandle, + protected val logger: WorkerLogger) { + + import WorkerSession.SessionState + + require(workerHandle != null, "workerHandle is required") + + // Single source of truth for the session: call-ordering contract, protocol + // phase, and terminal outcome. CAS transitions make the contract independent + // of caller threading, serialise the base's API-call edges against the + // subclass's protocol-event edges, and let the first terminal win. See the + // class-level "State machine" note. + private val state = new AtomicReference[SessionState](SessionState.Created) - private val initialized = new AtomicBoolean(false) - private val processed = new AtomicBoolean(false) + // close() is the finalizer; this records whether it has run so the worker + // handle is released exactly once. Kept separate from `state` because the + // terminal that close() settles describes the protocol outcome, not "close was + // called" -- a session can reach a terminal (finished / failed) before close(). + private val closeInvoked = new AtomicBoolean(false) /** * Initializes the UDF execution. Must be called exactly once before - * [[process]]. - * - * Throws `IllegalStateException` if called more than once. + * [[process]]. Returns the worker's [[InitResponse]] (carrying the + * init-callback result and metrics) on success; throws if the worker + * reports an init error or the stream fails before init completes. * * @param message the [[Init]] message carrying the UDF body, data * format, optional schemas, and any session context * the worker needs to start processing. */ - final def init(message: Init): Unit = { - if (!initialized.compareAndSet(false, true)) { - throw new IllegalStateException("init has already been called on this session") + final def init(message: Init): InitResponse = { + if (!state.compareAndSet(SessionState.Created, SessionState.Initializing)) { + throw new IllegalStateException( + s"init must be called exactly once before process (current state: ${state.get()})") } - doInit(message) + val response = doInit(message) + // doInit returned normally, so the worker accepted the session: ensure the + // state reflects that. A subclass MAY advance `Initializing -> Initialized` + // itself (e.g. to mark init resolved the moment `InitResponse` arrives), in + // which case this is a no-op; and a terminal that raced in keeps the CAS + // from clobbering it. On failure doInit throws and we never get here. + state.compareAndSet(SessionState.Initializing, SessionState.Initialized) + response } /** - * Processes input data through the worker and returns results. + * Processes input data through the worker and returns the result iterator. * * Follows Spark's Iterator-to-Iterator pattern: input batches are streamed * to the worker, and result batches are lazily pulled from the returned - * iterator. The session sends a finish signal to the worker when the input - * iterator is exhausted. + * iterator. When the input iterator is exhausted the session sends `Finish` + * -- built lazily from `finish` so the caller can carry a finish callback -- + * and the result iterator yields any trailing output before ending. A UDF or + * data-phase error surfaces as an exception while the result iterator is + * consumed. * - * Must be called after [[init]] and at most once per session. - * Throws `IllegalStateException` if called before [[init]] or more than once. + * Must be called after [[init]] has succeeded and at most once per session. + * Throws `IllegalStateException` if called before [[init]] completes, after + * the session has terminated (including a failed init), or more than once. * - * @param input iterator of raw input data batches (e.g., Arrow IPC) - * @return iterator of raw result data batches + * @param input iterator of input data batches ([[DataRequest]], each + * wrapping e.g. an Arrow IPC payload). + * @param finish supplies the [[Finish]] message sent once input is + * exhausted; invoked at most once. Defaults to an empty + * [[Finish]], which is sufficient when the caller has no extra + * information to convey to the worker on finish. + * @return iterator of result data batches ([[DataResponse]]). */ - final def process(input: Iterator[Array[Byte]]): Iterator[Array[Byte]] = { - if (!initialized.get()) { - throw new IllegalStateException("process called before init") + final def process( + input: Iterator[DataRequest], + finish: () => Finish = () => Finish.getDefaultInstance): Iterator[DataResponse] = { + if (!state.compareAndSet(SessionState.Initialized, SessionState.Streaming)) { + state.get() match { + case SessionState.Created | SessionState.Initializing => + throw new IllegalStateException("process called before init completed") + case s if s.isTerminal => + throw new IllegalStateException( + s"process called after the session terminated (state: $s)") + case _ => + throw new IllegalStateException("process has already been called on this session") + } } - if (!processed.compareAndSet(false, true)) { - throw new IllegalStateException("process has already been called on this session") + doProcess(input, finish) + } + + /** + * Finalizes the session and releases resources. Idempotent; safe to call + * from a `finally` block regardless of whether [[init]] or [[process]] ran. + * + * If the stream already finished (the result iterator was fully drained, so + * `Finish` has been sent and `FinishResponse` received), this returns the + * [[Termination.Finished]] terminator. Otherwise it cancels anything still in + * flight -- sending the [[Cancel]] produced by `cancel` so the cancel + * callback is carried -- waits for the terminator, and returns it. The + * `cancel` thunk is invoked only when a cancel actually needs to be sent. + * + * Does not raise application or transport failures -- they are reported through + * the returned [[Termination]], not thrown (though a thread interruption, or + * another fatal error, may still propagate). An execution error yields + * [[Termination.Failed]] and a transport failure / timeout yields + * [[Termination.TransportFailed]], each carrying the cause. A data-phase + * failure also surfaces while the [[process]] result iterator is consumed, but + * a failure that occurs during finish/close -- after + * the data is drained -- is reported only here, so the caller must inspect the + * [[Termination]]. After the terminator is settled the worker handle's ref + * count is decremented; if [[isWorkerSalvageable]] is false the handle is + * marked invalid so the dispatcher will not recycle the worker. + * + * @param cancel supplies the [[Cancel]] message used to abort in-flight work; + * invoked at most once, and only if the stream had not already + * finished. Defaults to an empty [[Cancel]], which is sufficient + * when the caller has no extra information to convey to the + * worker on cancellation (e.g. a reason). + * @return the [[Termination]] the stream settled on. + */ + final def close(cancel: () => Cancel = () => Cancel.getDefaultInstance): Termination = { + val firstClose = closeInvoked.compareAndSet(false, true) + try { + try { + doClose(cancel) + } catch { + case NonFatal(e) => + // doClose is contracted not to throw -- it reports its outcome via the + // settled terminal. A NonFatal escape is a subclass bug; uphold the + // finalizer's must-not-throw contract by settling a failure terminal + // here rather than letting it propagate out of close(). (Fatal errors + // and interrupts still propagate, per the contract.) + completeTerminal(Termination.TransportFailed(e)) + logger.warn("doClose threw unexpectedly; treating the worker as unsalvageable", e) + } + // Enforce the doClose post-condition: it must leave the session terminal, + // because isWorkerSalvageable (read below) classifies the worker off the + // settled SessionState. A non-terminal state here is a subclass contract + // violation. close() is the finalizer and must not throw, so rather than + // raise (or silently misclassify a healthy-looking worker as salvageable) + // we settle a failure terminal -- which makes the worker unsalvageable -- + // and log loudly. Terminals are sticky, so this is a no-op on the normal + // path where doClose already settled one. + if (!currentState.isTerminal) { + completeTerminal(Termination.TransportFailed(new IllegalStateException( + s"doClose returned without settling a terminal (state: ${currentState})"))) + logger.warn( + "doClose did not settle a terminal; treating the worker as unsalvageable") + } + // Derive the result from the settled terminal rather than trusting + // doClose's return value, so the Termination always agrees with the state + // close() classified the worker on -- in particular when the guard above + // had to override a misbehaving doClose. On the normal path this equals + // what doClose returned. + settledTermination + } finally { + if (firstClose) { + // close() is a finalizer (commonly invoked from a task-completion + // listener), so it avoids throwing: a misbehaving handle is logged via + // NonFatal, not propagated (a thread interrupt may still escape). + // isWorkerSalvageable reads the settled terminal (doClose, or the guard + // above, leaves the session terminal). releaseSession is still attempted + // even if markInvalid fails, so the ref count is not leaked; markInvalid + // runs first so a released worker is never recycled. + try { + if (!isWorkerSalvageable) workerHandle.markInvalid() + } catch { + case NonFatal(e) => + logger.warn(s"markInvalid for worker ${workerHandle.id} on close() failed", e) + } + try { + workerHandle.releaseSession() + } catch { + case NonFatal(e) => + logger.warn(s"releaseSession for worker ${workerHandle.id} on close() failed", e) + } + } } - doProcess(input) } /** - * Subclass hook for [[init]]. Called once, after the guard. - * Implementations MUST establish the worker connection here, not - * earlier. This ensures [[cancel]] before [[init]] is a no-op. + * Subclass hook for [[init]]. Called exactly once, after the state + * transition to `Initializing`. Implementations MUST establish the worker + * conversation here, not earlier, and advance `Initializing -> Initialized` + * (via [[compareAndSetState]]) once the worker accepts the session; on + * failure they settle a terminal (via [[completeTerminal]]) and throw. */ - protected def doInit(message: Init): Unit + protected def doInit(message: Init): InitResponse - /** Subclass hook for [[process]]. Called at most once, after the guard. */ - protected def doProcess(input: Iterator[Array[Byte]]): Iterator[Array[Byte]] + /** + * Subclass hook for [[process]]. Called at most once, after the state + * transition to `Streaming`. + */ + protected def doProcess( + input: Iterator[DataRequest], + finish: () => Finish): Iterator[DataResponse] /** - * Requests cancellation of the current UDF execution. + * Subclass hook for [[close]]. Settles the stream terminator (cancelling + * in-flight work via `cancel` if it has not already finished), releases any + * protocol resources, and returns the [[Termination]] (typically via + * [[settledTermination]]). * - * '''Thread-safety:''' [[cancel]] may be called concurrently with - * [[process]] from any execution context. + * '''Must settle a terminal before returning''' (via [[completeTerminal]] / + * [[settledTermination]]): [[close]] reads [[isWorkerSalvageable]] -- which + * inspects the settled [[SessionState]] -- the instant this returns, so a + * non-terminal state here would misclassify the worker. The base asserts this + * post-condition. * - * '''Lifecycle:''' [[cancel]] is idempotent and safe at any point in - * the session's life: - * - before [[init]] -- a no-op; the session may still be closed - * normally via [[close]]. - * - between [[init]] and [[process]] -- signals that the session - * should be terminated; the caller should not invoke [[process]] - * and should call [[close]] to release resources. - * Implementations SHOULD surface this as an error if [[process]] - * is subsequently invoked despite the cancellation. - * - during [[process]] (data flowing or awaiting completion) - * -- requests the worker to abort on a best-effort basis. - * - after [[process]] has returned (session already terminated) - * -- a no-op. + * '''Must be thread-safe and idempotent.''' [[close]] does not serialise calls + * to this hook: it can be invoked '''concurrently''' (e.g. a task-completion + * listener on one thread while the result iterator is consumed on another) and + * '''repeatedly''' (a second [[close]]). Implementations must tolerate both -- + * e.g. settle the terminal through the first-wins [[completeTerminal]] and + * guard any wire write under the same lock the data path uses. The `cancel` + * thunk must be evaluated at most once across all callers. * - * Implementations are responsible for the lifecycle-aware behavior - * described above so that the caller does not need to coordinate - * with the execution context driving [[process]]. + * As the cleanup path it does not re-throw a failed terminal -- it reports a + * best-effort [[Termination]] (the underlying failure surfaces through the + * result iterator). Should it nonetheless throw an unexpected (non-fatal) + * error, [[close]] converts that into a [[Termination.TransportFailed]] rather + * than propagating, to honour its must-not-throw contract. */ - def cancel(): Unit + protected def doClose(cancel: () => Cancel): Termination /** - * Closes this session and releases resources. Idempotent; safe to - * call from a `finally` block regardless of whether [[init]], - * [[process]], or [[cancel]] have been invoked. - * - * If [[init]] was called but [[process]] was not (e.g. an exception - * was thrown between the two), [[close]] signals cancellation to the - * worker before releasing resources so it can clean up - * deterministically. Subclasses implement [[doClose]] for resource - * teardown; the base class handles the cancel-before-close guarantee - * automatically. + * Whether the underlying worker is in a state safe to reuse after this + * session ends. The default treats only a dead or unknown transport as + * unsafe: a [[Termination.TransportFailed]] outcome (transport failure, + * timeout, or interrupt) -- or a session that never settled -- leaves the + * worker in an unknown state and is not salvageable. Every other terminal is + * salvageable: a clean [[Termination.Finished]], a cooperative + * [[Termination.Cancelled]], and also an execution [[Termination.Failed]], + * which is typically a user-code (UDF) error reported by a still-healthy + * worker rather than a worker fault. A `false` result tells [[close]] to mark + * `workerHandle` invalid so no reuse pool recycles the worker. Subclasses may + * override for protocol-specific nuances. + */ + protected def isWorkerSalvageable: Boolean = state.get() match { + case SessionState.Terminal(_: Termination.TransportFailed) => false + case t if t.isTerminal => true + case _ => false + } + + // ---- State-machine primitives for subclasses ------------------------------ + + /** The current [[SessionState]]. */ + protected final def currentState: SessionState = state.get() + + /** + * Atomically advances the state from `expect` to `update`, returning whether + * the transition was applied. Used by subclasses to drive the protocol-event + * edges (e.g. `Initializing -> Initialized`, `Streaming -> Finishing`). A + * terminal that arrived first makes the CAS fail, so the terminal wins. + * Terminals must be reached via [[completeTerminal]], not this method. + */ + protected final def compareAndSetState( + expect: SessionState, update: SessionState): Boolean = + state.compareAndSet(expect, update) + + /** + * Settles the session on its terminal [[Termination]] outcome. The first + * caller wins (an already-terminal state is left untouched); only the winner + * runs [[onTerminalSettled]] and this returns `true`. All later callers + * observe the existing terminal and return `false`. */ - final override def close(): Unit = { - if (initialized.get() && !processed.get()) { - cancel() + protected final def completeTerminal(termination: Termination): Boolean = { + val terminal = SessionState.Terminal(termination) + var cur = state.get() + while (!cur.isTerminal) { + if (state.compareAndSet(cur, terminal)) { + onTerminalSettled(termination) + return true + } + cur = state.get() } - doClose() + false + } + + /** + * Hook invoked exactly once, by the caller that settles the terminal in + * [[completeTerminal]], with the settled [[Termination]]. Subclasses override + * it to wake whatever is blocked on the terminator (output queues, latches). + * The default does nothing. + */ + protected def onTerminalSettled(termination: Termination): Unit = () + + /** + * The [[Termination]] the session settled on, for [[doClose]] to return. The + * clean terminators carry the worker's response proto (metrics + finish/cancel + * callback data/error); the failure terminations carry their cause and have no + * proto terminator, so no metrics. A failure also surfaces through the result + * iterator when it is consumed, but is carried here too because an error raised + * during finish/close -- after the data has been drained -- reaches no one + * through the iterator. Throws only if called before a terminal is settled, + * which [[doClose]] must not do. + */ + protected final def settledTermination: Termination = state.get() match { + case SessionState.Terminal(termination) => termination + case other => + throw new IllegalStateException(s"settledTermination called in non-terminal state: $other") } +} + +object WorkerSession { - /** Subclass hook for [[close]]. The base class guarantees that - * [[cancel]] has already been called if [[init]] was invoked but - * [[process]] was not. + /** + * The single state machine for a [[WorkerSession]]: the call-ordering + * contract, the protocol phase, and the terminal outcome. The normal + * progression is `Created -> Initializing -> Initialized -> Streaming -> + * Finishing -> Terminal`; `close`/cancel and errors fork to `Cancelling` + * before the `Terminal`. The single terminal state carries the public + * [[Termination]] outcome. See [[WorkerSession]] for the full diagram and + * which party drives each edge. + * + * Visible to the udf-worker modules (`private[worker]`) so concrete protocol + * subclasses can drive and inspect the state; it is an implementation detail, + * not part of the public API. */ - protected def doClose(): Unit + private[worker] sealed trait SessionState { + /** Whether this is a terminal state (the session is over). */ + def isTerminal: Boolean = false + } + + private[worker] object SessionState { + /** Constructed; [[WorkerSession.init]] not yet called. */ + case object Created extends SessionState + /** [[WorkerSession.init]] is in flight; awaiting the worker's `InitResponse`. */ + case object Initializing extends SessionState + /** `InitResponse` OK; [[WorkerSession.process]] not yet called. */ + case object Initialized extends SessionState + /** [[WorkerSession.process]] is active; input may still be sent. */ + case object Streaming extends SessionState + /** Input exhausted, `Finish` sent; awaiting trailing output + terminator. */ + case object Finishing extends SessionState + /** `Cancel` sent (engine cancel or post-error cleanup); awaiting terminator. */ + case object Cancelling extends SessionState + + /** + * The session is over; no further writes are valid. The single terminal + * state carries the public [[Termination]] outcome (`Finished` / + * `Cancelled` / `Failed` / `TransportFailed`), so those four outcome cases + * live once on [[Termination]] rather than being mirrored on the state. + */ + final case class Terminal(termination: Termination) extends SessionState { + override def isTerminal: Boolean = true + } + } } diff --git a/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/direct/DirectUnixSocketWorkerDispatcher.scala b/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/direct/DirectUnixSocketWorkerDispatcher.scala deleted file mode 100644 index 8da0354187e4f..0000000000000 --- a/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/direct/DirectUnixSocketWorkerDispatcher.scala +++ /dev/null @@ -1,145 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one or more - * contributor license agreements. See the NOTICE file distributed with - * this work for additional information regarding copyright ownership. - * The ASF licenses this file to You under the Apache License, Version 2.0 - * (the "License"); you may not use this file except in compliance with - * the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package org.apache.spark.udf.worker.core.direct - -import java.io.File -import java.nio.file.{Files, Path} -import java.nio.file.attribute.PosixFilePermissions - -import org.apache.spark.annotation.Experimental -import org.apache.spark.udf.worker.UDFWorkerSpecification -import org.apache.spark.udf.worker.core.{UnixSocketWorkerConnection, WorkerLogger} -import org.apache.spark.udf.worker.core.direct.DirectWorkerDispatcher.SOCKET_POLL_INTERVAL_MS - -/** - * :: Experimental :: - * A [[DirectWorkerDispatcher]] using Unix domain sockets as the worker - * transport. Allocates a private 0700 socket directory at construction; - * each worker is given a UDS path inside it. - * - * Concrete subclasses implement [[createConnection]] (with a UDS protocol - * of choice) and [[createSessionForWorker]]. - */ -@Experimental -abstract class DirectUnixSocketWorkerDispatcher( - workerSpec: UDFWorkerSpecification, - logger: WorkerLogger = WorkerLogger.NoOp) - extends DirectWorkerDispatcher(workerSpec, logger) { - - // Removed explicitly in closeTransport(). deleteOnExit is avoided because - // the JDK retains the path for the JVM lifetime, which leaks in - // long-lived drivers. - private val socketDir: Path = createPrivateTempDirectory() - - override protected def newEndpointAddress(workerId: String): String = - socketDir.resolve(s"worker-$workerId.sock").toString - - override protected def waitForReady( - address: String, - process: Process, - outputFile: File): Unit = { - val file = new File(address) - // At least one poll so very small initTimeouts don't trip a premature - // timeout before the worker has any chance to create the socket. - val maxAttempts = math.max(1, (initTimeoutMs / SOCKET_POLL_INTERVAL_MS).toInt) - var attempts = 0 - while (!file.exists() && attempts < maxAttempts) { - if (!process.isAlive) throwWorkerExitedBeforeSocket(process, address, outputFile) - Thread.sleep(SOCKET_POLL_INTERVAL_MS) - attempts += 1 - } - if (!file.exists()) { - if (process.isAlive) { - DirectWorkerDispatcher.destroyForciblyAndReap( - process, logger, s"init timeout $address") - val tail = readOutputTail(outputFile) - throw new DirectWorkerTimeoutException( - s"Worker did not create socket at $address within ${initTimeoutMs}ms\n$tail") - } else { - // Worker exited after the last poll without creating the socket; - // prefer the exit-code message over the ambiguous "did not create". - throwWorkerExitedBeforeSocket(process, address, outputFile) - } - } - } - - override protected def cleanupEndpointAddress(address: String): Unit = { - Files.deleteIfExists(new File(address).toPath) - } - - override protected def closeTransport(): Unit = { - val dir = socketDir.toFile - if (dir.exists()) { - val remaining = dir.listFiles() - if (remaining != null) remaining.foreach(_.delete()) - dir.delete() - } - } - - override protected def validateTransportSupport(): Unit = { - val props = workerSpec.getDirect.getProperties - require(props.hasConnection, - "DirectWorker.properties.connection must be set") - val conn = props.getConnection - require(conn.hasUnixDomainSocket, - "DirectUnixSocketWorkerDispatcher requires UNIX domain socket transport, " + - s"got ${conn.getTransportCase}") - } - - override protected def createConnection(address: String): UnixSocketWorkerConnection - - private def throwWorkerExitedBeforeSocket( - process: Process, - address: String, - outputFile: File): Nothing = { - val tail = readOutputTail(outputFile) - throw new DirectWorkerException( - s"Worker exited with code ${process.exitValue()} " + - s"before creating socket at $address\n$tail") - } - - /** - * Creates a temp directory with owner-only permissions (0700 on POSIX). - * On non-POSIX filesystems falls back to best-effort `File.setXxx`, - * which is TOCTOU-racy and weaker; a WARN surfaces if the platform - * refuses the setters. - */ - private def createPrivateTempDirectory(): Path = { - val attr = PosixFilePermissions.asFileAttribute( - PosixFilePermissions.fromString("rwx------")) - try { - Files.createTempDirectory("spark-udf-worker", attr) - } catch { - case _: UnsupportedOperationException => - val dir = Files.createTempDirectory("spark-udf-worker") - val f = dir.toFile - // `&` (non-short-circuiting) so every setter is attempted even if - // an earlier one refused. - val applied = - f.setReadable(false, false) & f.setWritable(false, false) & - f.setExecutable(false, false) & f.setReadable(true, true) & - f.setWritable(true, true) & f.setExecutable(true, true) - if (!applied) { - logger.warn( - s"Could not fully restrict permissions on $dir; socket " + - s"directory may be accessible to other local users on this " + - s"filesystem") - } - dir - } - } -} diff --git a/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/direct/DirectWorkerDispatcher.scala b/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/direct/DirectWorkerDispatcher.scala index 2438546ddafb8..db81b5265b794 100644 --- a/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/direct/DirectWorkerDispatcher.scala +++ b/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/direct/DirectWorkerDispatcher.scala @@ -29,7 +29,7 @@ import scala.util.control.NonFatal import org.apache.spark.annotation.Experimental import org.apache.spark.udf.worker.{ProcessCallable, UDFWorkerSpecification} -import org.apache.spark.udf.worker.core.{WorkerConnection, WorkerDispatcher, +import org.apache.spark.udf.worker.core.{WorkerConnection, WorkerDispatcher, WorkerHandle, WorkerLogger, WorkerSecurityScope, WorkerSession} import org.apache.spark.udf.worker.core.direct.DirectWorkerDispatcher.{CallableResult, DEFAULT_CALLABLE_TIMEOUT_MS, DEFAULT_GRACEFUL_TIMEOUT_MS, DEFAULT_INIT_TIMEOUT_MS, @@ -46,18 +46,17 @@ import org.apache.spark.udf.worker.core.direct.DirectWorkerDispatcher.{CallableR * currently gets a fresh worker that is terminated when the session closes * (the single-reference case of the future pooling policy). * - * Subclasses implement [[createConnection]] and [[createSessionForWorker]] - * to provide protocol-specific behavior (e.g., gRPC, raw sockets). + * Subclasses pick the transport and protocol: they allocate the endpoint + * address, wait for the worker to be reachable, build the + * [[WorkerConnection]] and the per-session [[WorkerSession]]. See + * [[newEndpointAddress]], [[newConnection]], [[newSession]]. * * For workers obtained through a provisioning service or daemon (indirect * creation), see the `indirect` package (TODO). * - * @param workerSpec worker specification (proto) - * @param logger [[WorkerLogger]] used for dispatcher-internal messages. - * The framework does not depend on any concrete logging - * backend; callers should pass an adapter that forwards - * to their preferred logger (Spark's `Logging` trait, - * SLF4J, etc.). Defaults to [[WorkerLogger.NoOp]]. + * @param workerSpec worker specification (proto). + * @param logger [[WorkerLogger]] for dispatcher-internal messages. + * Defaults to [[WorkerLogger.NoOp]]. */ @Experimental abstract class DirectWorkerDispatcher( @@ -68,8 +67,22 @@ abstract class DirectWorkerDispatcher( // TODO: Connection pooling -- reuse idle workers across sessions. // TODO: Security scope isolation -- partition pool by WorkerSecurityScope. - validateTransportSupport() + // Pre-flight spec validation. Per convention in this dispatcher, all + // spec-shape validation throws `IllegalArgumentException` (via `require` + // in the validators below); runtime failures during environment + // preparation throw `DirectWorkerException`. Callers can rely on this + // split to decide between programmer-error and operational-error paths. + // `workerSpec` is passed explicitly to `validateTransportSupport` so + // subclass implementations cannot read partially-initialized subclass + // fields: at this point in construction, the parent constructor body is + // still running and subclass `val`s have not yet been assigned. + validateTransportSupport(workerSpec) validateEnvironmentCallables() + // Transport-specific setup runs only after validation passes, so an + // invalid spec allocates no resources. Subclasses override [[initialize]] + // rather than relying on field-initialiser ordering relative to this + // constructor. + initialize() /** * Maximum time to wait for a setup/verify/cleanup callable to finish. @@ -82,7 +95,7 @@ abstract class DirectWorkerDispatcher( // dispatcher-internal callableTimeoutMs above is subclass-controlled and // not subject to the cap. // Package-private for test access. - private[core] val initTimeoutMs: Long = { + private[worker] val initTimeoutMs: Long = { val props = workerSpec.getDirect.getProperties val raw = if (props.hasInitializationTimeoutMs && props.getInitializationTimeoutMs > 0) { props.getInitializationTimeoutMs.toLong @@ -117,6 +130,10 @@ abstract class DirectWorkerDispatcher( private[this] val workers = new ConcurrentHashMap[String, DirectWorkerProcess]() private[this] val closed = new AtomicBoolean(false) + // TODO [SPARK-55278]: extract the env state machine + JVM shutdown hook + // into a standalone EnvironmentManager once the gRPC dispatcher work has + // landed; the env machinery is unrelated to the gRPC protocol and + // only lives here for historical reasons. @volatile private var environmentState: EnvironmentState = EnvironmentState.Pending private val environmentLock = new Object private[this] var cleanupHook: Option[Thread] = None @@ -150,17 +167,74 @@ abstract class DirectWorkerDispatcher( protected def closeTransport(): Unit /** - * Validates the worker spec's transport choice. Subclasses declare - * which transports they support. Called from the base constructor; - * implementations must only read base-class state (`workerSpec`). + * Validates the worker spec from the dispatcher's point of view -- both + * the transport choice (which transports can this dispatcher provision?) + * and the protocol's spec-level requirements (e.g. capabilities flags). + * + * Invoked from the base-class constructor BEFORE subclass `val`s have + * been initialised. Implementations MUST validate against `spec` only; + * reading subclass fields will see uninitialised values. + * + * Convention: throw `IllegalArgumentException` (via `require`) for spec + * problems. Operational/install failures are reported separately via + * `DirectWorkerException` from `ensureEnvironmentReady`. + * + * @param spec the worker specification to validate. Identical to the + * `workerSpec` field but passed explicitly to make the + * "constructor-time validation" contract obvious at call + * sites and prevent subclass-state reads. */ - protected def validateTransportSupport(): Unit + protected def validateTransportSupport(spec: UDFWorkerSpecification): Unit - /** Creates a protocol-specific connection to a worker at the given address. */ - protected def createConnection(address: String): WorkerConnection + /** + * Transport-specific one-time setup, invoked from the base-class + * constructor exactly once, AFTER spec validation has passed and BEFORE + * any session is created. Subclasses override this to allocate + * transport resources (e.g. a private socket directory) instead of + * relying on the ordering of subclass field initialisers relative to + * the base constructor. Overrides MUST call `super.initialize()`. + * + * The base implementation does nothing. + */ + protected def initialize(): Unit = () + + /** + * Opens the transport-level connection to a worker reachable at + * `address` (typically the value returned by [[newEndpointAddress]]). + * Called once per worker; the resulting connection is shared by every + * session opened against that worker. + */ + protected def newConnection(address: String): WorkerConnection - /** Creates a protocol-specific session for the given worker. */ - protected def createSessionForWorker(worker: DirectWorkerProcess): WorkerSession + /** + * Constructs the per-invocation [[WorkerSession]] for a worker. + * Subclasses build the concrete session implementation (e.g. + * `GrpcWorkerSession` for gRPC over UDS) using `workerHandle` for + * dispatcher-side cleanup and `connection` for the wire transport. + */ + protected def newSession( + workerHandle: WorkerHandle, + connection: WorkerConnection): WorkerSession + + /** + * Test-only hook: invoked once per [[createSession]] after the worker + * has been spawned and registered in the dispatcher's `workers` map, + * but before [[newSession]] has been called. The default is a no-op. + * + * Tests override to capture the worker reference or to inject a + * deliberate block to drive race conditions between `createSession` + * and `close`. Production code should not override this -- the + * production path provides no useful extension point here and the + * race-window control it offers is only meaningful for unit tests. + * + * Implementations MAY block. They run on the createSession caller's + * thread. + * + * Visible for testing only: `protected` so tests in this package can + * override it; production code should not. (The Guava annotation for this + * is banned by scalastyle per SPARK-11615, so the intent is noted here.) + */ + protected def afterWorkerRegistered(worker: DirectWorkerProcess): Unit = () override def createSession( securityScope: Option[WorkerSecurityScope]): WorkerSession = { @@ -180,7 +254,8 @@ abstract class DirectWorkerDispatcher( throwClosed() } try { - createSessionForWorker(worker) + afterWorkerRegistered(worker) + newSession(worker, worker.connection) } catch { case e: InterruptedException => Thread.currentThread().interrupt() @@ -343,7 +418,7 @@ abstract class DirectWorkerDispatcher( * Runs a [[ProcessCallable]] synchronously and returns the result. * Always throws on timeout; callers check `exitCode` for non-timeout failures. */ - private[core] def runCallable(callable: ProcessCallable): CallableResult = { + private[worker] def runCallable(callable: ProcessCallable): CallableResult = { val cmd = (callable.getCommandList.asScala ++ callable.getArgumentsList.asScala).toSeq require(cmd.nonEmpty, "ProcessCallable must have at least one entry in command or arguments") @@ -383,8 +458,11 @@ abstract class DirectWorkerDispatcher( try { waitForReady(address, process, outputFile.toFile) - val connection = createConnection(address) + val connection = newConnection(address) val artifacts = new WorkerArtifacts(process, connection, outputFile, logger) + // Remove the worker's endpoint artifact (its UDS socket file) on close; + // the worker creates it, the dispatcher owns its deletion. + artifacts.registerCleanup(() => cleanupEndpointAddress(address)) new DirectWorkerProcess( workerId, artifacts, gracefulTimeoutMs, logger, onLastSessionReleased = releaseWorker) @@ -471,8 +549,10 @@ abstract class DirectWorkerDispatcher( } } -private[direct] object DirectWorkerDispatcher { - private[direct] val SOCKET_POLL_INTERVAL_MS = 100L +// Visible to `core` (and the sibling `grpc` module) so concrete dispatchers in +// the `grpc` module can call shared helpers like `destroyForciblyAndReap`. +private[worker] object DirectWorkerDispatcher { + private[worker] val SOCKET_POLL_INTERVAL_MS = 100L private[direct] val DEFAULT_INIT_TIMEOUT_MS = 10000L private[direct] val DEFAULT_CALLABLE_TIMEOUT_MS = 120000L private[direct] val DEFAULT_GRACEFUL_TIMEOUT_MS = 5000L @@ -499,7 +579,7 @@ private[direct] object DirectWorkerDispatcher { * @param context short tag included in the timeout warning so operators * can correlate a stuck child with its source. */ - private[direct] def destroyForciblyAndReap( + private[worker] def destroyForciblyAndReap( process: Process, logger: WorkerLogger, context: String = ""): Unit = { @@ -522,7 +602,7 @@ private[direct] object DirectWorkerDispatcher { } /** Result of running a [[ProcessCallable]]. */ - private[core] case class CallableResult(exitCode: Int, outputTail: String) + private[worker] case class CallableResult(exitCode: Int, outputTail: String) private[direct] sealed trait EnvironmentState private[direct] object EnvironmentState { diff --git a/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/direct/DirectWorkerProcess.scala b/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/direct/DirectWorkerProcess.scala index f4b5c1df63193..52e7d5123d81b 100644 --- a/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/direct/DirectWorkerProcess.scala +++ b/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/direct/DirectWorkerProcess.scala @@ -17,13 +17,13 @@ package org.apache.spark.udf.worker.core.direct import java.nio.file.{Files, Path} -import java.util.concurrent.TimeUnit +import java.util.concurrent.{CopyOnWriteArrayList, TimeUnit} import java.util.concurrent.atomic.{AtomicBoolean, AtomicInteger} import scala.util.control.NonFatal import org.apache.spark.annotation.Experimental -import org.apache.spark.udf.worker.core.{WorkerConnection, WorkerLogger} +import org.apache.spark.udf.worker.core.{WorkerConnection, WorkerHandle, WorkerLogger} /** * :: Experimental :: @@ -33,7 +33,7 @@ import org.apache.spark.udf.worker.core.{WorkerConnection, WorkerLogger} * future pooling -- today one process per session. * * Closing sends SIGTERM, waits up to [[gracefulTimeoutMs]], then - * delegates connection close + forced kill + file cleanup to + * delegates connection close + forced kill + socket/log cleanup to * [[WorkerArtifacts.close]]. * * @param id stable worker identifier (UUID passed to the binary as `--id`). @@ -48,18 +48,27 @@ import org.apache.spark.udf.worker.core.{WorkerConnection, WorkerLogger} */ @Experimental class DirectWorkerProcess( - val id: String, + override val id: String, private[direct] val artifacts: WorkerArtifacts, val gracefulTimeoutMs: Long, protected val logger: WorkerLogger = WorkerLogger.NoOp, private[direct] val onLastSessionReleased: DirectWorkerProcess => Unit = _ => ()) - extends AutoCloseable { + extends AutoCloseable with WorkerHandle { // TODO: idle-timeout tracking and concurrent session capacity. private val activeSessionCount = new AtomicInteger(0) private val closed = new AtomicBoolean(false) + // Reuse-readiness flag. Sessions set this to `true` when they observe the + // worker in a state that is unsafe to recycle (transport error, hung + // worker, observed protocol violation). Pool-aware dispatchers consult + // [[isInvalid]] in [[onLastSessionReleased]] to decide between + // terminate-and-discard vs. return-to-pool. Today the non-pooling path + // always terminates, so this is forward-looking; setting it now is + // harmless and locks in the reuse contract. + private val invalid = new AtomicBoolean(false) + /** The OS process handle for this worker. */ def process: Process = artifacts.process @@ -72,20 +81,59 @@ class DirectWorkerProcess( /** Number of sessions currently using this worker. */ def activeSessions: Int = activeSessionCount.get() - /** Increments the active session count. */ - def acquireSession(): Unit = activeSessionCount.incrementAndGet() + /** + * Increments the active session count. + * + * Preconditions: must not be called after [[close]]. The dispatcher + * arbitrates acquire-vs-close ordering; calling this on a closed worker + * indicates a dispatcher-side bug. + * + * @throws IllegalStateException if the worker has already been closed. + */ + def acquireSession(): Unit = { + if (closed.get()) { + throw new IllegalStateException( + s"cannot acquire session: worker $id is already closed") + } + activeSessionCount.incrementAndGet() + } + + /** + * Returns true if a session has marked this worker as unsafe to reuse. + * Pool implementations MUST treat [[isInvalid]] as a hard rejection in + * [[onLastSessionReleased]] -- the worker must be torn down, not + * recycled. Once set, the flag is sticky for the worker's lifetime. + */ + def isInvalid: Boolean = invalid.get() + + /** + * Marks this worker as unsafe to return to a reuse pool. Idempotent; + * sticky once set. Sessions call this when they observe a transport + * failure, hung worker, or any termination path that leaves the worker + * in an unknown state. + */ + override def markInvalid(): Unit = invalid.set(true) /** * Decrements the active session count. Fires [[onLastSessionReleased]] - * on the 0-transition. A negative count indicates an unbalanced - * acquire/release; we log and reset to 0 rather than silently mask it. + * on the 0-transition. + * + * A negative count means this was called without a matching + * [[acquireSession]] -- an unbalanced acquire/release, which is a + * dispatcher-side bug -- and throws `IllegalStateException`. + * + * @throws IllegalStateException if the active session count goes negative. */ - def releaseSession(): Unit = { + override def releaseSession(): Unit = { val c = activeSessionCount.decrementAndGet() if (c < 0) { - logger.warn( + // A negative count means releaseSession was called without a matching + // acquireSession -- an unbalanced acquire/release, which is a + // dispatcher-side bug. Fail fast: a corrupt count could later let a + // release prematurely fire onLastSessionReleased and tear the worker + // down under a live session. + throw new IllegalStateException( s"releaseSession called without a matching acquireSession (count=$c)") - activeSessionCount.set(0) } else if (c == 0) { // Swallow callback errors so session.close cannot throw. try onLastSessionReleased(this) catch { @@ -125,10 +173,15 @@ class DirectWorkerProcess( /** * Closeable bundle of per-worker OS resources: the child [[Process]], its * transport [[WorkerConnection]], and its merged stdout/stderr log. - * [[close]] runs connection close (which for UDS removes the socket - * file), then SIGKILL-reaps the process, then deletes the output log. - * Graceful SIGTERM is the higher layer's responsibility (see - * [[DirectWorkerProcess#close]]). + * [[close]] tears down the connection, SIGKILL-reaps the process, runs any + * cleanup hooks registered via [[registerCleanup]] (e.g. removing the worker's + * UDS socket file), then deletes the output log. Graceful SIGTERM is the higher + * layer's responsibility (see [[DirectWorkerProcess#close]]). + * + * Transport-specific teardown is attached via [[registerCleanup]] rather than + * baked into this bundle, so `WorkerArtifacts` stays transport-agnostic: a + * dispatcher registers whatever per-worker resources it created without this + * class knowing what they are. */ private[direct] final class WorkerArtifacts( val process: Process, @@ -138,11 +191,37 @@ private[direct] final class WorkerArtifacts( private[this] val closed = new AtomicBoolean(false) + // Resource-cleanup callbacks run during close(), in registration order, after + // the process is reaped. A CopyOnWriteArrayList gives thread-safe registration + // and lock-free iteration at close without an explicit lock; in practice all + // hooks are registered at spawn time, before the bundle is exposed to any + // close path. + private[this] val cleanupHooks = new CopyOnWriteArrayList[() => Unit]() + /** - * Idempotently closes the connection (transport teardown + any - * transport-specific cleanup such as deleting a UDS socket file), - * SIGKILL-reaps the process, and deletes the output log. Each step - * is guarded so a failure in one does not skip the next. + * Registers a resource-cleanup callback to run during [[close]], in + * registration order, after the process is reaped. Lets a dispatcher attach + * the per-worker resources it created -- e.g. `() => cleanupEndpointAddress(address)` + * to delete the UDS socket file -- without `WorkerArtifacts` knowing what they + * are. Each hook is invoked at most once and guarded independently during + * close, so one failing hook does not skip the rest. + * + * Must be called before [[close]]; registering on an already-closed bundle is + * a dispatcher-side bug and throws `IllegalStateException`. + */ + def registerCleanup(hook: () => Unit): Unit = { + if (closed.get()) { + throw new IllegalStateException( + "cannot register a cleanup hook on an already-closed WorkerArtifacts") + } + cleanupHooks.add(hook) + } + + /** + * Idempotently closes the connection (transport teardown), + * SIGKILL-reaps the process, runs the dispatcher cleanup hooks, and deletes + * the output log. Each step (and each hook) is guarded so a failure in one + * does not skip the next. */ override def close(): Unit = { if (!closed.compareAndSet(false, true)) return @@ -154,6 +233,16 @@ private[direct] final class WorkerArtifacts( DirectWorkerDispatcher.destroyForciblyAndReap(process, logger, "worker artifacts") + // `closed` is already set, so registerCleanup rejects any further hook; + // every hook registered before close is iterated here (CopyOnWriteArrayList + // iteration needs no lock and no snapshot copy). + cleanupHooks.forEach { hook => + try hook() catch { + case NonFatal(e) => + logger.warn("Error running worker resource-cleanup hook", e) + } + } + try Files.deleteIfExists(outputFile) catch { case NonFatal(e) => logger.warn(s"Error cleaning up worker output file $outputFile", e) diff --git a/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/direct/DirectWorkerSession.scala b/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/direct/DirectWorkerSession.scala deleted file mode 100644 index de1dc45b5a8d3..0000000000000 --- a/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/direct/DirectWorkerSession.scala +++ /dev/null @@ -1,56 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one or more - * contributor license agreements. See the NOTICE file distributed with - * this work for additional information regarding copyright ownership. - * The ASF licenses this file to You under the Apache License, Version 2.0 - * (the "License"); you may not use this file except in compliance with - * the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package org.apache.spark.udf.worker.core.direct - -import java.util.concurrent.atomic.AtomicBoolean - -import org.apache.spark.annotation.Experimental -import org.apache.spark.udf.worker.core.{WorkerConnection, WorkerSession} - -/** - * :: Experimental :: - * A [[WorkerSession]] backed by a locally-spawned [[DirectWorkerProcess]]. - * - * This is the session type returned by [[DirectWorkerDispatcher]]. It ties - * the session lifecycle to the worker's ref-count: the dispatcher increments - * the count before construction, and [[doClose]] decrements it, so the - * dispatcher knows when a worker process is idle and can be terminated or - * reused. - * - * Subclasses implement the protocol-specific data transmission - * ([[init]], [[process]], [[cancel]]). - * - * @param workerProcess the direct worker process backing this session. - * Internal to the `core` package and test code -- the - * worker handle is a dispatcher implementation detail, - * not part of the public WorkerSession API. - */ -@Experimental -abstract class DirectWorkerSession( - private[core] val workerProcess: DirectWorkerProcess) extends WorkerSession { - - private val released = new AtomicBoolean(false) - - /** The connection to the worker for this session. */ - def connection: WorkerConnection = workerProcess.connection - - override protected def doClose(): Unit = { - if (released.compareAndSet(false, true)) { - workerProcess.releaseSession() - } - } -} diff --git a/udf/worker/core/src/test/scala/org/apache/spark/udf/worker/core/DirectWorkerDispatcherSuite.scala b/udf/worker/core/src/test/scala/org/apache/spark/udf/worker/core/DirectWorkerDispatcherSuite.scala index 06574fdf1013b..28ecc1179f6ac 100644 --- a/udf/worker/core/src/test/scala/org/apache/spark/udf/worker/core/DirectWorkerDispatcherSuite.scala +++ b/udf/worker/core/src/test/scala/org/apache/spark/udf/worker/core/DirectWorkerDispatcherSuite.scala @@ -27,11 +27,12 @@ import org.scalatest.BeforeAndAfterEach import org.scalatest.funsuite.AnyFunSuite import org.apache.spark.udf.worker.{ - DirectWorker, LocalTcpConnection, ProcessCallable, UDFWorkerProperties, - UDFWorkerSpecification, UnixDomainSocket, WorkerConnectionSpec, - WorkerEnvironment} -import org.apache.spark.udf.worker.core.direct.{DirectUnixSocketWorkerDispatcher, - DirectWorkerException, DirectWorkerProcess, + Cancel, DirectWorker, LocalTcpConnection, ProcessCallable, UDFProtoCommunicationPattern, + UDFWorkerProperties, UDFWorkerSpecification, UnixDomainSocket, WorkerCapabilities, + WorkerConnectionSpec, WorkerEnvironment} +import org.apache.spark.udf.worker.core.{WorkerConnection, WorkerHandle, WorkerSecurityScope, + WorkerSession} +import org.apache.spark.udf.worker.core.direct.{DirectWorkerException, DirectWorkerProcess, DirectWorkerTimeoutException} /** @@ -71,8 +72,14 @@ class DirectWorkerDispatcherSuite private def directWorker(runner: ProcessCallable): DirectWorker = DirectWorker.newBuilder().setRunner(runner).setProperties(udsProperties).build() + // TestDirectWorkerDispatcher requires the spec to advertise BIDIRECTIONAL_STREAMING. + private def bidiCapabilities: WorkerCapabilities = WorkerCapabilities.newBuilder() + .addSupportedCommunicationPatterns(UDFProtoCommunicationPattern.BIDIRECTIONAL_STREAMING) + .build() + private def specWithRunner(runner: ProcessCallable): UDFWorkerSpecification = UDFWorkerSpecification.newBuilder() + .setCapabilities(bidiCapabilities) .setDirect(directWorker(runner)) .build() @@ -81,11 +88,16 @@ class DirectWorkerDispatcherSuite env: WorkerEnvironment): UDFWorkerSpecification = UDFWorkerSpecification.newBuilder() .setEnvironment(env) + .setCapabilities(bidiCapabilities) .setDirect(directWorker(runner)) .build() private var dispatcher: TestDirectWorkerDispatcher = _ + // Finalization message for session.close(); these lifecycle-only tests never + // drive the protocol, so an empty Cancel suffices. + private val emptyCancel: () => Cancel = () => Cancel.getDefaultInstance + override def afterEach(): Unit = { if (dispatcher != null) { dispatcher.close() @@ -94,36 +106,46 @@ class DirectWorkerDispatcherSuite super.afterEach() } - // Narrow the publicly-typed WorkerSession returned by `createSession` back - // down to StubWorkerSession in one place, with a descriptive failure if - // the cast is ever wrong, so individual tests don't scatter `asInstanceOf` - // (which would throw ClassCastException rather than a useful message). - private def createStubSession(): StubWorkerSession = - dispatcher.createSession(None) match { - case stub: StubWorkerSession => stub + // The dispatcher returns a NoOpWorkerSession over the dispatcher's + // workerProcess; tests reach through to workerHandle (a + // DirectWorkerProcess in this configuration) via the session's + // package-private field directly. + private def createStubSession(): WorkerSession = + dispatcher.createSession(None) + + // Centralised cast: every test in this suite uses TestDirectWorkerDispatcher, + // which always returns a DirectWorkerProcess. Wrapping in a helper with a + // clear failure message (a) deduplicates the boilerplate cast at every + // assertion site, and (b) means a future dispatcher that returns a + // different WorkerHandle implementation fails with a recognisable test + // assertion rather than a bare ClassCastException. + private def workerProcess(s: WorkerSession): DirectWorkerProcess = + s.workerHandle match { + case dp: DirectWorkerProcess => dp case other => fail( - s"Expected StubWorkerSession, got ${other.getClass.getSimpleName}") + s"Expected DirectWorkerProcess, got ${other.getClass.getSimpleName}") } - // The whole suite uses UDS as the only transport, so reaching past the - // generic WorkerConnection abstraction to read the socket path is fine. + // The whole suite uses the SocketFileConnection test fixture, so reaching + // past the generic WorkerConnection abstraction to read the socket path + // is fine. private def udsPath(w: DirectWorkerProcess): String = w.connection match { - case uds: UnixSocketWorkerConnection => uds.socketPath + case sfc: SocketFileConnection => sfc.socketPath case other => fail( - s"Expected UnixSocketWorkerConnection, got ${other.getClass.getSimpleName}") + s"Expected SocketFileConnection, got ${other.getClass.getSimpleName}") } test("creates a worker and session") { dispatcher = new TestDirectWorkerDispatcher(specWithRunner(defaultRunner)) val session = createStubSession() - val worker = session.workerProcess + val worker = workerProcess(session) assert(worker.isAlive, "worker should be alive after creation") assert(worker.activeSessions == 1, "should have 1 active session") assert(new File(udsPath(worker)).exists(), "socket file should exist") - session.close() + session.close(emptyCancel) assert(worker.activeSessions == 0, "should have 0 sessions after close") } @@ -131,7 +153,7 @@ class DirectWorkerDispatcherSuite dispatcher = new TestDirectWorkerDispatcher(specWithRunner(defaultRunner)) val threads = 8 - val sessions = new java.util.concurrent.ConcurrentLinkedQueue[StubWorkerSession]() + val sessions = new java.util.concurrent.ConcurrentLinkedQueue[WorkerSession]() val startGate = new java.util.concurrent.CountDownLatch(1) val doneGate = new java.util.concurrent.CountDownLatch(threads) val errors = new java.util.concurrent.ConcurrentLinkedQueue[Throwable]() @@ -157,7 +179,7 @@ class DirectWorkerDispatcherSuite assert(sessions.size == threads, "expected one session per thread") val sessionList = sessions.asScala.toList - val workerObjects = sessionList.map(_.workerProcess) + val workerObjects = sessionList.map(workerProcess) assert(workerObjects.distinct.length == threads, "each session should have its own DirectWorkerProcess") // Object-identity is not sufficient on its own: a future regression @@ -168,7 +190,7 @@ class DirectWorkerDispatcherSuite assert(socketPaths.distinct.length == threads, s"each worker should have its own socket path, got $socketPaths") - sessionList.foreach(_.close()) + sessionList.foreach(_.close(emptyCancel)) } test("close shuts down all workers via SIGTERM") { @@ -177,11 +199,11 @@ class DirectWorkerDispatcherSuite val session1 = createStubSession() val session2 = createStubSession() - val worker1 = session1.workerProcess - val worker2 = session2.workerProcess + val worker1 = workerProcess(session1) + val worker2 = workerProcess(session2) - session1.close() - session2.close() + session1.close(emptyCancel) + session2.close(emptyCancel) dispatcher.close() dispatcher = null @@ -217,17 +239,18 @@ class DirectWorkerDispatcherSuite .setGracefulTerminationTimeoutMs(500) .build() val spec = UDFWorkerSpecification.newBuilder() + .setCapabilities(bidiCapabilities) .setDirect(DirectWorker.newBuilder() .setRunner(runner).setProperties(shortGracefulProps).build()) .build() dispatcher = new TestDirectWorkerDispatcher(spec) val session = createStubSession() - val worker = session.workerProcess + val worker = workerProcess(session) assert(worker.process.isAlive, "worker should be alive before close") val closeStart = System.nanoTime() - session.close() + session.close(emptyCancel) val closeElapsedMs = (System.nanoTime() - closeStart) / 1000000L assert(!worker.process.isAlive, @@ -241,13 +264,13 @@ class DirectWorkerDispatcherSuite dispatcher = new TestDirectWorkerDispatcher(specWithRunner(defaultRunner)) val session = createStubSession() - val worker = session.workerProcess + val worker = workerProcess(session) val socketFile = new File(udsPath(worker)) assert(worker.process.isAlive, "worker should be alive before session close") assert(socketFile.exists(), "socket file should exist before session close") - session.close() + session.close(emptyCancel) // The session-close path is synchronous: SIGTERM is sent and the process // is reaped before `close` returns. @@ -261,7 +284,7 @@ class DirectWorkerDispatcherSuite dispatcher = new TestDirectWorkerDispatcher(specWithRunner(defaultRunner)) val sessions = (1 to 4).map(_ => createStubSession()) - val workers = sessions.map(_.workerProcess) + val workers = sessions.map(workerProcess) val barrier = new java.util.concurrent.CyclicBarrier(sessions.size + 1) val errors = new java.util.concurrent.ConcurrentLinkedQueue[Throwable]() @@ -270,7 +293,7 @@ class DirectWorkerDispatcherSuite val t = new Thread(() => { try { barrier.await() - s.close() + s.close(emptyCancel) } catch { case t: Throwable => errors.add(t) } @@ -312,12 +335,8 @@ class DirectWorkerDispatcherSuite val releaseLatch = new java.util.concurrent.CountDownLatch(1) val capturedWorkers = new java.util.concurrent.ConcurrentLinkedQueue[DirectWorkerProcess]() - val racing = new DirectUnixSocketWorkerDispatcher(specWithRunner(defaultRunner)) { - override protected def createConnection( - socketPath: String): UnixSocketWorkerConnection = - new SocketFileConnection(socketPath) - override protected def createSessionForWorker( - worker: DirectWorkerProcess): WorkerSession = { + val racing = new TestDirectWorkerDispatcher(specWithRunner(defaultRunner)) { + override protected def afterWorkerRegistered(worker: DirectWorkerProcess): Unit = { capturedWorkers.add(worker) readyLatch.countDown() // Block here so dispatcher.close() runs while createSession is in @@ -325,7 +344,6 @@ class DirectWorkerDispatcherSuite if (!releaseLatch.await(30, java.util.concurrent.TimeUnit.SECONDS)) { fail("releaseLatch never fired -- test orchestration broken") } - new StubWorkerSession(worker) } } try { @@ -344,7 +362,7 @@ class DirectWorkerDispatcherSuite // Wait for thread A to have published the worker and entered the // blocking override. assert(readyLatch.await(10, java.util.concurrent.TimeUnit.SECONDS), - "createSession thread never reached createSessionForWorker") + "createSession thread never reached afterWorkerRegistered") val closeThread = new Thread(() => racing.close(), "close-racer") closeThread.start() @@ -406,16 +424,17 @@ class DirectWorkerDispatcherSuite .setGracefulTerminationTimeoutMs(60000) .build() val spec = UDFWorkerSpecification.newBuilder() + .setCapabilities(bidiCapabilities) .setDirect(DirectWorker.newBuilder() .setRunner(defaultRunner).setProperties(oversizedProps).build()) .build() dispatcher = new TestDirectWorkerDispatcher(spec) val session = createStubSession() - assert(session.workerProcess.gracefulTimeoutMs == 30000L, - s"graceful timeout should be capped at 30000ms, " + - s"got ${session.workerProcess.gracefulTimeoutMs}") - session.close() + val worker = workerProcess(session) + assert(worker.gracefulTimeoutMs == 30000L, + s"graceful timeout should be capped at 30000ms, got ${worker.gracefulTimeoutMs}") + session.close(emptyCancel) } test("worker-provided init timeout is capped at the engine-side maximum") { @@ -425,6 +444,7 @@ class DirectWorkerDispatcherSuite .setInitializationTimeoutMs(60000) .build() val spec = UDFWorkerSpecification.newBuilder() + .setCapabilities(bidiCapabilities) .setDirect(DirectWorker.newBuilder() .setRunner(defaultRunner).setProperties(oversizedProps).build()) .build() @@ -451,8 +471,8 @@ class DirectWorkerDispatcherSuite // Drive one createSession so a worker (and therefore the socket dir) is // observable via the UDS connection's path. val session = createStubSession() - val socketDir: Path = new File(udsPath(session.workerProcess)).toPath.getParent - session.close() + val socketDir: Path = new File(udsPath(workerProcess(session))).toPath.getParent + session.close(emptyCancel) val view = Files.getFileAttributeView(socketDir, classOf[PosixFileAttributeView]) // Skip explicitly on non-POSIX filesystems rather than silently pass, @@ -469,10 +489,10 @@ class DirectWorkerDispatcherSuite test("socket directory is removed after dispatcher.close") { dispatcher = new TestDirectWorkerDispatcher(specWithRunner(defaultRunner)) val session = createStubSession() - val socketDir = new File(udsPath(session.workerProcess)).toPath.getParent.toFile + val socketDir = new File(udsPath(workerProcess(session))).toPath.getParent.toFile assert(socketDir.exists(), s"socket directory $socketDir should exist while a session is open") - session.close() + session.close(emptyCancel) dispatcher.close() dispatcher = null @@ -483,19 +503,18 @@ class DirectWorkerDispatcherSuite // -- Error-path tests ------------------------------------------------------- - test("worker is cleaned up when createSessionForWorker throws") { - // A dispatcher whose createSessionForWorker always throws. The spawned - // worker must be terminated rather than leaked until dispatcher.close(). + test("worker is cleaned up when newSession throws") { + // A dispatcher whose newSession always throws. The spawned worker + // must be terminated rather than leaked until dispatcher.close(). var capturedWorker: DirectWorkerProcess = null val failingDispatcher = - new DirectUnixSocketWorkerDispatcher(specWithRunner(defaultRunner)) { - override protected def createConnection( - socketPath: String): UnixSocketWorkerConnection = - new SocketFileConnection(socketPath) - override protected def createSessionForWorker( - worker: DirectWorkerProcess): WorkerSession = { - capturedWorker = worker + new TestDirectWorkerDispatcher(specWithRunner(defaultRunner)) { + override protected def newSession( + workerHandle: WorkerHandle, + connection: WorkerConnection): WorkerSession = throw new RuntimeException("session creation failed") + override protected def afterWorkerRegistered(w: DirectWorkerProcess): Unit = { + capturedWorker = w } } @@ -541,25 +560,22 @@ class DirectWorkerDispatcherSuite s"expected UDS-only error, got: ${ex.getMessage}") } - test("socket file is cleaned up when createConnection throws") { + test("socket file is cleaned up when newConnection throws") { val capturedSocketPaths = new java.util.concurrent.ConcurrentLinkedQueue[String]() val failingDispatcher = - new DirectUnixSocketWorkerDispatcher(specWithRunner(defaultRunner)) { - override protected def createConnection( - socketPath: String): UnixSocketWorkerConnection = { - capturedSocketPaths.add(socketPath) + new TestDirectWorkerDispatcher(specWithRunner(defaultRunner)) { + override protected def newConnection(address: String): WorkerConnection = { + capturedSocketPaths.add(address) throw new RuntimeException("connection creation failed") } - override protected def createSessionForWorker( - worker: DirectWorkerProcess): WorkerSession = - new StubWorkerSession(worker) } try { val ex = intercept[RuntimeException] { failingDispatcher.createSession(None) } assert(ex.getMessage.contains("connection creation failed")) - assert(capturedSocketPaths.size == 1, "createConnection should have been called once") + assert(capturedSocketPaths.size == 1, + "newConnection should have been called once") val socketPath = capturedSocketPaths.peek() assert(!new File(socketPath).exists(), s"socket file $socketPath should have been cleaned up") @@ -608,6 +624,7 @@ class DirectWorkerDispatcherSuite .setInitializationTimeoutMs(500) .build() val spec = UDFWorkerSpecification.newBuilder() + .setCapabilities(bidiCapabilities) .setDirect(DirectWorker.newBuilder() .setRunner(hangingRunner).setProperties(shortInitProps).build()) .build() @@ -638,7 +655,7 @@ class DirectWorkerDispatcherSuite dispatcher = new TestDirectWorkerDispatcher(specWithEnv(env = env)) val session = dispatcher.createSession(None) - session.close() + session.close(emptyCancel) assert(!markerFile.exists(), "installation should not run when verification succeeds") @@ -658,7 +675,7 @@ class DirectWorkerDispatcherSuite dispatcher = new TestDirectWorkerDispatcher(specWithEnv(env = env)) val session = dispatcher.createSession(None) - session.close() + session.close(emptyCancel) assert(markerFile.exists(), "installation should run when verification fails") @@ -677,7 +694,7 @@ class DirectWorkerDispatcherSuite dispatcher = new TestDirectWorkerDispatcher(specWithEnv(env = env)) val session = dispatcher.createSession(None) - session.close() + session.close(emptyCancel) assert(markerFile.exists(), "installation should run when no verification is defined") @@ -710,14 +727,8 @@ class DirectWorkerDispatcherSuite .addCommand("sleep 30").build() val env = WorkerEnvironment.newBuilder().setInstallation(slowInstall).build() val shortTimeoutDispatcher = - new DirectUnixSocketWorkerDispatcher(specWithEnv(env = env)) { + new TestDirectWorkerDispatcher(specWithEnv(env = env)) { override protected def callableTimeoutMs: Long = 500L - override protected def createConnection( - socketPath: String): UnixSocketWorkerConnection = - new SocketFileConnection(socketPath) - override protected def createSessionForWorker( - worker: DirectWorkerProcess): WorkerSession = - new StubWorkerSession(worker) } try { val ex = intercept[DirectWorkerTimeoutException] { @@ -743,8 +754,8 @@ class DirectWorkerDispatcherSuite .build() dispatcher = new TestDirectWorkerDispatcher(specWithEnv(env = env)) - val s1 = dispatcher.createSession(None); s1.close() - val s2 = dispatcher.createSession(None); s2.close() + val s1 = dispatcher.createSession(None); s1.close(emptyCancel) + val s2 = dispatcher.createSession(None); s2.close(emptyCancel) val src = scala.io.Source.fromFile(counterFile) val lines = try src.getLines().toList finally src.close() @@ -799,7 +810,7 @@ class DirectWorkerDispatcherSuite s"installation should run exactly once under concurrent createSession, " + s"but ran ${lines.size} time(s)") - sessions.asScala.foreach(_.close()) + sessions.asScala.foreach(_.close(emptyCancel)) counterFile.delete() } @@ -847,14 +858,8 @@ class DirectWorkerDispatcherSuite s"echo invoked >> ${counterFile.getAbsolutePath}; sleep 30").build()) .build() val timeoutDispatcher = - new DirectUnixSocketWorkerDispatcher(specWithEnv(env = env)) { + new TestDirectWorkerDispatcher(specWithEnv(env = env)) { override protected def callableTimeoutMs: Long = 500L - override protected def createConnection( - socketPath: String): UnixSocketWorkerConnection = - new SocketFileConnection(socketPath) - override protected def createSessionForWorker( - worker: DirectWorkerProcess): WorkerSession = - new StubWorkerSession(worker) } try { val first = intercept[DirectWorkerTimeoutException] { @@ -916,7 +921,7 @@ class DirectWorkerDispatcherSuite dispatcher = new TestDirectWorkerDispatcher(specWithEnv(env = env)) val session = dispatcher.createSession(None) - session.close() + session.close(emptyCancel) assert(!cleanupMarker.exists(), "cleanup should not run until dispatcher is closed") diff --git a/udf/worker/core/src/test/scala/org/apache/spark/udf/worker/core/TestDirectWorkerHelpers.scala b/udf/worker/core/src/test/scala/org/apache/spark/udf/worker/core/TestDirectWorkerHelpers.scala index 1843024e9fb8c..ca7c86a82a26d 100644 --- a/udf/worker/core/src/test/scala/org/apache/spark/udf/worker/core/TestDirectWorkerHelpers.scala +++ b/udf/worker/core/src/test/scala/org/apache/spark/udf/worker/core/TestDirectWorkerHelpers.scala @@ -16,64 +16,226 @@ */ package org.apache.spark.udf.worker.core -import java.io.File +import java.io.{File, IOException} +import java.nio.file.{Files, FileVisitResult, Path, SimpleFileVisitor} +import java.nio.file.attribute.{BasicFileAttributes, PosixFilePermissions} -import org.apache.spark.udf.worker.{Init, UDFWorkerSpecification} -import org.apache.spark.udf.worker.core.direct.{ - DirectUnixSocketWorkerDispatcher, DirectWorkerProcess, - DirectWorkerSession} +import org.apache.spark.udf.worker.{Cancel, DataRequest, DataResponse, Finish, FinishResponse, + Init, InitResponse, UDFProtoCommunicationPattern, UDFWorkerSpecification, WorkerConnectionSpec} +import org.apache.spark.udf.worker.core.direct.{DirectWorkerDispatcher, DirectWorkerException, + DirectWorkerTimeoutException} +import org.apache.spark.udf.worker.core.direct.DirectWorkerDispatcher.SOCKET_POLL_INTERVAL_MS /** - * A [[WorkerConnection]] test implementation that considers the - * connection active as long as the socket file exists on disk. - * Inherits socket-file deletion from - * [[UnixSocketWorkerConnection.close]]. + * A [[WorkerConnection]] test implementation that treats the connection as + * active as long as the worker's UDS file exists on disk. The socket file is + * removed on close. + * + * Suitable for dispatcher-lifecycle tests that don't need to drive a wire + * protocol -- e.g. verifying that a worker spec spawns a real worker process + * that creates the expected socket. */ -class SocketFileConnection(socketPath: String) - extends UnixSocketWorkerConnection(socketPath) { +class SocketFileConnection(val socketPath: String) extends WorkerConnection { override def isActive: Boolean = new File(socketPath).exists() + override def close(): Unit = { + val f = new File(socketPath) + if (f.exists()) f.delete() + } } /** - * A stub [[DirectWorkerSession]] for process-lifecycle tests - * that don't need actual data transmission. - * - * TODO: [[cancel]] is a no-op here. Once a concrete - * [[DirectWorkerSession]] with real data-plane wiring lands, add - * tests exercising cancel() in particular: cancel from a - * different thread than process(), cancel after process() has - * returned, and cancel before init (should be a no-op). See the - * thread-safety contract in the docstring on - * [[org.apache.spark.udf.worker.core.WorkerSession.cancel]]. + * No-op [[WorkerSession]] for lifecycle-only tests. All protocol methods are + * inert (init/finish report empty responses); tests that exercise the actual + * wire protocol use a concrete transport-backed session. */ -class StubWorkerSession(workerProcess: DirectWorkerProcess) - extends DirectWorkerSession(workerProcess) { - - override protected def doInit(message: Init): Unit = {} +class NoOpWorkerSession( + workerHandle: WorkerHandle, + logger: WorkerLogger = WorkerLogger.NoOp) + extends WorkerSession(workerHandle, logger) { + override protected def doInit(message: Init): InitResponse = InitResponse.getDefaultInstance override protected def doProcess( - input: Iterator[Array[Byte]] - ): Iterator[Array[Byte]] = - Iterator.empty - - override def cancel(): Unit = {} + input: Iterator[DataRequest], + finish: () => Finish): Iterator[DataResponse] = + Iterator.empty[DataResponse] + override protected def doClose(cancel: () => Cancel): Termination = { + // Settle the clean terminal so close() does not fall through to its + // contract-violation recovery path. A no-op session has no in-flight work, + // so the cancel thunk is never needed. + completeTerminal(Termination.Finished(FinishResponse.getDefaultInstance)) + settledTermination + } } /** - * A [[DirectUnixSocketWorkerDispatcher]] subclass for testing - * that uses a socket-file connection and stub sessions instead - * of a real protocol implementation. + * A concrete [[DirectWorkerDispatcher]] for tests that spawns workers over a + * Unix domain socket and yields [[SocketFileConnection]]s / [[NoOpWorkerSession]]s, + * so lifecycle tests exercise the dispatcher's spawn / wait-for-ready / cleanup + * machinery without driving a wire protocol. Allocates a private 0700 socket + * directory at construction; each worker is given a UDS path inside it. + * + * Reusable across modules: callers in `sql/core` (or anywhere with a test-jar + * dependency on `udf-worker-core`) can drop this in for tests that only need to + * verify a worker spec produces a spawnable worker. */ -class TestDirectWorkerDispatcher(spec: UDFWorkerSpecification) - extends DirectUnixSocketWorkerDispatcher(spec) { - - override protected def createConnection( - socketPath: String - ): UnixSocketWorkerConnection = - new SocketFileConnection(socketPath) - - override protected def createSessionForWorker( - worker: DirectWorkerProcess - ): WorkerSession = - new StubWorkerSession(worker) +class TestDirectWorkerDispatcher( + workerSpec: UDFWorkerSpecification, + logger: WorkerLogger = WorkerLogger.NoOp) + extends DirectWorkerDispatcher(workerSpec, logger) { + + // Upper bound on the rename-on-collision retry loop in newEndpointAddress. + private val MAX_SOCKET_LEAF_RETRIES = 16 + + // The private 0700 socket directory, created in [[initialize]] (after the + // base class has validated the spec) and removed in [[closeTransport]]. + private lazy val socketDir: Path = createPrivateTempDirectory() + + override protected def initialize(): Unit = { + super.initialize() + // Force the lazy val now so the directory is created (and any failure + // surfaces) at construction time, after spec validation has passed. + socketDir + } + + /** + * Returns the UDS path the worker should bind. Uses a short 16-hex-char leaf + * (the worker's full UUID still travels via `--id`) to stay within the + * 108-byte UDS `sun_path` limit. + */ + override protected def newEndpointAddress(workerId: String): String = { + val short = workerId.replace("-", "").take(16) + var candidate = socketDir.resolve(s"w-$short.sock") + var suffix = 0 + while (Files.exists(candidate) && suffix < MAX_SOCKET_LEAF_RETRIES) { + suffix += 1 + candidate = socketDir.resolve(s"w-$short-$suffix.sock") + } + if (Files.exists(candidate)) { + throw new IllegalStateException( + s"could not allocate a free UDS path under $socketDir after " + + s"$MAX_SOCKET_LEAF_RETRIES retries (truncated id=$short)") + } + candidate.toString + } + + override protected def waitForReady( + address: String, + process: Process, + outputFile: File): Unit = { + val file = new File(address) + // At least one poll so very small initTimeouts don't trip a premature + // timeout before the worker has any chance to create the socket. + val maxAttempts = math.max(1, (initTimeoutMs / SOCKET_POLL_INTERVAL_MS).toInt) + var attempts = 0 + while (!file.exists() && attempts < maxAttempts) { + if (!process.isAlive) throwWorkerExitedBeforeSocket(process, address, outputFile) + Thread.sleep(SOCKET_POLL_INTERVAL_MS) + attempts += 1 + } + if (!file.exists()) { + if (process.isAlive) { + DirectWorkerDispatcher.destroyForciblyAndReap( + process, logger, s"init timeout $address") + val tail = readOutputTail(outputFile) + throw new DirectWorkerTimeoutException( + s"Worker did not create socket at $address within ${initTimeoutMs}ms\n$tail") + } else { + // Worker exited after the last poll without creating the socket; + // prefer the exit-code message over the ambiguous "did not create". + throwWorkerExitedBeforeSocket(process, address, outputFile) + } + } + } + + override protected def cleanupEndpointAddress(address: String): Unit = { + Files.deleteIfExists(new File(address).toPath) + } + + override protected def closeTransport(): Unit = { + if (!Files.exists(socketDir)) return + // Recursive post-order delete: today socketDir contains only socket files + // at the top level, but a future change that namespaces workers into + // subdirectories should not silently leak them. + Files.walkFileTree(socketDir, new SimpleFileVisitor[Path] { + override def visitFile(file: Path, attrs: BasicFileAttributes): FileVisitResult = { + try Files.deleteIfExists(file) catch { case _: IOException => () } + FileVisitResult.CONTINUE + } + override def postVisitDirectory(dir: Path, exc: IOException): FileVisitResult = { + try Files.deleteIfExists(dir) catch { case _: IOException => () } + FileVisitResult.CONTINUE + } + }) + } + + // `spec` is the same object as the `workerSpec` field but passed explicitly: + // at the point this runs (parent constructor body), `this` is only partially + // constructed and reading subclass fields is unsafe. + override protected def validateTransportSupport(spec: UDFWorkerSpecification): Unit = { + val props = spec.getDirect.getProperties + require(props.hasConnection, + "DirectWorker.properties.connection must be set") + val conn = props.getConnection + require(conn.getTransportCase == WorkerConnectionSpec.TransportCase.UNIX_DOMAIN_SOCKET, + "TestDirectWorkerDispatcher requires UNIX domain socket transport, " + + s"got ${conn.getTransportCase}") + require(spec.hasCapabilities, + "TestDirectWorkerDispatcher requires WorkerCapabilities declaring " + + "BIDIRECTIONAL_STREAMING in supported_communication_patterns") + val patterns = spec.getCapabilities.getSupportedCommunicationPatternsList + val supportsBidi = (0 until patterns.size()).exists { i => + patterns.get(i) == UDFProtoCommunicationPattern.BIDIRECTIONAL_STREAMING + } + require(supportsBidi, + "TestDirectWorkerDispatcher requires BIDIRECTIONAL_STREAMING " + + "in WorkerCapabilities.supported_communication_patterns") + } + + override protected def newConnection(address: String): WorkerConnection = + new SocketFileConnection(address) + + override protected def newSession( + workerHandle: WorkerHandle, + connection: WorkerConnection): WorkerSession = + new NoOpWorkerSession(workerHandle, logger) + + private def throwWorkerExitedBeforeSocket( + process: Process, + address: String, + outputFile: File): Nothing = { + val tail = readOutputTail(outputFile) + throw new DirectWorkerException( + s"Worker exited with code ${process.exitValue()} " + + s"before creating socket at $address\n$tail") + } + + /** + * Creates a private (owner-only, 0700) temp directory for worker sockets. + * On POSIX filesystems the permissions are applied atomically at creation; + * the non-POSIX branch tightens best-effort after creation. + */ + private def createPrivateTempDirectory(): Path = { + val attr = PosixFilePermissions.asFileAttribute( + PosixFilePermissions.fromString("rwx------")) + try { + Files.createTempDirectory("spark-udf-worker", attr) + } catch { + case _: UnsupportedOperationException => + val dir = Files.createTempDirectory("spark-udf-worker") + val f = dir.toFile + // Bit-wise AND (NOT &&): all six setters must run even if an earlier + // one returns false, so the final permission state matches owner-only. + val applied = + f.setReadable(false, false) & f.setWritable(false, false) & + f.setExecutable(false, false) & f.setReadable(true, true) & + f.setWritable(true, true) & f.setExecutable(true, true) + if (!applied) { + logger.warn( + s"Could not fully restrict permissions on $dir; socket " + + s"directory may be accessible to other local users on this " + + s"filesystem") + } + dir + } + } } diff --git a/udf/worker/core/src/test/scala/org/apache/spark/udf/worker/core/WorkerSessionSuite.scala b/udf/worker/core/src/test/scala/org/apache/spark/udf/worker/core/WorkerSessionSuite.scala new file mode 100644 index 0000000000000..37380d0d77655 --- /dev/null +++ b/udf/worker/core/src/test/scala/org/apache/spark/udf/worker/core/WorkerSessionSuite.scala @@ -0,0 +1,270 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.udf.worker.core + +// scalastyle:off funsuite +import org.scalatest.funsuite.AnyFunSuite + +import org.apache.spark.udf.worker.{Cancel, CancelResponse, DataRequest, DataResponse, + ExecutionError, Finish, FinishResponse, Init, InitResponse} +import org.apache.spark.udf.worker.core.WorkerSession.SessionState + +/** + * Unit tests for the [[WorkerSession]] state machine, exercised through a fake + * subclass that drives the protocol-event edges directly with no transport. + * + * The concrete transport-backed sessions test this same machine end-to-end over + * the wire; these tests pin the transport-agnostic base-class contract on its + * own: call ordering, terminal settling (first-wins), the [[Termination]] + * mapping, and the `close()` finalizer (release-once, invalidation, and the + * doClose post-condition guard). + */ +class WorkerSessionSuite extends AnyFunSuite { +// scalastyle:on funsuite + + /** A [[WorkerHandle]] that counts the lifecycle callbacks made against it. */ + private final class RecordingHandle extends WorkerHandle { + var released = 0 + var invalidated = 0 + override def id: String = "test-worker" + override def markInvalid(): Unit = invalidated += 1 + override def releaseSession(): Unit = released += 1 + } + + /** + * A [[WorkerSession]] whose protocol hooks are supplied by the test and whose + * protected state-machine primitives are re-exposed, so a test can drive the + * machine (settle terminals, CAS edges) without a real worker. + */ + private final class FakeWorkerSession( + handle: WorkerHandle = new RecordingHandle, + onInit: FakeWorkerSession => InitResponse = _ => InitResponse.getDefaultInstance, + onProcess: (FakeWorkerSession, Iterator[DataRequest], () => Finish) => + Iterator[DataResponse] = (_, _, _) => Iterator.empty[DataResponse], + onCloseHook: (FakeWorkerSession, () => Cancel) => Termination = + (self, _) => { + self.settle(Termination.Finished(FinishResponse.getDefaultInstance)) + self.term + }) + extends WorkerSession(handle, WorkerLogger.NoOp) { + + var terminalSettledCount = 0 + + override protected def doInit(message: Init): InitResponse = onInit(this) + override protected def doProcess( + input: Iterator[DataRequest], finish: () => Finish): Iterator[DataResponse] = + onProcess(this, input, finish) + override protected def doClose(cancel: () => Cancel): Termination = onCloseHook(this, cancel) + override protected def onTerminalSettled(termination: Termination): Unit = + terminalSettledCount += 1 + + // Re-expose the protected primitives so the test can drive the machine. + def state: SessionState = currentState + def cas(expect: SessionState, update: SessionState): Boolean = + compareAndSetState(expect, update) + def settle(t: Termination): Boolean = completeTerminal(t) + def term: Termination = settledTermination + } + + test("init returns the worker InitResponse and advances Created -> Initialized") { + val resp = InitResponse.getDefaultInstance + val s = new FakeWorkerSession(onInit = _ => resp) + assert(s.state == SessionState.Created) + assert(s.init(Init.getDefaultInstance) eq resp) + assert(s.state == SessionState.Initialized) + } + + test("init must be called exactly once") { + val s = new FakeWorkerSession() + s.init(Init.getDefaultInstance) + val ex = intercept[IllegalStateException](s.init(Init.getDefaultInstance)) + assert(ex.getMessage.contains("exactly once")) + } + + test("process before init is rejected") { + val s = new FakeWorkerSession() + val ex = intercept[IllegalStateException](s.process(Iterator.empty)) + assert(ex.getMessage.contains("before init")) + } + + test("process advances to Streaming and may only be called once") { + val s = new FakeWorkerSession() + s.init(Init.getDefaultInstance) + s.process(Iterator.empty) + assert(s.state == SessionState.Streaming) + val ex = intercept[IllegalStateException](s.process(Iterator.empty)) + assert(ex.getMessage.contains("already been called")) + } + + test("a failed doInit settles a terminal and rejects a later process") { + val err = ExecutionError.getDefaultInstance + val s = new FakeWorkerSession(onInit = self => { + self.settle(Termination.Failed(err)) + throw new RuntimeException("init boom") + }) + val initEx = intercept[RuntimeException](s.init(Init.getDefaultInstance)) + assert(initEx.getMessage.contains("init boom")) + assert(s.term == Termination.Failed(err)) + val ex = intercept[IllegalStateException](s.process(Iterator.empty)) + assert(ex.getMessage.contains("terminated")) + } + + test("close returns the Termination produced by doClose") { + val resp = FinishResponse.getDefaultInstance + val s = new FakeWorkerSession(onCloseHook = (self, _) => { + self.settle(Termination.Finished(resp)) + self.term + }) + assert(s.close() == Termination.Finished(resp)) + } + + test("close releases the worker exactly once and leaves a salvageable worker valid") { + val h = new RecordingHandle + val s = new FakeWorkerSession(handle = h) // default doClose settles Finished + s.close() + s.close() + assert(h.released == 1) + assert(h.invalidated == 0) + } + + test("close marks a transport-failed worker invalid") { + val h = new RecordingHandle + val s = new FakeWorkerSession(handle = h, onCloseHook = (self, _) => { + self.settle(Termination.TransportFailed(new RuntimeException("transport down"))) + self.term + }) + s.close() + assert(h.invalidated == 1) + assert(h.released == 1) + } + + test("close leaves a worker salvageable after an execution Failed") { + val h = new RecordingHandle + val err = ExecutionError.getDefaultInstance + // A Failed terminal is typically a user-code (UDF) error, not a worker fault, + // so the worker stays reusable: markInvalid must NOT be called. + val s = new FakeWorkerSession(handle = h, onCloseHook = (self, _) => { + self.settle(Termination.Failed(err)) + self.term + }) + assert(s.close() == Termination.Failed(err)) + assert(h.invalidated == 0) + assert(h.released == 1) + } + + test("close enforces the doClose terminal post-condition") { + val h = new RecordingHandle + // doClose returns a Termination without settling any terminal -- a subclass + // contract violation. close() must settle a TransportFailed terminal, treat + // the worker as unsalvageable, and return a Termination consistent with that + // settled state rather than the (untrustworthy) value doClose produced. + val s = new FakeWorkerSession(handle = h, + onCloseHook = (_, _) => Termination.Finished(FinishResponse.getDefaultInstance)) + val termination = s.close() + assert(s.term.isInstanceOf[Termination.TransportFailed]) + assert(termination.isInstanceOf[Termination.TransportFailed]) + assert(h.invalidated == 1) + } + + test("close converts an unexpected internal doClose failure into TransportFailed") { + val h = new RecordingHandle + val boom = new RuntimeException("internal doClose failure") + // doClose throws an unexpected internal / transport-level error (e.g. a + // failed wire write). close() must not propagate it: it settles + // TransportFailed, returns it, and marks the worker unsalvageable. + val s = new FakeWorkerSession(handle = h, onCloseHook = (_, _) => throw boom) + assert(s.close() == Termination.TransportFailed(boom)) + assert(h.invalidated == 1) + assert(h.released == 1) + } + + test("close swallows a non-fatal worker-handle failure") { + val throwingHandle = new WorkerHandle { + override def id: String = "boom" + override def markInvalid(): Unit = throw new RuntimeException("markInvalid boom") + override def releaseSession(): Unit = throw new RuntimeException("releaseSession boom") + } + val cause = new RuntimeException("transport down") + // A TransportFailed terminal makes the worker unsalvageable, so close() + // attempts both markInvalid() and releaseSession() -- here both throw. + // close() must swallow them and still return the settled termination. + val s = new FakeWorkerSession(handle = throwingHandle, onCloseHook = (self, _) => { + self.settle(Termination.TransportFailed(cause)) + self.term + }) + assert(s.close() == Termination.TransportFailed(cause)) + } + + test("close still releases the worker when markInvalid throws") { + // markInvalid throws but releaseSession succeeds: close() must swallow the + // markInvalid failure and still attempt releaseSession, so the ref count is + // not leaked. A TransportFailed terminal makes the worker unsalvageable, so + // markInvalid is reached on the way to releaseSession. + var released = 0 + val handle = new WorkerHandle { + override def id: String = "throws-on-invalidate" + override def markInvalid(): Unit = throw new RuntimeException("markInvalid boom") + override def releaseSession(): Unit = released += 1 + } + val cause = new RuntimeException("transport down") + val s = new FakeWorkerSession(handle = handle, onCloseHook = (self, _) => { + self.settle(Termination.TransportFailed(cause)) + self.term + }) + assert(s.close() == Termination.TransportFailed(cause)) + assert(released == 1, "releaseSession must be attempted even when markInvalid throws") + } + + test("completeTerminal is first-wins and runs onTerminalSettled once") { + val first = Termination.Finished(FinishResponse.getDefaultInstance) + val s = new FakeWorkerSession() + assert(s.settle(first)) + assert(!s.settle(Termination.Cancelled(CancelResponse.getDefaultInstance))) + assert(s.term == first) + assert(s.terminalSettledCount == 1) + } + + test("compareAndSetState drives non-terminal edges and loses to a settled terminal") { + val s = new FakeWorkerSession() + s.init(Init.getDefaultInstance) + assert(s.cas(SessionState.Initialized, SessionState.Streaming)) + assert(s.state == SessionState.Streaming) + assert(s.settle(Termination.Failed(ExecutionError.getDefaultInstance))) + assert(!s.cas(SessionState.Streaming, SessionState.Finishing)) + } + + test("settledTermination maps each terminal to its Termination") { + val fin = FinishResponse.getDefaultInstance + val can = CancelResponse.getDefaultInstance + val err = ExecutionError.getDefaultInstance + val cause = new RuntimeException("transport down") + def termFor(t: Termination): Termination = { + val s = new FakeWorkerSession() + s.settle(t) + s.term + } + assert(termFor(Termination.Finished(fin)) == Termination.Finished(fin)) + assert(termFor(Termination.Cancelled(can)) == Termination.Cancelled(can)) + assert(termFor(Termination.Failed(err)) == Termination.Failed(err)) + assert(termFor(Termination.TransportFailed(cause)) == Termination.TransportFailed(cause)) + } + + test("settledTermination throws before a terminal is settled") { + val s = new FakeWorkerSession() + intercept[IllegalStateException](s.term) + } +}