SemanticIndex | Add deterministic IndexStoreDB cleanup on shutdown

Package.swift

@@ -271,6 +271,8 @@ var targets: [Target] = [
     dependencies: [
       "SemanticIndex",
       "SKTestSupport",
+      "ToolchainRegistry",
+      .product(name: "IndexStoreDB", package: "indexstore-db"),
       .product(name: "SKLogging", package: "swift-tools-protocols"),
       .product(name: "ToolsProtocolsSwiftExtensions", package: "swift-tools-protocols"),
     ],

Sources/BuildServerIntegration/BuildServerManager.swift

@@ -641,6 +641,10 @@ package actor BuildServerManager: QueueBasedMessageHandler {
   /// which could result in the connection being reported as a leak. To avoid this problem, we want to explicitly shut
   /// down the build server when the `SourceKitLSPServer` gets shut down.
   package func shutdown() async {
+    // Close the index store before shutting down the build server so it is
+    // released deterministically rather than waiting for deallocation.
+    await self.mainFilesProvider?.value?.close()
+
     // Clear any pending work done progresses from the build server.
     self.workDoneProgressManagers.removeAll()
     guard let buildServerAdapter = try? await self.buildServerAdapterAfterInitialized else {

Sources/BuildServerIntegration/MainFilesProvider.swift

@@ -25,4 +25,11 @@ package protocol MainFilesProvider: Sendable {
   /// mainFilesContainingFile("foo.h") == Set(["foo.cpp", "bar.cpp"])
   /// ```
   func mainFiles(containing uri: DocumentURI, crossLanguage: Bool) async -> Set<DocumentURI>
+
+  /// Close and release any underlying resources (e.g. IndexStoreDB).
+  func close() async
+}
+
+extension MainFilesProvider {
+  package func close() async {}
 }

Sources/SemanticIndex/CheckedIndex.swift

@@ -310,7 +310,26 @@ package final class CheckedIndex {
 /// calling `underlyingIndexStoreDB`) and we don't accidentally call into the `IndexStoreDB` when we wanted a
 /// `CheckedIndex`.
 package final actor UncheckedIndex: Sendable {
-  package nonisolated let underlyingIndexStoreDB: IndexStoreDB
+  /// Wrapper to store `IndexStoreDB` in `ThreadSafeBox` so `UncheckedIndex` can explicitly set it to `nil` and
+  /// deterministically trigger deinitialization of the underlying `IndexStoreDB` instance.
+  ///
+  /// `IndexStoreDB` is not `Sendable`, but access to this value is synchronized by the lock in
+  /// `ThreadSafeBox`, so marking this wrapper as `@unchecked Sendable` is safe.
+  private struct UncheckedSendableOptionalIndexStoreDB: @unchecked Sendable {
+    var indexStoreDB: IndexStoreDB?
+  }
+
+  private let indexStoreDB: ThreadSafeBox<UncheckedSendableOptionalIndexStoreDB>
+
+  package nonisolated var underlyingIndexStoreDB: IndexStoreDB {
+    let optionalIndexStoreDB = indexStoreDB.withLock { indexStoreDB in
+      indexStoreDB
+    }
+    guard let indexStoreDB = optionalIndexStoreDB.indexStoreDB else {
+      preconditionFailure("Tried to access an IndexStoreDB that was already closed")
+    }
+    return indexStoreDB
+  }
 
   /// Whether the underlying `IndexStoreDB` uses has `useExplicitOutputUnits` enabled and thus needs to receive updates
   /// updates as output paths are added or removed from the project.
@@ -324,7 +343,7 @@ package final actor UncheckedIndex: Sendable {
       return nil
     }
     self.usesExplicitOutputPaths = usesExplicitOutputPaths
-    self.underlyingIndexStoreDB = index
+    self.indexStoreDB = ThreadSafeBox(initialValue: UncheckedSendableOptionalIndexStoreDB(indexStoreDB: index))
   }
 
   /// Update the set of output paths that should be considered visible in the project. For example, if a source file is
@@ -353,6 +372,13 @@ package final actor UncheckedIndex: Sendable {
   package nonisolated func processUnitsForOutputPathsAndWait(_ outputPaths: some Collection<String>) {
     self.underlyingIndexStoreDB.processUnitsForOutputPathsAndWait(outputPaths)
   }
+
+  /// Explicitly close the underlying `IndexStoreDB` instance and release it immediately.
+  package func close() {
+    indexStoreDB.withLock { indexStoreDB in
+      indexStoreDB.indexStoreDB = nil
+    }
+  }
 }
 
 /// Helper class to check if symbols from the index are up-to-date or if the source file has been modified after it was

Tests/SemanticIndexTests/TaskSchedulerTests.swift

@@ -10,10 +10,12 @@
 //
 //===----------------------------------------------------------------------===//
 
+@preconcurrency import IndexStoreDB
 @_spi(SourceKitLSP) import SKLogging
 import SKTestSupport
 import SemanticIndex
 import SwiftExtensions
+import ToolchainRegistry
 @_spi(SourceKitLSP) import ToolsProtocolsSwiftExtensions
 import XCTest
 
@@ -312,6 +314,39 @@ final class TaskSchedulerTests: SourceKitLSPTestCase {
     lowPriorityTaskFinished.waitOrXCTFail()
   }
 
+  func testIndexStoreDBClosing() async throws {
+    guard let libIndexStore = await ToolchainRegistry.forTesting.default?.libIndexStore else {
+      throw XCTSkip("libIndexStore not available")
+    }
+
+    try await withTestScratchDir { tempDir in
+      let storePath = tempDir.appending(component: "store")
+      let dbPath = tempDir.appending(component: "db")
+      try FileManager.default.createDirectory(at: storePath, withIntermediateDirectories: true)
+      try FileManager.default.createDirectory(at: dbPath, withIntermediateDirectories: true)
+
+      weak var weakDB: IndexStoreDB?
+      let uncheckedIndex: UncheckedIndex
+
+      do {
+        let db = try IndexStoreDB(
+          storePath: storePath.filePath,
+          databasePath: dbPath.filePath,
+          library: IndexStoreLibrary(dylibPath: libIndexStore.filePath),
+          listenToUnitEvents: false
+        )
+        weakDB = db
+        uncheckedIndex = UncheckedIndex(db, usesExplicitOutputPaths: false)!
+      }
+
+      XCTAssertNotNil(weakDB)
+
+      await uncheckedIndex.close()
+
+      XCTAssertNil(weakDB)
+    }
+  }
+
   func testScheduleTask() {
     XCTAssertFalse(
       TaskScheduler<ClosureTaskDescription>.canScheduleTask(