cancellation handler (#4211)

* Revert "Revert "cancellation handler (#4173)""

This reverts commit c948a86.

* windows fix

Author: tomerd

Sources/Basics/Archiver+Zip.swift

@@ -12,54 +12,83 @@ import TSCBasic
 import Dispatch
 
 /// An `Archiver` that handles ZIP archives using the command-line `zip` and `unzip` tools.
-public struct ZipArchiver: Archiver {
+public struct ZipArchiver: Archiver, Cancellable {
     public var supportedExtensions: Set<String> { ["zip"] }
 
     /// The file-system implementation used for various file-system operations and checks.
     private let fileSystem: FileSystem
 
+    /// Helper for cancelling in-fligh requests
+    private let cancellator: Cancellator
+
     /// Creates a `ZipArchiver`.
     ///
     /// - Parameters:
     ///   - fileSystem: The file-system to used by the `ZipArchiver`.
     public init(fileSystem: FileSystem) {
         self.fileSystem = fileSystem
+        self.cancellator = Cancellator(observabilityScope: .none)
     }
 
     public func extract(
         from archivePath: AbsolutePath,
         to destinationPath: AbsolutePath,
         completion: @escaping (Result<Void, Error>) -> Void
     ) {
-        guard fileSystem.exists(archivePath) else {
-            completion(.failure(FileSystemError(.noEntry, archivePath)))
-            return
-        }
+        do {
+            guard self.fileSystem.exists(archivePath) else {
+                throw FileSystemError(.noEntry, archivePath)
+            }
 
-        guard fileSystem.isDirectory(destinationPath) else {
-            completion(.failure(FileSystemError(.notDirectory, destinationPath)))
-            return
-        }
+            guard self.fileSystem.isDirectory(destinationPath) else {
+                throw FileSystemError(.notDirectory, destinationPath)
+            }
+
+            let process = Process(arguments: ["unzip", archivePath.pathString, "-d", destinationPath.pathString])
+            guard let registrationKey = self.cancellator.register(process) else {
+                throw StringError("cancellation")
+            }
 
-        Process.popen(arguments: ["unzip", archivePath.pathString, "-d", destinationPath.pathString], queue: .sharedConcurrent) { result in
-            completion(result.tryMap { processResult in
-                guard processResult.exitStatus == .terminated(code: 0) else {
-                    throw try StringError(processResult.utf8stderrOutput())
-                }
-            })
+            DispatchQueue.sharedConcurrent.async {
+                defer { self.cancellator.deregister(registrationKey) }
+                completion(.init(catching: {
+                    try process.launch()
+                    let processResult = try process.waitUntilExit()
+                    guard processResult.exitStatus == .terminated(code: 0) else {
+                        throw try StringError(processResult.utf8stderrOutput())
+                    }
+                }))
+            }
+        } catch {
+            return completion(.failure(error))
         }
     }
 
     public func validate(path: AbsolutePath, completion: @escaping (Result<Bool, Error>) -> Void) {
-        guard fileSystem.exists(path) else {
-            completion(.failure(FileSystemError(.noEntry, path)))
-            return
-        }
+        do {
+            guard self.fileSystem.exists(path) else {
+                throw FileSystemError(.noEntry, path)
+            }
+
+            let process = Process(arguments: ["unzip", "-t", path.pathString])
+            guard let registrationKey = self.cancellator.register(process) else {
+                throw StringError("cancellation")
+            }
 
-        Process.popen(arguments: ["unzip", "-t", path.pathString], queue: .sharedConcurrent) { result in
-            completion(result.tryMap { processResult in
-                return processResult.exitStatus == .terminated(code: 0)
-            })
+            DispatchQueue.sharedConcurrent.async {
+                defer { self.cancellator.deregister(registrationKey) }
+                completion(.init(catching: {
+                    try process.launch()
+                    let processResult = try process.waitUntilExit()
+                    return processResult.exitStatus == .terminated(code: 0)
+                }))
+            }
+        } catch {
+            return completion(.failure(error))
         }
     }
+
+    public func cancel(deadline: DispatchTime) throws {
+        try self.cancellator.cancel(deadline: deadline)
+    }
 }

Sources/Basics/CMakeLists.txt

@@ -11,6 +11,7 @@ add_library(Basics
   Archiver+Zip.swift
   AuthorizationProvider.swift
   ByteString+Extensions.swift
+  Cancellator.swift
   ConcurrencyHelpers.swift
   Dictionary+Extensions.swift
   DispatchTimeInterval+Extensions.swift

Sources/Basics/Cancellator.swift

@@ -0,0 +1,135 @@
+/*
+ This source file is part of the Swift.org open source project
+
+ Copyright (c) 2022 Apple Inc. and the Swift project authors
+ Licensed under Apache License v2.0 with Runtime Library Exception
+
+ See http://swift.org/LICENSE.txt for license information
+ See http://swift.org/CONTRIBUTORS.txt for Swift project authors
+ */
+
+import Dispatch
+import Foundation
+import TSCBasic
+
+public typealias CancellationHandler = (DispatchTime) throws -> Void
+
+public class Cancellator: Cancellable {
+    public typealias RegistrationKey = String
+
+    private let observabilityScope: ObservabilityScope?
+    private let registry = ThreadSafeKeyValueStore<String, (name: String, handler: CancellationHandler)>()
+    private let cancelationQueue = DispatchQueue(label: "org.swift.swiftpm.cancellator", qos: .userInteractive, attributes: .concurrent)
+    private let cancelling = ThreadSafeBox<Bool>(false)
+
+    public init(observabilityScope: ObservabilityScope?) {
+        self.observabilityScope = observabilityScope
+    }
+
+    @discardableResult
+    public func register(name: String, handler: @escaping CancellationHandler) -> RegistrationKey? {
+        if self.cancelling.get(default: false) {
+            self.observabilityScope?.emit(debug: "not registering '\(name)' with terminator, termination in progress")
+            return .none
+        }
+        let key = UUID().uuidString
+        self.observabilityScope?.emit(debug: "registering '\(name)' with terminator")
+        self.registry[key] = (name: name, handler: handler)
+        return key
+    }
+
+    @discardableResult
+    public func register(name: String, handler: Cancellable) -> RegistrationKey? {
+        self.register(name: name, handler: handler.cancel(deadline:))
+    }
+
+    @discardableResult
+    public func register(name: String, handler: @escaping () throws -> Void) -> RegistrationKey? {
+        self.register(name: name, handler: { _ in try handler() })
+    }
+
+    public func register(_ process: TSCBasic.Process) -> RegistrationKey? {
+        self.register(name: "\(process.arguments.joined(separator: " "))", handler:  process.terminate)
+    }
+
+    public func deregister(_ key: RegistrationKey) {
+        self.registry[key] = nil
+    }
+
+    public func cancel(deadline: DispatchTime) throws -> Void {
+        self._cancel(deadline: deadline)
+    }
+
+    // marked internal for testing
+    @discardableResult
+    internal func _cancel(deadline: DispatchTime? = .none)-> Int {
+        self.cancelling.put(true)
+
+        self.observabilityScope?.emit(info: "starting cancellation cycle with \(self.registry.count) cancellation handlers registered")
+
+        let deadline = deadline ?? .now() + .seconds(30)
+        // deadline for individual handlers set slightly before overall deadline
+        let delta: DispatchTimeInterval = .nanoseconds(abs(deadline.distance(to: .now()).nanoseconds() ?? 0)  / 5)
+        let handlersDeadline = deadline - delta
+
+        let cancellationHandlers = self.registry.get()
+        let cancelled = ThreadSafeArrayStore<String>()
+        let group = DispatchGroup()
+        for (_, (name, handler)) in cancellationHandlers {
+            self.cancelationQueue.async(group: group) {
+                do {
+                    self.observabilityScope?.emit(debug: "cancelling '\(name)'")
+                    try handler(handlersDeadline)
+                    cancelled.append(name)
+                } catch {
+                    self.observabilityScope?.emit(warning: "failed cancelling '\(name)': \(error)")
+                }
+            }
+        }
+
+        if case .timedOut = group.wait(timeout: deadline) {
+            self.observabilityScope?.emit(warning: "timeout waiting for cancellation with \(cancellationHandlers.count - cancelled.count) cancellation handlers remaining")
+        } else {
+            self.observabilityScope?.emit(info: "cancellation cycle completed successfully")
+        }
+
+        self.cancelling.put(false)
+
+        return cancelled.count
+    }
+}
+
+public protocol Cancellable {
+    func cancel(deadline: DispatchTime) throws -> Void
+}
+
+public struct CancellationError: Error, CustomStringConvertible {
+    public let description = "Operation cancelled"
+
+    public init() {}
+}
+
+extension TSCBasic.Process {
+    fileprivate func terminate(timeout: DispatchTime) {
+        // send graceful shutdown signal
+        self.signal(SIGINT)
+
+        // start a thread to see if we need to terminate more forcibly
+        let forceKillSemaphore = DispatchSemaphore(value: 0)
+        let forceKillThread = TSCBasic.Thread {
+            if case .timedOut = forceKillSemaphore.wait(timeout: timeout) {
+                // send a force-kill signal
+                #if os(Windows)
+                self.signal(SIGTERM)
+                #else
+                self.signal(SIGKILL)
+                #endif
+            }
+        }
+        forceKillThread.start()
+        _ = try? self.waitUntilExit()
+        forceKillSemaphore.signal() // let the force-kill thread know we do not need it any more
+        // join the force-kill thread thread so we don't exit before everything terminates
+        forceKillThread.join()
+    }
+}

Sources/Basics/ConcurrencyHelpers.swift

@@ -55,9 +55,12 @@ public final class ThreadSafeKeyValueStore<Key, Value> where Key: Hashable {
         }
     }
 
-    public func clear() {
+    @discardableResult
+    public func clear() -> [Key: Value] {
         self.lock.withLock {
+            let underlying = self.underlying
             self.underlying.removeAll()
+            return underlying
         }
     }
 
@@ -113,9 +116,12 @@ public final class ThreadSafeArrayStore<Value> {
         }
     }
 
-    public func clear() {
+    @discardableResult
+    public func clear() -> [Value] {
         self.lock.withLock {
-            self.underlying = []
+            let underlying = self.underlying
+            self.underlying.removeAll()
+            return underlying
         }
     }
 

Sources/Basics/DispatchTimeInterval+Extensions.swift

@@ -27,6 +27,21 @@ extension DispatchTimeInterval {
         }
     }
 
+    public func nanoseconds() -> Int? {
+        switch self {
+        case .seconds(let value):
+            return value.multipliedReportingOverflow(by: 1_000_000_000).partialValue
+        case .milliseconds(let value):
+            return value.multipliedReportingOverflow(by: 1_000_000).partialValue
+        case .microseconds(let value):
+            return value.multipliedReportingOverflow(by: 1000).partialValue
+        case .nanoseconds(let value):
+            return value
+        default:
+            return nil
+        }
+    }
+
     public func milliseconds() -> Int? {
         switch self {
         case .seconds(let value):
@@ -81,7 +96,7 @@ extension DispatchTimeInterval {
 #if os(Linux) || os(Windows) || os(Android) || os(OpenBSD)
 extension DispatchTime {
     public func distance(to: DispatchTime) -> DispatchTimeInterval {
-        let duration = to.uptimeNanoseconds - self.uptimeNanoseconds
+        let duration = to.uptimeNanoseconds.subtractingReportingOverflow(self.uptimeNanoseconds).partialValue
         return .nanoseconds(duration >= Int.max ? Int.max : Int(duration))
     }
 }