docs: improve withTaskCancellationHandler docs (#70035)

* docs: improve withTaskCancellationHandler docs

* prefer using non-deprecated withCancHandler method

* Apply suggestions from code review

Co-authored-by: Alex Martini <amartini@apple.com>

---------

Co-authored-by: Alex Martini <amartini@apple.com>

Author: ktoso

stdlib/public/Concurrency/AsyncStreamBuffer.swift

@@ -253,12 +253,12 @@ extension AsyncStream {
     }
 
     func next() async -> Element? {
-      await withTaskCancellationHandler { [cancel] in
-        cancel()
-      } operation: {
+      await withTaskCancellationHandler {
         await withUnsafeContinuation {
           next($0)
         }
+      } onCancel: { [cancel] in
+        cancel()
       }
     }
 
@@ -512,12 +512,12 @@ extension AsyncThrowingStream {
     }
 
     func next() async throws -> Element? {
-      try await withTaskCancellationHandler { [cancel] in
-        cancel()
-      } operation: {
+      try await withTaskCancellationHandler {
         try await withUnsafeThrowingContinuation {
           next($0)
         }
+      } onCancel: { [cancel] in
+        cancel()
       }
     }
 

stdlib/public/Concurrency/SourceCompatibilityShims.swift

@@ -70,7 +70,7 @@ extension Task where Success == Never, Failure == Never {
     handler: @Sendable () -> Void,
     operation: () async throws -> T
   ) async rethrows -> T {
-    try await withTaskCancellationHandler(handler: handler, operation: operation)
+    try await withTaskCancellationHandler(operation: operation, onCancel: handler)
   }
 }
 

stdlib/public/Concurrency/TaskCancellation.swift

@@ -22,27 +22,51 @@ import Swift
 /// and reacting to it in that the cancellation handler is _always_ and
 /// _immediately_ invoked when the task is canceled. For example, even if the
 /// operation is running code that never checks for cancellation, a cancellation
-/// handler still runs and provides a chance to run some cleanup code.
+/// handler still runs and provides a chance to run some cleanup code:
+///
+/// ```
+/// await withTaskCancellationHandler {
+///   var sum = 0
+///   while condition {
+///     sum += 1
+///   }
+///   return sum
+/// } onCancel: {
+///   // This onCancel closure might execute concurrently with the operation.
+///   condition.cancel()
+/// }
+/// ```
+///
+/// ### Execution order and semantics
+/// The `operation` closure is always invoked, even when the
+/// `withTaskCancellationHandler(operation:onCancel:)` method is called from a task
+/// that was already cancelled.
+///
+/// When `withTaskCancellationHandler(operation:onCancel:)` is used in a task that has already been
+/// cancelled, the cancellation handler will be executed
+/// immediately before the `operation` closure gets to execute.
+///
+/// This allows the cancellation handler to set some external "cancelled" flag
+/// that the operation may be *atomically* checking for in order to avoid
+/// performing any actual work once the operation gets to run.
+///
+/// The `operation` closure executes on the calling execution context, and doesn't
+/// suspend or change execution context unless code contained within the closure
+/// does so. In other words, the potential suspension point of the
+/// `withTaskCancellationHandler(operation:onCancel:)` never suspends by itself before
+/// executing the operation.
+///
+/// If cancellation occurs while the operation is running, the cancellation
+/// handler executes *concurrently* with the operation.
+///
+/// ### Cancellation handlers and locks
 ///
 /// Cancellation handlers which acquire locks must take care to avoid deadlock.
 /// The cancellation handler may be invoked while holding internal locks
 /// associated with the task or other tasks.  Other operations on the task, such
 /// as resuming a continuation, may acquire these same internal locks.
 /// Therefore, if a cancellation handler must acquire a lock, other code should
 /// not cancel tasks or resume continuations while holding that lock.
-///
-/// Doesn't check for cancellation, and always executes the passed `operation`.
-///
-/// The `operation` executes on the calling execution context and does not suspend by itself,
-/// unless the code contained within the closure does. If cancellation occurs while the
-/// operation is running, the cancellation `handler` will execute *concurrently* with the `operation`.
-///
-/// ### Already cancelled tasks
-/// When `withTaskCancellationHandler` is used in a `Task` that has already been cancelled,
-/// the `onCancel` cancellation ``handler`` will be executed immediately before operation gets
-/// to execute. This allows the cancellation handler to set some external "cancelled" flag that the
-/// operation may be *atomically* checking for in order to avoid performing any actual work once
-/// the operation gets to run.
 @_unsafeInheritExecutor // the operation runs on the same executor as we start out with
 @available(SwiftStdlib 5.1, *)
 @backDeployed(before: SwiftStdlib 5.8)