From f6d77010482c2cf3fcaaed64a131b5ee9671a2be Mon Sep 17 00:00:00 2001 From: Ben Langmuir Date: Tue, 2 Jun 2020 13:25:38 -0700 Subject: [PATCH] [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. --- Package.swift | 6 ++ .../SourceKit/Swift/SwiftLanguageServer.swift | 2 +- Sources/SourceKitD/SourceKitDImpl.swift | 7 +- Sources/SourceKitD/SourceKitDRegistry.swift | 82 +++++++++++++++++++ Sources/SourceKitD/sourcekitd_functions.swift | 2 +- .../SourceKitDRegistryTests.swift | 78 ++++++++++++++++++ 6 files changed, 174 insertions(+), 3 deletions(-) create mode 100644 Sources/SourceKitD/SourceKitDRegistry.swift create mode 100644 Tests/SourceKitDTests/SourceKitDRegistryTests.swift diff --git a/Package.swift b/Package.swift index 237927bc..c2567264 100644 --- a/Package.swift +++ b/Package.swift @@ -116,6 +116,12 @@ let package = Package( "SwiftToolsSupport-auto", ] ), + .testTarget( + name: "SourceKitDTests", + dependencies: [ + "SourceKitD", + ] + ), // Csourcekitd: C modules wrapper for sourcekitd. .target( diff --git a/Sources/SourceKit/Swift/SwiftLanguageServer.swift b/Sources/SourceKit/Swift/SwiftLanguageServer.swift index 93454023..66d87cd1 100644 --- a/Sources/SourceKit/Swift/SwiftLanguageServer.swift +++ b/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() diff --git a/Sources/SourceKitD/SourceKitDImpl.swift b/Sources/SourceKitD/SourceKitDImpl.swift index b8979f31..2a05b3cd 100644 --- a/Sources/SourceKitD/SourceKitDImpl.swift +++ b/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: []) diff --git a/Sources/SourceKitD/SourceKitDRegistry.swift b/Sources/SourceKitD/SourceKitDRegistry.swift new file mode 100644 index 00000000..0104a4b5 --- /dev/null +++ b/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? +} diff --git a/Sources/SourceKitD/sourcekitd_functions.swift b/Sources/SourceKitD/sourcekitd_functions.swift index bd172349..668241f2 100644 --- a/Sources/SourceKitD/sourcekitd_functions.swift +++ b/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() diff --git a/Tests/SourceKitDTests/SourceKitDRegistryTests.swift b/Tests/SourceKitDTests/SourceKitDRegistryTests.swift new file mode 100644 index 00000000..9507afed --- /dev/null +++ b/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() }) + } +}