Shutdown toolchain connections on exit

When using SourceKit-LSP in tests (or otherwise in a library), we do not
want to leak the toolchain connections.

Author: benlangmuir

Sources/LanguageServerProtocol/Connection.swift

@@ -2,7 +2,7 @@
 //
 // This source file is part of the Swift.org open source project
 //
-// Copyright (c) 2014 - 2018 Apple Inc. and the Swift project authors
+// Copyright (c) 2014 - 2020 Apple Inc. and the Swift project authors
 // Licensed under Apache License v2.0 with Runtime Library Exception
 //
 // See https://swift.org/LICENSE.txt for license information
@@ -104,14 +104,18 @@ public final class LocalConnection {
 
 extension LocalConnection: Connection {
   public func send<Notification>(_ notification: Notification) where Notification: NotificationType {
-    precondition(state == .started)
-    handler!.handle(notification, from: ObjectIdentifier(self))
+    handler?.handle(notification, from: ObjectIdentifier(self))
   }
 
   public func send<Request>(_ request: Request, queue: DispatchQueue, reply: @escaping (LSPResult<Request.Response>) -> Void) -> RequestID where Request: RequestType {
-    precondition(state == .started)
     let id = nextRequestID()
-    handler!.handle(request, id: id, from: ObjectIdentifier(self)) { result in
+    guard let handler = handler else {
+      queue.async { reply(.failure(.cancelled)) }
+      return id
+    }
+
+    precondition(state == .started)
+    handler.handle(request, id: id, from: ObjectIdentifier(self)) { result in
       queue.async {
         reply(result)
       }

Sources/SKTestSupport/TestServer.swift

@@ -19,7 +19,7 @@ import SourceKitLSP
 import class Foundation.Pipe
 import LSPTestSupport
 
-public struct TestSourceKitServer {
+public final class TestSourceKitServer {
   public enum ConnectionKind {
     case local, jsonrpc
   }
@@ -40,6 +40,8 @@ public struct TestSourceKitServer {
   public let client: TestClient
   let connImpl: ConnectionImpl
 
+  public var hasShutdown: Bool = false
+
   /// The server, if it is in the same process.
   public let server: SourceKitServer?
 
@@ -98,15 +100,15 @@ public struct TestSourceKitServer {
     }
   }
 
+  deinit {
+    close()
+  }
+
   func close() {
-    switch connImpl {
-      case .local(clientConnection: let cc, serverConnection: let sc):
-      cc.close()
-      sc.close()
-
-    case .jsonrpc(clientToServer: _, serverToClient: _, clientConnection: let cc, serverConnection: let sc):
-      cc.close()
-      sc.close()
+    if !hasShutdown {
+      hasShutdown = true
+      _ = try! self.client.sendSync(ShutdownRequest())
+      self.client.send(ExitNotification())
     }
   }
 }

Sources/SourceKitLSP/Clang/ClangLanguageServer.swift

@@ -26,16 +26,23 @@ final class ClangLanguageServerShim: ToolchainLanguageServer {
 
   let clangd: Connection
 
+  let client: LocalConnection
+
   var capabilities: ServerCapabilities? = nil
 
   let buildSystem: BuildSystem
 
   let clang: AbsolutePath?
 
   /// Creates a language server for the given client using the sourcekitd dylib at the specified path.
-  public init(client: Connection, clangd: Connection, buildSystem: BuildSystem,
-              clang: AbsolutePath?) throws {
+  public init(
+    client: LocalConnection,
+    clangd: Connection,
+    buildSystem: BuildSystem,
+    clang: AbsolutePath?
+  ) throws {
     self.clangd = clangd
+    self.client = client
     self.buildSystem = buildSystem
     self.clang = clang
   }
@@ -82,6 +89,13 @@ extension ClangLanguageServerShim {
     clangd.send(initialized)
   }
 
+  public func shutdown() {
+    _ = clangd.send(ShutdownRequest(), queue: queue) { [weak self] _ in
+      self?.clangd.send(ExitNotification())
+      self?.client.close()
+    }
+  }
+
   // MARK: - Text synchronization
 
   public func openDocument(_ note: DidOpenTextDocumentNotification) {

Sources/SourceKitLSP/SourceKitServer.swift

@@ -520,6 +520,10 @@ extension SourceKitServer {
 
   func shutdown(_ request: Request<ShutdownRequest>) {
     _prepareForExit()
+    for service in languageService.values {
+      service.shutdown()
+    }
+    languageService = [:]
     request.reply(VoidResponse())
   }
 

Sources/SourceKitLSP/Swift/SwiftLanguageServer.swift

@@ -2,7 +2,7 @@
 //
 // This source file is part of the Swift.org open source project
 //
-// Copyright (c) 2014 - 2019 Apple Inc. and the Swift project authors
+// Copyright (c) 2014 - 2020 Apple Inc. and the Swift project authors
 // Licensed under Apache License v2.0 with Runtime Library Exception
 //
 // See https://swift.org/LICENSE.txt for license information
@@ -81,7 +81,7 @@ public final class SwiftLanguageServer: ToolchainLanguageServer {
   /// The server's request queue, used to serialize requests and responses to `sourcekitd`.
   public let queue: DispatchQueue = DispatchQueue(label: "swift-language-server-queue", qos: .userInitiated)
 
-  let client: Connection
+  let client: LocalConnection
 
   let sourcekitd: SourceKitD
 
@@ -96,20 +96,23 @@ public final class SwiftLanguageServer: ToolchainLanguageServer {
 
   var commandsByFile: [DocumentURI: SwiftCompileCommand] = [:]
 
-  let onExit: () -> Void
-
   var keys: sourcekitd_keys { return sourcekitd.keys }
   var requests: sourcekitd_requests { return sourcekitd.requests }
   var values: sourcekitd_values { return sourcekitd.values }
 
-  /// Creates a language server for the given client using the sourcekitd dylib at the specified path.
-  public init(client: Connection, sourcekitd: AbsolutePath, buildSystem: BuildSystem, clientCapabilities: ClientCapabilities, onExit: @escaping () -> Void = {}) throws {
+  /// Creates a language server for the given client using the sourcekitd dylib at the specified
+  /// path.
+  public init(
+    client: LocalConnection,
+    sourcekitd: AbsolutePath,
+    buildSystem: BuildSystem,
+    clientCapabilities: ClientCapabilities
+  ) throws {
     self.client = client
     self.sourcekitd = try SourceKitDImpl.getOrCreate(dylibPath: sourcekitd)
     self.buildSystem = buildSystem
     self.clientCapabilities = clientCapabilities
     self.documentManager = DocumentManager()
-    self.onExit = onExit
   }
 
   /// Should be called on self.queue.
@@ -202,12 +205,9 @@ extension SwiftLanguageServer {
     // Nothing to do.
   }
 
-  func shutdown(_ request: Request<ShutdownRequest>) {
+  public func shutdown() {
     sourcekitd.removeNotificationHandler(self)
-  }
-
-  func exit(_ notification: Notification<ExitNotification>) {
-    onExit()
+    client.close()
   }
 
   // MARK: - Build System Integration

Sources/SourceKitLSP/ToolchainLanguageServer.swift

@@ -2,7 +2,7 @@
 //
 // This source file is part of the Swift.org open source project
 //
-// Copyright (c) 2014 - 2019 Apple Inc. and the Swift project authors
+// Copyright (c) 2014 - 2020 Apple Inc. and the Swift project authors
 // Licensed under Apache License v2.0 with Runtime Library Exception
 //
 // See https://swift.org/LICENSE.txt for license information
@@ -21,6 +21,7 @@ public protocol ToolchainLanguageServer: AnyObject {
 
   func initializeSync(_ initialize: InitializeRequest) throws -> InitializeResult
   func clientInitialized(_ initialized: InitializedNotification)
+  func shutdown()
 
   // MARK: - Text synchronization