From abd837a382cc71cd0a90f0cb2e9dffedcce7343a Mon Sep 17 00:00:00 2001 From: PoAn Yang Date: Thu, 30 Jul 2026 14:54:30 +0900 Subject: [PATCH] Java SDK: Change coordinator comm to use multiplexing (#69254) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- .../kotlin/org/apache/airflow/sdk/Server.kt | 29 ++- .../org/apache/airflow/sdk/execution/Comm.kt | 170 +++++++++----- .../org/apache/airflow/sdk/execution/Frame.kt | 18 +- .../apache/airflow/sdk/execution/Logger.kt | 10 +- .../org/apache/airflow/sdk/ServerTest.kt | 142 ++++++++++++ .../apache/airflow/sdk/execution/CommTest.kt | 210 ++++++++++++++++-- 6 files changed, 493 insertions(+), 86 deletions(-) create mode 100644 java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ServerTest.kt 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(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>() - 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(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 + val waiter = CompletableDeferred() - suspend fun handle(f: IncomingFrame) { - frame = f + stateMutex.withLock { + readError?.let { throw it } + pending[requestId] = waiter + if (!dispatcherStarted) { + dispatcherStarted = true + dispatcherScope.launch { readLoop() } } - 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 } + + // 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) + } 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,16 +198,16 @@ internal class Logger( arguments: Map = emptyMap(), ) = log(Level.DEBUG, message, arguments) - fun error( - message: String, - arguments: Map = emptyMap(), - ) = log(Level.ERROR, message, arguments) - fun warning( message: String, arguments: Map = emptyMap(), ) = log(Level.WARNING, message, arguments) + fun error( + message: String, + arguments: Map = emptyMap(), + ) = log(Level.ERROR, message, arguments) + private fun log( level: Level, event: String, 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(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 error = - Assertions.assertThrows(ApiError::class.java) { + 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(GetVariable().also { it.key = "k" }) + } + + 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(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() + val results = ConcurrentLinkedQueue() + val workers = + (1..n).map { + Thread { + try { + results.add(runBlocking { comm.communicate(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(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(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(GetVariable().also { it.key = "a" }) } + fromClient.readOneRequest() + val healthy = async { comm.communicate(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(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() } }