This is an automated email from the ASF dual-hosted git repository.
uranusjr pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 474e5b598f5 Harden frame processing on large data (#69125)
474e5b598f5 is described below
commit 474e5b598f54f45d7b189aa3ff08c132ec388683
Author: Tzu-ping Chung <[email protected]>
AuthorDate: Fri Aug 21 07:48:10 2026 +0800
Harden frame processing on large data (#69125)
---
.../main/kotlin/org/apache/airflow/sdk/Server.kt | 26 ++--
.../org/apache/airflow/sdk/execution/Comm.kt | 148 +++++++++++++++++++--
.../org/apache/airflow/sdk/execution/Frame.kt | 46 +++++--
.../org/apache/airflow/sdk/execution/MsgPack.kt | 15 ++-
.../org/apache/airflow/sdk/execution/Task.kt | 3 +
.../kotlin/org/apache/airflow/sdk/ServerTest.kt | 4 +-
.../org/apache/airflow/sdk/execution/CommTest.kt | 13 +-
7 files changed, 205 insertions(+), 50 deletions(-)
diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Server.kt
b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Server.kt
index 654baba944b..361fd0fe319 100644
--- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Server.kt
+++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Server.kt
@@ -147,13 +147,15 @@ class Server(
launch {
try {
- aSocket(SelectorManager(Dispatchers.IO)).tcp().connect(comm).use {
socket ->
- logger.debug("Connected comm", mapOf("addr" to comm))
- CoordinatorComm(
- socket.openReadChannel(),
- socket.openWriteChannel(autoFlush = true),
- ).use { coordinator ->
- dispatchTask(bundle, coordinator)
+ SelectorManager(Dispatchers.IO).use { selector ->
+ aSocket(selector).tcp().connect(comm).use { socket ->
+ logger.debug("Connected comm", mapOf("addr" to comm))
+ CoordinatorComm(
+ socket.openReadChannel(),
+ socket.openWriteChannel(autoFlush = true),
+ ).use { coordinator ->
+ dispatchTask(bundle, coordinator)
+ }
}
}
} finally {
@@ -161,10 +163,12 @@ class Server(
}
}
launch {
- aSocket(SelectorManager(Dispatchers.IO)).tcp().connect(logs).use {
socket ->
- logger.debug("Connected logs", mapOf("addr" to logs))
- LogSender.configure(socket.openWriteChannel(autoFlush = true))
- deferral.await()
+ SelectorManager(Dispatchers.IO).use { selector ->
+ aSocket(selector).tcp().connect(logs).use { socket ->
+ logger.debug("Connected logs", mapOf("addr" to logs))
+ LogSender.configure(socket.openWriteChannel(autoFlush = true))
+ deferral.await()
+ }
}
}
}
diff --git
a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Comm.kt
b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Comm.kt
index 34815d7b7f1..19bf7922462 100644
--- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Comm.kt
+++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Comm.kt
@@ -27,16 +27,66 @@ import kotlinx.coroutines.CancellationException
import kotlinx.coroutines.CompletableDeferred
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.Job
import kotlinx.coroutines.SupervisorJob
import kotlinx.coroutines.cancel
+import kotlinx.coroutines.currentCoroutineContext
+import kotlinx.coroutines.ensureActive
import kotlinx.coroutines.launch
+import kotlinx.coroutines.runBlocking
import kotlinx.coroutines.sync.Mutex
import kotlinx.coroutines.sync.withLock
import org.apache.airflow.sdk.ApiError
import org.apache.airflow.sdk.execution.comm.ErrorResponse
+import org.msgpack.core.buffer.MessageBuffer
+import org.msgpack.core.buffer.MessageBufferInput
+import java.io.IOException
import java.util.concurrent.ConcurrentHashMap
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.
+ *
+ * Up to [CHUNK_SIZE] bytes are read from [reader] per chunk, [declaredLength]
+ * bytes in total. This bounds only the transport read buffer so a frame larger
+ * than [Int.MAX_VALUE] can decode without one giant allocation.
+ *
+ * The MessageBufferInput contract is synchronous while the underlying read
+ * suspends, so each [next] bridges with [runBlocking]. This is fine since we
+ * only use this class with `Dispatchers.IO`, which is capable of blocking.
+ */
+private class ChannelFrameInput(
+ private val reader: ByteReadChannel,
+ declaredLength: UInt,
+ private val coroutineContext: CoroutineContext,
+) : MessageBufferInput {
+ private companion object {
+ const val CHUNK_SIZE = (64 * 1024).toLong()
+ }
+
+ var remaining = declaredLength.toLong()
+ private set
+
+ override fun next(): MessageBuffer? {
+ if (remaining <= 0L) return null
+ coroutineContext.ensureActive()
+ val array =
+ minOf(remaining, CHUNK_SIZE).toInt().let { want ->
+ runBlocking(coroutineContext[Job] ?: EmptyCoroutineContext) {
+ reader.readByteArray(want)
+ }.apply {
+ if (size != want) throw IOException("Truncated frame: expected $want
more bytes, got $size")
+ remaining -= want
+ }
+ }
+ return MessageBuffer.wrap(array)
+ }
+
+ override fun close() {} // No cleanup here. The caller owns the channel's
lifecycle.
+}
data class IncomingFrame(
val id: Int,
@@ -48,6 +98,11 @@ data class OutgoingFrame(
val body: Any,
)
+internal class FrameProcessingException(
+ message: String,
+ cause: Throwable? = null,
+) : Exception(message, cause)
+
@OptIn(ExperimentalAtomicApi::class)
class CoordinatorComm(
private val reader: ByteReadChannel,
@@ -56,6 +111,27 @@ class CoordinatorComm(
internal companion object {
private val logger = Logger(CoordinatorComm::class)
+ // A frame can at most use 1/8 of available memory. This is a magic number.
+ private const val MAX_HEAP_FRACTION_PER_FRAME = 8L
+
+ /**
+ * Upper bound on the wire size of a single inbound frame.
+ *
+ * This is used to reject oversized frames before [Frame.decodeRaw]
+ * allocates. An otherwise unrecoverable OOM can therefore be turned into a
+ * catchable [FrameProcessingException].
+ *
+ * Deriving from [Runtime.maxMemory] means this value respects `-Xmx`. The
+ * value is always clamped to [Frame.MAX_FRAME_LENGTH] from the protocol.
+ */
+ private val MAX_INBOUND_FRAME_SIZE: Long =
+ Runtime
+ .getRuntime()
+ .maxMemory()
+ .let {
+ if (it == Long.MAX_VALUE) Frame.MAX_FRAME_LENGTH else it /
MAX_HEAP_FRACTION_PER_FRAME
+ }.coerceAtMost(Frame.MAX_FRAME_LENGTH)
+
fun encode(outgoing: OutgoingFrame) = Frame.encodeRequest(outgoing.id,
outgoing.body)
fun decode(bytes: ByteArray) = Frame.decode(bytes)
@@ -82,10 +158,13 @@ class CoordinatorComm(
id: Int,
body: Any,
) {
- val data = encode(OutgoingFrame(id, body))
+ val buffers = encode(OutgoingFrame(id, body))
logger.debug("Sending", mapOf("id" to id, "body" to body))
writeMutex.withLock {
- writer.writeByteArray(Frame.lengthPrefix(data.size) + data)
+ writer.writeByteArray(Frame.lengthPrefix(Frame.payloadLength(buffers)))
+ for (buffer in buffers) {
+ writer.writeByteArray(buffer.toByteArray())
+ }
}
}
@@ -138,30 +217,75 @@ class CoordinatorComm(
failAllWaiters(readError ?: ApiError("Coordinator comm closed"))
}
- private suspend fun readPayload(): ByteArray {
+ /**
+ * Read one frame envelope off the wire, streaming its payload in bounded
+ * chunks so an oversized frame cannot exhaust the heap in a single read.
+ *
+ * The declared length is rejected up front when it exceeds
+ * [MAX_INBOUND_FRAME_SIZE], and an under-run payload (the prefix promised
+ * more than the encoded body actually consumed) is treated as a stream
+ * desync. Both surface as a [FrameProcessingException] rather than a silent
+ * shutdown.
+ */
+ private suspend fun readRawFrame(): RawFrame {
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 declaredLength = Frame.parseLengthPrefix(prefix)
+ if (declaredLength.toLong() > MAX_INBOUND_FRAME_SIZE) {
+ logger.error(
+ "Inbound frame exceeds size limit",
+ mapOf("length" to declaredLength, "limit" to MAX_INBOUND_FRAME_SIZE),
+ )
+ throw FrameProcessingException(
+ "Inbound frame of $declaredLength bytes exceeds limit of
$MAX_INBOUND_FRAME_SIZE bytes",
+ )
+ }
+
+ val input = ChannelFrameInput(reader, declaredLength,
currentCoroutineContext())
+ val raw =
+ try {
+ Frame.decodeRaw(input)
+ } catch (e: CancellationException) {
+ throw e // Let coroutine cancellation propagate so the task coroutine
unwinds.
+ } catch (e: IOException) {
+ logger.error(
+ "Failed to read frame",
+ mapOf("length" to declaredLength, "exception" to e),
+ )
+ throw FrameProcessingException("Failed to read frame of
$declaredLength bytes", e)
+ } catch (e: Exception) {
+ logger.error(
+ "Failed to decode frame",
+ mapOf("length" to declaredLength, "exception" to e),
+ )
+ throw FrameProcessingException("Failed to decode frame of
$declaredLength bytes", e)
+ }
+
+ if (input.remaining != 0L) {
+ logger.error(
+ "Frame length prefix overran the payload",
+ mapOf("length" to declaredLength, "undrained" to input.remaining),
+ )
+ throw FrameProcessingException("Frame declared $declaredLength bytes but
${input.remaining} left unread")
}
- return payload
+
+ return raw
}
private suspend fun readFrame(): IncomingFrame {
- val frame = decode(readPayload())
- logger.debug("Received", mapOf("id" to frame.id))
- return frame
+ val raw = readRawFrame()
+ logger.debug("Received", mapOf("id" to raw.id))
+ return IncomingFrame(raw.id, Frame.decodeBody(raw))
}
private suspend fun readLoop() {
while (true) {
val raw =
try {
- Frame.decodeRaw(readPayload())
+ readRawFrame()
} catch (e: CancellationException) {
// Coroutine cancellation is delivered by throwing this exception,
and
// cooperative cancellation requires rethrowing it so it propagates.
diff --git
a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Frame.kt
b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Frame.kt
index 4d488fae0fc..035cfb89114 100644
--- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Frame.kt
+++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Frame.kt
@@ -26,7 +26,10 @@ import com.fasterxml.jackson.databind.util.StdDateFormat
import com.fasterxml.jackson.datatype.jsr310.JavaTimeModule
import org.apache.airflow.sdk.execution.comm.Discriminator
import org.msgpack.core.MessagePack
-import java.io.ByteArrayOutputStream
+import org.msgpack.core.MessageUnpacker
+import org.msgpack.core.buffer.ArrayBufferInput
+import org.msgpack.core.buffer.MessageBuffer
+import org.msgpack.core.buffer.MessageBufferInput
data class RawFrame(
val id: Int,
@@ -35,6 +38,8 @@ data class RawFrame(
)
object Frame {
+ internal const val MAX_FRAME_LENGTH = 0xFFFF_FFFFL
+
private val mapper =
ObjectMapper().apply {
configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false)
@@ -47,29 +52,35 @@ object Frame {
fun encodeRequest(
id: Int,
body: Any,
- ): ByteArray = encodeFrame(id, body)
+ ): List<MessageBuffer> = encodeFrame(id, body)
+
+ /**
+ * Decode the structural envelope of a frame, streaming the payload from
+ * [input] so a large frame never has to be held in one contiguous array.
+ */
+ fun decodeRaw(input: MessageBufferInput): RawFrame =
MessagePack.newDefaultUnpacker(input).use { decodeRawFrom(it) }
- fun decodeRaw(bytes: ByteArray): RawFrame {
- val unpacker = MessagePack.newDefaultUnpacker(bytes)
+ private fun decodeRawFrom(unpacker: MessageUnpacker): RawFrame {
val headerSize = unpacker.unpackArrayHeader()
check(headerSize >= 1) { "Unexpected Task SDK frame arity $headerSize" }
val id = unpacker.unpackInt()
val rawBody = if (headerSize >= 2) unpacker.unpackAny() else null
val rawError = if (headerSize >= 3) unpacker.unpackAny() else null
- unpacker.close()
return RawFrame(id, rawBody, rawError)
}
fun decodeBody(raw: RawFrame): Any? = decodeMessage(raw.rawError) ?:
decodeMessage(raw.rawBody)
- fun decode(bytes: ByteArray): IncomingFrame {
- val raw = decodeRaw(bytes)
+ fun decode(input: MessageBufferInput): IncomingFrame {
+ val raw = decodeRaw(input)
return IncomingFrame(raw.id, decodeBody(raw))
}
- fun lengthPrefix(length: Int) =
+ fun decode(bytes: ByteArray): IncomingFrame = decode(ArrayBufferInput(bytes))
+
+ fun lengthPrefix(length: UInt) =
byteArrayOf(
(length shr 24).toByte(),
(length shr 16).toByte(),
@@ -77,22 +88,29 @@ object Frame {
length.toByte(),
)
- fun parseLengthPrefix(prefix: ByteArray): Int {
+ fun payloadLength(buffers: List<MessageBuffer>): UInt {
+ val total = buffers.sumOf { it.size().toLong() }
+ require(total <= MAX_FRAME_LENGTH) {
+ "Frame payload $total bytes exceeds protocol maximum $MAX_FRAME_LENGTH"
+ }
+ return total.toUInt()
+ }
+
+ fun parseLengthPrefix(prefix: ByteArray): UInt {
check(prefix.size == 4) { "Need 4 prefix bytes" }
- return prefix.fold(0) { acc, byte -> (acc shl 8) or (byte.toInt() and
0xff) }
+ return prefix.fold(0u) { acc, byte -> (acc shl 8) or (byte.toUInt() and
0xffu) }
}
private fun encodeFrame(
id: Int,
body: Any?,
- ): ByteArray {
- val payload = ByteArrayOutputStream()
- val packer = MessagePack.newDefaultPacker(payload)
+ ): List<MessageBuffer> {
+ val packer = MessagePack.newDefaultBufferPacker()
packer.packArrayHeader(2)
packer.packInt(id)
packer.packAny(body?.let(::toBody))
packer.close()
- return payload.toByteArray()
+ return packer.toBufferList()
}
private fun decodeMessage(raw: Any?): Any? {
diff --git
a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/MsgPack.kt
b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/MsgPack.kt
index 88fa3f632d9..48a9c259ba5 100644
--- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/MsgPack.kt
+++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/MsgPack.kt
@@ -92,10 +92,11 @@ private fun MapValue.decodeMap(): Map<*, *> =
private fun ExtensionValue.decodeExtensionValue(): Any? =
when (type) {
TimestampToJavaOffsetDateTimeModule.EXT_TYPE -> {
- MessagePack
- .newDefaultUnpacker(data)
- .unpackTimestamp(ExtensionTypeHeader(type, data.size))
- .atOffset(ZoneOffset.UTC)
+ MessagePack.newDefaultUnpacker(data).use { unpacker ->
+ unpacker
+ .unpackTimestamp(ExtensionTypeHeader(type, data.size))
+ .atOffset(ZoneOffset.UTC)
+ }
}
else -> throw IllegalArgumentException("Unsupported extension type: $this")
}
@@ -143,8 +144,10 @@ class TimestampToJavaOffsetDateTimeModule : SimpleModule()
{
if (ext.type != EXT_TYPE) {
return null
}
- val unpacker = MessagePack.newDefaultUnpacker(ext.data)
- val instant = unpacker.unpackTimestamp(ExtensionTypeHeader(EXT_TYPE,
ext.data.size))
+ val instant =
+ MessagePack.newDefaultUnpacker(ext.data).use { unpacker ->
+ unpacker.unpackTimestamp(ExtensionTypeHeader(EXT_TYPE,
ext.data.size))
+ }
return instant.atOffset(ZoneOffset.UTC)
}
}
diff --git
a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Task.kt
b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Task.kt
index 60258dacc62..3685644c0b7 100644
--- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Task.kt
+++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Task.kt
@@ -19,6 +19,7 @@
package org.apache.airflow.sdk.execution
+import kotlinx.coroutines.CancellationException
import org.apache.airflow.sdk.Bundle
import org.apache.airflow.sdk.Client
import org.apache.airflow.sdk.Context
@@ -74,6 +75,8 @@ internal object TaskRunner {
return try {
task.getDeclaredConstructor().newInstance().execute(Context.from(request),
client)
TaskResult.success()
+ } catch (e: CancellationException) {
+ throw e // Let coroutine cancellation propagate so the task coroutine
unwinds.
} catch (e: Throwable) {
logger.error("Error executing task", mapOf("ti" to request.ti, "error"
to e, "trace" to e.stackTraceToString()))
e.printStackTrace()
diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ServerTest.kt
b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ServerTest.kt
index 6cae7a4c27b..e436dd2fc61 100644
--- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ServerTest.kt
+++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ServerTest.kt
@@ -56,7 +56,7 @@ class ServerTest {
}
private suspend fun ByteChannel.writeFrame(payload: ByteArray) {
- writeByteArray(Frame.lengthPrefix(payload.size))
+ writeByteArray(Frame.lengthPrefix(payload.size.toUInt()))
writeByteArray(payload)
}
@@ -76,7 +76,7 @@ class ServerTest {
// Deliver the StartupDetails frame (id 2, dag_id "c", task_id "a").
toServer.writeFrame(hexToBytes(STARTUP_HEX))
val prefix = fromServer.readByteArray(4)
- val payload =
fromServer.readByteArray(Frame.parseLengthPrefix(prefix))
+ val payload =
fromServer.readByteArray(Frame.parseLengthPrefix(prefix).toInt())
val result = CoordinatorComm.decode(payload)
reported.put(result)
toServer.writeFrame(ackFrame(result.id))
diff --git
a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/CommTest.kt
b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/CommTest.kt
index 978e377f52d..ef0459b0922 100644
--- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/CommTest.kt
+++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/CommTest.kt
@@ -87,7 +87,10 @@ class CommsTest {
@DisplayName("Should serialize all fields")
fun shouldEncodeSucceedTask() {
val endDate = OffsetDateTime.of(2024, 12, 1, 1, 0, 0, 0, ZoneOffset.UTC)
- val bytes = CoordinatorComm.encode(OutgoingFrame(3,
TaskResult.success(endDate = endDate)))
+ val bytes =
+ CoordinatorComm
+ .encode(OutgoingFrame(3, TaskResult.success(endDate = endDate)))
+ .fold(ByteArray(0)) { acc, buffer -> acc + buffer.toByteArray() }
val actual = bytes.toHexString(HexFormat { bytes { byteSeparator = " " } })
val expected =
@@ -134,13 +137,13 @@ class CommsTest {
}
private suspend fun ByteChannel.writeFrame(payload: ByteArray) {
- writeByteArray(Frame.lengthPrefix(payload.size))
+ writeByteArray(Frame.lengthPrefix(payload.size.toUInt()))
writeByteArray(payload)
}
private suspend fun ByteChannel.readOneRequest() {
val prefix = readByteArray(4)
- readByteArray(Frame.parseLengthPrefix(prefix))
+ readByteArray(Frame.parseLengthPrefix(prefix).toInt())
}
@Test
@@ -180,7 +183,7 @@ class CommsTest {
val ids =
(0 until n).map {
val prefix = fromClient.readByteArray(4)
- val payload =
fromClient.readByteArray(Frame.parseLengthPrefix(prefix))
+ val payload =
fromClient.readByteArray(Frame.parseLengthPrefix(prefix).toInt())
CoordinatorComm.decode(payload).id
}
ids.reversed().forEach { toClient.writeFrame(responseFrame(it)) }
@@ -232,7 +235,7 @@ class CommsTest {
runBlocking {
repeat(n) {
val prefix = fromClient.readByteArray(4)
- val payload =
fromClient.readByteArray(Frame.parseLengthPrefix(prefix))
+ val payload =
fromClient.readByteArray(Frame.parseLengthPrefix(prefix).toInt())
toClient.writeFrame(responseFrame(CoordinatorComm.decode(payload).id))
}
}
