Add multi-response test, fix responder delivery race

Author: mattmassicotte

Sources/JSONRPC/JSONRPCSession.swift

@@ -250,25 +250,30 @@ extension JSONRPCSession {
 }
 
 extension JSONRPCSession {
-	private func sendDataRequest<Request>(_ params: Request, method: String, responseHandler: @escaping MessageResponder)
-	where Request: Encodable {
+	private func sendDataRequest<Request>(
+		_ params: Request, method: String,
+		responseHandler: @escaping MessageResponder
+	) where Request: Encodable {
 		let issuedId = generateID()
 
 		let request = JSONRPCRequest(id: issuedId, method: method, params: params)
 
+		// make sure to store the responser *first*, before sending the message. This prevents a race where the response comes in so fast we aren't yet waiting for it
+		let key = issuedId.description
+
+		precondition(responders[key] == nil)
+
+		self.responders[key] = responseHandler
+
 		Task {
 			do {
 				try await encodeAndWrite(request)
 			} catch {
 				responseHandler(.failure(error))
+
+				self.responders[key] = nil
 				return
 			}
-			
-			let key = issuedId.description
-			
-			precondition(responders[key] == nil)
-			
-			self.responders[key] = responseHandler
 		}
 	}
 }

Tests/JSONRPCTests/JSONRPCSessionTests.swift

@@ -3,11 +3,45 @@ import JSONRPC
 
 #if compiler(>=5.9)
 
+actor PredefinedMessageRelay {
+	private var messages: [Data]
+	private let continuation: DataChannel.DataSequence.Continuation
+	public nonisolated let sequence: DataChannel.DataSequence
+
+	init(messages: [Data]) {
+		self.messages = messages
+
+		(self.sequence, self.continuation) = DataChannel.DataSequence.makeStream()
+	}
+
+	init<T: Encodable>(content: [T]) throws {
+		let messages = try content.map { try JSONEncoder().encode($0) }
+
+		self.init(messages: messages)
+	}
+
+	func write() {
+		continuation.yield(messages.removeFirst())
+	}
+}
+
+extension DataChannel {
+	static func predefinedMessagesChannel<T: Encodable>(_ content: [T]) throws -> DataChannel {
+		let relay = try PredefinedMessageRelay(content: content)
+
+		return DataChannel(
+			writeHandler: { _ in
+				// strong-ref here to keep relay alive
+				await relay.write()
+			},
+			dataSequence: relay.sequence)
+	}
+}
+
 final class JSONRPCSessionTests: XCTestCase {
 	typealias TestResponse = JSONRPCResponse<String?>
-	typealias TestResult = Result<TestResponse, Error>
 
-	func testThing() async throws {
+	func testResultResponse() async throws {
 		let pair = DataChannel.DataSequence.makeStream()
 
 		let channel = DataChannel(writeHandler: { _ in },
@@ -27,6 +61,30 @@ final class JSONRPCSessionTests: XCTestCase {
 
 		XCTAssertEqual(response.result, "goodbye")
 	}
+
+	func testManySendRequestsWithResponses() async throws {
+		let iterations = 1000
+
+		// be sure to start at 1, to match id generation
+		let responses = (1...iterations).map { i in
+			let responseParam = "goodbye-\(i)"
+
+			return TestResponse(id: JSONId(i), result: responseParam)
+		}
+
+		let channel = try DataChannel.predefinedMessagesChannel(responses)
+		let session = JSONRPCSession(channel: channel)
+
+		let params = "hello"
+		
+		for i in 1...iterations {
+			let expectedResponse = "goodbye-\(i)"
+
+			let response: TestResponse = try await session.sendRequest(params, method: "mymethod")
+
+			XCTAssertEqual(try! response.content.get(), expectedResponse)
+		}
+	}
 }
 
 #endif