Fix a number of resource leaks when using SourceKit-LSP in tests or as a library

Sources/LSPTestSupport/TestJSONRPCConnection.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
@@ -18,7 +18,7 @@ import XCTest
 // Workaround ambiguity with Foundation.
 public typealias Notification = LanguageServerProtocol.Notification
 
-public struct TestJSONRPCConnection {
+public final class TestJSONRPCConnection {
   public let clientToServer: Pipe = Pipe()
   public let serverToClient: Pipe = Pipe()
   public let client: TestClient
@@ -27,11 +27,6 @@ public struct TestJSONRPCConnection {
   public let serverConnection: JSONRPCConnection
 
   public init() {
-    // 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 = JSONRPCConnection(
       protocol: testMessageRegistry,
       inFD: serverToClient.fileHandleForReading.fileDescriptor,
@@ -47,8 +42,16 @@ public struct TestJSONRPCConnection {
     client = TestClient(server: clientConnection)
     server = TestServer(client: serverConnection)
 
-    clientConnection.start(receiveHandler: client)
-    serverConnection.start(receiveHandler: server)
+    clientConnection.start(receiveHandler: client) {
+      // FIXME: keep the pipes alive until we close the connection. This
+      // should be fixed systemically.
+      withExtendedLifetime(self) {}
+    }
+    serverConnection.start(receiveHandler: server) {
+      // FIXME: keep the pipes alive until we close the connection. This
+      // should be fixed systemically.
+      withExtendedLifetime(self) {}
+    }
   }
 
   public func close() {

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/SKCore/BuildServerBuildSystem.swift

@@ -254,7 +254,11 @@ private func makeJSONRPCBuildServer(client: MessageHandler, serverPath: Absolute
     outFD: clientToServer.fileHandleForWriting.fileDescriptor
   )
 
-  connection.start(receiveHandler: client)
+  connection.start(receiveHandler: client) {
+    // FIXME: keep the pipes alive until we close the connection. This
+    // should be fixed systemically.
+    withExtendedLifetime((clientToServer, serverToClient)) {}
+  }
   let process = Foundation.Process()
 
   if #available(OSX 10.13, *) {

Sources/SKCore/BuildSystemManager.swift

@@ -110,7 +110,7 @@ public final class BuildSystemManager {
   weak var mainFilesProvider: MainFilesProvider?
 
   /// Build system delegate that will receive notifications about setting changes, etc.
-  var _delegate: BuildSystemDelegate?
+  weak var _delegate: BuildSystemDelegate?
 
   /// Create a BuildSystemManager that wraps the given build system. The new
   /// manager will modify the delegate of the underlying build system.

Sources/SKTestSupport/TestServer.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
@@ -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?
 
@@ -63,11 +65,6 @@ public struct TestSourceKitServer {
         let clientToServer: Pipe = Pipe()
         let serverToClient: Pipe = Pipe()
 
-        // 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)
-
         let clientConnection = JSONRPCConnection(
           protocol: MessageRegistry.lspProtocol,
           inFD: serverToClient.fileHandleForReading.fileDescriptor,
@@ -84,8 +81,16 @@ public struct TestSourceKitServer {
           serverConnection.close()
         })
 
-        clientConnection.start(receiveHandler: client)
-        serverConnection.start(receiveHandler: server!)
+        clientConnection.start(receiveHandler: client) {
+          // FIXME: keep the pipes alive until we close the connection. This
+          // should be fixed systemically.
+          withExtendedLifetime((clientToServer, serverToClient)) {}
+        }
+        serverConnection.start(receiveHandler: server!) {
+          // FIXME: keep the pipes alive until we close the connection. This
+          // should be fixed systemically.
+          withExtendedLifetime((clientToServer, serverToClient)) {}
+        }
 
         connImpl = .jsonrpc(
           clientToServer: clientToServer,
@@ -95,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/SourceKitD/SourceKitDRegistry.swift

@@ -69,12 +69,20 @@ public final class SourceKitDRegistry {
     lock.withLock {
       let existing = active.removeValue(forKey: key)
       if let existing = existing {
-        assert(self.cemetary[key] == nil)
+        assert(self.cemetary[key]?.value == nil)
         cemetary[key] = WeakSourceKitD(value: existing)
       }
       return existing
     }
   }
+
+  /// Remove all SourceKitD instances, including weak ones.
+  public func clear() {
+    lock.withLock {
+      active.removeAll()
+      cemetary.removeAll()
+    }
+  }
 }
 
 struct WeakSourceKitD {

Sources/SourceKitLSP/Clang/ClangLanguageServer.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
@@ -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) {
@@ -247,7 +261,11 @@ func makeJSONRPCClangServer(
     clang: toolchain.clang)
 
   connectionToClient.start(handler: client)
-  connection.start(receiveHandler: client)
+  connection.start(receiveHandler: client) {
+    // FIXME: keep the pipes alive until we close the connection. This
+    // should be fixed systemically.
+    withExtendedLifetime((clientToServer, serverToClient)) {}
+  }
 
   let process = Foundation.Process()
 

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
 

Tests/LanguageServerProtocolJSONRPCTests/ConnectionTests.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
@@ -237,6 +237,10 @@ class ConnectionTests: XCTestCase {
       conn.start(receiveHandler: DummyHandler(), closeHandler: {
         // We get an error from XCTest if this is fulfilled more than once.
         expectation.fulfill()
+
+        // FIXME: keep the pipes alive until we close the connection. This
+        // should be fixed systemically.
+        withExtendedLifetime((to, from)) {}
       })
 
       to.fileHandleForWriting.closeFile()