jason810496 commented on code in PR #69254:
URL: https://github.com/apache/airflow/pull/69254#discussion_r3533980025
##########
java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Comm.kt:
##########
@@ -138,4 +119,57 @@ class CoordinatorComm(
else -> throw ApiError("Unexpected response type
${response::class.java}")
}
}
+
+ override fun close() {
+ dispatcherScope.cancel()
+ }
+
+ private suspend fun readFrame(): IncomingFrame {
+ val prefix = reader.readByteArray(4) // First 4 bytes as length.
+ if (prefix.size != 4) {
+ throw ApiError("Coordinator socket closed while reading frame length")
+ }
+ val payloadLength = Frame.parseLengthPrefix(prefix)
+ val payload = reader.readByteArray(payloadLength)
+ if (payload.size != payloadLength) {
+ throw ApiError("Coordinator socket closed while reading frame payload")
+ }
+ val frame = decode(payload)
+ logger.debug("Received", mapOf("id" to frame.id))
+ return frame
+ }
+
+ private suspend fun readLoop() {
+ while (true) {
+ val frame =
+ try {
+ readFrame()
+ } catch (e: CancellationException) {
+ // Coroutine cancellation is delivered by throwing this exception,
and
+ // cooperative cancellation requires rethrowing it so it propagates.
+ throw e
+ } catch (e: Throwable) {
+ failAllWaiters(e)
Review Comment:
**One undecodable frame tears down every in-flight request (major).**
`readFrame()` -> `decode` -> `Discriminator.discriminate` throws
`IllegalStateException` for any unknown message `type` (in `Frame.kt`), which
lands in this `catch (Throwable)` and calls `failAllWaiters`, failing *all*
pending requests rather than just the offending id. Under the old one-shot
protocol a bad frame only killed a single exchange; under multiplexing one
forward-incompatible / unknown message type kills the whole coordinator
connection. Consider decoding the frame id first and scoping an undecodable
body to that id, or making an unknown `type` a non-fatal sentinel.
##########
java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Comm.kt:
##########
@@ -138,4 +119,57 @@ class CoordinatorComm(
else -> throw ApiError("Unexpected response type
${response::class.java}")
}
}
+
+ override fun close() {
+ dispatcherScope.cancel()
+ }
Review Comment:
**`close()` doesn't fail in-flight waiters -> caller hangs forever (major).**
`close()` only cancels `dispatcherScope`; `waiter.await()` runs in the
*caller's* scope, and `readLoop` rethrows `CancellationException` so
`failAllWaiters` never runs. A `close()` during an in-flight request (the exact
`.use { }` error path in `Server.serveAsync`) leaves the caller blocked and can
keep the JVM alive. Drain `pending` on close:
```suggestion
override fun close() {
dispatcherScope.cancel()
val error = readError ?: ApiError("Coordinator comm closed")
val waiters = pending.values.toList()
pending.clear()
waiters.forEach { it.completeExceptionally(error) }
}
```
Note: this drains `pending` outside `stateMutex` (close is not `suspend`),
which races with `readLoop`'s `pending.remove`. Recommend also backing
`pending` with a `java.util.concurrent.ConcurrentHashMap` (all current call
sites are compatible) so the drain is race-free.
##########
java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Comm.kt:
##########
@@ -138,4 +119,57 @@ class CoordinatorComm(
else -> throw ApiError("Unexpected response type
${response::class.java}")
}
}
+
+ override fun close() {
+ dispatcherScope.cancel()
+ }
+
+ private suspend fun readFrame(): IncomingFrame {
+ val prefix = reader.readByteArray(4) // First 4 bytes as length.
+ if (prefix.size != 4) {
+ throw ApiError("Coordinator socket closed while reading frame length")
+ }
+ val payloadLength = Frame.parseLengthPrefix(prefix)
+ val payload = reader.readByteArray(payloadLength)
+ if (payload.size != payloadLength) {
+ throw ApiError("Coordinator socket closed while reading frame payload")
+ }
+ val frame = decode(payload)
+ logger.debug("Received", mapOf("id" to frame.id))
+ return frame
+ }
+
+ private suspend fun readLoop() {
+ while (true) {
+ val frame =
+ try {
+ readFrame()
+ } catch (e: CancellationException) {
+ // Coroutine cancellation is delivered by throwing this exception,
and
+ // cooperative cancellation requires rethrowing it so it propagates.
+ throw e
+ } catch (e: Throwable) {
+ failAllWaiters(e)
+ return
+ }
+ val waiter = stateMutex.withLock { pending.remove(frame.id) }
+ if (waiter == null) {
+ logger.warning("Discarding response with no matching request",
mapOf("id" to frame.id))
+ continue
+ }
+ waiter.complete(frame)
+ }
+ }
+
+ private suspend fun failAllWaiters(cause: Throwable) {
+ val error = cause as? ApiError ?: ApiError("Coordinator comm closed:
${cause.message}")
Review Comment:
**Preserve the original cause (minor).**
Wrapping a non-`ApiError` cause into a fresh `ApiError` drops the original
exception and its stack trace, hurting debuggability of the concurrency paths.
`ApiError` has no cause-accepting constructor, but `initCause` works:
```suggestion
val error = cause as? ApiError ?: ApiError("Coordinator comm closed:
${cause.message}").apply { initCause(cause) }
```
##########
java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Comm.kt:
##########
@@ -57,77 +61,54 @@ class CoordinatorComm(
}
private val nextId = AtomicInt(0)
- private var shutDownRequested = false
- private val commMutex = Mutex()
-
- suspend fun startProcessing() {
- while (!shutDownRequested) {
- processOnce(::handleIncoming)
- }
- logger.debug("Goodbye")
- }
-
- private suspend fun processOnce(handle: suspend (IncomingFrame) -> Unit) {
- val prefix = reader.readByteArray(4) // First 4 bytes as length.
- if (prefix.size != 4) { // Something is terribly wrong. Let's bail.
- logger.error("Need 4 prefix bytes", mapOf("actual" to prefix.size))
- shutDownRequested = true
- return
- }
-
- val payloadLength = Frame.parseLengthPrefix(prefix)
- val payload = reader.readByteArray(payloadLength)
- if (payload.size != payloadLength) { // Something is terribly wrong. Let's
bail.
- logger.error(
- "Payload length not right",
- mapOf("expect" to payloadLength, "receive" to payload.size),
- )
- shutDownRequested = true
- return
+ private val writeMutex = Mutex()
+ private val stateMutex = Mutex()
+ private val pending = mutableMapOf<Int, CompletableDeferred<IncomingFrame>>()
+ private var dispatcherStarted = false
+ private var readError: ApiError? = null
+
+ private val dispatcherScope = CoroutineScope(Dispatchers.IO +
SupervisorJob())
+
+ suspend fun readMessage(): IncomingFrame =
+ stateMutex.withLock {
+ check(!dispatcherStarted) { "readMessage cannot be used after the
dispatcher has started" }
+ readFrame()
}
- val frame = decode(payload)
- logger.debug("Handling", mapOf("id" to frame.id))
- handle(frame)
- }
private suspend fun sendMessage(
id: Int,
body: Any,
) {
val data = encode(OutgoingFrame(id, body))
logger.debug("Sending", mapOf("id" to id, "body" to body))
- writer.writeByteArray(Frame.lengthPrefix(data.size))
- writer.writeByteArray(data)
- }
-
- suspend fun handleIncoming(frame: IncomingFrame) {
- when (val request = frame.body) {
- null -> {}
- is ErrorResponse -> throw ApiError("[${request.error}]
${request.detail}")
- is StartupDetails -> {
- communicate<Unit>(runTask(bundle, request, this))
- shutDownRequested = true
- }
+ writeMutex.withLock {
+ writer.writeByteArray(Frame.lengthPrefix(data.size))
+ writer.writeByteArray(data)
}
Review Comment:
**Partial write can corrupt the shared framed stream (major).**
The length prefix and payload are two separate `writeByteArray` calls; if
the second throws after the first succeeded, the shared writer is left
mid-frame and every other concurrent multiplexed request reads a misaligned
stream. Concatenate into a single write so the framed unit is emitted
atomically under `writeMutex`:
```suggestion
writeMutex.withLock {
writer.writeByteArray(Frame.lengthPrefix(data.size) + data)
}
```
##########
java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Server.kt:
##########
@@ -164,4 +168,25 @@ class Server(
}
}
}
+
+ internal suspend fun dispatchTask(
+ bundle: Bundle,
+ coordinator: CoordinatorComm,
+ ) {
+ val frame = coordinator.readMessage()
+ when (val body = frame.body) {
+ is StartupDetails -> runTaskAndReport(bundle, body, coordinator)
+ is ErrorResponse -> throw ApiError("[${body.error}] ${body.detail}")
+ else -> logger.debug("Ignoring unexpected initial frame", mapOf("id" to
frame.id))
Review Comment:
**Unexpected initial frame is silently swallowed (major).**
The `else` branch only logs at debug and returns, so `dispatchTask` ends
with no terminal state reported to the coordinator and the task is silently
lost. Surface it as an error, mirroring the `ErrorResponse` branch above:
```suggestion
else -> throw ApiError("Unexpected initial frame (id=${frame.id})")
```
--
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]
