This is an automated email from the ASF dual-hosted git repository.
jason810496 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 abd837a382c Java SDK: Change coordinator comm to use multiplexing
(#69254)
abd837a382c is described below
commit abd837a382cc71cd0a90f0cb2e9dffedcce7343a
Author: PoAn Yang <[email protected]>
AuthorDate: Thu Jul 30 14:54:30 2026 +0900
Java SDK: Change coordinator comm to use multiplexing (#69254)
Serializing each client's whole send → receive → deserialize round trip
behind one lock (#69080) allows only a single request in flight at a
time, so a task issuing concurrent calls over the shared comm socket
pays a full, serialized round trip per call.
Every frame already carries a request id, so a single background
dispatcher can read the socket and route each response by id, with the
lock guarding only the write.
Signed-off-by: PoAn Yang <[email protected]>
---
.../main/kotlin/org/apache/airflow/sdk/Server.kt | 29 ++-
.../org/apache/airflow/sdk/execution/Comm.kt | 170 +++++++++++------
.../org/apache/airflow/sdk/execution/Frame.kt | 18 +-
.../org/apache/airflow/sdk/execution/Logger.kt | 8 +-
.../kotlin/org/apache/airflow/sdk/ServerTest.kt | 142 ++++++++++++++
.../org/apache/airflow/sdk/execution/CommTest.kt | 210 ++++++++++++++++++---
6 files changed, 492 insertions(+), 85 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 6f82ca3c4c3..654baba944b 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
@@ -33,6 +33,9 @@ import kotlinx.coroutines.runBlocking
import org.apache.airflow.sdk.execution.CoordinatorComm
import org.apache.airflow.sdk.execution.LogSender
import org.apache.airflow.sdk.execution.Logger
+import org.apache.airflow.sdk.execution.comm.ErrorResponse
+import org.apache.airflow.sdk.execution.comm.StartupDetails
+import org.apache.airflow.sdk.execution.runTask
import kotlin.text.substringAfterLast
import kotlin.text.substringBeforeLast
@@ -147,10 +150,11 @@ class Server(
aSocket(SelectorManager(Dispatchers.IO)).tcp().connect(comm).use {
socket ->
logger.debug("Connected comm", mapOf("addr" to comm))
CoordinatorComm(
- bundle,
socket.openReadChannel(),
socket.openWriteChannel(autoFlush = true),
- ).startProcessing()
+ ).use { coordinator ->
+ dispatchTask(bundle, coordinator)
+ }
}
} finally {
deferral.complete(Unit)
@@ -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 -> throw ApiError("Unexpected initial frame (id=${frame.id})")
+ }
+ }
+
+ private suspend fun runTaskAndReport(
+ bundle: Bundle,
+ startup: StartupDetails,
+ coordinator: CoordinatorComm,
+ ) {
+ val result = runTask(bundle, startup, coordinator)
+ coordinator.communicate<Unit>(result)
+ }
}
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 3312d99919c..34815d7b7f1 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
@@ -23,12 +23,18 @@ 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.CompletableDeferred
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.SupervisorJob
+import kotlinx.coroutines.cancel
+import kotlinx.coroutines.launch
import kotlinx.coroutines.sync.Mutex
import kotlinx.coroutines.sync.withLock
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 java.util.concurrent.ConcurrentHashMap
import kotlin.concurrent.atomics.AtomicInt
import kotlin.concurrent.atomics.ExperimentalAtomicApi
@@ -44,10 +50,9 @@ data class OutgoingFrame(
@OptIn(ExperimentalAtomicApi::class)
class CoordinatorComm(
- private val bundle: Bundle,
private val reader: ByteReadChannel,
private val writer: ByteWriteChannel,
-) {
+) : AutoCloseable {
internal companion object {
private val logger = Logger(CoordinatorComm::class)
@@ -57,38 +62,21 @@ class CoordinatorComm(
}
private val nextId = AtomicInt(0)
- private var shutDownRequested = false
- private val commMutex = Mutex()
+ private val writeMutex = Mutex()
+ private val stateMutex = Mutex()
+ private val pending = ConcurrentHashMap<Int,
CompletableDeferred<IncomingFrame>>()
- suspend fun startProcessing() {
- while (!shutDownRequested) {
- processOnce(::handleIncoming)
- }
- logger.debug("Goodbye")
- }
+ @Volatile
+ private var readError: ApiError? = null
+ private var dispatcherStarted = false
- 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
- }
+ private val dispatcherScope = CoroutineScope(Dispatchers.IO +
SupervisorJob())
- 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
+ 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,
@@ -96,38 +84,41 @@ class CoordinatorComm(
) {
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) + data)
}
}
@Throws(ApiError::class)
suspend fun communicateImpl(body: Any): Any {
val requestId = nextId.fetchAndAdd(1)
- return commMutex.withLock {
- var frame: IncomingFrame? = null
-
- suspend fun handle(f: IncomingFrame) {
- frame = f
+ val waiter = CompletableDeferred<IncomingFrame>()
+
+ stateMutex.withLock {
+ readError?.let { throw it }
+ pending[requestId] = waiter
+ if (!dispatcherStarted) {
+ dispatcherStarted = true
+ dispatcherScope.launch { readLoop() }
}
+ }
+
+ // A close() concurrent with the block above either drains this waiter (it
+ // was registered first) or is visible here, so no caller is left awaiting
a
+ // comm that is already gone.
+ readError?.let {
+ pending.remove(requestId)
+ throw it
+ }
+
+ try {
sendMessage(requestId, body)
- processOnce(::handle)
- val received = frame ?: throw ApiError("No response received")
- if (received.id != requestId) {
- throw ApiError("response id ${received.id} does not match request id
$requestId")
- }
- received.body ?: Unit
+ } catch (e: Throwable) {
+ pending.remove(requestId)
+ throw e
}
+
+ return waiter.await().body ?: Unit
}
@Throws(ApiError::class)
@@ -138,4 +129,73 @@ class CoordinatorComm(
else -> throw ApiError("Unexpected response type
${response::class.java}")
}
}
+
+ /**
+ * Stop the dispatcher and fail anything still awaiting a response.
+ */
+ override fun close() {
+ dispatcherScope.cancel()
+ failAllWaiters(readError ?: ApiError("Coordinator comm closed"))
+ }
+
+ private suspend fun readPayload(): ByteArray {
+ 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")
+ }
+ return payload
+ }
+
+ private suspend fun readFrame(): IncomingFrame {
+ val frame = decode(readPayload())
+ logger.debug("Received", mapOf("id" to frame.id))
+ return frame
+ }
+
+ private suspend fun readLoop() {
+ while (true) {
+ val raw =
+ try {
+ Frame.decodeRaw(readPayload())
+ } 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
+ }
+ logger.debug("Received", mapOf("id" to raw.id))
+
+ val waiter = pending.remove(raw.id)
+ if (waiter == null) {
+ logger.warning("Discarding response with no matching request",
mapOf("id" to raw.id))
+ continue
+ }
+
+ try {
+ waiter.complete(IncomingFrame(raw.id, Frame.decodeBody(raw)))
+ } catch (e: CancellationException) {
+ throw e
+ } catch (e: Throwable) {
+ waiter.completeExceptionally(
+ ApiError("Cannot decode response for request ${raw.id}:
${e.message}")
+ .apply { initCause(e) },
+ )
+ }
+ }
+ }
+
+ private fun failAllWaiters(cause: Throwable) {
+ val error =
+ cause as? ApiError
+ ?: ApiError("Coordinator comm closed: ${cause.message}").apply {
initCause(cause) }
+ readError = error
+ pending.keys.forEach { id ->
pending.remove(id)?.completeExceptionally(error) }
+ }
}
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 a3815d61401..4d488fae0fc 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
@@ -28,6 +28,12 @@ import org.apache.airflow.sdk.execution.comm.Discriminator
import org.msgpack.core.MessagePack
import java.io.ByteArrayOutputStream
+data class RawFrame(
+ val id: Int,
+ val rawBody: Any?,
+ val rawError: Any?,
+)
+
object Frame {
private val mapper =
ObjectMapper().apply {
@@ -43,7 +49,7 @@ object Frame {
body: Any,
): ByteArray = encodeFrame(id, body)
- fun decode(bytes: ByteArray): IncomingFrame {
+ fun decodeRaw(bytes: ByteArray): RawFrame {
val unpacker = MessagePack.newDefaultUnpacker(bytes)
val headerSize = unpacker.unpackArrayHeader()
check(headerSize >= 1) { "Unexpected Task SDK frame arity $headerSize" }
@@ -53,8 +59,14 @@ object Frame {
val rawError = if (headerSize >= 3) unpacker.unpackAny() else null
unpacker.close()
- val body = decodeMessage(rawError) ?: decodeMessage(rawBody)
- return IncomingFrame(id, body)
+ 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)
+ return IncomingFrame(raw.id, decodeBody(raw))
}
fun lengthPrefix(length: Int) =
diff --git
a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Logger.kt
b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Logger.kt
index f39490eb0a9..228d59e7240 100644
--- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Logger.kt
+++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Logger.kt
@@ -198,15 +198,15 @@ internal class Logger(
arguments: Map<String, Any> = emptyMap(),
) = log(Level.DEBUG, message, arguments)
- fun error(
+ fun warning(
message: String,
arguments: Map<String, Any> = emptyMap(),
- ) = log(Level.ERROR, message, arguments)
+ ) = log(Level.WARNING, message, arguments)
- fun warning(
+ fun error(
message: String,
arguments: Map<String, Any> = emptyMap(),
- ) = log(Level.WARNING, message, arguments)
+ ) = log(Level.ERROR, message, arguments)
private fun log(
level: Level,
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
new file mode 100644
index 00000000000..6cae7a4c27b
--- /dev/null
+++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ServerTest.kt
@@ -0,0 +1,142 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.airflow.sdk
+
+import io.ktor.network.sockets.InetSocketAddress
+import io.ktor.utils.io.ByteChannel
+import io.ktor.utils.io.readByteArray
+import io.ktor.utils.io.writeByteArray
+import kotlinx.coroutines.runBlocking
+import org.apache.airflow.sdk.execution.CoordinatorComm
+import org.apache.airflow.sdk.execution.Frame
+import org.apache.airflow.sdk.execution.IncomingFrame
+import org.apache.airflow.sdk.execution.comm.TaskState
+import org.junit.jupiter.api.Assertions
+import org.junit.jupiter.api.DisplayName
+import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.Timeout
+import org.msgpack.core.MessagePack
+import java.io.ByteArrayOutputStream
+import java.util.concurrent.ArrayBlockingQueue
+import java.util.concurrent.TimeUnit
+
+class ServerTest {
+ private fun hexToBytes(hex: String): ByteArray =
+ hex
+ .split(' ', '\r', '\n')
+ .filter { it.isNotEmpty() }
+ .map { it.toUByte(16).toByte() }
+ .toByteArray()
+
+ private fun ackFrame(id: Int): ByteArray {
+ val out = ByteArrayOutputStream()
+ MessagePack.newDefaultPacker(out).use { packer ->
+ packer.packArrayHeader(2)
+ packer.packInt(id)
+ packer.packNil()
+ }
+ return out.toByteArray()
+ }
+
+ private suspend fun ByteChannel.writeFrame(payload: ByteArray) {
+ writeByteArray(Frame.lengthPrefix(payload.size))
+ writeByteArray(payload)
+ }
+
+ @Test
+ @DisplayName("Should run the task and report the result as a normal awaited
request")
+ @Timeout(value = 30, unit = TimeUnit.SECONDS)
+ fun runsTaskAndReportsResultAsAwaitedRequest() {
+ val toServer = ByteChannel(autoFlush = true)
+ val fromServer = ByteChannel(autoFlush = true)
+ val comm = CoordinatorComm(toServer, fromServer)
+ val server = Server(InetSocketAddress("localhost", 0),
InetSocketAddress("localhost", 0))
+
+ val reported = ArrayBlockingQueue<IncomingFrame>(1)
+ val supervisor =
+ Thread {
+ runBlocking {
+ // 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 result = CoordinatorComm.decode(payload)
+ reported.put(result)
+ toServer.writeFrame(ackFrame(result.id))
+ }
+ }
+ supervisor.start()
+
+ runBlocking { server.dispatchTask(Bundle(emptyList()), comm) }
+ supervisor.join()
+
+ val result = reported.take()
+ Assertions.assertEquals(0, result.id)
+ Assertions.assertInstanceOf(TaskState::class.java, result.body)
+ Assertions.assertEquals(TaskState.State.REMOVED, (result.body as
TaskState).state)
+ comm.close()
+ }
+
+ @Test
+ @DisplayName("Should fail when the initial frame is not startup details")
+ @Timeout(value = 30, unit = TimeUnit.SECONDS)
+ fun failsWhenInitialFrameIsNotStartupDetails() {
+ val toServer = ByteChannel(autoFlush = true)
+ val fromServer = ByteChannel(autoFlush = true)
+ val comm = CoordinatorComm(toServer, fromServer)
+ val server = Server(InetSocketAddress("localhost", 0),
InetSocketAddress("localhost", 0))
+
+ Assertions.assertThrows(ApiError::class.java) {
+ runBlocking {
+ toServer.writeFrame(ackFrame(7))
+ server.dispatchTask(Bundle(emptyList()), comm)
+ }
+ }
+ comm.close()
+ }
+
+ private companion object {
+ // [2, msg, null] with msg coming from
+ //
https://github.com/astronomer/airflow/blob/f39c8da8/task-sdk/tests/task_sdk/execution_time/test_comms.py#L73-L108
+ val STARTUP_HEX =
+ """
+ 92 02 88 a4 74 79 70 65 ae 53 74 61 72 74 75 70 44 65 74 61 69 6c 73 a2
74 69 86 a2 69 64 d9 24
+ 34 64 38 32 38 61 36 32 2d 61 34 31 37 2d 34 39 33 36 2d 61 37 61 36 2d
32 62 33 66 61 62 61 63
+ 65 63 61 62 a7 74 61 73 6b 5f 69 64 a1 61 aa 74 72 79 5f 6e 75 6d 62 65
72 01 a6 72 75 6e 5f 69
+ 64 a1 62 a6 64 61 67 5f 69 64 a1 63 ae 64 61 67 5f 76 65 72 73 69 6f 6e
5f 69 64 d9 24 34 64 38
+ 32 38 61 36 32 2d 61 34 31 37 2d 34 39 33 36 2d 61 37 61 36 2d 32 62 33
66 61 62 61 63 65 63 61
+ 62 aa 74 69 5f 63 6f 6e 74 65 78 74 85 a7 64 61 67 5f 72 75 6e 8c a6 64
61 67 5f 69 64 a1 63 a6
+ 72 75 6e 5f 69 64 a1 62 ac 6c 6f 67 69 63 61 6c 5f 64 61 74 65 b4 32 30
32 34 2d 31 32 2d 30 31
+ 54 30 31 3a 30 30 3a 30 30 5a b3 64 61 74 61 5f 69 6e 74 65 72 76 61 6c
5f 73 74 61 72 74 b4 32
+ 30 32 34 2d 31 32 2d 30 31 54 30 30 3a 30 30 3a 30 30 5a b1 64 61 74 61
5f 69 6e 74 65 72 76 61
+ 6c 5f 65 6e 64 b4 32 30 32 34 2d 31 32 2d 30 31 54 30 31 3a 30 30 3a 30
30 5a aa 73 74 61 72 74
+ 5f 64 61 74 65 b4 32 30 32 34 2d 31 32 2d 30 31 54 30 31 3a 30 30 3a 30
30 5a a9 72 75 6e 5f 61
+ 66 74 65 72 b4 32 30 32 34 2d 31 32 2d 30 31 54 30 31 3a 30 30 3a 30 30
5a a8 65 6e 64 5f 64 61
+ 74 65 c0 a8 72 75 6e 5f 74 79 70 65 a6 6d 61 6e 75 61 6c a5 73 74 61 74
65 a7 73 75 63 63 65 73
+ 73 a4 63 6f 6e 66 c0 b5 63 6f 6e 73 75 6d 65 64 5f 61 73 73 65 74 5f 65
76 65 6e 74 73 90 a9 6d
+ 61 78 5f 74 72 69 65 73 00 ac 73 68 6f 75 6c 64 5f 72 65 74 72 79 c2 a9
76 61 72 69 61 62 6c 65
+ 73 c0 ab 63 6f 6e 6e 65 63 74 69 6f 6e 73 c0 a4 66 69 6c 65 a9 2f 64 65
76 2f 6e 75 6c 6c aa 73
+ 74 61 72 74 5f 64 61 74 65 b4 32 30 32 34 2d 31 32 2d 30 31 54 30 31 3a
30 30 3a 30 30 5a ac 64
+ 61 67 5f 72 65 6c 5f 70 61 74 68 a9 2f 64 65 76 2f 6e 75 6c 6c ab 62 75
6e 64 6c 65 5f 69 6e 66
+ 6f 82 a4 6e 61 6d 65 a8 61 6e 79 2d 6e 61 6d 65 a7 76 65 72 73 69 6f 6e
ab 61 6e 79 2d 76 65 72
+ 73 69 6f 6e b2 73 65 6e 74 72 79 5f 69 6e 74 65 67 72 61 74 69 6f 6e a0
c0
+ """.trimIndent()
+ }
+}
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 ead1d1dbe51..978e377f52d 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
@@ -22,9 +22,10 @@ package org.apache.airflow.sdk.execution
import io.ktor.utils.io.ByteChannel
import io.ktor.utils.io.readByteArray
import io.ktor.utils.io.writeByteArray
+import kotlinx.coroutines.async
import kotlinx.coroutines.runBlocking
+import kotlinx.coroutines.supervisorScope
import org.apache.airflow.sdk.ApiError
-import org.apache.airflow.sdk.Bundle
import org.apache.airflow.sdk.execution.comm.GetVariable
import org.apache.airflow.sdk.execution.comm.StartupDetails
import org.apache.airflow.sdk.execution.comm.TaskInstance
@@ -99,7 +100,10 @@ class CommsTest {
Assertions.assertEquals(expected, actual)
}
- private fun responseFrame(id: Int): ByteArray {
+ private fun responseFrame(
+ id: Int,
+ key: String = "return_value",
+ ): ByteArray {
val out = ByteArrayOutputStream()
MessagePack.newDefaultPacker(out).use { packer ->
packer.packArrayHeader(3)
@@ -108,7 +112,7 @@ class CommsTest {
packer.packString("type")
packer.packString("XComResult")
packer.packString("key")
- packer.packString("return_value")
+ packer.packString(key)
packer.packString("value")
packer.packInt(1)
packer.packNil()
@@ -116,27 +120,93 @@ class CommsTest {
return out.toByteArray()
}
+ private fun unknownTypeFrame(id: Int): ByteArray {
+ val out = ByteArrayOutputStream()
+ MessagePack.newDefaultPacker(out).use { packer ->
+ packer.packArrayHeader(3)
+ packer.packInt(id)
+ packer.packMapHeader(1)
+ packer.packString("type")
+ packer.packString("SomeFutureMessage")
+ packer.packNil()
+ }
+ return out.toByteArray()
+ }
+
+ private suspend fun ByteChannel.writeFrame(payload: ByteArray) {
+ writeByteArray(Frame.lengthPrefix(payload.size))
+ writeByteArray(payload)
+ }
+
+ private suspend fun ByteChannel.readOneRequest() {
+ val prefix = readByteArray(4)
+ readByteArray(Frame.parseLengthPrefix(prefix))
+ }
+
@Test
- @DisplayName("Should reject a response whose id does not match the request")
- fun rejectsResponseWhoseIdDoesNotMatchRequest() {
+ @DisplayName("Should discard a frame with no matching waiter and still
deliver the correct response")
+ @Timeout(value = 30, unit = TimeUnit.SECONDS)
+ fun discardsUnmatchedFrameAndDeliversCorrectResponse() {
val toClient = ByteChannel(autoFlush = true)
val fromClient = ByteChannel(autoFlush = true)
- val comm = CoordinatorComm(Bundle(emptyList()), toClient, fromClient)
+ val comm = CoordinatorComm(toClient, fromClient)
+
+ val result =
+ runBlocking {
+ // A frame whose id (99) matches no in-flight request.
+ // The dispatcher must drop the stray frame and still hand id 0 to its
waiter.
+ toClient.writeFrame(responseFrame(99, key = "stray"))
+ toClient.writeFrame(responseFrame(0, key = "return_value"))
+ comm.communicate<XComResult>(GetVariable().also { it.key = "k" })
+ }
- val error =
- Assertions.assertThrows(ApiError::class.java) {
+ Assertions.assertEquals("return_value", result.key)
+ comm.close()
+ }
+
+ @Test
+ @DisplayName("Should keep many requests in flight and match out-of-order
responses")
+ @Timeout(value = 30, unit = TimeUnit.SECONDS)
+ fun allowsMultipleInFlightRequestsAnsweredOutOfOrder() {
+ val toClient = ByteChannel(autoFlush = true)
+ val fromClient = ByteChannel(autoFlush = true)
+ val comm = CoordinatorComm(toClient, fromClient)
+ val n = 10
+
+ // Collect every request before answering any, then reply in reverse order.
+ val server =
+ Thread {
runBlocking {
- // The first request is sent with id 0. The 99 doesn't match 0.
- val payload = responseFrame(99)
- toClient.writeByteArray(Frame.lengthPrefix(payload.size))
- toClient.writeByteArray(payload)
- comm.communicate<XComResult>(GetVariable().also { it.key = "k" })
+ val ids =
+ (0 until n).map {
+ val prefix = fromClient.readByteArray(4)
+ val payload =
fromClient.readByteArray(Frame.parseLengthPrefix(prefix))
+ CoordinatorComm.decode(payload).id
+ }
+ ids.reversed().forEach { toClient.writeFrame(responseFrame(it)) }
}
}
- Assertions.assertTrue(
- error.message!!.contains("does not match"),
- "expected an id-mismatch error, got: ${error.message}",
- )
+ server.start()
+
+ val errors = ConcurrentLinkedQueue<Throwable>()
+ val results = ConcurrentLinkedQueue<XComResult>()
+ val workers =
+ (1..n).map {
+ Thread {
+ try {
+ results.add(runBlocking {
comm.communicate<XComResult>(GetVariable().also { it.key = "k" }) })
+ } catch (e: Throwable) {
+ errors.add(e)
+ }
+ }
+ }
+ workers.forEach { it.start() }
+ workers.forEach { it.join() }
+ server.join()
+
+ Assertions.assertTrue(errors.isEmpty(), "concurrent in-flight calls
failed: $errors")
+ Assertions.assertEquals(n, results.size)
+ comm.close()
}
@Test
@@ -145,7 +215,7 @@ class CommsTest {
fun publicClientSurvivesConcurrentThreadCalls() {
val toClient = ByteChannel(autoFlush = true)
val fromClient = ByteChannel(autoFlush = true)
- val comm = CoordinatorComm(Bundle(emptyList()), toClient, fromClient)
+ val comm = CoordinatorComm(toClient, fromClient)
val details =
StartupDetails().also {
it.ti =
@@ -163,9 +233,7 @@ class CommsTest {
repeat(n) {
val prefix = fromClient.readByteArray(4)
val payload =
fromClient.readByteArray(Frame.parseLengthPrefix(prefix))
- val response = responseFrame(CoordinatorComm.decode(payload).id)
- toClient.writeByteArray(Frame.lengthPrefix(response.size))
- toClient.writeByteArray(response)
+
toClient.writeFrame(responseFrame(CoordinatorComm.decode(payload).id))
}
}
}
@@ -189,5 +257,105 @@ class CommsTest {
Assertions.assertTrue(errors.isEmpty(), "concurrent public-client calls
failed: $errors")
Assertions.assertEquals(n, results.size)
+ comm.close()
+ }
+
+ @Test
+ @DisplayName("Should fail a pending call when the coordinator socket closes")
+ @Timeout(value = 30, unit = TimeUnit.SECONDS)
+ fun failsPendingCallWhenSocketCloses() {
+ val toClient = ByteChannel(autoFlush = true)
+ val fromClient = ByteChannel(autoFlush = true)
+ val comm = CoordinatorComm(toClient, fromClient)
+
+ Assertions.assertThrows(ApiError::class.java) {
+ runBlocking {
+ val call = async { comm.communicate<XComResult>(GetVariable().also {
it.key = "k" }) }
+ // No response is ever written; closing the read side must surface an
+ // error to the waiting caller instead of hanging forever.
+ toClient.flushAndClose()
+ call.await()
+ }
+ }
+ comm.close()
+ }
+
+ @Test
+ @DisplayName("Should fail a call that is still in flight when the comm is
closed")
+ @Timeout(value = 30, unit = TimeUnit.SECONDS)
+ fun closeFailsInFlightCall() {
+ val toClient = ByteChannel(autoFlush = true)
+ val fromClient = ByteChannel(autoFlush = true)
+ val comm = CoordinatorComm(toClient, fromClient)
+
+ Assertions.assertThrows(ApiError::class.java) {
+ runBlocking {
+ val call = async { comm.communicate<XComResult>(GetVariable().also {
it.key = "k" }) }
+ fromClient.readOneRequest()
+ // The waiter awaits in the caller's own scope, so closing the comm
should make it stop waiting.
+ comm.close()
+ call.await()
+ }
+ }
+ }
+
+ @Test
+ @DisplayName("Should fail only the request whose response cannot be decoded")
+ @Timeout(value = 30, unit = TimeUnit.SECONDS)
+ fun undecodableResponseFailsOnlyItsOwnRequest() {
+ val toClient = ByteChannel(autoFlush = true)
+ val fromClient = ByteChannel(autoFlush = true)
+ val comm = CoordinatorComm(toClient, fromClient)
+
+ runBlocking {
+ supervisorScope {
+ // Start the calls one at a time so their ids are known: 0, then 1.
+ val undecodable = async {
comm.communicate<XComResult>(GetVariable().also { it.key = "a" }) }
+ fromClient.readOneRequest()
+ val healthy = async { comm.communicate<XComResult>(GetVariable().also
{ it.key = "b" }) }
+ fromClient.readOneRequest()
+
+ toClient.writeFrame(unknownTypeFrame(0))
+ toClient.writeFrame(responseFrame(1))
+
+ Assertions.assertInstanceOf(
+ ApiError::class.java,
+ runCatching { undecodable.await() }.exceptionOrNull(),
+ "a response naming an unknown type should fail its own request",
+ )
+ // The dispatcher survived the bad frame, so the other call still
lands.
+ Assertions.assertEquals("return_value", healthy.await().key)
+ }
+ }
+ comm.close()
+ }
+
+ @Test
+ @DisplayName("Should keep the original cause when a malformed frame stops
the dispatcher")
+ @Timeout(value = 30, unit = TimeUnit.SECONDS)
+ fun malformedFramePreservesOriginalCause() {
+ val toClient = ByteChannel(autoFlush = true)
+ val fromClient = ByteChannel(autoFlush = true)
+ val comm = CoordinatorComm(toClient, fromClient)
+
+ val failure =
+ runBlocking {
+ supervisorScope {
+ val call = async { comm.communicate<XComResult>(GetVariable().also {
it.key = "k" }) }
+ fromClient.readOneRequest()
+ // Msgpack nil where the frame envelope belongs: without an id there
is no
+ // request to attribute the failure to, so every waiter goes down
with it.
+ toClient.writeFrame(byteArrayOf(0xc0.toByte()))
+ runCatching { call.await() }.exceptionOrNull()
+ }
+ }
+
+ Assertions.assertInstanceOf(ApiError::class.java, failure)
+ val causes = generateSequence(failure) { it.cause }.toList()
+ Assertions.assertTrue(
+ causes.any { it !is ApiError },
+ "the underlying decode failure should be preserved, got: $causes",
+ )
+ comm.close()
}
}
