Explicitly close the index when shutting down SourceKit-LSP

Sources/DocumentationLanguageService/IndexStoreDB+Extensions.swift

@@ -32,7 +32,7 @@ extension CheckedIndex {
     }
     // Find all occurrences of the symbol by name alone
     var topLevelSymbolOccurrences: [SymbolOccurrence] = []
-    forEachCanonicalSymbolOccurrence(byName: topLevelSymbolName) { symbolOccurrence in
+    try forEachCanonicalSymbolOccurrence(byName: topLevelSymbolName) { symbolOccurrence in
       topLevelSymbolOccurrences.append(symbolOccurrence)
       return true  // continue
     }
@@ -61,14 +61,14 @@ extension CheckedIndex {
     ofUSR usr: String,
     fetchSymbolGraph: (SymbolLocation) async throws -> String?
   ) async throws -> DocCSymbolInformation {
-    guard let topLevelSymbolOccurrence = primaryDefinitionOrDeclarationOccurrence(ofUSR: usr) else {
+    guard let topLevelSymbolOccurrence = try primaryDefinitionOrDeclarationOccurrence(ofUSR: usr) else {
       throw DocCCheckedIndexError.emptyDocCSymbolLink
     }
     let moduleName = topLevelSymbolOccurrence.location.moduleName
     var symbols = [topLevelSymbolOccurrence]
     // Find any parent symbols
     var symbolOccurrence: SymbolOccurrence = topLevelSymbolOccurrence
-    while let parentSymbolOccurrence = symbolOccurrence.parent(self) {
+    while let parentSymbolOccurrence = try symbolOccurrence.parent(self) {
       symbols.insert(parentSymbolOccurrence, at: 0)
       symbolOccurrence = parentSymbolOccurrence
     }
@@ -106,7 +106,7 @@ enum DocCCheckedIndexError: LocalizedError {
 }
 
 extension SymbolOccurrence {
-  func parent(_ index: CheckedIndex) -> SymbolOccurrence? {
+  func parent(_ index: CheckedIndex) throws -> SymbolOccurrence? {
     let allParentRelations =
       relations
       .filter { $0.roles.contains(.childOf) }
@@ -118,13 +118,13 @@ extension SymbolOccurrence {
       return nil
     }
     if parentRelation.symbol.kind == .extension {
-      let allSymbolOccurrences = index.occurrences(relatedToUSR: parentRelation.symbol.usr, roles: .extendedBy)
+      let allSymbolOccurrences = try index.occurrences(relatedToUSR: parentRelation.symbol.usr, roles: .extendedBy)
         .sorted()
       if allSymbolOccurrences.count > 1 {
         logger.debug("Extension \(parentRelation.symbol.usr) extends multiple symbols")
       }
       return allSymbolOccurrences.first
     }
-    return index.primaryDefinitionOrDeclarationOccurrence(ofUSR: parentRelation.symbol.usr)
+    return try index.primaryDefinitionOrDeclarationOccurrence(ofUSR: parentRelation.symbol.usr)
   }
 }

Sources/SemanticIndex/CheckedIndex.swift

@@ -45,14 +45,25 @@ package enum IndexCheckLevel {
   case inMemoryModifiedFiles(any InMemoryDocumentManager)
 }
 
+struct IndexClosedError: Error, CustomStringConvertible {
+  var description: String { "Index has been closed" }
+}
+
 /// A wrapper around `IndexStoreDB` that checks if returned symbol occurrences are up-to-date with regard to a
 /// `IndexCheckLevel`.
 ///
 /// - SeeAlso: Comment on `IndexOutOfDateChecker`
 package final class CheckedIndex {
   private var checker: IndexOutOfDateChecker
   package let unchecked: UncheckedIndex
-  private var index: IndexStoreDB { unchecked.underlyingIndexStoreDB }
+  private var index: IndexStoreDB {
+    get throws {
+      guard let underlyingIndexStoreDB = unchecked.underlyingIndexStoreDB else {
+        throw IndexClosedError()
+      }
+      return underlyingIndexStoreDB
+    }
+  }
 
   /// Maps the USR of a symbol to its name and the name of all its containers, from outermost to innermost.
   ///
@@ -90,21 +101,21 @@ package final class CheckedIndex {
     byUSR usr: String,
     roles: SymbolRole,
     _ body: (SymbolOccurrence) -> Bool
-  ) -> Bool {
-    index.forEachSymbolOccurrence(byUSR: usr, roles: roles) { occurrence in
+  ) throws -> Bool {
+    try index.forEachSymbolOccurrence(byUSR: usr, roles: roles) { occurrence in
       guard self.checker.isUpToDate(occurrence.location) else {
         return true  // continue
       }
       return body(occurrence)
     }
   }
 
-  package func occurrences(ofUSR usr: String, roles: SymbolRole) -> [SymbolOccurrence] {
-    return index.occurrences(ofUSR: usr, roles: roles).filter { checker.isUpToDate($0.location) }
+  package func occurrences(ofUSR usr: String, roles: SymbolRole) throws -> [SymbolOccurrence] {
+    return try index.occurrences(ofUSR: usr, roles: roles).filter { checker.isUpToDate($0.location) }
   }
 
-  package func occurrences(relatedToUSR usr: String, roles: SymbolRole) -> [SymbolOccurrence] {
-    return index.occurrences(relatedToUSR: usr, roles: roles).filter { checker.isUpToDate($0.location) }
+  package func occurrences(relatedToUSR usr: String, roles: SymbolRole) throws -> [SymbolOccurrence] {
+    return try index.occurrences(relatedToUSR: usr, roles: roles).filter { checker.isUpToDate($0.location) }
   }
 
   @discardableResult package func forEachCanonicalSymbolOccurrence(
@@ -114,8 +125,8 @@ package final class CheckedIndex {
     subsequence: Bool,
     ignoreCase: Bool,
     body: (SymbolOccurrence) -> Bool
-  ) -> Bool {
-    index.forEachCanonicalSymbolOccurrence(
+  ) throws -> Bool {
+    try index.forEachCanonicalSymbolOccurrence(
       containing: pattern,
       anchorStart: anchorStart,
       anchorEnd: anchorEnd,
@@ -132,30 +143,30 @@ package final class CheckedIndex {
   @discardableResult package func forEachCanonicalSymbolOccurrence(
     byName name: String,
     body: (SymbolOccurrence) -> Bool
-  ) -> Bool {
-    index.forEachCanonicalSymbolOccurrence(byName: name) { occurrence in
+  ) throws -> Bool {
+    try index.forEachCanonicalSymbolOccurrence(byName: name) { occurrence in
       guard self.checker.isUpToDate(occurrence.location) else {
         return true  // continue
       }
       return body(occurrence)
     }
   }
 
-  package func symbols(inFilePath path: String) -> [Symbol] {
-    guard self.hasAnyUpToDateUnit(for: DocumentURI(filePath: path, isDirectory: false)) else {
+  package func symbols(inFilePath path: String) throws -> [Symbol] {
+    guard try self.hasAnyUpToDateUnit(for: DocumentURI(filePath: path, isDirectory: false)) else {
       return []
     }
-    return index.symbols(inFilePath: path)
+    return try index.symbols(inFilePath: path)
   }
 
   /// Returns all unit test symbol in unit files that reference one of the main files in `mainFilePaths`.
-  package func unitTests(referencedByMainFiles mainFilePaths: [String]) -> [SymbolOccurrence] {
-    return index.unitTests(referencedByMainFiles: mainFilePaths).filter { checker.isUpToDate($0.location) }
+  package func unitTests(referencedByMainFiles mainFilePaths: [String]) throws -> [SymbolOccurrence] {
+    return try index.unitTests(referencedByMainFiles: mainFilePaths).filter { checker.isUpToDate($0.location) }
   }
 
   /// Returns all unit test symbols in the index.
-  package func unitTests() -> [SymbolOccurrence] {
-    return index.unitTests().filter { checker.isUpToDate($0.location) }
+  package func unitTests() throws -> [SymbolOccurrence] {
+    return try index.unitTests().filter { checker.isUpToDate($0.location) }
   }
 
   /// Return `true` if a unit file has been indexed for the given file path after its last modification date.
@@ -165,20 +176,24 @@ package final class CheckedIndex {
   /// This means that at least a single build configuration of this file has been indexed since its last modification.
   /// This method does not care about which target (identified by output path in the index) produced the up-to-date
   /// unit.
-  package func hasAnyUpToDateUnit(for uri: DocumentURI, mainFile: DocumentURI? = nil) -> Bool {
-    return checker.indexHasUpToDateUnit(for: uri, mainFile: mainFile, index: index)
+  package func hasAnyUpToDateUnit(for uri: DocumentURI, mainFile: DocumentURI? = nil) throws -> Bool {
+    return try checker.indexHasUpToDateUnit(for: uri, mainFile: mainFile, index: index)
   }
 
   /// Return `true` if a unit file with the given output path has been indexed after its last modification date of
   /// `uri`.
   ///
   /// If `outputPath` is `notSupported`, this behaves the same as `hasAnyUpToDateUnit`.
-  package func hasUpToDateUnit(for uri: DocumentURI, mainFile: DocumentURI? = nil, outputPath: OutputPath) -> Bool {
+  package func hasUpToDateUnit(
+    for uri: DocumentURI,
+    mainFile: DocumentURI? = nil,
+    outputPath: OutputPath
+  ) throws -> Bool {
     switch outputPath {
     case .path(let outputPath):
-      return checker.indexHasUpToDateUnit(for: uri, outputPath: outputPath, index: index)
+      return try checker.indexHasUpToDateUnit(for: uri, outputPath: outputPath, index: index)
     case .notSupported:
-      return self.hasAnyUpToDateUnit(for: uri, mainFile: mainFile)
+      return try self.hasAnyUpToDateUnit(for: uri, mainFile: mainFile)
     }
   }
 
@@ -206,20 +221,20 @@ package final class CheckedIndex {
 
   /// If there are any definition occurrences of the given USR, return these.
   /// Otherwise return declaration occurrences.
-  package func definitionOrDeclarationOccurrences(ofUSR usr: String) -> [SymbolOccurrence] {
-    let definitions = occurrences(ofUSR: usr, roles: [.definition])
+  package func definitionOrDeclarationOccurrences(ofUSR usr: String) throws -> [SymbolOccurrence] {
+    let definitions = try occurrences(ofUSR: usr, roles: [.definition])
     if !definitions.isEmpty {
       return definitions
     }
-    return occurrences(ofUSR: usr, roles: [.declaration])
+    return try occurrences(ofUSR: usr, roles: [.declaration])
   }
 
   /// Find a `SymbolOccurrence` that is considered the primary definition of the symbol with the given USR.
   ///
   /// If the USR has an ambiguous definition, the most important role of this function is to deterministically return
   /// the same result every time.
-  package func primaryDefinitionOrDeclarationOccurrence(ofUSR usr: String) -> SymbolOccurrence? {
-    let result = definitionOrDeclarationOccurrences(ofUSR: usr).sorted().first
+  package func primaryDefinitionOrDeclarationOccurrence(ofUSR usr: String) throws -> SymbolOccurrence? {
+    let result = try definitionOrDeclarationOccurrences(ofUSR: usr).sorted().first
     if result == nil {
       logger.error("Failed to find definition of \(usr) in index")
     }
@@ -244,15 +259,15 @@ package final class CheckedIndex {
   ///   }
   /// }
   /// ```
-  package func containerNames(of symbol: SymbolOccurrence) -> [String] {
+  package func containerNames(of symbol: SymbolOccurrence) throws -> [String] {
     // The container name of accessors is the container of the surrounding variable.
     let accessorOf = symbol.relations.filter { $0.roles.contains(.accessorOf) }
     if let primaryVariable = accessorOf.sorted().first {
       if accessorOf.count > 1 {
         logger.fault("Expected an occurrence to an accessor of at most one symbol, not multiple")
       }
-      if let primaryVariable = primaryDefinitionOrDeclarationOccurrence(ofUSR: primaryVariable.symbol.usr) {
-        return containerNames(of: primaryVariable)
+      if let primaryVariable = try primaryDefinitionOrDeclarationOccurrence(ofUSR: primaryVariable.symbol.usr) {
+        return try containerNames(of: primaryVariable)
       }
     }
 
@@ -279,7 +294,7 @@ package final class CheckedIndex {
     }
 
     if containerSymbol.kind == .extension,
-      let extendedSymbol = self.occurrences(relatedToUSR: containerSymbol.usr, roles: .extendedBy).first?.symbol
+      let extendedSymbol = try self.occurrences(relatedToUSR: containerSymbol.usr, roles: .extendedBy).first?.symbol
     {
       containerSymbol = extendedSymbol
     }
@@ -291,12 +306,12 @@ package final class CheckedIndex {
     // of these files. But we expect all all of these declarations to have the same parent container names and we don't
     // care about locations here.
     var containerDefinition: SymbolOccurrence?
-    forEachSymbolOccurrence(byUSR: containerSymbol.usr, roles: [.definition, .declaration]) { occurrence in
+    try forEachSymbolOccurrence(byUSR: containerSymbol.usr, roles: [.definition, .declaration]) { occurrence in
       containerDefinition = occurrence
       return false  // stop iteration
     }
     if let containerDefinition {
-      result = self.containerNames(of: containerDefinition) + [containerSymbol.name]
+      result = try self.containerNames(of: containerDefinition) + [containerSymbol.name]
     } else {
       result = [containerSymbol.name]
     }
@@ -310,7 +325,12 @@ 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
+  // Ideally, this would be an isolated member instead of a `ThreadSafeBox`, but that causes issues with the workarounds
+  // around https://github.com/swiftlang/swift/issues/75600 when all functions become async.
+  private nonisolated let _underlyingIndexStoreDB: ThreadSafeBox<IndexStoreDB?>
+  package nonisolated var underlyingIndexStoreDB: IndexStoreDB? {
+    _underlyingIndexStoreDB.value
+  }
 
   /// 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 +344,14 @@ package final actor UncheckedIndex: Sendable {
       return nil
     }
     self.usesExplicitOutputPaths = usesExplicitOutputPaths
-    self.underlyingIndexStoreDB = index
+    self._underlyingIndexStoreDB = ThreadSafeBox(initialValue: index)
+  }
+
+  /// Close the index store, writing it to the `saved` directory on disk.
+  package func close() {
+    // IndexStoreDB writes the index to disk when the retain count of the `IndexStoreDB` object hits zero. We hope that
+    // nobody else still has a reference to `IndexStoreDB` here.
+    _underlyingIndexStoreDB.value = nil
   }
 
   /// Update the set of output paths that should be considered visible in the project. For example, if a source file is
@@ -333,6 +360,10 @@ package final actor UncheckedIndex: Sendable {
     guard usesExplicitOutputPaths else {
       return
     }
+    guard let underlyingIndexStoreDB else {
+      logger.error("Not setting unit output paths because the index was closed")
+      return
+    }
     let addedPaths = paths.filter { !unitOutputPaths.contains($0) }
     let removedPaths = unitOutputPaths.filter { !paths.contains($0) }
     underlyingIndexStoreDB.addUnitOutFilePaths(Array(addedPaths), waitForProcessing: false)
@@ -346,12 +377,22 @@ package final actor UncheckedIndex: Sendable {
 
   /// Wait for IndexStoreDB to be updated based on new unit files written to disk.
   package nonisolated func pollForUnitChangesAndWait() {
-    self.underlyingIndexStoreDB.pollForUnitChangesAndWait()
+    guard let underlyingIndexStoreDB else {
+      logger.error("Not polling for unit changes because the index was closed")
+      return
+    }
+
+    underlyingIndexStoreDB.pollForUnitChangesAndWait()
   }
 
   /// Import the units for the given output paths into indexstore-db. Returns after the import has finished.
   package nonisolated func processUnitsForOutputPathsAndWait(_ outputPaths: some Collection<String>) {
-    self.underlyingIndexStoreDB.processUnitsForOutputPathsAndWait(outputPaths)
+    guard let underlyingIndexStoreDB else {
+      logger.error("Not processing units for output paths because the index was closed")
+      return
+    }
+
+    underlyingIndexStoreDB.processUnitsForOutputPathsAndWait(outputPaths)
   }
 }
 

Sources/SemanticIndex/SemanticIndexManager.swift

@@ -504,7 +504,11 @@ package final actor SemanticIndexManager {
             logger.info("Not indexing \(uri.forLogging) because its output file could not be determined")
             continue
           }
-          if !indexFilesWithUpToDateUnits, modifiedFilesIndex.hasUpToDateUnit(for: uri, outputPath: outputPath) {
+          let hasUpToDateUnit =
+            orLog("Checking if modified file has up-to-date unit") {
+              try modifiedFilesIndex.hasUpToDateUnit(for: uri, outputPath: outputPath)
+            } ?? true
+          if !indexFilesWithUpToDateUnits, hasUpToDateUnit {
             continue
           }
           // If this is a source file, just index it.
@@ -549,9 +553,11 @@ package final actor SemanticIndexManager {
         )
         continue
       }
-      if !indexFilesWithUpToDateUnits,
-        modifiedFilesIndex.hasUpToDateUnit(for: uri, mainFile: mainFile, outputPath: outputPath)
-      {
+      let hasUpToDateUnit =
+        orLog("Checking if modified file has up-to-date unit") {
+          try modifiedFilesIndex.hasUpToDateUnit(for: uri, mainFile: mainFile, outputPath: outputPath)
+        } ?? true
+      if !indexFilesWithUpToDateUnits, hasUpToDateUnit {
         continue
       }
       guard let language = await buildServerManager.defaultLanguage(for: uri, in: targetAndOutputPath.key) else {

Sources/SemanticIndex/UpdateIndexStoreTaskDescription.swift

@@ -256,11 +256,14 @@ package struct UpdateIndexStoreTaskDescription: IndexTaskDescription {
       if indexFilesWithUpToDateUnit {
         return true
       }
-      let hasUpToDateUnit = index.checked(for: .modifiedFiles).hasUpToDateUnit(
-        for: fileInfo.sourceFile,
-        mainFile: fileInfo.mainFile,
-        outputPath: fileInfo.outputPath
-      )
+      let hasUpToDateUnit =
+        orLog("Checking if file has up-to-date unit") {
+          try index.checked(for: .modifiedFiles).hasUpToDateUnit(
+            for: fileInfo.sourceFile,
+            mainFile: fileInfo.mainFile,
+            outputPath: fileInfo.outputPath
+          )
+        } ?? true
       if hasUpToDateUnit {
         logger.debug("Not indexing \(fileInfo.file.forLogging) because index has an up-to-date unit")
         // We consider a file's index up-to-date if we have any up-to-date unit. Changing build settings does not

Sources/SourceKitLSP/DefinitionLocations.swift

@@ -104,7 +104,7 @@ package func definitionLocations(
   guard let usr = symbol.usr else { return DefinitionLocationsResult(locations: []) }
   logger.info("Performing indexed jump-to-definition with USR \(usr)")
 
-  let occurrences = index.definitionOrDeclarationOccurrences(ofUSR: usr)
+  let occurrences = try index.definitionOrDeclarationOccurrences(ofUSR: usr)
 
   if occurrences.isEmpty {
     if let bestLocalDeclaration = symbol.bestLocalDeclaration {

Sources/SourceKitLSP/IndexStoreDB+MainFilesProvider.swift

@@ -24,7 +24,11 @@ extension UncheckedIndex: BuildServerIntegration.MainFilesProvider {
   package func mainFiles(containing uri: DocumentURI, crossLanguage: Bool) -> Set<DocumentURI> {
     let mainFiles: Set<DocumentURI>
     if let filePath = orLog("File path to get main files", { try uri.fileURL?.filePath }) {
-      let mainFilePaths = self.underlyingIndexStoreDB.mainFilesContainingFile(
+      guard let underlyingIndexStoreDB = self.underlyingIndexStoreDB else {
+        logger.error("Not checking main files for URI because index has been closed")
+        return []
+      }
+      let mainFilePaths = underlyingIndexStoreDB.mainFilesContainingFile(
         path: filePath,
         crossLanguage: crossLanguage
       )