Merge pull request swiftlang#81205 from mikeash/fix-statusRecordLock

[Concurrency] Fix alreadyLocked in withStatusRecordLock.

Author: mikeash

stdlib/public/Concurrency/TaskPrivate.h

@@ -829,11 +829,12 @@ struct AsyncTask::PrivateStorage {
     // result of the task
     auto oldStatus = task->_private()._status().load(std::memory_order_relaxed);
     while (true) {
-      // Task is completing
       assert(oldStatus.getInnermostRecord() == NULL &&
              "Status records should have been removed by this time!");
-      assert(!oldStatus.isStatusRecordLocked() &&
-             "Task is completing, cannot be locked anymore!");
+
+      // Don't assert !isStatusRecordLocked(). It can be legitimately true here,
+      // for example if another thread is canceling this task right as it
+      // completes.
 
       assert(oldStatus.isRunning());
 

stdlib/public/Concurrency/TaskStatus.cpp

@@ -67,7 +67,19 @@ static void withStatusRecordLock(
   // see that the is-locked bit is now set, and then wait for the lock.
   task->_private().statusLock.lock();
 
-  bool alreadyLocked = status.isStatusRecordLocked();
+  bool alreadyLocked = false;
+
+  // `status` was loaded before we acquired the lock. If its is-locked bit is
+  // not set, then we know that this thread doesn't already hold the lock.
+  // However, if the is-locked bit is set, then we don't know if this thread
+  // held the lock or another thread did. In that case, we reload the status
+  // after acquiring the lock. If the reloaded status still has the is-locked
+  // bit set, then we know it's this thread. If it doesn't, then we know it was
+  // a different thread.
+  if (status.isStatusRecordLocked()) {
+    status = task->_private()._status().load(std::memory_order_relaxed);
+    alreadyLocked = status.isStatusRecordLocked();
+  }
 
   // If it's already locked then this thread is the thread that locked it, and
   // we can leave that bit alone here.

test/Concurrency/Runtime/cancellation_handler_concurrent.swift

@@ -0,0 +1,33 @@
+// RUN: %target-run-simple-swift( -target %target-swift-5.1-abi-triple %import-libdispatch)
+// REQUIRES: concurrency
+// REQUIRES: executable_test
+
+// REQUIRES: concurrency_runtime
+// UNSUPPORTED: freestanding
+
+func recurseABunch(_ call: () async throws -> Void, n: Int = 100) async throws {
+  try await withTaskCancellationHandler {
+    if n == 0 {
+      try await call()
+      return
+    }
+
+    try await recurseABunch(call, n: n - 1)
+  } onCancel: {
+  }
+}
+
+for _ in 0..<100_000 {
+    let task: Task<Void, any Error> = Task {
+        try await Task.sleep(nanoseconds: UInt64.random(in: 0..<1_000))
+        try await recurseABunch() {
+          await withTaskCancellationHandler {
+          } onCancel: {
+          }
+        }
+    }
+    Task {
+        try await Task.sleep(nanoseconds: UInt64.random(in: 0..<1_000))
+        task.cancel()
+    }
+}