[Concurrency] Fix alreadyLocked in withStatusRecordLock.

If the preloaded status is locked, then we need to reload it in order to distinguish between the current thread holding the lock and another thread holding the lock. Without this, if another thread holds the lock, then we won't set the is-locked bit. We'll still actually hold the lock, but other threads may perform operations locklessly if the bit is not set, which can cause a crash. By reloading status in that case, we ensure that the bit is always set correctly.

This manifested as crashes in task cancellation but could cause other task-related issues as well.

Also remove an assert of !isStatusRecordLocked() in AsyncTask::complete(). We allow other threads to access tasks and take the lock for things like cancellation, so the lock may legitimately be held at that point.

rdar://150327908

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()
+    }
+}