Merge pull request #2409 from loveucifer/fix-language-services-lifecycle

Fix language services lifecycle management in Workspace

Sources/SourceKitLSP/SourceKitLSPServer.swift

@@ -592,7 +592,7 @@ package actor SourceKitLSPServer {
       """
     )
 
-    return workspace.setLanguageServices(for: uri, languageServices)
+    return languageServices
   }
 
   /// The language service with the highest precedence that can handle the given document.
@@ -1350,6 +1350,7 @@ extension SourceKitLSPServer {
     let language = textDocument.language
 
     let languageServices = await languageServices(for: uri, language, in: workspace)
+    workspace.setLanguageServices(for: uri, languageServices)
 
     if languageServices.isEmpty {
       // If we can't create a service, this document is unsupported and we can bail here.
@@ -1404,6 +1405,8 @@ extension SourceKitLSPServer {
       await languageService.closeDocument(notification)
     }
 
+    workspace.removeLanguageServices(for: uri)
+
     workspaceQueue.async {
       self.workspaceForUri[notification.textDocument.uri] = nil
     }

Sources/SourceKitLSP/Workspace.swift

@@ -456,23 +456,23 @@ package final class Workspace: Sendable, BuildServerManagerDelegate {
     return languageServices(for: uri).first
   }
 
-  /// Set a language service for a document uri and returns if none exists already.
+  /// Set the language services for a document URI.
   ///
-  /// If language services already exist for this document, eg. because two requests start creating a language
-  /// service for a document and race, `newLanguageServices` is dropped and the existing language services for the
-  /// document are returned.
-  func setLanguageServices(for uri: DocumentURI, _ newLanguageService: [any LanguageService]) -> [any LanguageService] {
-    return languageServices.withLock { languageServices in
-      if let languageService = languageServices[uri] {
-        return languageService
-      }
-
-      languageServices[uri] = newLanguageService
-      return newLanguageService
+  /// This should only be called from `openDocument` to ensure there are no race conditions.
+  func setLanguageServices(for uri: DocumentURI, _ newLanguageService: [any LanguageService]) {
+    languageServices.withLock { languageServices in
+      languageServices[uri.buildSettingsFile] = newLanguageService
     }
   }
 
-  /// Handle a build settings change notification from the build serveer.
+  /// Remove the language services association for a document when it is closed.
+  func removeLanguageServices(for uri: DocumentURI) {
+    languageServices.withLock { languageServices in
+      languageServices[uri.buildSettingsFile] = nil
+    }
+  }
+
+  /// Handle a build settings change notification from the build server.
   /// This has two primary cases:
   /// - Initial settings reported for a given file, now we can fully open it
   /// - Changed settings for an already open file