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/UnixSocketWorkerConnection.scala b/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/UnixSocketWorkerConnection.scala deleted file mode 100644 index b3b40d16e7443..0000000000000 --- a/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/UnixSocketWorkerConnection.scala +++ /dev/null @@ -1,41 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one or more - * contributor license agreements. See the NOTICE file distributed with - * this work for additional information regarding copyright ownership. - * The ASF licenses this file to You under the Apache License, Version 2.0 - * (the "License"); you may not use this file except in compliance with - * the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package org.apache.spark.udf.worker.core - -import java.io.File - -import org.apache.spark.annotation.Experimental - -/** - * :: 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. - * - * [[close]] is idempotent: deleting an already-removed file is a no-op. - */ -@Experimental -abstract class UnixSocketWorkerConnection(val socketPath: String) - extends WorkerConnection { - - override def close(): Unit = { - val f = new File(socketPath) - if (f.exists()) f.delete() - } -} 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 db81b5265b794..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, @@ -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,9 @@ abstract class DirectWorkerDispatcher( private[this] val workers = new ConcurrentHashMap[String, DirectWorkerProcess]() private[this] val closed = new AtomicBoolean(false) + 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 @@ -145,18 +148,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 +204,13 @@ 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. + * `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 @@ -240,81 +236,163 @@ 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() + var creationInFlight = true try { - afterWorkerRegistered(worker) - newSession(worker, worker.connection) - } catch { - case e: InterruptedException => - Thread.currentThread().interrupt() - worker.releaseSession() - throw e - case NonFatal(e) => + 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 + throwClosed() + } + try { + afterWorkerRegistered(worker) + // The test hook may block, so close could have started since the + // first post-publication check. + if (closed.get()) throwClosed() + 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() + worker.releaseSession() + throw e + case NonFatal(e) => + worker.releaseSession() + throw e + } + } finally { + 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 = lifecycleLock.synchronized { + if (closed.get()) throwClosed() + inFlightSessionCreations += 1 + } + + /** Returns whether this creation linearized before dispatcher close. */ + private def endSessionCreation(): Boolean = lifecycleLock.synchronized { + inFlightSessionCreations -= 1 + 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 cannot race an in-flight lifecycle operation. + * Returns whether interruption must be restored after cleanup completes. + */ + private def awaitInFlightOperations(): Boolean = { + var interrupted = Thread.interrupted() + lifecycleLock.synchronized { + while (inFlightSessionCreations != 0 || inFlightWorkerReleases != 0) { + try { + lifecycleLock.wait() + } catch { + case _: InterruptedException => interrupted = true + } + } + } + interrupted + } + /** * 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 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 } - // 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 = awaitInFlightOperations() + 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 ------------------------------------------------- @@ -455,10 +533,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 +547,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 +638,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..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/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 28ecc1179f6ac..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,10 +30,9 @@ 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} + DirectWorkerTimeoutException, UnixDomainSocketEndpointDirectory} /** * Tests for [[DirectWorkerDispatcher]] process lifecycle: spawning workers @@ -135,6 +134,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 lifecycle barrier; state=${thread.getState}") + } + test("creates a worker and session") { dispatcher = new TestDirectWorkerDispatcher(specWithRunner(defaultRunner)) @@ -324,26 +333,30 @@ 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. + 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 = { + transportClosed.set(true) + super.closeTransport() } } try { @@ -359,15 +372,17 @@ 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() - // Give close() time to flip `closed` and iterate workers. - Thread.sleep(200) + 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") // Now release the in-flight createSession. releaseLatch.countDown() @@ -384,18 +399,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 newSession") } // Whichever path won, the worker must not still be running and the @@ -415,6 +423,150 @@ 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() + 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. @@ -509,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 7ffe72f917573..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 @@ -16,15 +16,13 @@ */ package org.apache.spark.udf.worker.core -import java.io.{File, IOException} -import java.nio.file.{Files, FileVisitResult, Path, Paths, SimpleFileVisitor} -import java.nio.file.attribute.{BasicFileAttributes, PosixFilePermissions} +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} -import org.apache.spark.udf.worker.core.direct.DirectWorkerDispatcher.SOCKET_POLL_INTERVAL_MS + DirectWorkerProcess, DirectWorkerTimeoutException, UnixDomainSocketEndpointDirectory} +import org.apache.spark.udf.worker.core.direct.DirectWorkerDispatcher.READY_POLL_INTERVAL_MS /** * A [[WorkerConnection]] test implementation that treats the connection as @@ -83,53 +81,28 @@ 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 - - // 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)") - } - candidate.toString - } + override protected def newEndpointAddress(workerId: String): String = + endpointDirectory.newEndpointAddress(workerId) - override protected def waitForReady( + override protected def connectWorker( address: String, process: Process, - outputFile: File): Unit = { + 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 / SOCKET_POLL_INTERVAL_MS).toInt) + 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(SOCKET_POLL_INTERVAL_MS) + Thread.sleep(READY_POLL_INTERVAL_MS) attempts += 1 } if (!file.exists()) { @@ -145,28 +118,13 @@ class TestDirectWorkerDispatcher( throwWorkerExitedBeforeSocket(process, address, outputFile) } } + 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 - // Recursive post-order delete: today socketDir contains only socket files - // at the top level, but a future change that namespaces workers into - // subdirectories should not silently leak them. - Files.walkFileTree(socketDir, new SimpleFileVisitor[Path] { - override def visitFile(file: Path, attrs: BasicFileAttributes): FileVisitResult = { - try Files.deleteIfExists(file) catch { case _: IOException => () } - FileVisitResult.CONTINUE - } - override def postVisitDirectory(dir: Path, exc: IOException): FileVisitResult = { - try Files.deleteIfExists(dir) catch { case _: IOException => () } - FileVisitResult.CONTINUE - } - }) - } + 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 @@ -191,13 +149,11 @@ class TestDirectWorkerDispatcher( "in WorkerCapabilities.supported_communication_patterns") } - override protected def newConnection(address: String): WorkerConnection = + 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, @@ -208,75 +164,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)) - // 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 } - usable - .sortBy(realLen) - .headOption - .getOrElse(Paths.get(System.getProperty("java.io.tmpdir"))) - } } diff --git a/udf/worker/grpc/pom.xml b/udf/worker/grpc/pom.xml index aa563d157d9c0..c915e819410ea 100644 --- a/udf/worker/grpc/pom.xml +++ b/udf/worker/grpc/pom.xml @@ -91,6 +91,45 @@ io.grpc grpc-stub + + + io.grpc + grpc-netty + ${io.grpc.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/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 new file mode 100644 index 0000000000000..2aed6b68246be --- /dev/null +++ b/udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/DirectGrpcDispatcher.scala @@ -0,0 +1,191 @@ +/* + * 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.{Files, Paths} +import java.util.concurrent.{CountDownLatch, TimeUnit} + +import scala.util.control.NonFatal + +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, WorkerLogger, WorkerSession} +import org.apache.spark.udf.worker.core.direct.{DirectWorkerDispatcher, DirectWorkerException, + DirectWorkerProcess, DirectWorkerTimeoutException, UnixDomainSocketEndpointDirectory} +import org.apache.spark.udf.worker.core.direct.DirectWorkerDispatcher.READY_POLL_INTERVAL_MS + +/** + * :: 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. + * + * @param workerSpec worker specification used to launch each worker. + * @param logger logger for lifecycle diagnostics. + */ +@Experimental +class DirectGrpcDispatcher( + workerSpec: UDFWorkerSpecification, + logger: WorkerLogger = WorkerLogger.NoOp) + extends DirectWorkerDispatcher(workerSpec, logger) { + + // 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() + endpointDirectory + } + + override protected def newEndpointAddress(workerId: String): String = + endpointDirectory.newEndpointAddress(workerId) + + 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 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}") + } + val deadlineNanos = System.nanoTime() + TimeUnit.MILLISECONDS.toNanos(initTimeoutMs) + 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 become reachable at $address within ${initTimeoutMs}ms\n$tail") + } + + // 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() + } + 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) + } + } + + override protected def cleanupEndpointAddress(address: String): Unit = + endpointDirectory.cleanupEndpointAddress(address) + + 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 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, + "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, + "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, + "DirectGrpcDispatcher requires BIDIRECTIONAL_STREAMING " + + "in WorkerCapabilities.supported_communication_patterns") + } + + protected def newConnection(address: String): WorkerConnection = + new GrpcWorkerChannel(address, logger) + + override protected def newSession(worker: DirectWorkerProcess): WorkerSession = + worker.connection match { + case g: GrpcWorkerChannel => + new GrpcWorkerSession(worker, g.channel, logger) + case other => + throw new IllegalStateException( + s"DirectGrpcDispatcher.newConnection should have produced a " + + s"GrpcWorkerChannel but got ${other.getClass.getName}") + } + + 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 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) + } + } + +} + +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 new file mode 100644 index 0000000000000..220a31d633312 --- /dev/null +++ b/udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/GrpcWorkerChannel.scala @@ -0,0 +1,136 @@ +/* + * 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 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 + * 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. + */ +@Experimental +class GrpcWorkerChannel( + val socketPath: String, + logger: WorkerLogger = WorkerLogger.NoOp) + extends WorkerConnection { + + 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(GrpcWorkerChannel.MAX_INBOUND_MESSAGE_SIZE) + .usePlaintext() + .build() + } catch { + case NonFatal(e) => + UnixDomainSocketTransport.shutdown( + eventLoopGroup, GrpcWorkerChannel.SHUTDOWN_PHASE_TIMEOUT_MS) + throw e + } + + private val closed = new AtomicBoolean(false) + + override def isActive: Boolean = { + if (closed.get()) return false + 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 + * removed by the dispatcher, not here. + */ + override def close(): Unit = { + if (!closed.compareAndSet(false, true)) return + + try { + channel.shutdown() + if (!channel.awaitTermination( + GrpcWorkerChannel.SHUTDOWN_PHASE_TIMEOUT_MS, TimeUnit.MILLISECONDS)) { + channel.shutdownNow() + channel.awaitTermination( + GrpcWorkerChannel.SHUTDOWN_PHASE_TIMEOUT_MS, 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, GrpcWorkerChannel.SHUTDOWN_PHASE_TIMEOUT_MS) + } 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. */ + private[grpc] val MAX_INBOUND_MESSAGE_SIZE: Int = 128 * 1024 * 1024 + + 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 new file mode 100644 index 0000000000000..6057ec8482db7 --- /dev/null +++ b/udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/UnixDomainSocketTransport.scala @@ -0,0 +1,158 @@ +/* + * 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} +import io.netty.util.concurrent.DefaultThreadFactory + +/** + * 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 + * [[UnixDomainSocketTransport.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 { + + 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] = + classOf[io.netty.channel.epoll.EpollDomainSocketChannel] + override val serverChannelType: Class[_ <: ServerChannel] = + classOf[io.netty.channel.epoll.EpollServerDomainSocketChannel] + override def newEventLoopGroup(): EventLoopGroup = + new MultiThreadIoEventLoopGroup( + 1, + new DefaultThreadFactory(EVENT_LOOP_THREAD_NAME_PREFIX, true), + 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, + new DefaultThreadFactory(EVENT_LOOP_THREAD_NAME_PREFIX, true), + 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..370bf907414c9 --- /dev/null +++ b/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/DirectGrpcDispatcherIntegrationSuite.scala @@ -0,0 +1,216 @@ +/* + * 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.concurrent.{Callable, TimeUnit} + +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, DirectWorker, Finish, Init, + ProcessCallable, UdfPayload, UDFProtoCommunicationPattern, UDFWorkerDataFormat, + UDFWorkerProperties, UDFWorkerSpecification, UnixDomainSocket, WorkerCapabilities, + WorkerConnectionSpec} +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 + +/** + * 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 + + private val javaClasspath = System.getProperty("java.class.path") + private val javaExecutable = + Paths.get(System.getProperty("java.home"), "bin", "java").toString + private val workerMainClass = + 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") + } + + override def afterEach(): Unit = { + if (dispatcher != null) { + try dispatcher.close() finally dispatcher = null + } + super.afterEach() + } + + private def echoRunner: ProcessCallable = ProcessCallable.newBuilder() + .addCommand(javaExecutable) + .addCommand("-cp") + .addCommand(javaClasspath) + .addCommand(workerMainClass) + .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 basicInit: Init = Init.newBuilder() + .setProtocolVersion(1) + .setDataFormat(UDFWorkerDataFormat.ARROW) + .setUdf(UdfPayload.newBuilder() + .setPayload(ByteString.copyFromUtf8("echo")) + .setFormat("echo") + .build()) + .build() + + private val emptyFinish: () => Finish = () => Finish.getDefaultInstance + private val emptyCancel: () => Cancel = () => Cancel.getDefaultInstance + + private def workerProcess(session: WorkerSession): DirectWorkerProcess = + session.workerHandle match { + case process: DirectWorkerProcess => process + case other => fail(s"Expected DirectWorkerProcess, got ${other.getClass.getSimpleName}") + } + + private def grpcChannel(process: DirectWorkerProcess): GrpcWorkerChannel = + process.connection match { + case channel: GrpcWorkerChannel => channel + 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) + 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)) + + try { + 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) + } + + assert(!process.process.isAlive, "session close should terminate the worker") + assert(!socketFile.exists(), "session close should remove the worker socket") + dispatcher.close() + dispatcher = null + assert(!socketDir.exists(), "dispatcher close should remove its socket directory") + } + + 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() + 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/EchoProtocolSuite.scala b/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/EchoProtocolSuite.scala index 4de8342208399..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 @@ -20,21 +20,17 @@ import java.util.concurrent.{CountDownLatch, LinkedBlockingQueue, TimeUnit} import java.util.concurrent.atomic.AtomicBoolean import com.google.protobuf.ByteString -import io.grpc.{ManagedChannel, Server, Status} +import io.grpc.{ManagedChannel, Server} import io.grpc.inprocess.{InProcessChannelBuilder, InProcessServerBuilder} import io.grpc.stub.StreamObserver import org.scalatest.BeforeAndAfterEach // scalastyle:off funsuite import org.scalatest.funsuite.AnyFunSuite -import org.apache.spark.udf.worker.{ - Cancel, CancelResponse, DataRequest, DataResponse, ErrorResponse, ExecutionError, - Finish, FinishResponse, Heartbeat, HeartbeatResponse, Init, InitResponse, - PayloadChunk, ProtocolError, ShutdownRequest, ShutdownResponse, - UdfControlRequest, UdfControlResponse, UdfPayload, UdfRequest, UdfResponse, - UDFWorkerDataFormat, UdfWorkerGrpc, UserError, WorkerError, WorkerRequest, - WorkerResponse -} +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. @@ -48,22 +44,23 @@ import org.apache.spark.udf.worker.{ 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") + // 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 = _ - private var stub: UdfWorkerGrpc.UdfWorkerStub = _ + 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(new EchoWorkerService) + .addService(workerService) .build() .start() channel = InProcessChannelBuilder.forName(serverName).directExecutor().build() @@ -75,382 +72,6 @@ class EchoProtocolSuite extends AnyFunSuite with BeforeAndAfterEach { 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) // =========================================================================== @@ -463,7 +84,7 @@ class EchoProtocolSuite extends AnyFunSuite with BeforeAndAfterEach { * 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) { + 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 @@ -541,7 +162,7 @@ class EchoProtocolSuite extends AnyFunSuite with BeforeAndAfterEach { def sendInit( payloadBytes: Array[Byte], sendChunked: Boolean = false, - protocolVersion: Int = SUPPORTED_VERSION): Unit = { + protocolVersion: Int = SupportedVersion): Unit = { if (sendChunked) { requestObserver.onNext(UdfRequest.newBuilder() .setControl(UdfControlRequest.newBuilder() @@ -591,7 +212,7 @@ class EchoProtocolSuite extends AnyFunSuite with BeforeAndAfterEach { // 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 = { + def sendInitChunked(protocolVersion: Int = SupportedVersion): Unit = { requestObserver.onNext(UdfRequest.newBuilder() .setControl(UdfControlRequest.newBuilder() .setInit(Init.newBuilder() @@ -755,7 +376,7 @@ class EchoProtocolSuite extends AnyFunSuite with BeforeAndAfterEach { reqObs.onNext(UdfRequest.newBuilder() .setControl(UdfControlRequest.newBuilder() .setInit(Init.newBuilder() - .setProtocolVersion(SUPPORTED_VERSION) + .setProtocolVersion(SupportedVersion) .setDataFormat(UDFWorkerDataFormat.ARROW) .setUdf(UdfPayload.newBuilder() .setPayload(ByteString.copyFromUtf8("payload")) @@ -832,7 +453,7 @@ class EchoProtocolSuite extends AnyFunSuite with BeforeAndAfterEach { 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) + client.sendData(ErrorTrigger.toByteArray) assert(client.awaitDone()) assert(client.streamError.isEmpty, s"expected no gRPC error, got ${client.streamError}") @@ -854,7 +475,7 @@ class EchoProtocolSuite extends AnyFunSuite with BeforeAndAfterEach { test("init error: inline payload triggers init failure") { val client = new EngineClient(stub) - client.sendInit(INIT_ERROR_TRIGGER.toByteArray) + client.sendInit(InitErrorTrigger.toByteArray) assert(client.awaitDone()) assert(client.streamError.isEmpty, s"expected no gRPC error, got ${client.streamError}") @@ -865,7 +486,7 @@ class EchoProtocolSuite extends AnyFunSuite with BeforeAndAfterEach { test("init error: chunked payload assembled across chunks triggers init failure") { val client = new EngineClient(stub) - client.sendInit(INIT_ERROR_TRIGGER.toByteArray, sendChunked = true) + client.sendInit(InitErrorTrigger.toByteArray, sendChunked = true) assert(client.awaitDone()) assert(client.streamError.isEmpty, s"expected no gRPC error, got ${client.streamError}") @@ -875,7 +496,7 @@ class EchoProtocolSuite extends AnyFunSuite with BeforeAndAfterEach { test("init error: unsupported protocol version triggers init failure") { val client = new EngineClient(stub) - client.sendInit("payload".getBytes, protocolVersion = SUPPORTED_VERSION + 999) + client.sendInit("payload".getBytes, protocolVersion = SupportedVersion + 999) assert(client.awaitDone()) assert(client.streamError.isEmpty, s"expected no gRPC error, got ${client.streamError}") 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..b48eb57f361c9 --- /dev/null +++ b/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/testing/EchoGrpcWorkerMain.scala @@ -0,0 +1,165 @@ +/* + * 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 java.util.concurrent.atomic.AtomicBoolean + +import scala.util.control.NonFatal + +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 sockFile = Paths.get(socketPath) + val cleanupStarted = new AtomicBoolean(false) + var bossElg: EventLoopGroup = null + var workerElg: EventLoopGroup = null + var server: Server = null + var shutdownHook: Thread = null + + def cleanup(): Unit = { + if (!cleanupStarted.compareAndSet(false, true)) return + try { + 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 { + 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 => () } + } + } + + 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"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}") + } + } + } + } + // 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..cd7eba73f0ee0 --- /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`. + * - A 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() + } + } +}