[sourcekitd] Add a registry for sourcekitd instances

Protect ourselves from ever having multiple sourcekitd instances for the
same path live at once, which is not safe.

Author: benlangmuir

Package.swift

@@ -116,6 +116,12 @@ let package = Package(
           "SwiftToolsSupport-auto",
         ]
       ),
+      .testTarget(
+        name: "SourceKitDTests",
+        dependencies: [
+          "SourceKitD",
+        ]
+      ),
 
       // Csourcekitd: C modules wrapper for sourcekitd.
       .target(

Sources/SourceKit/Swift/SwiftLanguageServer.swift

@@ -105,7 +105,7 @@ public final class SwiftLanguageServer: ToolchainLanguageServer {
   /// 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 {
     self.client = client
-    self.sourcekitd = try SourceKitDImpl(dylib: sourcekitd)
+    self.sourcekitd = try SourceKitDImpl.getOrCreate(dylibPath: sourcekitd)
     self.buildSystem = buildSystem
     self.clientCapabilities = clientCapabilities
     self.documentManager = DocumentManager()

Sources/SourceKitD/SourceKitDImpl.swift

@@ -51,7 +51,12 @@ public final class SourceKitDImpl: SourceKitD {
     }
   }
 
-  public init(dylib path: AbsolutePath) throws {
+  public static func getOrCreate(dylibPath: AbsolutePath) throws -> SourceKitD {
+    try SourceKitDRegistry.shared
+      .getOrAdd(dylibPath, create: { try SourceKitDImpl(dylib: dylibPath) })
+  }
+
+  init(dylib path: AbsolutePath) throws {
     self.path = path
     #if os(Windows)
     self.dylib = try dlopen(path.pathString, mode: [])

Sources/SourceKitD/SourceKitDRegistry.swift

@@ -0,0 +1,82 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// 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
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+import TSCBasic
+
+/// The set of known SourceKitD instances, uniqued by path.
+///
+/// It is not generally safe to have two instances of SourceKitD for the same libsourcekitd, so
+/// care is taken to ensure that there is only ever one instance per path.
+///
+/// * To get a new instance, use `getOrAdd("path", create: { NewSourceKitD() })`.
+/// * To remove an existing instance, use `remove("path")`, but be aware that if there are any other
+///   references to the instances in the program, it can be resurrected if `getOrAdd` is called with
+///   the same path. See note on `remove(_:)`
+public final class SourceKitDRegistry {
+
+  /// Mutex protecting mutable state in the registry.
+  let lock: Lock = Lock()
+
+  /// Mapping from path to active SourceKitD instance.
+  var active: [AbsolutePath: SourceKitD] = [:]
+
+  /// Instances that have been unregistered, but may be resurrected if accessed before destruction.
+  var cemetary: [AbsolutePath: WeakSourceKitD] = [:]
+
+  /// Initialize an empty registry.
+  public init() {}
+
+  /// The global shared SourceKitD registry.
+  public static var shared: SourceKitDRegistry = SourceKitDRegistry()
+
+  /// Returns the existing SourceKitD for the given path, or creates it and registers it.
+  public func getOrAdd(
+    _ key: AbsolutePath,
+    create: () throws -> SourceKitD
+  ) rethrows -> SourceKitD {
+    try lock.withLock {
+      if let existing = active[key] {
+        return existing
+      }
+      if let resurrected = cemetary[key]?.value {
+        cemetary[key] = nil
+        active[key] = resurrected
+        return resurrected
+      }
+      let newValue = try create()
+      active[key] = newValue
+      return newValue
+    }
+  }
+
+  /// Removes the SourceKitD instance registered for the given path, if any, from the set of active
+  /// instances.
+  ///
+  /// Since it is not generally safe to have two sourcekitd connections at once, the existing value
+  /// is converted to a weak reference until it is no longer referenced anywhere by the program. If
+  /// the same path is looked up again before the original service is deinitialized, the original
+  /// service is resurrected rather than creating a new instance.
+  public func remove(_ key: AbsolutePath) -> SourceKitD? {
+    lock.withLock {
+      let existing = active.removeValue(forKey: key)
+      if let existing = existing {
+        assert(self.cemetary[key] == nil)
+        cemetary[key] = WeakSourceKitD(value: existing)
+      }
+      return existing
+    }
+  }
+}
+
+struct WeakSourceKitD {
+  weak var value: SourceKitD?
+}

Sources/SourceKitD/sourcekitd_functions.swift

@@ -14,7 +14,7 @@ import Csourcekitd
 import SKSupport
 
 extension sourcekitd_functions_t {
-  init(_ sourcekitd: DLHandle) throws {
+  public init(_ sourcekitd: DLHandle) throws {
     // Zero-initialize
     self.init()
 

Tests/SourceKitDTests/SourceKitDRegistryTests.swift

@@ -0,0 +1,78 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// 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
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+import SourceKitD
+import TSCBasic
+import XCTest
+
+final class SourceKitDRegistryTests: XCTestCase {
+
+  func testAdd() {
+    let registry = SourceKitDRegistry()
+
+    let a = FakeSourceKitD.getOrCreate(AbsolutePath("/a"), in: registry)
+    let b = FakeSourceKitD.getOrCreate(AbsolutePath("/b"), in: registry)
+    let a2 = FakeSourceKitD.getOrCreate(AbsolutePath("/a"), in: registry)
+
+    XCTAssert(a === a2)
+    XCTAssert(a !== b)
+  }
+
+  func testRemove() {
+    let registry = SourceKitDRegistry()
+
+    let a = FakeSourceKitD.getOrCreate(AbsolutePath("/a"), in: registry)
+    XCTAssert(registry.remove(AbsolutePath("/a")) === a)
+    XCTAssertNil(registry.remove(AbsolutePath("/a")))
+  }
+
+  func testRemoveResurrect() {
+    let registry = SourceKitDRegistry()
+
+    @inline(never)
+    func scope(registry: SourceKitDRegistry) -> Int {
+      let a = FakeSourceKitD.getOrCreate(AbsolutePath("/a"), in: registry)
+
+      XCTAssert(a === FakeSourceKitD.getOrCreate(AbsolutePath("/a"), in: registry))
+      XCTAssert(registry.remove(AbsolutePath("/a")) === a)
+      // Resurrected.
+      XCTAssert(a === FakeSourceKitD.getOrCreate(AbsolutePath("/a"), in: registry))
+      // Remove again.
+      XCTAssert(registry.remove(AbsolutePath("/a")) === a)
+      return (a as! FakeSourceKitD).token
+    }
+
+    let id = scope(registry: registry)
+    let a2 = FakeSourceKitD.getOrCreate(AbsolutePath("/a"), in: registry)
+    XCTAssertNotEqual(id, (a2 as! FakeSourceKitD).token)
+  }
+}
+
+private var nextToken = 0
+
+final class FakeSourceKitD: SourceKitD {
+  let token: Int
+  var api: sourcekitd_functions_t { fatalError() }
+  var keys: sourcekitd_keys { fatalError() }
+  var requests: sourcekitd_requests { fatalError() }
+  var values: sourcekitd_values { fatalError() }
+  func addNotificationHandler(_ handler: SKDNotificationHandler) { fatalError() }
+  func removeNotificationHandler(_ handler: SKDNotificationHandler) { fatalError() }
+  private init() {
+    token = nextToken
+    nextToken += 1
+  }
+
+  static func getOrCreate(_ path: AbsolutePath, in registry: SourceKitDRegistry) -> SourceKitD {
+    return registry.getOrAdd(path, create: { Self.init() })
+  }
+}