Merge pull request swiftlang#144 from rmaz/separateregistries

Pass in MessageRegistry instead of using a singleton

Author: benlangmuir

Sources/LanguageServerProtocol/MessageRegistry.swift

@@ -10,22 +10,19 @@
 //
 //===----------------------------------------------------------------------===//
 
-import Dispatch
 
 public final class MessageRegistry {
 
-  private(set)
-  var methodToRequest: [String: _RequestType.Type] = {
-    Dictionary(uniqueKeysWithValues: builtinRequests.map { ($0.method, $0) })
-  }()
+  public static let lspProtocol: MessageRegistry =
+    MessageRegistry(requests: builtinRequests, notifications: builtinNotifications)
 
-  private(set)
-  var methodToNotification: [String: NotificationType.Type] = {
-    Dictionary(uniqueKeysWithValues: builtinNotifications.map { ($0.method, $0) })
-  }()
+  private let methodToRequest: [String: _RequestType.Type]
+  private let methodToNotification: [String: NotificationType.Type]
 
-  /// The global message registry.
-  public static let shared: MessageRegistry = .init()
+  public init(requests: [_RequestType.Type], notifications: [NotificationType.Type]) {
+    self.methodToRequest = Dictionary(uniqueKeysWithValues: requests.map { ($0.method, $0) })
+    self.methodToNotification = Dictionary(uniqueKeysWithValues: notifications.map { ($0.method, $0) })
+  }
 
   /// Returns the type of the message named `method`, or nil if it is unknown.
   public func requestType(for method: String) -> _RequestType.Type? {
@@ -37,15 +34,4 @@ public final class MessageRegistry {
     return methodToNotification[method]
   }
 
-  /// Adds a new message type to the registry. **For test messages only; not thread-safe!**.
-  public func _register(_ type: _RequestType.Type) {
-    precondition(methodToRequest[type.method] == nil)
-    methodToRequest[type.method] = type
-  }
-
-  /// Adds a new message type to the registry. **For test messages only; not thread-safe!**.
-  public func _register(_ type: NotificationType.Type) {
-    precondition(methodToNotification[type.method] == nil)
-    methodToNotification[type.method] = type
-  }
 }

Sources/LanguageServerProtocolJSONRPC/JSONRPCConnection.swift

@@ -25,6 +25,7 @@ public final class JSONRPCConection {
   let sendQueue: DispatchQueue = DispatchQueue(label: "jsonrpc-send-queue", qos: .userInitiated)
   let receiveIO: DispatchIO
   let sendIO: DispatchIO
+  let messageRegistry: MessageRegistry
 
   enum State {
     case created, running, closed
@@ -50,9 +51,10 @@ public final class JSONRPCConection {
 
   var closeHandler: () -> Void
 
-  public init(inFD: Int32, outFD: Int32, closeHandler: @escaping () -> Void = {}) {
+  public init(protocol messageRegistry: MessageRegistry, inFD: Int32, outFD: Int32, closeHandler: @escaping () -> Void = {}) {
     state = .created
     self.closeHandler = closeHandler
+    self.messageRegistry = messageRegistry
 
     receiveIO = DispatchIO(type: .stream, fileDescriptor: inFD, queue: queue) { (error: Int32) in
       if error != 0 {
@@ -137,6 +139,9 @@ public final class JSONRPCConection {
 
     let decoder = JSONDecoder()
 
+    // Set message registry to use for model decoding.
+    decoder.userInfo[.messageRegistryKey] = messageRegistry
+
     // Setup callback for response type.
     decoder.userInfo[.responseTypeCallbackKey] = { id in
       guard let outstanding = self.outstandingRequests[id] else {

Sources/LanguageServerProtocolJSONRPC/MessageCoding.swift

@@ -22,6 +22,7 @@ public enum JSONRPCMessage {
 
 extension CodingUserInfoKey {
   public static let responseTypeCallbackKey: CodingUserInfoKey = CodingUserInfoKey(rawValue: "lsp.jsonrpc.responseTypeCallback")!
+  public static let messageRegistryKey: CodingUserInfoKey = CodingUserInfoKey(rawValue: "lsp.jsonrpc.messageRegistry")!
 }
 
 extension JSONRPCMessage: Codable {
@@ -39,6 +40,9 @@ extension JSONRPCMessage: Codable {
 
   public init(from decoder: Decoder) throws {
 
+    guard let messageRegistry = decoder.userInfo[.messageRegistryKey] as? MessageRegistry else {
+      fatalError("missing or invalid messageRegistryKey on decoder")
+    }
     let container = try decoder.container(keyedBy: CodingKeys.self)
     let jsonrpc = try container.decodeIfPresent(String.self, forKey: .jsonrpc)
     let id = try container.decodeIfPresent(RequestID.self, forKey: .id)
@@ -58,7 +62,7 @@ extension JSONRPCMessage: Codable {
       case (nil, let method?, _, nil):
         msgKind = .notification
 
-        guard let messageType = MessageRegistry.shared.notificationType(for: method) else {
+        guard let messageType = messageRegistry.notificationType(for: method) else {
           throw MessageDecodingError.methodNotFound(method)
         }
 
@@ -69,7 +73,7 @@ extension JSONRPCMessage: Codable {
       case (let id?, let method?, _, nil):
         msgKind = .request
 
-        guard let messageType = MessageRegistry.shared.requestType(for: method) else {
+        guard let messageType = messageRegistry.requestType(for: method) else {
           throw MessageDecodingError.methodNotFound(method)
         }
 

Sources/SKTestSupport/TestJSONRPCConnection.swift

@@ -19,8 +19,6 @@ import XCTest
 // Workaround ambiguity with Foundation.
 public typealias Notification = LanguageServerProtocol.Notification
 
-public let initRequestsOnce: Void = registerTestRequests()
-
 public struct TestJSONRPCConnection {
   public let clientToServer: Pipe = Pipe()
   public let serverToClient: Pipe = Pipe()
@@ -30,19 +28,19 @@ public struct TestJSONRPCConnection {
   public let serverConnection: JSONRPCConection
 
   public init() {
-    _ = initRequestsOnce
-
     // FIXME: DispatchIO doesn't like when the Pipes close behind its back even after the tests
     // finish. Until we fix the lifetime, leak.
     _ = Unmanaged.passRetained(clientToServer)
     _ = Unmanaged.passRetained(serverToClient)
 
     clientConnection = JSONRPCConection(
+      protocol: testMessageRegistry,
       inFD: serverToClient.fileHandleForReading.fileDescriptor,
       outFD: clientToServer.fileHandleForWriting.fileDescriptor
     )
 
     serverConnection = JSONRPCConection(
+      protocol: testMessageRegistry,
       inFD: clientToServer.fileHandleForReading.fileDescriptor,
       outFD: serverToClient.fileHandleForWriting.fileDescriptor
     )
@@ -67,8 +65,6 @@ public struct TestLocalConnection {
   public let serverConnection: LocalConnection = .init()
 
   public init() {
-    _ = initRequestsOnce
-
     client = TestClient(server: serverConnection)
     server = TestServer(client: clientConnection)
 
@@ -230,16 +226,9 @@ public final class TestServer: LanguageServer {
 
 // MARK: Test requests.
 
-public func registerTestRequests() {
-  [
-    EchoRequest.self,
-    EchoError.self,
-  ].forEach { MessageRegistry.shared._register($0) }
-
-  [
-    EchoNotification.self,
-  ].forEach { MessageRegistry.shared._register($0) }
-}
+private let testMessageRegistry = MessageRegistry(
+  requests: [EchoRequest.self, EchoError.self],
+  notifications: [EchoNotification.self])
 
 extension String: ResponseType {}
 

Sources/SKTestSupport/TestServer.swift

@@ -45,7 +45,6 @@ public struct TestSourceKitServer {
   public let server: SourceKitServer?
 
   public init(connectionKind: ConnectionKind = .local) {
-     _ = initRequestsOnce
 
     switch connectionKind {
       case .local:
@@ -71,10 +70,12 @@ public struct TestSourceKitServer {
         _ = Unmanaged.passRetained(serverToClient)
 
         let clientConnection = JSONRPCConection(
+          protocol: MessageRegistry.lspProtocol,
           inFD: serverToClient.fileHandleForReading.fileDescriptor,
           outFD: clientToServer.fileHandleForWriting.fileDescriptor
         )
         let serverConnection = JSONRPCConection(
+          protocol: MessageRegistry.lspProtocol,
           inFD: clientToServer.fileHandleForReading.fileDescriptor,
           outFD: serverToClient.fileHandleForWriting.fileDescriptor
         )

Sources/SourceKit/clangd/ClangLanguageServer.swift

@@ -134,6 +134,7 @@ func makeJSONRPCClangServer(client: MessageHandler, toolchain: Toolchain, buildS
   let serverToClient: Pipe = Pipe()
 
   let connection = JSONRPCConection(
+    protocol: MessageRegistry.lspProtocol,
     inFD: serverToClient.fileHandleForReading.fileDescriptor,
     outFD: clientToServer.fileHandleForWriting.fileDescriptor
   )

Sources/sourcekit-lsp/main.swift

@@ -53,7 +53,11 @@ func parseArguments() throws -> BuildSetup {
     flags: buildFlags)
 }
 
-let clientConnection = JSONRPCConection(inFD: STDIN_FILENO, outFD: STDOUT_FILENO, closeHandler: {
+let clientConnection = JSONRPCConection(
+  protocol: MessageRegistry.lspProtocol,
+  inFD: STDIN_FILENO,
+  outFD: STDOUT_FILENO,
+  closeHandler: {
   exit(0)
 })
 

Tests/LanguageServerProtocolJSONRPCTests/CodingTests.swift

@@ -191,7 +191,7 @@ final class CodingTests: XCTestCase {
       return $0 == .string("unknown") ? nil : InitializeResult.self
     }
 
-    let info = [CodingUserInfoKey.responseTypeCallbackKey: responseTypeCallback]
+    let info = defaultCodingInfo.merging([CodingUserInfoKey.responseTypeCallbackKey: responseTypeCallback]) { (_, new) in new }
 
     checkMessageDecodingError(MessageDecodingError.invalidRequest("message not recognized as request, response or notification", id: .number(2)), json: """
     {"jsonrpc":"2.0","id":2,"params":{}}
@@ -215,8 +215,10 @@ final class CodingTests: XCTestCase {
   }
 }
 
+let defaultCodingInfo: [CodingUserInfoKey: Any] = [CodingUserInfoKey.messageRegistryKey:MessageRegistry.lspProtocol]
+
 private func checkMessageCoding<Request>(_ value: Request, id: RequestID, json: String, file: StaticString = #file, line: UInt = #line) where Request: RequestType & Equatable {
-  checkCoding(JSONRPCMessage.request(value, id: id), json: json, file: file, line: line) {
+  checkCoding(JSONRPCMessage.request(value, id: id), json: json, userInfo: defaultCodingInfo, file: file, line: line) {
 
     guard case JSONRPCMessage.request(let decodedValueOpaque, let decodedID) = $0, let decodedValue = decodedValueOpaque as? Request else {
       XCTFail("decodedValue \($0) does not match expected \(value)", file: file, line: line)
@@ -229,7 +231,7 @@ private func checkMessageCoding<Request>(_ value: Request, id: RequestID, json:
 }
 
 private func checkMessageCoding<Notification>(_ value: Notification, json: String, file: StaticString = #file, line: UInt = #line) where Notification: NotificationType & Equatable {
-  checkCoding(JSONRPCMessage.notification(value), json: json, file: file, line: line) {
+  checkCoding(JSONRPCMessage.notification(value), json: json, userInfo: defaultCodingInfo, file: file, line: line) {
 
     guard case JSONRPCMessage.notification(let decodedValueOpaque) = $0, let decodedValue = decodedValueOpaque as? Notification else {
       XCTFail("decodedValue \($0) does not match expected \(value)", file: file, line: line)
@@ -246,7 +248,10 @@ private func checkMessageCoding<Response>(_ value: Response, id: RequestID, json
     return $0 == .string("unknown") ? nil : Response.self
   }
 
-  checkCoding(JSONRPCMessage.response(value, id: id), json: json, userInfo: [.responseTypeCallbackKey: callback], file: file, line: line) {
+  var codingInfo = defaultCodingInfo
+  codingInfo[.responseTypeCallbackKey] = callback
+
+  checkCoding(JSONRPCMessage.response(value, id: id), json: json, userInfo: codingInfo, file: file, line: line) {
 
     guard case JSONRPCMessage.response(let decodedValueOpaque, let decodedID) = $0, let decodedValue = decodedValueOpaque as? Response else {
       XCTFail("decodedValue \($0) does not match expected \(value)", file: file, line: line)
@@ -259,7 +264,7 @@ private func checkMessageCoding<Response>(_ value: Response, id: RequestID, json
 }
 
 private func checkMessageCoding(_ value: ResponseError, id: RequestID, json: String, file: StaticString = #file, line: UInt = #line) {
-  checkCoding(JSONRPCMessage.errorResponse(value, id: id), json: json, file: file, line: line) {
+  checkCoding(JSONRPCMessage.errorResponse(value, id: id), json: json, userInfo: defaultCodingInfo, file: file, line: line) {
 
     guard case JSONRPCMessage.errorResponse(let decodedValue, let decodedID) = $0 else {
       XCTFail("decodedValue \($0) does not match expected \(value)", file: file, line: line)
@@ -271,7 +276,7 @@ private func checkMessageCoding(_ value: ResponseError, id: RequestID, json: Str
   }
 }
 
-private func checkMessageDecodingError(_ expected: MessageDecodingError, json: String, userInfo: [CodingUserInfoKey: Any] = [:], file: StaticString = #file, line: UInt = #line) {
+private func checkMessageDecodingError(_ expected: MessageDecodingError, json: String, userInfo: [CodingUserInfoKey: Any] = defaultCodingInfo, file: StaticString = #file, line: UInt = #line) {
   let data = json.data(using: .utf8)!
   let decoder = JSONDecoder()
   decoder.userInfo = userInfo