Make SourceKitDRegistry an actor

Author: ahoppen

Sources/Diagnose/SourcekitdRequestCommand.swift

@@ -44,7 +44,7 @@ public struct SourceKitdRequestCommand: AsyncParsableCommand {
   public func run() async throws {
     var requestString = try String(contentsOf: URL(fileURLWithPath: sourcekitdRequestPath))
 
-    let sourcekitd = try DynamicallyLoadedSourceKitD.getOrCreate(
+    let sourcekitd = try await DynamicallyLoadedSourceKitD.getOrCreate(
       dylibPath: try! AbsolutePath(validating: sourcekitdPath)
     )
 

Sources/SourceKitD/DynamicallyLoadedSourceKitD.swift

@@ -47,8 +47,8 @@ public final class DynamicallyLoadedSourceKitD: SourceKitD {
   /// List of notification handlers that will be called for each notification.
   private var _notificationHandlers: [WeakSKDNotificationHandler] = []
 
-  public static func getOrCreate(dylibPath: AbsolutePath) throws -> SourceKitD {
-    try SourceKitDRegistry.shared
+  public static func getOrCreate(dylibPath: AbsolutePath) async throws -> SourceKitD {
+    try await SourceKitDRegistry.shared
       .getOrAdd(dylibPath, create: { try DynamicallyLoadedSourceKitD(dylib: dylibPath) })
   }
 

Sources/SourceKitD/SourceKitDRegistry.swift

@@ -32,10 +32,7 @@ extension NSLock {
 /// * 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: NSLock = NSLock()
+public actor SourceKitDRegistry {
 
   /// Mapping from path to active SourceKitD instance.
   var active: [AbsolutePath: SourceKitD] = [:]
@@ -47,26 +44,24 @@ public final class SourceKitDRegistry {
   public init() {}
 
   /// The global shared SourceKitD registry.
-  public static var shared: SourceKitDRegistry = SourceKitDRegistry()
+  public static let 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
+    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
@@ -77,22 +72,18 @@ public final class SourceKitDRegistry {
   /// 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]?.value == nil)
-        cemetary[key] = WeakSourceKitD(value: existing)
-      }
-      return existing
+    let existing = active.removeValue(forKey: key)
+    if let existing = existing {
+      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()
-    }
+    active.removeAll()
+    cemetary.removeAll()
   }
 }
 

Sources/SourceKitLSP/Swift/SwiftLanguageServer.swift

@@ -179,7 +179,7 @@ public actor SwiftLanguageServer: ToolchainLanguageServer {
     guard let sourcekitd = toolchain.sourcekitd else { return nil }
     self.sourceKitServer = sourceKitServer
     self.swiftFormat = toolchain.swiftFormat
-    self.sourcekitd = try DynamicallyLoadedSourceKitD.getOrCreate(dylibPath: sourcekitd)
+    self.sourcekitd = try await DynamicallyLoadedSourceKitD.getOrCreate(dylibPath: sourcekitd)
     self.capabilityRegistry = workspace.capabilityRegistry
     self.serverOptions = options
     self.documentManager = DocumentManager()

Tests/DiagnoseTests/DiagnoseTests.swift

@@ -211,7 +211,7 @@ private struct InProcessSourceKitRequestExecutor: SourceKitRequestExecutor {
   func run(request requestYaml: String) async throws -> SourceKitDRequestResult {
     logger.info("Sending request: \(requestYaml)")
 
-    let sourcekitd = try DynamicallyLoadedSourceKitD.getOrCreate(
+    let sourcekitd = try await DynamicallyLoadedSourceKitD.getOrCreate(
       dylibPath: try! AbsolutePath(validating: sourcekitd.path)
     )
     let response = try await sourcekitd.run(requestYaml: requestYaml)

Tests/SourceKitDTests/SourceKitDRegistryTests.swift

@@ -10,49 +10,50 @@
 //
 //===----------------------------------------------------------------------===//
 
+import LSPTestSupport
 import SourceKitD
 import TSCBasic
 import XCTest
 
 final class SourceKitDRegistryTests: XCTestCase {
 
-  func testAdd() throws {
+  func testAdd() async throws {
     let registry = SourceKitDRegistry()
 
-    let a = try FakeSourceKitD.getOrCreate(AbsolutePath(validating: "/a"), in: registry)
-    let b = try FakeSourceKitD.getOrCreate(AbsolutePath(validating: "/b"), in: registry)
-    let a2 = try FakeSourceKitD.getOrCreate(AbsolutePath(validating: "/a"), in: registry)
+    let a = try await FakeSourceKitD.getOrCreate(AbsolutePath(validating: "/a"), in: registry)
+    let b = try await FakeSourceKitD.getOrCreate(AbsolutePath(validating: "/b"), in: registry)
+    let a2 = try await FakeSourceKitD.getOrCreate(AbsolutePath(validating: "/a"), in: registry)
 
     XCTAssert(a === a2)
     XCTAssert(a !== b)
   }
 
-  func testRemove() throws {
+  func testRemove() async throws {
     let registry = SourceKitDRegistry()
 
-    let a = FakeSourceKitD.getOrCreate(try AbsolutePath(validating: "/a"), in: registry)
-    XCTAssert(registry.remove(try AbsolutePath(validating: "/a")) === a)
-    XCTAssertNil(registry.remove(try AbsolutePath(validating: "/a")))
+    let a = await FakeSourceKitD.getOrCreate(try AbsolutePath(validating: "/a"), in: registry)
+    await assertTrue(registry.remove(try AbsolutePath(validating: "/a")) === a)
+    await assertNil(registry.remove(try AbsolutePath(validating: "/a")))
   }
 
-  func testRemoveResurrect() throws {
+  func testRemoveResurrect() async throws {
     let registry = SourceKitDRegistry()
 
     @inline(never)
-    func scope(registry: SourceKitDRegistry) throws -> Int {
-      let a = FakeSourceKitD.getOrCreate(try AbsolutePath(validating: "/a"), in: registry)
+    func scope(registry: SourceKitDRegistry) async throws -> Int {
+      let a = await FakeSourceKitD.getOrCreate(try AbsolutePath(validating: "/a"), in: registry)
 
-      XCTAssert(a === FakeSourceKitD.getOrCreate(try AbsolutePath(validating: "/a"), in: registry))
-      XCTAssert(registry.remove(try AbsolutePath(validating: "/a")) === a)
+      await assertTrue(a === FakeSourceKitD.getOrCreate(try AbsolutePath(validating: "/a"), in: registry))
+      await assertTrue(registry.remove(try AbsolutePath(validating: "/a")) === a)
       // Resurrected.
-      XCTAssert(a === FakeSourceKitD.getOrCreate(try AbsolutePath(validating: "/a"), in: registry))
+      await assertTrue(a === FakeSourceKitD.getOrCreate(try AbsolutePath(validating: "/a"), in: registry))
       // Remove again.
-      XCTAssert(registry.remove(try AbsolutePath(validating: "/a")) === a)
+      await assertTrue(registry.remove(try AbsolutePath(validating: "/a")) === a)
       return (a as! FakeSourceKitD).token
     }
 
-    let id = try scope(registry: registry)
-    let a2 = FakeSourceKitD.getOrCreate(try AbsolutePath(validating: "/a"), in: registry)
+    let id = try await scope(registry: registry)
+    let a2 = await FakeSourceKitD.getOrCreate(try AbsolutePath(validating: "/a"), in: registry)
     XCTAssertNotEqual(id, (a2 as! FakeSourceKitD).token)
   }
 }
@@ -72,8 +73,8 @@ final class FakeSourceKitD: SourceKitD {
     nextToken += 1
   }
 
-  static func getOrCreate(_ path: AbsolutePath, in registry: SourceKitDRegistry) -> SourceKitD {
-    return registry.getOrAdd(path, create: { Self.init() })
+  static func getOrCreate(_ path: AbsolutePath, in registry: SourceKitDRegistry) async -> SourceKitD {
+    return await registry.getOrAdd(path, create: { Self.init() })
   }
 
   public func log(request: SKDRequestDictionary) {}

Tests/SourceKitDTests/SourceKitDTests.swift

@@ -28,7 +28,7 @@ import class TSCBasic.Process
 final class SourceKitDTests: XCTestCase {
   func testMultipleNotificationHandlers() async throws {
     let sourcekitdPath = await ToolchainRegistry.forTesting.default!.sourcekitd!
-    let sourcekitd = try DynamicallyLoadedSourceKitD.getOrCreate(dylibPath: sourcekitdPath)
+    let sourcekitd = try await DynamicallyLoadedSourceKitD.getOrCreate(dylibPath: sourcekitdPath)
     let keys = sourcekitd.keys
     let path = DocumentURI.for(.swift).pseudoPath