diff --git a/driver-multiplatform-wire/src/commonMain/kotlin/MongoWireClient.kt b/driver-multiplatform-wire/src/commonMain/kotlin/MongoWireClient.kt index 04544398..1b351a1b 100644 --- a/driver-multiplatform-wire/src/commonMain/kotlin/MongoWireClient.kt +++ b/driver-multiplatform-wire/src/commonMain/kotlin/MongoWireClient.kt @@ -31,6 +31,7 @@ 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 { @@ -39,6 +40,13 @@ interface MongoWireClient : AutoCloseable { message: Message, ): ReceiveChannel + /** + * Sends a [message] that expects a single response. + */ + suspend fun sendSingle( + message: Message, + ): Message + companion object } @@ -65,11 +73,13 @@ private class SocketWireClient( private class Request( val data: Buffer, val output: Channel, + val expectsMultipleResponses: Boolean, ) private class SentMessage( val requestId: Int, val output: Channel, + val expectsMultipleResponses: Boolean, ) private class Response( @@ -159,7 +169,7 @@ private class SocketWireClient( writeSocket.flush() log("$requestId was sent") - sentChannel.send(SentMessage(requestId, request.output)) + sentChannel.send(SentMessage(requestId, request.output, expectsMultipleResponses = request.expectsMultipleResponses)) } } @@ -197,6 +207,7 @@ private class SocketWireClient( triagedChannel: SendChannel, ) { val waiting = HashMap>() + val requestsExpectingMultipleResponses = HashSet() while (currentCoroutineContext().isActive && socket.isActive) { select { @@ -207,12 +218,18 @@ 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) + waiting.remove(response.responseTo) + } } } } @@ -380,10 +397,20 @@ private class SocketWireClient( val output = Channel() log("Preparing to write $message…") val buffer = writeMessage(message) - requestChannel.send(Request(buffer, output)) + requestChannel.send(Request(buffer, output, expectsMultipleResponses = true)) return output } + override suspend fun sendSingle(message: Message): Message { + val output = Channel() + 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.")) + return message + } + override fun close() { socket.close() } diff --git a/driver-multiplatform/src/commonMain/kotlin/MongoCollectionImpl.kt b/driver-multiplatform/src/commonMain/kotlin/MongoCollectionImpl.kt index 455587b4..2ababffe 100644 --- a/driver-multiplatform/src/commonMain/kotlin/MongoCollectionImpl.kt +++ b/driver-multiplatform/src/commonMain/kotlin/MongoCollectionImpl.kt @@ -16,7 +16,6 @@ package opensavvy.ktmongo.multiplatform -import kotlinx.coroutines.CancellationException import opensavvy.ktmongo.dsl.LowLevelApi import opensavvy.ktmongo.dsl.command.InsertOne import opensavvy.ktmongo.dsl.command.InsertOneOptions @@ -48,7 +47,7 @@ internal class MongoCollectionImpl( } } - val responses = database.client.wire.send( + val message = database.client.wire.sendSingle( Message.OpMsg( body = MessageSection.Body( command, @@ -56,9 +55,6 @@ internal class MongoCollectionImpl( ) ) - val message = responses.receive() - responses.cancel(CancellationException("insertOne expects a single response")) - check(message is Message.OpMsg) check(message.body.document["ok"]?.decodeDouble() == 1.0) }