Merge pull request #2422 from loveucifer/fix-language-service-shutdown-2211

Sources/ClangLanguageService/ClangLanguageService.swift

@@ -393,7 +393,9 @@ extension ClangLanguageService {
       guard let clangd else { return }
       // Give clangd 2 seconds to shut down by itself. If it doesn't shut down within that time, terminate it.
       try await withTimeout(.seconds(2)) {
-        _ = try await clangd.send(ShutdownRequest())
+        let shutdownRequest = ShutdownRequest()
+        await self.sourceKitLSPServer?.hooks.preForwardRequestToClangd?(shutdownRequest)
+        _ = try await clangd.send(shutdownRequest)
         clangd.send(ExitNotification())
       }
     }

Sources/SourceKitLSP/LanguageService.swift

@@ -126,6 +126,14 @@ package protocol LanguageService: AnyObject, Sendable {
   /// Experimental capabilities that should be reported to the client if this language service is enabled.
   static var experimentalCapabilities: [String: LSPAny] { get }
 
+  /// Whether this language service should be kept alive when its workspace is closed.
+  ///
+  /// When `true`, the language service will not be shut down even if all workspaces referencing it are closed.
+  /// This is useful for language services that use global state (like sourcekitd) where shutting down and
+  /// restarting would cause unnecessary overhead since the new instance will just reinitialize the same
+  /// global state.
+  static var isImmortal: Bool { get }
+
   // MARK: - Lifetime
 
   func initialize(_ initialize: InitializeRequest) async throws -> InitializeResult
@@ -351,6 +359,8 @@ package extension LanguageService {
 
   static var experimentalCapabilities: [String: LSPAny] { [:] }
 
+  static var isImmortal: Bool { false }
+
   func clientInitialized(_ initialized: InitializedNotification) async {}
 
   func openOnDiskDocument(snapshot: DocumentSnapshot, buildSettings: FileBuildSettings) async throws {

Sources/SourceKitLSP/SourceKitLSPServer.swift

@@ -1529,6 +1529,50 @@ extension SourceKitLSPServer {
         self.workspacesAndIsImplicit += newWorkspaces.map { (workspace: $0, isImplicit: false) }
       }
     }.value
+
+    // Shut down any language services that are no longer referenced by any workspace.
+    await self.shutdownOrphanedLanguageServices()
+  }
+
+  /// Shuts down any language services that are no longer referenced by any open workspace.
+  ///
+  /// This method gathers all language services that are currently referenced by open workspaces
+  /// and shuts down any language services that are not in that set.
+  private func shutdownOrphanedLanguageServices() async {
+    // Gather all language services referenced by open workspaces
+    var referencedServices: Set<ObjectIdentifier> = []
+    for workspace in workspaces {
+      for languageService in workspace.allLanguageServices {
+        referencedServices.insert(ObjectIdentifier(languageService))
+      }
+    }
+
+    // Find and remove orphaned language services, skipping immortal ones
+    var orphanedServices: [any LanguageService] = []
+    for (serviceType, services) in languageServices {
+      var remainingServices: [any LanguageService] = []
+      for service in services {
+        if referencedServices.contains(ObjectIdentifier(service)) || type(of: service).isImmortal {
+          remainingServices.append(service)
+        } else {
+          orphanedServices.append(service)
+        }
+      }
+      if remainingServices.count != services.count {
+        languageServices[serviceType] = remainingServices.isEmpty ? nil : remainingServices
+      }
+    }
+
+    // Shut down orphaned services in a background task to avoid blocking other requests.
+
+    if !orphanedServices.isEmpty {
+      Task {
+        for service in orphanedServices {
+          logger.info("Shutting down orphaned language service: \(type(of: service))")
+          await service.shutdown()
+        }
+      }
+    }
   }
 
   func didChangeWatchedFiles(_ notification: DidChangeWatchedFilesNotification) async {

Sources/SourceKitLSP/Workspace.swift

@@ -197,6 +197,11 @@ package final class Workspace: Sendable, BuildServerManagerDelegate {
   /// Language service for an open document, if available.
   private let languageServices: ThreadSafeBox<[DocumentURI: [any LanguageService]]> = ThreadSafeBox(initialValue: [:])
 
+  /// All language services that are registered with this workspace.
+  var allLanguageServices: [any LanguageService] {
+    return languageServices.value.values.flatMap { $0 }
+  }
+
   /// The task that constructs the `SemanticIndexManager`, which keeps track of whose file's index is up-to-date in the
   /// workspace and schedules indexing and preparation tasks for files with out-of-date index.
   ///

Sources/SwiftLanguageService/SwiftCommand.swift

@@ -55,4 +55,11 @@ extension SwiftLanguageService {
       command.identifier
     }
   }
+
+  /// `SwiftLanguageService` is immortal because sourcekitd uses global state.
+  ///
+  /// Since all instances of `SwiftLanguageService` share the same underlying sourcekitd process,
+  /// shutting down and restarting would cause unnecessary overhead as the new instance would
+  /// just reinitialize the same global state. Instead, we keep the service alive.
+  package static var isImmortal: Bool { true }
 }

Tests/SourceKitLSPTests/WorkspaceTests.swift

@@ -1386,6 +1386,180 @@ final class WorkspaceTests: SourceKitLSPTestCase {
     )
     XCTAssertEqual(outputPaths.outputPaths.map { $0.suffix(13) }.sorted(), ["FileA.swift.o", "FileB.swift.o"])
   }
