diff --git a/sql/core/pom.xml b/sql/core/pom.xml index d96b71a9d8058..47117c7350351 100644 --- a/sql/core/pom.xml +++ b/sql/core/pom.xml @@ -308,9 +308,20 @@ spark-udf-worker-core_${scala.binary.version} ${project.version} + org.apache.spark - spark-udf-worker-core_${scala.binary.version} + spark-udf-worker-grpc_${scala.binary.version} + ${project.version} + test + + + org.apache.spark + spark-udf-worker-grpc_${scala.binary.version} ${project.version} test-jar test 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 bbd47741aba1c..bf8e6bc14b52c 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 @@ -29,7 +29,8 @@ 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.{Cancel, UDFWorkerSpecification} -import org.apache.spark.udf.worker.core.{TestDirectWorkerDispatcher, WorkerConnection} +import org.apache.spark.udf.worker.core.WorkerConnection +import org.apache.spark.udf.worker.grpc.testing.TestDirectGrpcDispatcher /** * A [[WorkerConnection]] that opens a real Unix-domain-socket channel to the @@ -52,12 +53,12 @@ private object ConnectingSocketConnection { } /** - * A [[TestDirectWorkerDispatcher]] whose connection actually connects to the + * A [[TestDirectGrpcDispatcher]] 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 ConnectingTestDispatcher(spec: UDFWorkerSpecification) - extends TestDirectWorkerDispatcher(spec) { + extends TestDirectGrpcDispatcher(spec) { override protected def newConnection(address: String): WorkerConnection = ConnectingSocketConnection.connect(address) } @@ -74,8 +75,8 @@ private class ConnectingTestDispatcher(spec: UDFWorkerSpecification) * 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 + * Reuses [[TestDirectGrpcDispatcher]]'s spawn / wait-for-ready machinery + * (shared from `udf-worker-grpc` 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. diff --git a/udf/worker/README.md b/udf/worker/README.md index 2250c8df10535..f82e89a7ccfd4 100644 --- a/udf/worker/README.md +++ b/udf/worker/README.md @@ -29,6 +29,84 @@ framework currently provides **direct worker creation** (local OS processes) and is designed for future **indirect creation** (via a provisioning service or daemon). +## Architecture + +The framework factors into three layers, separating *what the engine sees* +(transport-agnostic API), *how a worker is provisioned* (local subprocess +today, pooled / remote later), and *how bytes are moved* (gRPC over UDS +today). + +``` + Engine + │ + ▼ + ┌──────────────────────┐ + │ UDFDispatcherManager│ caches one dispatcher per worker spec + └──────────┬───────────┘ + │ + ▼ + ┌──────────────────────┐ engine-facing API + │ WorkerDispatcher │ (transport- and provisioning-agnostic) + └──────────┬───────────┘ + │ creates + ▼ + ┌──────────────────────┐ holds ┌─────────────────┐ + │ WorkerSession │ ─────────▶ │ WorkerHandle │ + │ (one UDF execution) │ │ (release hook) │ + └──────────┬───────────┘ └─────────────────┘ + │ uses + ▼ + ┌──────────────────────┐ 1 per worker process; + │ WorkerConnection │ shared across all sessions on that worker + └──────────────────────┘ + + ── Direct mode (core.direct) ───────────────────────────────────────────── + DirectWorkerDispatcher (abstract: spawn / wait / SIGTERM-SIGKILL) + │ + └─▶ DirectWorkerProcess (live subprocess; implements WorkerHandle) + + ── gRPC over UDS (grpc) ───────────────────────────────────────────── + DirectGrpcDispatcher (extends DirectWorkerDispatcher) + │ + ├─▶ GrpcWorkerSession (drives bidi Execute stream) + ├─▶ GrpcWorkerChannel (ManagedChannel + Netty event loop) + └─▶ UnixDomainSocketTransport (epoll / kqueue detection) +``` + +### Cardinality + +| Relationship | Cardinality | Notes | +|---|---|---| +| `UDFDispatcherManager` ↔ `WorkerDispatcher` | 1 ↔ N | One dispatcher per worker spec, cached | +| `WorkerDispatcher` ↔ `WorkerSession` | 1 ↔ N | Sessions are short-lived; dispatcher outlives them | +| `WorkerDispatcher` ↔ worker process | 1 ↔ N | In direct mode the dispatcher manages a fleet of subprocesses | +| worker process ↔ `WorkerConnection` | 1 ↔ 1 | One transport per worker | +| worker process ↔ `WorkerSession` | 1 ↔ N | Future via gRPC stream multiplexing; today effectively 1 ↔ 1 (no pooling yet) | +| `WorkerSession` ↔ `WorkerHandle` | 1 ↔ 1 | Each session holds one handle; released exactly once on session close | +| `GrpcWorkerSession` ↔ `Execute` stream | 1 ↔ 1 | One bidi gRPC stream per UDF execution | + +### Key ideas + +- **Three independent lifecycle scopes.** Dispatcher (per spec), connection + (per worker process), session (per UDF execution). The split lets a worker + be reused across many UDFs without conflating the transport with the + per-execution state machine. +- **Engine doesn't know the transport.** It sees only `WorkerSession`. A new + transport (TCP, named pipe, future remote workers) plugs in via fresh + `WorkerDispatcher` / `WorkerConnection` / `WorkerSession` implementations; + engine code is unchanged. +- **Engine doesn't know how the worker was provisioned.** Sessions hold a + `WorkerHandle`, not a worker. A future pooled or remotely-leased worker + model supplies a different `WorkerHandle` without touching session code. +- **Lifecycle guards live in base classes; protocol lives in subclasses.** + `WorkerSession.init / process / close` are `final` and own the AtomicBoolean + guards; subclasses fill `doInit / doProcess / doClose` and do not re-check + the contract. +- **The wire protocol is language-agnostic.** `udf_message.proto` (message + types) and `udf_service.proto` (the gRPC service) define it; any worker + (Python, Scala, Rust, ...) implementing the proto is consumable by this + stack. + ## Sub-packages ``` @@ -39,26 +117,34 @@ udf/worker/ │ udf_service.proto -- UdfWorker gRPC service (Execute, Manage) │ common.proto -- shared enums (UDFWorkerDataFormat, etc.) │ -├── core/ -- abstract interfaces +├── core/ -- abstract interfaces (no gRPC dependency) │ WorkerDispatcher.scala -- creates sessions, manages worker lifecycle │ WorkerSession.scala -- per-UDF init/process/cancel/close │ WorkerConnection.scala -- transport channel abstraction +│ WorkerHandle.scala -- release hook used by sessions │ WorkerSecurityScope.scala -- security boundary for worker pooling +│ UDFDispatcherManager.scala -- caches one dispatcher per worker spec │ │ │ └── direct/ -- "direct" creation: local OS processes │ DirectWorkerDispatcher.scala -- spawns processes, env lifecycle -│ DirectWorkerProcess.scala -- OS process + connection + UDS socket -│ DirectWorkerSession.scala -- session backed by a direct process +│ DirectWorkerProcess.scala -- OS process + connection + WorkerHandle │ -└── grpc/ -- gRPC transport (gRPC runtime confined here) - (generated) -- UdfWorkerGrpc service stubs from proto/udf_service.proto +└── grpc/ -- gRPC over UDS transport (gRPC runtime confined here) + (generated) -- UdfWorkerGrpc service stubs from proto/udf_service.proto + DirectGrpcDispatcher.scala -- direct dispatcher speaking gRPC/UDS + GrpcWorkerSession.scala -- drives one bidi Execute stream + GrpcWorkerChannel.scala -- ManagedChannel + Netty event loop + UnixDomainSocketTransport.scala -- epoll/kqueue detection ``` The `core/` package defines abstract interfaces that are independent of how -workers are created. The `core/direct/` sub-package implements "direct" -worker creation where Spark spawns local OS processes. Future packages -(e.g., `core/indirect/`) can implement alternative creation modes such as -obtaining workers from a provisioning service or daemon. +workers are created and what transport carries the protocol. `core/direct/` +implements "direct" worker creation where Spark spawns local OS processes. +The `grpc/` module provides the gRPC-over-UDS transport on top of the direct +provisioning model. Future packages (e.g., `core/indirect/`) can implement +alternative creation modes such as obtaining workers from a provisioning +service or daemon, and additional transports can be added alongside `grpc/` +without touching the engine-facing API. The `grpc/` module owns the gRPC service-stub generation (from `proto/`'s `udf_service.proto`) and the gRPC runtime dependencies. Keeping @@ -130,9 +216,10 @@ val spec = UDFWorkerSpecification.newBuilder() .build()) .build() -// 2. Create a dispatcher. Use a protocol-specific subclass of -// DirectWorkerDispatcher (e.g., gRPC over UDS). -val dispatcher: WorkerDispatcher = ... +// 2. Create a dispatcher. DirectGrpcDispatcher is the shipped concrete +// implementation (gRPC over UDS). +val dispatcher: WorkerDispatcher = + new org.apache.spark.udf.worker.grpc.DirectGrpcDispatcher(spec) // 3. Create a session for one UDF execution. val session = dispatcher.createSession(securityScope = None) @@ -184,14 +271,17 @@ build/sbt "udf-worker-core/test" "udf-worker-grpc/test" ## Current status -This is the **first MVP** providing the core abstraction layer and the -direct worker dispatcher. +The framework ships the core abstraction layer, the direct worker dispatcher, +and the first concrete transport (**gRPC over UDS** via `DirectGrpcDispatcher`). The following are left as TODOs: - **Connection pooling** -- reuse workers across sessions - **Security scope isolation** -- partition pools by `WorkerSecurityScope` - **Indirect worker creation** -- obtain workers from a service or daemon -- **Protocol-specific implementations** -- e.g., gRPC over UDS +- **Protocol-level flow control** -- complement gRPC's transport-level flow + control with worker-side acks so the engine can bound in-flight input + batches (useful for reduce-style UDFs that hold input without producing + output) ## Design references 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 index c66cb5689b875..821b00be2abdd 100644 --- 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 @@ -63,4 +63,13 @@ trait WorkerHandle { * implementations need not be idempotent. */ def releaseSession(): Unit + + /** + * The transport [[WorkerConnection]] used to talk to this worker. The + * dispatcher attaches it to the handle so the session-construction seam + * (the dispatcher's `newSession`) takes a single object rather than a + * separate handle + connection. Its lifecycle is owned by the dispatcher, + * not the session: it is torn down when the worker is, not on session close. + */ + def connection: WorkerConnection } 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 db81b5265b794..5a24c25750628 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 @@ -47,9 +47,9 @@ import org.apache.spark.udf.worker.core.direct.DirectWorkerDispatcher.{CallableR * (the single-reference case of the future pooling policy). * * 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]]. + * address, establish a verified [[WorkerConnection]], and build the + * per-session [[WorkerSession]]. See [[newEndpointAddress]], + * [[connectWorker]], and [[newSession]]. * * For workers obtained through a provisioning service or daemon (indirect * creation), see the `indirect` package (TODO). @@ -129,6 +129,8 @@ abstract class DirectWorkerDispatcher( private[this] val workers = new ConcurrentHashMap[String, DirectWorkerProcess]() private[this] val closed = new AtomicBoolean(false) + private[this] val sessionCreationLock = new Object + private[this] var inFlightSessionCreations = 0 // TODO [SPARK-55278]: extract the env state machine + JVM shutdown hook // into a standalone EnvironmentManager once the gRPC dispatcher work has @@ -145,18 +147,21 @@ abstract class DirectWorkerDispatcher( protected def newEndpointAddress(workerId: String): String /** - * Waits for the worker process to be ready to accept connections at - * `address`. Throws [[DirectWorkerTimeoutException]] on timeout, or - * [[DirectWorkerException]] if the process exits early. + * Establishes and verifies the transport connection to a spawned worker. + * Implementations may construct the connection before or after the worker + * binds its endpoint, as required by the transport. + * + * @throws DirectWorkerTimeoutException if the worker does not become ready. + * @throws DirectWorkerException if the worker exits before becoming ready. */ - protected def waitForReady( + protected def connectWorker( address: String, process: Process, - outputFile: File): Unit + outputFile: File): WorkerConnection /** * Best-effort per-endpoint cleanup, called from the spawn-failure path - * before any [[WorkerArtifacts]] / [[WorkerConnection]] exists. + * before a [[WorkerArtifacts]] bundle owns the connection. */ protected def cleanupEndpointAddress(address: String): Unit @@ -198,23 +203,14 @@ abstract class DirectWorkerDispatcher( */ 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 - /** * 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. + * dispatcher-side cleanup and `workerHandle.connection` for the wire + * transport. */ - protected def newSession( - workerHandle: WorkerHandle, - connection: WorkerConnection): WorkerSession + protected def newSession(workerHandle: WorkerHandle): WorkerSession /** * Test-only hook: invoked once per [[createSession]] after the worker @@ -240,30 +236,38 @@ abstract class DirectWorkerDispatcher( securityScope: Option[WorkerSecurityScope]): WorkerSession = { require(securityScope.isEmpty, "securityScope is not supported yet; pass None until pooling lands") - if (closed.get()) throwClosed() - ensureEnvironmentReady() - val worker = spawnWorker() - // Acquire before publish: a concurrent close() iterating `workers` must - // not tear down this worker before we hand it to the caller. - worker.acquireSession() - workers.put(worker.id, worker) - // Re-check for close() that ran concurrently. Releasing fires the - // ref-count callback, which removes and tears down the worker. - if (closed.get()) { - worker.releaseSession() - throwClosed() - } + beginSessionCreation() try { - afterWorkerRegistered(worker) - newSession(worker, worker.connection) - } catch { - case e: InterruptedException => - Thread.currentThread().interrupt() + ensureEnvironmentReady() + if (closed.get()) throwClosed() + val worker = spawnWorker() + // Acquire before publish: a concurrent close() iterating `workers` must + // not tear down this worker before we hand it to the caller. + worker.acquireSession() + workers.put(worker.id, worker) + // Re-check for close() that ran concurrently. Releasing fires the + // ref-count callback, which removes and tears down the worker. + if (closed.get()) { worker.releaseSession() - throw e - case NonFatal(e) => - worker.releaseSession() - throw e + throwClosed() + } + try { + afterWorkerRegistered(worker) + // The test hook may block, so close could have started since the + // first post-publication check. + if (closed.get()) throwClosed() + newSession(worker) + } catch { + case e: InterruptedException => + Thread.currentThread().interrupt() + worker.releaseSession() + throw e + case NonFatal(e) => + worker.releaseSession() + throw e + } + } finally { + endSessionCreation() } } @@ -286,17 +290,42 @@ abstract class DirectWorkerDispatcher( private def throwClosed(): Nothing = throw new IllegalStateException("Dispatcher is closed") + private def beginSessionCreation(): Unit = sessionCreationLock.synchronized { + if (closed.get()) throwClosed() + inFlightSessionCreations += 1 + } + + private def endSessionCreation(): Unit = sessionCreationLock.synchronized { + inFlightSessionCreations -= 1 + if (inFlightSessionCreations == 0) sessionCreationLock.notifyAll() + } + + /** Waits uninterruptibly so cleanup can never race an in-flight creation. */ + private def awaitSessionCreations(): Unit = { + var interrupted = false + sessionCreationLock.synchronized { + while (inFlightSessionCreations != 0) { + try { + sessionCreationLock.wait() + } catch { + case _: InterruptedException => interrupted = true + } + } + } + if (interrupted) Thread.currentThread().interrupt() + } + /** * Terminates tracked workers, removes the socket directory, and runs - * environment cleanup. Idempotent via CAS. Does not drain in-flight - * createSession calls -- a worker spawned racing with close tears - * itself down through the ref-count callback, which may outlive this - * method. + * environment cleanup. Idempotent via CAS. In-flight [[createSession]] + * calls are drained before cleanup, so they cannot publish a worker after + * the transport directory or environment has been removed. */ override def close(): Unit = { if (!closed.compareAndSet(false, true)) { return } + awaitSessionCreations() // TODO [SPARK-55278]: Cleanup sessions as well? // TODO [SPARK-55278]: close workers in parallel -- today shutdown is serialised at // N * gracefulTimeoutMs worst case. @@ -455,10 +484,10 @@ abstract class DirectWorkerDispatcher( val env = runner.getEnvironmentVariablesMap.asScala.toMap val outputFile = Files.createTempFile("udf-worker-", ".log") val process = launchProcess(cmd, env, outputFile.toFile) + var connection: WorkerConnection = null try { - waitForReady(address, process, outputFile.toFile) - val connection = newConnection(address) + connection = connectWorker(address, process, outputFile.toFile) 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. @@ -469,17 +498,25 @@ abstract class DirectWorkerDispatcher( } catch { case e: InterruptedException => Thread.currentThread().interrupt() - cleanupRawSpawn(process, address, outputFile) + cleanupRawSpawn(process, connection, address, outputFile) throw e case NonFatal(e) => - cleanupRawSpawn(process, address, outputFile) + cleanupRawSpawn(process, connection, address, outputFile) throw e } } - // Pre-WorkerArtifacts cleanup: the connection has not been built yet, - // so we have no bundle to close(). Each step is independent. - private def cleanupRawSpawn(p: Process, address: String, outputFile: Path): Unit = { + // Pre-WorkerArtifacts cleanup. Each step is independent. + private def cleanupRawSpawn( + p: Process, + connection: WorkerConnection, + address: String, + outputFile: Path): Unit = { + if (connection != null) { + try connection.close() catch { + case NonFatal(e) => logger.debug("Failed to close worker connection", e) + } + } DirectWorkerDispatcher.destroyForciblyAndReap(p, logger, "failed spawn") try cleanupEndpointAddress(address) catch { case NonFatal(e) => @@ -552,7 +589,7 @@ abstract class DirectWorkerDispatcher( // 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[worker] val READY_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 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 52e7d5123d81b..9ff5cc07d2e92 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 @@ -72,8 +72,8 @@ class DirectWorkerProcess( /** The OS process handle for this worker. */ def process: Process = artifacts.process - /** The transport connection for this worker. */ - def connection: WorkerConnection = artifacts.connection + /** The transport connection for this worker (see [[WorkerHandle.connection]]). */ + override def connection: WorkerConnection = artifacts.connection /** Path to the merged stdout/stderr log for this worker. */ def outputFile: Path = artifacts.outputFile 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 index c8bbcfffd7bb2..002dad7ee1158 100644 --- 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 @@ -36,6 +36,12 @@ import org.apache.spark.udf.worker.core.WorkerSession.SessionState class WorkerSessionSuite extends AnyFunSuite { // scalastyle:on funsuite + /** A no-op [[WorkerConnection]] for handles whose transport is never used. */ + private val stubConnection: WorkerConnection = new WorkerConnection { + override def isActive: Boolean = false + override def close(): Unit = () + } + /** A [[WorkerHandle]] that counts the lifecycle callbacks made against it. */ private final class RecordingHandle extends WorkerHandle { var released = 0 @@ -43,6 +49,7 @@ class WorkerSessionSuite extends AnyFunSuite { override def id: String = "test-worker" override def markInvalid(): Unit = invalidated += 1 override def releaseSession(): Unit = released += 1 + override def connection: WorkerConnection = stubConnection } /** @@ -224,6 +231,7 @@ class WorkerSessionSuite extends AnyFunSuite { override def id: String = "boom" override def markInvalid(): Unit = throw new RuntimeException("markInvalid boom") override def releaseSession(): Unit = throw new RuntimeException("releaseSession boom") + override def connection: WorkerConnection = stubConnection } val cause = new RuntimeException("transport down") // A TransportFailed terminal makes the worker unsalvageable, so close() @@ -246,6 +254,7 @@ class WorkerSessionSuite extends AnyFunSuite { override def id: String = "throws-on-invalidate" override def markInvalid(): Unit = throw new RuntimeException("markInvalid boom") override def releaseSession(): Unit = released += 1 + override def connection: WorkerConnection = stubConnection } val cause = new RuntimeException("transport down") val s = new FakeWorkerSession(handle = handle, onCloseHook = (self, _) => { diff --git a/udf/worker/grpc/pom.xml b/udf/worker/grpc/pom.xml index aa563d157d9c0..6d6323444b31c 100644 --- a/udf/worker/grpc/pom.xml +++ b/udf/worker/grpc/pom.xml @@ -91,6 +91,55 @@ io.grpc grpc-stub + + + io.grpc + grpc-netty + ${io.grpc.version} + + + io.netty + netty-codec-http2 + ${netty.version} + + + io.netty + netty-handler-proxy + ${netty.version} + + + + io.netty + netty-transport-native-unix-common + ${netty.version} + + + + io.netty + netty-transport-native-epoll + linux-x86_64 + + + io.netty + netty-transport-native-epoll + linux-aarch_64 + + + io.netty + netty-transport-native-kqueue + osx-x86_64 + + + io.netty + netty-transport-native-kqueue + osx-aarch_64 + io.grpc grpc-inprocess diff --git a/udf/worker/core/src/test/scala/org/apache/spark/udf/worker/core/TestDirectWorkerHelpers.scala b/udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/DirectGrpcDispatcher.scala similarity index 55% rename from udf/worker/core/src/test/scala/org/apache/spark/udf/worker/core/TestDirectWorkerHelpers.scala rename to udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/DirectGrpcDispatcher.scala index 7ffe72f917573..1462c4e83ac8f 100644 --- a/udf/worker/core/src/test/scala/org/apache/spark/udf/worker/core/TestDirectWorkerHelpers.scala +++ b/udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/DirectGrpcDispatcher.scala @@ -14,80 +14,56 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.spark.udf.worker.core +package org.apache.spark.udf.worker.grpc import java.io.{File, IOException} import java.nio.file.{Files, FileVisitResult, Path, Paths, SimpleFileVisitor} import java.nio.file.attribute.{BasicFileAttributes, PosixFilePermissions} +import java.util.concurrent.{CountDownLatch, TimeUnit} -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 +import scala.util.control.NonFatal -/** - * 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(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() - } -} +import io.grpc.ConnectivityState -/** - * 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 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[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 - } -} +import org.apache.spark.annotation.Experimental +import org.apache.spark.udf.worker.{UDFProtoCommunicationPattern, UDFWorkerSpecification, + WorkerConnectionSpec} +import org.apache.spark.udf.worker.core.{WorkerConnection, WorkerHandle, WorkerLogger, + WorkerSession} +import org.apache.spark.udf.worker.core.direct.{DirectWorkerDispatcher, DirectWorkerException, + DirectWorkerTimeoutException} +import org.apache.spark.udf.worker.core.direct.DirectWorkerDispatcher.READY_POLL_INTERVAL_MS /** - * 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. + * :: Experimental :: + * A concrete [[DirectWorkerDispatcher]] that spawns workers and talks to + * them over the UDF gRPC protocol on a Unix domain socket. 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. + * @param workerSpec worker specification used to launch each worker. + * @param logger logger for lifecycle diagnostics. + * @param grpcMaxInboundMessageSize maximum serialized response size accepted + * from a worker. Defaults to 128 MiB. */ -class TestDirectWorkerDispatcher( +@Experimental +class DirectGrpcDispatcher( workerSpec: UDFWorkerSpecification, - logger: WorkerLogger = WorkerLogger.NoOp) + logger: WorkerLogger = WorkerLogger.NoOp, + grpcMaxInboundMessageSize: Int = GrpcWorkerChannel.DEFAULT_MAX_INBOUND_MESSAGE_SIZE) extends DirectWorkerDispatcher(workerSpec, logger) { // Upper bound on the rename-on-collision retry loop in newEndpointAddress. + // 16 is well above any realistic concurrent-worker count and keeps the + // failure mode bounded if something pathological occupies the directory. 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]]. + // `lazy val` + force-in-initialize keeps creation deterministic without + // depending on subclass field-initialiser ordering. deleteOnExit is + // avoided because the JDK retains the path for the JVM lifetime, which + // leaks in long-lived drivers. private lazy val socketDir: Path = createPrivateTempDirectory() override protected def initialize(): Unit = { @@ -99,8 +75,12 @@ class TestDirectWorkerDispatcher( /** * 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. + * (the worker's full UUID still travels via `--id`) to stay within the 108-byte + * UDS `sun_path` limit: on macOS `$TMPDIR` (~47 chars) plus the private dir + * (~29) plus a full-UUID leaf (49) overflows it and `bind(2)` fails with + * `ENAMETOOLONG`. 64 bits makes collisions negligible; the retry loop is a + * defensive guard so a near-impossible collision surfaces as a clear error + * here rather than an opaque "worker exited" from [[waitForReady]]. */ override protected def newEndpointAddress(workerId: String): String = { val short = workerId.replace("-", "").take(16) @@ -118,32 +98,62 @@ class TestDirectWorkerDispatcher( candidate.toString } - override protected def waitForReady( + override protected def connectWorker( address: String, process: Process, + outputFile: File): WorkerConnection = { + val connection = newConnection(address) + try { + waitForReady(address, connection, process, outputFile) + connection + } catch { + case e: InterruptedException => + closeFailedConnection(connection) + throw e + case NonFatal(e) => + closeFailedConnection(connection) + throw e + } + } + + private def waitForReady( + address: String, + connection: WorkerConnection, + 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 + val grpc = connection match { + case channel: GrpcWorkerChannel => channel + case other => + throw new IllegalStateException( + s"DirectGrpcDispatcher.newConnection should have produced a " + + s"GrpcWorkerChannel but got ${other.getClass.getName}") } - if (!file.exists()) { - if (process.isAlive) { - DirectWorkerDispatcher.destroyForciblyAndReap( - process, logger, s"init timeout $address") + val deadlineNanos = System.nanoTime() + TimeUnit.MILLISECONDS.toNanos(initTimeoutMs) + var socketTriggeredReconnect = false + var state = grpc.channel.getState(true) + while (state != ConnectivityState.READY) { + if (!process.isAlive) throwWorkerExitedBeforeReady(process, address, outputFile) + + val remainingNanos = deadlineNanos - System.nanoTime() + if (remainingNanos <= 0L) { 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) + s"Worker did not become reachable at $address within ${initTimeoutMs}ms\n$tail") } + + // A failed pre-bind connection attempt enters gRPC backoff. Once the + // worker creates its fresh socket, reset that backoff exactly once so + // the channel can prove readiness without waiting for a long retry. + if (state == ConnectivityState.TRANSIENT_FAILURE && + !socketTriggeredReconnect && Files.exists(Paths.get(address))) { + grpc.channel.resetConnectBackoff() + socketTriggeredReconnect = true + } + val stateChanged = new CountDownLatch(1) + grpc.channel.notifyWhenStateChanged(state, () => stateChanged.countDown()) + val pollNanos = TimeUnit.MILLISECONDS.toNanos(READY_POLL_INTERVAL_MS) + stateChanged.await(math.min(remainingNanos, pollNanos), TimeUnit.NANOSECONDS) + state = grpc.channel.getState(true) } } @@ -168,51 +178,79 @@ class TestDirectWorkerDispatcher( }) } - // `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. + // `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. + // See the contract on the abstract method in [[DirectWorkerDispatcher]]. 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, " + + "DirectGrpcDispatcher requires UNIX domain socket transport, " + s"got ${conn.getTransportCase}") + // BIDIRECTIONAL_STREAMING is the only pattern the gRPC `Execute` RPC + // speaks, so the spec MUST advertise it. We require the capabilities block + // and the pattern explicitly rather than treating an unset/empty block as + // "no constraint": a spec that does not declare bidi gives no evidence the + // worker can speak this transport, and accepting it would only defer the + // failure to stream time. require(spec.hasCapabilities, - "TestDirectWorkerDispatcher requires WorkerCapabilities declaring " + + "DirectGrpcDispatcher 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 " + + "DirectGrpcDispatcher requires BIDIRECTIONAL_STREAMING " + "in WorkerCapabilities.supported_communication_patterns") } - override protected def newConnection(address: String): WorkerConnection = - new SocketFileConnection(address) + protected def newConnection(address: String): WorkerConnection = + new GrpcWorkerChannel( + address, logger, maxInboundMessageSize = grpcMaxInboundMessageSize) - override protected def newSession( - workerHandle: WorkerHandle, - connection: WorkerConnection): WorkerSession = - new NoOpWorkerSession(workerHandle, logger) + override protected def newSession(workerHandle: WorkerHandle): WorkerSession = + workerHandle.connection match { + case g: GrpcWorkerChannel => + new GrpcWorkerSession(workerHandle, g.channel, logger) + case other => + throw new IllegalStateException( + s"DirectGrpcDispatcher.newConnection should have produced a " + + s"GrpcWorkerChannel but got ${other.getClass.getName}") + } - private def throwWorkerExitedBeforeSocket( + private def throwWorkerExitedBeforeReady( 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") + s"before becoming reachable at $address\n$tail") + } + + private def closeFailedConnection(connection: WorkerConnection): Unit = { + try connection.close() catch { + case NonFatal(e) => logger.debug("Failed to close worker connection", e) + } } /** * 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. + * + * On POSIX filesystems the permissions are applied atomically at creation via + * a file attribute, so there is '''no''' TOCTOU window. The non-POSIX branch + * cannot do that: `Files.createTempDirectory` first creates the directory with + * the platform default mask, then `File.setXxx` tightens it, leaving a brief + * window where the directory may be group/other-accessible. That race is an + * accepted limitation of the best-effort fallback -- Spark UDF workers run on + * POSIX in practice, the directory lives under the JVM temp dir, and a WARN is + * logged if the platform refuses the setters outright. Further hardening of + * the non-POSIX path (e.g. creating under an already-restricted parent) is out + * of scope here. * * The directory is anchored under a short base path (see [[shortTempBase]]) * rather than the configured `java.io.tmpdir`. Builds point the latter at a @@ -233,6 +271,7 @@ class TestDirectWorkerDispatcher( 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. + // && would short-circuit and silently leave permissions partially open. val applied = f.setReadable(false, false) & f.setWritable(false, false) & f.setExecutable(false, false) & f.setReadable(true, true) & diff --git a/udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/GrpcWorkerChannel.scala b/udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/GrpcWorkerChannel.scala new file mode 100644 index 0000000000000..07ae58415f3d0 --- /dev/null +++ b/udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/GrpcWorkerChannel.scala @@ -0,0 +1,148 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.udf.worker.grpc + +import java.util.concurrent.TimeUnit +import java.util.concurrent.atomic.AtomicBoolean + +import scala.util.control.NonFatal + +import io.grpc.{ConnectivityState, ManagedChannel} +import io.grpc.netty.NettyChannelBuilder +import io.netty.channel.EventLoopGroup +import io.netty.channel.unix.DomainSocketAddress + +import org.apache.spark.annotation.Experimental +import org.apache.spark.udf.worker.core.{WorkerConnection, WorkerLogger} + +/** + * :: Experimental :: + * A [[WorkerConnection]] backed by a gRPC `ManagedChannel` over a Netty + * Unix domain socket transport. + * + * Owns: + * - the [[ManagedChannel]] used by sessions to open Execute streams, + * - a dedicated Netty [[EventLoopGroup]] for that channel (one per worker so + * a slow worker cannot starve others sharing a global loop). + * + * Does NOT own the socket file at `socketPath`: the worker creates it and the + * dispatcher removes it (`cleanupEndpointAddress` / `closeTransport`). Keeping + * a single creation/deletion site avoids double-cleanup and ambiguous + * ownership. + * + * Lifecycle: + * - [[isActive]] is true while the channel reports any non-`SHUTDOWN` state. + * Transient failures keep the connection usable -- gRPC will reconnect on + * the next RPC. + * - [[close]] shuts the channel down and drains the event loop. Idempotent. + * + * Construction does NOT block on the worker actually being reachable; the + * dispatcher proves liveness separately via its + * `waitForReady`/init-timeout machinery before handing the connection to a + * session. + * + * @param socketPath path of the UDS the worker server binds. + * @param logger logger for shutdown diagnostics; defaults to NoOp. + * @param shutdownPhaseTimeoutMs upper bound on each individual phase of + * close() -- graceful channel drain, forced + * channel drain, and event-loop drain are + * three separate phases that each get this + * budget. Total worst-case close time is up + * to three times this value. Defaults to 5s + * per phase. + * @param maxInboundMessageSize maximum serialized response size accepted by + * the gRPC channel. Defaults to 128 MiB. + */ +@Experimental +class GrpcWorkerChannel( + val socketPath: String, + logger: WorkerLogger = WorkerLogger.NoOp, + shutdownPhaseTimeoutMs: Long = GrpcWorkerChannel.DEFAULT_SHUTDOWN_PHASE_TIMEOUT_MS, + maxInboundMessageSize: Int = GrpcWorkerChannel.DEFAULT_MAX_INBOUND_MESSAGE_SIZE) + extends WorkerConnection { + + require(maxInboundMessageSize > 0, "maxInboundMessageSize must be positive") + + private val transport: UnixDomainSocketTransport = UnixDomainSocketTransport.detect() + private val eventLoopGroup: EventLoopGroup = transport.newEventLoopGroup() + + /** + * Underlying gRPC channel. Sessions open one bidirectional Execute stream + * per session on this channel. + */ + val channel: ManagedChannel = try { + NettyChannelBuilder + .forAddress(new DomainSocketAddress(socketPath)) + .channelType(transport.channelType) + .eventLoopGroup(eventLoopGroup) + .maxInboundMessageSize(maxInboundMessageSize) + .usePlaintext() + .build() + } catch { + case NonFatal(e) => + UnixDomainSocketTransport.shutdown(eventLoopGroup, shutdownPhaseTimeoutMs) + throw e + } + + private val closed = new AtomicBoolean(false) + + override def isActive: Boolean = { + if (closed.get()) return false + channel.getState(false) != ConnectivityState.SHUTDOWN + } + + /** + * Idempotently tears down the channel and event loop. Steps are guarded so + * a failure in one does not skip the others. The socket file is owned and + * removed by the dispatcher, not here. + */ + override def close(): Unit = { + if (!closed.compareAndSet(false, true)) return + + try { + channel.shutdown() + if (!channel.awaitTermination(shutdownPhaseTimeoutMs, TimeUnit.MILLISECONDS)) { + channel.shutdownNow() + channel.awaitTermination(shutdownPhaseTimeoutMs, TimeUnit.MILLISECONDS) + } + } catch { + case _: InterruptedException => + Thread.currentThread().interrupt() + case NonFatal(e) => + logger.warn(s"Error shutting down gRPC channel for $socketPath", e) + } + + try { + UnixDomainSocketTransport.shutdown(eventLoopGroup, shutdownPhaseTimeoutMs) + } catch { + case NonFatal(e) => + logger.warn(s"Error shutting down event loop for $socketPath", e) + } + } +} + +object GrpcWorkerChannel { + /** Matches Spark Connect's default and accommodates normal Arrow batches. */ + val DEFAULT_MAX_INBOUND_MESSAGE_SIZE: Int = 128 * 1024 * 1024 + + /** + * Default per-phase upper bound on close() drain time. Each of the three + * drain phases (graceful channel, forced channel, event loop) is capped + * by this value, so total worst-case close time is up to 3x. + */ + val DEFAULT_SHUTDOWN_PHASE_TIMEOUT_MS: Long = 5000L +} diff --git a/udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/UnixDomainSocketTransport.scala b/udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/UnixDomainSocketTransport.scala new file mode 100644 index 0000000000000..ae2e2850d7ef4 --- /dev/null +++ b/udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/UnixDomainSocketTransport.scala @@ -0,0 +1,148 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.udf.worker.grpc + +import java.util.Locale +import java.util.concurrent.TimeUnit + +import io.netty.channel.{Channel, EventLoopGroup, MultiThreadIoEventLoopGroup, ServerChannel} + +/** + * Selects the Netty native transport that supports Unix domain sockets on the + * current OS. + * + * gRPC's `NettyChannelBuilder` / `NettyServerBuilder` can talk over UDS only when + * paired with a native transport that exposes a `DomainSocketChannel`: + * `epoll` on Linux and `kqueue` on macOS. NIO does not support UDS in a way + * gRPC's Netty transport can use. + * + * Only `epoll` and `kqueue` are wired up here. Netty also ships an `io_uring` + * transport (`io.netty.channel.uring`) that supports UDS on recent Linux + * kernels; it could be added as a preferred Linux option in [[detect]] once its + * native artifact is on the classpath, but is intentionally left out for now to + * keep the dependency set minimal. + * + * Callers obtain a [[UnixDomainSocketTransport]] via [[UnixDomainSocketTransport.detect]], + * then ask it for matching `channelType` / `serverChannelType` / `EventLoopGroup` + * objects to configure the builder. + */ +private[grpc] sealed trait UnixDomainSocketTransport { + def channelType: Class[_ <: Channel] + def serverChannelType: Class[_ <: ServerChannel] + def newEventLoopGroup(): EventLoopGroup + def name: String +} + +private[grpc] object UnixDomainSocketTransport { + + /** Linux epoll. */ + private case object EpollTransport extends UnixDomainSocketTransport { + override val channelType: Class[_ <: Channel] = + classOf[io.netty.channel.epoll.EpollDomainSocketChannel] + override val serverChannelType: Class[_ <: ServerChannel] = + classOf[io.netty.channel.epoll.EpollServerDomainSocketChannel] + override def newEventLoopGroup(): EventLoopGroup = + new MultiThreadIoEventLoopGroup(1, io.netty.channel.epoll.EpollIoHandler.newFactory()) + override val name: String = "epoll" + } + + /** macOS/BSD kqueue. */ + private case object KQueueTransport extends UnixDomainSocketTransport { + override val channelType: Class[_ <: Channel] = + classOf[io.netty.channel.kqueue.KQueueDomainSocketChannel] + override val serverChannelType: Class[_ <: ServerChannel] = + classOf[io.netty.channel.kqueue.KQueueServerDomainSocketChannel] + override def newEventLoopGroup(): EventLoopGroup = + new MultiThreadIoEventLoopGroup(1, io.netty.channel.kqueue.KQueueIoHandler.newFactory()) + override val name: String = "kqueue" + } + + // NOTE: io.netty.channel.epoll.Epoll.isAvailable triggers a one-shot + // native library load attempt the first time it is called. On platforms + // that don't ship epoll (e.g. macOS) the load fails silently inside + // Netty and is cached. The OS-aware paths in `detect()` below short- + // circuit on os.name first so we don't pay this on the dominant happy + // path -- only the fallback paths (wrong os.name string, container with + // mismatched native libs) reach this side-effect. + private def epollAvailable: Boolean = io.netty.channel.epoll.Epoll.isAvailable + private def epollUnavailable: Throwable = + io.netty.channel.epoll.Epoll.unavailabilityCause() + private def kqueueAvailable: Boolean = io.netty.channel.kqueue.KQueue.isAvailable + private def kqueueUnavailable: Throwable = + io.netty.channel.kqueue.KQueue.unavailabilityCause() + + private def isLinux(os: String): Boolean = os.contains("linux") + private def isMacOrBsd(os: String): Boolean = + os.contains("mac") || os.contains("darwin") || os.contains("bsd") + + /** + * Returns the transport for the current OS, or throws + * [[UnsupportedOperationException]] if neither epoll nor kqueue is + * available. + * + * Resolution order: + * 1. Linux + epoll available (the dominant happy path). + * 2. macOS/BSD + kqueue available. + * 3. Either native transport available regardless of `os.name` + * (covers misreported os.name, containers that ship the wrong + * native libs, etc.). The `&&` short-circuits in steps 1 and 2 + * mean we only trigger the "wrong-OS" native load attempt in + * this fallback branch. + */ + def detect(): UnixDomainSocketTransport = { + val os = System.getProperty("os.name", "").toLowerCase(Locale.ROOT) + if (isLinux(os) && epollAvailable) { + EpollTransport + } else if (isMacOrBsd(os) && kqueueAvailable) { + KQueueTransport + } else if (epollAvailable) { + EpollTransport + } else if (kqueueAvailable) { + KQueueTransport + } else { + // No socket-based UDS fallback exists: gRPC's Netty transport can only + // speak UDS through a native epoll/kqueue DomainSocketChannel. (Java NIO + // exposes UDS via java.nio.channels.SocketChannel, but Netty/gRPC cannot + // use it as the channel type for a managed channel.) When neither native + // transport is present the only alternative is a different transport + // entirely -- e.g. TCP loopback -- which is a deployment choice, not a + // fallback we silently substitute here, so we fail loudly instead. + val epollErr = Option(epollUnavailable).map(_.getMessage).getOrElse("unknown") + val kqueueErr = Option(kqueueUnavailable).map(_.getMessage).getOrElse("unknown") + throw new UnsupportedOperationException( + s"No Netty native UDS transport available on os=$os " + + s"(epoll: $epollErr, kqueue: $kqueueErr). " + + "UDS-backed gRPC requires netty-transport-native-epoll on Linux or " + + "netty-transport-native-kqueue on macOS.") + } + } + + /** + * Shut an event loop group down. Best-effort: preserves the interrupt + * flag on `InterruptedException` (does not propagate the exception) so + * callers in a cleanup path are not forced to handle it. + */ + def shutdown(elg: EventLoopGroup, timeoutMs: Long): Unit = { + if (elg == null || elg.isShuttingDown) return + val future = elg.shutdownGracefully(0L, timeoutMs, TimeUnit.MILLISECONDS) + try { + future.await(timeoutMs, TimeUnit.MILLISECONDS) + } catch { + case _: InterruptedException => Thread.currentThread().interrupt() + } + } +} diff --git a/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/DirectGrpcDispatcherIntegrationSuite.scala b/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/DirectGrpcDispatcherIntegrationSuite.scala new file mode 100644 index 0000000000000..ee07fd73ef30f --- /dev/null +++ b/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/DirectGrpcDispatcherIntegrationSuite.scala @@ -0,0 +1,530 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.udf.worker.grpc + +import java.io.File +import java.nio.file.Paths +import java.util.Locale + +import com.google.protobuf.ByteString +import org.scalatest.BeforeAndAfterEach +// scalastyle:off funsuite +import org.scalatest.funsuite.AnyFunSuite + +import org.apache.spark.udf.worker.{Cancel, DataRequest, DataResponse, DirectWorker, Finish, Init, + ProcessCallable, UdfPayload, UDFProtoCommunicationPattern, UDFWorkerDataFormat, + UDFWorkerProperties, UDFWorkerSpecification, UnixDomainSocket, WorkerCapabilities, + WorkerConnectionSpec} +import org.apache.spark.udf.worker.core.{Termination, WorkerSession} +import org.apache.spark.udf.worker.core.direct.DirectWorkerProcess +import org.apache.spark.udf.worker.grpc.testing.EchoGrpcWorkerMain + +/** + * Integration tests that spawn the [[EchoGrpcWorkerMain]] subprocess via a + * [[DirectGrpcDispatcher]] and exercise the engine-side + * dispatcher / session lifecycle end to end (real gRPC, real OS process). + * + * These tests run only on platforms where a Netty UDS transport is + * available (Linux epoll or macOS kqueue); other platforms `assume`-skip. + */ +class DirectGrpcDispatcherIntegrationSuite + extends AnyFunSuite with BeforeAndAfterEach { +// scalastyle:on funsuite + + /** + * Full Java classpath inherited from the SBT/Maven test runner. The + * dispatcher spawns `java -cp EchoGrpcWorkerMain` with this so the + * child JVM can resolve the test sources and all transitive deps. + */ + private val javaClasspath: String = System.getProperty("java.class.path") + + private val javaExecutable: String = + Paths.get(System.getProperty("java.home"), "bin", "java").toString + + private val workerMainClass: String = + classOf[EchoGrpcWorkerMain.type].getName.stripSuffix("$") + + private var dispatcher: DirectGrpcDispatcher = _ + + override def beforeEach(): Unit = { + val supported = try { + UnixDomainSocketTransport.detect(); true + } catch { + case _: UnsupportedOperationException => false + } + assume(supported, + "Netty UDS native transport (epoll on Linux or kqueue on macOS) is " + + "required for the gRPC direct dispatcher") + } + + override def afterEach(): Unit = { + if (dispatcher != null) { + try dispatcher.close() finally dispatcher = null + } + super.afterEach() + } + + // -------------------------------------------------------------------------- + // Spec builders + // -------------------------------------------------------------------------- + + private def echoRunner: ProcessCallable = ProcessCallable.newBuilder() + .addCommand(javaExecutable) + .addCommand("-cp").addCommand(javaClasspath) + .addCommand(workerMainClass) + .build() + + private def udsProperties: UDFWorkerProperties = UDFWorkerProperties.newBuilder() + .setConnection(WorkerConnectionSpec.newBuilder() + .setUnixDomainSocket(UnixDomainSocket.getDefaultInstance) + .build()) + // Generous init timeout: child JVM startup includes Netty native lib + // loading, several seconds on first launch. + .setInitializationTimeoutMs(30000) + .setGracefulTerminationTimeoutMs(10000) + .build() + + private def capabilities: WorkerCapabilities = WorkerCapabilities.newBuilder() + .addSupportedDataFormats(UDFWorkerDataFormat.ARROW) + .addSupportedCommunicationPatterns(UDFProtoCommunicationPattern.BIDIRECTIONAL_STREAMING) + .build() + + private def echoSpec: UDFWorkerSpecification = UDFWorkerSpecification.newBuilder() + .setCapabilities(capabilities) + .setDirect(DirectWorker.newBuilder() + .setRunner(echoRunner) + .setProperties(udsProperties) + .build()) + .build() + + private def newDispatcher(): DirectGrpcDispatcher = { + dispatcher = new DirectGrpcDispatcher(echoSpec) + dispatcher + } + + private def basicInit(payload: Array[Byte] = "echo".getBytes): Init = + Init.newBuilder() + .setProtocolVersion(1) + .setDataFormat(UDFWorkerDataFormat.ARROW) + .setUdf(UdfPayload.newBuilder() + .setPayload(ByteString.copyFrom(payload)) + .setFormat("echo") + .build()) + .build() + + // Default lifecycle messages for the data phase and finalization. + private val emptyFinish: () => Finish = () => Finish.getDefaultInstance + private val emptyCancel: () => Cancel = () => Cancel.getDefaultInstance + + /** Wraps strings as input [[DataRequest]] batches. */ + private def echoIn(batches: String*): Iterator[DataRequest] = + batches.iterator.map(s => DataRequest.newBuilder() + .setData(ByteString.copyFrom(s.getBytes)).build()) + + /** Decodes a result [[DataResponse]] batch to a string. */ + private def outStr(d: DataResponse): String = new String(d.getData.toByteArray) + + /** Runs `f` with a fresh session, ensuring it is closed afterwards. */ + private def withSession[T](f: WorkerSession => T): T = { + val session = dispatcher.createSession(None) + try f(session) finally session.close(emptyCancel) + } + + /** + * Centralised cast for tests that need to reach past the + * [[org.apache.spark.udf.worker.core.WorkerHandle]] interface to the + * concrete [[DirectWorkerProcess]] (e.g. to assert process liveness + * or read the underlying [[GrpcWorkerChannel]]). Fails with a clear + * test failure if the type is unexpected, rather than the opaque + * `ClassCastException` we'd get from a bare `asInstanceOf`. + */ + private def workerProcessOf(s: WorkerSession): DirectWorkerProcess = + s.workerHandle match { + case dp: DirectWorkerProcess => dp + case other => fail( + s"Expected DirectWorkerProcess, got ${other.getClass.getSimpleName}") + } + + // -------------------------------------------------------------------------- + // Tests + // -------------------------------------------------------------------------- + + test("dispatcher construction: smoke test") { + newDispatcher() + assert(dispatcher.getClass == classOf[DirectGrpcDispatcher], + s"expected a DirectGrpcDispatcher, got ${dispatcher.getClass}") + } + + test("dispatcher spawns a worker process and a session round-trips an echo batch") { + newDispatcher() + + withSession { session => + session.init(basicInit()) + val results = session.process(echoIn("hello"), emptyFinish).toList + assert(results.map(outStr) == List("hello")) + } + } + + test("responses larger than gRPC's default inbound limit round-trip successfully") { + newDispatcher() + val payload = ByteString.copyFrom(Array.fill[Byte](5 * 1024 * 1024)(7)) + + withSession { session => + session.init(basicInit()) + val input = DataRequest.newBuilder().setData(payload).build() + val results = session.process(Iterator.single(input), emptyFinish).toList + assert(results.map(_.getData) == List(payload)) + } + } + + test("multiple batches are echoed in order on a single session") { + newDispatcher() + + withSession { session => + session.init(basicInit()) + val inputs = Seq("alpha", "beta", "gamma", "delta") + val results = session.process(echoIn(inputs: _*), emptyFinish).map(outStr).toList + assert(results == inputs) + } + } + + test("two sequential sessions each spawn and clean up their own worker") { + newDispatcher() + + val s1 = dispatcher.createSession(None) + try { + s1.init(basicInit()) + val r1 = s1.process(echoIn("one"), emptyFinish).map(outStr).toList + assert(r1 == List("one")) + } finally s1.close(emptyCancel) + + val s2 = dispatcher.createSession(None) + try { + s2.init(basicInit()) + val r2 = s2.process(echoIn("two"), emptyFinish).map(outStr).toList + assert(r2 == List("two")) + } finally s2.close(emptyCancel) + } + + test("worker process is terminated when session closes") { + newDispatcher() + + val session = dispatcher.createSession(None) + val workerProcess = workerProcessOf(session) + try { + session.init(basicInit()) + session.process(echoIn("ping"), emptyFinish).toList + assert(workerProcess.process.isAlive, "worker should be alive while session is open") + } finally session.close(emptyCancel) + + assert(!workerProcess.process.isAlive, + "worker process should be terminated after session.close()") + val socketFile = workerProcess.connection match { + case g: GrpcWorkerChannel => new File(g.socketPath) + case _ => null + } + if (socketFile != null) { + assert(!socketFile.exists(), + s"socket file ${socketFile.getAbsolutePath} should be removed after close") + } + } + + test("dispatcher.close terminates any remaining worker processes") { + newDispatcher() + + val session = dispatcher.createSession(None) + session.init(basicInit()) + val workerProcess = workerProcessOf(session) + // Drive one batch through so we are sure the worker is fully up. + val it = session.process(echoIn("once"), emptyFinish) + assert(outStr(it.next()) == "once") + + // Deliberately do NOT close the session before closing the dispatcher. + dispatcher.close() + dispatcher = null + + val deadline = System.currentTimeMillis() + 5000 + while (workerProcess.process.isAlive && System.currentTimeMillis() < deadline) { + Thread.sleep(50) + } + assert(!workerProcess.process.isAlive, + "worker process should be terminated by dispatcher.close") + } + + test("worker user-error is surfaced as GrpcWorkerSessionException") { + newDispatcher() + + withSession { session => + session.init(basicInit()) + // The echo worker emits ErrorResponse(UserError) when it sees this batch. + val it = session.process(echoIn("ERROR"), emptyFinish) + val ex = intercept[GrpcWorkerSessionException] { it.toList } + assert(ex.executionError != null, "execution error must be attached") + assert(ex.executionError.hasUser, s"expected UserError, got $ex") + assert(ex.executionError.getUser.getErrorClass == "SimulatedError") + } + } + + test("init failure surfaces from session.init()") { + newDispatcher() + + val session = dispatcher.createSession(None) + try { + val init = Init.newBuilder(basicInit()) + // The echo worker treats an init payload of "INIT_ERROR" as a + // simulated init failure (InitResponse with WorkerError set). + .setUdf(UdfPayload.newBuilder() + .setPayload(ByteString.copyFromUtf8("INIT_ERROR")) + .setFormat("echo") + .build()) + .build() + val ex = intercept[GrpcWorkerSessionException] { session.init(init) } + assert(ex.executionError != null && ex.executionError.hasWorker) + assert(ex.executionError.getWorker.getMessage == "simulated init failure") + } finally session.close(emptyCancel) + } + + test("close after init but before process cancels the worker cleanly") { + newDispatcher() + + val session = dispatcher.createSession(None) + session.init(basicInit()) + // No process(): finalizing now must cancel the in-flight session and + // return the CancelResponse terminator. + val start = System.currentTimeMillis() + val termination = session.close(emptyCancel) + val elapsed = System.currentTimeMillis() - start + assert(termination.isInstanceOf[Termination.Cancelled], + s"expected a Cancelled termination, got $termination") + // A genuine cooperative cancel returns as soon as the worker replies with + // CancelResponse. A broken handshake would instead block until the + // terminalTimeoutMs (default 30s) and then settle a Termination.TransportFailed, + // so the Cancelled assertion above plus a prompt return distinguish the clean + // case from a timeout. + assert(elapsed < 10000, + s"close() should return promptly on a clean cancel, took ${elapsed}ms") + } + + test("close before init does not leak the worker process") { + newDispatcher() + + val session = dispatcher.createSession(None) + val workerProcess = workerProcessOf(session) + session.close(emptyCancel) + val deadline = System.currentTimeMillis() + 5000 + while (workerProcess.process.isAlive && System.currentTimeMillis() < deadline) { + Thread.sleep(50) + } + assert(!workerProcess.process.isAlive, + "worker should be terminated by close() even when init was never called") + } + + test("backend.validateSpec rejects spec without BIDIRECTIONAL_STREAMING when " + + "capabilities are explicitly declared") { + val incompatibleCaps = WorkerCapabilities.newBuilder() + .addSupportedDataFormats(UDFWorkerDataFormat.ARROW) + .addSupportedCommunicationPatterns( + UDFProtoCommunicationPattern.UDF_PROTO_COMMUNICATION_PATTERN_UNSPECIFIED) + .build() + val spec = UDFWorkerSpecification.newBuilder() + .setCapabilities(incompatibleCaps) + .setDirect(DirectWorker.newBuilder() + .setRunner(echoRunner).setProperties(udsProperties).build()) + .build() + val ex = intercept[IllegalArgumentException] { + new DirectGrpcDispatcher(spec) + } + assert(ex.getMessage.contains("BIDIRECTIONAL_STREAMING"), + s"expected pattern-required error, got: ${ex.getMessage}") + } + + test("backend.validateSpec rejects a spec that declares no capabilities at all") { + val spec = UDFWorkerSpecification.newBuilder() + .setDirect(DirectWorker.newBuilder() + .setRunner(echoRunner).setProperties(udsProperties).build()) + .build() + val ex = intercept[IllegalArgumentException] { + new DirectGrpcDispatcher(spec) + } + assert(ex.getMessage.contains("BIDIRECTIONAL_STREAMING"), + s"expected capabilities-required error, got: ${ex.getMessage}") + } + + test("a socket path alone does not make a worker ready") { + val socketOnlyWorker = + """ + |socket_path="" + |while [[ $# -gt 0 ]]; do + | case "$1" in + | --connection) socket_path="$2"; shift 2 ;; + | *) shift ;; + | esac + |done + |trap 'rm -f "$socket_path"; exit 0' SIGTERM + |touch "$socket_path" + |echo socket-created + |while true; do sleep 1; done + """.stripMargin.trim + val runner = ProcessCallable.newBuilder() + .addCommand("bash").addCommand("-c").addCommand(socketOnlyWorker).addCommand("--") + .build() + val properties = UDFWorkerProperties.newBuilder() + .setConnection(WorkerConnectionSpec.newBuilder() + .setUnixDomainSocket(UnixDomainSocket.getDefaultInstance) + .build()) + .setInitializationTimeoutMs(1000) + .setGracefulTerminationTimeoutMs(1000) + .build() + val spec = UDFWorkerSpecification.newBuilder() + .setCapabilities(capabilities) + .setDirect(DirectWorker.newBuilder().setRunner(runner).setProperties(properties).build()) + .build() + dispatcher = new DirectGrpcDispatcher(spec) + + val error = intercept[org.apache.spark.udf.worker.core.direct.DirectWorkerTimeoutException] { + dispatcher.createSession(None) + } + assert(error.getMessage.contains("did not become reachable")) + assert(error.getMessage.contains("socket-created"), + s"expected worker output in timeout, got: ${error.getMessage}") + } + + test("dispatcher cleans up socket directory on close") { + newDispatcher() + val session = dispatcher.createSession(None) + val workerProcess = workerProcessOf(session) + val socketPath = workerProcess.connection match { + case g: GrpcWorkerChannel => g.socketPath + case _ => fail("expected GrpcWorkerChannel") + } + val socketDir = new File(socketPath).getParentFile + assert(socketDir.exists(), s"socket dir $socketDir should exist while session is open") + session.close(emptyCancel) + dispatcher.close() + dispatcher = null + assert(!socketDir.exists(), + s"socket dir $socketDir should be removed when dispatcher closes") + } + + test("concurrent sessions on a single dispatcher each get their own worker " + + "and round-trip an echo batch in parallel") { + newDispatcher() + val threads = 4 + val errors = new java.util.concurrent.ConcurrentLinkedQueue[Throwable]() + val results = new java.util.concurrent.ConcurrentLinkedQueue[String]() + val barrier = new java.util.concurrent.CyclicBarrier(threads) + val handles = (1 to threads).map { i => + val t = new Thread(() => { + try { + val session = dispatcher.createSession(None) + try { + session.init(basicInit()) + barrier.await(30, java.util.concurrent.TimeUnit.SECONDS) + val out = session.process( + echoIn(s"thread-$i"), emptyFinish).toList + require(out.size == 1) + results.add(outStr(out.head)) + } finally session.close(emptyCancel) + } catch { + case t: Throwable => errors.add(t) + } + }, s"concurrent-session-$i") + t.start() + t + } + handles.foreach(_.join(30000)) + assert(errors.isEmpty, + s"unexpected errors: ${errors.toArray.mkString(", ")}") + assert(results.size == threads, s"expected $threads results, got ${results.size}") + import scala.jdk.CollectionConverters._ + val echoed = results.asScala.toSet + val expected = (1 to threads).map(i => s"thread-$i").toSet + assert(echoed == expected, s"echoed=$echoed expected=$expected") + } + + test("worker crash mid-process surfaces a transport error and marks worker invalid") { + newDispatcher() + val session = dispatcher.createSession(None) + val workerProc = workerProcessOf(session) + try { + session.init(basicInit()) + val it = session.process(echoIn("hello"), emptyFinish) + assert(outStr(it.next()) == "hello", "first batch must echo cleanly") + + // SIGKILL the worker subprocess to simulate an unexpected crash. + workerProc.process.destroyForcibly() + // Wait for the OS to reap it so the gRPC channel actually notices. + val killed = workerProc.process.waitFor( + 5, java.util.concurrent.TimeUnit.SECONDS) + assert(killed, "worker process should have died after destroyForcibly") + + // The Java iterator must surface the transport failure. + val ex = intercept[GrpcWorkerSessionException] { + while (it.hasNext) it.next() + } + assert(ex.getMessage.toLowerCase(Locale.ROOT).contains("worker") || + ex.getMessage.toLowerCase(Locale.ROOT).contains("stream"), + s"expected transport-failure message, got: ${ex.getMessage}") + } finally session.close(emptyCancel) + // Reuse-readiness: the session should have flagged the worker invalid so + // a future pool will not recycle it. + assert(workerProc.isInvalid, + "worker should be flagged invalid after a transport failure") + } + + test("spec-driven smoke: start from worker spec, get a working session in one factory call") { + // The canonical "from a spec, get a dispatcher, get a session, run a UDF" + // path that documentation can point to. + val spec = UDFWorkerSpecification.newBuilder() + .setCapabilities(WorkerCapabilities.newBuilder() + .addSupportedDataFormats(UDFWorkerDataFormat.ARROW) + .addSupportedCommunicationPatterns( + UDFProtoCommunicationPattern.BIDIRECTIONAL_STREAMING) + .build()) + .setDirect(DirectWorker.newBuilder() + .setRunner(echoRunner) + .setProperties(udsProperties) + .build()) + .build() + val d = new DirectGrpcDispatcher(spec) + try { + val s = d.createSession(None) + try { + s.init(basicInit()) + val out = s.process(echoIn("hello"), emptyFinish) + .map(outStr).toList + assert(out == List("hello")) + } finally s.close(emptyCancel) + } finally d.close() + } + + test("clean session end leaves worker salvageable (reuse-readiness invariant)") { + // Future pooling relies on workerProcess.isInvalid being false for a + // session that completed without errors. Pin this contract now. + newDispatcher() + val session = dispatcher.createSession(None) + val workerProc = workerProcessOf(session) + try { + session.init(basicInit()) + val out = session.process(echoIn("ping"), emptyFinish).toList + assert(out.map(outStr) == List("ping")) + } finally session.close(emptyCancel) + assert(!workerProc.isInvalid, + "worker must remain salvageable after a clean Finish") + } + +} diff --git a/udf/worker/core/src/test/scala/org/apache/spark/udf/worker/core/DirectWorkerDispatcherSuite.scala b/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/DirectWorkerDispatcherSuite.scala similarity index 90% rename from udf/worker/core/src/test/scala/org/apache/spark/udf/worker/core/DirectWorkerDispatcherSuite.scala rename to udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/DirectWorkerDispatcherSuite.scala index 28ecc1179f6ac..86aa443803516 100644 --- a/udf/worker/core/src/test/scala/org/apache/spark/udf/worker/core/DirectWorkerDispatcherSuite.scala +++ b/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/DirectWorkerDispatcherSuite.scala @@ -14,7 +14,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.spark.udf.worker.core +package org.apache.spark.udf.worker.grpc import java.io.File import java.nio.file.{Files, Path} @@ -34,6 +34,7 @@ import org.apache.spark.udf.worker.core.{WorkerConnection, WorkerHandle, WorkerS WorkerSession} import org.apache.spark.udf.worker.core.direct.{DirectWorkerException, DirectWorkerProcess, DirectWorkerTimeoutException} +import org.apache.spark.udf.worker.grpc.testing.{SocketFileConnection, TestDirectGrpcDispatcher} /** * Tests for [[DirectWorkerDispatcher]] process lifecycle: spawning workers @@ -72,7 +73,7 @@ class DirectWorkerDispatcherSuite private def directWorker(runner: ProcessCallable): DirectWorker = DirectWorker.newBuilder().setRunner(runner).setProperties(udsProperties).build() - // TestDirectWorkerDispatcher requires the spec to advertise BIDIRECTIONAL_STREAMING. + // DirectGrpcDispatcher requires the spec to advertise BIDIRECTIONAL_STREAMING. private def bidiCapabilities: WorkerCapabilities = WorkerCapabilities.newBuilder() .addSupportedCommunicationPatterns(UDFProtoCommunicationPattern.BIDIRECTIONAL_STREAMING) .build() @@ -92,7 +93,7 @@ class DirectWorkerDispatcherSuite .setDirect(directWorker(runner)) .build() - private var dispatcher: TestDirectWorkerDispatcher = _ + private var dispatcher: TestDirectGrpcDispatcher = _ // Finalization message for session.close(); these lifecycle-only tests never // drive the protocol, so an empty Cancel suffices. @@ -113,7 +114,7 @@ class DirectWorkerDispatcherSuite private def createStubSession(): WorkerSession = dispatcher.createSession(None) - // Centralised cast: every test in this suite uses TestDirectWorkerDispatcher, + // Centralised cast: every test in this suite uses TestDirectGrpcDispatcher, // 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 @@ -136,7 +137,7 @@ class DirectWorkerDispatcherSuite } test("creates a worker and session") { - dispatcher = new TestDirectWorkerDispatcher(specWithRunner(defaultRunner)) + dispatcher = new TestDirectGrpcDispatcher(specWithRunner(defaultRunner)) val session = createStubSession() val worker = workerProcess(session) @@ -150,7 +151,7 @@ class DirectWorkerDispatcherSuite } test("concurrent createSession calls produce distinct workers") { - dispatcher = new TestDirectWorkerDispatcher(specWithRunner(defaultRunner)) + dispatcher = new TestDirectGrpcDispatcher(specWithRunner(defaultRunner)) val threads = 8 val sessions = new java.util.concurrent.ConcurrentLinkedQueue[WorkerSession]() @@ -194,7 +195,7 @@ class DirectWorkerDispatcherSuite } test("close shuts down all workers via SIGTERM") { - dispatcher = new TestDirectWorkerDispatcher(specWithRunner(defaultRunner)) + dispatcher = new TestDirectGrpcDispatcher(specWithRunner(defaultRunner)) val session1 = createStubSession() val session2 = createStubSession() @@ -243,7 +244,7 @@ class DirectWorkerDispatcherSuite .setDirect(DirectWorker.newBuilder() .setRunner(runner).setProperties(shortGracefulProps).build()) .build() - dispatcher = new TestDirectWorkerDispatcher(spec) + dispatcher = new TestDirectGrpcDispatcher(spec) val session = createStubSession() val worker = workerProcess(session) @@ -261,7 +262,7 @@ class DirectWorkerDispatcherSuite } test("closing a session terminates its worker") { - dispatcher = new TestDirectWorkerDispatcher(specWithRunner(defaultRunner)) + dispatcher = new TestDirectGrpcDispatcher(specWithRunner(defaultRunner)) val session = createStubSession() val worker = workerProcess(session) @@ -281,7 +282,7 @@ class DirectWorkerDispatcherSuite } test("concurrent session.close and dispatcher.close do not double-close the worker") { - dispatcher = new TestDirectWorkerDispatcher(specWithRunner(defaultRunner)) + dispatcher = new TestDirectGrpcDispatcher(specWithRunner(defaultRunner)) val sessions = (1 to 4).map(_ => createStubSession()) val workers = sessions.map(workerProcess) @@ -325,17 +326,14 @@ class DirectWorkerDispatcherSuite } test("close racing with in-flight createSession does not leak the worker") { - // The acquire-before-publish + post-publish closed re-check pattern in - // createSession is designed for this race: thread A is mid-spawn when - // thread B calls close(). Thread A must either throw IllegalStateException - // (post-publish check caught the close) or receive a session whose worker - // is reaped by close()'s iteration. No orphan process or socket file - // should remain in either case. + // close() must wait for thread A to leave createSession before tearing down + // workers or dispatcher-level transport state. val readyLatch = new java.util.concurrent.CountDownLatch(1) val releaseLatch = new java.util.concurrent.CountDownLatch(1) val capturedWorkers = new java.util.concurrent.ConcurrentLinkedQueue[DirectWorkerProcess]() - val racing = new TestDirectWorkerDispatcher(specWithRunner(defaultRunner)) { + val transportClosed = new java.util.concurrent.atomic.AtomicBoolean(false) + val racing = new TestDirectGrpcDispatcher(specWithRunner(defaultRunner)) { override protected def afterWorkerRegistered(worker: DirectWorkerProcess): Unit = { capturedWorkers.add(worker) readyLatch.countDown() @@ -345,6 +343,11 @@ class DirectWorkerDispatcherSuite fail("releaseLatch never fired -- test orchestration broken") } } + + override protected def closeTransport(): Unit = { + transportClosed.set(true) + super.closeTransport() + } } try { val outcome = @@ -366,8 +369,11 @@ class DirectWorkerDispatcherSuite val closeThread = new Thread(() => racing.close(), "close-racer") closeThread.start() - // Give close() time to flip `closed` and iterate workers. + // Give close() time to flip `closed` and reach the creation barrier. Thread.sleep(200) + assert(closeThread.isAlive, "close should wait for the in-flight createSession") + assert(!transportClosed.get(), + "transport cleanup must not run while createSession is in flight") // Now release the in-flight createSession. releaseLatch.countDown() @@ -384,18 +390,11 @@ class DirectWorkerDispatcherSuite outcome.get() match { case Left(e: IllegalStateException) => - // Contractually allowed, but unreachable with this orchestration: - // readyLatch only fires after createSession has cleared both - // `closed` checks, so B's close cannot flip `closed` in time for - // A to observe it. Kept defensive so a future internal change - // that introduces a new window is still covered. assert(e.getMessage.contains("closed"), s"expected dispatcher-closed error, got: ${e.getMessage}") case Left(other) => fail(s"unexpected exception from racing createSession: $other") - case Right(_) => - // close() iterated the published worker and tore it down; the - // returned session points at a worker that should now be dead. + case Right(_) => fail("createSession should fail when close wins during the test hook") } // Whichever path won, the worker must not still be running and the @@ -428,7 +427,7 @@ class DirectWorkerDispatcherSuite .setDirect(DirectWorker.newBuilder() .setRunner(defaultRunner).setProperties(oversizedProps).build()) .build() - dispatcher = new TestDirectWorkerDispatcher(spec) + dispatcher = new TestDirectGrpcDispatcher(spec) val session = createStubSession() val worker = workerProcess(session) @@ -448,14 +447,14 @@ class DirectWorkerDispatcherSuite .setDirect(DirectWorker.newBuilder() .setRunner(defaultRunner).setProperties(oversizedProps).build()) .build() - dispatcher = new TestDirectWorkerDispatcher(spec) + dispatcher = new TestDirectGrpcDispatcher(spec) assert(dispatcher.initTimeoutMs == 30000L, s"init timeout should be capped at 30000ms, got ${dispatcher.initTimeoutMs}") } test("createSession after close is rejected") { - dispatcher = new TestDirectWorkerDispatcher(specWithRunner(defaultRunner)) + dispatcher = new TestDirectGrpcDispatcher(specWithRunner(defaultRunner)) dispatcher.close() val ex = intercept[IllegalStateException] { @@ -467,7 +466,7 @@ class DirectWorkerDispatcherSuite } test("socket directory is owner-only (0700) on POSIX") { - dispatcher = new TestDirectWorkerDispatcher(specWithRunner(defaultRunner)) + dispatcher = new TestDirectGrpcDispatcher(specWithRunner(defaultRunner)) // Drive one createSession so a worker (and therefore the socket dir) is // observable via the UDS connection's path. val session = createStubSession() @@ -487,7 +486,7 @@ class DirectWorkerDispatcherSuite } test("socket directory is removed after dispatcher.close") { - dispatcher = new TestDirectWorkerDispatcher(specWithRunner(defaultRunner)) + dispatcher = new TestDirectGrpcDispatcher(specWithRunner(defaultRunner)) val session = createStubSession() val socketDir = new File(udsPath(workerProcess(session))).toPath.getParent.toFile assert(socketDir.exists(), @@ -508,10 +507,8 @@ class DirectWorkerDispatcherSuite // must be terminated rather than leaked until dispatcher.close(). var capturedWorker: DirectWorkerProcess = null val failingDispatcher = - new TestDirectWorkerDispatcher(specWithRunner(defaultRunner)) { - override protected def newSession( - workerHandle: WorkerHandle, - connection: WorkerConnection): WorkerSession = + new TestDirectGrpcDispatcher(specWithRunner(defaultRunner)) { + override protected def newSession(workerHandle: WorkerHandle): WorkerSession = throw new RuntimeException("session creation failed") override protected def afterWorkerRegistered(w: DirectWorkerProcess): Unit = { capturedWorker = w @@ -538,7 +535,7 @@ class DirectWorkerDispatcherSuite .setDirect(DirectWorker.newBuilder().setRunner(defaultRunner).build()) .build() val ex = intercept[IllegalArgumentException] { - new TestDirectWorkerDispatcher(badSpec) + new TestDirectGrpcDispatcher(badSpec) } assert(ex.getMessage.contains("connection must be set"), s"expected missing-connection error, got: ${ex.getMessage}") @@ -554,7 +551,7 @@ class DirectWorkerDispatcherSuite .setRunner(defaultRunner).setProperties(tcpProperties).build()) .build() val ex = intercept[IllegalArgumentException] { - new TestDirectWorkerDispatcher(badSpec) + new TestDirectGrpcDispatcher(badSpec) } assert(ex.getMessage.contains("UNIX domain socket"), s"expected UDS-only error, got: ${ex.getMessage}") @@ -563,7 +560,7 @@ class DirectWorkerDispatcherSuite test("socket file is cleaned up when newConnection throws") { val capturedSocketPaths = new java.util.concurrent.ConcurrentLinkedQueue[String]() val failingDispatcher = - new TestDirectWorkerDispatcher(specWithRunner(defaultRunner)) { + new TestDirectGrpcDispatcher(specWithRunner(defaultRunner)) { override protected def newConnection(address: String): WorkerConnection = { capturedSocketPaths.add(address) throw new RuntimeException("connection creation failed") @@ -586,7 +583,7 @@ class DirectWorkerDispatcherSuite test("empty ProcessCallable command is rejected with a clear error") { val emptyRunner = ProcessCallable.newBuilder().build() - dispatcher = new TestDirectWorkerDispatcher(specWithRunner(emptyRunner)) + dispatcher = new TestDirectGrpcDispatcher(specWithRunner(emptyRunner)) val ex = intercept[IllegalArgumentException] { dispatcher.createSession(None) } @@ -599,7 +596,7 @@ class DirectWorkerDispatcherSuite .addCommand("bash").addCommand("-c") .addCommand("echo 'fatal: bad config' >&2; exit 42").addCommand("--") .build() - dispatcher = new TestDirectWorkerDispatcher(specWithRunner(runner)) + dispatcher = new TestDirectGrpcDispatcher(specWithRunner(runner)) val ex = intercept[RuntimeException] { dispatcher.createSession(None) @@ -628,7 +625,7 @@ class DirectWorkerDispatcherSuite .setDirect(DirectWorker.newBuilder() .setRunner(hangingRunner).setProperties(shortInitProps).build()) .build() - dispatcher = new TestDirectWorkerDispatcher(spec) + dispatcher = new TestDirectGrpcDispatcher(spec) val ex = intercept[DirectWorkerTimeoutException] { dispatcher.createSession(None) @@ -652,7 +649,7 @@ class DirectWorkerDispatcherSuite .addCommand("bash").addCommand("-c") .addCommand(s"touch ${markerFile.getAbsolutePath}").build()) .build() - dispatcher = new TestDirectWorkerDispatcher(specWithEnv(env = env)) + dispatcher = new TestDirectGrpcDispatcher(specWithEnv(env = env)) val session = dispatcher.createSession(None) session.close(emptyCancel) @@ -672,7 +669,7 @@ class DirectWorkerDispatcherSuite .addCommand("bash").addCommand("-c") .addCommand(s"touch ${markerFile.getAbsolutePath}").build()) .build() - dispatcher = new TestDirectWorkerDispatcher(specWithEnv(env = env)) + dispatcher = new TestDirectGrpcDispatcher(specWithEnv(env = env)) val session = dispatcher.createSession(None) session.close(emptyCancel) @@ -691,7 +688,7 @@ class DirectWorkerDispatcherSuite .addCommand("bash").addCommand("-c") .addCommand(s"touch ${markerFile.getAbsolutePath}").build()) .build() - dispatcher = new TestDirectWorkerDispatcher(specWithEnv(env = env)) + dispatcher = new TestDirectGrpcDispatcher(specWithEnv(env = env)) val session = dispatcher.createSession(None) session.close(emptyCancel) @@ -707,7 +704,7 @@ class DirectWorkerDispatcherSuite .addCommand("bash").addCommand("-c") .addCommand("echo 'missing dependency: libfoo' >&2; exit 7").build()) .build() - dispatcher = new TestDirectWorkerDispatcher(specWithEnv(env = env)) + dispatcher = new TestDirectGrpcDispatcher(specWithEnv(env = env)) val ex = intercept[RuntimeException] { dispatcher.createSession(None) @@ -727,7 +724,7 @@ class DirectWorkerDispatcherSuite .addCommand("sleep 30").build() val env = WorkerEnvironment.newBuilder().setInstallation(slowInstall).build() val shortTimeoutDispatcher = - new TestDirectWorkerDispatcher(specWithEnv(env = env)) { + new TestDirectGrpcDispatcher(specWithEnv(env = env)) { override protected def callableTimeoutMs: Long = 500L } try { @@ -752,7 +749,7 @@ class DirectWorkerDispatcherSuite .addCommand("bash").addCommand("-c") .addCommand(s"echo invoked >> ${counterFile.getAbsolutePath}").build()) .build() - dispatcher = new TestDirectWorkerDispatcher(specWithEnv(env = env)) + dispatcher = new TestDirectGrpcDispatcher(specWithEnv(env = env)) val s1 = dispatcher.createSession(None); s1.close(emptyCancel) val s2 = dispatcher.createSession(None); s2.close(emptyCancel) @@ -778,7 +775,7 @@ class DirectWorkerDispatcherSuite .addCommand( s"sleep 0.2; echo invoked >> ${counterFile.getAbsolutePath}").build()) .build() - dispatcher = new TestDirectWorkerDispatcher(specWithEnv(env = env)) + dispatcher = new TestDirectGrpcDispatcher(specWithEnv(env = env)) val threads = 4 val startGate = new java.util.concurrent.CountDownLatch(1) @@ -827,7 +824,7 @@ class DirectWorkerDispatcherSuite .addCommand( s"echo invoked >> ${counterFile.getAbsolutePath}; exit 1").build()) .build() - dispatcher = new TestDirectWorkerDispatcher(specWithEnv(env = env)) + dispatcher = new TestDirectGrpcDispatcher(specWithEnv(env = env)) val first = intercept[RuntimeException] { dispatcher.createSession(None) } assert(first.getMessage.contains("installation failed"), @@ -858,7 +855,7 @@ class DirectWorkerDispatcherSuite s"echo invoked >> ${counterFile.getAbsolutePath}; sleep 30").build()) .build() val timeoutDispatcher = - new TestDirectWorkerDispatcher(specWithEnv(env = env)) { + new TestDirectGrpcDispatcher(specWithEnv(env = env)) { override protected def callableTimeoutMs: Long = 500L } try { @@ -885,7 +882,7 @@ class DirectWorkerDispatcherSuite } test("non-None securityScope is rejected until pooling lands") { - dispatcher = new TestDirectWorkerDispatcher(specWithRunner(defaultRunner)) + dispatcher = new TestDirectGrpcDispatcher(specWithRunner(defaultRunner)) val scope = new WorkerSecurityScope { override def equals(obj: Any): Boolean = obj.isInstanceOf[this.type] override def hashCode(): Int = 0 @@ -903,7 +900,7 @@ class DirectWorkerDispatcherSuite .addCommand("bash").addCommand("-c").addCommand("exit 0").build()) .build() val ex = intercept[IllegalArgumentException] { - new TestDirectWorkerDispatcher(specWithEnv(env = env)) + new TestDirectGrpcDispatcher(specWithEnv(env = env)) } assert(ex.getMessage.contains("installation"), s"expected installation-required error, got: ${ex.getMessage}") @@ -918,7 +915,7 @@ class DirectWorkerDispatcherSuite .addCommand("bash").addCommand("-c") .addCommand(s"touch ${cleanupMarker.getAbsolutePath}").build()) .build() - dispatcher = new TestDirectWorkerDispatcher(specWithEnv(env = env)) + dispatcher = new TestDirectGrpcDispatcher(specWithEnv(env = env)) val session = dispatcher.createSession(None) session.close(emptyCancel) diff --git a/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/EchoProtocolSuite.scala b/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/EchoProtocolSuite.scala index eae6ef1a639a9..f26f1dc2c582b 100644 --- a/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/EchoProtocolSuite.scala +++ b/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/EchoProtocolSuite.scala @@ -16,930 +16,11 @@ */ package org.apache.spark.udf.worker.grpc -import java.util.concurrent.{CountDownLatch, LinkedBlockingQueue, TimeUnit} -import java.util.concurrent.atomic.AtomicBoolean +import org.apache.spark.udf.worker.grpc.testing.{EchoWorkerTestImplementation, + UdfWorkerTestImplementation} -import com.google.protobuf.ByteString -import io.grpc.{ManagedChannel, Server, Status} -import io.grpc.inprocess.{InProcessChannelBuilder, InProcessServerBuilder} -import io.grpc.stub.StreamObserver -import org.scalatest.BeforeAndAfterEach -// scalastyle:off funsuite -import org.scalatest.funsuite.AnyFunSuite - -import org.apache.spark.udf.worker.{ - Cancel, CancelResponse, DataRequest, DataResponse, ErrorResponse, ExecutionError, - Finish, FinishResponse, Heartbeat, HeartbeatResponse, Init, InitResponse, - PayloadChunk, ProtocolError, ShutdownRequest, ShutdownResponse, - UdfControlRequest, UdfControlResponse, UdfPayload, UdfRequest, UdfResponse, - UDFWorkerDataFormat, UdfWorkerGrpc, UserError, WorkerError, WorkerRequest, - WorkerResponse -} - -/** - * Protocol validation test for the UDF gRPC execution protocol. - * - * Implements a minimal echo worker (gRPC server) and engine client to verify - * the full Execute stream lifecycle: init, data streaming, finish, cancel, - * error handling, and the Manage RPC. The worker echoes each DataRequest - * batch back as a DataResponse; error paths are triggered by a sentinel - * payload value. - */ -class EchoProtocolSuite extends AnyFunSuite with BeforeAndAfterEach { -// scalastyle:on funsuite - - private val SUPPORTED_VERSION: Int = 1 - // A DataRequest whose payload equals this value triggers an ErrorResponse. - private val ERROR_TRIGGER: ByteString = ByteString.copyFromUtf8("ERROR") - // An init payload whose value equals this triggers an init failure - // (InitResponse with error set). - private val INIT_ERROR_TRIGGER: ByteString = ByteString.copyFromUtf8("INIT_ERROR") - - private var server: Server = _ - private var channel: ManagedChannel = _ - private var stub: UdfWorkerGrpc.UdfWorkerStub = _ - - override def beforeEach(): Unit = { - val serverName = InProcessServerBuilder.generateName() - server = InProcessServerBuilder.forName(serverName) - .directExecutor() - .addService(new EchoWorkerService) - .build() - .start() - channel = InProcessChannelBuilder.forName(serverName).directExecutor().build() - stub = UdfWorkerGrpc.newStub(channel) - } - - override def afterEach(): Unit = { - channel.shutdownNow() - server.shutdownNow() - } - - // =========================================================================== - // WORKER SIDE (gRPC server) - // =========================================================================== - - /** - * Worker state machine for one Execute stream. - * - * AwaitingInit --> AwaitingChunks? --> Data --> Draining --> Drained --> Done - * | - * +--> PostError --> Cancelling --> Cancelled --> Done - * - * `[process done]` marks an event (not a state): the asynchronous - * completion notification of in-flight work. - * - * AwaitingInit - * Init(inline) --> Data (send InitResponse) - * Init(chunked) --> AwaitingChunks - * Init(failed) --> PostError (send InitResponse with error) - * Cancel --> Cancelling --[process done]--> Cancelled --> Done (send CR) - * - * AwaitingChunks - * PayloadChunk(last=false) --> AwaitingChunks (accumulate) - * PayloadChunk(last=true) --> Data (send InitResponse) - * Cancel --> Cancelling --[process done]--> Cancelled --> Done - * Finish --> protocol error (engine must wait for - * InitResponse first) - * - * Data - * ErrorResponse sent --> PostError - * Finish --> Draining --[process done]--> Drained --> Done (send FR) - * Cancel --> Cancelling --[process done]--> Cancelled --> Done (send CR) - * - * PostError - * Finish --> PostError (no-op; engine MUST follow up with Cancel) - * Cancel --> Cancelling --[process done]--> Cancelled --> Done - * - * Draining (in-flight work running) - * [process done] --> Drained - * Cancel --> Cancelling (the pending [process done] then sees - * Cancelling and routes to CancelResponse) - * - * Drained (work done; optional post-work cleanup hook may run here; - * any error from the hook is reported via FinishResponse.error) - * --> Done (send FinishResponse) - * - * Cancelling (in-flight work being cancelled) - * [process done] --> Cancelled - * - * Cancelled (cleanup done; optional post-work cleanup hook may run here; - * any error from the hook is reported via CancelResponse.error) - * --> Done (send CancelResponse) - * - * Cross-cutting: - * - Protocol violation in any active state: send ErrorResponse(ProtocolError) - * followed by CancelResponse, transition to Done. - * - gRPC transport error (onError): transition to Done, no response sent. - */ - private sealed trait WorkerState - private case object AwaitingInit extends WorkerState - // Chunked init handshake in progress; `accumulated` holds the inline - // portion of Init.udf.payload plus all chunks received so far. - private case class AwaitingChunks(accumulated: ByteString) extends WorkerState - private case object Data extends WorkerState - private case object PostError extends WorkerState - // Finish received; in-flight finish-callback / drain work is running. - private case object Draining extends WorkerState - // Drain complete; FinishResponse not yet sent. Post-work cleanup hook - // (if any) runs in this state before the terminator is emitted. - private case object Drained extends WorkerState - // Cancel received; cancel-callback / cleanup work is running. - private case object Cancelling extends WorkerState - // Cleanup complete; CancelResponse not yet sent. Post-work cleanup hook - // (if any) runs in this state before the terminator is emitted. - private case object Cancelled extends WorkerState - private case object Done extends WorkerState - - private class EchoWorkerService extends UdfWorkerGrpc.UdfWorkerImplBase { - - override def execute( - responseObserver: StreamObserver[UdfResponse]): StreamObserver[UdfRequest] = - new ExecuteStreamHandler(responseObserver) - - override def manage( - request: WorkerRequest, - responseObserver: StreamObserver[WorkerResponse]): Unit = { - request.getManageCase match { - case WorkerRequest.ManageCase.HEARTBEAT => - responseObserver.onNext(WorkerResponse.newBuilder() - .setHeartbeat(HeartbeatResponse.getDefaultInstance) - .build()) - responseObserver.onCompleted() - - case WorkerRequest.ManageCase.SHUTDOWN => - responseObserver.onNext(WorkerResponse.newBuilder() - .setShutdown(ShutdownResponse.newBuilder().setSessionsSettled(true).build()) - .build()) - responseObserver.onCompleted() - - case _ => - responseObserver.onError( - Status.INVALID_ARGUMENT.withDescription("empty manage request") - .asRuntimeException()) - } - } - } - - private class ExecuteStreamHandler( - responseObserver: StreamObserver[UdfResponse]) extends StreamObserver[UdfRequest] { - - // State mutations go through `matchUpdateThen`: under stateLock, the - // caller-supplied function inspects the current state, returns the - // next state and a non-blocking follow-up callback; the helper writes - // the new state and releases the lock before invoking the callback, - // so I/O does not extend the critical section. - @volatile private var state: WorkerState = AwaitingInit - private val stateLock = new Object - - private def matchUpdateThen( - transition: WorkerState => (WorkerState, () => Unit)): Unit = { - val followUp = stateLock.synchronized { - val (next, callback) = transition(state) - state = next - callback - } - followUp() - } - - // gRPC does not permit concurrent calls to the response StreamObserver; - // all writes are serialized through this lock. - private val responseLock = new Object - - override def onNext(request: UdfRequest): Unit = { - request.getRequestCase match { - case UdfRequest.RequestCase.CONTROL => handleControl(request.getControl) - case UdfRequest.RequestCase.DATA => handleDataRequest(request.getData) - case _ => closeWithProtocolError("empty request oneof") - } - } - - private def handleControl(ctrl: UdfControlRequest): Unit = { - ctrl.getControlCase match { - case UdfControlRequest.ControlCase.INIT => handleInit(ctrl.getInit) - case UdfControlRequest.ControlCase.PAYLOAD => handleChunk(ctrl.getPayload) - case UdfControlRequest.ControlCase.FINISH => handleFinish() - case UdfControlRequest.ControlCase.CANCEL => handleCancel(ctrl.getCancel) - case _ => closeWithProtocolError("empty control oneof") - } - } - - private def handleInit(init: Init): Unit = matchUpdateThen { - case AwaitingInit => - if (init.hasProtocolVersion && - init.getProtocolVersion != SUPPORTED_VERSION) { - val err = ExecutionError.newBuilder() - .setProtocol(ProtocolError.newBuilder() - .setMessage(s"unsupported protocol version: ${init.getProtocolVersion}") - .build()) - .build() - (PostError, () => sendControl(UdfControlResponse.newBuilder() - .setInit(InitResponse.newBuilder().setError(err).build()) - .build())) - } else if (init.getIsChunkingPayload) { - // Payload arrives via PayloadChunk messages; defer init - // processing until the last chunk has been received. - (AwaitingChunks(init.getUdf.getPayload), () => ()) - } else { - // Payload is fully inline; process init outside the lock. - // finalizeInit performs its own CAS on entry. - val payload = init.getUdf.getPayload - (AwaitingInit, () => finalizeInit(payload)) - } - case other => - (other, () => closeWithProtocolError(s"Init received in state $other")) - } - - private def handleChunk(chunk: PayloadChunk): Unit = matchUpdateThen { - case AwaitingChunks(existing) => - val updated = existing.concat(chunk.getData) - if (chunk.hasLast && chunk.getLast) { - // Stay in AwaitingChunks until finalizeInit's CAS transitions us. - (AwaitingChunks(existing), () => finalizeInit(updated)) - } else { - (AwaitingChunks(updated), () => ()) - } - case other => - (other, () => closeWithProtocolError(s"PayloadChunk received in state $other")) - } - - // Init processing hook: invoked once with the complete assembled UDF - // payload (inline + all chunks, if any). A real worker would deserialize - // the UDF, run validation, set up runtime resources here. The echo worker - // succeeds for any payload other than INIT_ERROR_TRIGGER, which simulates - // an init-time failure (e.g. deserialization error, missing dependency). - private def finalizeInit(payload: ByteString): Unit = { - val initError: Option[ExecutionError] = if (payload == INIT_ERROR_TRIGGER) { - Some(ExecutionError.newBuilder() - .setWorker(WorkerError.newBuilder() - .setMessage("simulated init failure") - .build()) - .build()) - } else { - None - } - matchUpdateThen { - case AwaitingInit | AwaitingChunks(_) => - initError match { - case Some(err) => - (PostError, () => sendControl(UdfControlResponse.newBuilder() - .setInit(InitResponse.newBuilder().setError(err).build()) - .build())) - case None => - (Data, () => sendInitResponse()) - } - // Concurrent Cancel / transport error moved state past the init - // phase; the cancel path owns the terminator. - case other @ (Cancelling | Cancelled | Done) => - (other, () => ()) - case other => - (other, () => closeWithProtocolError(s"finalizeInit invoked in state $other")) - } - } - - private def handleDataRequest(data: DataRequest): Unit = state match { - case Data => processEcho(data) - - case _ => closeWithProtocolError(s"DataRequest received in state $state") - } - - // Echo "processing" runs inline on the gRPC callback thread for test - // simplicity. Workers that offload to a thread pool (the typical - // approach for non-trivial UDFs) must apply back-pressure via a - // bounded queue and serialize state mutations across threads. - private def processEcho(data: DataRequest): Unit = { - if (data.getData == ERROR_TRIGGER) { - // Data-phase error: emit ErrorResponse and enter PostError so the - // terminator becomes CancelResponse after the engine's Cancel. - // Only transition if we are still in Data: a concurrent Cancel - // may have moved us to Cancelling, in which case the cancel path - // owns the terminator. - val errEnvelope = UdfControlResponse.newBuilder() - .setError(ErrorResponse.newBuilder() - .setError(ExecutionError.newBuilder() - .setUser(UserError.newBuilder() - .setMessage("simulated user-code error") - .setErrorClass("SimulatedError") - .build()) - .build()) - .build()) - .build() - matchUpdateThen { - case Data => (PostError, () => sendControl(errEnvelope)) - // Concurrent Cancel / transport error already moved past data - // phase; the cancel path owns the terminator. - case other @ (Cancelling | Cancelled | Done) => (other, () => ()) - case other => - (other, () => closeWithProtocolError(s"processEcho invoked in state $other")) - } - } else { - responseLock.synchronized { - responseObserver.onNext(UdfResponse.newBuilder() - .setData(DataResponse.newBuilder().setData(data.getData).build()) - .build()) - } - } - } - - private def handleFinish(): Unit = matchUpdateThen { - case Data => - (Draining, () => onWorkComplete()) - case PostError => - // ErrorResponse already sent; this Finish was in flight before the - // engine learned about the error. The engine MUST follow up with - // Cancel; wait for it. - (PostError, () => ()) - // Finish in AwaitingInit or AwaitingChunks is a protocol error: - // the engine MUST wait for InitResponse before sending Finish. - case other => - (other, () => closeWithProtocolError(s"Finish received in state $other")) - } - - // Lazy-cancel: transition to Cancelling and let any in-flight work run - // to natural completion; the pending onWorkComplete (or this method's - // own follow-up call when no work is in flight) sees Cancelling and - // routes to CancelResponse. - private def handleCancel(cancel: Cancel): Unit = matchUpdateThen { - case AwaitingInit | AwaitingChunks(_) | Data | PostError | Draining | Drained => - (Cancelling, () => onWorkComplete()) - case other @ (Cancelling | Cancelled | Done) => - // Already cancelling or terminated; ignore duplicate Cancel. - (other, () => ()) - } - - // Called when in-flight work (finish callback, cancel cleanup, or - // batch processing) completes. The current state decides the - // terminator: - // Draining -> Drained -> send FinishResponse - // Cancelling -> Cancelled -> send CancelResponse - // - // An optional post-work cleanup hook (release file handles, flush - // metrics) belongs between the state transition and the terminator - // send. Any error from the hook is reported via FinishResponse.error - // or CancelResponse.error. - private def onWorkComplete(): Unit = matchUpdateThen { - case Draining => - (Drained, () => sendFinishResponseAndFinalize()) - case Cancelling => - (Cancelled, () => sendCancelResponseAndFinalize()) - // Stream already finalized (e.g. onError fired before this - // completion notification arrived) -- nothing to do. - case Done => (Done, () => ()) - case other => - (other, () => closeWithProtocolError(s"onWorkComplete invoked in state $other")) - } - - private def sendFinishResponseAndFinalize(): Unit = { - sendControl(UdfControlResponse.newBuilder() - .setFinish(FinishResponse.newBuilder() - .putMetrics("status", "ok") - .build()) - .build()) - matchUpdateThen { _ => - (Done, () => responseLock.synchronized { responseObserver.onCompleted() }) - } - } - - private def sendCancelResponseAndFinalize(): Unit = { - sendControl(UdfControlResponse.newBuilder() - .setCancel(CancelResponse.getDefaultInstance) - .build()) - matchUpdateThen { _ => - (Done, () => responseLock.synchronized { responseObserver.onCompleted() }) - } - } - - // gRPC transport error: the connection dropped. The stream is dead, - // so no response can be sent. The worker MUST still run the cleanup - // it would perform on an explicit Cancel (stop in-progress work, - // release resources, free buffers). The echo worker has nothing to - // release; only the state is updated. - override def onError(t: Throwable): Unit = matchUpdateThen { _ => - (Done, () => ()) - } - - override def onCompleted(): Unit = state match { - case Done => // normal: engine half-closed after session terminated - case _ => - closeWithProtocolError( - s"request stream closed by engine in unexpected state $state") - } - - private def sendInitResponse(): Unit = - sendControl(UdfControlResponse.newBuilder() - .setInit(InitResponse.getDefaultInstance) - .build()) - - private def sendControl(ctrl: UdfControlResponse): Unit = - responseLock.synchronized { - responseObserver.onNext( - UdfResponse.newBuilder().setControl(ctrl).build()) - } - - // Emit ErrorResponse(ProtocolError) followed immediately by - // CancelResponse. No in-flight work to drain, so the Cancelling / - // Cancelled intermediate states are bypassed. - private def closeWithProtocolError(msg: String): Unit = { - sendControl(UdfControlResponse.newBuilder() - .setError(ErrorResponse.newBuilder() - .setError(ExecutionError.newBuilder() - .setProtocol(ProtocolError.newBuilder().setMessage(msg).build()) - .build()) - .build()) - .build()) - sendCancelResponseAndFinalize() - } - } - - // =========================================================================== - // ENGINE SIDE (gRPC client) - // =========================================================================== - - /** - * Minimal engine client that drives the Execute stream and collects results. - * - * The request stream is half-closed (onCompleted) only after the session - * outcome is known from the server: on receiving FinishResponse, - * CancelResponse, or a gRPC error. This keeps the stream open long enough - * for Cancel to follow Finish when needed. - */ - private class EngineClient(stub: UdfWorkerGrpc.UdfWorkerStub) { - private val results = new LinkedBlockingQueue[Array[Byte]]() - private val done = new CountDownLatch(1) - @volatile var executionError: Option[ExecutionError] = None - @volatile var streamError: Option[Throwable] = None - private val requestCompleted = new AtomicBoolean(false) - // Counted down on InitResponse (success or failure) or on terminal error. - // The engine MUST wait for this before sending any DataRequest or Finish. - private val initResponseLatch = new CountDownLatch(1) - - private val responseObserver = new StreamObserver[UdfResponse] { - override def onNext(response: UdfResponse): Unit = { - response.getResponseCase match { - case UdfResponse.ResponseCase.DATA => - results.add(response.getData.getData.toByteArray) - - case UdfResponse.ResponseCase.CONTROL => - val ctrl = response.getControl - ctrl.getControlCase match { - case UdfControlResponse.ControlCase.INIT => - // InitResponse received. If error is set, init failed. - val resp = ctrl.getInit - if (resp.hasError) { - executionError = Some(resp.getError) - if (!requestCompleted.get()) sendCancel("aborting after init error") - } - initResponseLatch.countDown() - // Data phase begins only on success (no error). - - case UdfControlResponse.ControlCase.ERROR => - // Data-phase error. Send Cancel so the worker can abort cleanly; - // the error is surfaced after CancelResponse arrives. - executionError = Some(ctrl.getError.getError) - if (!requestCompleted.get()) { - sendCancel("aborting after ErrorResponse") - } - - case UdfControlResponse.ControlCase.FINISH => - completeRequestStream() - done.countDown() - - case UdfControlResponse.ControlCase.CANCEL => - completeRequestStream() - done.countDown() - - case unexpected => - throw new IllegalStateException( - s"unexpected control response: $unexpected") - } - - case unexpected => - throw new IllegalStateException( - s"unexpected response type: $unexpected") - } - } - - override def onError(t: Throwable): Unit = { - streamError = Some(t) - completeRequestStream() - initResponseLatch.countDown() - done.countDown() - } - - override def onCompleted(): Unit = { - initResponseLatch.countDown() - done.countDown() - } - } - - private val requestObserver: StreamObserver[UdfRequest] = stub.execute(responseObserver) - - def sendInit( - payloadBytes: Array[Byte], - sendChunked: Boolean = false, - protocolVersion: Int = SUPPORTED_VERSION): Unit = { - if (sendChunked) { - requestObserver.onNext(UdfRequest.newBuilder() - .setControl(UdfControlRequest.newBuilder() - .setInit(Init.newBuilder() - .setProtocolVersion(protocolVersion) - .setIsChunkingPayload(true) - .setDataFormat(UDFWorkerDataFormat.ARROW) - .setUdf(UdfPayload.newBuilder() - .setPayload(ByteString.EMPTY) - .setFormat("echo") - .build()) - .build()) - .build()) - .build()) - requestObserver.onNext(UdfRequest.newBuilder() - .setControl(UdfControlRequest.newBuilder() - .setPayload(PayloadChunk.newBuilder() - .setData(ByteString.copyFrom(payloadBytes)) - .setLast(true) - .build()) - .build()) - .build()) - } else { - requestObserver.onNext(UdfRequest.newBuilder() - .setControl(UdfControlRequest.newBuilder() - .setInit(Init.newBuilder() - .setProtocolVersion(protocolVersion) - .setDataFormat(UDFWorkerDataFormat.ARROW) - .setUdf(UdfPayload.newBuilder() - .setPayload(ByteString.copyFrom(payloadBytes)) - .setFormat("echo") - .build()) - .build()) - .build()) - .build()) - } - } - - def sendData(data: Array[Byte]): Unit = { - awaitInitResponse() - requestObserver.onNext(UdfRequest.newBuilder() - .setData(DataRequest.newBuilder() - .setData(ByteString.copyFrom(data)) - .build()) - .build()) - } - - // Sends Init with is_chunking_payload=true but no chunks. Tests then - // drive the chunks themselves via sendPayloadChunk. - def sendInitChunked(protocolVersion: Int = SUPPORTED_VERSION): Unit = { - requestObserver.onNext(UdfRequest.newBuilder() - .setControl(UdfControlRequest.newBuilder() - .setInit(Init.newBuilder() - .setProtocolVersion(protocolVersion) - .setIsChunkingPayload(true) - .setDataFormat(UDFWorkerDataFormat.ARROW) - .setUdf(UdfPayload.newBuilder() - .setPayload(ByteString.EMPTY) - .setFormat("echo") - .build()) - .build()) - .build()) - .build()) - } - - // Sends a single PayloadChunk. Does not wait for InitResponse -- - // chunks are part of the init handshake itself. - def sendPayloadChunk(data: Array[Byte], last: Boolean): Unit = { - requestObserver.onNext(UdfRequest.newBuilder() - .setControl(UdfControlRequest.newBuilder() - .setPayload(PayloadChunk.newBuilder() - .setData(ByteString.copyFrom(data)) - .setLast(last) - .build()) - .build()) - .build()) - } - - def sendFinish(): Unit = { - awaitInitResponse() - if (requestCompleted.get()) return - requestObserver.onNext(UdfRequest.newBuilder() - .setControl(UdfControlRequest.newBuilder() - .setFinish(Finish.getDefaultInstance) - .build()) - .build()) - // Request stream stays open: Cancel may still follow Finish. - // completeRequestStream() is called by the response observer. - } - - // The engine MUST wait for InitResponse before sending any DataRequest - // or Finish. Under directExecutor this returns immediately because - // sendInit's InitResponse callback runs synchronously. - private def awaitInitResponse(): Unit = { - if (!initResponseLatch.await(5, TimeUnit.SECONDS)) { - throw new IllegalStateException("InitResponse not received within timeout") - } - } - - def sendCancel(reason: String = ""): Unit = { - // If a terminator already arrived (FinishResponse / CancelResponse), - // the request stream has been half-closed and Cancel arrives too - // late -- silently ignore, matching the proto's Cancel-after-Finish - // contract. - if (requestCompleted.get()) return - requestObserver.onNext(UdfRequest.newBuilder() - .setControl(UdfControlRequest.newBuilder() - .setCancel(Cancel.newBuilder().setReason(reason).build()) - .build()) - .build()) - // Request stream stays open until the response terminator arrives; - // completeRequestStream() is called by the response observer. - } - - def completeRequestStream(): Unit = { - if (requestCompleted.compareAndSet(false, true)) { - requestObserver.onCompleted() - } - } - - def awaitDone(timeoutMs: Long = 5000): Boolean = - done.await(timeoutMs, TimeUnit.MILLISECONDS) - - def drainResults(): Seq[Array[Byte]] = { - val buf = new java.util.ArrayList[Array[Byte]]() - results.drainTo(buf) - import scala.jdk.CollectionConverters._ - buf.asScala.toSeq - } - } - - // =========================================================================== - // TESTS - // =========================================================================== - - test("echo: single DataRequest round-trip") { - val client = new EngineClient(stub) - client.sendInit("dummy-payload".getBytes) - client.sendData("hello".getBytes) - client.sendFinish() - - assert(client.awaitDone(), "stream did not complete in time") - assert(client.streamError.isEmpty, s"unexpected stream error: ${client.streamError}") - assert(client.executionError.isEmpty, s"unexpected execution error: ${client.executionError}") - val results = client.drainResults() - assert(results.length == 1) - assert(new String(results.head) == "hello") - } - - test("echo: multiple DataRequest batches are all echoed") { - val client = new EngineClient(stub) - client.sendInit("dummy-payload".getBytes) - Seq("batch1", "batch2", "batch3").foreach(b => client.sendData(b.getBytes)) - client.sendFinish() - - assert(client.awaitDone()) - assert(client.streamError.isEmpty) - val results = client.drainResults().map(new String(_)) - assert(results == Seq("batch1", "batch2", "batch3")) - } - - // The engine drives the request side from a producer thread while the - // response observer fires on a gRPC-managed callback thread. gRPC's - // bidirectional streaming and HTTP/2 flow control manage the interleaving; - // no explicit coordination is needed beyond the protocol ordering invariants. - test("echo: concurrent sending and receiving (producer/consumer pattern)") { - val asyncStub = UdfWorkerGrpc.newStub(channel) - - val receivedCount = new java.util.concurrent.atomic.AtomicInteger(0) - val doneLatch = new CountDownLatch(1) - @volatile var streamErr: Option[Throwable] = None - val requestCompleted = new AtomicBoolean(false) - // reqObs is assigned after responseObs is created. AtomicReference - // gives the response observer (which may run on a gRPC callback - // thread) a safe view of the assignment made by the test thread. - val reqObsRef = new java.util.concurrent.atomic.AtomicReference[StreamObserver[UdfRequest]]() - - val responseObs = new StreamObserver[UdfResponse] { - private def completeRequestStream(): Unit = - if (requestCompleted.compareAndSet(false, true)) reqObsRef.get().onCompleted() - - override def onNext(r: UdfResponse): Unit = r.getResponseCase match { - case UdfResponse.ResponseCase.DATA => receivedCount.incrementAndGet() - case UdfResponse.ResponseCase.CONTROL => - val c = r.getControl - c.getControlCase match { - case UdfControlResponse.ControlCase.INIT => // InitResponse: data phase can proceed - case UdfControlResponse.ControlCase.FINISH => - completeRequestStream() - doneLatch.countDown() - case unexpected => - throw new IllegalStateException( - s"unexpected control response: $unexpected") - } - case unexpected => - throw new IllegalStateException( - s"unexpected response type: $unexpected") - } - override def onError(t: Throwable): Unit = { - streamErr = Some(t) - completeRequestStream() - doneLatch.countDown() - } - override def onCompleted(): Unit = doneLatch.countDown() - } - reqObsRef.set(asyncStub.execute(responseObs)) - - val producer = new Thread(() => { - val reqObs = reqObsRef.get() - reqObs.onNext(UdfRequest.newBuilder() - .setControl(UdfControlRequest.newBuilder() - .setInit(Init.newBuilder() - .setProtocolVersion(SUPPORTED_VERSION) - .setDataFormat(UDFWorkerDataFormat.ARROW) - .setUdf(UdfPayload.newBuilder() - .setPayload(ByteString.copyFromUtf8("payload")) - .setFormat("echo").build()) - .build()) - .build()) - .build()) - (1 to 5).foreach { i => - reqObs.onNext(UdfRequest.newBuilder() - .setData(DataRequest.newBuilder() - .setData(ByteString.copyFromUtf8(s"batch-$i")).build()) - .build()) - } - reqObs.onNext(UdfRequest.newBuilder() - .setControl(UdfControlRequest.newBuilder() - .setFinish(Finish.getDefaultInstance).build()) - .build()) - // Request stream stays open; completeRequestStream() is called by - // the response observer on FinishResponse or gRPC error. - }, "producer") - producer.start() - - assert(doneLatch.await(10, TimeUnit.SECONDS), "stream did not complete") - assert(streamErr.isEmpty, s"unexpected error: $streamErr") - assert(receivedCount.get() == 5, s"expected 5 echoes, got ${receivedCount.get()}") - } - - test("echo: chunked payload delivery") { - val client = new EngineClient(stub) - client.sendInit("chunked-payload".getBytes, sendChunked = true) - client.sendData("data".getBytes) - client.sendFinish() - - assert(client.awaitDone()) - assert(client.streamError.isEmpty) - assert(new String(client.drainResults().head) == "data") - } - - test("echo: generator-style UDF (zero DataRequests, engine sends Finish after Init)") { - val client = new EngineClient(stub) - client.sendInit("generator-payload".getBytes) - client.sendFinish() - - assert(client.awaitDone()) - assert(client.streamError.isEmpty) - assert(client.drainResults().isEmpty) - } - - test("cancel: engine cancels mid-stream before sending Finish") { - val client = new EngineClient(stub) - client.sendInit("dummy-payload".getBytes) - client.sendData("batch1".getBytes) - client.sendCancel("task interrupted") - - assert(client.awaitDone()) - assert(client.streamError.isEmpty) - } - - // Cancel MAY follow Finish. The worker sends CancelResponse if Cancel arrives - // before FinishResponse is sent, or FinishResponse if it arrived too late. - // The engine must accept either outcome. - test("cancel: engine sends Cancel after Finish -- accepts FinishResponse or CancelResponse") { - val client = new EngineClient(stub) - client.sendInit("dummy-payload".getBytes) - client.sendData("data".getBytes) - client.sendFinish() - client.sendCancel("task interrupted after finish") - - assert(client.awaitDone(), "stream did not complete") - assert(client.streamError.isEmpty, - s"Cancel-after-Finish must not cause a gRPC error: ${client.streamError}") - } - - test("ErrorResponse: worker signals UserError, engine sends Cancel and receives CancelResponse") { - val client = new EngineClient(stub) - client.sendInit("dummy-payload".getBytes) - client.sendData(ERROR_TRIGGER.toByteArray) - - assert(client.awaitDone()) - assert(client.streamError.isEmpty, s"expected no gRPC error, got ${client.streamError}") - assert(client.executionError.isDefined, "expected an ExecutionError") - assert(client.executionError.get.hasUser, "expected UserError kind") - assert(client.executionError.get.getUser.getErrorClass == "SimulatedError") - } - - test("protocol error: second Init is rejected with ProtocolError + CancelResponse") { - val client = new EngineClient(stub) - client.sendInit("payload".getBytes) - client.sendInit("second-init".getBytes) - - assert(client.awaitDone()) - assert(client.streamError.isEmpty, "expected ProtocolError, not a gRPC stream error") - assert(client.executionError.isDefined, "expected an ExecutionError") - assert(client.executionError.get.hasProtocol, "expected ProtocolError kind") - } - - test("init error: inline payload triggers init failure") { - val client = new EngineClient(stub) - client.sendInit(INIT_ERROR_TRIGGER.toByteArray) - - assert(client.awaitDone()) - assert(client.streamError.isEmpty, s"expected no gRPC error, got ${client.streamError}") - assert(client.executionError.isDefined, "expected an init error") - assert(client.executionError.get.hasWorker, "expected WorkerError kind") - assert(client.executionError.get.getWorker.getMessage == "simulated init failure") - } - - test("init error: chunked payload assembled across chunks triggers init failure") { - val client = new EngineClient(stub) - client.sendInit(INIT_ERROR_TRIGGER.toByteArray, sendChunked = true) - - assert(client.awaitDone()) - assert(client.streamError.isEmpty, s"expected no gRPC error, got ${client.streamError}") - assert(client.executionError.isDefined, "expected an init error") - assert(client.executionError.get.hasWorker, "expected WorkerError kind") - } - - test("init error: unsupported protocol version triggers init failure") { - val client = new EngineClient(stub) - client.sendInit("payload".getBytes, protocolVersion = SUPPORTED_VERSION + 999) - - assert(client.awaitDone()) - assert(client.streamError.isEmpty, s"expected no gRPC error, got ${client.streamError}") - assert(client.executionError.isDefined, "expected an init error") - assert(client.executionError.get.hasProtocol, "expected ProtocolError kind") - } - - test("cancel: Cancel before Init is accepted (AwaitingInit state)") { - val client = new EngineClient(stub) - client.sendCancel("aborting before init") - - assert(client.awaitDone()) - assert(client.streamError.isEmpty, s"expected no gRPC error, got ${client.streamError}") - assert(client.executionError.isEmpty, "Cancel before Init is normal abort, not an error") - } - - test("cancel: Cancel during chunked payload delivery (AwaitingChunks state)") { - val client = new EngineClient(stub) - client.sendInitChunked() - client.sendPayloadChunk("partial".getBytes, last = false) - // No final chunk: worker is still accumulating when Cancel arrives. - client.sendCancel("aborting mid-chunking") - - assert(client.awaitDone()) - assert(client.streamError.isEmpty) - assert(client.executionError.isEmpty) - } - - test("echo: chunked payload assembled from multiple non-final chunks") { - val client = new EngineClient(stub) - client.sendInitChunked() - client.sendPayloadChunk("part1".getBytes, last = false) - client.sendPayloadChunk("part2".getBytes, last = false) - client.sendPayloadChunk("final".getBytes, last = true) - // The accumulator state machine should have produced InitResponse after - // the last=true chunk; the data-phase round-trip below verifies the - // worker correctly transitioned to Data. - client.sendData("after-chunks".getBytes) - client.sendFinish() - - assert(client.awaitDone()) - assert(client.streamError.isEmpty) - assert(new String(client.drainResults().head) == "after-chunks") - } - - test("protocol error: PayloadChunk in Data state (no chunking flag on Init)") { - val client = new EngineClient(stub) - client.sendInit("payload".getBytes) // non-chunked init transitions to Data - client.sendPayloadChunk("oops".getBytes, last = true) - - assert(client.awaitDone()) - assert(client.streamError.isEmpty) - assert(client.executionError.isDefined, "expected a protocol error") - assert(client.executionError.get.hasProtocol) - } - - test("Manage: heartbeat is acknowledged") { - val blockingStub = UdfWorkerGrpc.newBlockingStub(channel) - val resp = blockingStub.manage(WorkerRequest.newBuilder() - .setHeartbeat(Heartbeat.getDefaultInstance) - .build()) - assert(resp.hasHeartbeat, "expected HeartbeatResponse") - } - - test("Manage: ShutdownRequest is acknowledged") { - val blockingStub = UdfWorkerGrpc.newBlockingStub(channel) - val resp = blockingStub.manage(WorkerRequest.newBuilder() - .setShutdown(ShutdownRequest.newBuilder().setReason("test done").build()) - .build()) - assert(resp.hasShutdown, "expected ShutdownResponse") - } +/** Runs the UDF worker protocol contract against the minimal echo implementation. */ +class EchoProtocolSuite extends UdfWorkerProtocolSuite { + override protected def implementation: UdfWorkerTestImplementation = + EchoWorkerTestImplementation } diff --git a/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/GrpcWorkerSessionConcurrencySuite.scala b/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/GrpcWorkerSessionConcurrencySuite.scala index fb0eeae9ed376..215900b7262ad 100644 --- a/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/GrpcWorkerSessionConcurrencySuite.scala +++ b/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/GrpcWorkerSessionConcurrencySuite.scala @@ -35,7 +35,7 @@ import org.apache.spark.udf.worker.{Cancel, CancelResponse, DataRequest, DataRes ErrorResponse, ExecutionError, Finish, FinishResponse, Init, InitResponse, UdfControlResponse, UdfPayload, UdfRequest, UdfResponse, UDFWorkerDataFormat, UdfWorkerGrpc, UserError, WorkerRequest, WorkerResponse} -import org.apache.spark.udf.worker.core.{Termination, WorkerHandle, WorkerLogger} +import org.apache.spark.udf.worker.core.{Termination, WorkerConnection, WorkerHandle, WorkerLogger} /** * Concurrency tests for [[GrpcWorkerSession]] that pin the wire-ordering and @@ -115,12 +115,18 @@ class GrpcWorkerSessionConcurrencySuite // Helpers // --------------------------------------------------------------------------- + private val stubConnection = new WorkerConnection { + override def isActive: Boolean = true + override def close(): Unit = () + } + private class TestWorkerHandle extends WorkerHandle { val invalidated = new AtomicBoolean(false) val released = new AtomicBoolean(false) override def id: String = "test-worker" override def markInvalid(): Unit = invalidated.set(true) override def releaseSession(): Unit = released.set(true) + override def connection: WorkerConnection = stubConnection } /** diff --git a/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/UdfWorkerProtocolSuite.scala b/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/UdfWorkerProtocolSuite.scala new file mode 100644 index 0000000000000..ad7853aaa135d --- /dev/null +++ b/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/UdfWorkerProtocolSuite.scala @@ -0,0 +1,300 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.udf.worker.grpc + +import java.util.concurrent.TimeUnit + +import com.google.protobuf.ByteString +import io.grpc.{ManagedChannel, Server} +import io.grpc.inprocess.{InProcessChannelBuilder, InProcessServerBuilder} +import org.scalatest.BeforeAndAfterEach +// scalastyle:off funsuite +import org.scalatest.funsuite.AnyFunSuite + +import org.apache.spark.udf.worker.{Cancel, DataResponse, Finish, Heartbeat, Init, PayloadChunk, + ShutdownRequest, UdfWorkerGrpc, WorkerRequest} +import org.apache.spark.udf.worker.grpc.testing.{SuccessfulExecution, + UdfWorkerProtocolClient, UdfWorkerTestImplementation} + +/** + * Implementation-neutral contract for servers that implement the UDF worker protocol. + * + * Concrete suites provide deterministic requests and expected responses through + * [[implementation]]. Protocol sequencing, chunking, cancellation, and management RPC + * behavior stay here so another server implementation can run the same contract unchanged. + */ +private[grpc] abstract class UdfWorkerProtocolSuite + extends AnyFunSuite with BeforeAndAfterEach { +// scalastyle:on funsuite + + protected def implementation: UdfWorkerTestImplementation + + private var server: Server = _ + private var channel: ManagedChannel = _ + private var stub: UdfWorkerGrpc.UdfWorkerStub = _ + + override def beforeEach(): Unit = { + super.beforeEach() + val serverName = InProcessServerBuilder.generateName() + server = InProcessServerBuilder.forName(serverName) + .addService(implementation.newService()) + .build() + .start() + channel = InProcessChannelBuilder.forName(serverName).build() + stub = UdfWorkerGrpc.newStub(channel) + } + + override def afterEach(): Unit = { + if (channel != null) { + channel.shutdownNow().awaitTermination(5, TimeUnit.SECONDS) + channel = null + } + if (server != null) { + server.shutdownNow().awaitTermination(5, TimeUnit.SECONDS) + server = null + } + super.afterEach() + } + + test("single request completes with the expected response") { + runSuccessful(implementation.singleRequest) + } + + test("multiple requests preserve the implementation's response ordering") { + runSuccessful(implementation.multipleRequests) + } + + test("requests and responses can progress on separate threads") { + val scenario = implementation.multipleRequests + val client = new UdfWorkerProtocolClient(stub) + client.sendInit(scenario.init) + + val producer = new Thread(() => { + scenario.requests.foreach(client.sendData) + client.sendFinish() + }, "udf-protocol-producer") + producer.start() + + assert(client.awaitDone(10000L), "stream did not complete") + producer.join(10000L) + assert(!producer.isAlive, "producer did not finish") + assertSuccessfulResult(client, scenario.expectedResponses) + } + + test("a UDF payload can be delivered in one chunk") { + val scenario = implementation.singleRequest + val client = new UdfWorkerProtocolClient(stub) + val (init, chunks) = chunked(scenario.init, 1) + client.sendInit(init) + chunks.foreach(client.sendPayloadChunk) + scenario.requests.foreach(client.sendData) + client.sendFinish() + + assert(client.awaitDone()) + assertSuccessfulResult(client, scenario.expectedResponses) + } + + test("finish immediately after initialization is supported") { + runSuccessful(implementation.noInput) + } + + test("cancel before finish terminates the stream cleanly") { + val scenario = implementation.singleRequest + val client = new UdfWorkerProtocolClient(stub) + client.sendInit(scenario.init) + scenario.requests.headOption.foreach(client.sendData) + client.sendCancel(Cancel.newBuilder().setReason("task interrupted").build()) + + assert(client.awaitDone()) + assert(client.streamError.isEmpty, s"unexpected stream error: ${client.streamError}") + } + + test("cancel after finish accepts either terminal response") { + val scenario = implementation.singleRequest + val client = new UdfWorkerProtocolClient(stub) + client.sendInit(scenario.init) + scenario.requests.foreach(client.sendData) + client.sendFinish() + client.sendCancel(Cancel.newBuilder().setReason("interrupted after finish").build()) + + assert(client.awaitDone()) + assert(client.streamError.isEmpty, s"unexpected stream error: ${client.streamError}") + } + + test("a user-code error is returned as an execution error") { + val scenario = implementation.userError + val client = new UdfWorkerProtocolClient(stub) + client.sendInit(scenario.init) + client.sendData(scenario.request) + + assert(client.awaitDone()) + assert(client.streamError.isEmpty, s"unexpected stream error: ${client.streamError}") + assert(client.executionError.contains(scenario.expectedError), + s"expected ${scenario.expectedError}, got ${client.executionError}") + } + + test("a second Init is rejected with a protocol error") { + val scenario = implementation.singleRequest + val client = new UdfWorkerProtocolClient(stub) + client.sendInit(scenario.init) + assert(client.awaitInitResponse(), "first InitResponse not received") + client.sendInit(scenario.init) + + assert(client.awaitDone()) + assert(client.streamError.isEmpty) + assert(client.executionError.exists(_.hasProtocol), + s"expected ProtocolError, got ${client.executionError}") + } + + test("an inline initialization failure is returned in InitResponse") { + val scenario = implementation.initError + val client = new UdfWorkerProtocolClient(stub) + client.sendInit(scenario.init) + + assert(client.awaitDone()) + assert(client.streamError.isEmpty, s"unexpected stream error: ${client.streamError}") + assert(client.executionError.contains(scenario.expectedError), + s"expected ${scenario.expectedError}, got ${client.executionError}") + } + + test("a chunked initialization failure is returned in InitResponse") { + val scenario = implementation.initError + val client = new UdfWorkerProtocolClient(stub) + val (init, chunks) = chunked(scenario.init, 3) + client.sendInit(init) + chunks.foreach(client.sendPayloadChunk) + + assert(client.awaitDone()) + assert(client.streamError.isEmpty, s"unexpected stream error: ${client.streamError}") + assert(client.executionError.contains(scenario.expectedError), + s"expected ${scenario.expectedError}, got ${client.executionError}") + } + + test("an unsupported protocol version fails initialization") { + val valid = implementation.singleRequest.init + val invalid = Init.newBuilder(valid) + .setProtocolVersion(implementation.supportedProtocolVersion + 999) + .build() + val client = new UdfWorkerProtocolClient(stub) + client.sendInit(invalid) + + assert(client.awaitDone()) + assert(client.streamError.isEmpty) + assert(client.executionError.exists(_.hasProtocol), + s"expected ProtocolError, got ${client.executionError}") + } + + test("cancel before Init is accepted") { + val client = new UdfWorkerProtocolClient(stub) + client.sendCancel(Cancel.newBuilder().setReason("aborting before init").build()) + + assert(client.awaitDone()) + assert(client.streamError.isEmpty) + assert(client.executionError.isEmpty) + } + + test("cancel during chunked initialization is accepted") { + val client = new UdfWorkerProtocolClient(stub) + val (init, chunks) = chunked(implementation.singleRequest.init, 3) + client.sendInit(init) + client.sendPayloadChunk(chunks.head.toBuilder.setLast(false).build()) + client.sendCancel(Cancel.newBuilder().setReason("aborting during chunking").build()) + + assert(client.awaitDone()) + assert(client.streamError.isEmpty) + assert(client.executionError.isEmpty) + } + + test("a UDF payload can be assembled from multiple chunks") { + val scenario = implementation.singleRequest + val client = new UdfWorkerProtocolClient(stub) + val (init, chunks) = chunked(scenario.init, 3) + client.sendInit(init) + chunks.foreach(client.sendPayloadChunk) + scenario.requests.foreach(client.sendData) + client.sendFinish() + + assert(client.awaitDone()) + assertSuccessfulResult(client, scenario.expectedResponses) + } + + test("PayloadChunk in the data phase is rejected with a protocol error") { + val client = new UdfWorkerProtocolClient(stub) + client.sendInit(implementation.singleRequest.init) + assert(client.awaitInitResponse(), "InitResponse not received") + client.sendPayloadChunk(PayloadChunk.newBuilder() + .setData(ByteString.copyFromUtf8("unexpected")) + .setLast(true) + .build()) + + assert(client.awaitDone()) + assert(client.streamError.isEmpty) + assert(client.executionError.exists(_.hasProtocol), + s"expected ProtocolError, got ${client.executionError}") + } + + test("Manage acknowledges heartbeat") { + val response = UdfWorkerGrpc.newBlockingStub(channel).manage(WorkerRequest.newBuilder() + .setHeartbeat(Heartbeat.getDefaultInstance) + .build()) + assert(response.hasHeartbeat, "expected HeartbeatResponse") + } + + test("Manage acknowledges shutdown") { + val response = UdfWorkerGrpc.newBlockingStub(channel).manage(WorkerRequest.newBuilder() + .setShutdown(ShutdownRequest.newBuilder().setReason("test complete").build()) + .build()) + assert(response.hasShutdown, "expected ShutdownResponse") + } + + private def runSuccessful(scenario: SuccessfulExecution): Unit = { + val client = new UdfWorkerProtocolClient(stub) + client.sendInit(scenario.init) + scenario.requests.foreach(client.sendData) + client.sendFinish(Finish.getDefaultInstance) + + assert(client.awaitDone(), "stream did not complete") + assertSuccessfulResult(client, scenario.expectedResponses) + } + + private def assertSuccessfulResult( + client: UdfWorkerProtocolClient, + expected: Seq[DataResponse]): Unit = { + assert(client.streamError.isEmpty, s"unexpected stream error: ${client.streamError}") + assert(client.executionError.isEmpty, + s"unexpected execution error: ${client.executionError}") + assert(client.drainResults() == expected) + } + + private def chunked(init: Init, parts: Int): (Init, Seq[PayloadChunk]) = { + require(parts > 0, "parts must be positive") + val payload = init.getUdf.getPayload + val chunkedInit = Init.newBuilder(init) + .setIsChunkingPayload(true) + .setUdf(init.getUdf.toBuilder.setPayload(ByteString.EMPTY).build()) + .build() + val chunks = (0 until parts).map { index => + val start = payload.size() * index / parts + val end = payload.size() * (index + 1) / parts + PayloadChunk.newBuilder() + .setData(payload.substring(start, end)) + .setLast(index == parts - 1) + .build() + } + (chunkedInit, chunks) + } +} diff --git a/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/testing/EchoGrpcWorkerMain.scala b/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/testing/EchoGrpcWorkerMain.scala new file mode 100644 index 0000000000000..c22f0e8137e72 --- /dev/null +++ b/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/testing/EchoGrpcWorkerMain.scala @@ -0,0 +1,152 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.udf.worker.grpc.testing + +import java.io.IOException +import java.nio.file.{Files, Paths} +import java.util.concurrent.TimeUnit + +import io.grpc.Server +import io.grpc.netty.NettyServerBuilder +import io.netty.channel.EventLoopGroup +import io.netty.channel.unix.DomainSocketAddress + +import org.apache.spark.udf.worker.grpc.{GrpcWorkerChannel, UnixDomainSocketTransport} + +/** + * Test-only worker binary. Spawned by integration tests via the standard + * worker spec callable contract: + * + * java -cp org.apache.spark.udf.worker.grpc.testing.EchoGrpcWorkerMain + * --id --connection + * + * Hosts an [[EchoWorkerService]] on the supplied Unix domain socket and + * blocks until the JVM receives `SIGTERM` (then shuts down gracefully). + * + * The engine injects `--id` and `--connection` automatically, matching the + * contract documented in `worker_spec.proto`. + */ +object EchoGrpcWorkerMain { + + // scalastyle:off println + private def stderrln(s: String): Unit = System.err.println(s"[echo-worker] $s") + + def main(args: Array[String]): Unit = { + val parsed = parseArgs(args) + val socketPath = parsed.getOrElse("connection", + throw new IllegalArgumentException("--connection is required")) + val workerId = parsed.getOrElse("id", "") + + val transport = UnixDomainSocketTransport.detect() + val bossElg: EventLoopGroup = transport.newEventLoopGroup() + val workerElg: EventLoopGroup = transport.newEventLoopGroup() + + // Ensure the parent directory exists; the engine should already have + // created it, but be defensive in case the test points at an unusual + // path. + val sockFile = Paths.get(socketPath) + val parent = sockFile.getParent + if (parent != null && !Files.exists(parent)) { + Files.createDirectories(parent) + } + // The OS rejects bind() if the path already exists. Best-effort remove + // (e.g., a stale file from a previous run). + try { Files.deleteIfExists(sockFile) } catch { case _: IOException => () } + + val server: Server = NettyServerBuilder + .forAddress(new DomainSocketAddress(socketPath)) + .channelType(transport.serverChannelType) + .bossEventLoopGroup(bossElg) + .workerEventLoopGroup(workerElg) + .maxInboundMessageSize(GrpcWorkerChannel.DEFAULT_MAX_INBOUND_MESSAGE_SIZE) + .addService(new EchoWorkerService) + .build() + + val shutdownLatch = new java.util.concurrent.CountDownLatch(1) + + // SIGTERM is the engine's graceful-stop signal per the worker spec. + // SIGINT is added so manual testing (Ctrl-C) also shuts down cleanly. + val shutdownHook = new Thread(() => { + stderrln(s"shutdown signal received (worker=$workerId)") + try { + server.shutdown() + if (!server.awaitTermination(5, TimeUnit.SECONDS)) { + server.shutdownNow() + } + } catch { + case _: InterruptedException => Thread.currentThread().interrupt() + } finally { + UnixDomainSocketTransport.shutdown(bossElg, 5000L) + UnixDomainSocketTransport.shutdown(workerElg, 5000L) + // The dispatcher will also delete the socket file, but we delete it + // here so a worker that ran without an engine cleans up after itself. + try { Files.deleteIfExists(sockFile) } catch { case _: IOException => () } + shutdownLatch.countDown() + } + }, s"echo-worker-shutdown-$workerId") + // This is a standalone test-worker main (no Spark on its classpath), so it + // cannot use Spark's ShutdownHookManager and registers the JVM hook directly. + // scalastyle:off runtimeaddshutdownhook + Runtime.getRuntime.addShutdownHook(shutdownHook) + // scalastyle:on runtimeaddshutdownhook + + try { + server.start() + stderrln(s"listening on $socketPath (transport=${transport.name})") + } catch { + case e: Throwable => + stderrln(s"failed to start gRPC server: ${e.getMessage}") + // Clean up event loops we created; the JVM is about to exit anyway, + // but skipping the cleanup leaks fds on test reruns. + UnixDomainSocketTransport.shutdown(bossElg, 5000L) + UnixDomainSocketTransport.shutdown(workerElg, 5000L) + // Remove the hook explicitly: addShutdownHook keeps the JVM alive + // until the hook returns, and we already did its job. + try { Runtime.getRuntime.removeShutdownHook(shutdownHook) } catch { + case _: IllegalStateException => () // JVM shutting down already + } + sys.exit(2) + } + + // Park until the shutdown hook fires. + server.awaitTermination() + // The hook's own latch is the post-shutdown gate -- wait briefly for it + // so the worker process exits only after fd cleanup is done. + shutdownLatch.await(10, TimeUnit.SECONDS) + } + // scalastyle:on println + + /** + * Parses `--key value` flags. The engine guarantees `--id` and + * `--connection` are present; unknown flags are ignored so future + * engine-side additions don't break older workers. + */ + private def parseArgs(args: Array[String]): Map[String, String] = { + val out = scala.collection.mutable.Map.empty[String, String] + var i = 0 + while (i < args.length) { + val a = args(i) + if (a.startsWith("--") && i + 1 < args.length) { + out(a.stripPrefix("--")) = args(i + 1) + i += 2 + } else { + i += 1 + } + } + out.toMap + } +} diff --git a/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/testing/EchoWorkerService.scala b/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/testing/EchoWorkerService.scala new file mode 100644 index 0000000000000..fafee8f0f3d99 --- /dev/null +++ b/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/testing/EchoWorkerService.scala @@ -0,0 +1,309 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.udf.worker.grpc.testing + +import com.google.protobuf.ByteString +import io.grpc.Status +import io.grpc.stub.StreamObserver + +import org.apache.spark.udf.worker.{Cancel, CancelResponse, DataRequest, DataResponse, + ErrorResponse, ExecutionError, FinishResponse, HeartbeatResponse, Init, InitResponse, + PayloadChunk, ProtocolError, ShutdownResponse, UdfControlRequest, UdfControlResponse, + UdfRequest, UdfResponse, UdfWorkerGrpc, UserError, WorkerError, WorkerRequest, + WorkerResponse} + +/** + * Minimal echo worker implementing the UDF protocol. + * + * Used by: + * - `EchoProtocolSuite` (in-process), which exercises only the wire + * protocol via an in-process gRPC channel. + * - [[EchoGrpcWorkerMain]], which hosts this service over a UDS so the + * `DirectGrpcDispatcher` can spawn it as a separate + * worker process for integration tests. + * + * Behavior: + * - Each `DataRequest` is echoed back verbatim as a `DataResponse`. + * - A `DataRequest` whose payload is `"ERROR"` triggers an `ErrorResponse` + * carrying a `UserError`. The engine MUST follow with `Cancel`. + * - An `Init` whose assembled payload is `"INIT_ERROR"` triggers an + * `InitResponse` with a `WorkerError`. The engine MUST follow with `Cancel`. + * - Protocol version other than [[SupportedVersion]] triggers an init-time + * `ProtocolError`. + * - Chunked payloads accumulate inline + chunks before init finalises. + */ +object EchoWorkerService { + /** Protocol version the echo worker advertises. */ + val SupportedVersion: Int = 1 + + /** A `DataRequest` whose payload equals this triggers an `ErrorResponse`. */ + val ErrorTrigger: ByteString = ByteString.copyFromUtf8("ERROR") + + /** An `Init` whose assembled payload equals this triggers an init failure. */ + val InitErrorTrigger: ByteString = ByteString.copyFromUtf8("INIT_ERROR") +} + +class EchoWorkerService extends UdfWorkerGrpc.UdfWorkerImplBase { + import EchoWorkerService._ + + override def execute( + responseObserver: StreamObserver[UdfResponse]): StreamObserver[UdfRequest] = + new ExecuteStreamHandler(responseObserver) + + override def manage( + request: WorkerRequest, + responseObserver: StreamObserver[WorkerResponse]): Unit = { + request.getManageCase match { + case WorkerRequest.ManageCase.HEARTBEAT => + responseObserver.onNext(WorkerResponse.newBuilder() + .setHeartbeat(HeartbeatResponse.getDefaultInstance) + .build()) + responseObserver.onCompleted() + + case WorkerRequest.ManageCase.SHUTDOWN => + responseObserver.onNext(WorkerResponse.newBuilder() + .setShutdown(ShutdownResponse.newBuilder().setSessionsSettled(true).build()) + .build()) + responseObserver.onCompleted() + + case _ => + responseObserver.onError( + Status.INVALID_ARGUMENT.withDescription("empty manage request") + .asRuntimeException()) + } + } + + // State machine matches the protocol comment in udf_message.proto: + // AwaitingInit -> AwaitingChunks? -> Data -> Draining -> Drained -> Done + // | + // +--> PostError -> Cancelling -> Cancelled -> Done + private sealed trait WorkerState + private case object AwaitingInit extends WorkerState + private case class AwaitingChunks(accumulated: ByteString) extends WorkerState + private case object Data extends WorkerState + private case object PostError extends WorkerState + private case object Draining extends WorkerState + private case object Drained extends WorkerState + private case object Cancelling extends WorkerState + private case object Cancelled extends WorkerState + private case object Done extends WorkerState + + private class ExecuteStreamHandler( + responseObserver: StreamObserver[UdfResponse]) extends StreamObserver[UdfRequest] { + + @volatile private var state: WorkerState = AwaitingInit + private val stateLock = new Object + private val responseLock = new Object + + private def matchUpdateThen( + transition: WorkerState => (WorkerState, () => Unit)): Unit = { + val followUp = stateLock.synchronized { + val (next, callback) = transition(state) + state = next + callback + } + followUp() + } + + override def onNext(request: UdfRequest): Unit = request.getRequestCase match { + case UdfRequest.RequestCase.CONTROL => handleControl(request.getControl) + case UdfRequest.RequestCase.DATA => handleDataRequest(request.getData) + case _ => closeWithProtocolError("empty request oneof") + } + + private def handleControl(ctrl: UdfControlRequest): Unit = ctrl.getControlCase match { + case UdfControlRequest.ControlCase.INIT => handleInit(ctrl.getInit) + case UdfControlRequest.ControlCase.PAYLOAD => handleChunk(ctrl.getPayload) + case UdfControlRequest.ControlCase.FINISH => handleFinish() + case UdfControlRequest.ControlCase.CANCEL => handleCancel(ctrl.getCancel) + case _ => closeWithProtocolError("empty control oneof") + } + + private def handleInit(init: Init): Unit = matchUpdateThen { + case AwaitingInit => + if (init.hasProtocolVersion && init.getProtocolVersion != SupportedVersion) { + val err = ExecutionError.newBuilder() + .setProtocol(ProtocolError.newBuilder() + .setMessage(s"unsupported protocol version: ${init.getProtocolVersion}") + .build()) + .build() + (PostError, () => sendControl(UdfControlResponse.newBuilder() + .setInit(InitResponse.newBuilder().setError(err).build()) + .build())) + } else if (init.getIsChunkingPayload) { + (AwaitingChunks(init.getUdf.getPayload), () => ()) + } else { + val payload = init.getUdf.getPayload + (AwaitingInit, () => finalizeInit(payload)) + } + case other => + (other, () => closeWithProtocolError(s"Init received in state $other")) + } + + private def handleChunk(chunk: PayloadChunk): Unit = matchUpdateThen { + case AwaitingChunks(existing) => + val updated = existing.concat(chunk.getData) + if (chunk.hasLast && chunk.getLast) { + (AwaitingChunks(existing), () => finalizeInit(updated)) + } else { + (AwaitingChunks(updated), () => ()) + } + case other => + (other, () => closeWithProtocolError(s"PayloadChunk received in state $other")) + } + + private def finalizeInit(payload: ByteString): Unit = { + val initError: Option[ExecutionError] = if (payload == InitErrorTrigger) { + Some(ExecutionError.newBuilder() + .setWorker(WorkerError.newBuilder() + .setMessage("simulated init failure") + .build()) + .build()) + } else { + None + } + matchUpdateThen { + case AwaitingInit | AwaitingChunks(_) => + initError match { + case Some(err) => + (PostError, () => sendControl(UdfControlResponse.newBuilder() + .setInit(InitResponse.newBuilder().setError(err).build()) + .build())) + case None => + (Data, () => sendInitResponse()) + } + case other @ (Cancelling | Cancelled | Done) => (other, () => ()) + case other => + (other, () => closeWithProtocolError(s"finalizeInit invoked in state $other")) + } + } + + private def handleDataRequest(data: DataRequest): Unit = state match { + case Data => processEcho(data) + case _ => closeWithProtocolError(s"DataRequest received in state $state") + } + + private def processEcho(data: DataRequest): Unit = { + if (data.getData == ErrorTrigger) { + val errEnvelope = UdfControlResponse.newBuilder() + .setError(ErrorResponse.newBuilder() + .setError(ExecutionError.newBuilder() + .setUser(UserError.newBuilder() + .setMessage("simulated user-code error") + .setErrorClass("SimulatedError") + .build()) + .build()) + .build()) + .build() + matchUpdateThen { + case Data => (PostError, () => sendControl(errEnvelope)) + case other @ (Cancelling | Cancelled | Done) => (other, () => ()) + case other => + (other, () => closeWithProtocolError(s"processEcho invoked in state $other")) + } + } else { + responseLock.synchronized { + responseObserver.onNext(UdfResponse.newBuilder() + .setData(DataResponse.newBuilder().setData(data.getData).build()) + .build()) + } + } + } + + private def handleFinish(): Unit = matchUpdateThen { + case Data => (Draining, () => onWorkComplete()) + case PostError => (PostError, () => ()) + case other => + (other, () => closeWithProtocolError(s"Finish received in state $other")) + } + + private def handleCancel(@scala.annotation.unused cancel: Cancel): Unit = matchUpdateThen { + case AwaitingInit | AwaitingChunks(_) | Data | PostError | Draining | Drained => + (Cancelling, () => onWorkComplete()) + case other @ (Cancelling | Cancelled | Done) => (other, () => ()) + } + + private def onWorkComplete(): Unit = matchUpdateThen { + case Draining => + (Drained, () => { + sendControl(UdfControlResponse.newBuilder() + .setFinish(FinishResponse.newBuilder() + .putMetrics("status", "ok") + .build()) + .build()) + finalizeDone() + }) + case Cancelling => + (Cancelled, () => { + sendControl(UdfControlResponse.newBuilder() + .setCancel(CancelResponse.getDefaultInstance) + .build()) + finalizeDone() + }) + case Done => (Done, () => ()) + case other => + (other, () => closeWithProtocolError(s"onWorkComplete invoked in state $other")) + } + + /** + * Transitions to `Done` and half-closes the response stream. Shared by the + * FinishResponse, CancelResponse, and protocol-error paths, which all end + * the conversation with `onCompleted()`. + */ + private def finalizeDone(): Unit = matchUpdateThen { _ => + (Done, () => responseLock.synchronized { responseObserver.onCompleted() }) + } + + override def onError(t: Throwable): Unit = matchUpdateThen { _ => + (Done, () => ()) + } + + override def onCompleted(): Unit = state match { + case Done => // normal: engine half-closed after session terminated + case _ => + closeWithProtocolError( + s"request stream closed by engine in unexpected state $state") + } + + private def sendInitResponse(): Unit = + sendControl(UdfControlResponse.newBuilder() + .setInit(InitResponse.getDefaultInstance) + .build()) + + private def sendControl(ctrl: UdfControlResponse): Unit = + responseLock.synchronized { + responseObserver.onNext(UdfResponse.newBuilder().setControl(ctrl).build()) + } + + private def closeWithProtocolError(msg: String): Unit = { + sendControl(UdfControlResponse.newBuilder() + .setError(ErrorResponse.newBuilder() + .setError(ExecutionError.newBuilder() + .setProtocol(ProtocolError.newBuilder().setMessage(msg).build()) + .build()) + .build()) + .build()) + // The protocol requires a Cancel/Finish/Error envelope plus + // onCompleted to fully close the stream. After ErrorResponse the + // canonical terminator is CancelResponse. + sendControl(UdfControlResponse.newBuilder() + .setCancel(CancelResponse.getDefaultInstance) + .build()) + finalizeDone() + } + } +} diff --git a/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/testing/NoOpWorkerSession.scala b/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/testing/NoOpWorkerSession.scala new file mode 100644 index 0000000000000..276217e9a57d4 --- /dev/null +++ b/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/testing/NoOpWorkerSession.scala @@ -0,0 +1,46 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.udf.worker.grpc.testing + +import org.apache.spark.udf.worker.{Cancel, DataRequest, DataResponse, Finish, FinishResponse, + Init, InitResponse} +import org.apache.spark.udf.worker.core.{Termination, WorkerHandle, WorkerLogger, WorkerSession} + +/** + * 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 `GrpcWorkerSession` against a real or in-process + * echo worker. + */ +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[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 + } +} diff --git a/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/UnixSocketWorkerConnection.scala b/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/testing/SocketFileConnection.scala similarity index 60% rename from udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/UnixSocketWorkerConnection.scala rename to udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/testing/SocketFileConnection.scala index b3b40d16e7443..7e81ad74d6987 100644 --- a/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/UnixSocketWorkerConnection.scala +++ b/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/testing/SocketFileConnection.scala @@ -14,26 +14,23 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.spark.udf.worker.core +package org.apache.spark.udf.worker.grpc.testing import java.io.File -import org.apache.spark.annotation.Experimental +import org.apache.spark.udf.worker.core.WorkerConnection /** - * :: Experimental :: - * A [[WorkerConnection]] over a Unix domain socket. Owns the socket - * path and removes the socket file on [[close]]. Subclasses provide the - * protocol-specific channel (e.g. gRPC over UDS) and may override - * [[close]] to add transport-level shutdown -- they should call - * `super.close()` to ensure the socket file is removed. + * 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. * - * [[close]] is idempotent: deleting an already-removed file is a no-op. + * Suitable for dispatcher-lifecycle tests that don't need to drive the + * gRPC protocol -- e.g. verifying that a worker spec spawns a real + * worker process that creates the expected socket. */ -@Experimental -abstract class UnixSocketWorkerConnection(val socketPath: String) - extends WorkerConnection { - +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() diff --git a/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/testing/TestDirectGrpcDispatcher.scala b/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/testing/TestDirectGrpcDispatcher.scala new file mode 100644 index 0000000000000..fb75e3f0a570b --- /dev/null +++ b/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/testing/TestDirectGrpcDispatcher.scala @@ -0,0 +1,75 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.udf.worker.grpc.testing + +import java.io.File +import java.util.concurrent.TimeUnit + +import org.apache.spark.udf.worker.UDFWorkerSpecification +import org.apache.spark.udf.worker.core.{WorkerConnection, WorkerHandle, WorkerSession} +import org.apache.spark.udf.worker.core.direct.{DirectWorkerException, DirectWorkerTimeoutException} +import org.apache.spark.udf.worker.core.direct.DirectWorkerDispatcher.READY_POLL_INTERVAL_MS +import org.apache.spark.udf.worker.grpc.DirectGrpcDispatcher + +/** + * A [[DirectGrpcDispatcher]] convenience for tests: overrides the + * transport and session hooks to yield [[SocketFileConnection]]s and + * [[NoOpWorkerSession]]s, so lifecycle tests exercise the dispatcher's + * spawn / wait-for-ready / cleanup machinery without driving the gRPC + * protocol. + * + * Reusable across modules: callers in `sql/core` (or anywhere with a + * test-jar dependency on `udf-worker-grpc`) can drop this in for tests + * that only need to verify a worker spec produces a spawnable worker. + */ +class TestDirectGrpcDispatcher(spec: UDFWorkerSpecification) + extends DirectGrpcDispatcher(spec) { + + override protected def newConnection(address: String): WorkerConnection = + new SocketFileConnection(address) + + override protected def connectWorker( + address: String, + process: Process, + outputFile: File): WorkerConnection = { + val deadlineNanos = System.nanoTime() + TimeUnit.MILLISECONDS.toNanos(initTimeoutMs) + val socket = new File(address) + while (!socket.exists() && System.nanoTime() < deadlineNanos) { + if (!process.isAlive) throwWorkerExited(process, address, outputFile) + Thread.sleep(READY_POLL_INTERVAL_MS) + } + if (!socket.exists()) { + if (!process.isAlive) throwWorkerExited(process, address, outputFile) + throw new DirectWorkerTimeoutException( + s"Worker did not create socket at $address within ${initTimeoutMs}ms\n" + + readOutputTail(outputFile)) + } + newConnection(address) + } + + override protected def newSession(workerHandle: WorkerHandle): WorkerSession = + new NoOpWorkerSession(workerHandle, logger) + + private def throwWorkerExited( + process: Process, + address: String, + outputFile: File): Nothing = { + throw new DirectWorkerException( + s"Worker exited with code ${process.exitValue()} before creating socket at $address\n" + + readOutputTail(outputFile)) + } +} diff --git a/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/testing/UdfWorkerProtocolClient.scala b/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/testing/UdfWorkerProtocolClient.scala new file mode 100644 index 0000000000000..f0b6bef398ad6 --- /dev/null +++ b/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/testing/UdfWorkerProtocolClient.scala @@ -0,0 +1,148 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.udf.worker.grpc.testing + +import java.util.concurrent.{CountDownLatch, LinkedBlockingQueue, TimeUnit} +import java.util.concurrent.atomic.AtomicBoolean + +import scala.jdk.CollectionConverters._ + +import io.grpc.stub.StreamObserver + +import org.apache.spark.udf.worker.{Cancel, DataRequest, DataResponse, ExecutionError, Finish, + Init, PayloadChunk, UdfControlRequest, UdfControlResponse, UdfRequest, UdfResponse, + UdfWorkerGrpc} + +/** Engine-side protocol driver shared by UDF worker protocol contract suites. */ +private[grpc] final class UdfWorkerProtocolClient(stub: UdfWorkerGrpc.UdfWorkerStub) { + private val results = new LinkedBlockingQueue[DataResponse]() + private val done = new CountDownLatch(1) + private val initResponse = new CountDownLatch(1) + private val requestCompleted = new AtomicBoolean(false) + private val requestLock = new Object + + @volatile var executionError: Option[ExecutionError] = None + @volatile var streamError: Option[Throwable] = None + + private val responseObserver = new StreamObserver[UdfResponse] { + override def onNext(response: UdfResponse): Unit = { + response.getResponseCase match { + case UdfResponse.ResponseCase.DATA => + results.add(response.getData) + + case UdfResponse.ResponseCase.CONTROL => + val control = response.getControl + control.getControlCase match { + case UdfControlResponse.ControlCase.INIT => + if (control.getInit.hasError) { + executionError = Some(control.getInit.getError) + sendCancel(Cancel.newBuilder().setReason("aborting after init error").build()) + } + initResponse.countDown() + + case UdfControlResponse.ControlCase.ERROR => + executionError = Some(control.getError.getError) + sendCancel(Cancel.newBuilder().setReason("aborting after execution error").build()) + + case UdfControlResponse.ControlCase.FINISH | + UdfControlResponse.ControlCase.CANCEL => + completeRequestStream() + done.countDown() + + case unexpected => + throw new IllegalStateException(s"unexpected control response: $unexpected") + } + + case unexpected => + throw new IllegalStateException(s"unexpected response type: $unexpected") + } + } + + override def onError(t: Throwable): Unit = { + streamError = Some(t) + completeRequestStream() + initResponse.countDown() + done.countDown() + } + + override def onCompleted(): Unit = { + initResponse.countDown() + done.countDown() + } + } + + private val requestObserver: StreamObserver[UdfRequest] = stub.execute(responseObserver) + + def sendInit(init: Init): Unit = sendRaw(UdfRequest.newBuilder() + .setControl(UdfControlRequest.newBuilder().setInit(init).build()) + .build()) + + def sendPayloadChunk(chunk: PayloadChunk): Unit = sendRaw(UdfRequest.newBuilder() + .setControl(UdfControlRequest.newBuilder().setPayload(chunk).build()) + .build()) + + def sendData(data: DataRequest): Unit = { + requireInitResponse() + sendRaw(UdfRequest.newBuilder().setData(data).build()) + } + + def sendFinish(finish: Finish = Finish.getDefaultInstance): Unit = { + requireInitResponse() + if (!requestCompleted.get()) { + sendRaw(UdfRequest.newBuilder() + .setControl(UdfControlRequest.newBuilder().setFinish(finish).build()) + .build()) + } + } + + def sendCancel(cancel: Cancel = Cancel.getDefaultInstance): Unit = { + if (!requestCompleted.get()) { + sendRaw(UdfRequest.newBuilder() + .setControl(UdfControlRequest.newBuilder().setCancel(cancel).build()) + .build()) + } + } + + /** Sends an unsequenced request for tests that intentionally violate ordering. */ + def sendRaw(request: UdfRequest): Unit = requestLock.synchronized { + if (!requestCompleted.get()) requestObserver.onNext(request) + } + + def awaitInitResponse(timeoutMs: Long = 5000L): Boolean = + initResponse.await(timeoutMs, TimeUnit.MILLISECONDS) + + def awaitDone(timeoutMs: Long = 5000L): Boolean = + done.await(timeoutMs, TimeUnit.MILLISECONDS) + + def drainResults(): Seq[DataResponse] = { + val buffer = new java.util.ArrayList[DataResponse]() + results.drainTo(buffer) + buffer.asScala.toSeq + } + + def completeRequestStream(): Unit = requestLock.synchronized { + if (requestCompleted.compareAndSet(false, true)) { + requestObserver.onCompleted() + } + } + + private def requireInitResponse(): Unit = { + if (!awaitInitResponse()) { + throw new IllegalStateException("InitResponse not received within timeout") + } + } +} diff --git a/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/testing/UdfWorkerTestImplementation.scala b/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/testing/UdfWorkerTestImplementation.scala new file mode 100644 index 0000000000000..3effdda60b995 --- /dev/null +++ b/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/testing/UdfWorkerTestImplementation.scala @@ -0,0 +1,99 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.udf.worker.grpc.testing + +import com.google.protobuf.ByteString + +import org.apache.spark.udf.worker.{DataRequest, DataResponse, ExecutionError, Init, UdfPayload, + UDFWorkerDataFormat, UdfWorkerGrpc, UserError, WorkerError} + +private[grpc] final case class SuccessfulExecution( + init: Init, + requests: Seq[DataRequest], + expectedResponses: Seq[DataResponse]) + +private[grpc] final case class UserErrorExecution( + init: Init, + request: DataRequest, + expectedError: ExecutionError) + +private[grpc] final case class InitErrorExecution( + init: Init, + expectedError: ExecutionError) + +/** Implementation-specific scenarios consumed by [[UdfWorkerProtocolSuite]]. */ +private[grpc] trait UdfWorkerTestImplementation { + def newService(): UdfWorkerGrpc.UdfWorkerImplBase + def supportedProtocolVersion: Int + def singleRequest: SuccessfulExecution + def multipleRequests: SuccessfulExecution + def noInput: SuccessfulExecution + def userError: UserErrorExecution + def initError: InitErrorExecution +} + +/** Protocol scenarios for the minimal echo worker test implementation. */ +private[grpc] object EchoWorkerTestImplementation extends UdfWorkerTestImplementation { + import EchoWorkerService.{ErrorTrigger, InitErrorTrigger, SupportedVersion} + + override def newService(): UdfWorkerGrpc.UdfWorkerImplBase = new EchoWorkerService + + override val supportedProtocolVersion: Int = SupportedVersion + + override val singleRequest: SuccessfulExecution = SuccessfulExecution( + init("single-request"), + Seq(data("hello")), + Seq(response("hello"))) + + override val multipleRequests: SuccessfulExecution = SuccessfulExecution( + init("multiple-requests"), + Seq("batch-1", "batch-2", "batch-3").map(data), + Seq("batch-1", "batch-2", "batch-3").map(response)) + + override val noInput: SuccessfulExecution = + SuccessfulExecution(init("no-input"), Seq.empty, Seq.empty) + + override val userError: UserErrorExecution = UserErrorExecution( + init("user-error"), + DataRequest.newBuilder().setData(ErrorTrigger).build(), + ExecutionError.newBuilder() + .setUser(UserError.newBuilder() + .setMessage("simulated user-code error") + .setErrorClass("SimulatedError") + .build()) + .build()) + + override val initError: InitErrorExecution = InitErrorExecution( + init(InitErrorTrigger), + ExecutionError.newBuilder() + .setWorker(WorkerError.newBuilder().setMessage("simulated init failure").build()) + .build()) + + private def init(payload: String): Init = init(ByteString.copyFromUtf8(payload)) + + private def init(payload: ByteString): Init = Init.newBuilder() + .setProtocolVersion(SupportedVersion) + .setDataFormat(UDFWorkerDataFormat.ARROW) + .setUdf(UdfPayload.newBuilder().setPayload(payload).setFormat("echo").build()) + .build() + + private def data(value: String): DataRequest = + DataRequest.newBuilder().setData(ByteString.copyFromUtf8(value)).build() + + private def response(value: String): DataResponse = + DataResponse.newBuilder().setData(ByteString.copyFromUtf8(value)).build() +}