#1597 - Migrate DocumentManager to use ThreadSafeBox

Sources/SourceKitD/SourceKitD.swift

@@ -23,8 +23,8 @@ extension sourcekitd_api_values: @unchecked Sendable {}
 fileprivate extension ThreadSafeBox {
   /// If the wrapped value is `nil`, run `compute` and store the computed value. If it is not `nil`, return the stored
   /// value.
-  func computeIfNil<WrappedValue>(compute: () -> WrappedValue) -> WrappedValue where T == WrappedValue? {
-    return withLock { value in
+  func computeIfNil<WrappedValue: Sendable>(compute: () -> WrappedValue) -> WrappedValue where T == WrappedValue? {
+    return withLock { (value: inout WrappedValue?) -> WrappedValue in
       if let value {
         return value
       }

Sources/SourceKitLSP/DocumentManager.swift

@@ -10,12 +10,12 @@
 //
 //===----------------------------------------------------------------------===//
 
-import Dispatch
 import Foundation
 @_spi(SourceKitLSP) package import LanguageServerProtocol
 @_spi(SourceKitLSP) import SKLogging
 package import SKUtilities
 import SemanticIndex
+import SwiftExtensions
 package import SwiftSyntax
 
 /// An immutable snapshot of a document at a given time.
@@ -92,18 +92,15 @@ package final class DocumentManager: InMemoryDocumentManager, Sendable {
     case missingDocument(DocumentURI)
   }
 
-  // TODO: Migrate this to be an AsyncQueue (https://github.com/swiftlang/sourcekit-lsp/issues/1597)
-  private let queue: DispatchQueue = DispatchQueue(label: "document-manager-queue")
-
-  // `nonisolated(unsafe)` is fine because `documents` is guarded by queue.
-  private nonisolated(unsafe) var documents: [DocumentURI: Document] = [:]
+  // Documents storage, protected by a `ThreadSafeBox` to ensure thread safety without making APIs async.
+  private let documents: ThreadSafeBox<[DocumentURI: Document]> = ThreadSafeBox(initialValue: [:])
 
   package init() {}
 
   /// All currently opened documents.
   package var openDocuments: Set<DocumentURI> {
-    return queue.sync {
-      return Set(documents.keys)
+    return documents.withLock { docs in
+      return Set(docs.keys)
     }
   }
 
@@ -113,9 +110,9 @@ package final class DocumentManager: InMemoryDocumentManager, Sendable {
   /// - throws: Error.alreadyOpen if the document is already open.
   @discardableResult
   package func open(_ uri: DocumentURI, language: Language, version: Int, text: String) throws -> DocumentSnapshot {
-    return try queue.sync {
+    return try documents.withLock { docs in
       let document = Document(uri: uri, language: language, version: version, text: text)
-      if nil != documents.updateValue(document, forKey: uri) {
+      if nil != docs.updateValue(document, forKey: uri) {
         throw Error.alreadyOpen(uri)
       }
       return document.latestSnapshot
@@ -127,8 +124,8 @@ package final class DocumentManager: InMemoryDocumentManager, Sendable {
   /// - returns: The initial contents of the file.
   /// - throws: Error.missingDocument if the document is not open.
   package func close(_ uri: DocumentURI) throws {
-    try queue.sync {
-      if nil == documents.removeValue(forKey: uri) {
+    try documents.withLock { docs in
+      if nil == docs.removeValue(forKey: uri) {
         throw Error.missingDocument(uri)
       }
     }
@@ -151,8 +148,8 @@ package final class DocumentManager: InMemoryDocumentManager, Sendable {
     newVersion: Int,
     edits: [TextDocumentContentChangeEvent]
   ) throws -> (preEditSnapshot: DocumentSnapshot, postEditSnapshot: DocumentSnapshot, edits: [SourceEdit]) {
-    return try queue.sync {
-      guard let document = documents[uri] else {
+    return try documents.withLock { docs in
+      guard let document = docs[uri] else {
         throw Error.missingDocument(uri)
       }
       let preEditSnapshot = document.latestSnapshot
@@ -184,8 +181,8 @@ package final class DocumentManager: InMemoryDocumentManager, Sendable {
   }
 
   package func latestSnapshot(_ uri: DocumentURI) throws -> DocumentSnapshot {
-    return try queue.sync {
-      guard let document = documents[uri] else {
+    return try documents.withLock { docs in
+      guard let document = docs[uri] else {
         throw ResponseError.unknown("Failed to find snapshot for '\(uri)'")
       }
       return document.latestSnapshot

Sources/SwiftExtensions/ThreadSafeBox.swift

@@ -15,12 +15,26 @@ import Foundation
 /// A thread safe container that contains a value of type `T`.
 ///
 /// - Note: Unchecked sendable conformance because value is guarded by a lock.
-package class ThreadSafeBox<T: Sendable>: @unchecked Sendable {
+package final class ThreadSafeBox<T>: Sendable {
   /// Lock guarding `_value`.
   private let lock = NSLock()
 
-  private var _value: T
+  private nonisolated(unsafe) var _value: T
 
+  package init(initialValue: T) {
+    _value = initialValue
+  }
+
+  /// Restrict the result of the body to `Sendable` so callers can safely use
+  /// the returned value outside the lock without requiring `T: Sendable`.
+  package func withLock<Result: Sendable>(_ body: (inout T) throws -> Result) rethrows -> Result {
+    return try lock.withLock {
+      return try body(&_value)
+    }
+  }
+}
+
+extension ThreadSafeBox where T: Sendable {
   package var value: T {
     get {
       return lock.withLock {
@@ -39,16 +53,6 @@ package class ThreadSafeBox<T: Sendable>: @unchecked Sendable {
     }
   }
 
-  package init(initialValue: T) {
-    _value = initialValue
-  }
-
-  package func withLock<Result>(_ body: (inout T) throws -> Result) rethrows -> Result {
-    return try lock.withLock {
-      return try body(&_value)
-    }
-  }
-
   /// If the value in the box is an optional, return it and reset it to `nil`
   /// in an atomic operation.
   package func takeValue<U>() -> T where U? == T {