From fc0ccc94a3ee010db57a005cda17e728045c9eae Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ivan=20=E2=80=9CCLOVIS=E2=80=9D=20Canet?=
 <ivan.canet@gmail.com>
Date: Sun, 2 Aug 2026 11:56:13 +0200
Subject: [PATCH] refactor(driver-multiplatform-wire): Use a simple Deferred
 instead of a Channel in the common case

---
 .../src/commonMain/kotlin/MongoWireClient.kt  | 48 +++++++++----------
 1 file changed, 24 insertions(+), 24 deletions(-)

diff --git a/driver-multiplatform-wire/src/commonMain/kotlin/MongoWireClient.kt b/driver-multiplatform-wire/src/commonMain/kotlin/MongoWireClient.kt
index 6715b11d..b380c88c 100644
--- a/driver-multiplatform-wire/src/commonMain/kotlin/MongoWireClient.kt
+++ b/driver-multiplatform-wire/src/commonMain/kotlin/MongoWireClient.kt
@@ -31,7 +31,6 @@ import opensavvy.ktmongo.bson.multiplatform.BsonDocument
 import opensavvy.ktmongo.bson.multiplatform.BsonFactory
 import opensavvy.ktmongo.dsl.LowLevelApi
 import kotlin.coroutines.CoroutineContext
-import kotlin.coroutines.cancellation.CancellationException
 
 @LowLevelApi
 interface MongoWireClient : AutoCloseable {
@@ -71,16 +70,19 @@ private class SocketWireClient(
 	coroutineScope: CoroutineScope,
 ) : MongoWireClient {
 
+	private sealed class ResponseHandler {
+		data class Single(val result: CompletableDeferred<Message>) : ResponseHandler()
+		data class Multiple(val result: SendChannel<Message>) : ResponseHandler()
+	}
+
 	private class Request(
 		val data: Buffer,
-		val output: Channel<Message>,
-		val expectsMultipleResponses: Boolean,
+		val output: ResponseHandler,
 	)
 
 	private class SentMessage(
 		val requestId: Int,
-		val output: Channel<Message>,
-		val expectsMultipleResponses: Boolean,
+		val output: ResponseHandler,
 	)
 
 	private class Response(
@@ -91,7 +93,7 @@ private class SocketWireClient(
 
 	private class ResponseWithHandler(
 		val response: Response,
-		val output: SendChannel<Message>,
+		val output: ResponseHandler,
 	)
 
 	/**
@@ -170,7 +172,7 @@ private class SocketWireClient(
 			writeSocket.flush()
 
 			log("$requestId was sent")
-			sentChannel.send(SentMessage(requestId, request.output, expectsMultipleResponses = request.expectsMultipleResponses))
+			sentChannel.send(SentMessage(requestId, request.output))
 		}
 	}
 
@@ -207,8 +209,7 @@ private class SocketWireClient(
 		receivedChannel: ReceiveChannel<Response>,
 		triagedChannel: SendChannel<ResponseWithHandler>,
 	) {
-		val waiting = HashMap<Int, SendChannel<Message>>()
-		val requestsExpectingMultipleResponses = HashSet<Int>()
+		val waiting = HashMap<Int, ResponseHandler>()
 
 		while (currentCoroutineContext().isActive && socket.isActive) {
 			select {
@@ -219,16 +220,13 @@ private class SocketWireClient(
 				sentChannel.onReceive { message ->
 					log("${message.requestId} expects an answer")
 					waiting[message.requestId] = message.output
-					if (message.expectsMultipleResponses)
-						requestsExpectingMultipleResponses.add(message.requestId)
 				}
 
 				receivedChannel.onReceive { response ->
 					val handler = waiting[response.responseTo]
 						?: error("Received the message ${response.requestId} in response to ${response.responseTo}, but no known message with ID ${response.responseTo} has been sent by this client.\nCurrently in-flight requests: ${waiting.keys.sorted()}")
 					triagedChannel.send(ResponseWithHandler(response, handler))
-					if (response.responseTo !in requestsExpectingMultipleResponses) {
-						requestsExpectingMultipleResponses.remove(response.responseTo)
+					if (handler is ResponseHandler.Single) {
 						waiting.remove(response.responseTo)
 					}
 				}
@@ -291,13 +289,16 @@ private class SocketWireClient(
 			val body = sections.singleOrNull { it is MessageSection.Body } as? MessageSection.Body
 				?: error("An OP_MSG message must have a single body section, found: $sections")
 
-			received.output.send(
-				Message.OpMsg(
-					body,
-					sections.asSequence()
-						.filterIsInstance<MessageSection.DocumentSequence>(),
-				)
+			val response = Message.OpMsg(
+				body,
+				sections.asSequence()
+					.filterIsInstance<MessageSection.DocumentSequence>(),
 			)
+
+			when (received.output) {
+				is ResponseHandler.Single -> received.output.result.complete(response)
+				is ResponseHandler.Multiple -> received.output.result.send(response)
+			}
 		}
 	}
 
@@ -398,17 +399,16 @@ private class SocketWireClient(
 		val output = Channel<Message>()
 		log("Preparing to write $message…")
 		val buffer = writeMessage(message)
-		requestChannel.send(Request(buffer, output, expectsMultipleResponses = true))
+		requestChannel.send(Request(buffer, ResponseHandler.Multiple(output)))
 		return output
 	}
 
 	override suspend fun sendSingle(message: Message): Message {
-		val output = Channel<Message>()
+		val output = CompletableDeferred<Message>()
 		log("Preparing to write $message…")
 		val buffer = writeMessage(message)
-		requestChannel.send(Request(buffer, output, expectsMultipleResponses = false))
-		val message = output.receive()
-		output.close(CancellationException("We expected a single response, and we received it, so this channel was closed."))
+		requestChannel.send(Request(buffer, ResponseHandler.Single(output)))
+		val message = output.await()
 		return message
 	}
 
-- 
2.51.2