haiyangsun-db commented on code in PR #55683: URL: https://github.com/apache/spark/pull/55683#discussion_r3376149941
########## udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/grpc/DirectGrpcDispatcher.scala: ########## @@ -0,0 +1,237 @@ +/* + * 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. + * + * gRPC over UDS is the only protocol/transport combination shipped today. + * 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 + + // Removed explicitly in closeTransport(). deleteOnExit is avoided because + // the JDK retains the path for the JVM lifetime, which leaks in + // long-lived drivers. + // + // Initialisation order: the parent constructor runs validateTransportSupport + // (against the spec passed as a parameter, not subclass state) BEFORE this + // val is evaluated. If validation throws, this field is never assigned and + // no temp directory is created, which is the desired behaviour. Keep + // future field initialisers below this line aware of the same ordering + // contract -- the parent constructor cannot read them, and they cannot + // assume validateTransportSupport has not already failed. + private val socketDir: Path = createPrivateTempDirectory() + + /** + * 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}") + // BIDI is the only pattern the gRPC `Execute` RPC speaks. If the spec Review Comment: fixed ########## udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/DirectGrpcDispatcher.scala: ########## @@ -0,0 +1,237 @@ +/* + * 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. + * + * gRPC over UDS is the only protocol/transport combination shipped today. + * 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 + + // Removed explicitly in closeTransport(). deleteOnExit is avoided because + // the JDK retains the path for the JVM lifetime, which leaks in + // long-lived drivers. + // + // Initialisation order: the parent constructor runs validateTransportSupport + // (against the spec passed as a parameter, not subclass state) BEFORE this + // val is evaluated. If validation throws, this field is never assigned and + // no temp directory is created, which is the desired behaviour. Keep + // future field initialisers below this line aware of the same ordering + // contract -- the parent constructor cannot read them, and they cannot + // assume validateTransportSupport has not already failed. + private val socketDir: Path = createPrivateTempDirectory() + + /** + * 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}") + // BIDI is the only pattern the gRPC `Execute` RPC speaks. If the spec + // declares capabilities, it MUST include BIDI; an unset capabilities + // block is treated as "no constraint" and accepted. + if (spec.hasCapabilities) { Review Comment: fixed -- 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]
