Allow prefixing of the token for a WorkDoneProgress with a custom string

This allows VS Code to detect when sourcekitd has crashed and prompt the user to gather a diagnostic report + file an issue about the crash.

rdar://129678779
Fixes #1476

Author: ahoppen

Sources/SourceKitLSP/IndexProgressManager.swift

@@ -127,6 +127,7 @@ actor IndexProgressManager {
     } else {
       workDoneProgress = await WorkDoneProgressManager(
         server: sourceKitLSPServer,
+        tokenPrefix: "indexing",
         initialDebounce: sourceKitLSPServer.options.workDoneProgressDebounceDuration,
         title: "Indexing",
         message: message,

Sources/SourceKitLSP/SourceKitLSPServer.swift

@@ -126,10 +126,17 @@ public actor SourceKitLSPServer {
   /// initializer.
   private nonisolated(unsafe) var indexProgressManager: IndexProgressManager!
 
-  /// Number of workspaces that are currently reloading swift package. When this is not 0, a
-  /// `packageLoadingWorkDoneProgress` is created to show a work done progress indicator in the client.
-  private var inProgressPackageLoadingOperations = 0
-  private var packageLoadingWorkDoneProgress: WorkDoneProgressManager?
+  /// Implicitly unwrapped optional so we can create an `SharedWorkDoneProgressManager` that has a weak reference to
+  /// `SourceKitLSPServer`.
+  /// `nonisolated(unsafe)` because `packageLoadingWorkDoneProgress` will not be modified after it is assigned from the
+  /// initializer.
+  private nonisolated(unsafe) var packageLoadingWorkDoneProgress: SharedWorkDoneProgressManager!
+
+  /// Implicitly unwrapped optional so we can create an `SharedWorkDoneProgressManager` that has a weak reference to
+  /// `SourceKitLSPServer`.
+  /// `nonisolated(unsafe)` because `sourcekitdCrashedWorkDoneProgress` will not be modified after it is assigned from
+  /// the initializer.
+  nonisolated(unsafe) var sourcekitdCrashedWorkDoneProgress: SharedWorkDoneProgressManager!
 
   /// Caches which workspace a document with the given URI should be opened in.
   ///
@@ -210,6 +217,17 @@ public actor SourceKitLSPServer {
     ])
     self.indexProgressManager = nil
     self.indexProgressManager = IndexProgressManager(sourceKitLSPServer: self)
+    self.packageLoadingWorkDoneProgress = SharedWorkDoneProgressManager(
+      sourceKitLSPServer: self,
+      tokenPrefix: "package-reloading",
+      title: "SourceKit-LSP: Reloading Package"
+    )
+    self.sourcekitdCrashedWorkDoneProgress = SharedWorkDoneProgressManager(
+      sourceKitLSPServer: self,
+      tokenPrefix: "sourcekitd-crashed",
+      title: "SourceKit-LSP: Restoring functionality",
+      message: "Please run 'sourcekit-lsp diagnose' to file an issue"
+    )
   }
 
   /// Await until the server has send the reply to the initialize request.
@@ -885,21 +903,9 @@ extension SourceKitLSPServer {
   private func reloadPackageStatusCallback(_ status: ReloadPackageStatus) async {
     switch status {
     case .start:
-      inProgressPackageLoadingOperations += 1
-      if let capabilityRegistry, packageLoadingWorkDoneProgress == nil {
-        packageLoadingWorkDoneProgress = WorkDoneProgressManager(
-          server: self,
-          capabilityRegistry: capabilityRegistry,
-          initialDebounce: options.workDoneProgressDebounceDuration,
-          title: "SourceKit-LSP: Reloading Package"
-        )
-      }
+      await packageLoadingWorkDoneProgress.start()
     case .end:
-      inProgressPackageLoadingOperations -= 1
-      if inProgressPackageLoadingOperations == 0, let packageLoadingWorkDoneProgress {
-        self.packageLoadingWorkDoneProgress = nil
-        await packageLoadingWorkDoneProgress.end()
-      }
+      await packageLoadingWorkDoneProgress.end()
     }
   }
 

Sources/SourceKitLSP/Swift/SwiftLanguageService.swift

@@ -132,40 +132,8 @@ public actor SwiftLanguageService: LanguageService, Sendable {
   nonisolated var requests: sourcekitd_api_requests { return sourcekitd.requests }
   nonisolated var values: sourcekitd_api_values { return sourcekitd.values }
 
-  /// When sourcekitd is crashed, a `WorkDoneProgressManager` that display the sourcekitd crash status in the client.
-  private var sourcekitdCrashedWorkDoneProgress: WorkDoneProgressManager?
-
-  private var state: LanguageServerState {
-    didSet {
-      for handler in stateChangeHandlers {
-        handler(oldValue, state)
-      }
-
-      guard let sourceKitLSPServer else {
-        Task {
-          await sourcekitdCrashedWorkDoneProgress?.end()
-        }
-        sourcekitdCrashedWorkDoneProgress = nil
-        return
-      }
-      switch state {
-      case .connected:
-        Task {
-          await sourcekitdCrashedWorkDoneProgress?.end()
-        }
-        sourcekitdCrashedWorkDoneProgress = nil
-      case .connectionInterrupted, .semanticFunctionalityDisabled:
-        if sourcekitdCrashedWorkDoneProgress == nil {
-          sourcekitdCrashedWorkDoneProgress = WorkDoneProgressManager(
-            server: sourceKitLSPServer,
-            capabilityRegistry: capabilityRegistry,
-            title: "SourceKit-LSP: Restoring functionality",
-            message: "Please run 'sourcekit-lsp diagnose' to file an issue"
-          )
-        }
-      }
-    }
-  }
+  /// - Important: Use `setState` to change the state, which notifies the state change handlers
+  private var state: LanguageServerState
 
   private var stateChangeHandlers: [(_ oldState: LanguageServerState, _ newState: LanguageServerState) -> Void] = []
 
@@ -257,6 +225,30 @@ public actor SwiftLanguageService: LanguageService, Sendable {
     return true
   }
 
+  private func setState(_ newState: LanguageServerState) async {
+    let oldState = state
+    state = newState
+    for handler in stateChangeHandlers {
+      handler(oldState, newState)
+    }
+
+    guard let sourceKitLSPServer else {
+      return
+    }
+    switch (oldState, newState) {
+    case (.connected, .connectionInterrupted), (.connected, .semanticFunctionalityDisabled):
+      await sourceKitLSPServer.sourcekitdCrashedWorkDoneProgress.start()
+    case (.connectionInterrupted, .connected), (.semanticFunctionalityDisabled, .connected):
+      await sourceKitLSPServer.sourcekitdCrashedWorkDoneProgress.end()
+    case (.connected, .connected),
+      (.connectionInterrupted, .connectionInterrupted),
+      (.connectionInterrupted, .semanticFunctionalityDisabled),
+      (.semanticFunctionalityDisabled, .connectionInterrupted),
+      (.semanticFunctionalityDisabled, .semanticFunctionalityDisabled):
+      break
+    }
+  }
+
   public func addStateChangeHandler(
     handler: @Sendable @escaping (_ oldState: LanguageServerState, _ newState: LanguageServerState) -> Void
   ) {
@@ -982,13 +974,14 @@ extension SwiftLanguageService: SKDNotificationHandler {
     )
     // Check if we need to update our `state` based on the contents of the notification.
     if notification.value?[self.keys.notification] == self.values.semaEnabledNotification {
-      self.state = .connected
+      await self.setState(.connected)
+      return
     }
 
     if self.state == .connectionInterrupted {
       // If we get a notification while we are restoring the connection, it means that the server has restarted.
       // We still need to wait for semantic functionality to come back up.
-      self.state = .semanticFunctionalityDisabled
+      await self.setState(.semanticFunctionalityDisabled)
 
       // Ask our parent to re-open all of our documents.
       if let sourceKitLSPServer {
@@ -999,7 +992,7 @@ extension SwiftLanguageService: SKDNotificationHandler {
     }
 
     if notification.error == .connectionInterrupted {
-      self.state = .connectionInterrupted
+      await self.setState(.connectionInterrupted)
 
       // We don't have any open documents anymore after sourcekitd crashed.
       // Reset the document manager to reflect that.

Sources/SourceKitLSP/WorkDoneProgressManager.swift

@@ -20,7 +20,7 @@ import SwiftExtensions
 ///
 /// The work done progress is started when the object is created and ended when the object is destroyed.
 /// In between, updates can be sent to the client.
-final actor WorkDoneProgressManager {
+actor WorkDoneProgressManager {
   private enum Status: Equatable {
     case inProgress(message: String?, percentage: Int?)
     case done
@@ -36,6 +36,15 @@ final actor WorkDoneProgressManager {
 
   private weak var server: SourceKitLSPServer?
 
+  /// A string with which the `token` of the generated `WorkDoneProgress` sent to the client starts.
+  ///
+  /// A UUID will be appended to this prefix to make the token unique. The token prefix can be used to classify the work
+  /// done progress into a category, which makes debugging easier because the tokens have semantic meaning and also
+  /// allows clients to interpret what the `WorkDoneProgress` represents (for example Swift for VS Code explicitly
+  /// recognizes work done progress that indicates that sourcekitd has crashed to offer a diagnostic bundle to be
+  /// generated).
+  private let tokenPrefix: String
+
   private let title: String
 
   /// The next status that should be sent to the client by `sendProgressUpdateImpl`.
@@ -58,6 +67,7 @@ final actor WorkDoneProgressManager {
 
   init?(
     server: SourceKitLSPServer,
+    tokenPrefix: String,
     initialDebounce: Duration? = nil,
     title: String,
     message: String? = nil,
@@ -69,6 +79,7 @@ final actor WorkDoneProgressManager {
     self.init(
       server: server,
       capabilityRegistry: capabilityRegistry,
+      tokenPrefix: tokenPrefix,
       initialDebounce: initialDebounce,
       title: title,
       message: message,
@@ -79,6 +90,7 @@ final actor WorkDoneProgressManager {
   init?(
     server: SourceKitLSPServer,
     capabilityRegistry: CapabilityRegistry,
+    tokenPrefix: String,
     initialDebounce: Duration? = nil,
     title: String,
     message: String? = nil,
@@ -87,6 +99,7 @@ final actor WorkDoneProgressManager {
     guard capabilityRegistry.clientCapabilities.window?.workDoneProgress ?? false else {
       return nil
     }
+    self.tokenPrefix = tokenPrefix
     self.server = server
     self.title = title
     self.pendingStatus = .inProgress(message: message, percentage: percentage)
@@ -121,7 +134,7 @@ final actor WorkDoneProgressManager {
           )
         )
       } else {
-        let token = ProgressToken.string(UUID().uuidString)
+        let token = ProgressToken.string("\(tokenPrefix).\(UUID().uuidString)")
         do {
           _ = try await server.client.send(CreateWorkDoneProgressRequest(token: token))
         } catch {
@@ -177,3 +190,68 @@ final actor WorkDoneProgressManager {
     }
   }
 }
+
+/// A `WorkDoneProgressManager` that essentially has two states. If any operation tracked by this type is currently
+/// running, it displays a work done progress in the client. If multiple operations are running at the same time, it
+/// doesn't show multiple work done progress in the client. For example, we only want to show one progress indicator
+/// when sourcekitd has crashed, not one per `SwiftLanguageService`.
+actor SharedWorkDoneProgressManager {
+  private weak var sourceKitLSPServer: SourceKitLSPServer?
+
+  /// The number of in-progress operations. When greater than 0 `workDoneProgress` non-nil and a work done progress is
+  /// displayed to the user.
+  private var inProgressOperations = 0
+  private var workDoneProgress: WorkDoneProgressManager?
+
+  private let tokenPrefix: String
+  private let title: String
+  private let message: String?
+
+  public init(
+    sourceKitLSPServer: SourceKitLSPServer,
+    tokenPrefix: String,
+    title: String,
+    message: String? = nil
+  ) {
+    self.sourceKitLSPServer = sourceKitLSPServer
+    self.tokenPrefix = tokenPrefix
+    self.title = title
+    self.message = message
+  }
+
+  func start() async {
+    guard let sourceKitLSPServer else {
+      return
+    }
+    // Do all asynchronous operations up-front so that incrementing `inProgressOperations` and setting `workDoneProgress`
+    // cannot be interrupted by an `await` call
+    let initialDebounce = await sourceKitLSPServer.options.workDoneProgressDebounceDuration
+    let capabilityRegistry = await sourceKitLSPServer.capabilityRegistry
+
+    inProgressOperations += 1
+    if let capabilityRegistry, workDoneProgress == nil {
+      workDoneProgress = WorkDoneProgressManager(
+        server: sourceKitLSPServer,
+        capabilityRegistry: capabilityRegistry,
+        tokenPrefix: tokenPrefix,
+        initialDebounce: initialDebounce,
+        title: title,
+        message: message
+      )
+    }
+  }
+
+  func end() async {
+    if inProgressOperations > 0 {
+      inProgressOperations -= 1
+    } else {
+      logger.fault(
+        "Unbalanced calls to SharedWorkDoneProgressManager.start and end for \(self.tokenPrefix, privacy: .public)"
+      )
+    }
+    if inProgressOperations == 0, let workDoneProgress {
+      self.workDoneProgress = nil
+      await workDoneProgress.end()
+    }
+  }
+}