Add a send method to InProcessSourceKitLSPClient and Connection in which the client specifies the request ID

Author: ahoppen

Sources/InProcessClient/InProcessSourceKitLSPClient.swift

@@ -14,6 +14,7 @@ import BuildSystemIntegration
 public import Foundation
 public import LanguageServerProtocol
 import LanguageServerProtocolExtensions
+import SKLogging
 import SwiftExtensions
 import TSCExtensions
 
@@ -124,11 +125,28 @@ public final class InProcessSourceKitLSPClient: Sendable {
     _ request: R,
     reply: @Sendable @escaping (LSPResult<R.Response>) -> Void
   ) -> RequestID {
-    let requestID = RequestID.number(Int(nextRequestID.fetchAndIncrement()))
+    let requestID = RequestID.string("sk-\(Int(nextRequestID.fetchAndIncrement()))")
     server.handle(request, id: requestID, reply: reply)
     return requestID
   }
 
+  /// Send the request to `server` and return the request result via a completion handler.
+  ///
+  /// The request ID must not start with `sk-` to avoid conflicting with the request IDs that are created by
+  /// `send(:reply:)`.
+  public func send<R: RequestType>(
+    _ request: R,
+    id: RequestID,
+    reply: @Sendable @escaping (LSPResult<R.Response>) -> Void
+  ) {
+    if case .string(let string) = id {
+      if string.starts(with: "sk-") {
+        logger.fault("Manually specified request ID must not have reserved prefix 'sk-'")
+      }
+    }
+    server.handle(request, id: id, reply: reply)
+  }
+
   /// Send the notification to `server`.
   public func send(_ notification: some NotificationType) {
     server.handle(notification)

Sources/LanguageServerProtocol/Connection.swift

@@ -15,11 +15,36 @@ public protocol Connection: Sendable {
   /// Send a notification without a reply.
   func send(_ notification: some NotificationType)
 
-  /// Send a request and (asynchronously) receive a reply.
+  /// Generate a new request ID to be used in the `send` method that does not take an explicit request ID.
+  ///
+  /// These request IDs need to be unique and must not conflict with any request ID that clients might manually specify
+  /// to `send(_:id:reply:)`.
+  ///
+  /// To allow, this request IDs starting with `sk-` are reserved to only be generated by this method and are not
+  /// allowed to be passed directly to `send(_:id:reply:)`. Thus, generating request IDs prefixed with `sk-` here is
+  /// safe. Similarly returning UUID-based requests IDs is safe because UUIDs are already unique.
+  func nextRequestID() -> RequestID
+
+  /// Send a request with a pre-defined request ID and (asynchronously) receive a reply.
+  ///
+  /// The request ID must not conflict with any request ID generated by `nextRequestID()`.
   func send<Request: RequestType>(
     _ request: Request,
+    id: RequestID,
     reply: @escaping @Sendable (LSPResult<Request.Response>) -> Void
-  ) -> RequestID
+  )
+}
+
+extension Connection {
+  /// Send a request and (asynchronously) receive a reply.
+  public func send<Request: RequestType>(
+    _ request: Request,
+    reply: @escaping @Sendable (LSPResult<Request.Response>) -> Void
+  ) -> RequestID {
+    let id = nextRequestID()
+    self.send(request, id: id, reply: reply)
+    return id
+  }
 }
 
 /// An abstract message handler, such as a language server or client.

Sources/LanguageServerProtocolExtensions/LocalConnection.swift

@@ -10,16 +10,15 @@
 //
 //===----------------------------------------------------------------------===//
 
-#if compiler(>=6)
 import Dispatch
-package import LanguageServerProtocol
 import LanguageServerProtocolJSONRPC
 import SKLogging
+import SwiftExtensions
+
+#if compiler(>=6)
+package import LanguageServerProtocol
 #else
-import Dispatch
 import LanguageServerProtocol
-import LanguageServerProtocolJSONRPC
-import SKLogging
 #endif
 
 /// A connection between two message handlers in the same process.
@@ -45,8 +44,7 @@ package final class LocalConnection: Connection, Sendable {
   /// The queue guarding `_nextRequestID`.
   private let queue: DispatchQueue = DispatchQueue(label: "local-connection-queue")
 
-  /// - Important: Must only be accessed from `queue`
-  nonisolated(unsafe) private var _nextRequestID: Int = 0
+  private let _nextRequestID = AtomicUInt32(initialValue: 0)
 
   /// - Important: Must only be accessed from `queue`
   nonisolated(unsafe) private var state: State = .ready
@@ -88,11 +86,8 @@ package final class LocalConnection: Connection, Sendable {
     }
   }
 
