Wait until initialization has finished before starting a work done progress

Author: ahoppen

Sources/SourceKitLSP/SourceKitLSPServer.swift

@@ -115,6 +115,10 @@ final actor WorkDoneProgressState {
     case progressCreationFailed
   }
 
+  /// A queue so we can have synchronous `startProgress` and `endProgress` functions that don't need to wait for the
+  /// work done progress to be started or ended.
+  private let queue = AsyncQueue<Serial>()
+
   /// How many active tasks are running.
   ///
   /// A work done progress should be displayed if activeTasks > 0
@@ -135,13 +139,16 @@ final actor WorkDoneProgressState {
   /// 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
-  func startProgress(server: SourceKitLSPServer) async {
+  nonisolated func startProgress(server: SourceKitLSPServer) {
+    queue.async {
+      await self.startProgressImpl(server: server)
+    }
+  }
+
+  func startProgressImpl(server: SourceKitLSPServer) async {
+    await server.waitUntilInitialized()
     activeTasks += 1
     guard await server.capabilityRegistry?.clientCapabilities.window?.workDoneProgress ?? false else {
-      // If the client doesn't support workDoneProgress, keep track of the active task count but don't update the state.
-      // That way, if we call `startProgress` before initialization finishes, we won't send the
-      // `CreateWorkDoneProgressRequest` but if we call `startProgress` again after initialization finished (and thus
-      // the capability is set), we will create the work done progress.
       return
     }
     if state == .noProgress {
@@ -180,7 +187,13 @@ final actor WorkDoneProgressState {
   /// 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
-  func endProgress(server: SourceKitLSPServer) async {
+  nonisolated func endProgress(server: SourceKitLSPServer) {
+    queue.async {
+      await self.endProgressImpl(server: server)
+    }
+  }
+
+  func endProgressImpl(server: SourceKitLSPServer) async {
     assert(activeTasks > 0, "Unbalanced startProgress/endProgress calls")
     activeTasks -= 1
     guard await server.capabilityRegistry?.clientCapabilities.window?.workDoneProgress ?? false else {
@@ -440,11 +453,16 @@ public actor SourceKitLSPServer {
   /// The connection to the editor.
   public let client: Connection
 
+  /// Set to `true` after the `SourceKitLSPServer` has send the reply to the `InitializeRequest`.
+  ///
+  /// Initialization can be awaited using `waitUntilInitialized`.
+  private var initialized: Bool = false
+
   var options: Options
 
   let toolchainRegistry: ToolchainRegistry
 
-  var capabilityRegistry: CapabilityRegistry?
+  public var capabilityRegistry: CapabilityRegistry?
 
   var languageServices: [LanguageServerType: [LanguageService]] = [:]
 
@@ -508,19 +526,15 @@ public actor SourceKitLSPServer {
     self.inProgressRequests[id] = task
   }
 
-  let fs: FileSystem
-
   var onExit: () -> Void
 
   /// Creates a language server for the given client.
   public init(
     client: Connection,
-    fileSystem: FileSystem = localFileSystem,
     toolchainRegistry: ToolchainRegistry,
     options: Options,
     onExit: @escaping () -> Void = {}
   ) {
-    self.fs = fileSystem
     self.toolchainRegistry = toolchainRegistry
     self.options = options
     self.onExit = onExit
@@ -534,6 +548,22 @@ public actor SourceKitLSPServer {
     ])
   }
 
+  /// Await until the server has send the reply to the initialize request.
+  func waitUntilInitialized() async {
+    // The polling of `initialized` is not perfect but it should be OK, because
+    //  - In almost all cases the server should already be initialized.
+    //  - If it's not initialized, we expect initialization to finish fairly quickly. Even if initialization takes 5s
+    //    this only results in 50 polls, which is acceptable.
+    // Alternative solutions that signal via an async sequence seem overkill here.
+    while !initialized {
+      do {
+        try await Task.sleep(for: .seconds(0.1))
+      } catch {
+        break
+      }
+    }
+  }
+
   /// Search through all the parent directories of `uri` and check if any of these directories contain a workspace
   /// capable of handling `uri`.
   ///
@@ -958,6 +988,8 @@ extension SourceKitLSPServer: MessageHandler {
     switch request {
     case let request as RequestAndReply<InitializeRequest>:
       await request.reply { try await initialize(request.params) }
+      // Only set `initialized` to `true` after we have sent the response to the initialize request to the client.
+      initialized = true
     case let request as RequestAndReply<ShutdownRequest>:
       await request.reply { try await shutdown(request.params) }
     case let request as RequestAndReply<WorkspaceSymbolsRequest>:

Sources/SourceKitLSP/WorkDoneProgressManager.swift

@@ -23,20 +23,37 @@ final class WorkDoneProgressManager {
   private let queue = AsyncQueue<Serial>()
   private let server: SourceKitLSPServer
 
-  init?(server: SourceKitLSPServer, capabilityRegistry: CapabilityRegistry, title: String, message: String? = nil) {
+  convenience init?(server: SourceKitLSPServer, 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)
+  }
+
+  init?(
+    server: SourceKitLSPServer,
+    capabilityRegistry: CapabilityRegistry,
+    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] in
+      await server.waitUntilInitialized()
       await withCheckedContinuation { (continuation: CheckedContinuation<Void, Never>) in
         _ = server.client.send(CreateWorkDoneProgressRequest(token: token)) { result in
           continuation.resume()
         }
       }
       await server.sendNotificationToClient(
-        WorkDoneProgress(token: token, value: .begin(WorkDoneProgressBegin(title: title, message: message)))
+        WorkDoneProgress(
+          token: token,
+          value: .begin(WorkDoneProgressBegin(title: title, message: message, percentage: percentage))
+        )
       )
     }
   }