Show a progress indicator in the editor if SourceKit-LSP is reloading packages

I noticed that the initial package loading can take ~5s. It’s good behavior to inform the client that sourcekit-lsp is busy reloading the package, showing the user that semantic functionality might not be ready yet.

#620
rdar://111917300

Author: ahoppen

Sources/LSPTestSupport/TestJSONRPCConnection.swift

@@ -137,6 +137,13 @@ public final class TestClient: MessageHandler {
   }
 
   public func handle<R: RequestType>(_ params: R, id: RequestID, from clientID: ObjectIdentifier, reply: @escaping (LSPResult<R.Response>) -> Void) {
+    if R.self == CreateWorkDoneProgressRequest.self {
+      // We don’t want to require tests to specify request handlers for work done progress.
+      // Simply ignore requests to create WorkDoneProgress for now.
+      reply(.failure(.unknown("WorkDone not supported in TestClient")))
+      return
+    }
+
     let cancellationToken = CancellationToken()
 
     let request = Request(params, id: id, clientID: clientID, cancellation: cancellationToken, reply: reply)

Sources/SKSwiftPMWorkspace/SwiftPMWorkspace.swift

@@ -35,6 +35,14 @@ import func TSCBasic.resolveSymlinks
 import protocol TSCBasic.FileSystem
 import var TSCBasic.localFileSystem
 
+/// Parameter of `reloadPackageStatusCallback` in ``SwiftPMWorkspace``.
+///
+/// Informs the callback about whether `reloadPackage` started or finished executing.
+public enum ReloadPackageStatus {
+  case start
+  case end
+}
+
 /// Swift Package Manager build system and workspace support.
 ///
 /// This class implements the `BuildSystem` interface to provide the build settings for a Swift
@@ -78,18 +86,25 @@ public final class SwiftPMWorkspace {
   /// - `sourceDirToTarget`
   let queue: DispatchQueue = .init(label: "SwiftPMWorkspace.queue", qos: .utility)
 
+  /// This callback is informed when `reloadPackage` starts and ends executing.
+  var reloadPackageStatusCallback: (ReloadPackageStatus) -> Void
+
+
   /// Creates a build system using the Swift Package Manager, if this workspace is a package.
   ///
   /// - Parameters:
   ///   - workspace: The workspace root path.
   ///   - toolchainRegistry: The toolchain registry to use to provide the Swift compiler used for
   ///     manifest parsing and runtime support.
+  ///   - reloadPackageStatusCallback: Will be informed when `reloadPackage` starts and ends executing.
   /// - Throws: If there is an error loading the package, or no manifest is found.
   public init(
     workspacePath: TSCAbsolutePath,
     toolchainRegistry: ToolchainRegistry,
     fileSystem: FileSystem = localFileSystem,
-    buildSetup: BuildSetup) throws
+    buildSetup: BuildSetup,
+    reloadPackageStatusCallback: @escaping (ReloadPackageStatus) -> Void = { _ in }
+  ) throws
   {
     self.workspacePath = workspacePath
     self.fileSystem = fileSystem
@@ -142,23 +157,31 @@ public final class SwiftPMWorkspace {
     )
 
     self.packageGraph = try PackageGraph(rootPackages: [], dependencies: [], binaryArtifacts: [:])
+    self.reloadPackageStatusCallback = reloadPackageStatusCallback
 
     try reloadPackage()
   }
 
   /// Creates a build system using the Swift Package Manager, if this workspace is a package.
-  ///
+  /// 
+  /// - Parameters:
+  ///   - reloadPackageStatusCallback: Will be informed when `reloadPackage` starts and ends executing.
   /// - Returns: nil if `workspacePath` is not part of a package or there is an error.
-  public convenience init?(url: URL,
-                           toolchainRegistry: ToolchainRegistry,
-                           buildSetup: BuildSetup)
+  public convenience init?(
+    url: URL,
+    toolchainRegistry: ToolchainRegistry,
+    buildSetup: BuildSetup,
+    reloadPackageStatusCallback: @escaping (ReloadPackageStatus) -> Void
+  )
   {
     do {
       try self.init(
         workspacePath: try TSCAbsolutePath(validating: url.path),
         toolchainRegistry: toolchainRegistry,
         fileSystem: localFileSystem,
-        buildSetup: buildSetup)
+        buildSetup: buildSetup,
+        reloadPackageStatusCallback: reloadPackageStatusCallback
+      )
     } catch Error.noManifest(let path) {
       log("could not find manifest, or not a SwiftPM package: \(path)", level: .warning)
       return nil
@@ -175,6 +198,11 @@ extension SwiftPMWorkspace {
   /// dependencies.
   /// Must only be called on `queue` or from the initializer.
   func reloadPackage() throws {
+    reloadPackageStatusCallback(.start)
+    defer {
+      reloadPackageStatusCallback(.end)
+    }
+
 
     let observabilitySystem = ObservabilitySystem({ scope, diagnostic in
         log(diagnostic.description, level: diagnostic.severity.asLogLevel)

Sources/SourceKitLSP/SourceKitServer.swift

@@ -58,6 +58,84 @@ enum LanguageServerType: Hashable {
   }
 }
 
+/// Keeps track of the state to send work done progress updates to the client
+final class WorkDoneProgressState {
+  private enum State {
+    /// No `WorkDoneProgress` has been created.
+    case noProgress
+    /// We have sent the request to create a `WorkDoneProgress` but haven’t received a respose yet.
+    case creating
+    /// A `WorkDoneProgress` has been created.
+    case created
+    /// The creation of a `WorkDoneProgress has failed`.
+    ///
+    /// This causes us to just give up creating any more `WorkDoneProgress` in
+    /// the future as those will most likely also fail.
+    case progressCreationFailed
+  }
+
+  /// How many active tasks are running.
+  ///
+  /// A work done progress should be displayed if activeTasks > 0
+  private var activeTasks: Int = 0
+  private var state: State = .noProgress
+
+  /// The token by which we track the `WorkDoneProgress`.
+  private let token: ProgressToken
+
+  /// The title that should be displayed to the user in the UI.
+  private let title: String
+
+  init(_ token: String, title: String) {
+    self.token = ProgressToken.string(token)
+    self.title = title
+  }
+  
+  /// Start a new task, creating a new `WorkDoneProgress` if none is running right now.
+  ///
+  /// - Parameter server: The server that is used to create the `WorkDoneProgress` on the client
+  ///
+  /// - Important: Must be called on `server.queue`.
+  func startProgress(server: SourceKitServer) {
+    dispatchPrecondition(condition: .onQueue(server.queue))
+    activeTasks += 1
+    if state == .noProgress {
+      state = .creating
+      // Discard the handle. We don't support cancellation of the creation of a work done progress.
+      _ = server.client.send(CreateWorkDoneProgressRequest(token: token), queue: server.queue) { result in
+        if result.success != nil {
+          if self.activeTasks == 0 {
+            // ActiveTasks might have been decreased while we created the `WorkDoneProgress`
+            self.state = .noProgress
+            server.client.send(WorkDoneProgress(token: self.token, value: .end(WorkDoneProgressEnd())))
+          } else {
+            self.state = .created
+            server.client.send(WorkDoneProgress(token: self.token, value: .begin(WorkDoneProgressBegin(title: self.title))))
+          }
+        } else {
+          self.state = .progressCreationFailed
+        }
+      }
+    }
+  }
+
+  /// End a new task stated using `startProgress`.
+  ///
+  /// If this drops the active task count to 0, the work done progress is ended on the client.
+  ///
+  /// - Parameter server: The server that is used to send and update of the `WorkDoneProgress` to the client
+  ///
+  /// - Important: Must be called on `server.queue`.
+  func endProgress(server: SourceKitServer) {
+    dispatchPrecondition(condition: .onQueue(server.queue))
+    assert(activeTasks > 0, "Unbalanced startProgress/endProgress calls")
+    activeTasks -= 1
+    if state == .created && activeTasks == 0 {
+      server.client.send(WorkDoneProgress(token: token, value: .end(WorkDoneProgressEnd())))
+    }
+  }
+}
+
 /// The SourceKit language server.
 ///
 /// This is the client-facing language server implementation, providing indexing, multiple-toolchain
@@ -80,6 +158,8 @@ public final class SourceKitServer: LanguageServer {
 
   private let documentManager = DocumentManager()
 
+  private var packageLoadingWorkDoneProgress = WorkDoneProgressState("SourceKitLSP.SoruceKitServer.reloadPackage", title: "Reloading Package")
+
   /// **Public for testing**
   public var _documentManager: DocumentManager {
     return documentManager
@@ -559,7 +639,18 @@ extension SourceKitServer {
       capabilityRegistry: capabilityRegistry,
       toolchainRegistry: self.toolchainRegistry,
       buildSetup: self.options.buildSetup,
-      indexOptions: self.options.indexOptions)
+      indexOptions: self.options.indexOptions,
+      reloadPackageStatusCallback: { status in
+        self.queue.async {
+          switch status {
+          case .start:
+            self.packageLoadingWorkDoneProgress.startProgress(server: self)
+          case .end:
+            self.packageLoadingWorkDoneProgress.endProgress(server: self)
+          }
+        }
+      }
+    )
   }
 
   func initialize(_ req: Request<InitializeRequest>) {

Sources/SourceKitLSP/Workspace.swift

@@ -84,15 +84,19 @@ public final class Workspace {
     capabilityRegistry: CapabilityRegistry,
     toolchainRegistry: ToolchainRegistry,
     buildSetup: BuildSetup,
-    indexOptions: IndexOptions = IndexOptions()
+    indexOptions: IndexOptions = IndexOptions(),
+    reloadPackageStatusCallback: @escaping (ReloadPackageStatus) -> Void
   ) throws {
     var buildSystem: BuildSystem? = nil
     if let rootUrl = rootUri.fileURL, let rootPath = try? AbsolutePath(validating: rootUrl.path) {
       if let buildServer = BuildServerBuildSystem(projectRoot: rootPath, buildSetup: buildSetup) {
         buildSystem = buildServer
-      } else if let swiftpm = SwiftPMWorkspace(url: rootUrl,
-                                          toolchainRegistry: toolchainRegistry,
-                                          buildSetup: buildSetup) {
+      } else if let swiftpm = SwiftPMWorkspace(
+        url: rootUrl,
+        toolchainRegistry: toolchainRegistry,
+        buildSetup: buildSetup,
+        reloadPackageStatusCallback: reloadPackageStatusCallback
+      ) {
         buildSystem = swiftpm
       } else {
         buildSystem = CompilationDatabaseBuildSystem(projectRoot: rootPath)