Merge pull request #36126 from ktoso/wip-task-tl

[Concurrency] Taks APIs fully implemented via ActiveTask internal TL

Author: DougGregor

include/swift/Runtime/Concurrency.h

@@ -107,6 +107,10 @@ void swift_task_dealloc(AsyncTask *task, void *ptr);
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
 void swift_task_cancel(AsyncTask *task);
 
+/// Get 'active' AsyncTask, depending on platform this may use thread local storage.
+SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
+AsyncTask* swift_task_get_active();
+
 /// Escalate the priority of a task and all of its child tasks.
 ///
 /// This can be called from any thread.

stdlib/public/Concurrency/Actor.cpp

@@ -134,6 +134,11 @@ SWIFT_RUNTIME_DECLARE_THREAD_LOCAL(
 
 } // end anonymous namespace
 
+AsyncTask*
+swift::swift_task_get_active() {
+  return ActiveTask::get();
+}
+
 void swift::swift_job_run(Job *job, ExecutorRef executor) {
   ExecutorTrackingInfo trackingInfo;
   trackingInfo.enterAndShadow(executor);

stdlib/public/Concurrency/Task.swift

@@ -49,10 +49,8 @@ extension Task {
   ///
   /// All functions available on the Task
   // TODO: once we can have async properties land make this computed property
-  @available(*, deprecated, message: "Please use Builtin.getCurrentAsyncTask() or Task.__unsafeCurrentAsync() until this function becomes implemented.")
-  public static func current(file: StaticString = #file, line: UInt = #line) async -> Task {
-    fatalError("Task.current() is not implemented yet!", file: file, line: line)
-    Task.unsafeCurrent!.task // !-safe, guaranteed to have a Task available within an async function.
+  public static func current() async -> Task {
+    Task.unsafeCurrent!.task // !-safe, we are guaranteed to have a task available within an async function
   }
 
 }
@@ -67,7 +65,6 @@ extension Task {
   ///
   /// - SeeAlso: `Task.Priority`
   /// - SeeAlso: `Task.priority`
-  @available(*, deprecated, message: "Not implemented yet, until unsafeCurrent is ready. Please use Task.__unsafeCurrentAsync().priority instead.")
   public static var currentPriority: Priority {
     Task.unsafeCurrent?.priority ?? Priority.default
   }
@@ -140,7 +137,7 @@ extension Task {
   /// Dropping a handle however means losing the ability to await on the task's result
   /// and losing the ability to cancel it.
   public struct Handle<Success, Failure: Error> {
-    private let _task: Builtin.NativeObject
+    internal let _task: Builtin.NativeObject
 
     internal init(_ task: Builtin.NativeObject) {
       self._task = task
@@ -510,22 +507,17 @@ extension Task {
   /// asynchronous function present in this functions call stack.
   ///
   /// The returned value must not be accessed from tasks other than the current one.
-  @available(*, deprecated, message: "Not implemented yet, use Builtin.getCurrentAsyncTask() or Task.___unsafeCurrentAsync() until this function is implemented.")
   public static var unsafeCurrent: UnsafeCurrentTask? {
-    // FIXME: rdar://70546948 implement this once getCurrentAsyncTask can be called from sync funcs
-    //    guard let _task = Builtin.getCurrentAsyncTask() else {
-    //      return nil
-    //    }
-    //    return UnsafeCurrentTask(_task)
-    fatalError("\(#function) is not implemented yet")
+    guard let _task = _getActiveAsyncTask() else {
+      return nil
+    }
+    // FIXME: This retain seems pretty wrong, however if we don't we WILL crash
+    //        with "destroying a task that never completed" in the task's destroy.
+    //        How do we solve this properly?
+    _swiftRetain(_task)
+    return UnsafeCurrentTask(_task)
   }
 
-  @available(*, deprecated, message: "This will be removed, and replaced by unsafeCurrent().", renamed: "unsafeCurrent()")
-  public static func __unsafeCurrentAsync() async -> UnsafeCurrentTask {
-    let task = Builtin.getCurrentAsyncTask()
-    _swiftRetain(task)
-    return UnsafeCurrentTask(task)
-  }
 }
 
 /// An *unsafe* 'current' task handle.
@@ -542,7 +534,7 @@ extension Task {
 /// and most certainly will break invariants in other places of the program
 /// actively running on this task.
 public struct UnsafeCurrentTask {
-  private let _task: Builtin.NativeObject
+  internal let _task: Builtin.NativeObject
 
   // May only be created by the standard library.
   internal init(_ task: Builtin.NativeObject) {
@@ -576,8 +568,24 @@ public struct UnsafeCurrentTask {
 
 }
 
+extension UnsafeCurrentTask: Hashable {
+  public func hash(into hasher: inout Hasher) {
+    UnsafeRawPointer(Builtin.bridgeToRawPointer(_task)).hash(into: &hasher)
+  }
+}
+
+extension UnsafeCurrentTask: Equatable {
+  public static func ==(lhs: Self, rhs: Self) -> Bool {
+    UnsafeRawPointer(Builtin.bridgeToRawPointer(lhs._task)) ==
+      UnsafeRawPointer(Builtin.bridgeToRawPointer(rhs._task))
+  }
+}
+
 // ==== Internal ---------------------------------------------------------------
 
+@_silgen_name("swift_task_get_active")
+func _getActiveAsyncTask() -> Builtin.NativeObject?
+
 @_silgen_name("swift_task_getJobFlags")
 func getJobFlags(_ task: Builtin.NativeObject) -> Task.JobFlags
 

stdlib/public/Concurrency/TaskLocal.swift

@@ -70,10 +70,13 @@ extension Task {
   /// - Returns: the value bound to the key, or its default value it if was not
   ///            bound in the current (or any parent) tasks.
   public static func local<Key>(_ keyPath: KeyPath<TaskLocalValues, Key>)
-    async -> Key.Value where Key: TaskLocalKey {
-    let task = Builtin.getCurrentAsyncTask()
+    -> Key.Value where Key: TaskLocalKey {
+    guard let task = Task.unsafeCurrent else {
+      return Key.defaultValue
+    }
 
-    let value = _taskLocalValueGet(task, keyType: Key.self, inheritance: Key.inherit.rawValue)
+    let value = _taskLocalValueGet(
+      task._task, keyType: Key.self, inheritance: Key.inherit.rawValue)
     guard let rawValue = value else {
       return Key.defaultValue
     }

test/Concurrency/Runtime/async_task_equals_hashCode.swift

@@ -1,27 +1,57 @@
-// RUN: %target-run-simple-swift(-Xfrontend -enable-experimental-concurrency -parse-as-library) | %FileCheck %s --dump-input=always
+// RUN: %target-run-simple-swift(-Xfrontend -enable-experimental-concurrency %import-libdispatch -parse-as-library) | %FileCheck %s --dump-input=always
+
 // REQUIRES: executable_test
 // REQUIRES: concurrency
-// REQUIRES: OS=macosx
-// REQUIRES: CPU=x86_64
 
-import Dispatch
 #if canImport(Darwin)
 import Darwin
 #elseif canImport(Glibc)
 import Glibc
 #endif
 
-@main struct Main {
-  static func main() async {
-    let one = await Task.__unsafeCurrentAsync().task // FIXME: replace with Task.current
-    let two = await Task.__unsafeCurrentAsync().task // FIXME: replace with Task.current
-    print("same equal: \(one == two)") // CHECK: same equal: true
-    print("hashes equal: \(one.hashValue == two.hashValue)") // CHECK: hashes equal: true
+func simple() async {
+  print("\(#function) -----------------------")
+  let one = await Task.current()
+  let two = await Task.current()
+  print("same equal: \(one == two)") // CHECK: same equal: true
+  print("hashes equal: \(one.hashValue == two.hashValue)") // CHECK: hashes equal: true
+
+  async let x = Task.current()
+  let three = await x
+
+  print("parent/child equal: \(three == two)") // CHECK: parent/child equal: false
+  print("parent/child hashes equal: \(three.hashValue == two.hashValue)") // CHECK: parent/child hashes equal: false
+}
+
+func unsafe() async {
+  print("\(#function) -----------------------")
+  let one = Task.unsafeCurrent!
+  let two = Task.unsafeCurrent!
+  print("unsafe same equal: \(one == two)") // CHECK: same equal: true
+  print("unsafe hashes equal: \(one.hashValue == two.hashValue)") // CHECK: hashes equal: true
 
-    async let x = Task.__unsafeCurrentAsync().task // FIXME: replace with Task.current
+  async let x = Task.unsafeCurrent!
+  let three = await x
 
-    let three = await x
-    print("parent/child equal: \(three == two)") // CHECK: parent/child equal: false
-    print("parent/child hashes equal: \(three.hashValue == two.hashValue)") // CHECK: parent/child hashes equal: false
+  print("unsafe parent/child equal: \(three == two)") // CHECK: parent/child equal: false
+  print("unsafe parent/child hashes equal: \(three.hashValue == two.hashValue)") // CHECK: parent/child hashes equal: false
+
+  print("unsafe.task parent/child equal: \(three.task == two.task)") // CHECK: parent/child equal: false
+  print("unsafe.task parent/child hashes equal: \(three.task.hashValue == two.task.hashValue)") // CHECK: parent/child hashes equal: false
+}
+
+func unsafeSync() {
+  print("\(#function) -----------------------")
+  let one = Task.unsafeCurrent!
+  let two = Task.unsafeCurrent!
+  print("unsafe same equal: \(one == two)") // CHECK: same equal: true
+  print("unsafe hashes equal: \(one.hashValue == two.hashValue)") // CHECK: hashes equal: true
+}
+
+@main struct Main {
+  static func main() async {
+    await simple()
+    await unsafe()
+    unsafeSync()
   }
 }

test/Concurrency/Runtime/async_task_handle_cancellation.swift

@@ -0,0 +1,21 @@
+// RUN: %target-run-simple-swift(-Xfrontend -enable-experimental-concurrency %import-libdispatch -parse-as-library) | %FileCheck %s --dump-input=always
+
+// REQUIRES: executable_test
+// REQUIRES: concurrency
+
+@main struct Main {
+  static func main() async {
+    let handle = Task.runDetached {
+      while (!Task.isCancelled) { // no need for await here, yay
+        print("waiting")
+      }
+
+      print("done")
+    }
+
+    handle.cancel()
+
+    // CHECK: done
+    await handle.get()
+  }
+}

test/Concurrency/Runtime/async_task_priority_current.swift

@@ -13,28 +13,26 @@ import WinSDK
 #error("Unsupported platform")
 #endif
 
-// FIXME: use `Task.currentPriority` once unsafeCurrent works in all these
-
 func test_detach() async {
-  let a1 = await Task.__unsafeCurrentAsync().task.priority
+  let a1 = Task.currentPriority
   print("a1: \(a1)") // CHECK: a1: default
 
   // Note: remember to detach using a higher priority, otherwise a lower one
   // might be escalated by the get() and we could see `default` in the detached
   // task.
   await Task.runDetached(priority: .userInitiated) {
-    let a2 = await Task.__unsafeCurrentAsync().task.priority
+    let a2 = Task.currentPriority
     print("a2: \(a2)") // CHECK: a2: userInitiated
   }.get()
 
-  let a3 = await Task.__unsafeCurrentAsync().task.priority
+  let a3 = Task.currentPriority
   print("a3: \(a3)") // CHECK: a3: default
 }
 
 func test_multiple_lo_indirectly_escalated() async {
   @concurrent
   func loopUntil(priority: Task.Priority) async {
-    while (await Task.__unsafeCurrentAsync().task.priority != priority) {
+    while (Task.currentPriority != priority) {
 #if os(Windows)
       Sleep(1)
 #else

test/Concurrency/Runtime/async_taskgroup_next_not_invoked_cancelAll.swift

@@ -16,7 +16,7 @@ func test_skipCallingNext_butInvokeCancelAll() async {
       await group.add { () async -> Int in
         sleep(1)
         print("  inside group.add { \(n) }")
-        let cancelled = await Task.__unsafeCurrentAsync().isCancelled
+        let cancelled = Task.isCancelled
         print("  inside group.add { \(n) } (canceled: \(cancelled))")
         return n
       }
@@ -25,7 +25,7 @@ func test_skipCallingNext_butInvokeCancelAll() async {
     group.cancelAll()
 
     // return immediately; the group should wait on the tasks anyway
-    let c = await Task.__unsafeCurrentAsync().isCancelled
+    let c = Task.isCancelled
     print("return immediately 0 (canceled: \(c))")
     return 0
   }

test/Concurrency/Runtime/async_taskgroup_next_not_invoked_without_cancelAll.swift

@@ -15,14 +15,14 @@ func test_skipCallingNext() async {
       print("group.add { \(n) }")
       await group.add { () async -> Int in
         sleep(1)
-        let c = await Task.__unsafeCurrentAsync().isCancelled
+        let c = Task.isCancelled
         print("  inside group.add { \(n) } (canceled: \(c))")
         return n
       }
     }
 
     // return immediately; the group should wait on the tasks anyway
-    let c = await Task.__unsafeCurrentAsync().isCancelled
+    let c = Task.isCancelled
     print("return immediately 0 (canceled: \(c))")
     return 0
   }

test/Concurrency/Runtime/async_taskgroup_throw_recover.swift

@@ -31,7 +31,7 @@ func test_taskGroup_throws() async {
         print("error caught in group: \(error)")
 
         await group.add { () async -> Int in
-          let c = await Task.__unsafeCurrentAsync().isCancelled
+          let c = Task.isCancelled
           print("task 3 (cancelled: \(c))")
           return 3
         }