Ensure that unit file generated by background indexing are immediately loaded into indextore-db

Otherwise, we can end up in a situation where we declare indexing as done but haven’t loaded the new units into indexstore-db yet.

Author: ahoppen

Sources/SemanticIndex/CheckedIndex.swift

@@ -336,6 +336,11 @@ package final actor UncheckedIndex: Sendable {
   package nonisolated func pollForUnitChangesAndWait() {
     self.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)
+  }
 }
 
 /// Helper class to check if symbols from the index are up-to-date or if the source file has been modified after it was

Sources/SemanticIndex/UpdateIndexStoreTaskDescription.swift

@@ -188,6 +188,17 @@ package struct UpdateIndexStoreTaskDescription: IndexTaskDescription {
           outputPath: fileIndexInfo.outputPath
         )
       }
+      // If we know the output paths, make sure that we load their units into indexstore-db. We would eventually also
+      // pick the units up through file watching but that would leave a short time period in which we think that
+      // indexing has finished (because the index process has terminated) but when the new symbols aren't present in
+      // indexstore-db.
+      let outputPaths = filesToIndex.compactMap { fileToIndex in
+        switch fileToIndex.outputPath {
+        case .path(let string): return string
+        case .notSupported: return nil
+        }
+      }
+      index.processUnitsForOutputPathsAndWait(outputPaths)
       await hooks.updateIndexStoreTaskDidFinish?(self)
       logger.log(
         "Finished updating index store in \(Date().timeIntervalSince(startDate) * 1000, privacy: .public)ms: \(filesToIndexDescription)"

Tests/SourceKitLSPTests/BackgroundIndexingTests.swift

@@ -2585,6 +2585,48 @@ final class BackgroundIndexingTests: XCTestCase {
     let symbols = try await project.testClient.send(WorkspaceSymbolsRequest(query: "myTestFu"))
     XCTAssertEqual(symbols?.compactMap(\.symbolInformation?.name), ["myTestFunc()"])
   }
+
+  func testEnsureSymbolsLoadedIntoIndexstoreDbWhenIndexingHasFinished() async throws {
+    let testSetupComplete = AtomicBool(initialValue: false)
+    let updateIndexStoreStarted = self.expectation(description: "Update index store started")
+    let project = try await SwiftPMTestProject(
+      files: [
+        "Test.swift": ""
+      ],
+      options: .testDefault(experimentalFeatures: [.isIndexingRequest]),
+      hooks: Hooks(
+        indexHooks: IndexHooks(updateIndexStoreTaskDidStart: { _ in
+          guard testSetupComplete.value else {
+            return
+          }
+          updateIndexStoreStarted.fulfill()
+        })
+      ),
+      enableBackgroundIndexing: true,
+      pollIndex: false
+    )
+
+    try await project.changeFileOnDisk("Test.swift", newMarkedContents: "")
+    try await project.testClient.send(SynchronizeRequest(index: true))
+    let symbolsBeforeUpdate = try await project.testClient.send(WorkspaceSymbolsRequest(query: "myTestFu"))
+    XCTAssertEqual(symbolsBeforeUpdate, [])
+
+    testSetupComplete.value = true
+    try await project.changeFileOnDisk(
+      "Test.swift",
+      newMarkedContents: """
+        func myTestFunc() {}
+        """
+    )
+    try await fulfillmentOfOrThrow([updateIndexStoreStarted])
+    try await repeatUntilExpectedResult(sleepInterval: .milliseconds(2)) {
+      try await !project.testClient.send(IsIndexingRequest()).indexing
+    }
+    // Check that the newly added function has been registered in indexstore-db once indexing is done and that there is
+    // no time gap in which indexing has finished but the new unit has not been loaded into indexstore-db yet.
+    let symbols = try await project.testClient.send(WorkspaceSymbolsRequest(query: "myTestFu"))
+    XCTAssertEqual(symbols?.count, 1)
+  }
 }
 
 extension HoverResponseContents {