Merge pull request #1120 from ahoppen/ahoppen/fix-tsan-issue

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