From 804ec40cfcf22173637b907e16106d549a077975 Mon Sep 17 00:00:00 2001 From: Haiyang Sun Date: Mon, 3 Aug 2026 15:38:39 +0000 Subject: [PATCH 1/5] [SPARK-56413][SQL] Add gRPC UDF worker protocol dispatcher --- sql/core/pom.xml | 13 +- .../PythonUDFWorkerSpecificationSuite.scala | 11 +- udf/worker/README.md | 120 ++- .../spark/udf/worker/core/WorkerHandle.scala | 9 + .../core/direct/DirectWorkerDispatcher.scala | 147 ++- .../core/direct/DirectWorkerProcess.scala | 4 +- .../udf/worker/core/WorkerSessionSuite.scala | 9 + udf/worker/grpc/pom.xml | 49 + .../worker/grpc/DirectGrpcDispatcher.scala} | 223 +++-- .../udf/worker/grpc/GrpcWorkerChannel.scala | 148 +++ .../grpc/UnixDomainSocketTransport.scala | 148 +++ ...DirectGrpcDispatcherIntegrationSuite.scala | 530 ++++++++++ .../grpc}/DirectWorkerDispatcherSuite.scala | 103 +- .../udf/worker/grpc/EchoProtocolSuite.scala | 936 +----------------- .../GrpcWorkerSessionConcurrencySuite.scala | 8 +- .../worker/grpc/UdfWorkerProtocolSuite.scala | 300 ++++++ .../grpc/testing/EchoGrpcWorkerMain.scala | 152 +++ .../grpc/testing/EchoWorkerService.scala | 309 ++++++ .../grpc/testing/NoOpWorkerSession.scala | 46 + .../grpc/testing/SocketFileConnection.scala} | 23 +- .../testing/TestDirectGrpcDispatcher.scala | 75 ++ .../testing/UdfWorkerProtocolClient.scala | 148 +++ .../testing/UdfWorkerTestImplementation.scala | 99 ++ 23 files changed, 2443 insertions(+), 1167 deletions(-) rename udf/worker/{core/src/test/scala/org/apache/spark/udf/worker/core/TestDirectWorkerHelpers.scala => grpc/src/main/scala/org/apache/spark/udf/worker/grpc/DirectGrpcDispatcher.scala} (55%) create mode 100644 udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/GrpcWorkerChannel.scala create mode 100644 udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/UnixDomainSocketTransport.scala create mode 100644 udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/DirectGrpcDispatcherIntegrationSuite.scala rename udf/worker/{core/src/test/scala/org/apache/spark/udf/worker/core => grpc/src/test/scala/org/apache/spark/udf/worker/grpc}/DirectWorkerDispatcherSuite.scala (90%) create mode 100644 udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/UdfWorkerProtocolSuite.scala create mode 100644 udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/testing/EchoGrpcWorkerMain.scala create mode 100644 udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/testing/EchoWorkerService.scala create mode 100644 udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/testing/NoOpWorkerSession.scala rename udf/worker/{core/src/main/scala/org/apache/spark/udf/worker/core/UnixSocketWorkerConnection.scala => grpc/src/test/scala/org/apache/spark/udf/worker/grpc/testing/SocketFileConnection.scala} (60%) create mode 100644 udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/testing/TestDirectGrpcDispatcher.scala create mode 100644 udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/testing/UdfWorkerProtocolClient.scala create mode 100644 udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/testing/UdfWorkerTestImplementation.scala diff --git a/sql/core/pom.xml b/sql/core/pom.xml index 281546b901a7a..b393e1f2aa6f2 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 4de8342208399..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,935 +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) - // gRPC forbids concurrent calls to the request StreamObserver too: the response - // observer half-closes it (completeRequestStream) on a callback thread while the - // test thread may still be sending a trailing Cancel. Serialize both through this. - private val requestLock = new Object - // 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 = requestLock.synchronized { - // 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. requestLock makes this check - // and the send atomic against completeRequestStream's half-close. - if (!requestCompleted.get()) { - 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 = requestLock.synchronized { - 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() +} From dee7c0c25509bc3b2a5960ad2341b070f5fea7f8 Mon Sep 17 00:00:00 2001 From: Haiyang Sun Date: Mon, 10 Aug 2026 11:37:29 +0000 Subject: [PATCH 2/5] [SPARK-56413][SQL] Simplify gRPC dispatcher changes --- sql/core/pom.xml | 13 +- .../PythonUDFWorkerSpecificationSuite.scala | 11 +- udf/worker/README.md | 120 +--- .../spark/udf/worker/core/WorkerHandle.scala | 9 - .../core/direct/DirectWorkerDispatcher.scala | 16 +- .../core/direct/DirectWorkerProcess.scala | 4 +- .../core}/DirectWorkerDispatcherSuite.scala | 82 +-- .../worker/core/TestDirectWorkerHelpers.scala | 306 ++++++++++ .../udf/worker/core/WorkerSessionSuite.scala | 9 - udf/worker/grpc/pom.xml | 10 - .../worker/grpc/DirectGrpcDispatcher.scala | 55 +- ...DirectGrpcDispatcherIntegrationSuite.scala | 495 +++------------- .../udf/worker/grpc/EchoProtocolSuite.scala | 552 +++++++++++++++++- .../GrpcWorkerSessionConcurrencySuite.scala | 8 +- .../worker/grpc/UdfWorkerProtocolSuite.scala | 300 ---------- .../grpc/testing/EchoGrpcWorkerMain.scala | 4 +- .../grpc/testing/NoOpWorkerSession.scala | 46 -- .../grpc/testing/SocketFileConnection.scala | 38 -- .../testing/TestDirectGrpcDispatcher.scala | 75 --- .../testing/UdfWorkerProtocolClient.scala | 148 ----- .../testing/UdfWorkerTestImplementation.scala | 99 ---- 21 files changed, 1045 insertions(+), 1355 deletions(-) rename udf/worker/{grpc/src/test/scala/org/apache/spark/udf/worker/grpc => core/src/test/scala/org/apache/spark/udf/worker/core}/DirectWorkerDispatcherSuite.scala (92%) create mode 100644 udf/worker/core/src/test/scala/org/apache/spark/udf/worker/core/TestDirectWorkerHelpers.scala delete mode 100644 udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/UdfWorkerProtocolSuite.scala delete mode 100644 udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/testing/NoOpWorkerSession.scala delete mode 100644 udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/testing/SocketFileConnection.scala delete mode 100644 udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/testing/TestDirectGrpcDispatcher.scala delete mode 100644 udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/testing/UdfWorkerProtocolClient.scala delete mode 100644 udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/testing/UdfWorkerTestImplementation.scala diff --git a/sql/core/pom.xml b/sql/core/pom.xml index b393e1f2aa6f2..281546b901a7a 100644 --- a/sql/core/pom.xml +++ b/sql/core/pom.xml @@ -308,20 +308,9 @@ spark-udf-worker-core_${scala.binary.version} ${project.version} - - - org.apache.spark - spark-udf-worker-grpc_${scala.binary.version} - ${project.version} - test - org.apache.spark - spark-udf-worker-grpc_${scala.binary.version} + spark-udf-worker-core_${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 bf8e6bc14b52c..bbd47741aba1c 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/externalUDF/PythonUDFWorkerSpecificationSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/externalUDF/PythonUDFWorkerSpecificationSuite.scala @@ -29,8 +29,7 @@ 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.WorkerConnection -import org.apache.spark.udf.worker.grpc.testing.TestDirectGrpcDispatcher +import org.apache.spark.udf.worker.core.{TestDirectWorkerDispatcher, WorkerConnection} /** * A [[WorkerConnection]] that opens a real Unix-domain-socket channel to the @@ -53,12 +52,12 @@ private object ConnectingSocketConnection { } /** - * A [[TestDirectGrpcDispatcher]] whose connection actually connects to the + * A [[TestDirectWorkerDispatcher]] whose connection actually connects to the * worker socket (rather than just checking file existence), so the test * verifies the spawned Python worker is reachable. */ private class ConnectingTestDispatcher(spec: UDFWorkerSpecification) - extends TestDirectGrpcDispatcher(spec) { + extends TestDirectWorkerDispatcher(spec) { override protected def newConnection(address: String): WorkerConnection = ConnectingSocketConnection.connect(address) } @@ -75,8 +74,8 @@ private class ConnectingTestDispatcher(spec: UDFWorkerSpecification) * for SIGTERM, verifying that pythonExec, PYTHONPATH, env vars, * and command construction all work end-to-end. * - * Reuses [[TestDirectGrpcDispatcher]]'s spawn / wait-for-ready machinery - * (shared from `udf-worker-grpc` test sources) but overrides the connection + * Reuses [[TestDirectWorkerDispatcher]]'s spawn / wait-for-ready machinery + * (shared from `udf-worker-core` test sources) but overrides the connection * to open a real UDS channel -- the Python test worker only binds a raw * socket, it does not host a gRPC server, so a full gRPC session is not * exercised here. diff --git a/udf/worker/README.md b/udf/worker/README.md index f82e89a7ccfd4..2250c8df10535 100644 --- a/udf/worker/README.md +++ b/udf/worker/README.md @@ -29,84 +29,6 @@ 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 ``` @@ -117,34 +39,26 @@ udf/worker/ │ udf_service.proto -- UdfWorker gRPC service (Execute, Manage) │ common.proto -- shared enums (UDFWorkerDataFormat, etc.) │ -├── core/ -- abstract interfaces (no gRPC dependency) +├── core/ -- abstract interfaces │ 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 + WorkerHandle +│ DirectWorkerProcess.scala -- OS process + connection + UDS socket +│ DirectWorkerSession.scala -- session backed by a direct process │ -└── 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 +└── grpc/ -- gRPC transport (gRPC runtime confined here) + (generated) -- UdfWorkerGrpc service stubs from proto/udf_service.proto ``` The `core/` package defines abstract interfaces that are independent of how -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. +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. The `grpc/` module owns the gRPC service-stub generation (from `proto/`'s `udf_service.proto`) and the gRPC runtime dependencies. Keeping @@ -216,10 +130,9 @@ val spec = UDFWorkerSpecification.newBuilder() .build()) .build() -// 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) +// 2. Create a dispatcher. Use a protocol-specific subclass of +// DirectWorkerDispatcher (e.g., gRPC over UDS). +val dispatcher: WorkerDispatcher = ... // 3. Create a session for one UDF execution. val session = dispatcher.createSession(securityScope = None) @@ -271,17 +184,14 @@ build/sbt "udf-worker-core/test" "udf-worker-grpc/test" ## Current status -The framework ships the core abstraction layer, the direct worker dispatcher, -and the first concrete transport (**gRPC over UDS** via `DirectGrpcDispatcher`). +This is the **first MVP** providing the core abstraction layer and the +direct worker dispatcher. 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-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) +- **Protocol-specific implementations** -- e.g., gRPC over UDS ## 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 821b00be2abdd..c66cb5689b875 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,13 +63,4 @@ 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 5a24c25750628..a2a8197a6b979 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 @@ -207,10 +207,11 @@ abstract class DirectWorkerDispatcher( * 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 `workerHandle.connection` for the wire - * transport. + * dispatcher-side cleanup and `connection` for the wire transport. */ - protected def newSession(workerHandle: WorkerHandle): WorkerSession + protected def newSession( + workerHandle: WorkerHandle, + connection: WorkerConnection): WorkerSession /** * Test-only hook: invoked once per [[createSession]] after the worker @@ -232,6 +233,12 @@ abstract class DirectWorkerDispatcher( */ protected def afterWorkerRegistered(worker: DirectWorkerProcess): Unit = () + /** + * Test-only hook invoked after [[close]] wins its idempotence CAS and before + * it waits for in-flight session creation. The default is a no-op. + */ + protected def afterCloseStarted(): Unit = () + override def createSession( securityScope: Option[WorkerSecurityScope]): WorkerSession = { require(securityScope.isEmpty, @@ -256,7 +263,7 @@ abstract class DirectWorkerDispatcher( // The test hook may block, so close could have started since the // first post-publication check. if (closed.get()) throwClosed() - newSession(worker) + newSession(worker, worker.connection) } catch { case e: InterruptedException => Thread.currentThread().interrupt() @@ -325,6 +332,7 @@ abstract class DirectWorkerDispatcher( if (!closed.compareAndSet(false, true)) { return } + afterCloseStarted() awaitSessionCreations() // TODO [SPARK-55278]: Cleanup sessions as well? // TODO [SPARK-55278]: close workers in parallel -- today shutdown is serialised at 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 9ff5cc07d2e92..52e7d5123d81b 100644 --- a/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/direct/DirectWorkerProcess.scala +++ b/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/direct/DirectWorkerProcess.scala @@ -72,8 +72,8 @@ class DirectWorkerProcess( /** The OS process handle for this worker. */ def process: Process = artifacts.process - /** The transport connection for this worker (see [[WorkerHandle.connection]]). */ - override def connection: WorkerConnection = artifacts.connection + /** The transport connection for this worker. */ + 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/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/DirectWorkerDispatcherSuite.scala b/udf/worker/core/src/test/scala/org/apache/spark/udf/worker/core/DirectWorkerDispatcherSuite.scala similarity index 92% rename from udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/DirectWorkerDispatcherSuite.scala rename to udf/worker/core/src/test/scala/org/apache/spark/udf/worker/core/DirectWorkerDispatcherSuite.scala index 86aa443803516..397a7b7dcab4f 100644 --- a/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/DirectWorkerDispatcherSuite.scala +++ b/udf/worker/core/src/test/scala/org/apache/spark/udf/worker/core/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.grpc +package org.apache.spark.udf.worker.core import java.io.File import java.nio.file.{Files, Path} @@ -34,7 +34,6 @@ 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 @@ -73,7 +72,7 @@ class DirectWorkerDispatcherSuite private def directWorker(runner: ProcessCallable): DirectWorker = DirectWorker.newBuilder().setRunner(runner).setProperties(udsProperties).build() - // DirectGrpcDispatcher requires the spec to advertise BIDIRECTIONAL_STREAMING. + // TestDirectWorkerDispatcher requires the spec to advertise BIDIRECTIONAL_STREAMING. private def bidiCapabilities: WorkerCapabilities = WorkerCapabilities.newBuilder() .addSupportedCommunicationPatterns(UDFProtoCommunicationPattern.BIDIRECTIONAL_STREAMING) .build() @@ -93,7 +92,7 @@ class DirectWorkerDispatcherSuite .setDirect(directWorker(runner)) .build() - private var dispatcher: TestDirectGrpcDispatcher = _ + private var dispatcher: TestDirectWorkerDispatcher = _ // Finalization message for session.close(); these lifecycle-only tests never // drive the protocol, so an empty Cancel suffices. @@ -114,7 +113,7 @@ class DirectWorkerDispatcherSuite private def createStubSession(): WorkerSession = dispatcher.createSession(None) - // Centralised cast: every test in this suite uses TestDirectGrpcDispatcher, + // Centralised cast: every test in this suite uses TestDirectWorkerDispatcher, // which always returns a DirectWorkerProcess. Wrapping in a helper with a // clear failure message (a) deduplicates the boilerplate cast at every // assertion site, and (b) means a future dispatcher that returns a @@ -137,7 +136,7 @@ class DirectWorkerDispatcherSuite } test("creates a worker and session") { - dispatcher = new TestDirectGrpcDispatcher(specWithRunner(defaultRunner)) + dispatcher = new TestDirectWorkerDispatcher(specWithRunner(defaultRunner)) val session = createStubSession() val worker = workerProcess(session) @@ -151,7 +150,7 @@ class DirectWorkerDispatcherSuite } test("concurrent createSession calls produce distinct workers") { - dispatcher = new TestDirectGrpcDispatcher(specWithRunner(defaultRunner)) + dispatcher = new TestDirectWorkerDispatcher(specWithRunner(defaultRunner)) val threads = 8 val sessions = new java.util.concurrent.ConcurrentLinkedQueue[WorkerSession]() @@ -195,7 +194,7 @@ class DirectWorkerDispatcherSuite } test("close shuts down all workers via SIGTERM") { - dispatcher = new TestDirectGrpcDispatcher(specWithRunner(defaultRunner)) + dispatcher = new TestDirectWorkerDispatcher(specWithRunner(defaultRunner)) val session1 = createStubSession() val session2 = createStubSession() @@ -244,7 +243,7 @@ class DirectWorkerDispatcherSuite .setDirect(DirectWorker.newBuilder() .setRunner(runner).setProperties(shortGracefulProps).build()) .build() - dispatcher = new TestDirectGrpcDispatcher(spec) + dispatcher = new TestDirectWorkerDispatcher(spec) val session = createStubSession() val worker = workerProcess(session) @@ -262,7 +261,7 @@ class DirectWorkerDispatcherSuite } test("closing a session terminates its worker") { - dispatcher = new TestDirectGrpcDispatcher(specWithRunner(defaultRunner)) + dispatcher = new TestDirectWorkerDispatcher(specWithRunner(defaultRunner)) val session = createStubSession() val worker = workerProcess(session) @@ -282,7 +281,7 @@ class DirectWorkerDispatcherSuite } test("concurrent session.close and dispatcher.close do not double-close the worker") { - dispatcher = new TestDirectGrpcDispatcher(specWithRunner(defaultRunner)) + dispatcher = new TestDirectWorkerDispatcher(specWithRunner(defaultRunner)) val sessions = (1 to 4).map(_ => createStubSession()) val workers = sessions.map(workerProcess) @@ -330,10 +329,11 @@ class DirectWorkerDispatcherSuite // workers or dispatcher-level transport state. val readyLatch = new java.util.concurrent.CountDownLatch(1) val releaseLatch = new java.util.concurrent.CountDownLatch(1) + val closeStartedLatch = new java.util.concurrent.CountDownLatch(1) val capturedWorkers = new java.util.concurrent.ConcurrentLinkedQueue[DirectWorkerProcess]() val transportClosed = new java.util.concurrent.atomic.AtomicBoolean(false) - val racing = new TestDirectGrpcDispatcher(specWithRunner(defaultRunner)) { + val racing = new TestDirectWorkerDispatcher(specWithRunner(defaultRunner)) { override protected def afterWorkerRegistered(worker: DirectWorkerProcess): Unit = { capturedWorkers.add(worker) readyLatch.countDown() @@ -348,6 +348,8 @@ class DirectWorkerDispatcherSuite transportClosed.set(true) super.closeTransport() } + + override protected def afterCloseStarted(): Unit = closeStartedLatch.countDown() } try { val outcome = @@ -369,8 +371,8 @@ class DirectWorkerDispatcherSuite val closeThread = new Thread(() => racing.close(), "close-racer") closeThread.start() - // Give close() time to flip `closed` and reach the creation barrier. - Thread.sleep(200) + assert(closeStartedLatch.await(10, java.util.concurrent.TimeUnit.SECONDS), + "close thread never flipped the dispatcher to closed") assert(closeThread.isAlive, "close should wait for the in-flight createSession") assert(!transportClosed.get(), "transport cleanup must not run while createSession is in flight") @@ -427,7 +429,7 @@ class DirectWorkerDispatcherSuite .setDirect(DirectWorker.newBuilder() .setRunner(defaultRunner).setProperties(oversizedProps).build()) .build() - dispatcher = new TestDirectGrpcDispatcher(spec) + dispatcher = new TestDirectWorkerDispatcher(spec) val session = createStubSession() val worker = workerProcess(session) @@ -447,14 +449,14 @@ class DirectWorkerDispatcherSuite .setDirect(DirectWorker.newBuilder() .setRunner(defaultRunner).setProperties(oversizedProps).build()) .build() - dispatcher = new TestDirectGrpcDispatcher(spec) + dispatcher = new TestDirectWorkerDispatcher(spec) assert(dispatcher.initTimeoutMs == 30000L, s"init timeout should be capped at 30000ms, got ${dispatcher.initTimeoutMs}") } test("createSession after close is rejected") { - dispatcher = new TestDirectGrpcDispatcher(specWithRunner(defaultRunner)) + dispatcher = new TestDirectWorkerDispatcher(specWithRunner(defaultRunner)) dispatcher.close() val ex = intercept[IllegalStateException] { @@ -466,7 +468,7 @@ class DirectWorkerDispatcherSuite } test("socket directory is owner-only (0700) on POSIX") { - dispatcher = new TestDirectGrpcDispatcher(specWithRunner(defaultRunner)) + dispatcher = new TestDirectWorkerDispatcher(specWithRunner(defaultRunner)) // Drive one createSession so a worker (and therefore the socket dir) is // observable via the UDS connection's path. val session = createStubSession() @@ -486,7 +488,7 @@ class DirectWorkerDispatcherSuite } test("socket directory is removed after dispatcher.close") { - dispatcher = new TestDirectGrpcDispatcher(specWithRunner(defaultRunner)) + dispatcher = new TestDirectWorkerDispatcher(specWithRunner(defaultRunner)) val session = createStubSession() val socketDir = new File(udsPath(workerProcess(session))).toPath.getParent.toFile assert(socketDir.exists(), @@ -507,8 +509,10 @@ class DirectWorkerDispatcherSuite // must be terminated rather than leaked until dispatcher.close(). var capturedWorker: DirectWorkerProcess = null val failingDispatcher = - new TestDirectGrpcDispatcher(specWithRunner(defaultRunner)) { - override protected def newSession(workerHandle: WorkerHandle): WorkerSession = + new TestDirectWorkerDispatcher(specWithRunner(defaultRunner)) { + override protected def newSession( + workerHandle: WorkerHandle, + connection: WorkerConnection): WorkerSession = throw new RuntimeException("session creation failed") override protected def afterWorkerRegistered(w: DirectWorkerProcess): Unit = { capturedWorker = w @@ -535,7 +539,7 @@ class DirectWorkerDispatcherSuite .setDirect(DirectWorker.newBuilder().setRunner(defaultRunner).build()) .build() val ex = intercept[IllegalArgumentException] { - new TestDirectGrpcDispatcher(badSpec) + new TestDirectWorkerDispatcher(badSpec) } assert(ex.getMessage.contains("connection must be set"), s"expected missing-connection error, got: ${ex.getMessage}") @@ -551,7 +555,7 @@ class DirectWorkerDispatcherSuite .setRunner(defaultRunner).setProperties(tcpProperties).build()) .build() val ex = intercept[IllegalArgumentException] { - new TestDirectGrpcDispatcher(badSpec) + new TestDirectWorkerDispatcher(badSpec) } assert(ex.getMessage.contains("UNIX domain socket"), s"expected UDS-only error, got: ${ex.getMessage}") @@ -560,7 +564,7 @@ class DirectWorkerDispatcherSuite test("socket file is cleaned up when newConnection throws") { val capturedSocketPaths = new java.util.concurrent.ConcurrentLinkedQueue[String]() val failingDispatcher = - new TestDirectGrpcDispatcher(specWithRunner(defaultRunner)) { + new TestDirectWorkerDispatcher(specWithRunner(defaultRunner)) { override protected def newConnection(address: String): WorkerConnection = { capturedSocketPaths.add(address) throw new RuntimeException("connection creation failed") @@ -583,7 +587,7 @@ class DirectWorkerDispatcherSuite test("empty ProcessCallable command is rejected with a clear error") { val emptyRunner = ProcessCallable.newBuilder().build() - dispatcher = new TestDirectGrpcDispatcher(specWithRunner(emptyRunner)) + dispatcher = new TestDirectWorkerDispatcher(specWithRunner(emptyRunner)) val ex = intercept[IllegalArgumentException] { dispatcher.createSession(None) } @@ -596,7 +600,7 @@ class DirectWorkerDispatcherSuite .addCommand("bash").addCommand("-c") .addCommand("echo 'fatal: bad config' >&2; exit 42").addCommand("--") .build() - dispatcher = new TestDirectGrpcDispatcher(specWithRunner(runner)) + dispatcher = new TestDirectWorkerDispatcher(specWithRunner(runner)) val ex = intercept[RuntimeException] { dispatcher.createSession(None) @@ -625,7 +629,7 @@ class DirectWorkerDispatcherSuite .setDirect(DirectWorker.newBuilder() .setRunner(hangingRunner).setProperties(shortInitProps).build()) .build() - dispatcher = new TestDirectGrpcDispatcher(spec) + dispatcher = new TestDirectWorkerDispatcher(spec) val ex = intercept[DirectWorkerTimeoutException] { dispatcher.createSession(None) @@ -649,7 +653,7 @@ class DirectWorkerDispatcherSuite .addCommand("bash").addCommand("-c") .addCommand(s"touch ${markerFile.getAbsolutePath}").build()) .build() - dispatcher = new TestDirectGrpcDispatcher(specWithEnv(env = env)) + dispatcher = new TestDirectWorkerDispatcher(specWithEnv(env = env)) val session = dispatcher.createSession(None) session.close(emptyCancel) @@ -669,7 +673,7 @@ class DirectWorkerDispatcherSuite .addCommand("bash").addCommand("-c") .addCommand(s"touch ${markerFile.getAbsolutePath}").build()) .build() - dispatcher = new TestDirectGrpcDispatcher(specWithEnv(env = env)) + dispatcher = new TestDirectWorkerDispatcher(specWithEnv(env = env)) val session = dispatcher.createSession(None) session.close(emptyCancel) @@ -688,7 +692,7 @@ class DirectWorkerDispatcherSuite .addCommand("bash").addCommand("-c") .addCommand(s"touch ${markerFile.getAbsolutePath}").build()) .build() - dispatcher = new TestDirectGrpcDispatcher(specWithEnv(env = env)) + dispatcher = new TestDirectWorkerDispatcher(specWithEnv(env = env)) val session = dispatcher.createSession(None) session.close(emptyCancel) @@ -704,7 +708,7 @@ class DirectWorkerDispatcherSuite .addCommand("bash").addCommand("-c") .addCommand("echo 'missing dependency: libfoo' >&2; exit 7").build()) .build() - dispatcher = new TestDirectGrpcDispatcher(specWithEnv(env = env)) + dispatcher = new TestDirectWorkerDispatcher(specWithEnv(env = env)) val ex = intercept[RuntimeException] { dispatcher.createSession(None) @@ -724,7 +728,7 @@ class DirectWorkerDispatcherSuite .addCommand("sleep 30").build() val env = WorkerEnvironment.newBuilder().setInstallation(slowInstall).build() val shortTimeoutDispatcher = - new TestDirectGrpcDispatcher(specWithEnv(env = env)) { + new TestDirectWorkerDispatcher(specWithEnv(env = env)) { override protected def callableTimeoutMs: Long = 500L } try { @@ -749,7 +753,7 @@ class DirectWorkerDispatcherSuite .addCommand("bash").addCommand("-c") .addCommand(s"echo invoked >> ${counterFile.getAbsolutePath}").build()) .build() - dispatcher = new TestDirectGrpcDispatcher(specWithEnv(env = env)) + dispatcher = new TestDirectWorkerDispatcher(specWithEnv(env = env)) val s1 = dispatcher.createSession(None); s1.close(emptyCancel) val s2 = dispatcher.createSession(None); s2.close(emptyCancel) @@ -775,7 +779,7 @@ class DirectWorkerDispatcherSuite .addCommand( s"sleep 0.2; echo invoked >> ${counterFile.getAbsolutePath}").build()) .build() - dispatcher = new TestDirectGrpcDispatcher(specWithEnv(env = env)) + dispatcher = new TestDirectWorkerDispatcher(specWithEnv(env = env)) val threads = 4 val startGate = new java.util.concurrent.CountDownLatch(1) @@ -824,7 +828,7 @@ class DirectWorkerDispatcherSuite .addCommand( s"echo invoked >> ${counterFile.getAbsolutePath}; exit 1").build()) .build() - dispatcher = new TestDirectGrpcDispatcher(specWithEnv(env = env)) + dispatcher = new TestDirectWorkerDispatcher(specWithEnv(env = env)) val first = intercept[RuntimeException] { dispatcher.createSession(None) } assert(first.getMessage.contains("installation failed"), @@ -855,7 +859,7 @@ class DirectWorkerDispatcherSuite s"echo invoked >> ${counterFile.getAbsolutePath}; sleep 30").build()) .build() val timeoutDispatcher = - new TestDirectGrpcDispatcher(specWithEnv(env = env)) { + new TestDirectWorkerDispatcher(specWithEnv(env = env)) { override protected def callableTimeoutMs: Long = 500L } try { @@ -882,7 +886,7 @@ class DirectWorkerDispatcherSuite } test("non-None securityScope is rejected until pooling lands") { - dispatcher = new TestDirectGrpcDispatcher(specWithRunner(defaultRunner)) + dispatcher = new TestDirectWorkerDispatcher(specWithRunner(defaultRunner)) val scope = new WorkerSecurityScope { override def equals(obj: Any): Boolean = obj.isInstanceOf[this.type] override def hashCode(): Int = 0 @@ -900,7 +904,7 @@ class DirectWorkerDispatcherSuite .addCommand("bash").addCommand("-c").addCommand("exit 0").build()) .build() val ex = intercept[IllegalArgumentException] { - new TestDirectGrpcDispatcher(specWithEnv(env = env)) + new TestDirectWorkerDispatcher(specWithEnv(env = env)) } assert(ex.getMessage.contains("installation"), s"expected installation-required error, got: ${ex.getMessage}") @@ -915,7 +919,7 @@ class DirectWorkerDispatcherSuite .addCommand("bash").addCommand("-c") .addCommand(s"touch ${cleanupMarker.getAbsolutePath}").build()) .build() - dispatcher = new TestDirectGrpcDispatcher(specWithEnv(env = env)) + dispatcher = new TestDirectWorkerDispatcher(specWithEnv(env = env)) val session = dispatcher.createSession(None) session.close(emptyCancel) diff --git a/udf/worker/core/src/test/scala/org/apache/spark/udf/worker/core/TestDirectWorkerHelpers.scala b/udf/worker/core/src/test/scala/org/apache/spark/udf/worker/core/TestDirectWorkerHelpers.scala new file mode 100644 index 0000000000000..01fee8761aee6 --- /dev/null +++ b/udf/worker/core/src/test/scala/org/apache/spark/udf/worker/core/TestDirectWorkerHelpers.scala @@ -0,0 +1,306 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.udf.worker.core + +import java.io.{File, IOException} +import java.nio.charset.StandardCharsets +import java.nio.file.{Files, FileVisitResult, Path, Paths, SimpleFileVisitor} +import java.nio.file.attribute.{BasicFileAttributes, PosixFilePermissions} +import java.util.Locale + +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.READY_POLL_INTERVAL_MS + +/** + * 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() + } +} + +/** + * 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 + } +} + +/** + * A concrete [[DirectWorkerDispatcher]] for tests that spawns workers over a + * Unix domain socket and yields [[SocketFileConnection]]s / [[NoOpWorkerSession]]s, + * so lifecycle tests exercise the dispatcher's spawn / wait-for-ready / cleanup + * machinery without driving a wire protocol. Allocates a private 0700 socket + * directory at construction; each worker is given a UDS path inside it. + * + * Reusable across modules: callers in `sql/core` (or anywhere with a test-jar + * dependency on `udf-worker-core`) can drop this in for tests that only need to + * verify a worker spec produces a spawnable worker. + */ +class TestDirectWorkerDispatcher( + workerSpec: UDFWorkerSpecification, + logger: WorkerLogger = WorkerLogger.NoOp) + extends DirectWorkerDispatcher(workerSpec, logger) { + + // Upper bound on the rename-on-collision retry loop in newEndpointAddress. + private val MAX_SOCKET_LEAF_RETRIES = 16 + + private val MAX_UDS_PATH_BYTES = { + val osName = System.getProperty("os.name", "").toLowerCase(Locale.ROOT) + if (osName.contains("mac") || osName.contains("darwin")) 103 else 107 + } + + // The private 0700 socket directory, created in [[initialize]] (after the + // base class has validated the spec) and removed in [[closeTransport]]. + private lazy val socketDir: Path = createPrivateTempDirectory() + + override protected def initialize(): Unit = { + super.initialize() + // Force the lazy val now so the directory is created (and any failure + // surfaces) at construction time, after spec validation has passed. + socketDir + } + + /** + * Returns the UDS path the worker should bind. Uses a short 16-hex-char leaf + * (the worker's full UUID still travels via `--id`) to stay within the + * 108-byte UDS `sun_path` limit. + */ + override protected def newEndpointAddress(workerId: String): String = { + val short = workerId.replace("-", "").take(16) + var candidate = socketDir.resolve(s"w-$short.sock") + var suffix = 0 + while (Files.exists(candidate) && suffix < MAX_SOCKET_LEAF_RETRIES) { + suffix += 1 + candidate = socketDir.resolve(s"w-$short-$suffix.sock") + } + if (Files.exists(candidate)) { + throw new IllegalStateException( + s"could not allocate a free UDS path under $socketDir after " + + s"$MAX_SOCKET_LEAF_RETRIES retries (truncated id=$short)") + } + val address = candidate.toString + val pathBytes = address.getBytes(StandardCharsets.UTF_8).length + if (pathBytes > MAX_UDS_PATH_BYTES) { + throw new IllegalStateException( + s"UDS path requires $pathBytes UTF-8 bytes but this platform allows " + + s"at most $MAX_UDS_PATH_BYTES: $address") + } + address + } + + override protected def connectWorker( + address: String, + process: Process, + outputFile: File): WorkerConnection = { + 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 / READY_POLL_INTERVAL_MS).toInt) + var attempts = 0 + while (!file.exists() && attempts < maxAttempts) { + if (!process.isAlive) throwWorkerExitedBeforeSocket(process, address, outputFile) + Thread.sleep(READY_POLL_INTERVAL_MS) + attempts += 1 + } + if (!file.exists()) { + if (process.isAlive) { + DirectWorkerDispatcher.destroyForciblyAndReap( + process, logger, s"init timeout $address") + val tail = readOutputTail(outputFile) + throw new DirectWorkerTimeoutException( + s"Worker did not create socket at $address within ${initTimeoutMs}ms\n$tail") + } else { + // Worker exited after the last poll without creating the socket; + // prefer the exit-code message over the ambiguous "did not create". + throwWorkerExitedBeforeSocket(process, address, outputFile) + } + } + newConnection(address) + } + + override protected def cleanupEndpointAddress(address: String): Unit = { + Files.deleteIfExists(new File(address).toPath) + } + + override protected def closeTransport(): Unit = { + if (!Files.exists(socketDir)) return + var firstFailure: IOException = null + def recordFailure(e: IOException): Unit = { + if (firstFailure == null) firstFailure = e else firstFailure.addSuppressed(e) + } + def delete(path: Path): Unit = { + try Files.deleteIfExists(path) catch { case e: IOException => recordFailure(e) } + } + // Recursive post-order delete: today socketDir contains only socket files + // at the top level, but a future change that namespaces workers into + // subdirectories should not silently leak them. + Files.walkFileTree(socketDir, new SimpleFileVisitor[Path] { + override def visitFile(file: Path, attrs: BasicFileAttributes): FileVisitResult = { + delete(file) + FileVisitResult.CONTINUE + } + override def visitFileFailed(file: Path, exc: IOException): FileVisitResult = { + recordFailure(exc) + FileVisitResult.CONTINUE + } + override def postVisitDirectory(dir: Path, exc: IOException): FileVisitResult = { + if (exc != null) recordFailure(exc) + delete(dir) + FileVisitResult.CONTINUE + } + }) + if (firstFailure != null) throw firstFailure + } + + // `spec` is the same object as the `workerSpec` field but passed explicitly: + // at the point this runs (parent constructor body), `this` is only partially + // constructed and reading subclass fields is unsafe. + override protected def validateTransportSupport(spec: UDFWorkerSpecification): Unit = { + val props = spec.getDirect.getProperties + require(props.hasConnection, + "DirectWorker.properties.connection must be set") + val conn = props.getConnection + require(conn.getTransportCase == WorkerConnectionSpec.TransportCase.UNIX_DOMAIN_SOCKET, + "TestDirectWorkerDispatcher requires UNIX domain socket transport, " + + s"got ${conn.getTransportCase}") + require(spec.hasCapabilities, + "TestDirectWorkerDispatcher requires WorkerCapabilities declaring " + + "BIDIRECTIONAL_STREAMING in supported_communication_patterns") + val patterns = spec.getCapabilities.getSupportedCommunicationPatternsList + val supportsBidi = (0 until patterns.size()).exists { i => + patterns.get(i) == UDFProtoCommunicationPattern.BIDIRECTIONAL_STREAMING + } + require(supportsBidi, + "TestDirectWorkerDispatcher requires BIDIRECTIONAL_STREAMING " + + "in WorkerCapabilities.supported_communication_patterns") + } + + protected def newConnection(address: String): WorkerConnection = + new SocketFileConnection(address) + + override protected def newSession( + workerHandle: WorkerHandle, + connection: WorkerConnection): WorkerSession = + new NoOpWorkerSession(workerHandle, logger) + + private def throwWorkerExitedBeforeSocket( + process: Process, + address: String, + outputFile: File): Nothing = { + val tail = readOutputTail(outputFile) + throw new DirectWorkerException( + s"Worker exited with code ${process.exitValue()} " + + s"before creating socket at $address\n$tail") + } + + /** + * Creates a private (owner-only, 0700) temp directory for worker sockets. + * On POSIX filesystems the permissions are applied atomically at creation; + * the non-POSIX branch tightens best-effort after creation. + * + * The directory is anchored under a short base path (see [[shortTempBase]]) + * rather than the configured `java.io.tmpdir`. Builds point the latter at a + * deep location (e.g. `/target/tmp`), which combined with the + * generated socket leaf would push the worker's Unix-domain socket path past + * the platform `sun_path` limit (108 bytes on Linux, 104 on macOS) and fail + * with "AF_UNIX path too long". + */ + private def createPrivateTempDirectory(): Path = { + val attr = PosixFilePermissions.asFileAttribute( + PosixFilePermissions.fromString("rwx------")) + val base = shortTempBase() + try { + Files.createTempDirectory(base, "spark-udf-worker", attr) + } catch { + case _: UnsupportedOperationException => + val dir = Files.createTempDirectory(base, "spark-udf-worker") + val f = dir.toFile + // Bit-wise AND (NOT &&): all six setters must run even if an earlier + // one returns false, so the final permission state matches owner-only. + val applied = + f.setReadable(false, false) & f.setWritable(false, false) & + f.setExecutable(false, false) & f.setReadable(true, true) & + f.setWritable(true, true) & f.setExecutable(true, true) + if (!applied) { + logger.warn( + s"Could not fully restrict permissions on $dir; socket " + + s"directory may be accessible to other local users on this " + + s"filesystem") + } + dir + } + } + + /** + * Picks the shortest usable base directory for the socket temp dir so the + * resulting Unix-domain socket path stays within the platform `sun_path` + * limit (108 bytes on Linux, 104 on macOS). + * + * This mirrors how PySpark already chooses its UDS directory in + * `python/run-tests.py` (`spark.python.unix.domain.socket.dir`): prefer the + * OS temp-dir env vars (`TMPDIR`/`TEMP`/`TMP`) and fall back to `/tmp`. Builds + * point `java.io.tmpdir` at a deep `/target/tmp`, which combined with + * the generated socket leaf would overflow `sun_path`; the OS temp dir is + * short (e.g. a per-user `/tmp/...` on Linux runners, `/var/folders/...` via + * `$TMPDIR` on macOS). `java.io.tmpdir` remains the last-resort fallback. + */ + private def shortTempBase(): Path = { + val candidates = + Seq("TMPDIR", "TEMP", "TMP").flatMap(sys.env.get).filter(_.nonEmpty) :+ "/tmp" + val usable = candidates + .map(Paths.get(_)) + .filter(p => Files.isDirectory(p) && Files.isWritable(p)) + // Netty passes the unresolved pathname spelling to the native bind call. + // Rank that exact spelling in UTF-8 bytes, then validate the complete + // generated path in newEndpointAddress. + def encodedLength(p: Path): Int = p.toString.getBytes(StandardCharsets.UTF_8).length + usable + .sortBy(encodedLength) + .headOption + .getOrElse(Paths.get(System.getProperty("java.io.tmpdir"))) + } +} 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 002dad7ee1158..c8bbcfffd7bb2 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,12 +36,6 @@ 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 @@ -49,7 +43,6 @@ 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 } /** @@ -231,7 +224,6 @@ 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() @@ -254,7 +246,6 @@ 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 6d6323444b31c..c915e819410ea 100644 --- a/udf/worker/grpc/pom.xml +++ b/udf/worker/grpc/pom.xml @@ -100,16 +100,6 @@ grpc-netty ${io.grpc.version} - - io.netty - netty-codec-http2 - ${netty.version} - - - io.netty - netty-handler-proxy - ${netty.version} - io.netty diff --git a/udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/DirectGrpcDispatcher.scala b/udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/DirectGrpcDispatcher.scala index 1462c4e83ac8f..c35166535d4ac 100644 --- a/udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/DirectGrpcDispatcher.scala +++ b/udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/DirectGrpcDispatcher.scala @@ -17,8 +17,10 @@ package org.apache.spark.udf.worker.grpc import java.io.{File, IOException} +import java.nio.charset.StandardCharsets import java.nio.file.{Files, FileVisitResult, Path, Paths, SimpleFileVisitor} import java.nio.file.attribute.{BasicFileAttributes, PosixFilePermissions} +import java.util.Locale import java.util.concurrent.{CountDownLatch, TimeUnit} import scala.util.control.NonFatal @@ -58,6 +60,13 @@ class DirectGrpcDispatcher( // failure mode bounded if something pathological occupies the directory. private val MAX_SOCKET_LEAF_RETRIES = 16 + // sockaddr_un.sun_path includes its terminating NUL. Linux reserves 108 + // bytes and macOS 104, leaving at most 107 or 103 pathname bytes. + private val MAX_UDS_PATH_BYTES = { + val osName = System.getProperty("os.name", "").toLowerCase(Locale.ROOT) + if (osName.contains("mac") || osName.contains("darwin")) 103 else 107 + } + // 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 @@ -95,7 +104,14 @@ class DirectGrpcDispatcher( s"could not allocate a free UDS path under $socketDir after " + s"$MAX_SOCKET_LEAF_RETRIES retries (truncated id=$short)") } - candidate.toString + val address = candidate.toString + val pathBytes = address.getBytes(StandardCharsets.UTF_8).length + if (pathBytes > MAX_UDS_PATH_BYTES) { + throw new IllegalStateException( + s"UDS path requires $pathBytes UTF-8 bytes but this platform allows " + + s"at most $MAX_UDS_PATH_BYTES: $address") + } + address } override protected def connectWorker( @@ -163,19 +179,32 @@ class DirectGrpcDispatcher( override protected def closeTransport(): Unit = { if (!Files.exists(socketDir)) return + var firstFailure: IOException = null + def recordFailure(e: IOException): Unit = { + if (firstFailure == null) firstFailure = e else firstFailure.addSuppressed(e) + } + def delete(path: Path): Unit = { + try Files.deleteIfExists(path) catch { case e: IOException => recordFailure(e) } + } // Recursive post-order delete: today socketDir contains only socket files // at the top level, but a future change that namespaces workers into // subdirectories should not silently leak them. Files.walkFileTree(socketDir, new SimpleFileVisitor[Path] { override def visitFile(file: Path, attrs: BasicFileAttributes): FileVisitResult = { - try Files.deleteIfExists(file) catch { case _: IOException => () } + delete(file) + FileVisitResult.CONTINUE + } + override def visitFileFailed(file: Path, exc: IOException): FileVisitResult = { + recordFailure(exc) FileVisitResult.CONTINUE } override def postVisitDirectory(dir: Path, exc: IOException): FileVisitResult = { - try Files.deleteIfExists(dir) catch { case _: IOException => () } + if (exc != null) recordFailure(exc) + delete(dir) FileVisitResult.CONTINUE } }) + if (firstFailure != null) throw firstFailure } // `spec` is the same object as the `workerSpec` field but passed @@ -212,8 +241,10 @@ class DirectGrpcDispatcher( new GrpcWorkerChannel( address, logger, maxInboundMessageSize = grpcMaxInboundMessageSize) - override protected def newSession(workerHandle: WorkerHandle): WorkerSession = - workerHandle.connection match { + override protected def newSession( + workerHandle: WorkerHandle, + connection: WorkerConnection): WorkerSession = + connection match { case g: GrpcWorkerChannel => new GrpcWorkerSession(workerHandle, g.channel, logger) case other => @@ -305,16 +336,12 @@ class DirectGrpcDispatcher( val usable = candidates .map(Paths.get(_)) .filter(p => Files.isDirectory(p) && Files.isWritable(p)) - // Pick the SHORTEST usable base, measured by its *real* (symlink-resolved) path, - // because that is what actually counts against the `sun_path` limit when the socket - // is bound. On macOS `$TMPDIR` is a long `/var/folders/<...>/T` path while `/tmp` - // resolves to the short `/private/tmp`; picking the first writable candidate (the old - // behavior) chose the long `$TMPDIR` and overflowed the 104-byte macOS limit - // ("AF_UNIX path too long"). Choosing the shortest real path avoids that. - def realLen(p: Path): Int = - try p.toRealPath().toString.length catch { case _: IOException => p.toString.length } + // Netty passes the unresolved pathname spelling to the native bind call. + // Rank that exact spelling in UTF-8 bytes, then validate the complete + // generated path in newEndpointAddress. + def encodedLength(p: Path): Int = p.toString.getBytes(StandardCharsets.UTF_8).length usable - .sortBy(realLen) + .sortBy(encodedLength) .headOption .getOrElse(Paths.get(System.getProperty("java.io.tmpdir"))) } 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 index ee07fd73ef30f..41d3bb4d46b3c 100644 --- 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 @@ -18,57 +18,46 @@ 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, +import org.apache.spark.udf.worker.{Cancel, DataRequest, 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.core.WorkerSession +import org.apache.spark.udf.worker.core.direct.{DirectWorkerProcess, DirectWorkerTimeoutException} 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. + * End-to-end coverage for the integration points unique to [[DirectGrpcDispatcher]]: + * spawning a real gRPC worker over a Unix domain socket and waiting for gRPC + * readiness rather than only for the socket file. */ 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 = + private val javaClasspath = System.getProperty("java.class.path") + private val javaExecutable = Paths.get(System.getProperty("java.home"), "bin", "java").toString - - private val workerMainClass: String = + private val workerMainClass = classOf[EchoGrpcWorkerMain.type].getName.stripSuffix("$") private var dispatcher: DirectGrpcDispatcher = _ override def beforeEach(): Unit = { val supported = try { - UnixDomainSocketTransport.detect(); true + 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") + "Netty UDS native transport (epoll on Linux or kqueue on macOS) is required") } override def afterEach(): Unit = { @@ -78,290 +67,80 @@ class DirectGrpcDispatcherIntegrationSuite super.afterEach() } - // -------------------------------------------------------------------------- - // Spec builders - // -------------------------------------------------------------------------- - private def echoRunner: ProcessCallable = ProcessCallable.newBuilder() .addCommand(javaExecutable) - .addCommand("-cp").addCommand(javaClasspath) + .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 workerSpec( + runner: ProcessCallable = echoRunner, + initTimeoutMs: Int = 30000): UDFWorkerSpecification = + UDFWorkerSpecification.newBuilder() + .setCapabilities(WorkerCapabilities.newBuilder() + .addSupportedDataFormats(UDFWorkerDataFormat.ARROW) + .addSupportedCommunicationPatterns(UDFProtoCommunicationPattern.BIDIRECTIONAL_STREAMING) + .build()) + .setDirect(DirectWorker.newBuilder() + .setRunner(runner) + .setProperties(UDFWorkerProperties.newBuilder() + .setConnection(WorkerConnectionSpec.newBuilder() + .setUnixDomainSocket(UnixDomainSocket.getDefaultInstance) + .build()) + .setInitializationTimeoutMs(initTimeoutMs) + .setGracefulTerminationTimeoutMs(10000) + .build()) + .build()) + .build() - private def echoSpec: UDFWorkerSpecification = UDFWorkerSpecification.newBuilder() - .setCapabilities(capabilities) - .setDirect(DirectWorker.newBuilder() - .setRunner(echoRunner) - .setProperties(udsProperties) + private def basicInit: Init = Init.newBuilder() + .setProtocolVersion(1) + .setDataFormat(UDFWorkerDataFormat.ARROW) + .setUdf(UdfPayload.newBuilder() + .setPayload(ByteString.copyFromUtf8("echo")) + .setFormat("echo") .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}") + private def workerProcess(session: WorkerSession): DirectWorkerProcess = + session.workerHandle match { + case process: DirectWorkerProcess => process + 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")) + private def grpcChannel(process: DirectWorkerProcess): GrpcWorkerChannel = + process.connection match { + case channel: GrpcWorkerChannel => channel + case other => fail(s"Expected GrpcWorkerChannel, got ${other.getClass.getSimpleName}") } - } - test("responses larger than gRPC's default inbound limit round-trip successfully") { - newDispatcher() + test("spawns a gRPC worker and round-trips a response above gRPC's default limit") { + dispatcher = new DirectGrpcDispatcher(workerSpec()) + val session = dispatcher.createSession(None) + val process = workerProcess(session) + val channel = grpcChannel(process) + val socketFile = new File(channel.socketPath) + val socketDir = socketFile.getParentFile 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") + session.init(basicInit) + val input = DataRequest.newBuilder().setData(payload).build() + assert(session.process(Iterator.single(input), emptyFinish).map(_.getData).toList == + List(payload)) + } finally { + session.close(emptyCancel) } - } - - 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. + assert(!process.process.isAlive, "session close should terminate the worker") + assert(!socketFile.exists(), "session close should remove the worker socket") 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}") + assert(!socketDir.exists(), "dispatcher close should remove its socket directory") } test("a socket path alone does not make a worker ready") { @@ -380,151 +159,17 @@ class DirectGrpcDispatcherIntegrationSuite |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()) + .addCommand("bash") + .addCommand("-c") + .addCommand(socketOnlyWorker) + .addCommand("--") .build() - dispatcher = new DirectGrpcDispatcher(spec) + dispatcher = new DirectGrpcDispatcher(workerSpec(runner, initTimeoutMs = 1000)) - val error = intercept[org.apache.spark.udf.worker.core.direct.DirectWorkerTimeoutException] { + val error = intercept[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}") + assert(error.getMessage.contains("socket-created")) } - - 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/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 f26f1dc2c582b..26922185bb7a6 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,11 +16,551 @@ */ package org.apache.spark.udf.worker.grpc -import org.apache.spark.udf.worker.grpc.testing.{EchoWorkerTestImplementation, - UdfWorkerTestImplementation} +import java.util.concurrent.{CountDownLatch, LinkedBlockingQueue, TimeUnit} +import java.util.concurrent.atomic.AtomicBoolean -/** Runs the UDF worker protocol contract against the minimal echo implementation. */ -class EchoProtocolSuite extends UdfWorkerProtocolSuite { - override protected def implementation: UdfWorkerTestImplementation = - EchoWorkerTestImplementation +import com.google.protobuf.ByteString +import io.grpc.{ManagedChannel, Server} +import io.grpc.inprocess.{InProcessChannelBuilder, InProcessServerBuilder} +import io.grpc.stub.StreamObserver +import org.scalatest.BeforeAndAfterEach +// scalastyle:off funsuite +import org.scalatest.funsuite.AnyFunSuite + +import org.apache.spark.udf.worker.{Cancel, DataRequest, ExecutionError, Finish, Heartbeat, + Init, PayloadChunk, ShutdownRequest, UdfControlRequest, UdfControlResponse, UdfPayload, + UdfRequest, UdfResponse, UDFWorkerDataFormat, UdfWorkerGrpc, WorkerRequest} +import org.apache.spark.udf.worker.grpc.testing.EchoWorkerService + +/** + * 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 + + // Constants reused across protocol tests. Defined on the shared + // [[EchoWorkerService]] so this suite and the cross-process integration + // tests speak the same triggers. + import EchoWorkerService.{SupportedVersion, ErrorTrigger, InitErrorTrigger} + + private var server: Server = _ + private var channel: ManagedChannel = _ + protected var stub: UdfWorkerGrpc.UdfWorkerStub = _ + + /** Overridable so another implementation can run the shared protocol cases. */ + protected def workerService: UdfWorkerGrpc.UdfWorkerImplBase = new EchoWorkerService + + override def beforeEach(): Unit = { + val serverName = InProcessServerBuilder.generateName() + server = InProcessServerBuilder.forName(serverName) + .directExecutor() + .addService(workerService) + .build() + .start() + channel = InProcessChannelBuilder.forName(serverName).directExecutor().build() + stub = UdfWorkerGrpc.newStub(channel) + } + + override def afterEach(): Unit = { + channel.shutdownNow() + server.shutdownNow() + } + + // =========================================================================== + // 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. + */ + protected 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 = SupportedVersion): 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 = SupportedVersion): 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(SupportedVersion) + .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(ErrorTrigger.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(InitErrorTrigger.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(InitErrorTrigger.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 = SupportedVersion + 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") + } } 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 215900b7262ad..fb0eeae9ed376 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, WorkerConnection, WorkerHandle, WorkerLogger} +import org.apache.spark.udf.worker.core.{Termination, WorkerHandle, WorkerLogger} /** * Concurrency tests for [[GrpcWorkerSession]] that pin the wire-ordering and @@ -115,18 +115,12 @@ 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 deleted file mode 100644 index ad7853aaa135d..0000000000000 --- a/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/UdfWorkerProtocolSuite.scala +++ /dev/null @@ -1,300 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one or more - * contributor license agreements. See the NOTICE file distributed with - * this work for additional information regarding copyright ownership. - * The ASF licenses this file to You under the Apache License, Version 2.0 - * (the "License"); you may not use this file except in compliance with - * the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package org.apache.spark.udf.worker.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 index c22f0e8137e72..e78f02908eae4 100644 --- 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 @@ -20,6 +20,8 @@ import java.io.IOException import java.nio.file.{Files, Paths} import java.util.concurrent.TimeUnit +import scala.util.control.NonFatal + import io.grpc.Server import io.grpc.netty.NettyServerBuilder import io.netty.channel.EventLoopGroup @@ -108,7 +110,7 @@ object EchoGrpcWorkerMain { server.start() stderrln(s"listening on $socketPath (transport=${transport.name})") } catch { - case e: Throwable => + case NonFatal(e) => 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. 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 deleted file mode 100644 index 276217e9a57d4..0000000000000 --- a/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/testing/NoOpWorkerSession.scala +++ /dev/null @@ -1,46 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one or more - * contributor license agreements. See the NOTICE file distributed with - * this work for additional information regarding copyright ownership. - * The ASF licenses this file to You under the Apache License, Version 2.0 - * (the "License"); you may not use this file except in compliance with - * the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package org.apache.spark.udf.worker.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/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/testing/SocketFileConnection.scala b/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/testing/SocketFileConnection.scala deleted file mode 100644 index 7e81ad74d6987..0000000000000 --- a/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/testing/SocketFileConnection.scala +++ /dev/null @@ -1,38 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one or more - * contributor license agreements. See the NOTICE file distributed with - * this work for additional information regarding copyright ownership. - * The ASF licenses this file to You under the Apache License, Version 2.0 - * (the "License"); you may not use this file except in compliance with - * the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package org.apache.spark.udf.worker.grpc.testing - -import java.io.File - -import org.apache.spark.udf.worker.core.WorkerConnection - -/** - * 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 the - * gRPC 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() - } -} 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 deleted file mode 100644 index fb75e3f0a570b..0000000000000 --- a/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/testing/TestDirectGrpcDispatcher.scala +++ /dev/null @@ -1,75 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one or more - * contributor license agreements. See the NOTICE file distributed with - * this work for additional information regarding copyright ownership. - * The ASF licenses this file to You under the Apache License, Version 2.0 - * (the "License"); you may not use this file except in compliance with - * the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package org.apache.spark.udf.worker.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 deleted file mode 100644 index f0b6bef398ad6..0000000000000 --- a/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/testing/UdfWorkerProtocolClient.scala +++ /dev/null @@ -1,148 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one or more - * contributor license agreements. See the NOTICE file distributed with - * this work for additional information regarding copyright ownership. - * The ASF licenses this file to You under the Apache License, Version 2.0 - * (the "License"); you may not use this file except in compliance with - * the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package org.apache.spark.udf.worker.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 deleted file mode 100644 index 3effdda60b995..0000000000000 --- a/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/testing/UdfWorkerTestImplementation.scala +++ /dev/null @@ -1,99 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one or more - * contributor license agreements. See the NOTICE file distributed with - * this work for additional information regarding copyright ownership. - * The ASF licenses this file to You under the Apache License, Version 2.0 - * (the "License"); you may not use this file except in compliance with - * the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package org.apache.spark.udf.worker.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() -} From df6980327d8dd8c594df8a263f0a5598399e19de Mon Sep 17 00:00:00 2001 From: Haiyang Sun Date: Mon, 10 Aug 2026 16:06:01 +0000 Subject: [PATCH 3/5] [SPARK-56413][SQL] Address gRPC dispatcher review feedback --- udf/worker/README.md | 27 ++- .../core/direct/DirectWorkerDispatcher.scala | 54 ++--- .../UnixDomainSocketEndpointDirectory.scala | 157 ++++++++++++++ .../core/DirectWorkerDispatcherSuite.scala | 99 ++++++++- .../worker/core/TestDirectWorkerHelpers.scala | 155 +------------ .../worker/grpc/DirectGrpcDispatcher.scala | 204 +++--------------- .../udf/worker/grpc/GrpcWorkerChannel.scala | 46 ++-- .../grpc/UnixDomainSocketTransport.scala | 3 +- ...DirectGrpcDispatcherIntegrationSuite.scala | 28 ++- .../grpc/DirectGrpcDispatcherSuite.scala | 92 ++++++++ .../grpc/testing/EchoGrpcWorkerMain.scala | 119 +++++----- .../grpc/testing/EchoWorkerService.scala | 2 +- 12 files changed, 530 insertions(+), 456 deletions(-) create mode 100644 udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/direct/UnixDomainSocketEndpointDirectory.scala create mode 100644 udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/DirectGrpcDispatcherSuite.scala diff --git a/udf/worker/README.md b/udf/worker/README.md index 2250c8df10535..9e364bdb3c4c4 100644 --- a/udf/worker/README.md +++ b/udf/worker/README.md @@ -47,11 +47,16 @@ udf/worker/ │ │ │ └── direct/ -- "direct" creation: local OS processes │ DirectWorkerDispatcher.scala -- spawns processes, env lifecycle -│ DirectWorkerProcess.scala -- OS process + connection + UDS socket +│ DirectWorkerProcess.scala -- process, connection, cleanup artifacts │ DirectWorkerSession.scala -- session backed by a direct process +│ UnixDomainSocketEndpointDirectory.scala -- private UDS path lifecycle │ └── grpc/ -- gRPC transport (gRPC runtime confined here) - (generated) -- UdfWorkerGrpc service stubs from proto/udf_service.proto + DirectGrpcDispatcher.scala -- launches direct workers over gRPC on UDS + GrpcWorkerChannel.scala -- owns a Netty gRPC channel and event loop + GrpcWorkerSession.scala -- one UDF execution over an Execute stream + UnixDomainSocketTransport.scala -- selects native epoll/kqueue transport + (generated) -- UdfWorkerGrpc stubs from proto/udf_service.proto ``` The `core/` package defines abstract interfaces that are independent of how @@ -60,9 +65,10 @@ 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. -The `grpc/` module owns the gRPC service-stub generation (from -`proto/`'s `udf_service.proto`) and the gRPC runtime dependencies. Keeping -gRPC here means `proto/`, `core/`, and their consumers (`core`, `catalyst`, +The `grpc/` module owns the concrete direct dispatcher, channel and session +implementations, native Unix-domain-socket transport selection, gRPC service-stub +generation (from `proto/`'s `udf_service.proto`), and the gRPC runtime dependencies. +Keeping gRPC here means `proto/`, `core/`, and their consumers (`core`, `catalyst`, `sql/core`) carry no gRPC dependency on their classpath. ## Wire protocol @@ -105,6 +111,7 @@ import org.apache.spark.udf.worker.{ UDFWorkerSpecification, UnixDomainSocket, WorkerCapabilities, WorkerConnectionSpec, WorkerEnvironment} import org.apache.spark.udf.worker.core._ +import org.apache.spark.udf.worker.grpc.DirectGrpcDispatcher import com.google.protobuf.ByteString // 1. Define a worker spec (direct creation mode). @@ -130,9 +137,8 @@ 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 that launches workers and connects over gRPC on UDS. +val dispatcher: WorkerDispatcher = new DirectGrpcDispatcher(spec) // 3. Create a session for one UDF execution. val session = dispatcher.createSession(securityScope = None) @@ -184,14 +190,13 @@ 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. +This is the **first MVP** providing the core abstraction layer, direct worker +creation, and a gRPC-over-UDS dispatcher and session implementation. 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 ## Design references 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 a2a8197a6b979..0a1eb2c70f197 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 @@ -233,12 +233,6 @@ abstract class DirectWorkerDispatcher( */ protected def afterWorkerRegistered(worker: DirectWorkerProcess): Unit = () - /** - * Test-only hook invoked after [[close]] wins its idempotence CAS and before - * it waits for in-flight session creation. The default is a no-op. - */ - protected def afterCloseStarted(): Unit = () - override def createSession( securityScope: Option[WorkerSecurityScope]): WorkerSession = { require(securityScope.isEmpty, @@ -307,9 +301,12 @@ abstract class DirectWorkerDispatcher( if (inFlightSessionCreations == 0) sessionCreationLock.notifyAll() } - /** Waits uninterruptibly so cleanup can never race an in-flight creation. */ - private def awaitSessionCreations(): Unit = { - var interrupted = false + /** + * Waits uninterruptibly so cleanup can never race an in-flight creation. + * Returns whether interruption must be restored after cleanup completes. + */ + private def awaitSessionCreations(): Boolean = { + var interrupted = Thread.interrupted() sessionCreationLock.synchronized { while (inFlightSessionCreations != 0) { try { @@ -319,7 +316,7 @@ abstract class DirectWorkerDispatcher( } } } - if (interrupted) Thread.currentThread().interrupt() + interrupted } /** @@ -332,26 +329,29 @@ abstract class DirectWorkerDispatcher( if (!closed.compareAndSet(false, true)) { return } - afterCloseStarted() - awaitSessionCreations() - // TODO [SPARK-55278]: Cleanup sessions as well? - // TODO [SPARK-55278]: close workers in parallel -- today shutdown is serialised at - // N * gracefulTimeoutMs worst case. - workers.values().iterator().asScala.foreach { w => - try { - w.close() - } catch { + val interrupted = awaitSessionCreations() + try { + // TODO [SPARK-55278]: Cleanup sessions as well? + // TODO [SPARK-55278]: close workers in parallel -- today shutdown is serialised at + // N * gracefulTimeoutMs worst case. + workers.values().iterator().asScala.foreach { w => + try { + w.close() + } catch { + case NonFatal(e) => + logger.warn(s"Error closing worker ${w.id}", e) + } + } + workers.clear() + try closeTransport() catch { case NonFatal(e) => - logger.warn(s"Error closing worker ${w.id}", e) + logger.warn("Error cleaning up transport state", e) } + deregisterEnvironmentCleanupHook() + runEnvironmentCleanup() + } finally { + if (interrupted) Thread.currentThread().interrupt() } - workers.clear() - try closeTransport() catch { - case NonFatal(e) => - logger.warn("Error cleaning up transport state", e) - } - deregisterEnvironmentCleanupHook() - runEnvironmentCleanup() } // -- Environment lifecycle ------------------------------------------------- diff --git a/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/direct/UnixDomainSocketEndpointDirectory.scala b/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/direct/UnixDomainSocketEndpointDirectory.scala new file mode 100644 index 0000000000000..45c57c9aac299 --- /dev/null +++ b/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/direct/UnixDomainSocketEndpointDirectory.scala @@ -0,0 +1,157 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.udf.worker.core.direct + +import java.io.IOException +import java.nio.charset.StandardCharsets +import java.nio.file.{Files, FileVisitResult, Path, Paths, SimpleFileVisitor} +import java.nio.file.attribute.{BasicFileAttributes, PosixFilePermissions} +import java.util.Locale + +import org.apache.spark.udf.worker.core.WorkerLogger + +/** + * Owns the filesystem lifecycle for direct-worker Unix domain sockets. + * + * The directory is private to its owner and anchored under a short temp path + * so generated socket addresses fit the platform `sockaddr_un.sun_path` limit. + * Individual socket files are created by workers and removed through + * [[cleanupEndpointAddress]]; [[close]] recursively removes the directory. + */ +private[worker] final class UnixDomainSocketEndpointDirectory( + logger: WorkerLogger, + osName: String = System.getProperty("os.name", "")) extends AutoCloseable { + + import UnixDomainSocketEndpointDirectory._ + + private val socketDir: Path = createPrivateTempDirectory() + private val maxPathBytes: Int = maxPathBytesForOs(osName) + + /** Returns a collision-checked socket path for `workerId`. */ + def newEndpointAddress(workerId: String): String = { + val shortId = workerId.replace("-", "").take(WORKER_ID_PREFIX_LENGTH) + var candidate = socketDir.resolve(s"w-$shortId.sock") + var suffix = 0 + while (Files.exists(candidate) && suffix < MAX_SOCKET_LEAF_RETRIES) { + suffix += 1 + candidate = socketDir.resolve(s"w-$shortId-$suffix.sock") + } + if (Files.exists(candidate)) { + throw new IllegalStateException( + s"could not allocate a free UDS path under $socketDir after " + + s"$MAX_SOCKET_LEAF_RETRIES retries (truncated id=$shortId)") + } + val address = candidate.toString + val pathBytes = address.getBytes(StandardCharsets.UTF_8).length + if (pathBytes > maxPathBytes) { + throw new IllegalStateException( + s"UDS path requires $pathBytes UTF-8 bytes but this platform allows " + + s"at most $maxPathBytes: $address") + } + address + } + + /** Removes the socket file created by a worker. */ + def cleanupEndpointAddress(address: String): Unit = { + Files.deleteIfExists(Paths.get(address)) + } + + /** Recursively removes all endpoint artifacts and the private directory. */ + override def close(): Unit = { + if (!Files.exists(socketDir)) return + var firstFailure: Option[IOException] = None + def recordFailure(e: IOException): Unit = firstFailure match { + case Some(first) => first.addSuppressed(e) + case None => firstFailure = Some(e) + } + def delete(path: Path): Unit = { + try Files.deleteIfExists(path) catch { case e: IOException => recordFailure(e) } + } + Files.walkFileTree(socketDir, new SimpleFileVisitor[Path] { + override def visitFile(file: Path, attrs: BasicFileAttributes): FileVisitResult = { + delete(file) + FileVisitResult.CONTINUE + } + override def visitFileFailed(file: Path, exc: IOException): FileVisitResult = { + recordFailure(exc) + FileVisitResult.CONTINUE + } + override def postVisitDirectory(dir: Path, exc: IOException): FileVisitResult = { + if (exc != null) recordFailure(exc) + delete(dir) + FileVisitResult.CONTINUE + } + }) + firstFailure.foreach(e => throw e) + } + + private def createPrivateTempDirectory(): Path = { + val attr = PosixFilePermissions.asFileAttribute( + PosixFilePermissions.fromString("rwx------")) + val base = shortTempBase() + try { + Files.createTempDirectory(base, "spark-udf-worker", attr) + } catch { + case _: UnsupportedOperationException => + val dir = Files.createTempDirectory(base, "spark-udf-worker") + val file = dir.toFile + // Non-short-circuiting Boolean AND ensures every setter runs. + val applied = + file.setReadable(false, false) & file.setWritable(false, false) & + file.setExecutable(false, false) & file.setReadable(true, true) & + file.setWritable(true, true) & file.setExecutable(true, true) + if (!applied) { + logger.warn( + s"Could not fully restrict permissions on $dir; socket " + + s"directory may be accessible to other local users on this " + + s"filesystem") + } + dir + } + } + + private def shortTempBase(): Path = { + val candidates = + Seq("TMPDIR", "TEMP", "TMP").flatMap(sys.env.get).filter(_.nonEmpty) :+ "/tmp" + val usable = candidates + .map(Paths.get(_)) + .filter(path => Files.isDirectory(path) && Files.isWritable(path)) + def encodedLength(path: Path): Int = + path.toString.getBytes(StandardCharsets.UTF_8).length + usable + .sortBy(encodedLength) + .headOption + .getOrElse(Paths.get(System.getProperty("java.io.tmpdir"))) + } +} + +private[worker] object UnixDomainSocketEndpointDirectory { + private val MAX_SOCKET_LEAF_RETRIES = 16 + private val WORKER_ID_PREFIX_LENGTH = 16 + private val LINUX_PATH_BYTES = 107 + private val KQUEUE_PATH_BYTES = 103 + + private[worker] def maxPathBytesForOs(osName: String): Int = { + val normalized = osName.toLowerCase(Locale.ROOT) + if (normalized.contains("mac") || normalized.contains("darwin") || + normalized.contains("bsd")) { + KQUEUE_PATH_BYTES + } else { + LINUX_PATH_BYTES + } + } +} diff --git a/udf/worker/core/src/test/scala/org/apache/spark/udf/worker/core/DirectWorkerDispatcherSuite.scala b/udf/worker/core/src/test/scala/org/apache/spark/udf/worker/core/DirectWorkerDispatcherSuite.scala index 397a7b7dcab4f..dec934da39005 100644 --- a/udf/worker/core/src/test/scala/org/apache/spark/udf/worker/core/DirectWorkerDispatcherSuite.scala +++ b/udf/worker/core/src/test/scala/org/apache/spark/udf/worker/core/DirectWorkerDispatcherSuite.scala @@ -33,7 +33,7 @@ import org.apache.spark.udf.worker.{ import org.apache.spark.udf.worker.core.{WorkerConnection, WorkerHandle, WorkerSecurityScope, WorkerSession} import org.apache.spark.udf.worker.core.direct.{DirectWorkerException, DirectWorkerProcess, - DirectWorkerTimeoutException} + DirectWorkerTimeoutException, UnixDomainSocketEndpointDirectory} /** * Tests for [[DirectWorkerDispatcher]] process lifecycle: spawning workers @@ -135,6 +135,16 @@ class DirectWorkerDispatcherSuite s"Expected SocketFileConnection, got ${other.getClass.getSimpleName}") } + private def awaitThreadWaiting(thread: Thread, description: String): Unit = { + val deadline = System.nanoTime() + java.util.concurrent.TimeUnit.SECONDS.toNanos(10) + while (thread.isAlive && thread.getState != Thread.State.WAITING && + System.nanoTime() < deadline) { + Thread.sleep(10) + } + assert(thread.getState == Thread.State.WAITING, + s"$description did not enter the creation barrier; state=${thread.getState}") + } + test("creates a worker and session") { dispatcher = new TestDirectWorkerDispatcher(specWithRunner(defaultRunner)) @@ -329,7 +339,6 @@ class DirectWorkerDispatcherSuite // workers or dispatcher-level transport state. val readyLatch = new java.util.concurrent.CountDownLatch(1) val releaseLatch = new java.util.concurrent.CountDownLatch(1) - val closeStartedLatch = new java.util.concurrent.CountDownLatch(1) val capturedWorkers = new java.util.concurrent.ConcurrentLinkedQueue[DirectWorkerProcess]() val transportClosed = new java.util.concurrent.atomic.AtomicBoolean(false) @@ -348,8 +357,6 @@ class DirectWorkerDispatcherSuite transportClosed.set(true) super.closeTransport() } - - override protected def afterCloseStarted(): Unit = closeStartedLatch.countDown() } try { val outcome = @@ -371,8 +378,7 @@ class DirectWorkerDispatcherSuite val closeThread = new Thread(() => racing.close(), "close-racer") closeThread.start() - assert(closeStartedLatch.await(10, java.util.concurrent.TimeUnit.SECONDS), - "close thread never flipped the dispatcher to closed") + awaitThreadWaiting(closeThread, "close thread") assert(closeThread.isAlive, "close should wait for the in-flight createSession") assert(!transportClosed.get(), "transport cleanup must not run while createSession is in flight") @@ -416,6 +422,87 @@ class DirectWorkerDispatcherSuite } } + test("interrupted close completes cleanup before restoring the interrupt") { + val cleanupMarker = Files.createTempFile("interrupted-close-cleanup", ".txt").toFile + cleanupMarker.delete() + val env = WorkerEnvironment.newBuilder() + .setEnvironmentCleanup(ProcessCallable.newBuilder() + .addCommand("bash").addCommand("-c") + .addCommand(s"sleep 0.2; touch ${cleanupMarker.getAbsolutePath}").build()) + .build() + val workerRegistered = new java.util.concurrent.CountDownLatch(1) + val releaseCreation = new java.util.concurrent.CountDownLatch(1) + val racing = new TestDirectWorkerDispatcher(specWithEnv(env = env)) { + override protected def afterWorkerRegistered(worker: DirectWorkerProcess): Unit = { + workerRegistered.countDown() + if (!releaseCreation.await(30, java.util.concurrent.TimeUnit.SECONDS)) { + fail("releaseCreation never fired -- test orchestration broken") + } + } + } + + val createOutcome = + new java.util.concurrent.atomic.AtomicReference[Either[Throwable, WorkerSession]]() + val closeError = new java.util.concurrent.atomic.AtomicReference[Throwable]() + val closeWasInterrupted = new java.util.concurrent.atomic.AtomicBoolean(false) + val createThread = new Thread(() => { + try { + createOutcome.set(Right(racing.createSession(None))) + } catch { + case t: Throwable => createOutcome.set(Left(t)) + } + }, "interrupted-close-create") + val closeThread = new Thread(() => { + try { + racing.close() + } catch { + case t: Throwable => closeError.set(t) + } finally { + closeWasInterrupted.set(Thread.currentThread().isInterrupted) + } + }, "interrupted-close") + + try { + createThread.start() + assert(workerRegistered.await(10, java.util.concurrent.TimeUnit.SECONDS), + "createSession thread never registered its worker") + + closeThread.start() + awaitThreadWaiting(closeThread, "interrupted close thread") + closeThread.interrupt() + assert(closeThread.isAlive, "close should keep waiting after interruption") + assert(!cleanupMarker.exists(), + "environment cleanup must not run while createSession is in flight") + + releaseCreation.countDown() + createThread.join(10000) + closeThread.join(10000) + assert(!createThread.isAlive, "createSession thread did not finish") + assert(!closeThread.isAlive, "close thread did not finish") + assert(closeError.get() == null, + s"close should not fail after interruption: ${closeError.get()}") + assert(cleanupMarker.exists(), "environment cleanup should complete before close returns") + assert(closeWasInterrupted.get(), "close should restore the interrupt flag after cleanup") + createOutcome.get() match { + case Left(e: IllegalStateException) => + assert(e.getMessage.contains("closed"), + s"expected dispatcher-closed error, got: ${e.getMessage}") + case Left(other) => fail(s"unexpected exception from createSession: $other") + case Right(_) => fail("createSession should fail when close wins") + } + } finally { + releaseCreation.countDown() + racing.close() + cleanupMarker.delete() + } + } + + test("uses the kqueue UDS path limit on BSD") { + assert(UnixDomainSocketEndpointDirectory.maxPathBytesForOs("Linux") == 107) + assert(UnixDomainSocketEndpointDirectory.maxPathBytesForOs("Mac OS X") == 103) + assert(UnixDomainSocketEndpointDirectory.maxPathBytesForOs("FreeBSD") == 103) + } + test("worker-provided graceful timeout is capped at the engine-side maximum") { // The proto documents an engine-configurable maximum (fixed at 30s today). // A 60s spec value should be clamped down. diff --git a/udf/worker/core/src/test/scala/org/apache/spark/udf/worker/core/TestDirectWorkerHelpers.scala b/udf/worker/core/src/test/scala/org/apache/spark/udf/worker/core/TestDirectWorkerHelpers.scala index 01fee8761aee6..32e5001160285 100644 --- a/udf/worker/core/src/test/scala/org/apache/spark/udf/worker/core/TestDirectWorkerHelpers.scala +++ b/udf/worker/core/src/test/scala/org/apache/spark/udf/worker/core/TestDirectWorkerHelpers.scala @@ -16,16 +16,12 @@ */ package org.apache.spark.udf.worker.core -import java.io.{File, IOException} -import java.nio.charset.StandardCharsets -import java.nio.file.{Files, FileVisitResult, Path, Paths, SimpleFileVisitor} -import java.nio.file.attribute.{BasicFileAttributes, PosixFilePermissions} -import java.util.Locale +import java.io.File 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} + DirectWorkerTimeoutException, UnixDomainSocketEndpointDirectory} import org.apache.spark.udf.worker.core.direct.DirectWorkerDispatcher.READY_POLL_INTERVAL_MS /** @@ -85,52 +81,15 @@ class TestDirectWorkerDispatcher( logger: WorkerLogger = WorkerLogger.NoOp) extends DirectWorkerDispatcher(workerSpec, logger) { - // Upper bound on the rename-on-collision retry loop in newEndpointAddress. - private val MAX_SOCKET_LEAF_RETRIES = 16 - - private val MAX_UDS_PATH_BYTES = { - val osName = System.getProperty("os.name", "").toLowerCase(Locale.ROOT) - if (osName.contains("mac") || osName.contains("darwin")) 103 else 107 - } - - // The private 0700 socket directory, created in [[initialize]] (after the - // base class has validated the spec) and removed in [[closeTransport]]. - private lazy val socketDir: Path = createPrivateTempDirectory() + private lazy val endpointDirectory = new UnixDomainSocketEndpointDirectory(logger) override protected def initialize(): Unit = { super.initialize() - // Force the lazy val now so the directory is created (and any failure - // surfaces) at construction time, after spec validation has passed. - socketDir + endpointDirectory } - /** - * Returns the UDS path the worker should bind. Uses a short 16-hex-char leaf - * (the worker's full UUID still travels via `--id`) to stay within the - * 108-byte UDS `sun_path` limit. - */ - override protected def newEndpointAddress(workerId: String): String = { - val short = workerId.replace("-", "").take(16) - var candidate = socketDir.resolve(s"w-$short.sock") - var suffix = 0 - while (Files.exists(candidate) && suffix < MAX_SOCKET_LEAF_RETRIES) { - suffix += 1 - candidate = socketDir.resolve(s"w-$short-$suffix.sock") - } - if (Files.exists(candidate)) { - throw new IllegalStateException( - s"could not allocate a free UDS path under $socketDir after " + - s"$MAX_SOCKET_LEAF_RETRIES retries (truncated id=$short)") - } - val address = candidate.toString - val pathBytes = address.getBytes(StandardCharsets.UTF_8).length - if (pathBytes > MAX_UDS_PATH_BYTES) { - throw new IllegalStateException( - s"UDS path requires $pathBytes UTF-8 bytes but this platform allows " + - s"at most $MAX_UDS_PATH_BYTES: $address") - } - address - } + override protected def newEndpointAddress(workerId: String): String = + endpointDirectory.newEndpointAddress(workerId) override protected def connectWorker( address: String, @@ -162,39 +121,10 @@ class TestDirectWorkerDispatcher( newConnection(address) } - override protected def cleanupEndpointAddress(address: String): Unit = { - Files.deleteIfExists(new File(address).toPath) - } + override protected def cleanupEndpointAddress(address: String): Unit = + endpointDirectory.cleanupEndpointAddress(address) - override protected def closeTransport(): Unit = { - if (!Files.exists(socketDir)) return - var firstFailure: IOException = null - def recordFailure(e: IOException): Unit = { - if (firstFailure == null) firstFailure = e else firstFailure.addSuppressed(e) - } - def delete(path: Path): Unit = { - try Files.deleteIfExists(path) catch { case e: IOException => recordFailure(e) } - } - // Recursive post-order delete: today socketDir contains only socket files - // at the top level, but a future change that namespaces workers into - // subdirectories should not silently leak them. - Files.walkFileTree(socketDir, new SimpleFileVisitor[Path] { - override def visitFile(file: Path, attrs: BasicFileAttributes): FileVisitResult = { - delete(file) - FileVisitResult.CONTINUE - } - override def visitFileFailed(file: Path, exc: IOException): FileVisitResult = { - recordFailure(exc) - FileVisitResult.CONTINUE - } - override def postVisitDirectory(dir: Path, exc: IOException): FileVisitResult = { - if (exc != null) recordFailure(exc) - delete(dir) - FileVisitResult.CONTINUE - } - }) - if (firstFailure != null) throw firstFailure - } + override protected def closeTransport(): Unit = endpointDirectory.close() // `spec` is the same object as the `workerSpec` field but passed explicitly: // at the point this runs (parent constructor body), `this` is only partially @@ -236,71 +166,4 @@ class TestDirectWorkerDispatcher( s"Worker exited with code ${process.exitValue()} " + s"before creating socket at $address\n$tail") } - - /** - * Creates a private (owner-only, 0700) temp directory for worker sockets. - * On POSIX filesystems the permissions are applied atomically at creation; - * the non-POSIX branch tightens best-effort after creation. - * - * The directory is anchored under a short base path (see [[shortTempBase]]) - * rather than the configured `java.io.tmpdir`. Builds point the latter at a - * deep location (e.g. `/target/tmp`), which combined with the - * generated socket leaf would push the worker's Unix-domain socket path past - * the platform `sun_path` limit (108 bytes on Linux, 104 on macOS) and fail - * with "AF_UNIX path too long". - */ - private def createPrivateTempDirectory(): Path = { - val attr = PosixFilePermissions.asFileAttribute( - PosixFilePermissions.fromString("rwx------")) - val base = shortTempBase() - try { - Files.createTempDirectory(base, "spark-udf-worker", attr) - } catch { - case _: UnsupportedOperationException => - val dir = Files.createTempDirectory(base, "spark-udf-worker") - val f = dir.toFile - // Bit-wise AND (NOT &&): all six setters must run even if an earlier - // one returns false, so the final permission state matches owner-only. - val applied = - f.setReadable(false, false) & f.setWritable(false, false) & - f.setExecutable(false, false) & f.setReadable(true, true) & - f.setWritable(true, true) & f.setExecutable(true, true) - if (!applied) { - logger.warn( - s"Could not fully restrict permissions on $dir; socket " + - s"directory may be accessible to other local users on this " + - s"filesystem") - } - dir - } - } - - /** - * Picks the shortest usable base directory for the socket temp dir so the - * resulting Unix-domain socket path stays within the platform `sun_path` - * limit (108 bytes on Linux, 104 on macOS). - * - * This mirrors how PySpark already chooses its UDS directory in - * `python/run-tests.py` (`spark.python.unix.domain.socket.dir`): prefer the - * OS temp-dir env vars (`TMPDIR`/`TEMP`/`TMP`) and fall back to `/tmp`. Builds - * point `java.io.tmpdir` at a deep `/target/tmp`, which combined with - * the generated socket leaf would overflow `sun_path`; the OS temp dir is - * short (e.g. a per-user `/tmp/...` on Linux runners, `/var/folders/...` via - * `$TMPDIR` on macOS). `java.io.tmpdir` remains the last-resort fallback. - */ - private def shortTempBase(): Path = { - val candidates = - Seq("TMPDIR", "TEMP", "TMP").flatMap(sys.env.get).filter(_.nonEmpty) :+ "/tmp" - val usable = candidates - .map(Paths.get(_)) - .filter(p => Files.isDirectory(p) && Files.isWritable(p)) - // Netty passes the unresolved pathname spelling to the native bind call. - // Rank that exact spelling in UTF-8 bytes, then validate the complete - // generated path in newEndpointAddress. - def encodedLength(p: Path): Int = p.toString.getBytes(StandardCharsets.UTF_8).length - usable - .sortBy(encodedLength) - .headOption - .getOrElse(Paths.get(System.getProperty("java.io.tmpdir"))) - } } diff --git a/udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/DirectGrpcDispatcher.scala b/udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/DirectGrpcDispatcher.scala index c35166535d4ac..4b33366ef73d1 100644 --- a/udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/DirectGrpcDispatcher.scala +++ b/udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/DirectGrpcDispatcher.scala @@ -16,11 +16,8 @@ */ package org.apache.spark.udf.worker.grpc -import java.io.{File, IOException} -import java.nio.charset.StandardCharsets -import java.nio.file.{Files, FileVisitResult, Path, Paths, SimpleFileVisitor} -import java.nio.file.attribute.{BasicFileAttributes, PosixFilePermissions} -import java.util.Locale +import java.io.File +import java.nio.file.{Files, Paths} import java.util.concurrent.{CountDownLatch, TimeUnit} import scala.util.control.NonFatal @@ -33,7 +30,7 @@ import org.apache.spark.udf.worker.{UDFProtoCommunicationPattern, UDFWorkerSpeci import org.apache.spark.udf.worker.core.{WorkerConnection, WorkerHandle, WorkerLogger, WorkerSession} import org.apache.spark.udf.worker.core.direct.{DirectWorkerDispatcher, DirectWorkerException, - DirectWorkerTimeoutException} + DirectWorkerTimeoutException, UnixDomainSocketEndpointDirectory} import org.apache.spark.udf.worker.core.direct.DirectWorkerDispatcher.READY_POLL_INTERVAL_MS /** @@ -45,74 +42,24 @@ import org.apache.spark.udf.worker.core.direct.DirectWorkerDispatcher.READY_POLL * * @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. */ @Experimental class DirectGrpcDispatcher( workerSpec: UDFWorkerSpecification, - logger: WorkerLogger = WorkerLogger.NoOp, - grpcMaxInboundMessageSize: Int = GrpcWorkerChannel.DEFAULT_MAX_INBOUND_MESSAGE_SIZE) + logger: WorkerLogger = WorkerLogger.NoOp) 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 - - // sockaddr_un.sun_path includes its terminating NUL. Linux reserves 108 - // bytes and macOS 104, leaving at most 107 or 103 pathname bytes. - private val MAX_UDS_PATH_BYTES = { - val osName = System.getProperty("os.name", "").toLowerCase(Locale.ROOT) - if (osName.contains("mac") || osName.contains("darwin")) 103 else 107 - } - - // 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() + // Laziness avoids reading subclass state during the base constructor's + // validation phase. initialize() forces deterministic allocation afterward. + private lazy val endpointDirectory = new UnixDomainSocketEndpointDirectory(logger) override protected def initialize(): Unit = { super.initialize() - // Force the lazy val now so the directory is created (and any failure - // surfaces) at construction time, after spec validation has passed. - socketDir + endpointDirectory } - /** - * 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: 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) - var candidate = socketDir.resolve(s"w-$short.sock") - var suffix = 0 - while (Files.exists(candidate) && suffix < MAX_SOCKET_LEAF_RETRIES) { - suffix += 1 - candidate = socketDir.resolve(s"w-$short-$suffix.sock") - } - if (Files.exists(candidate)) { - throw new IllegalStateException( - s"could not allocate a free UDS path under $socketDir after " + - s"$MAX_SOCKET_LEAF_RETRIES retries (truncated id=$short)") - } - val address = candidate.toString - val pathBytes = address.getBytes(StandardCharsets.UTF_8).length - if (pathBytes > MAX_UDS_PATH_BYTES) { - throw new IllegalStateException( - s"UDS path requires $pathBytes UTF-8 bytes but this platform allows " + - s"at most $MAX_UDS_PATH_BYTES: $address") - } - address - } + override protected def newEndpointAddress(workerId: String): String = + endpointDirectory.newEndpointAddress(workerId) override protected def connectWorker( address: String, @@ -145,7 +92,6 @@ class DirectGrpcDispatcher( s"GrpcWorkerChannel but got ${other.getClass.getName}") } 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) @@ -157,13 +103,12 @@ class DirectGrpcDispatcher( 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))) { + // A failed connection attempt enters gRPC backoff. While the worker's + // fresh socket exists, reset each observed backoff so a server that is + // still starting does not consume the initialization budget. + if (DirectGrpcDispatcher.shouldResetConnectBackoff( + state, Files.exists(Paths.get(address)))) { grpc.channel.resetConnectBackoff() - socketTriggeredReconnect = true } val stateChanged = new CountDownLatch(1) grpc.channel.notifyWhenStateChanged(state, () => stateChanged.countDown()) @@ -173,39 +118,10 @@ class DirectGrpcDispatcher( } } - override protected def cleanupEndpointAddress(address: String): Unit = { - Files.deleteIfExists(new File(address).toPath) - } + override protected def cleanupEndpointAddress(address: String): Unit = + endpointDirectory.cleanupEndpointAddress(address) - override protected def closeTransport(): Unit = { - if (!Files.exists(socketDir)) return - var firstFailure: IOException = null - def recordFailure(e: IOException): Unit = { - if (firstFailure == null) firstFailure = e else firstFailure.addSuppressed(e) - } - def delete(path: Path): Unit = { - try Files.deleteIfExists(path) catch { case e: IOException => recordFailure(e) } - } - // Recursive post-order delete: today socketDir contains only socket files - // at the top level, but a future change that namespaces workers into - // subdirectories should not silently leak them. - Files.walkFileTree(socketDir, new SimpleFileVisitor[Path] { - override def visitFile(file: Path, attrs: BasicFileAttributes): FileVisitResult = { - delete(file) - FileVisitResult.CONTINUE - } - override def visitFileFailed(file: Path, exc: IOException): FileVisitResult = { - recordFailure(exc) - FileVisitResult.CONTINUE - } - override def postVisitDirectory(dir: Path, exc: IOException): FileVisitResult = { - if (exc != null) recordFailure(exc) - delete(dir) - FileVisitResult.CONTINUE - } - }) - if (firstFailure != null) throw firstFailure - } + override protected def closeTransport(): Unit = endpointDirectory.close() // `spec` is the same object as the `workerSpec` field but passed // explicitly: at the point this runs (parent constructor body), `this` @@ -238,8 +154,7 @@ class DirectGrpcDispatcher( } protected def newConnection(address: String): WorkerConnection = - new GrpcWorkerChannel( - address, logger, maxInboundMessageSize = grpcMaxInboundMessageSize) + new GrpcWorkerChannel(address, logger) override protected def newSession( workerHandle: WorkerHandle, @@ -269,80 +184,11 @@ class DirectGrpcDispatcher( } } - /** - * Creates a private (owner-only, 0700) temp directory for worker sockets. - * - * 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 - * deep location (e.g. `/target/tmp`), which combined with the - * generated socket leaf would push the worker's Unix-domain socket path past - * the platform `sun_path` limit (108 bytes on Linux, 104 on macOS) and fail - * with "AF_UNIX path too long". - */ - private def createPrivateTempDirectory(): Path = { - val attr = PosixFilePermissions.asFileAttribute( - PosixFilePermissions.fromString("rwx------")) - val base = shortTempBase() - try { - Files.createTempDirectory(base, "spark-udf-worker", attr) - } catch { - case _: UnsupportedOperationException => - val dir = Files.createTempDirectory(base, "spark-udf-worker") - val f = dir.toFile - // Bit-wise AND (NOT &&): all six setters must run even if an earlier - // one returns false, so the final permission state matches owner-only. - // && 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) & - f.setWritable(true, true) & f.setExecutable(true, true) - if (!applied) { - logger.warn( - s"Could not fully restrict permissions on $dir; socket " + - s"directory may be accessible to other local users on this " + - s"filesystem") - } - dir - } - } +} - /** - * Picks the shortest usable base directory for the socket temp dir so the - * resulting Unix-domain socket path stays within the platform `sun_path` - * limit (108 bytes on Linux, 104 on macOS). - * - * This mirrors how PySpark already chooses its UDS directory in - * `python/run-tests.py` (`spark.python.unix.domain.socket.dir`): prefer the - * OS temp-dir env vars (`TMPDIR`/`TEMP`/`TMP`) and fall back to `/tmp`. Builds - * point `java.io.tmpdir` at a deep `/target/tmp`, which combined with - * the generated socket leaf would overflow `sun_path`; the OS temp dir is - * short (e.g. a per-user `/tmp/...` on Linux runners, `/var/folders/...` via - * `$TMPDIR` on macOS). `java.io.tmpdir` remains the last-resort fallback. - */ - private def shortTempBase(): Path = { - val candidates = - Seq("TMPDIR", "TEMP", "TMP").flatMap(sys.env.get).filter(_.nonEmpty) :+ "/tmp" - val usable = candidates - .map(Paths.get(_)) - .filter(p => Files.isDirectory(p) && Files.isWritable(p)) - // Netty passes the unresolved pathname spelling to the native bind call. - // Rank that exact spelling in UTF-8 bytes, then validate the complete - // generated path in newEndpointAddress. - def encodedLength(p: Path): Int = p.toString.getBytes(StandardCharsets.UTF_8).length - usable - .sortBy(encodedLength) - .headOption - .getOrElse(Paths.get(System.getProperty("java.io.tmpdir"))) - } +private[grpc] object DirectGrpcDispatcher { + private[grpc] def shouldResetConnectBackoff( + state: ConnectivityState, + endpointExists: Boolean): Boolean = + endpointExists && state == ConnectivityState.TRANSIENT_FAILURE } 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 index 07ae58415f3d0..220a31d633312 100644 --- 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 @@ -45,9 +45,9 @@ import org.apache.spark.udf.worker.core.{WorkerConnection, WorkerLogger} * 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. + * - [[isActive]] is true until this wrapper is closed and while the channel + * reports a 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 @@ -57,26 +57,13 @@ import org.apache.spark.udf.worker.core.{WorkerConnection, WorkerLogger} * * @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) + logger: WorkerLogger = WorkerLogger.NoOp) extends WorkerConnection { - require(maxInboundMessageSize > 0, "maxInboundMessageSize must be positive") - private val transport: UnixDomainSocketTransport = UnixDomainSocketTransport.detect() private val eventLoopGroup: EventLoopGroup = transport.newEventLoopGroup() @@ -89,12 +76,13 @@ class GrpcWorkerChannel( .forAddress(new DomainSocketAddress(socketPath)) .channelType(transport.channelType) .eventLoopGroup(eventLoopGroup) - .maxInboundMessageSize(maxInboundMessageSize) + .maxInboundMessageSize(GrpcWorkerChannel.MAX_INBOUND_MESSAGE_SIZE) .usePlaintext() .build() } catch { case NonFatal(e) => - UnixDomainSocketTransport.shutdown(eventLoopGroup, shutdownPhaseTimeoutMs) + UnixDomainSocketTransport.shutdown( + eventLoopGroup, GrpcWorkerChannel.SHUTDOWN_PHASE_TIMEOUT_MS) throw e } @@ -105,6 +93,8 @@ class GrpcWorkerChannel( channel.getState(false) != ConnectivityState.SHUTDOWN } + private[grpc] def isEventLoopTerminated: Boolean = eventLoopGroup.isTerminated + /** * 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 @@ -115,9 +105,11 @@ class GrpcWorkerChannel( try { channel.shutdown() - if (!channel.awaitTermination(shutdownPhaseTimeoutMs, TimeUnit.MILLISECONDS)) { + if (!channel.awaitTermination( + GrpcWorkerChannel.SHUTDOWN_PHASE_TIMEOUT_MS, TimeUnit.MILLISECONDS)) { channel.shutdownNow() - channel.awaitTermination(shutdownPhaseTimeoutMs, TimeUnit.MILLISECONDS) + channel.awaitTermination( + GrpcWorkerChannel.SHUTDOWN_PHASE_TIMEOUT_MS, TimeUnit.MILLISECONDS) } } catch { case _: InterruptedException => @@ -127,7 +119,8 @@ class GrpcWorkerChannel( } try { - UnixDomainSocketTransport.shutdown(eventLoopGroup, shutdownPhaseTimeoutMs) + UnixDomainSocketTransport.shutdown( + eventLoopGroup, GrpcWorkerChannel.SHUTDOWN_PHASE_TIMEOUT_MS) } catch { case NonFatal(e) => logger.warn(s"Error shutting down event loop for $socketPath", e) @@ -137,12 +130,7 @@ class GrpcWorkerChannel( object GrpcWorkerChannel { /** Matches Spark Connect's default and accommodates normal Arrow batches. */ - val DEFAULT_MAX_INBOUND_MESSAGE_SIZE: Int = 128 * 1024 * 1024 + private[grpc] val 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 + private val 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 index ae2e2850d7ef4..d07ab0c320def 100644 --- 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 @@ -32,7 +32,8 @@ import io.netty.channel.{Channel, EventLoopGroup, MultiThreadIoEventLoopGroup, S * * 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 + * kernels; it could be added as a preferred Linux option in + * [[UnixDomainSocketTransport.detect]] once its * native artifact is on the classpath, but is intentionally left out for now to * keep the dependency set minimal. * 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 index 41d3bb4d46b3c..dc82bcfe205e7 100644 --- 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 @@ -28,7 +28,7 @@ import org.apache.spark.udf.worker.{Cancel, DataRequest, DirectWorker, Finish, I ProcessCallable, UdfPayload, UDFProtoCommunicationPattern, UDFWorkerDataFormat, UDFWorkerProperties, UDFWorkerSpecification, UnixDomainSocket, WorkerCapabilities, WorkerConnectionSpec} -import org.apache.spark.udf.worker.core.WorkerSession +import org.apache.spark.udf.worker.core.{WorkerConnection, WorkerSession} import org.apache.spark.udf.worker.core.direct.{DirectWorkerProcess, DirectWorkerTimeoutException} import org.apache.spark.udf.worker.grpc.testing.EchoGrpcWorkerMain @@ -164,12 +164,36 @@ class DirectGrpcDispatcherIntegrationSuite .addCommand(socketOnlyWorker) .addCommand("--") .build() - dispatcher = new DirectGrpcDispatcher(workerSpec(runner, initTimeoutMs = 1000)) + var failedProcess: Process = null + var failedConnection: GrpcWorkerChannel = null + dispatcher = new DirectGrpcDispatcher(workerSpec(runner, initTimeoutMs = 1000)) { + override protected def connectWorker( + address: String, + process: Process, + outputFile: File): WorkerConnection = { + failedProcess = process + super.connectWorker(address, process, outputFile) + } + + override protected def newConnection(address: String): WorkerConnection = { + val connection = super.newConnection(address) + failedConnection = connection.asInstanceOf[GrpcWorkerChannel] + connection + } + } val error = intercept[DirectWorkerTimeoutException] { dispatcher.createSession(None) } assert(error.getMessage.contains("did not become reachable")) assert(error.getMessage.contains("socket-created")) + assert(failedProcess != null, "the failed worker process should have been captured") + assert(!failedProcess.isAlive, "the failed worker process should be reaped") + assert(failedConnection != null, "the failed worker connection should have been captured") + assert(failedConnection.channel.isShutdown, "the failed gRPC channel should be shut down") + assert(failedConnection.isEventLoopTerminated, + "the failed gRPC channel's event loop should be terminated") + assert(!new File(failedConnection.socketPath).exists(), + "the failed worker socket should be removed") } } diff --git a/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/DirectGrpcDispatcherSuite.scala b/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/DirectGrpcDispatcherSuite.scala new file mode 100644 index 0000000000000..70640c720d5a3 --- /dev/null +++ b/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/DirectGrpcDispatcherSuite.scala @@ -0,0 +1,92 @@ +/* + * 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 io.grpc.ConnectivityState +// scalastyle:off funsuite +import org.scalatest.funsuite.AnyFunSuite + +import org.apache.spark.udf.worker.{DirectWorker, LocalTcpConnection, + UDFProtoCommunicationPattern, UDFWorkerProperties, UDFWorkerSpecification, + UnixDomainSocket, WorkerCapabilities, WorkerConnectionSpec} + +class DirectGrpcDispatcherSuite extends AnyFunSuite { +// scalastyle:on funsuite + + private def directSpec( + properties: UDFWorkerProperties, + capabilities: Option[WorkerCapabilities] = None): UDFWorkerSpecification = { + val builder = UDFWorkerSpecification.newBuilder() + .setDirect(DirectWorker.newBuilder().setProperties(properties).build()) + capabilities.foreach(builder.setCapabilities) + builder.build() + } + + private def udsProperties: UDFWorkerProperties = UDFWorkerProperties.newBuilder() + .setConnection(WorkerConnectionSpec.newBuilder() + .setUnixDomainSocket(UnixDomainSocket.getDefaultInstance) + .build()) + .build() + + test("requires a connection") { + val error = intercept[IllegalArgumentException] { + new DirectGrpcDispatcher(directSpec(UDFWorkerProperties.getDefaultInstance)) + } + assert(error.getMessage.contains("connection must be set")) + } + + test("requires Unix domain socket transport") { + val properties = UDFWorkerProperties.newBuilder() + .setConnection(WorkerConnectionSpec.newBuilder() + .setTcp(LocalTcpConnection.getDefaultInstance) + .build()) + .build() + val error = intercept[IllegalArgumentException] { + new DirectGrpcDispatcher(directSpec(properties)) + } + assert(error.getMessage.contains("requires UNIX domain socket transport")) + } + + test("requires capabilities") { + val error = intercept[IllegalArgumentException] { + new DirectGrpcDispatcher(directSpec(udsProperties)) + } + assert(error.getMessage.contains("requires WorkerCapabilities")) + } + + test("requires bidirectional streaming capability") { + val capabilities = WorkerCapabilities.newBuilder() + .addSupportedCommunicationPatterns( + UDFProtoCommunicationPattern.UDF_PROTO_COMMUNICATION_PATTERN_UNSPECIFIED) + .build() + val error = intercept[IllegalArgumentException] { + new DirectGrpcDispatcher(directSpec(udsProperties, Some(capabilities))) + } + assert(error.getMessage.contains("requires BIDIRECTIONAL_STREAMING")) + } + + test("requests a backoff reset for every transient failure after endpoint creation") { + assert((1 to 2).forall { _ => + DirectGrpcDispatcher.shouldResetConnectBackoff( + ConnectivityState.TRANSIENT_FAILURE, endpointExists = true) + }) + assert(!DirectGrpcDispatcher.shouldResetConnectBackoff( + ConnectivityState.CONNECTING, endpointExists = true)) + assert(!DirectGrpcDispatcher.shouldResetConnectBackoff( + ConnectivityState.TRANSIENT_FAILURE, endpointExists = false)) + } +} 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 index e78f02908eae4..b48eb57f361c9 100644 --- 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 @@ -19,6 +19,7 @@ package org.apache.spark.udf.worker.grpc.testing import java.io.IOException import java.nio.file.{Files, Paths} import java.util.concurrent.TimeUnit +import java.util.concurrent.atomic.AtomicBoolean import scala.util.control.NonFatal @@ -54,81 +55,91 @@ object EchoGrpcWorkerMain { 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) + val cleanupStarted = new AtomicBoolean(false) + var bossElg: EventLoopGroup = null + var workerElg: EventLoopGroup = null + var server: Server = null + var shutdownHook: Thread = null - // 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)") + def cleanup(): Unit = { + if (!cleanupStarted.compareAndSet(false, true)) return try { - server.shutdown() - if (!server.awaitTermination(5, TimeUnit.SECONDS)) { - server.shutdownNow() + if (server != null) { + server.shutdown() + if (!server.awaitTermination(5, TimeUnit.SECONDS)) { + server.shutdownNow() + } } } catch { case _: InterruptedException => Thread.currentThread().interrupt() + case NonFatal(e) => stderrln(s"failed to shut down gRPC server: ${e.getMessage}") } finally { - UnixDomainSocketTransport.shutdown(bossElg, 5000L) - UnixDomainSocketTransport.shutdown(workerElg, 5000L) + try UnixDomainSocketTransport.shutdown(bossElg, 5000L) catch { + case NonFatal(e) => stderrln(s"failed to shut down boss event loop: ${e.getMessage}") + } + try UnixDomainSocketTransport.shutdown(workerElg, 5000L) catch { + case NonFatal(e) => stderrln(s"failed to shut down worker event loop: ${e.getMessage}") + } // 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 { + bossElg = transport.newEventLoopGroup() + workerElg = 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 parent = sockFile.getParent + if (parent != null && !Files.exists(parent)) { + Files.createDirectories(parent) + } + // The OS rejects bind() if the path already exists. Remove any stale file + // from a previous run on a best-effort basis. + try { Files.deleteIfExists(sockFile) } catch { case _: IOException => () } + + server = NettyServerBuilder + .forAddress(new DomainSocketAddress(socketPath)) + .channelType(transport.serverChannelType) + .bossEventLoopGroup(bossElg) + .workerEventLoopGroup(workerElg) + .maxInboundMessageSize(GrpcWorkerChannel.MAX_INBOUND_MESSAGE_SIZE) + .addService(new EchoWorkerService) + .build() + + // SIGTERM is the engine's graceful-stop signal per the worker spec. + // SIGINT is added so manual testing (Ctrl-C) also shuts down cleanly. + shutdownHook = new Thread(() => { + stderrln(s"shutdown signal received (worker=$workerId)") + cleanup() + }, s"echo-worker-shutdown-$workerId") + // This standalone test-worker main has no Spark on its classpath, so it + // cannot use Spark's ShutdownHookManager and registers directly. + // scalastyle:off runtimeaddshutdownhook + Runtime.getRuntime.addShutdownHook(shutdownHook) + // scalastyle:on runtimeaddshutdownhook + server.start() stderrln(s"listening on $socketPath (transport=${transport.name})") + // Park until the shutdown hook fires. + server.awaitTermination() } catch { case NonFatal(e) => - 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. + stderrln(s"gRPC worker failed: ${e.getMessage}") + throw e + } finally { + cleanup() + if (shutdownHook != null) { try { Runtime.getRuntime.removeShutdownHook(shutdownHook) } catch { case _: IllegalStateException => () // JVM shutting down already + case NonFatal(e) => stderrln(s"failed to remove shutdown hook: ${e.getMessage}") } - 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 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 index fafee8f0f3d99..cd7eba73f0ee0 100644 --- 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 @@ -42,7 +42,7 @@ import org.apache.spark.udf.worker.{Cancel, CancelResponse, DataRequest, DataRes * 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 + * - A protocol version other than [[SupportedVersion]] triggers an init-time * `ProtocolError`. * - Chunked payloads accumulate inline + chunks before init finalises. */ From 8c61f5a4a6008ffd33a7ba4808067e47777bd3ea Mon Sep 17 00:00:00 2001 From: Haiyang Sun Date: Thu, 3 Sep 2026 16:21:47 +0000 Subject: [PATCH 4/5] [SPARK-56413][SQL] Tighten gRPC worker lifecycle --- .../udf/worker/core/WorkerConnection.scala | 9 +- .../core/direct/DirectWorkerDispatcher.scala | 101 ++++++++++++------ .../core/direct/DirectWorkerProcess.scala | 2 +- .../core/DirectWorkerDispatcherSuite.scala | 94 +++++++++++++--- .../worker/core/TestDirectWorkerHelpers.scala | 8 +- .../worker/grpc/DirectGrpcDispatcher.scala | 13 +-- .../grpc/UnixDomainSocketTransport.scala | 13 ++- ...DirectGrpcDispatcherIntegrationSuite.scala | 17 +++ 8 files changed, 193 insertions(+), 64 deletions(-) diff --git a/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/WorkerConnection.scala b/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/WorkerConnection.scala index 921dc910fd34e..8b7f2185501a9 100644 --- a/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/WorkerConnection.scala +++ b/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/WorkerConnection.scala @@ -34,7 +34,7 @@ import org.apache.spark.annotation.Experimental * * Implementations expose only lifecycle. Data transmission happens at * the [[WorkerSession]] level -- this trait is solely about whether the - * channel is open. + * channel remains open or recoverable. * * '''Relationship to other classes (direct creation mode):''' * {{{ @@ -44,6 +44,11 @@ import org.apache.spark.annotation.Experimental */ @Experimental trait WorkerConnection extends AutoCloseable { - /** Returns true if the underlying transport channel is still usable. */ + /** + * Returns true while the connection remains open or can recover for a future + * session. An implementation may return true during a transient transport + * failure if it reconnects automatically. False means the worker must not be + * reused through this connection. + */ def isActive: Boolean } 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 0a1eb2c70f197..dd880a620467c 100644 --- a/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/direct/DirectWorkerDispatcher.scala +++ b/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/direct/DirectWorkerDispatcher.scala @@ -29,8 +29,8 @@ import scala.util.control.NonFatal import org.apache.spark.annotation.Experimental import org.apache.spark.udf.worker.{ProcessCallable, UDFWorkerSpecification} -import org.apache.spark.udf.worker.core.{WorkerConnection, WorkerDispatcher, WorkerHandle, - WorkerLogger, WorkerSecurityScope, WorkerSession} +import org.apache.spark.udf.worker.core.{WorkerConnection, WorkerDispatcher, WorkerLogger, + WorkerSecurityScope, WorkerSession} import org.apache.spark.udf.worker.core.direct.DirectWorkerDispatcher.{CallableResult, DEFAULT_CALLABLE_TIMEOUT_MS, DEFAULT_GRACEFUL_TIMEOUT_MS, DEFAULT_INIT_TIMEOUT_MS, ENGINE_MAX_TIMEOUT_MS, EnvironmentState, MAX_OUTPUT_SCAN_BYTES, @@ -129,8 +129,9 @@ 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] val lifecycleLock = new Object private[this] var inFlightSessionCreations = 0 + private[this] var inFlightWorkerReleases = 0 // TODO [SPARK-55278]: extract the env state machine + JVM shutdown hook // into a standalone EnvironmentManager once the gRPC dispatcher work has @@ -206,12 +207,10 @@ abstract class DirectWorkerDispatcher( /** * 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. + * `GrpcWorkerSession` for gRPC over UDS) from the process and its + * worker-owned transport connection. */ - protected def newSession( - workerHandle: WorkerHandle, - connection: WorkerConnection): WorkerSession + protected def newSession(worker: DirectWorkerProcess): WorkerSession /** * Test-only hook: invoked once per [[createSession]] after the worker @@ -238,6 +237,7 @@ abstract class DirectWorkerDispatcher( require(securityScope.isEmpty, "securityScope is not supported yet; pass None until pooling lands") beginSessionCreation() + var creationInFlight = true try { ensureEnvironmentReady() if (closed.get()) throwClosed() @@ -257,7 +257,14 @@ abstract class DirectWorkerDispatcher( // The test hook may block, so close could have started since the // first post-publication check. if (closed.get()) throwClosed() - newSession(worker, worker.connection) + val session = newSession(worker) + // Atomically publish the completed session relative to close(). If + // close started while newSession was constructing it, reject the + // session and let the catch block release its worker. + val dispatcherOpen = endSessionCreation() + creationInFlight = false + if (!dispatcherOpen) throwClosed() + session } catch { case e: InterruptedException => Thread.currentThread().interrupt() @@ -268,49 +275,82 @@ abstract class DirectWorkerDispatcher( throw e } } finally { - endSessionCreation() + if (creationInFlight) endSessionCreation() } } /** * Invoked when a worker's last session closes. Terminates the worker - * today; future pooling can reuse it here instead. Safe to call after - * dispatcher close -- the worker's own CAS-idempotent close makes a - * second teardown a no-op. + * today; future pooling can reuse it here instead. Once dispatcher close + * starts, it owns every worker still in the map and performs the teardown. */ private def releaseWorker(worker: DirectWorkerProcess): Unit = { - workers.remove(worker.id) + // Once close starts, it owns every worker still in the map. A release that + // started earlier is tracked so close waits for its teardown to finish. + if (!beginWorkerRelease()) return try { - worker.close() - } catch { - case NonFatal(e) => - logger.warn(s"Error closing worker ${worker.id}", e) + workers.remove(worker.id) + try { + worker.close() + } catch { + case NonFatal(e) => + logger.warn(s"Error closing worker ${worker.id}", e) + } + } finally { + endWorkerRelease() } } private def throwClosed(): Nothing = throw new IllegalStateException("Dispatcher is closed") - private def beginSessionCreation(): Unit = sessionCreationLock.synchronized { + private def beginSessionCreation(): Unit = lifecycleLock.synchronized { if (closed.get()) throwClosed() inFlightSessionCreations += 1 } - private def endSessionCreation(): Unit = sessionCreationLock.synchronized { + /** Returns whether this creation linearized before dispatcher close. */ + private def endSessionCreation(): Boolean = lifecycleLock.synchronized { inFlightSessionCreations -= 1 - if (inFlightSessionCreations == 0) sessionCreationLock.notifyAll() + val dispatcherOpen = !closed.get() + notifyIfDrained() + dispatcherOpen + } + + private def beginWorkerRelease(): Boolean = lifecycleLock.synchronized { + if (closed.get()) { + false + } else { + inFlightWorkerReleases += 1 + true + } + } + + private def endWorkerRelease(): Unit = lifecycleLock.synchronized { + inFlightWorkerReleases -= 1 + notifyIfDrained() + } + + private def notifyIfDrained(): Unit = { + if (inFlightSessionCreations == 0 && inFlightWorkerReleases == 0) { + lifecycleLock.notifyAll() + } + } + + private def beginClose(): Boolean = lifecycleLock.synchronized { + closed.compareAndSet(false, true) } /** - * Waits uninterruptibly so cleanup can never race an in-flight creation. + * Waits uninterruptibly so cleanup cannot race an in-flight lifecycle operation. * Returns whether interruption must be restored after cleanup completes. */ - private def awaitSessionCreations(): Boolean = { + private def awaitInFlightOperations(): Boolean = { var interrupted = Thread.interrupted() - sessionCreationLock.synchronized { - while (inFlightSessionCreations != 0) { + lifecycleLock.synchronized { + while (inFlightSessionCreations != 0 || inFlightWorkerReleases != 0) { try { - sessionCreationLock.wait() + lifecycleLock.wait() } catch { case _: InterruptedException => interrupted = true } @@ -322,14 +362,15 @@ abstract class DirectWorkerDispatcher( /** * Terminates tracked workers, removes the socket directory, and runs * 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. + * calls and worker releases are drained before cleanup, so they cannot + * publish or tear down a worker concurrently with transport or environment + * cleanup. */ override def close(): Unit = { - if (!closed.compareAndSet(false, true)) { + if (!beginClose()) { return } - val interrupted = awaitSessionCreations() + val interrupted = awaitInFlightOperations() try { // TODO [SPARK-55278]: Cleanup sessions as well? // TODO [SPARK-55278]: close workers in parallel -- today shutdown is serialised at 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..117d07ac2311e 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 @@ -143,7 +143,7 @@ class DirectWorkerProcess( } } - /** Returns true if the OS process is running and the connection is usable. */ + /** Returns true if the OS process is running and the connection is open or recoverable. */ def isAlive: Boolean = process.isAlive && connection.isActive /** diff --git a/udf/worker/core/src/test/scala/org/apache/spark/udf/worker/core/DirectWorkerDispatcherSuite.scala b/udf/worker/core/src/test/scala/org/apache/spark/udf/worker/core/DirectWorkerDispatcherSuite.scala index dec934da39005..d9228088f7128 100644 --- a/udf/worker/core/src/test/scala/org/apache/spark/udf/worker/core/DirectWorkerDispatcherSuite.scala +++ b/udf/worker/core/src/test/scala/org/apache/spark/udf/worker/core/DirectWorkerDispatcherSuite.scala @@ -30,8 +30,7 @@ import org.apache.spark.udf.worker.{ Cancel, DirectWorker, LocalTcpConnection, ProcessCallable, UDFProtoCommunicationPattern, UDFWorkerProperties, UDFWorkerSpecification, UnixDomainSocket, WorkerCapabilities, WorkerConnectionSpec, WorkerEnvironment} -import org.apache.spark.udf.worker.core.{WorkerConnection, WorkerHandle, WorkerSecurityScope, - WorkerSession} +import org.apache.spark.udf.worker.core.{WorkerConnection, WorkerSecurityScope, WorkerSession} import org.apache.spark.udf.worker.core.direct.{DirectWorkerException, DirectWorkerProcess, DirectWorkerTimeoutException, UnixDomainSocketEndpointDirectory} @@ -142,7 +141,7 @@ class DirectWorkerDispatcherSuite Thread.sleep(10) } assert(thread.getState == Thread.State.WAITING, - s"$description did not enter the creation barrier; state=${thread.getState}") + s"$description did not enter the lifecycle barrier; state=${thread.getState}") } test("creates a worker and session") { @@ -334,23 +333,25 @@ class DirectWorkerDispatcherSuite } } - test("close racing with in-flight createSession does not leak the worker") { - // close() must wait for thread A to leave createSession before tearing down - // workers or dispatcher-level transport state. + test("close racing with newSession construction does not return a torn-down session") { + // close() must wait for thread A to construct the session, then make the + // final publication check reject it before tearing down dispatcher 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 transportClosed = new java.util.concurrent.atomic.AtomicBoolean(false) val racing = new TestDirectWorkerDispatcher(specWithRunner(defaultRunner)) { - override protected def afterWorkerRegistered(worker: DirectWorkerProcess): Unit = { + override protected def newSession(worker: DirectWorkerProcess): WorkerSession = { + val session = super.newSession(worker) capturedWorkers.add(worker) readyLatch.countDown() - // Block here so dispatcher.close() runs while createSession is in - // flight. Use a generous wait so a slow CI doesn't time out. + // Block after construction but before returning the session so close + // starts in the interval covered by the final publication check. if (!releaseLatch.await(30, java.util.concurrent.TimeUnit.SECONDS)) { fail("releaseLatch never fired -- test orchestration broken") } + session } override protected def closeTransport(): Unit = { @@ -371,10 +372,10 @@ class DirectWorkerDispatcherSuite }, "createSession-racer") createThread.start() - // Wait for thread A to have published the worker and entered the - // blocking override. + // Wait for thread A to construct the session and enter the blocking + // override before returning it to createSession. assert(readyLatch.await(10, java.util.concurrent.TimeUnit.SECONDS), - "createSession thread never reached afterWorkerRegistered") + "createSession thread never constructed the session") val closeThread = new Thread(() => racing.close(), "close-racer") closeThread.start() @@ -402,7 +403,7 @@ class DirectWorkerDispatcherSuite s"expected dispatcher-closed error, got: ${e.getMessage}") case Left(other) => fail(s"unexpected exception from racing createSession: $other") - case Right(_) => fail("createSession should fail when close wins during the test hook") + case Right(_) => fail("createSession should fail when close wins during newSession") } // Whichever path won, the worker must not still be running and the @@ -422,6 +423,69 @@ class DirectWorkerDispatcherSuite } } + test("close waits for an in-flight worker release before transport cleanup") { + val connectionCloseStarted = new java.util.concurrent.CountDownLatch(1) + val allowConnectionClose = new java.util.concurrent.CountDownLatch(1) + val transportClosed = new java.util.concurrent.atomic.AtomicBoolean(false) + val errors = new java.util.concurrent.ConcurrentLinkedQueue[Throwable]() + val racing = new TestDirectWorkerDispatcher(specWithRunner(defaultRunner)) { + override protected def newConnection(address: String): WorkerConnection = + new SocketFileConnection(address) { + override def close(): Unit = { + connectionCloseStarted.countDown() + if (!allowConnectionClose.await(30, java.util.concurrent.TimeUnit.SECONDS)) { + throw new IllegalStateException("allowConnectionClose never fired") + } + super.close() + } + } + + override protected def closeTransport(): Unit = { + transportClosed.set(true) + super.closeTransport() + } + } + try { + val session = racing.createSession(None) + val worker = workerProcess(session) + val releaseThread = new Thread(() => { + try session.close(emptyCancel) catch { + case t: Throwable => errors.add(t) + } + }, "session-release-racer") + releaseThread.start() + + assert(connectionCloseStarted.await(10, java.util.concurrent.TimeUnit.SECONDS), + "session release never started closing the worker connection") + + val closeThread = new Thread(() => { + try racing.close() catch { + case t: Throwable => errors.add(t) + } + }, "dispatcher-close-racer") + closeThread.start() + awaitThreadWaiting(closeThread, "dispatcher close thread") + + assert(closeThread.isAlive, "dispatcher close should wait for worker release") + assert(!transportClosed.get(), + "transport cleanup must wait for worker release teardown") + + allowConnectionClose.countDown() + releaseThread.join(10000) + closeThread.join(10000) + + assert(!releaseThread.isAlive, "session release thread did not finish") + assert(!closeThread.isAlive, "dispatcher close thread did not finish") + assert(errors.isEmpty, + s"unexpected close errors: ${errors.toArray.mkString(", ")}") + assert(transportClosed.get(), "dispatcher close should clean up transport state") + assert(!worker.process.isAlive, "worker process should be terminated") + } finally { + allowConnectionClose.countDown() + racing.close() + } + } + test("interrupted close completes cleanup before restoring the interrupt") { val cleanupMarker = Files.createTempFile("interrupted-close-cleanup", ".txt").toFile cleanupMarker.delete() @@ -597,9 +661,7 @@ class DirectWorkerDispatcherSuite var capturedWorker: DirectWorkerProcess = null val failingDispatcher = new TestDirectWorkerDispatcher(specWithRunner(defaultRunner)) { - override protected def newSession( - workerHandle: WorkerHandle, - connection: WorkerConnection): WorkerSession = + override protected def newSession(worker: DirectWorkerProcess): WorkerSession = throw new RuntimeException("session creation failed") override protected def afterWorkerRegistered(w: DirectWorkerProcess): Unit = { capturedWorker = w diff --git a/udf/worker/core/src/test/scala/org/apache/spark/udf/worker/core/TestDirectWorkerHelpers.scala b/udf/worker/core/src/test/scala/org/apache/spark/udf/worker/core/TestDirectWorkerHelpers.scala index 32e5001160285..242c6462768be 100644 --- a/udf/worker/core/src/test/scala/org/apache/spark/udf/worker/core/TestDirectWorkerHelpers.scala +++ b/udf/worker/core/src/test/scala/org/apache/spark/udf/worker/core/TestDirectWorkerHelpers.scala @@ -21,7 +21,7 @@ import java.io.File 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, UnixDomainSocketEndpointDirectory} + DirectWorkerProcess, DirectWorkerTimeoutException, UnixDomainSocketEndpointDirectory} import org.apache.spark.udf.worker.core.direct.DirectWorkerDispatcher.READY_POLL_INTERVAL_MS /** @@ -152,10 +152,8 @@ class TestDirectWorkerDispatcher( protected def newConnection(address: String): WorkerConnection = new SocketFileConnection(address) - override protected def newSession( - workerHandle: WorkerHandle, - connection: WorkerConnection): WorkerSession = - new NoOpWorkerSession(workerHandle, logger) + override protected def newSession(worker: DirectWorkerProcess): WorkerSession = + new NoOpWorkerSession(worker, logger) private def throwWorkerExitedBeforeSocket( process: Process, diff --git a/udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/DirectGrpcDispatcher.scala b/udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/DirectGrpcDispatcher.scala index 4b33366ef73d1..2aed6b68246be 100644 --- a/udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/DirectGrpcDispatcher.scala +++ b/udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/DirectGrpcDispatcher.scala @@ -27,10 +27,9 @@ import io.grpc.ConnectivityState 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.{WorkerConnection, WorkerLogger, WorkerSession} import org.apache.spark.udf.worker.core.direct.{DirectWorkerDispatcher, DirectWorkerException, - DirectWorkerTimeoutException, UnixDomainSocketEndpointDirectory} + DirectWorkerProcess, DirectWorkerTimeoutException, UnixDomainSocketEndpointDirectory} import org.apache.spark.udf.worker.core.direct.DirectWorkerDispatcher.READY_POLL_INTERVAL_MS /** @@ -156,12 +155,10 @@ class DirectGrpcDispatcher( protected def newConnection(address: String): WorkerConnection = new GrpcWorkerChannel(address, logger) - override protected def newSession( - workerHandle: WorkerHandle, - connection: WorkerConnection): WorkerSession = - connection match { + override protected def newSession(worker: DirectWorkerProcess): WorkerSession = + worker.connection match { case g: GrpcWorkerChannel => - new GrpcWorkerSession(workerHandle, g.channel, logger) + new GrpcWorkerSession(worker, g.channel, logger) case other => throw new IllegalStateException( s"DirectGrpcDispatcher.newConnection should have produced a " + 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 index d07ab0c320def..6057ec8482db7 100644 --- 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 @@ -20,6 +20,7 @@ import java.util.Locale import java.util.concurrent.TimeUnit import io.netty.channel.{Channel, EventLoopGroup, MultiThreadIoEventLoopGroup, ServerChannel} +import io.netty.util.concurrent.DefaultThreadFactory /** * Selects the Netty native transport that supports Unix domain sockets on the @@ -50,6 +51,8 @@ private[grpc] sealed trait UnixDomainSocketTransport { private[grpc] object UnixDomainSocketTransport { + private[grpc] val EVENT_LOOP_THREAD_NAME_PREFIX = "spark-udf-worker-grpc" + /** Linux epoll. */ private case object EpollTransport extends UnixDomainSocketTransport { override val channelType: Class[_ <: Channel] = @@ -57,7 +60,10 @@ private[grpc] object UnixDomainSocketTransport { override val serverChannelType: Class[_ <: ServerChannel] = classOf[io.netty.channel.epoll.EpollServerDomainSocketChannel] override def newEventLoopGroup(): EventLoopGroup = - new MultiThreadIoEventLoopGroup(1, io.netty.channel.epoll.EpollIoHandler.newFactory()) + new MultiThreadIoEventLoopGroup( + 1, + new DefaultThreadFactory(EVENT_LOOP_THREAD_NAME_PREFIX, true), + io.netty.channel.epoll.EpollIoHandler.newFactory()) override val name: String = "epoll" } @@ -68,7 +74,10 @@ private[grpc] object UnixDomainSocketTransport { override val serverChannelType: Class[_ <: ServerChannel] = classOf[io.netty.channel.kqueue.KQueueServerDomainSocketChannel] override def newEventLoopGroup(): EventLoopGroup = - new MultiThreadIoEventLoopGroup(1, io.netty.channel.kqueue.KQueueIoHandler.newFactory()) + new MultiThreadIoEventLoopGroup( + 1, + new DefaultThreadFactory(EVENT_LOOP_THREAD_NAME_PREFIX, true), + io.netty.channel.kqueue.KQueueIoHandler.newFactory()) override val name: String = "kqueue" } 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 index dc82bcfe205e7..370bf907414c9 100644 --- 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 @@ -18,6 +18,7 @@ package org.apache.spark.udf.worker.grpc import java.io.File import java.nio.file.Paths +import java.util.concurrent.{Callable, TimeUnit} import com.google.protobuf.ByteString import org.scalatest.BeforeAndAfterEach @@ -118,6 +119,22 @@ class DirectGrpcDispatcherIntegrationSuite case other => fail(s"Expected GrpcWorkerChannel, got ${other.getClass.getSimpleName}") } + test("event-loop threads are named daemons and terminate on shutdown") { + val eventLoopGroup = UnixDomainSocketTransport.detect().newEventLoopGroup() + try { + val eventLoopThread = eventLoopGroup.next().submit(new Callable[Thread] { + override def call(): Thread = Thread.currentThread() + }).get(10, TimeUnit.SECONDS) + + assert(eventLoopThread.isDaemon) + assert(eventLoopThread.getName.startsWith( + UnixDomainSocketTransport.EVENT_LOOP_THREAD_NAME_PREFIX)) + } finally { + UnixDomainSocketTransport.shutdown(eventLoopGroup, 5000L) + } + assert(eventLoopGroup.terminationFuture().isDone) + } + test("spawns a gRPC worker and round-trips a response above gRPC's default limit") { dispatcher = new DirectGrpcDispatcher(workerSpec()) val session = dispatcher.createSession(None) From 32708dcf177cade8b06e0a35cf36d197e5275395 Mon Sep 17 00:00:00 2001 From: Haiyang Sun Date: Thu, 3 Sep 2026 23:40:29 +0000 Subject: [PATCH 5/5] [SPARK-56413][SQL] Preserve echo client request synchronization --- .../udf/worker/grpc/EchoProtocolSuite.scala | 27 +++++++++++-------- 1 file changed, 16 insertions(+), 11 deletions(-) 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 26922185bb7a6..3072c57d5bb57 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 @@ -90,6 +90,10 @@ class EchoProtocolSuite extends AnyFunSuite with BeforeAndAfterEach { @volatile var executionError: Option[ExecutionError] = None @volatile var streamError: Option[Throwable] = None private val requestCompleted = new AtomicBoolean(false) + // gRPC forbids concurrent calls to the request StreamObserver too: the response + // observer half-closes it (completeRequestStream) on a callback thread while the + // test thread may still be sending a trailing Cancel. Serialize both through this. + private val requestLock = new Object // 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) @@ -258,22 +262,23 @@ class EchoProtocolSuite extends AnyFunSuite with BeforeAndAfterEach { } } - 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()) + def sendCancel(reason: String = ""): Unit = requestLock.synchronized { + // 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. requestLock makes this check + // and the send atomic against completeRequestStream's half-close. + if (!requestCompleted.get()) { + requestObserver.onNext(UdfRequest.newBuilder() + .setControl(UdfControlRequest.newBuilder() + .setCancel(Cancel.newBuilder().setReason(reason).build()) + .build()) .build()) - .build()) + } // Request stream stays open until the response terminator arrives; // completeRequestStream() is called by the response observer. } - def completeRequestStream(): Unit = { + def completeRequestStream(): Unit = requestLock.synchronized { if (requestCompleted.compareAndSet(false, true)) { requestObserver.onCompleted() }