Re: [PR] Harden frame processing on large data [airflow]

["jason810496 (via GitHub)" <gi...@apache.org>]

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: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org