Merge pull request #1467 from ahoppen/debounce-index-progress

Debounce the creation of work done progress for indexing and package reloading

Author: ahoppen

Sources/SourceKitLSP/CMakeLists.txt

@@ -19,7 +19,6 @@ add_library(SourceKitLSP STATIC
   TestDiscovery.swift
   TextEdit+IsNoop.swift
   WorkDoneProgressManager.swift
-  WorkDoneProgressState.swift
   Workspace.swift
 )
 target_sources(SourceKitLSP PRIVATE

Sources/SourceKitLSP/IndexProgressManager.swift

@@ -78,6 +78,7 @@ actor IndexProgressManager {
 
   private func indexProgressStatusDidChangeImpl() async {
     guard let sourceKitLSPServer else {
+      await workDoneProgress?.end()
       workDoneProgress = nil
       return
     }
@@ -116,15 +117,17 @@ actor IndexProgressManager {
     case .upToDate:
       // Nothing left to index. Reset the target count and dismiss the work done progress.
       queuedIndexTasks = 0
+      await workDoneProgress?.end()
       workDoneProgress = nil
       return
     }
 
     if let workDoneProgress {
-      workDoneProgress.update(message: message, percentage: percentage)
+      await workDoneProgress.update(message: message, percentage: percentage)
     } else {
       workDoneProgress = await WorkDoneProgressManager(
         server: sourceKitLSPServer,
+        initialDebounce: sourceKitLSPServer.options.workDoneProgressDebounceDuration,
         title: "Indexing",
         message: message,
         percentage: percentage

Sources/SourceKitLSP/SourceKitLSPServer+Options.swift

@@ -48,6 +48,11 @@ extension SourceKitLSPServer {
     /// notification when running unit tests.
     public var swiftPublishDiagnosticsDebounceDuration: TimeInterval
 
+    /// When a task is started that should be displayed to the client as a work done progress, how many milliseconds to
+    /// wait before actually starting the work done progress. This prevents flickering of the work done progress in the
+    /// client for short-lived index tasks which end within this duration.
+    public var workDoneProgressDebounceDuration: Duration
+
     /// Experimental features that are enabled.
     public var experimentalFeatures: Set<ExperimentalFeature>
 
@@ -61,6 +66,7 @@ extension SourceKitLSPServer {
       completionOptions: SKCompletionOptions = .init(),
       generatedInterfacesPath: AbsolutePath = defaultDirectoryForGeneratedInterfaces,
       swiftPublishDiagnosticsDebounceDuration: TimeInterval = 2, /* 2s */
+      workDoneProgressDebounceDuration: Duration = .seconds(0),
       experimentalFeatures: Set<ExperimentalFeature> = [],
       indexTestHooks: IndexTestHooks = IndexTestHooks()
     ) {
@@ -72,6 +78,7 @@ extension SourceKitLSPServer {
       self.generatedInterfacesPath = generatedInterfacesPath
       self.swiftPublishDiagnosticsDebounceDuration = swiftPublishDiagnosticsDebounceDuration
       self.experimentalFeatures = experimentalFeatures
+      self.workDoneProgressDebounceDuration = workDoneProgressDebounceDuration
       self.indexTestHooks = indexTestHooks
     }
   }

Sources/SourceKitLSP/SourceKitLSPServer.swift

@@ -126,10 +126,10 @@ public actor SourceKitLSPServer {
   /// initializer.
   private nonisolated(unsafe) var indexProgressManager: IndexProgressManager!
 
-  private var packageLoadingWorkDoneProgress = WorkDoneProgressState(
-    "SourceKitLSP.SourceKitLSPServer.reloadPackage",
-    title: "SourceKit-LSP: Reloading Package"
-  )
+  /// 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?
 
   /// **Public for testing**
   public var _documentManager: DocumentManager {
@@ -896,6 +896,27 @@ 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"
+        )
+      }
+    case .end:
+      inProgressPackageLoadingOperations -= 1
+      if inProgressPackageLoadingOperations == 0, let packageLoadingWorkDoneProgress {
+        self.packageLoadingWorkDoneProgress = nil
+        await packageLoadingWorkDoneProgress.end()
+      }
+    }
+  }
+
   /// Creates a workspace at the given `uri`.
   ///
   /// If the build system that was determined for the workspace does not satisfy `condition`, `nil` is returned.
@@ -914,13 +935,7 @@ extension SourceKitLSPServer {
       options: options,
       toolchainRegistry: toolchainRegistry,
       reloadPackageStatusCallback: { [weak self] status in
-        guard let self else { return }
-        switch status {
-        case .start:
-          await self.packageLoadingWorkDoneProgress.startProgress(server: self)
-        case .end:
-          await self.packageLoadingWorkDoneProgress.endProgress(server: self)
-        }
+        await self?.reloadPackageStatusCallback(status)
       }
     )
     guard await condition(buildSystem) else {

Sources/SourceKitLSP/Swift/SwiftLanguageService.swift

@@ -142,11 +142,17 @@ public actor SwiftLanguageService: LanguageService, Sendable {
       }
 
       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 {

Sources/SourceKitLSP/WorkDoneProgressManager.swift

@@ -11,6 +11,7 @@
 //===----------------------------------------------------------------------===//
 
 import Foundation
+import LSPLogging
 import LanguageServerProtocol
 import SKSupport
 import SwiftExtensions
@@ -19,89 +20,155 @@ 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 class WorkDoneProgressManager {
-  private let token: ProgressToken
-  private let queue = AsyncQueue<Serial>()
-  private let server: SourceKitLSPServer
-  /// `true` if the client returned without an error from the `CreateWorkDoneProgressRequest`.
+final actor WorkDoneProgressManager {
+  private enum Status: Equatable {
+    case inProgress(message: String?, percentage: Int?)
+    case done
+  }
+
+  /// The token with which the work done progress has been created. `nil` if no work done progress has been created yet,
+  /// either because we didn't send the `WorkDoneProgress` request yet, because the work done progress creation failed,
+  /// or because the work done progress has been ended.
+  private var token: ProgressToken?
+
+  /// The queue on which progress updates are sent to the client.
+  private let progressUpdateQueue = AsyncQueue<Serial>()
+
+  private weak var server: SourceKitLSPServer?
+
+  private let title: String
+
+  /// The next status that should be sent to the client by `sendProgressUpdateImpl`.
   ///
-  /// Since all work done progress reports are being sent on `queue`, we never access it in a state where the
-  /// `CreateWorkDoneProgressRequest` is still in progress.
+  /// While progress updates are being queued in `progressUpdateQueue` this status can evolve. The next
+  /// `sendProgressUpdateImpl` call will pick up the latest status.
   ///
-  /// Must be a reference because `deinit` captures it and wants to observe changes to it from `init` eg. in the
-  /// following:
-  ///  - `init` is called
-  ///  - `deinit` is called
-  ///  - The task from `init` gets executed
-  ///  - The task from `deinit` gets executed
-  ///    - This should have `workDoneProgressCreated == true` so that it can send the work progress end.
-  private let workDoneProgressCreated: ThreadSafeBox<Bool> & AnyObject = ThreadSafeBox<Bool>(initialValue: false)
+  /// For example, if we receive two update calls to 25% and 50% in quick succession the `sendProgressUpdateImpl`
+  /// scheduled from the 25% update will already pick up the new 50% status. The `sendProgressUpdateImpl` call scheduled
+  /// from the 50% update will then realize that the `lastStatus` is already up-to-date and be a no-op.
+  private var pendingStatus: Status
 
-  /// The last message and percentage so we don't send a new report notification to the client if `update` is called
-  /// without any actual change.
-  private var lastStatus: (message: String?, percentage: Int?)
+  /// The last status that was sent to the client. Used so we don't send no-op updates to the client.
+  private var lastStatus: Status? = nil
 
-  convenience init?(server: SourceKitLSPServer, title: String, message: String? = nil, percentage: Int? = nil) async {
+  init?(
+    server: SourceKitLSPServer,
+    initialDebounce: Duration? = nil,
+    title: String,
+    message: String? = nil,
+    percentage: Int? = nil
+  ) async {
     guard let capabilityRegistry = await server.capabilityRegistry else {
       return nil
     }
-    self.init(server: server, capabilityRegistry: capabilityRegistry, title: title, message: message)
+    self.init(
+      server: server,
+      capabilityRegistry: capabilityRegistry,
+      initialDebounce: initialDebounce,
+      title: title,
+      message: message,
+      percentage: percentage
+    )
   }
 
   init?(
     server: SourceKitLSPServer,
     capabilityRegistry: CapabilityRegistry,
+    initialDebounce: Duration? = nil,
     title: String,
     message: String? = nil,
     percentage: Int? = nil
   ) {
     guard capabilityRegistry.clientCapabilities.window?.workDoneProgress ?? false else {
       return nil
     }
-    self.token = .string("WorkDoneProgress-\(UUID())")
     self.server = server
-    queue.async { [server, token, workDoneProgressCreated] in
-      await server.waitUntilInitialized()
-      do {
-        _ = try await server.client.send(CreateWorkDoneProgressRequest(token: token))
-      } catch {
-        return
+    self.title = title
+    self.pendingStatus = .inProgress(message: message, percentage: percentage)
+    progressUpdateQueue.async {
+      if let initialDebounce {
+        try? await Task.sleep(for: initialDebounce)
       }
-      server.sendNotificationToClient(
-        WorkDoneProgress(
-          token: token,
-          value: .begin(WorkDoneProgressBegin(title: title, message: message, percentage: percentage))
-        )
-      )
-      workDoneProgressCreated.value = true
-      self.lastStatus = (message, percentage)
+      await self.sendProgressUpdateAssumingOnProgressUpdateQueue()
     }
   }
 
-  func update(message: String? = nil, percentage: Int? = nil) {
-    queue.async { [server, token, workDoneProgressCreated] in
-      guard workDoneProgressCreated.value else {
-        return
+  /// Send the necessary messages to the client to update the work done progress to `status`.
+  ///
+  /// Must be called on `progressUpdateQueue`
+  private func sendProgressUpdateAssumingOnProgressUpdateQueue() async {
+    let statusToSend = pendingStatus
+    guard statusToSend != lastStatus else {
+      return
+    }
+    guard let server else {
+      // SourceKitLSPServer has been destroyed, we don't have a way to send notifications to the client anymore.
+      return
+    }
+    await server.waitUntilInitialized()
+    switch statusToSend {
+    case .inProgress(message: let message, percentage: let percentage):
+      if let token {
+        server.sendNotificationToClient(
+          WorkDoneProgress(
+            token: token,
+            value: .report(WorkDoneProgressReport(cancellable: false, message: message, percentage: percentage))
+          )
+        )
+      } else {
+        let token = ProgressToken.string(UUID().uuidString)
+        do {
+          _ = try await server.client.send(CreateWorkDoneProgressRequest(token: token))
+        } catch {
+          return
+        }
+        server.sendNotificationToClient(
+          WorkDoneProgress(
+            token: token,
+            value: .begin(WorkDoneProgressBegin(title: title, message: message, percentage: percentage))
+          )
+        )
+        self.token = token
       }
-      guard (message, percentage) != self.lastStatus else {
-        return
+    case .done:
+      if let token {
+        server.sendNotificationToClient(WorkDoneProgress(token: token, value: .end(WorkDoneProgressEnd())))
+        self.token = nil
       }
-      self.lastStatus = (message, percentage)
-      server.sendNotificationToClient(
-        WorkDoneProgress(
-          token: token,
-          value: .report(WorkDoneProgressReport(cancellable: false, message: message, percentage: percentage))
-        )
-      )
+    }
+    lastStatus = statusToSend
+  }
+
+  func update(message: String? = nil, percentage: Int? = nil) {
+    pendingStatus = .inProgress(message: message, percentage: percentage)
+    progressUpdateQueue.async {
+      await self.sendProgressUpdateAssumingOnProgressUpdateQueue()
+    }
+  }
+
+  /// Ends the work done progress. Any further update calls are no-ops.
+  ///
+  /// `end` must be should be called before the `WorkDoneProgressManager` is deallocated.
+  func end() {
+    pendingStatus = .done
+    progressUpdateQueue.async {
+      await self.sendProgressUpdateAssumingOnProgressUpdateQueue()
     }
   }
 
   deinit {
-    queue.async { [server, token, workDoneProgressCreated] in
-      guard workDoneProgressCreated.value else {
-        return
+    if pendingStatus != .done {
+      // If there is still a pending work done progress, end it. We know that we don't have any pending updates on
+      // `progressUpdateQueue` because they would capture `self` strongly and thus we wouldn't be deallocating this
+      // object.
+      // This is a fallback logic to ensure we don't leave pending work done progresses in the editor if the
+      // `WorkDoneProgressManager` is destroyed without a call to `end` (eg. because its owning object is destroyed).
+      // Calling `end()` is preferred because it ends the work done progress even if there are pending status updates
+      // in `progressUpdateQueue`, which keep the `WorkDoneProgressManager` alive and thus prevent the work done
+      // progress to be implicitly ended by the deinitializer.
+      if let token {
+        server?.sendNotificationToClient(WorkDoneProgress(token: token, value: .end(WorkDoneProgressEnd())))
       }
-      server.sendNotificationToClient(WorkDoneProgress(token: token, value: .end(WorkDoneProgressEnd())))
     }
   }
 }