Fix a race condition in TestMessageHandler

Not entirely sure what introduced the race condition. Found by thread sanitizer. Probably wouldn’t have been an issue if we had migrated this code to strict concurrency checking.

Author: ahoppen

Package.swift

@@ -155,6 +155,7 @@ let package = Package(
       dependencies: [
         "LanguageServerProtocol",
         "LanguageServerProtocolJSONRPC",
+        "SKSupport",
       ]
     ),
 

Sources/LSPTestSupport/TestJSONRPCConnection.swift

@@ -12,6 +12,7 @@
 
 import LanguageServerProtocol
 import LanguageServerProtocolJSONRPC
+import SKSupport
 import XCTest
 
 import class Foundation.Pipe
@@ -24,7 +25,7 @@ public final class TestJSONRPCConnection {
   public let server: TestServer
   public let serverConnection: JSONRPCConnection
 
-  public init() {
+  public init(allowUnexpectedNotification: Bool = true) {
     clientConnection = JSONRPCConnection(
       name: "client",
       protocol: testMessageRegistry,
@@ -39,7 +40,7 @@ public final class TestJSONRPCConnection {
       outFD: serverToClient.fileHandleForWriting
     )
 
-    client = TestMessageHandler(server: clientConnection)
+    client = TestMessageHandler(server: clientConnection, allowUnexpectedNotification: allowUnexpectedNotification)
     server = TestServer(client: serverConnection)
 
     clientConnection.start(receiveHandler: client) {
@@ -66,8 +67,8 @@ public struct TestLocalConnection {
   public let server: TestServer
   public let serverConnection: LocalConnection = .init()
 
-  public init() {
-    client = TestMessageHandler(server: serverConnection)
+  public init(allowUnexpectedNotification: Bool = true) {
+    client = TestMessageHandler(server: serverConnection, allowUnexpectedNotification: allowUnexpectedNotification)
     server = TestServer(client: clientConnection)
 
     clientConnection.start(handler: client)
@@ -80,17 +81,20 @@ public struct TestLocalConnection {
   }
 }
 
-public final class TestMessageHandler: MessageHandler {
+public actor TestMessageHandler: MessageHandler {
   /// The connection to the language client.
   public let server: Connection
 
-  public init(server: Connection) {
-    self.server = server
-  }
+  private let messageHandlingQueue = AsyncQueue<Serial>()
 
-  var oneShotNotificationHandlers: [((Any) -> Void)] = []
+  private var oneShotNotificationHandlers: [((Any) -> Void)] = []
 
-  public var allowUnexpectedNotification: Bool = true
+  private let allowUnexpectedNotification: Bool
+
+  public init(server: Connection, allowUnexpectedNotification: Bool = true) {
+    self.server = server
+    self.allowUnexpectedNotification = allowUnexpectedNotification
+  }
 
   public func appendOneShotNotificationHandler<N: NotificationType>(_ handler: @escaping (N) -> Void) {
     oneShotNotificationHandlers.append({ anyNote in
@@ -101,7 +105,14 @@ public final class TestMessageHandler: MessageHandler {
     })
   }
 
-  public func handle(_ notification: some NotificationType, from clientID: ObjectIdentifier) {
+  /// The LSP server sent a notification to the client. Handle it.
+  public nonisolated func handle(_ notification: some NotificationType, from clientID: ObjectIdentifier) {
+    messageHandlingQueue.async {
+      await self.handleNotificationImpl(notification)
+    }
+  }
+
+  public func handleNotificationImpl(_ notification: some NotificationType) {
     guard !oneShotNotificationHandlers.isEmpty else {
       if allowUnexpectedNotification { return }
       fatalError("unexpected notification \(notification)")
@@ -110,25 +121,25 @@ public final class TestMessageHandler: MessageHandler {
     handler(notification)
   }
 
-  public func handle<R: RequestType>(
-    _ params: R,
+  /// The LSP server sent a request to the client. Handle it.
+  public nonisolated func handle<Request: RequestType>(
+    _ request: Request,
     id: RequestID,
     from clientID: ObjectIdentifier,
-    reply: @escaping (LSPResult<R.Response>) -> Void
+    reply: @escaping (LSPResult<Request.Response>) -> Void
   ) {
-    reply(.failure(.methodNotFound(R.method)))
+    reply(.failure(.methodNotFound(Request.method)))
   }
 }
 
 extension TestMessageHandler: Connection {
-
-  /// Send a notification to the language server.
-  public func send(_ notification: some NotificationType) {
+  /// Send a notification to the LSP server.
+  public nonisolated func send(_ notification: some NotificationType) {
     server.send(notification)
   }
 
-  /// Send a request to the language server and (asynchronously) receive a reply.
-  public func send<Request: RequestType>(
+  /// Send a request to the LSP server and (asynchronously) receive a reply.
+  public nonisolated func send<Request: RequestType>(
     _ request: Request,
     reply: @escaping (LSPResult<Request.Response>) -> Void
   ) -> RequestID {

Tests/LanguageServerProtocolJSONRPCTests/ConnectionTests.swift

@@ -24,8 +24,7 @@ class ConnectionTests: XCTestCase {
   var connection: TestJSONRPCConnection! = nil
 
   override func setUp() {
-    connection = TestJSONRPCConnection()
-    connection.client.allowUnexpectedNotification = false
+    connection = TestJSONRPCConnection(allowUnexpectedNotification: false)
   }
 
   override func tearDown() {
@@ -52,12 +51,12 @@ class ConnectionTests: XCTestCase {
     waitForExpectations(timeout: defaultTimeout)
   }
 
-  func testMessageBuffer() throws {
+  func testMessageBuffer() async throws {
     let client = connection.client
     let clientConnection = connection.clientConnection
     let expectation = self.expectation(description: "note received")
 
-    client.appendOneShotNotificationHandler { (note: EchoNotification) in
+    await client.appendOneShotNotificationHandler { (note: EchoNotification) in
       XCTAssertEqual(note.string, "hello!")
       expectation.fulfill()
     }
@@ -76,11 +75,11 @@ class ConnectionTests: XCTestCase {
       _rawData: [note1Str.utf8.last!, note2Str.utf8.first!].withUnsafeBytes { DispatchData(bytes: $0) }
     )
 
-    waitForExpectations(timeout: defaultTimeout)
+    try await fulfillmentOfOrThrow([expectation])
 
     let expectation2 = self.expectation(description: "note received")
 
-    client.appendOneShotNotificationHandler { (note: EchoNotification) in
+    await client.appendOneShotNotificationHandler { (note: EchoNotification) in
       XCTAssertEqual(note.string, "no way!")
       expectation2.fulfill()
     }
@@ -89,7 +88,7 @@ class ConnectionTests: XCTestCase {
       clientConnection.send(_rawData: [b].withUnsafeBytes { DispatchData(bytes: $0) })
     }
 
-    waitForExpectations(timeout: defaultTimeout)
+    try await fulfillmentOfOrThrow([expectation2])
 
     // Close the connection before accessing _requestBuffer, which ensures we don't race.
     connection.serverConnection.close()
@@ -118,18 +117,18 @@ class ConnectionTests: XCTestCase {
     waitForExpectations(timeout: defaultTimeout)
   }
 
-  func testEchoNote() {
+  func testEchoNote() async throws {
     let client = connection.client
     let expectation = self.expectation(description: "note received")
 
-    client.appendOneShotNotificationHandler { (note: EchoNotification) in
+    await client.appendOneShotNotificationHandler { (note: EchoNotification) in
       XCTAssertEqual(note.string, "hello!")
       expectation.fulfill()
     }
 
     client.send(EchoNotification(string: "hello!"))
 
-    waitForExpectations(timeout: defaultTimeout)
+    try await fulfillmentOfOrThrow([expectation])
   }
 
   func testUnknownRequest() {
@@ -210,33 +209,33 @@ class ConnectionTests: XCTestCase {
     waitForExpectations(timeout: defaultTimeout)
   }
 
-  func testSendBeforeClose() {
+  func testSendBeforeClose() async throws {
     let client = connection.client
     let server = connection.server
 
     let expectation = self.expectation(description: "received notification")
-    client.appendOneShotNotificationHandler { (note: EchoNotification) in
+    await client.appendOneShotNotificationHandler { (note: EchoNotification) in
       expectation.fulfill()
     }
 
     server.client.send(EchoNotification(string: "about to close!"))
     connection.serverConnection.close()
 
-    waitForExpectations(timeout: defaultTimeout)
+    try await fulfillmentOfOrThrow([expectation])
   }
 
-  func testSendSynchronouslyBeforeClose() {
+  func testSendSynchronouslyBeforeClose() async throws {
     let client = connection.client
 
     let expectation = self.expectation(description: "received notification")
-    client.appendOneShotNotificationHandler { (note: EchoNotification) in
+    await client.appendOneShotNotificationHandler { (note: EchoNotification) in
       expectation.fulfill()
     }
     let notification = EchoNotification(string: "about to close!")
     connection.serverConnection._send(.notification(notification), async: false)
     connection.serverConnection.close()
 
-    waitForExpectations(timeout: defaultTimeout)
+    try await fulfillmentOfOrThrow([expectation])
   }
 
   /// We can explicitly close a connection, but the connection also

Tests/LanguageServerProtocolTests/ConnectionTests.swift

@@ -19,8 +19,7 @@ class ConnectionTests: XCTestCase {
   var connection: TestLocalConnection! = nil
 
   override func setUp() {
-    connection = TestLocalConnection()
-    connection.client.allowUnexpectedNotification = false
+    connection = TestLocalConnection(allowUnexpectedNotification: false)
   }
 
   override func tearDown() {
@@ -61,17 +60,17 @@ class ConnectionTests: XCTestCase {
     waitForExpectations(timeout: defaultTimeout)
   }
 
-  func testEchoNote() {
+  func testEchoNote() async throws {
     let client = connection.client
     let expectation = self.expectation(description: "note received")
 
-    client.appendOneShotNotificationHandler { (note: EchoNotification) in
+    await client.appendOneShotNotificationHandler { (note: EchoNotification) in
       XCTAssertEqual(note.string, "hello!")
       expectation.fulfill()
     }
 
     client.send(EchoNotification(string: "hello!"))
 
-    waitForExpectations(timeout: defaultTimeout)
+    try await fulfillmentOfOrThrow([expectation])
   }
 }