Merge pull request swiftlang#36689 from apple/tsan-task-groups

Author: swift-ci

stdlib/public/Concurrency/TaskGroup.cpp

@@ -611,6 +611,8 @@ void TaskGroupImpl::offer(AsyncTask *completedTask, AsyncContext *context) {
 
         fillGroupNextResult(waitingContext, result);
 
+        _swift_tsan_acquire(static_cast<Job *>(waitingTask));
+
         // TODO: allow the caller to suggest an executor
         swift_task_enqueueGlobal(waitingTask);
         return;
@@ -740,6 +742,7 @@ PollResult TaskGroupImpl::poll(AsyncTask *waitingTask) {
           result.status = PollStatus::Success;
           result.storage = futureFragment->getStoragePtr();
           assert(result.retainedTask && "polled a task, it must be not null");
+          _swift_tsan_acquire(static_cast<Job *>(result.retainedTask));
           mutex.unlock(); // TODO: remove fragment lock, and use status for synchronization
           return result;
 
@@ -749,6 +752,7 @@ PollResult TaskGroupImpl::poll(AsyncTask *waitingTask) {
           result.storage =
               reinterpret_cast<OpaqueValue *>(futureFragment->getError());
           assert(result.retainedTask && "polled a task, it must be not null");
+          _swift_tsan_acquire(static_cast<Job *>(result.retainedTask));
           mutex.unlock(); // TODO: remove fragment lock, and use status for synchronization
           return result;
 
@@ -765,6 +769,7 @@ PollResult TaskGroupImpl::poll(AsyncTask *waitingTask) {
 
   // ==== 3) Add to wait queue -------------------------------------------------
   assert(assumed.readyTasks() == 0);
+  _swift_tsan_release(static_cast<Job *>(waitingTask));
   while (true) {
     // Put the waiting task at the beginning of the wait queue.
     if (waitQueue.compare_exchange_weak(

test/Sanitizers/tsan/async_taskgroup_next.swift

@@ -0,0 +1,61 @@
+// RUN: %target-run-simple-swift(-Xfrontend -enable-experimental-concurrency %import-libdispatch -parse-as-library -sanitize=thread)
+
+// REQUIRES: executable_test
+// REQUIRES: concurrency
+// REQUIRES: libdispatch
+// REQUIRES: tsan_runtime
+
+var scratchBuffer: UnsafeMutableBufferPointer<Int> = .allocate(capacity: 1000)
+
+@available(macOS 9999, iOS 9999, watchOS 9999, tvOS 9999, *)
+func completeFastOrSlow(n: Int) async -> Int {
+  if n % 2 == 0 {
+    await Task.sleep(2_000_000_000)
+  }
+  assert(scratchBuffer[n] == 6)
+  scratchBuffer[n] = 7
+  return n
+}
+
+@available(macOS 9999, iOS 9999, watchOS 9999, tvOS 9999, *)
+func test_sum_nextOnCompletedOrPending() async {
+  scratchBuffer.initialize(repeating: 0)
+
+  let numbers = 0..<1000
+  let expected = 499_500
+
+  let sum = await withTaskGroup(of: Int.self) { (group) async -> Int in
+    for n in numbers {
+      scratchBuffer[n] = 6
+      await group.spawn {
+        let res = await completeFastOrSlow(n: n)
+        return res
+      }
+    }
+
+    // We want to await on completed and pending child tasks.  This gives the
+    // fast tasks some time to complete before we call group.next().
+    await Task.sleep(1_000_000_000)
+
+    var sum = 0
+    while let r = try! await group.next() {
+      assert(scratchBuffer[r] == 7)
+      sum += r
+    }
+
+    assert(group.isEmpty, "Group must be empty after we consumed all tasks")
+
+    print("task group returning: \(sum)")
+    return sum
+  }
+
+  print("result: \(sum)")
+  assert(sum == expected, "Expected: \(expected), got: \(sum)")
+}
+
+@available(macOS 9999, iOS 9999, watchOS 9999, tvOS 9999, *)
+@main struct Main {
+  static func main() async {
+    await test_sum_nextOnCompletedOrPending()
+  }
+}