hvanhovell commented on code in PR #55683: URL: https://github.com/apache/spark/pull/55683#discussion_r3342517076
########## udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/grpc/DirectGrpcDispatcher.scala: ########## @@ -0,0 +1,238 @@ +/* + * 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.grpc + +import java.io.{File, IOException} +import java.nio.file.{FileVisitResult, Files, Path, SimpleFileVisitor} +import java.nio.file.attribute.{BasicFileAttributes, PosixFilePermissions} + +import org.apache.spark.annotation.Experimental +import org.apache.spark.udf.worker.{UDFProtoCommunicationPattern, UDFWorkerSpecification, + WorkerConnectionSpec} +import org.apache.spark.udf.worker.core.{WorkerConnection, WorkerHandle, WorkerLogger, + WorkerSession} +import org.apache.spark.udf.worker.core.direct.{DirectWorkerDispatcher, DirectWorkerException, + DirectWorkerTimeoutException} +import org.apache.spark.udf.worker.core.direct.DirectWorkerDispatcher.SOCKET_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. + * + * Lives in the `core.grpc` package alongside [[GrpcWorkerSession]] and + * [[GrpcWorkerChannel]] -- the two collaborators it composes. + */ +@Experimental +class DirectGrpcDispatcher( + workerSpec: UDFWorkerSpecification, + logger: WorkerLogger = WorkerLogger.NoOp) + extends DirectWorkerDispatcher(workerSpec, logger) { + + // Upper bound on the rename-on-collision retry loop in newEndpointAddress. + // 16 is well above any realistic concurrent-worker count and keeps the + // failure mode bounded if something pathological occupies the directory. + private val MAX_SOCKET_LEAF_RETRIES = 16 + + // The private 0700 socket directory, created in [[initialize]] (after the + // base class has validated the spec) and removed in [[closeTransport]]. + // `lazy val` + force-in-initialize keeps creation deterministic without + // depending on subclass field-initialiser ordering. deleteOnExit is + // avoided because the JDK retains the path for the JVM lifetime, which + // leaks in long-lived drivers. + private lazy val socketDir: Path = createPrivateTempDirectory() + + override protected def initialize(): Unit = { + super.initialize() + // Force the lazy val now so the directory is created (and any failure + // surfaces) at construction time, after spec validation has passed. + socketDir + } + + /** + * Returns the UDS path the worker should bind. Uses a short leaf so the + * full path stays within the 108-byte UDS `sun_path` limit; the worker's + * stable id (full UUID) is still passed via `--id` and logged separately. + * + * '''Why the truncation:''' macOS resolves `$TMPDIR` to + * `/var/folders/xx/yyyy/T/` (~47 chars) and `Files.createTempDirectory` + * appends another ~29 chars for the dispatcher's private dir. With a + * full UUID leaf (`worker-<36>.sock` = 49 chars) the total path is + * ~126 chars and `bind(2)` would silently fail with `ENAMETOOLONG`. + * 16 hex chars (64 bits) keeps the leaf at ~25 chars (total ~100) while + * making collisions negligible at any realistic concurrent-worker count. + * + * '''Defensive collision check:''' a 64-bit hash collision is + * astronomically unlikely (~2^-32 per pair) but if it ever fired the + * second worker's `bind(2)` would surface as an opaque "Worker exited + * with code N" via [[waitForReady]] rather than a clear filename + * collision. We probe the path here and append a counter suffix on + * the off-chance the leaf already exists so the error message is + * recognisable in the wild. + */ + override protected def newEndpointAddress(workerId: String): String = { + val short = workerId.replace("-", "").take(16) + var candidate = socketDir.resolve(s"w-$short.sock") + var suffix = 0 + while (Files.exists(candidate) && suffix < MAX_SOCKET_LEAF_RETRIES) { + suffix += 1 + candidate = socketDir.resolve(s"w-$short-$suffix.sock") + } + if (Files.exists(candidate)) { + throw new IllegalStateException( + s"could not allocate a free UDS path under $socketDir after " + + s"$MAX_SOCKET_LEAF_RETRIES retries (truncated id=$short)") + } + candidate.toString + } + + override protected def waitForReady( + address: String, + process: Process, + outputFile: File): Unit = { + val file = new File(address) + // At least one poll so very small initTimeouts don't trip a premature + // timeout before the worker has any chance to create the socket. + val maxAttempts = math.max(1, (initTimeoutMs / SOCKET_POLL_INTERVAL_MS).toInt) + var attempts = 0 + while (!file.exists() && attempts < maxAttempts) { + if (!process.isAlive) throwWorkerExitedBeforeSocket(process, address, outputFile) + Thread.sleep(SOCKET_POLL_INTERVAL_MS) + attempts += 1 + } + if (!file.exists()) { + if (process.isAlive) { + DirectWorkerDispatcher.destroyForciblyAndReap( + process, logger, s"init timeout $address") + val tail = readOutputTail(outputFile) + throw new DirectWorkerTimeoutException( + s"Worker did not create socket at $address within ${initTimeoutMs}ms\n$tail") + } else { + // Worker exited after the last poll without creating the socket; + // prefer the exit-code message over the ambiguous "did not create". + throwWorkerExitedBeforeSocket(process, address, outputFile) + } + } + } + + override protected def cleanupEndpointAddress(address: String): Unit = { + Files.deleteIfExists(new File(address).toPath) + } + + override protected def closeTransport(): Unit = { + if (!Files.exists(socketDir)) return + // Recursive post-order delete: today socketDir contains only socket files + // at the top level, but a future change that namespaces workers into + // subdirectories should not silently leak them. + Files.walkFileTree(socketDir, new SimpleFileVisitor[Path] { + override def visitFile(file: Path, attrs: BasicFileAttributes): FileVisitResult = { + try Files.deleteIfExists(file) catch { case _: IOException => () } + FileVisitResult.CONTINUE + } + override def postVisitDirectory(dir: Path, exc: IOException): FileVisitResult = { + try Files.deleteIfExists(dir) catch { case _: IOException => () } + FileVisitResult.CONTINUE + } + }) + } + + // `spec` is the same object as the `workerSpec` field but passed + // explicitly: at the point this runs (parent constructor body), `this` + // is only partially constructed and reading subclass fields is unsafe. + // 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. If the spec declares capabilities, it MUST include it; an + // unset capabilities block is treated as "no constraint" and accepted. + if (spec.hasCapabilities) { + val patterns = spec.getCapabilities.getSupportedCommunicationPatternsList + if (!patterns.isEmpty) { + 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") + } + } + } + + override protected def newConnection(address: String): WorkerConnection = + new GrpcWorkerChannel(address, logger) + + override protected def newSession( + workerHandle: WorkerHandle, + connection: WorkerConnection): WorkerSession = connection match { + case g: GrpcWorkerChannel => + new GrpcWorkerSession(workerHandle, g.channel, logger) + case other => + throw new IllegalStateException( + s"DirectGrpcDispatcher.newConnection should have produced a " + + s"GrpcWorkerChannel but got ${other.getClass.getName}") + } + + private def throwWorkerExitedBeforeSocket( + process: Process, + address: String, + outputFile: File): Nothing = { + val tail = readOutputTail(outputFile) + throw new DirectWorkerException( + s"Worker exited with code ${process.exitValue()} " + + s"before creating socket at $address\n$tail") + } + + /** + * Creates a temp directory with owner-only permissions (0700 on POSIX). + * On non-POSIX filesystems falls back to best-effort `File.setXxx`, + * which is TOCTOU-racy and weaker; a WARN surfaces if the platform Review Comment: TOCTOU? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