+
+  func testOrphanedClangLanguageServiceShutdown() async throws {
+    // test that when we remove a workspace, the ClangLanguageService for that workspace is shut down.
+    // verify this by checking that clangd receives a ShutdownRequest.
+
+    let clangdReceivedShutdown = self.expectation(description: "clangd received shutdown request")
+    clangdReceivedShutdown.assertForOverFulfill = false
+
+    let project = try await MultiFileTestProject(
+      files: [
+        "WorkspaceA/compile_flags.txt": "",
+        "WorkspaceA/dummy.c": "",
+        "WorkspaceB/main.c": """
+        int main() { return 0; }
+        """,
+        "WorkspaceB/compile_flags.txt": "",
+      ],
+      workspaces: { scratchDir in
+        return [
+          WorkspaceFolder(uri: DocumentURI(scratchDir.appending(component: "WorkspaceA"))),
+          WorkspaceFolder(uri: DocumentURI(scratchDir.appending(component: "WorkspaceB"))),
+        ]
+      },
+      hooks: Hooks(preForwardRequestToClangd: { request in
+        if request is ShutdownRequest {
+          clangdReceivedShutdown.fulfill()
+        }
+      })
+    )
+
+    // open a .c file in WorkspaceB to launch clangd
+    let (mainUri, _) = try project.openDocument("main.c")
+
+    // send a request to ensure clangd is up and running
+    _ = try await project.testClient.send(
+      DocumentSymbolRequest(textDocument: TextDocumentIdentifier(mainUri))
+    )
+
+    // Get the language service for WorkspaceB before closing
+    let clangLanguageServiceBeforeClose = try await project.testClient.server.primaryLanguageService(
+      for: mainUri,
+      .c,
+      in: unwrap(project.testClient.server.workspaceForDocument(uri: mainUri))
+    )
+
+    // close the document
+    project.testClient.send(DidCloseTextDocumentNotification(textDocument: TextDocumentIdentifier(mainUri)))
+
+    // remove WorkspaceB
+    let workspaceBUri = DocumentURI(project.scratchDirectory.appending(component: "WorkspaceB"))
+    project.testClient.send(
+      DidChangeWorkspaceFoldersNotification(
+        event: WorkspaceFoldersChangeEvent(removed: [WorkspaceFolder(uri: workspaceBUri)])
+      )
+    )
+    _ = try await project.testClient.send(SynchronizeRequest())
+    // wait for clangd to receive the shutdown request
+    try await fulfillmentOfOrThrow(clangdReceivedShutdown)
+
+    let workspaceAfterRemoval = await project.testClient.server.workspaceForDocument(uri: mainUri)
+    XCTAssertNotEqual(
+      try XCTUnwrap(workspaceAfterRemoval?.rootUri?.fileURL?.lastPathComponent),
+      "WorkspaceB",
+      "WorkspaceB should have been removed"
+    )
+
+    // verify the language service is orphaned - opening a file in WorkspaceA should get a different language service
+    let (dummyUri, _) = try project.openDocument("dummy.c")
+    _ = try await project.testClient.send(
+      DocumentSymbolRequest(textDocument: TextDocumentIdentifier(dummyUri))
+    )
+
+    let clangLanguageServiceForWorkspaceA = try await project.testClient.server.primaryLanguageService(
+      for: dummyUri,
+      .c,
+      in: unwrap(project.testClient.server.workspaceForDocument(uri: dummyUri))
+    )
+
+    XCTAssertFalse(
+      clangLanguageServiceBeforeClose === clangLanguageServiceForWorkspaceA,
+      "WorkspaceB's ClangLanguageService should have been shut down and a new one created for WorkspaceA"
+    )
+  }
+
+  func testOrphanedSwiftLanguageServiceIsImmortal() async throws {
+    try await SkipUnless.sourcekitdSupportsPlugin()
+
+    let project = try await MultiFileTestProject(
+      files: [
+        "WorkspaceA/Sources/LibA/LibA.swift": """
+        public struct LibA {
+          public func 1️⃣foo() {}
+          public init() {}
+        }
+        """,
+        "WorkspaceA/Package.swift": """
+        // swift-tools-version: 5.7
+        import PackageDescription
+        let package = Package(
+          name: "LibA",
+          targets: [.target(name: "LibA")]
+        )
+        """,
+        "WorkspaceB/Sources/LibB/LibB.swift": """
+        public struct LibB {
+          public func 2️⃣bar() {}
+          public init() {}
+        }
+        """,
+        "WorkspaceB/Package.swift": """
+        // swift-tools-version: 5.7
+        import PackageDescription
+        let package = Package(
+          name: "LibB",
+          targets: [.target(name: "LibB")]
+        )
+        """,
+      ],
+      workspaces: { scratchDir in
+        return [
+          WorkspaceFolder(uri: DocumentURI(scratchDir.appending(component: "WorkspaceA"))),
+          WorkspaceFolder(uri: DocumentURI(scratchDir.appending(component: "WorkspaceB"))),
+        ]
+      }
+    )
+
+    let (libBUri, libBPositions) = try project.openDocument("LibB.swift")
+
+    let initialHover = try await project.testClient.send(
+      HoverRequest(textDocument: TextDocumentIdentifier(libBUri), position: libBPositions["2️⃣"])
+    )
+    XCTAssertNotNil(initialHover, "Should get hover response for LibB.swift")
+
+    // Get the SwiftLanguageService before closing WorkspaceB
+    let swiftLanguageServiceBeforeClose = try await project.testClient.server.primaryLanguageService(
+      for: libBUri,
+      .swift,
+      in: unwrap(project.testClient.server.workspaceForDocument(uri: libBUri))
+    )
+
+    // Close the document in WorkspaceB
+    project.testClient.send(DidCloseTextDocumentNotification(textDocument: TextDocumentIdentifier(libBUri)))
+
+    // Remove WorkspaceB
+    let workspaceBUri = DocumentURI(project.scratchDirectory.appending(component: "WorkspaceB"))
+    project.testClient.send(
+      DidChangeWorkspaceFoldersNotification(
+        event: WorkspaceFoldersChangeEvent(removed: [WorkspaceFolder(uri: workspaceBUri)])
+      )
+    )
+    _ = try await project.testClient.send(SynchronizeRequest())
+
+    // Open a file in WorkspaceA
+    let (libAUri, libAPositions) = try project.openDocument("LibA.swift")
+
+    // Verify that the language service in WorkspaceA still works correctly
+    let hover = try await project.testClient.send(
+      HoverRequest(textDocument: TextDocumentIdentifier(libAUri), position: libAPositions["1️⃣"])
+    )
+    XCTAssertNotNil(hover, "Should still get hover response after removing WorkspaceB")
+    assertContains(hover?.contents.markupContent?.value ?? "", "foo")
+
+    // Verify that the same SwiftLanguageService is reused (immortal, not shut down)
+    let swiftLanguageServiceForWorkspaceA = try await project.testClient.server.primaryLanguageService(
+      for: libAUri,
+      .swift,
+      in: unwrap(project.testClient.server.workspaceForDocument(uri: libAUri))
+    )
+
+    XCTAssertTrue(
+      swiftLanguageServiceBeforeClose === swiftLanguageServiceForWorkspaceA,
+      "SwiftLanguageService should be immortal and reused across workspaces"
+    )
+  }
 }
 
 private let defaultSDKArgs: String = {