Add a maximum duration for sourcekitd requests

VS Code does not cancel semantic tokens requests. If a source file gets into a state where an AST build takes very long, this can cause us to wait for the semantic tokens from sourcekitd for a few minutes, effectively blocking all other semantic functionality in that file.

To circumvent this problem (or any other problem where an editor might not be cancelling requests they are no longer interested in) add a maximum request duration for SourceKitD requests, defaulting to 2 minutes.

rdar://130948453

Author: ahoppen

Sources/Diagnose/RunSourcekitdRequestCommand.swift

@@ -96,6 +96,9 @@ public struct RunSourceKitdRequestCommand: AsyncParsableCommand {
     case .requestCancelled:
       print("request cancelled")
       throw ExitCode(1)
+    case .timedOut:
+      print("request timed out")
+      throw ExitCode(1)
     case .missingRequiredSymbol:
       print("missing required symbol")
       throw ExitCode(1)

Sources/LanguageServerProtocol/Error.swift

@@ -81,7 +81,7 @@ public struct ErrorCode: RawRepresentable, Codable, Hashable, Sendable {
   /// It doesn't denote a real error code.
   public static let lspReservedErrorRangeEnd = ErrorCode(rawValue: -32800)
 
-  // MARK: SourceKit-LSP specifiic eror codes
+  // MARK: SourceKit-LSP specific error codes
   public static let workspaceNotOpen: ErrorCode = ErrorCode(rawValue: -32003)
 }
 

Sources/SKCore/SourceKitLSPOptions.swift

@@ -191,6 +191,10 @@ public struct SourceKitLSPOptions: Sendable, Codable {
   /// Whether background indexing is enabled.
   public var backgroundIndexing: Bool?
 
+  public var backgroundIndexingOrDefault: Bool {
+    return backgroundIndexing ?? false
+  }
+
   /// Experimental features that are enabled.
   public var experimentalFeatures: Set<ExperimentalFeature>? = nil
 
@@ -220,8 +224,21 @@ public struct SourceKitLSPOptions: Sendable, Codable {
     return .seconds(1)
   }
 
-  public var backgroundIndexingOrDefault: Bool {
-    return backgroundIndexing ?? false
+  /// The maximum duration that a sourcekitd request should be allowed to execute before being declared as timed out.
+  ///
+  /// In general, editors should cancel requests that they are no longer interested in, but in case editors don't cancel
+  /// requests, this ensures that a long-running non-cancelled request is not blocking sourcekitd and thus most semantic
+  /// functionality.
+  ///
+  /// In particular, VS Code does not cancel the semantic tokens request, which can cause a long-running AST build that
+  /// blocks sourcekitd.
+  public var sourcekitdRequestTimeout: Double? = nil
+
+  public var sourcekitdRequestTimeoutOrDefault: Duration {
+    if let sourcekitdRequestTimeout {
+      return .seconds(sourcekitdRequestTimeout)
+    }
+    return .seconds(120)
   }
 
   public init(
@@ -235,7 +252,8 @@ public struct SourceKitLSPOptions: Sendable, Codable {
     backgroundIndexing: Bool? = nil,
     experimentalFeatures: Set<ExperimentalFeature>? = nil,
     swiftPublishDiagnosticsDebounceDuration: Double? = nil,
-    workDoneProgressDebounceDuration: Double? = nil
+    workDoneProgressDebounceDuration: Double? = nil,
+    sourcekitdRequestTimeout: Double? = nil
   ) {
     self.swiftPM = swiftPM
     self.fallbackBuildSystem = fallbackBuildSystem
@@ -248,6 +266,7 @@ public struct SourceKitLSPOptions: Sendable, Codable {
     self.experimentalFeatures = experimentalFeatures
     self.swiftPublishDiagnosticsDebounceDuration = swiftPublishDiagnosticsDebounceDuration
     self.workDoneProgressDebounceDuration = workDoneProgressDebounceDuration
+    self.sourcekitdRequestTimeout = sourcekitdRequestTimeout
   }
 
   public init?(fromLSPAny lspAny: LSPAny?) throws {
@@ -300,7 +319,8 @@ public struct SourceKitLSPOptions: Sendable, Codable {
       swiftPublishDiagnosticsDebounceDuration: override?.swiftPublishDiagnosticsDebounceDuration
         ?? base.swiftPublishDiagnosticsDebounceDuration,
       workDoneProgressDebounceDuration: override?.workDoneProgressDebounceDuration
-        ?? base.workDoneProgressDebounceDuration
+        ?? base.workDoneProgressDebounceDuration,
+      sourcekitdRequestTimeout: override?.sourcekitdRequestTimeout ?? base.sourcekitdRequestTimeout
     )
   }
 

Sources/SourceKitD/SourceKitD.swift

@@ -84,31 +84,41 @@ public enum SKDError: Error, Equatable {
   /// The request was cancelled.
   case requestCancelled
 
+  /// The request exceeded the maximum allowed duration.
+  case timedOut
+
   /// Loading a required symbol from the sourcekitd library failed.
   case missingRequiredSymbol(String)
 }
 
 extension SourceKitD {
-
   // MARK: - Convenience API for requests.
 
   /// - Parameters:
-  ///   - req: The request to send to sourcekitd.
+  ///   - request: The request to send to sourcekitd.
+  ///   - timeout: The maximum duration how long to wait for a response. If no response is returned within this time,
+  ///     declare the request as having timed out.
   ///   - fileContents: The contents of the file that the request operates on. If sourcekitd crashes, the file contents
   ///     will be logged.
-  public func send(_ request: SKDRequestDictionary, fileContents: String?) async throws -> SKDResponseDictionary {
+  public func send(
+    _ request: SKDRequestDictionary,
+    timeout: Duration,
+    fileContents: String?
+  ) async throws -> SKDResponseDictionary {
     log(request: request)
 
-    let sourcekitdResponse: SKDResponse = try await withCancellableCheckedThrowingContinuation { continuation in
-      var handle: sourcekitd_api_request_handle_t? = nil
-      api.send_request(request.dict, &handle) { response in
-        continuation.resume(returning: SKDResponse(response!, sourcekitd: self))
-      }
-      return handle
-    } cancel: { handle in
-      if let handle {
-        logRequestCancellation(request: request)
-        api.cancel_request(handle)
+    let sourcekitdResponse = try await withTimeout(timeout) {
+      return try await withCancellableCheckedThrowingContinuation { continuation in
+        var handle: sourcekitd_api_request_handle_t? = nil
+        self.api.send_request(request.dict, &handle) { response in
+          continuation.resume(returning: SKDResponse(response!, sourcekitd: self))
+        }
+        return handle
+      } cancel: { handle in
+        if let handle {
+          self.logRequestCancellation(request: request)
+          self.api.cancel_request(handle)
+        }
       }
     }
 
@@ -118,6 +128,9 @@ extension SourceKitD {
       if sourcekitdResponse.error == .connectionInterrupted {
         log(crashedRequest: request, fileContents: fileContents)
       }
+      if sourcekitdResponse.error == .requestCancelled && !Task.isCancelled {
+        throw SKDError.timedOut
+      }
       throw sourcekitdResponse.error!
     }
 

Sources/SourceKitLSP/Rename.swift

@@ -359,7 +359,7 @@ extension SwiftLanguageService {
       keys.argNames: sourcekitd.array(name.parameters.map { $0.stringOrWildcard }),
     ])
 
-    let response = try await sourcekitd.send(req, fileContents: snapshot.text)
+    let response = try await sendSourcekitdRequest(req, fileContents: snapshot.text)
 
     guard let isZeroArgSelector: Int = response[keys.isZeroArgSelector],
       let selectorPieces: SKDResponseArray = response[keys.selectorPieces]
@@ -416,7 +416,7 @@ extension SwiftLanguageService {
       req.set(keys.baseName, to: name)
     }
 
-    let response = try await sourcekitd.send(req, fileContents: snapshot.text)
+    let response = try await sendSourcekitdRequest(req, fileContents: snapshot.text)
 
     guard let baseName: String = response[keys.baseName] else {
       throw NameTranslationError.malformedClangToSwiftTranslateNameResponse(response)
@@ -914,7 +914,7 @@ extension SwiftLanguageService {
       keys.renameLocations: locations,
     ])
 
-    let syntacticRenameRangesResponse = try await sourcekitd.send(skreq, fileContents: snapshot.text)
+    let syntacticRenameRangesResponse = try await sendSourcekitdRequest(skreq, fileContents: snapshot.text)
     guard let categorizedRanges: SKDResponseArray = syntacticRenameRangesResponse[keys.categorizedRanges] else {
       throw ResponseError.internalError("sourcekitd did not return categorized ranges")
     }

Sources/SourceKitLSP/Swift/CodeCompletion.swift

@@ -33,6 +33,7 @@ extension SwiftLanguageService {
     return try await CodeCompletionSession.completionList(
       sourcekitd: sourcekitd,
       snapshot: snapshot,
+      options: options,
       indentationWidth: inferredIndentationWidth,
       completionPosition: completionPos,
       completionUtf8Offset: offset,

Sources/SourceKitLSP/Swift/CodeCompletionSession.swift

@@ -13,6 +13,7 @@
 import Dispatch
 import LSPLogging
 import LanguageServerProtocol
+import SKCore
 import SKSupport
 import SourceKitD
 import SwiftExtensions
@@ -92,6 +93,7 @@ class CodeCompletionSession {
   static func completionList(
     sourcekitd: any SourceKitD,
     snapshot: DocumentSnapshot,
+    options: SourceKitLSPOptions,
     indentationWidth: Trivia?,
     completionPosition: Position,
     completionUtf8Offset: Int,
@@ -122,6 +124,7 @@ class CodeCompletionSession {
       let session = CodeCompletionSession(
         sourcekitd: sourcekitd,
         snapshot: snapshot,
+        options: options,
         indentationWidth: indentationWidth,
         utf8Offset: completionUtf8Offset,
         position: completionPosition,
@@ -139,6 +142,7 @@ class CodeCompletionSession {
 
   private let sourcekitd: any SourceKitD
   private let snapshot: DocumentSnapshot
+  private let options: SourceKitLSPOptions
   /// The inferred indentation width of the source file the completion is being performed in
   private let indentationWidth: Trivia?
   private let utf8StartOffset: Int
@@ -158,13 +162,15 @@ class CodeCompletionSession {
   private init(
     sourcekitd: any SourceKitD,
     snapshot: DocumentSnapshot,
+    options: SourceKitLSPOptions,
     indentationWidth: Trivia?,
     utf8Offset: Int,
     position: Position,
     compileCommand: SwiftCompileCommand?,
     clientSupportsSnippets: Bool
   ) {
     self.sourcekitd = sourcekitd
+    self.options = options
     self.indentationWidth = indentationWidth
     self.snapshot = snapshot
     self.utf8StartOffset = utf8Offset
@@ -193,7 +199,11 @@ class CodeCompletionSession {
       keys.compilerArgs: compileCommand?.compilerArgs as [SKDRequestValue]?,
     ])
 
-    let dict = try await sourcekitd.send(req, fileContents: snapshot.text)
+    let dict = try await sourcekitd.send(
+      req,
+      timeout: options.sourcekitdRequestTimeoutOrDefault,
+      fileContents: snapshot.text
+    )
     self.state = .open
 
     guard let completions: SKDResponseArray = dict[keys.results] else {
@@ -226,7 +236,11 @@ class CodeCompletionSession {
       keys.codeCompleteOptions: optionsDictionary(filterText: filterText),
     ])
 
-    let dict = try await sourcekitd.send(req, fileContents: snapshot.text)
+    let dict = try await sourcekitd.send(
+      req,
+      timeout: options.sourcekitdRequestTimeoutOrDefault,
+      fileContents: snapshot.text
+    )
     guard let completions: SKDResponseArray = dict[keys.results] else {
       return CompletionList(isIncomplete: false, items: [])
     }
@@ -269,7 +283,7 @@ class CodeCompletionSession {
         keys.name: snapshot.uri.pseudoPath,
       ])
       logger.info("Closing code completion session: \(self.description)")
-      _ = try? await sourcekitd.send(req, fileContents: nil)
+      _ = try? await sourcekitd.send(req, timeout: options.sourcekitdRequestTimeoutOrDefault, fileContents: nil)
       self.state = .closed
     }
   }

Sources/SourceKitLSP/Swift/CursorInfo.swift

@@ -159,7 +159,7 @@ extension SwiftLanguageService {
 
     appendAdditionalParameters?(skreq)
 
-    let dict = try await self.sourcekitd.send(skreq, fileContents: snapshot.text)
+    let dict = try await sendSourcekitdRequest(skreq, fileContents: snapshot.text)
 
     var cursorInfoResults: [CursorInfo] = []
     if let cursorInfo = CursorInfo(dict, sourcekitd: sourcekitd) {

Sources/SourceKitLSP/Swift/DiagnosticReportManager.swift

@@ -12,6 +12,7 @@
 
 import LSPLogging
 import LanguageServerProtocol
+import SKCore
 import SKSupport
 import SourceKitD
 import SwiftExtensions
@@ -22,6 +23,7 @@ actor DiagnosticReportManager {
   private typealias ReportTask = RefCountedCancellableTask<RelatedFullDocumentDiagnosticReport>
 
   private let sourcekitd: SourceKitD
+  private let options: SourceKitLSPOptions
   private let syntaxTreeManager: SyntaxTreeManager
   private let documentManager: DocumentManager
   private let clientHasDiagnosticsCodeDescriptionSupport: Bool
@@ -48,11 +50,13 @@ actor DiagnosticReportManager {
 
   init(
     sourcekitd: SourceKitD,
+    options: SourceKitLSPOptions,
     syntaxTreeManager: SyntaxTreeManager,
     documentManager: DocumentManager,
     clientHasDiagnosticsCodeDescriptionSupport: Bool
   ) {
     self.sourcekitd = sourcekitd
+    self.options = options
     self.syntaxTreeManager = syntaxTreeManager
     self.documentManager = documentManager
     self.clientHasDiagnosticsCodeDescriptionSupport = clientHasDiagnosticsCodeDescriptionSupport
@@ -104,7 +108,11 @@ actor DiagnosticReportManager {
       keys.compilerArgs: compilerArgs as [SKDRequestValue],
     ])
 
-    let dict = try await self.sourcekitd.send(skreq, fileContents: snapshot.text)
+    let dict = try await self.sourcekitd.send(
+      skreq,
+      timeout: options.sourcekitdRequestTimeoutOrDefault,
+      fileContents: snapshot.text
+    )
 
     try Task.checkCancellation()
 

Sources/SourceKitLSP/Swift/OpenInterface.swift

@@ -52,7 +52,7 @@ extension SwiftLanguageService {
         symbol: symbol
       )
       _ = await orLog("Closing generated interface") {
-        try await self.sourcekitd.send(closeDocumentSourcekitdRequest(uri: interfaceDocURI), fileContents: nil)
+        try await sendSourcekitdRequest(closeDocumentSourcekitdRequest(uri: interfaceDocURI), fileContents: nil)
       }
       return result
     }
@@ -80,7 +80,7 @@ extension SwiftLanguageService {
       keys.compilerArgs: await self.buildSettings(for: request.textDocument.uri)?.compilerArgs as [SKDRequestValue]?,
     ])
 
-    let dict = try await self.sourcekitd.send(skreq, fileContents: nil)
+    let dict = try await sendSourcekitdRequest(skreq, fileContents: nil)
     return GeneratedInterfaceInfo(contents: dict[keys.sourceText] ?? "")
   }
 
@@ -101,7 +101,7 @@ extension SwiftLanguageService {
         keys.usr: symbol,
       ])
 
-      let dict = try await self.sourcekitd.send(skreq, fileContents: snapshot.text)
+      let dict = try await sendSourcekitdRequest(skreq, fileContents: snapshot.text)
       if let offset: Int = dict[keys.offset] {
         return GeneratedInterfaceDetails(uri: uri, position: snapshot.positionOf(utf8Offset: offset))
       } else {