From 4dfed47766f27796f5f258ff978031c651ab0941 Mon Sep 17 00:00:00 2001 From: PoAn Yang Date: Fri, 16 Oct 2026 07:07:21 +0800 Subject: [PATCH] [v3-4-test] Java SDK: Reject byte arrays instead of letting the task hang (#74218) (cherry picked from commit 665b94999c38426f794cc95d8644028189882ead) Co-authored-by: PoAn Yang Signed-off-by: PoAn Yang --- .../kotlin/org/apache/airflow/sdk/Client.kt | 5 ++- .../apache/airflow/sdk/execution/MsgPack.kt | 13 +++---- .../apache/airflow/sdk/execution/CommTest.kt | 34 +++++++++++++++++++ 3 files changed, 45 insertions(+), 7 deletions(-) diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Client.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Client.kt index c394e270436d8..fa3355fc9615a 100644 --- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Client.kt +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Client.kt @@ -179,8 +179,11 @@ class Client internal constructor( * Pushes an XCom value for downstream tasks to read. * * @param key XCom key; defaults to [XCOM_RETURN_KEY]. - * @param value Value to push. Must be JSON-serializable. + * @param value Value to push. Must be JSON-serializable. The SDK rejects byte + * arrays, so encode binary data as a string first, for example with Base64. * @throws ApiError if the API call fails. + * @throws IllegalArgumentException if [value] is or contains a type that the + * SDK cannot send, such as a byte array. */ @JvmOverloads fun setXCom( key: String = XCOM_RETURN_KEY, 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 48a9c259ba5b9..31764e28732ca 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 @@ -41,11 +41,6 @@ import java.math.BigInteger import java.time.OffsetDateTime import java.time.ZoneOffset -private fun MessagePacker.packByteArray(data: ByteArray) { - packBinaryHeader(data.size) - data.forEach { packByte(it) } -} - private fun MessagePacker.packMap(data: Map<*, *>) { packMapHeader(data.size) data.forEach { (k, v) -> @@ -71,7 +66,13 @@ fun MessagePacker.packAny(data: Any?) { is BigInteger -> packBigInteger(data) is Float -> packFloat(data) is Double -> packDouble(data) - is ByteArray -> packByteArray(data) + // The supervisor rejects bytes in values such as XCom values. It does not reply + // to a request it cannot decode, so the task that sent the request would wait forever. + is ByteArray -> + throw IllegalArgumentException( + "Airflow does not accept byte arrays because they are not JSON values. " + + "Encode the bytes as a string first, for example with java.util.Base64.", + ) is String -> packString(data) is Map<*, *> -> packMap(data) is Collection<*> -> packCollection(data) 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 20fba213948f5..b6c12c36a572f 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 @@ -20,6 +20,7 @@ package org.apache.airflow.sdk.execution import io.ktor.utils.io.ByteChannel +import io.ktor.utils.io.exhausted import io.ktor.utils.io.readByteArray import io.ktor.utils.io.writeByteArray import kotlinx.coroutines.async @@ -679,6 +680,39 @@ class CommsTest { comm.close() } + @Test + @DisplayName("setXCom rejects a byte array nested in the value without writing anything to the supervisor") + @Timeout(value = 30, unit = TimeUnit.SECONDS) + fun setXComRejectsByteArrayBeforeWriting() { + val toClient = ByteChannel(autoFlush = true) + val fromClient = ByteChannel(autoFlush = true) + val comm = CoordinatorComm(toClient, fromClient) + val details = + StartupDetails().also { + it.ti = + TaskInstance().also { ti -> + ti.dagId = "d" + ti.taskId = "t" + ti.runId = "r" + } + } + val client = PublicClient(details, CoordinatorClient(comm)) + // Answer the first request in advance. If setXCom sent the byte array anyway, + // this reply would let setXCom return. Then assertThrows would fail at once + // instead of the test waiting for its 30-second timeout. + runBlocking { toClient.writeFrame(emptyResponseFrame(0)) } + + val error = + Assertions.assertThrows(IllegalArgumentException::class.java) { + client.setXCom(value = mapOf("blob" to byteArrayOf(1))) + } + + Assertions.assertTrue("Base64" in error.message.orEmpty(), "the error should suggest an encoding: $error") + runBlocking { fromClient.flushAndClose() } + Assertions.assertTrue(runBlocking { fromClient.exhausted() }, "nothing should be written to the supervisor") + comm.close() + } + @Test @DisplayName("Should fail a pending call when the coordinator socket closes") @Timeout(value = 30, unit = TimeUnit.SECONDS)