jason810496 commented on code in PR #69125:
URL: https://github.com/apache/airflow/pull/69125#discussion_r3503398587
##########
java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Comm.kt:
##########
@@ -72,17 +115,20 @@ class CoordinatorComm(
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
- }
- val frame = decode(payload)
+ val declaredLength = Frame.parseLengthPrefix(prefix)
+ val frame =
+ try {
+ Frame.decode(ChannelFrameInput(reader, declaredLength,
currentCoroutineContext()))
+ } catch (e: CancellationException) {
+ throw e // Let coroutine cancellation propagate so the task coroutine
unwinds.
+ } catch (e: Exception) {
+ logger.error(
+ "Failed to read or decode frame",
+ mapOf("length" to declaredLength, "exception" to e),
+ )
+ shutDownRequested = true
+ return
+ }
Review Comment:
IIUC, would it be better to propagate the decode failures or force a
non-zero exit?
Decode failure on the initial supervisor frame (unsupported schema /
malformed msgpack) sets `shutDownRequested` and returns, so `startProcessing`
exits normally with code 0 before any terminal state is sent.
https://github.com/apache/airflow/blob/f24c5b0bc7933cc4856bebf7e3bc78af4db9e84e/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Server.kt#L145-L158
Though the scheduler will still mark the TI state correctly as fail / up for
retry during reconciliation, so it's also okey to keep it as-is (but the exit
code might be confusing).
##########
java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Comm.kt:
##########
@@ -23,23 +23,70 @@ import io.ktor.utils.io.ByteReadChannel
import io.ktor.utils.io.ByteWriteChannel
import io.ktor.utils.io.readByteArray
import io.ktor.utils.io.writeByteArray
+import kotlinx.coroutines.CancellationException
+import kotlinx.coroutines.Job
+import kotlinx.coroutines.currentCoroutineContext
+import kotlinx.coroutines.ensureActive
+import kotlinx.coroutines.runBlocking
import org.apache.airflow.sdk.ApiError
import org.apache.airflow.sdk.Bundle
import org.apache.airflow.sdk.execution.comm.ErrorResponse
import org.apache.airflow.sdk.execution.comm.StartupDetails
+import org.msgpack.core.buffer.MessageBuffer
+import org.msgpack.core.buffer.MessageBufferInput
+import java.io.IOException
import kotlin.concurrent.atomics.AtomicInt
import kotlin.concurrent.atomics.ExperimentalAtomicApi
+import kotlin.coroutines.CoroutineContext
+import kotlin.coroutines.EmptyCoroutineContext
+
+/**
+ * A [MessageBufferInput] that feeds a MessageUnpacker in chunks.
+ *
+ * Exactly [declaredLength] bytes is fed from [reader] in each chunk.
+ * Heap use is bounded by [CHUNK_SIZE], so a frame larger than
Review Comment:
I think it makes sense to respect the `maxMemory()` even we restrict there
at most 4GB memory usage (2^32). Though it _likely_ not happen at all.
Here's the context that I discuss with CC:
---
This "heap bounded by `CHUNK_SIZE`" only covers the transport read buffer.
`decodeFrom` still calls `unpackAny()` (`MsgPack.kt`,
`unpackValue().decode()`), which fully materializes the body into a
`Map`/`List` graph. And since `parseLengthPrefix` accepts the full `UInt` range
(~4 GiB) with no cap on the receive side (`MAX_FRAME_LENGTH` is only enforced
on the send side in `payloadLength()`), decode can still pull in up to ~4 GiB
of wire data - and the decoded object graph is a multiple of that (boxing,
`HashMap`/`ArrayList` overhead, UTF-16 strings). Chunking only avoids
allocating one contiguous `ByteArray` (its length is an `Int`, hence ~2 GiB
max); it does not bound total heap. So an oversized
XCom/Variable/StartupDetails response OOMs the worker JVM instead of failing
with a bounded framing error.
Could we gate on a max inbound body size and reject before decode? Rather
than a hardcoded constant, we could derive it from the runtime - a conservative
fraction of `Runtime.getRuntime().maxMemory()` - so it respects the operator's
`-Xmx` and turns an unrecoverable OOM into a catchable framing error. Declared
length is only a proxy (the decoded graph is larger), so the fraction should be
conservative, and `maxMemory()` is more reliable to reason against than
`freeMemory()` (which is racy and only reflects the currently-committed heap).
Note msgpack-core 0.9.11 only exposes `withStringSizeLimit` (default ~2 GiB) -
no array/map/binary limits - so the library won't bound total heap on its own.
Separately, the KDoc could say the *transport buffer* (not heap) is bounded
by `CHUNK_SIZE`.
##########
java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Comm.kt:
##########
@@ -72,17 +115,20 @@ class CoordinatorComm(
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
- }
- val frame = decode(payload)
+ val declaredLength = Frame.parseLengthPrefix(prefix)
+ val frame =
+ try {
+ Frame.decode(ChannelFrameInput(reader, declaredLength,
currentCoroutineContext()))
Review Comment:
When the 4-byte prefix declares more bytes than the encoded msgpack object,
`ChannelFrameInput` only pulls the chunks the unpacker actually consumes, and
`close()` is a no-op that doesn't drain the rest. So when `declaredLength >
CHUNK_SIZE` (64 KiB) and the declared length overruns the real body, the
leftover bytes stay in the ktor channel and the next `processOnce` reads them
as a bogus length prefix -> stream desync.
The previous code read exactly `payloadLength` bytes, so it always advanced
the channel by the full declared length. This looks like a regression precisely
on the malformed/oversized path this PR is meant to harden. Could we drain (or
validate) the full declared length before returning to the prefix loop? Normal
frames (prefix == encoded size) and <=64 KiB frames aren't affected.
---
Drafted-by: Claude Code (Opus 4.8); reviewed by @jason810496 before posting
--
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]
