Make SwiftLanguageServer and ClangLanguageServerShim call directly into SourceKitServer instead of through a LocalConnection

`LocalConnection` with its dynamic registration of a message handler made the overall design unnecessarily complicated. If we just call `SourceKitServer` from `ClangLanguageServerShim` and `SwiftLanguageServer` directly, it’s a lot more obvious, what’s going on, IMO.

Author: ahoppen

Sources/SourceKitLSP/Clang/ClangLanguageServer.swift

@@ -45,10 +45,10 @@ actor ClangLanguageServerShim: ToolchainLanguageServer, MessageHandler {
   /// The queue on which clangd calls us back.
   public let clangdCommunicationQueue: DispatchQueue = DispatchQueue(label: "language-server-queue", qos: .userInitiated)
 
-  /// The connection to the client. In the case of `ClangLanguageServerShim`,
-  /// the client is always a ``SourceKitServer``, which will forward the request
-  /// to the editor.
-  public let client: Connection
+  /// The ``SourceKitServer`` instance that created this `ClangLanguageServerShim`.
+  ///
+  /// Used to send requests and notifications to the editor.
+  private weak var sourceKitServer: SourceKitServer?
 
   /// The connection to the clangd LSP. `nil` until `startClangdProcesss` has been called.
   var clangd: Connection!
@@ -93,9 +93,6 @@ actor ClangLanguageServerShim: ToolchainLanguageServer, MessageHandler {
   /// `clangd` doesn't have support for multi-root workspaces, so we need to start a separate `clangd` instance for every workspace root.
   private let workspace: WeakWorkspace
 
-  /// A callback with which `ClangLanguageServer` can request its owner to reopen all documents in case it has crashed.
-  private let reopenDocuments: (ToolchainLanguageServer) async -> Void
-
   /// The documents that have been opened and which language they have been
   /// opened with.
   private var openDocuments: [DocumentURI: Language] = [:]
@@ -110,12 +107,10 @@ actor ClangLanguageServerShim: ToolchainLanguageServer, MessageHandler {
   /// Creates a language server for the given client referencing the clang binary specified in `toolchain`.
   /// Returns `nil` if `clangd` can't be found.
   public init?(
-    client: LocalConnection,
+    sourceKitServer: SourceKitServer,
     toolchain: Toolchain,
     options: SourceKitServer.Options,
-    workspace: Workspace,
-    reopenDocuments: @escaping (ToolchainLanguageServer) async -> Void,
-    workspaceForDocument: @escaping (DocumentURI) async -> Workspace?
+    workspace: Workspace
   ) async throws {
     guard let clangdPath = toolchain.clangd else {
       return nil
@@ -124,9 +119,8 @@ actor ClangLanguageServerShim: ToolchainLanguageServer, MessageHandler {
     self.clangdPath = clangdPath
     self.clangdOptions = options.clangdOptions
     self.workspace = WeakWorkspace(workspace)
-    self.reopenDocuments = reopenDocuments
     self.state = .connected
-    self.client = client
+    self.sourceKitServer = sourceKitServer
     try startClangdProcesss()
   }
 
@@ -245,7 +239,11 @@ actor ClangLanguageServerShim: ToolchainLanguageServer, MessageHandler {
         // But since SourceKitServer more or less ignores them right now anyway, this should be fine for now.
         _ = try self.initializeSync(initializeRequest)
         self.clientInitialized(InitializedNotification())
-        await self.reopenDocuments(self)
+        if let sourceKitServer {
+          await sourceKitServer.reopenDocuments(for: self)
+        } else {
+          log("Cannot reopen documents because SourceKitServer is no longer alive", level: .error)
+        }
         self.state = .connected
       } catch {
         log("Failed to restart clangd after a crash.", level: .error)
@@ -256,12 +254,15 @@ actor ClangLanguageServerShim: ToolchainLanguageServer, MessageHandler {
   /// Handler for notifications received **from** clangd, ie. **clangd** is
   /// sending a notification that's intended for the editor.
   ///
-  /// We should either handle it ourselves or forward it to the client.
+  /// We should either handle it ourselves or forward it to the editor.
   func handle(_ params: some NotificationType, from clientID: ObjectIdentifier) async {
-    if let publishDiags = params as? PublishDiagnosticsNotification {
+    switch params {
+    case let publishDiags as PublishDiagnosticsNotification:
       await self.publishDiagnostics(Notification(publishDiags, clientID: clientID))
-    } else if clientID == ObjectIdentifier(self.clangd) {
-      self.client.send(params)
+    default:
+      // We don't know how to handle any other notifications and ignore them.
+      log("Ignoring unknown notification \(type(of: params))", level: .warning)
+      break
     }
   }
 
@@ -274,13 +275,19 @@ actor ClangLanguageServerShim: ToolchainLanguageServer, MessageHandler {
     id: RequestID,
     from clientID: ObjectIdentifier,
     reply: @escaping (LSPResult<R.Response>) -> Void
-  ) {
+  ) async {
     let request = Request(params, id: id, clientID: clientID, cancellation: CancellationToken(), reply: { result in
       reply(result)
     })
+    guard let sourceKitServer else {
+      // `SourceKitServer` has been destructed. We are tearing down the language
+      // server. Nothing left to do.
+      request.reply(.failure(.unknown("Connection to the editor closed")))
+      return
+    }
 
     if request.clientID == ObjectIdentifier(self.clangd) {
-      self.forwardRequest(request, to: self.client)
+      await sourceKitServer.sendRequestToClient(request.params, reply: request.reply)
     } else {
       request.reply(.failure(ResponseError.methodNotFound(R.method)))
     }
@@ -354,12 +361,21 @@ extension ClangLanguageServerShim {
     // non-fallback settings very shortly after, which will override the
     // incorrect result, making it very temporary.
     let buildSettings = await self.buildSettings(for: params.uri)
+    guard let sourceKitServer else {
+      log("Cannot publish diagnostics because SourceKitServer has been destroyed", level: .error)
+      return
+    }
     if buildSettings?.isFallback ?? true {
       // Fallback: send empty publish notification instead.
-      client.send(PublishDiagnosticsNotification(
-        uri: params.uri, version: params.version, diagnostics: []))
+      await sourceKitServer.sendNotificationToClient(
+        PublishDiagnosticsNotification(
+          uri: params.uri,
+          version: params.version,
+          diagnostics: []
+        )
+      )
     } else {
-      client.send(note.params)
+      await sourceKitServer.sendNotificationToClient(note.params)
     }
   }
 
@@ -388,9 +404,6 @@ extension ClangLanguageServerShim {
         guard let self else { return }
         Task {
           await self.clangd.send(ExitNotification())
-          if let localConnection = self.client as? LocalConnection {
-            localConnection.close()
-          }
           continuation.resume()
         }
       }

Sources/SourceKitLSP/SourceKitServer.swift

@@ -292,12 +292,23 @@ public actor SourceKitServer {
     }
 
     // Unknown requests from a language server are passed on to the client.
-    let id = client.send(req.params, queue: clientCommunicationQueue) { result in
-      req.reply(result)
-    }
-    req.cancellationToken.addCancellationHandler {
-      self.client.send(CancelRequestNotification(id: id))
+    sendRequestToClient(req.params, reply: req.reply)
+  }
+
+  /// Send the given notification to the editor.
+  public func sendNotificationToClient(_ notification: some NotificationType) {
+    client.send(notification)
+  }
+
+  /// Send the given request to the editor.
+  ///
+  /// Return the result via the `reply` completion handler. `reply` is guaranteed
+  /// to be called exactly once.
+  public func sendRequestToClient<R: RequestType>(_ request: R, reply: @escaping (LSPResult<R.Response>) -> Void) {
+    _ = client.send(request, queue: clientCommunicationQueue) { result in
+      reply(result)
     }
+    // FIXME: (async) Handle cancellation
   }
 
   /// Handle an unknown notification.
@@ -388,19 +399,11 @@ public actor SourceKitServer {
 
     // Start a new service.
     return await orLog("failed to start language service", level: .error) {
-      let service = try await SourceKitLSP.languageService(
-        for: toolchain,
-        serverType,
+      let service = try await serverType.serverType.init(
+        sourceKitServer: self,
+        toolchain: toolchain,
         options: options,
-        client: self,
-        in: workspace,
-        reopenDocuments: { [weak self] toolchainLanguageServer in
-          await self?.reopenDocuments(for: toolchainLanguageServer)
-        },
-        workspaceForDocument: { [weak self] document in
-          guard let self else { return nil }
-          return await self.workspaceForDocument(uri: document)
-        }
+        workspace: workspace
       )
 
       guard let service else {
@@ -1989,34 +1992,6 @@ extension SourceKitServer {
   }
 }
 
-/// Creates a new connection from `client` to a service for `language` if available, and launches
-/// the service. Does *not* send the initialization request.
-///
-/// - returns: The connection, if a suitable language service is available; otherwise nil.
-/// - throws: If there is a suitable service but it fails to launch, throws an error.
-func languageService(
-  for toolchain: Toolchain,
-  _ languageServerType: LanguageServerType,
-  options: SourceKitServer.Options,
-  client: MessageHandler,
-  in workspace: Workspace,
-  reopenDocuments: @escaping (ToolchainLanguageServer) async -> Void,
-  workspaceForDocument: @escaping (DocumentURI) async -> Workspace?
-) async throws -> ToolchainLanguageServer? {
-  let connectionToClient = LocalConnection()
-
-  let server = try await languageServerType.serverType.init(
-    client: connectionToClient,
-    toolchain: toolchain,
-    options: options,
-    workspace: workspace,
-    reopenDocuments: reopenDocuments,
-    workspaceForDocument: workspaceForDocument
-  )
-  connectionToClient.start(handler: client)
-  return server
-}
-
 private func languageClass(for language: Language) -> [Language] {
   switch language {
   case .c, .cpp, .objective_c, .objective_cpp: