Define commands as asynchronous and use Task for preview cancellation (#1049)

* Define commands/actions as asynchronous

* Use `Task` for preview action cancellation

* Move mutable logHandle variables in preview tests inside Tasks

* Fix flaky test that sometimes hung due to blocking the async tasks

* Fix memory leak where DiagnosticEngine retains itself

* Fix PreviewServer retain cycle

* Fix another flaky test that blocked the async tasks

* Fix leak of NIO thread pools and event loops during tests

Author: d-ronnqvist

Sources/SwiftDocC/Infrastructure/Diagnostics/DiagnosticEngine.swift

@@ -29,7 +29,9 @@ public final class DiagnosticEngine {
     /// diagnostics with a severity up to and including `.information` will be printed.
     public var filterLevel: DiagnosticSeverity {
         didSet {
-            self.filter = { $0.diagnostic.severity.rawValue <= self.filterLevel.rawValue }
+            self.filter = { [filterLevel] in
+                $0.diagnostic.severity.rawValue <= filterLevel.rawValue
+            }
         }
     }
 

Sources/SwiftDocC/Infrastructure/DocumentationContext.swift

@@ -1331,6 +1331,7 @@ public class DocumentationContext: DocumentationContextDataProviderDelegate {
     }
 
     private func shouldContinueRegistration() throws {
+        try Task.checkCancellation()
         guard isRegistrationEnabled.sync({ $0 }) else {
             throw ContextError.registrationDisabled
         }
@@ -1773,6 +1774,7 @@ public class DocumentationContext: DocumentationContextDataProviderDelegate {
     ///
     /// When given `false` the context will try to cancel as quick as possible
     /// any ongoing bundle registrations.
+    @available(*, deprecated, message: "This deprecated API will be removed after 6.2 is released")
     public func setRegistrationEnabled(_ value: Bool) {
         isRegistrationEnabled.sync({ $0 = value })
     }

Sources/SwiftDocC/Infrastructure/DocumentationConverter.swift

@@ -50,6 +50,7 @@ public struct DocumentationConverter: DocumentationConverterProtocol {
     private var dataProvider: DocumentationWorkspaceDataProvider
 
     /// An optional closure that sets up a context before the conversion begins.
+    @available(*, deprecated, message: "This deprecated API will be removed after 6.2 is released")
     public var setupContext: ((inout DocumentationContext) -> Void)?
 
     /// Conversion batches should be big enough to keep all cores busy but small enough not to keep
@@ -189,49 +190,17 @@ public struct DocumentationConverter: DocumentationConverterProtocol {
         if let dataProvider = self.currentDataProvider {
             try workspace.unregisterProvider(dataProvider)
         }
-        
-        // Do additional context setup.
-        setupContext?(&context)
 
-        /*
-           Asynchronously cancel registration if necessary.
-           We spawn a timer that periodically checks `isCancelled` and if necessary
-           disables registration in `DocumentationContext` as registration being
-           the largest part of a documentation conversion.
-        */
         let context = self.context
-        let isCancelled = self.isCancelled
-        
-        // `true` if the `isCancelled` flag is set.
-        func isConversionCancelled() -> Bool {
-            return isCancelled?.sync({ $0 }) == true
-        }
-
-        // Run a timer that synchronizes the cancelled state between the converter and the context directly.
-        // We need a timer on a separate dispatch queue because `workspace.registerProvider()` blocks
-        // the current thread until it loads all symbol graphs, markdown files, and builds the topic graph
-        // so in order to be able to update the context cancellation flag we need to run on a different thread.
-        var cancelTimerQueue: DispatchQueue? = DispatchQueue(label: "org.swift.docc.ConvertActionCancelTimer", qos: .unspecified, attributes: .concurrent)
-        let cancelTimer = DispatchSource.makeTimerSource(queue: cancelTimerQueue)
-        cancelTimer.schedule(deadline: .now(), repeating: .milliseconds(500), leeway: .milliseconds(50))
-        cancelTimer.setEventHandler {
-            if isConversionCancelled() {
-                cancelTimer.cancel()
-                context.setRegistrationEnabled(false)
-            }
-        }
-        cancelTimer.resume()
 
         // Start bundle registration
         try workspace.registerProvider(dataProvider, options: bundleDiscoveryOptions)
         self.currentDataProvider = dataProvider
-
-        // Bundle registration is finished - stop the timer and reset the context cancellation state.
-        cancelTimer.cancel()
-        cancelTimerQueue = nil
-        context.setRegistrationEnabled(true)
 
         // If cancelled, return early before we emit diagnostics.
+        func isConversionCancelled() -> Bool {
+            Task.isCancelled || isCancelled?.sync({ $0 }) == true
+        }
         guard !isConversionCancelled() else { return ([], []) }
 
         processingDurationMetric = benchmark(begin: Benchmark.Duration(id: "documentation-processing"))

Sources/SwiftDocCUtilities/Action/Action.swift

@@ -1,7 +1,7 @@
 /*
  This source file is part of the Swift.org open source project
 
- Copyright (c) 2021 Apple Inc. and the Swift project authors
+ Copyright (c) 2021-2024 Apple Inc. and the Swift project authors
  Licensed under Apache License v2.0 with Runtime Library Exception
 
  See https://swift.org/LICENSE.txt for license information
@@ -13,18 +13,16 @@ import SwiftDocC
 
 /// An independent unit of work in the command-line workflow.
 ///
-/// An `Action` represents a discrete documentation task; it takes options and inputs, 
-/// performs its work, reports any problems it encounters, and outputs it generates.
-public protocol Action {
+/// An action represents a discrete documentation task; it takes options and inputs, performs its work, reports any problems it encounters, and outputs it generates.
+package protocol AsyncAction {
     /// Performs the action and returns an ``ActionResult``.
-    mutating func perform(logHandle: LogHandle) throws -> ActionResult
+    mutating func perform(logHandle: inout LogHandle) async throws -> ActionResult
 }
 
-/// An action for which you can optionally customize the documentation context.
-public protocol RecreatingContext: Action {
-    /// A closure that an action calls with the action's context for built documentation,
-    /// before the action performs work.
-    ///
-    /// Use this closure to set the action's context to a certain state before the action runs.
-    var setupContext: ((inout DocumentationContext) -> Void)? { get set }
+package extension AsyncAction {
+    mutating func perform(logHandle: LogHandle) async throws -> ActionResult {
+        var logHandle = logHandle
+        return try await perform(logHandle: &logHandle)
+    }
 }
+

Sources/SwiftDocCUtilities/Action/Actions/Action+MoveOutput.swift

@@ -11,7 +11,7 @@
 import Foundation
 import SwiftDocC
 
-extension Action {
+extension AsyncAction {
 
     /// Creates a new unique directory, with an optional template, inside of specified container.
     /// - Parameters:

Sources/SwiftDocCUtilities/Action/Actions/Convert/ConvertAction.swift

@@ -14,7 +14,7 @@ import Foundation
 import SwiftDocC
 
 /// An action that converts a source bundle into compiled documentation.
-public struct ConvertAction: Action, RecreatingContext {
+public struct ConvertAction: AsyncAction {
     enum Error: DescribedError {
         case doesNotContainBundle(url: URL)
         case cancelPending
@@ -59,12 +59,6 @@ public struct ConvertAction: Action, RecreatingContext {
     private var fileManager: FileManagerProtocol
     private let temporaryDirectory: URL
 
-    public var setupContext: ((inout DocumentationContext) -> Void)? {
-        didSet {
-            converter.setupContext = setupContext
-        }
-    }
-    
     var converter: DocumentationConverter
 
     private var durationMetric: Benchmark.Duration?
@@ -239,52 +233,17 @@ public struct ConvertAction: Action, RecreatingContext {
             dataProvider: dataProvider,
             bundleDiscoveryOptions: bundleDiscoveryOptions,
             sourceRepository: sourceRepository,
-            isCancelled: isCancelled,
             diagnosticEngine: self.diagnosticEngine,
             experimentalModifyCatalogWithGeneratedCuration: experimentalModifyCatalogWithGeneratedCuration
         )
     }
-
-    /// `true` if the convert action is cancelled.
-    private let isCancelled = Synchronized<Bool>(false)
 
-    /// `true` if the convert action is currently running.
-    let isPerforming = Synchronized<Bool>(false)
-    
-    /// A block to execute when conversion has finished.
-    /// It's used as a "future" for when the action is cancelled.
-    var didPerformFuture: (()->Void)?
-    
-    /// A block to execute when conversion has started.
-    var willPerformFuture: (()->Void)?
-
-    /// Cancels the action.
-    ///
-    /// The method blocks until the action has completed cancelling.
-    mutating func cancel() throws {
-        /// If the action is not running, there is nothing to cancel
-        guard isPerforming.sync({ $0 }) == true else { return }
-        
-        /// If the action is already cancelled throw `cancelPending`.
-        if isCancelled.sync({ $0 }) == true {
-            throw Error.cancelPending
-        }
-
-        /// Set the cancelled flag.
-        isCancelled.sync({ $0 = true })
-        
-        /// Wait for the `perform(logHandle:)` method to call `didPerformFuture()`
-        let waitGroup = DispatchGroup()
-        waitGroup.enter()
-        didPerformFuture = {
-            waitGroup.leave()
-        }
-        waitGroup.wait()
-    }
+    /// A block of extra work that tests perform to affect the time it takes to convert documentation
+    var _extraTestWork: (() async -> Void)?
 
     /// Converts each eligible file from the source documentation bundle,
     /// saves the results in the given output alongside the template files.
-    mutating public func perform(logHandle: LogHandle) throws -> ActionResult {
+    public mutating func perform(logHandle: inout LogHandle) async throws -> ActionResult {
         // Add the default diagnostic console writer now that we know what log handle it should write to.
         if !diagnosticEngine.hasConsumer(matching: { $0 is DiagnosticConsoleWriter }) {
             diagnosticEngine.add(
@@ -302,15 +261,13 @@ public struct ConvertAction: Action, RecreatingContext {
         var postConversionProblems: [Problem] = []
         let totalTimeMetric = benchmark(begin: Benchmark.Duration(id: "convert-total-time"))
 
-        // While running this method keep the `isPerforming` flag up.
-        isPerforming.sync({ $0 = true })
-        willPerformFuture?()
         defer {
-            didPerformFuture?()
-            isPerforming.sync({ $0 = false })
             diagnosticEngine.flush()
         }
 
+        // Run any extra work that the test may have injected
+        await _extraTestWork?()
+        
         let temporaryFolder = try createTempFolder(with: htmlTemplateDirectory)
 
         defer {
@@ -451,7 +408,7 @@ public struct ConvertAction: Action, RecreatingContext {
         benchmark(end: totalTimeMetric)
 
         if !didEncounterError {
-            let coverageResults = try coverageAction.perform(logHandle: logHandle)
+            let coverageResults = try await coverageAction.perform(logHandle: &logHandle)
             postConversionProblems.append(contentsOf: coverageResults.problems)
         }
 

Sources/SwiftDocCUtilities/Action/Actions/CoverageAction.swift

@@ -12,24 +12,24 @@ import Foundation
 import SwiftDocC
 
 /// An action that creates documentation coverage info for a documentation bundle.
-public struct CoverageAction: Action {
-    internal init(
+public struct CoverageAction: AsyncAction {
+    init(
         documentationCoverageOptions: DocumentationCoverageOptions,
         workingDirectory: URL,
-        fileManager: FileManagerProtocol) {
+        fileManager: FileManagerProtocol
+    ) {
         self.documentationCoverageOptions = documentationCoverageOptions
         self.workingDirectory = workingDirectory
         self.fileManager = fileManager
     }
 
     public let documentationCoverageOptions: DocumentationCoverageOptions
-    internal let workingDirectory: URL
+    let workingDirectory: URL
     private let fileManager: FileManagerProtocol
 
-    public mutating func perform(logHandle: LogHandle) throws -> ActionResult {
+    public mutating func perform(logHandle: inout LogHandle) async throws -> ActionResult {
         switch documentationCoverageOptions.level {
         case .brief, .detailed:
-            var logHandle = logHandle
             print("   --- Experimental coverage output enabled. ---", to: &logHandle)
 
             let summaryString = try CoverageDataEntry.generateSummary(

Sources/SwiftDocCUtilities/Action/Actions/EmitGeneratedCurationAction.swift

@@ -12,7 +12,7 @@ import Foundation
 import SwiftDocC
 
 /// An action that emits documentation extension files that reflect the auto-generated curation.
-struct EmitGeneratedCurationAction: Action {
+struct EmitGeneratedCurationAction: AsyncAction {
     let catalogURL: URL?
     let additionalSymbolGraphDirectory: URL?
     let outputURL: URL
@@ -41,7 +41,7 @@ struct EmitGeneratedCurationAction: Action {
         self.fileManager = fileManager
     }
 
-    mutating func perform(logHandle: LogHandle) throws -> ActionResult {
+    mutating func perform(logHandle: inout LogHandle) async throws -> ActionResult {
         let workspace = DocumentationWorkspace()
         let context = try DocumentationContext(dataProvider: workspace)
 

Sources/SwiftDocCUtilities/Action/Actions/IndexAction.swift

@@ -12,7 +12,7 @@ import Foundation
 import SwiftDocC
 
 /// An action that creates an index of a documentation bundle.
-public struct IndexAction: Action {
+public struct IndexAction: AsyncAction {
     let rootURL: URL
     let outputURL: URL
     let bundleIdentifier: String
@@ -22,8 +22,7 @@ public struct IndexAction: Action {
     private var dataProvider: LocalFileSystemDataProvider!
 
     /// Initializes the action with the given validated options, creates or uses the given action workspace & context.
-    public init(documentationBundleURL: URL, outputURL: URL, bundleIdentifier: String, diagnosticEngine: DiagnosticEngine = .init()) throws
-    {
+    public init(documentationBundleURL: URL, outputURL: URL, bundleIdentifier: String, diagnosticEngine: DiagnosticEngine = .init()) throws {
         // Initialize the action context.
         self.rootURL = documentationBundleURL
         self.outputURL = outputURL
@@ -35,7 +34,7 @@ public struct IndexAction: Action {
 
     /// Converts each eligible file from the source documentation bundle,
     /// saves the results in the given output alongside the template files.
-    mutating public func perform(logHandle: LogHandle) throws -> ActionResult {
+    mutating public func perform(logHandle: inout LogHandle) async throws -> ActionResult {
         let problems = try buildIndex()
         diagnosticEngine.emit(problems)
 

Sources/SwiftDocCUtilities/Action/Actions/Init/InitAction.swift

@@ -12,8 +12,8 @@ import Foundation
 import SwiftDocC
 
 /// An action that generates a documentation catalog from a template seed.
-public struct InitAction: Action {
-    
+public struct InitAction: AsyncAction {
+
     enum Error: DescribedError {
         case catalogAlreadyExists
         var errorDescription: String {
@@ -72,8 +72,7 @@ public struct InitAction: Action {
     /// Generates a documentation catalog from a catalog template.
     ///
     /// - Parameter logHandle: The file handle that the convert and preview actions will print debug messages to.
-    public mutating func perform(logHandle: SwiftDocC.LogHandle) throws -> ActionResult {
-        
+    public mutating func perform(logHandle: inout LogHandle) async throws -> ActionResult {
         let diagnosticEngine: DiagnosticEngine = DiagnosticEngine(treatWarningsAsErrors: false)
         diagnosticEngine.filterLevel = .warning
         diagnosticEngine.add(DiagnosticConsoleWriter(formattingOptions: []))