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 <payang@apache.org>

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)
+  }
 }

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
+    val waiter = CompletableDeferred<IncomingFrame>()
 
-      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) }
+  }
 }

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) =

java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Logger.kt

@@ -198,16 +198,16 @@ internal class Logger(
     arguments: Map<String, Any> = emptyMap(),
   ) = log(Level.DEBUG, message, arguments)
 
-  fun error(
-    message: String,
-    arguments: Map<String, Any> = emptyMap(),
-  ) = log(Level.ERROR, message, arguments)
-
   fun warning(
     message: String,
     arguments: Map<String, Any> = emptyMap(),
   ) = log(Level.WARNING, message, arguments)
 
+  fun error(
+    message: String,
+    arguments: Map<String, Any> = emptyMap(),
+  ) = log(Level.ERROR, message, arguments)
+
   private fun log(
     level: Level,
     event: String,

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()
+  }
+}

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<XComResult>(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<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()
   }
 }