Use trylock to eliminate the remaining race condition in Test.cancel(). (#1415)

This PR fixes the race condition in `Test.cancel()` that could occur if
an unstructured task, created from within a test's task, called
`Test.cancel()` at just the right moment. The order of events for the
race is:

- Unstructured task is created and inherits task-locals including the
reference to the test's unsafe current task;
- Test's task starts tearing down;
- Unstructured task calls `takeUnsafeCurrentTask()` and gets a reference
to the unsafe current task;
- Test's task finishes tearing down;
- Unstructured task calls `UnsafeCurrentTask.cancel()`.

The fix is to use `trylock` semantics when cancelling the unsafe current
task. If the test's task is still alive, the task is cancelled while the
lock is held, which will block the test's task from being torn down as
it has a lock-guarded call to clear the unsafe current task reference.
If the test's task is no longer alive, the reference is already `nil` by
the time the unstructured task acquires the lock and it bails early. If
we recursively call `cancel()` (which can happen via the
concurrency-level cancellation handler), the `trylock` means we won't
acquire the lock a second time, so we won't end up deadlocking or
aborting (which is what prevents calling `cancel()` while holding the
lock in the current implementation).

It is possible for `cancel()` to trigger user code, especially if the
user has set up a cancellation handler, but there is no code path that
can then lead to a deadlock because the only user-accessible calls that
might touch this lock use `trylock`.

I hope some part of that made sense.

### Checklist:

- [x] Code and documentation should follow the style of the [Style
Guide](https://github.com/apple/swift-testing/blob/main/Documentation/StyleGuide.md).
- [x] If public symbols are renamed or modified, DocC references should
be updated.

Author: grynspan

Sources/Testing/Support/Locked.swift

@@ -26,6 +26,14 @@ struct Locked<T> {
   /// A type providing storage for the underlying lock and wrapped value.
 #if SWT_TARGET_OS_APPLE && canImport(os)
   private typealias _Storage = ManagedBuffer<T, os_unfair_lock_s>
+#elseif !SWT_FIXED_85448 && (os(Linux) || os(Android))
+  private final class _Storage: ManagedBuffer<T, pthread_mutex_t> {
+    deinit {
+      withUnsafeMutablePointerToElements { lock in
+        _ = pthread_mutex_destroy(lock)
+      }
+    }
+  }
 #else
   private final class _Storage {
     let mutex: Mutex<T>
@@ -49,6 +57,11 @@ extension Locked: RawRepresentable {
     _storage.withUnsafeMutablePointerToElements { lock in
       lock.initialize(to: .init())
     }
+#elseif !SWT_FIXED_85448 && (os(Linux) || os(Android))
+    _storage = _Storage.create(minimumCapacity: 1, makingHeaderWith: { _ in rawValue }) as! _Storage
+    _storage.withUnsafeMutablePointerToElements { lock in
+      _ = pthread_mutex_init(lock, nil)
+    }
 #else
     nonisolated(unsafe) let rawValue = rawValue
     _storage = _Storage(rawValue)
@@ -77,20 +90,72 @@ extension Locked {
   /// synchronous caller. Wherever possible, use actor isolation or other Swift
   /// concurrency tools.
   func withLock<R>(_ body: (inout T) throws -> sending R) rethrows -> sending R where R: ~Copyable {
+    nonisolated(unsafe) let result: R
 #if SWT_TARGET_OS_APPLE && canImport(os)
-    nonisolated(unsafe) let result = try _storage.withUnsafeMutablePointers { rawValue, lock in
+    result = try _storage.withUnsafeMutablePointers { rawValue, lock in
       os_unfair_lock_lock(lock)
       defer {
         os_unfair_lock_unlock(lock)
       }
       return try body(&rawValue.pointee)
     }
-    return result
+#elseif !SWT_FIXED_85448 && (os(Linux) || os(Android))
+     result = try _storage.withUnsafeMutablePointers { rawValue, lock in
+      pthread_mutex_lock(lock)
+      defer {
+        pthread_mutex_unlock(lock)
+      }
+      return try body(&rawValue.pointee)
+    }
 #else
-    try _storage.mutex.withLock { rawValue in
+    result = try _storage.mutex.withLock { rawValue in
       try body(&rawValue)
     }
 #endif
+    return result
+  }
+
+  /// Try to acquire the lock and invoke a function while it is held.
+  ///
+  /// - Parameters:
+  ///   - body: A closure to invoke while the lock is held.
+  ///
+  /// - Returns: Whatever is returned by `body`, or `nil` if the lock could not
+  ///   be acquired.
+  ///
+  /// - Throws: Whatever is thrown by `body`.
+  ///
+  /// This function can be used to synchronize access to shared data from a
+  /// synchronous caller. Wherever possible, use actor isolation or other Swift
+  /// concurrency tools.
+  func withLockIfAvailable<R>(_ body: (inout T) throws -> sending R) rethrows -> sending R? where R: ~Copyable {
+    nonisolated(unsafe) let result: R?
+#if SWT_TARGET_OS_APPLE && canImport(os)
+    result = try _storage.withUnsafeMutablePointers { rawValue, lock in
+      guard os_unfair_lock_trylock(lock) else {
+        return nil
+      }
+      defer {
+        os_unfair_lock_unlock(lock)
+      }
+      return try body(&rawValue.pointee)
+    }
+#elseif !SWT_FIXED_85448 && (os(Linux) || os(Android))
+    result = try _storage.withUnsafeMutablePointers { rawValue, lock in
+      guard 0 == pthread_mutex_trylock(lock) else {
+        return nil
+      }
+      defer {
+        pthread_mutex_unlock(lock)
+      }
+      return try body(&rawValue.pointee)
+    }
+#else
+    result = try _storage.mutex.withLockIfAvailable { rawValue in
+      return try body(&rawValue)
+    }
+#endif
+    return result
   }
 }
 

Sources/Testing/Test+Cancellation.swift

@@ -25,9 +25,8 @@ protocol TestCancellable: Sendable {
 
 // MARK: - Tracking the current task
 
-/// A structure describing a reference to a task that is associated with some
-/// ``TestCancellable`` value.
-private struct _TaskReference: Sendable {
+/// A structure that is able to cancel a task.
+private struct _TaskCanceller: Sendable {
   /// The unsafe underlying reference to the associated task.
   private nonisolated(unsafe) var _unsafeCurrentTask = Locked<UnsafeCurrentTask?>()
 
@@ -45,25 +44,46 @@ private struct _TaskReference: Sendable {
     _unsafeCurrentTask = withUnsafeCurrentTask { Locked(rawValue: $0) }
   }
 
-  /// Take this instance's reference to its associated task.
-  ///
-  /// - Returns: An `UnsafeCurrentTask` instance, or `nil` if it was already
-  ///   taken or if it was never available.
-  ///
-  /// This function consumes the reference to the task. After the first call,
-  /// subsequent calls on the same instance return `nil`.
-  func takeUnsafeCurrentTask() -> UnsafeCurrentTask? {
+  /// Clear this instance's reference to its associated task without first
+  /// cancelling it.
+  func clear() {
     _unsafeCurrentTask.withLock { unsafeCurrentTask in
-      let result = unsafeCurrentTask
       unsafeCurrentTask = nil
-      return result
     }
   }
+
+  /// Cancel this instance's associated task and clear the reference to it.
+  ///
+  /// - Returns: Whether or not this instance's task was cancelled.
+  ///
+  /// After the first call to this function _starts_, subsequent calls on the
+  /// same instance return `false`. In other words, if another thread calls this
+  /// function before it has returned (or the same thread calls it recursively),
+  /// it returns `false` without cancelling the task a second time.
+  func cancel(with skipInfo: SkipInfo) -> Bool {
+    // trylock means a recursive call to this function won't ruin our day, nor
+    // should interleaving locks.
+    _unsafeCurrentTask.withLockIfAvailable { unsafeCurrentTask in
+      defer {
+        unsafeCurrentTask = nil
+      }
+      if let unsafeCurrentTask {
+        // The task is still valid, so we'll cancel it.
+        $_currentSkipInfo.withValue(skipInfo) {
+          unsafeCurrentTask.cancel()
+        }
+        return true
+      }
+
+      // The task has already been cancelled and/or cleared.
+      return false
+    } ?? false
+  }
 }
 
-/// A dictionary of tracked tasks, keyed by types that conform to
+/// A dictionary of cancellable tasks keyed by types that conform to
 /// ``TestCancellable``.
-@TaskLocal private var _currentTaskReferences = [ObjectIdentifier: _TaskReference]()
+@TaskLocal private var _currentTaskCancellers = [ObjectIdentifier: _TaskCanceller]()
 
 /// The instance of ``SkipInfo`` to propagate to children of the current task.
 ///
@@ -87,16 +107,15 @@ extension TestCancellable {
   /// the current task, test, or test case is cancelled, it records a
   /// corresponding cancellation event.
   func withCancellationHandling<R>(_ body: () async throws -> R) async rethrows -> R {
-    let taskReference = _TaskReference()
-    var currentTaskReferences = _currentTaskReferences
-    currentTaskReferences[ObjectIdentifier(Self.self)] = taskReference
-    return try await $_currentTaskReferences.withValue(currentTaskReferences) {
-      // Before returning, explicitly clear the stored task. This minimizes
-      // the potential race condition that can occur if test code creates an
-      // unstructured task and calls `Test.cancel()` in it after the test body
-      // has finished.
+    let taskCanceller = _TaskCanceller()
+    var currentTaskCancellers = _currentTaskCancellers
+    currentTaskCancellers[ObjectIdentifier(Self.self)] = taskCanceller
+    return try await $_currentTaskCancellers.withValue(currentTaskCancellers) {
+      // Before returning, explicitly clear the stored task so that an
+      // unstructured task that inherits the task local isn't able to
+      // accidentally cancel the task after it has been deallocated.
       defer {
-        _ = taskReference.takeUnsafeCurrentTask()
+        taskCanceller.clear()
       }
 
       return try await withTaskCancellationHandler {
@@ -121,18 +140,16 @@ extension TestCancellable {
 ///   - testAndTestCase: The test and test case to use when posting an event.
 ///   - skipInfo: Information about the cancellation event.
 private func _cancel<T>(_ cancellableValue: T?, for testAndTestCase: (Test?, Test.Case?), skipInfo: SkipInfo) where T: TestCancellable {
-  if cancellableValue != nil {
-    // If the current test case is still running, take its task property (which
-    // signals to subsequent callers that it has been cancelled.)
-    let task = _currentTaskReferences[ObjectIdentifier(T.self)]?.takeUnsafeCurrentTask()
-
-    // If we just cancelled the current test case's task, post a corresponding
-    // event with the relevant skip info.
-    if let task {
-      $_currentSkipInfo.withValue(skipInfo) {
-        task.cancel()
-      }
+  if cancellableValue != nil, let taskCanceller = _currentTaskCancellers[ObjectIdentifier(T.self)] {
+    // Try to cancel the task associated with `T`, if any. If we succeed, post a
+    // corresponding event with the relevant skip info. If we fail, we still
+    // attempt to cancel the current *task* in order to honor our API contract.
+    if taskCanceller.cancel(with: skipInfo) {
       Event.post(T.makeCancelledEventKind(with: skipInfo), for: testAndTestCase)
+    } else {
+      withUnsafeCurrentTask { task in
+        task?.cancel()
+      }
     }
   } else {
     // The current task isn't associated with a test/case, so just cancel the