Debounce the creation of work done progress for indexing

This avoids us from showing a work done progress for short-lived index tasks that take less than a second.

Fixes #1298
rdar://128071328

Author: ahoppen

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.indexProgressDebounceDuration,
         title: "Indexing",
         message: message,
         percentage: percentage

Sources/SourceKitLSP/SourceKitLSPServer+Options.swift

@@ -48,6 +48,10 @@ extension SourceKitLSPServer {
     /// notification when running unit tests.
     public var swiftPublishDiagnosticsDebounceDuration: TimeInterval
 
+    /// When an index task is started, how many milliseconds to wait before sending a work done progress to the client.
+    // This prevents flickering of the work done progress in the client for short-lived index tasks
+    public var indexProgressDebounceDuration: Duration
+
     /// Experimental features that are enabled.
     public var experimentalFeatures: Set<ExperimentalFeature>
 
@@ -61,6 +65,7 @@ extension SourceKitLSPServer {
       completionOptions: SKCompletionOptions = .init(),
       generatedInterfacesPath: AbsolutePath = defaultDirectoryForGeneratedInterfaces,
       swiftPublishDiagnosticsDebounceDuration: TimeInterval = 2, /* 2s */
+      indexProgressDebounceDuration: Duration = .seconds(0),
       experimentalFeatures: Set<ExperimentalFeature> = [],
       indexTestHooks: IndexTestHooks = IndexTestHooks()
     ) {
@@ -72,6 +77,7 @@ extension SourceKitLSPServer {
       self.generatedInterfacesPath = generatedInterfacesPath
       self.swiftPublishDiagnosticsDebounceDuration = swiftPublishDiagnosticsDebounceDuration
       self.experimentalFeatures = experimentalFeatures
+      self.indexProgressDebounceDuration = indexProgressDebounceDuration
       self.indexTestHooks = indexTestHooks
     }
   }

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,29 +20,44 @@ 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
     }
@@ -51,57 +67,101 @@ final class WorkDoneProgressManager {
   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("WorkDoneProgress-\(UUID())")
+        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())))
     }
   }
 }

Sources/sourcekit-lsp/SourceKitLSP.swift

@@ -193,11 +193,6 @@ struct SourceKitLSP: AsyncParsableCommand {
   )
   var generatedInterfacesPath = defaultDirectoryForGeneratedInterfaces
 
-  @Option(
-    help: "When server-side filtering is enabled, the maximum number of results to return"
-  )
-  var completionMaxResults = 200
-
   @Option(
     name: .customLong("experimental-feature"),
     help: """
@@ -207,6 +202,14 @@ struct SourceKitLSP: AsyncParsableCommand {
   )
   var experimentalFeatures: [ExperimentalFeature] = []
 
+  // MARK: Configuration options intended for SourceKit-LSP developers.
+
+  @Option(help: .hidden)
+  var completionMaxResults = 200
+
+  @Option(help: .hidden)
+  var indexProgressDebounceDuration: Int = 1_000
+
   func mapOptions() -> SourceKitLSPServer.Options {
     var serverOptions = SourceKitLSPServer.Options()
 
@@ -222,9 +225,10 @@ struct SourceKitLSP: AsyncParsableCommand {
     serverOptions.indexOptions.indexStorePath = indexStorePath
     serverOptions.indexOptions.indexDatabasePath = indexDatabasePath
     serverOptions.indexOptions.indexPrefixMappings = indexPrefixMappings
-    serverOptions.completionOptions.maxResults = completionMaxResults
     serverOptions.generatedInterfacesPath = generatedInterfacesPath
     serverOptions.experimentalFeatures = Set(experimentalFeatures)
+    serverOptions.completionOptions.maxResults = completionMaxResults
+    serverOptions.indexProgressDebounceDuration = .milliseconds(indexProgressDebounceDuration)
 
     return serverOptions
   }