-  func nextRequestID() -> RequestID {
-    return queue.sync {
-      _nextRequestID += 1
-      return .number(_nextRequestID)
-    }
+  public func nextRequestID() -> RequestID {
+    return .string("sk-\(_nextRequestID.fetchAndIncrement())")
   }
 
   package func send<Notification: NotificationType>(_ notification: Notification) {
@@ -110,10 +105,9 @@ package final class LocalConnection: Connection, Sendable {
 
   package func send<Request: RequestType>(
     _ request: Request,
+    id: RequestID,
     reply: @Sendable @escaping (LSPResult<Request.Response>) -> Void
-  ) -> RequestID {
-    let id = nextRequestID()
-
+  ) {
     logger.info(
       """
       Sending request to \(self.name, privacy: .public) (id: \(id, privacy: .public)):
@@ -128,7 +122,7 @@ package final class LocalConnection: Connection, Sendable {
         """
       )
       reply(.failure(.serverCancelled))
-      return id
+      return
     }
 
     precondition(self.state == .started)
@@ -153,7 +147,5 @@ package final class LocalConnection: Connection, Sendable {
       }
       reply(result)
     }
-
-    return id
   }
 }

Sources/LanguageServerProtocolJSONRPC/JSONRPCConnection.swift

@@ -97,10 +97,7 @@ public final class JSONRPCConnection: Connection {
   }
 
   /// An integer that hasn't been used for a request ID yet.
-  ///
-  /// Access to this must be be guaranteed to be sequential to avoid data races. Currently, all access are
-  ///  - `nextRequestID()`: This is synchronized on `queue`.
-  private nonisolated(unsafe) var nextRequestIDStorage: Int = 0
+  let nextRequestIDStorage = AtomicUInt32(initialValue: 0)
 
   struct OutstandingRequest: Sendable {
     var responseType: ResponseType.Type
@@ -663,12 +660,8 @@ public final class JSONRPCConnection: Connection {
   }
 
   /// Request id for the next outgoing request.
-  ///
-  /// - Important: Must be called on `queue`
-  private func nextRequestID() -> RequestID {
-    dispatchPrecondition(condition: .onQueue(queue))
-    nextRequestIDStorage += 1
-    return .number(nextRequestIDStorage)
+  public func nextRequestID() -> RequestID {
+    return .string("sk-\(nextRequestIDStorage.fetchAndIncrement())")
   }
 
   // MARK: Connection interface
@@ -691,14 +684,13 @@ public final class JSONRPCConnection: Connection {
   /// When the receiving end replies to the request, execute `reply` with the response.
   public func send<Request: RequestType>(
     _ request: Request,
+    id: RequestID,
     reply: @escaping @Sendable (LSPResult<Request.Response>) -> Void
-  ) -> RequestID {
-    let id: RequestID = self.queue.sync {
-      let id = nextRequestID()
-
+  ) {
+    self.queue.sync {
       guard readyToSend() else {
         reply(.failure(.serverCancelled))
-        return id
+        return
       }
 
       outstandingRequests[id] = OutstandingRequest(
@@ -734,10 +726,8 @@ public final class JSONRPCConnection: Connection {
       )
 
       send(.request(request, id: id))
-      return id
+      return
     }
-
-    return id
   }
 
   /// After the remote side of the connection sent a request to us, return a reply to the remote side.

Sources/SKTestSupport/DummyBuildSystemManagerConnectionToClient.swift

@@ -29,12 +29,16 @@ package struct DummyBuildSystemManagerConnectionToClient: BuildSystemManagerConn
 
   package func send(_ notification: some LanguageServerProtocol.NotificationType) {}
 
+  package func nextRequestID() -> RequestID {
+    return .string(UUID().uuidString)
+  }
+
   package func send<Request: RequestType>(
     _ request: Request,
+    id: RequestID,
     reply: @escaping @Sendable (LSPResult<Request.Response>) -> Void
-  ) -> RequestID {
+  ) {
     reply(.failure(ResponseError.unknown("Not implemented")))
-    return .string(UUID().uuidString)
   }
 
   package func watchFiles(_ fileWatchers: [LanguageServerProtocol.FileSystemWatcher]) async {}

Sources/SourceKitLSP/Workspace.swift

@@ -230,17 +230,22 @@ package final class Workspace: Sendable, BuildSystemManagerDelegate {
         sourceKitLSPServer.sendNotificationToClient(notification)
       }
 
+      func nextRequestID() -> RequestID {
+        return .string(UUID().uuidString)
+      }
+
       func send<Request: RequestType>(
         _ request: Request,
+        id: RequestID,
         reply: @escaping @Sendable (LSPResult<Request.Response>) -> Void
-      ) -> RequestID {
+      ) {
         guard let sourceKitLSPServer else {
           // `SourceKitLSPServer` has been destructed. We are tearing down the
           // language server. Nothing left to do.
           reply(.failure(ResponseError.unknown("Connection to the editor closed")))
-          return .string(UUID().uuidString)
+          return
         }
-        return sourceKitLSPServer.client.send(request, reply: reply)
+        sourceKitLSPServer.client.send(request, id: id, reply: reply)
       }
 
       /// Whether the client can handle `WorkDoneProgress` requests.