Concurrency: Factor atomic operations in Task.sleep() into a Sendable wrapper.

When building the Swift standard library, the compiler warns:

```
warning: capture of 'wordPtr' with non-sendable type 'UnsafeMutablePointer<Builtin.Word>' in a `@Sendable` closure
```

This diagnostic will become an error in the Swift 6 language mode, so it needs
to be addressed. This PR factors the atomic operations used by the
implementation of `Task.sleep()` into an `@unchecked Sendable` wrapper in order
to convince the compiler that these operations are thread-safe. As an added
benefit, the code is more readable when the atomic operatios are abstracted
away. This refactor intentionally translates the existing implementation into a
wrapper as faithfully as possible and does not attempt to improve on any other
aspects of the implementation, such as the unsafe manual memory allocation and
deallocation.

Author: tshortli

stdlib/public/Concurrency/TaskSleep.swift

@@ -105,25 +105,55 @@ extension Task where Success == Never, Failure == Never {
     }
   }
 
+  /// A simple wrapper for a pointer to heap allocated storage of a `SleepState`
+  /// value. This wrapper is `Sendable` because it facilitates atomic load and
+  /// exchange operations on the underlying storage. However, this wrapper is also
+  /// _unsafe_ because the owner must manually deallocate the token once it is no
+  /// longer needed.
+  struct UnsafeSleepStateToken: @unchecked Sendable {
+    let wordPtr: UnsafeMutablePointer<Builtin.Word>
+
+    /// Allocates the underlying storage and sets the value to `.notStarted`.
+    init() {
+      wordPtr = .allocate(capacity: 1)
+      Builtin.atomicstore_seqcst_Word(
+          wordPtr._rawValue, SleepState.notStarted.word._builtinWordValue)
+    }
+
+    /// Atomically loads the current state.
+    func load() -> SleepState {
+      return SleepState(word: Builtin.atomicload_seqcst_Word(wordPtr._rawValue))
+    }
+
+    /// Attempts to atomically set the stored value to `desired` if the current
+    /// value is equal to `expected`. Returns true if the exchange was successful.
+    func exchange(expected: SleepState, desired: SleepState) -> Bool {
+      let (_, won) = Builtin.cmpxchg_seqcst_seqcst_Word(
+          wordPtr._rawValue,
+          expected.word._builtinWordValue,
+          desired.word._builtinWordValue)
+      return Bool(_builtinBooleanLiteral: won)
+    }
+
+    /// Deallocates the underlying storage.
+    func deallocate() {
+      wordPtr.deallocate()
+    }
+  }
+
   /// Called when the sleep(nanoseconds:) operation woke up without being
   /// canceled.
-  static func onSleepWake(
-      _ wordPtr: UnsafeMutablePointer<Builtin.Word>
-  ) {
+  static func onSleepWake(_ token: UnsafeSleepStateToken) {
     while true {
-      let state = SleepState(loading: wordPtr)
+      let state = token.load()
       switch state {
       case .notStarted:
         fatalError("Cannot wake before we even started")
 
       case .activeContinuation(let continuation):
         // We have an active continuation, so try to transition to the
         // "finished" state.
-        let (_, won) = Builtin.cmpxchg_seqcst_seqcst_Word(
-            wordPtr._rawValue,
-            state.word._builtinWordValue,
-            SleepState.finished.word._builtinWordValue)
-        if Bool(_builtinBooleanLiteral: won) {
+        if token.exchange(expected: state, desired: .finished) {
           // The sleep finished, so invoke the continuation: we're done.
           continuation.resume()
           return
@@ -137,9 +167,9 @@ extension Task where Success == Never, Failure == Never {
 
       case .cancelled:
         // The task was cancelled, which means the continuation was
-        // called by the cancellation handler. We need to deallocate the flag
-        // word, because it was left over for this task to complete.
-        wordPtr.deallocate()
+        // called by the cancellation handler. We need to deallocate the token
+        // because it was left over for this task to complete.
+        token.deallocate()
         return
 
       case .cancelledBeforeStarted:
@@ -151,20 +181,14 @@ extension Task where Success == Never, Failure == Never {
 
   /// Called when the sleep(nanoseconds:) operation has been canceled before
   /// the sleep completed.
-  static func onSleepCancel(
-      _ wordPtr: UnsafeMutablePointer<Builtin.Word>
-  ) {
+  static func onSleepCancel(_ token: UnsafeSleepStateToken) {
     while true {
-      let state = SleepState(loading: wordPtr)
+      let state = token.load()
       switch state {
       case .notStarted:
         // We haven't started yet, so try to transition to the cancelled-before
         // started state.
-        let (_, won) = Builtin.cmpxchg_seqcst_seqcst_Word(
-            wordPtr._rawValue,
-            state.word._builtinWordValue,
-            SleepState.cancelledBeforeStarted.word._builtinWordValue)
-        if Bool(_builtinBooleanLiteral: won) {
+        if token.exchange(expected: state, desired: .cancelledBeforeStarted) {
           return
         }
 
@@ -174,11 +198,7 @@ extension Task where Success == Never, Failure == Never {
       case .activeContinuation(let continuation):
         // We have an active continuation, so try to transition to the
         // "cancelled" state.
-        let (_, won) = Builtin.cmpxchg_seqcst_seqcst_Word(
-            wordPtr._rawValue,
-            state.word._builtinWordValue,
-            SleepState.cancelled.word._builtinWordValue)
-        if Bool(_builtinBooleanLiteral: won) {
+        if token.exchange(expected: state, desired: .cancelled) {
           // We recorded the task cancellation before the sleep finished, so
           // invoke the continuation with the cancellation error.
           continuation.resume(throwing: _Concurrency.CancellationError())
@@ -203,33 +223,22 @@ extension Task where Success == Never, Failure == Never {
   ///
   /// This function doesn't block the underlying thread.
   public static func sleep(nanoseconds duration: UInt64) async throws {
-    // Allocate storage for the storage word.
-    let wordPtr = UnsafeMutablePointer<Builtin.Word>.allocate(capacity: 1)
-
-    // Initialize the flag word to "not started", which means the continuation
-    // has neither been created nor completed.
-    Builtin.atomicstore_seqcst_Word(
-        wordPtr._rawValue, SleepState.notStarted.word._builtinWordValue)
+    // Create a token which will initially have the value "not started", which
+    // means the continuation has neither been created nor completed.
+    let token = UnsafeSleepStateToken()
 
     do {
       // Install a cancellation handler to resume the continuation by
       // throwing CancellationError.
       try await withTaskCancellationHandler {
         let _: () = try await withUnsafeThrowingContinuation { continuation in
           while true {
-            let state = SleepState(loading: wordPtr)
+            let state = token.load()
             switch state {
             case .notStarted:
-              // The word that describes the active continuation state.
-              let continuationWord =
-                SleepState.activeContinuation(continuation).word
-
-              // Try to swap in the continuation word.
-              let (_, won) = Builtin.cmpxchg_seqcst_seqcst_Word(
-                  wordPtr._rawValue,
-                  state.word._builtinWordValue,
-                  continuationWord._builtinWordValue)
-              if !Bool(_builtinBooleanLiteral: won) {
+              // Try to swap in the continuation state.
+              let newState = SleepState.activeContinuation(continuation)
+              if !token.exchange(expected: state, desired: newState) {
                 // Keep trying!
                 continue
               }
@@ -243,7 +252,7 @@ extension Task where Success == Never, Failure == Never {
                 addPendingGroupTaskUnconditionally: false,
                 isDiscardingTask: false)
               let (sleepTask, _) = Builtin.createAsyncTask(sleepTaskFlags) {
-                onSleepWake(wordPtr)
+                onSleepWake(token)
               }
               _enqueueJobGlobalWithDelay(
                   duration, Builtin.convertTaskToJob(sleepTask))
@@ -264,12 +273,12 @@ extension Task where Success == Never, Failure == Never {
         }
         }
       } onCancel: {
-        onSleepCancel(wordPtr)
+        onSleepCancel(token)
       }
 
       // Determine whether we got cancelled before we even started.
       let cancelledBeforeStarted: Bool
-      switch SleepState(loading: wordPtr) {
+      switch token.load() {
       case .notStarted, .activeContinuation, .cancelled:
         fatalError("Invalid state for non-cancelled sleep task")
 
@@ -282,7 +291,7 @@ extension Task where Success == Never, Failure == Never {
 
       // We got here without being cancelled, so deallocate the storage for
       // the flag word and continuation.
-      wordPtr.deallocate()
+      token.deallocate()
 
       // If we got cancelled before we even started, through the cancellation
       // error now.

stdlib/public/Concurrency/TaskSleepDuration.swift

@@ -21,33 +21,22 @@ extension Task where Success == Never, Failure == Never {
     tolerance: Duration?,
     clock: _ClockID
   ) async throws {
-    // Allocate storage for the storage word.
-    let wordPtr = UnsafeMutablePointer<Builtin.Word>.allocate(capacity: 1)
-
-    // Initialize the flag word to "not started", which means the continuation
-    // has neither been created nor completed.
-    Builtin.atomicstore_seqcst_Word(
-        wordPtr._rawValue, SleepState.notStarted.word._builtinWordValue)
+    // Create a token which will initially have the value "not started", which
+    // means the continuation has neither been created nor completed.
+    let token = UnsafeSleepStateToken()
 
     do {
       // Install a cancellation handler to resume the continuation by
       // throwing CancellationError.
       try await withTaskCancellationHandler {
         let _: () = try await withUnsafeThrowingContinuation { continuation in
           while true {
-            let state = SleepState(loading: wordPtr)
+            let state = token.load()
             switch state {
             case .notStarted:
-              // The word that describes the active continuation state.
-              let continuationWord =
-                SleepState.activeContinuation(continuation).word
-
               // Try to swap in the continuation word.
-              let (_, won) = Builtin.cmpxchg_seqcst_seqcst_Word(
-                  wordPtr._rawValue,
-                  state.word._builtinWordValue,
-                  continuationWord._builtinWordValue)
-              if !Bool(_builtinBooleanLiteral: won) {
+              let newState = SleepState.activeContinuation(continuation)
+              if !token.exchange(expected: state, desired: newState) {
                 // Keep trying!
                 continue
               }
@@ -61,7 +50,7 @@ extension Task where Success == Never, Failure == Never {
                 addPendingGroupTaskUnconditionally: false,
                 isDiscardingTask: false)
               let (sleepTask, _) = Builtin.createAsyncTask(sleepTaskFlags) {
-                onSleepWake(wordPtr)
+                onSleepWake(token)
               }
               let toleranceSeconds: Int64
               let toleranceNanoseconds: Int64
@@ -94,12 +83,12 @@ extension Task where Success == Never, Failure == Never {
         }
         }
       } onCancel: {
-        onSleepCancel(wordPtr)
+        onSleepCancel(token)
       }
 
       // Determine whether we got cancelled before we even started.
       let cancelledBeforeStarted: Bool
-      switch SleepState(loading: wordPtr) {
+      switch token.load() {
       case .notStarted, .activeContinuation, .cancelled:
         fatalError("Invalid state for non-cancelled sleep task")
 
@@ -112,7 +101,7 @@ extension Task where Success == Never, Failure == Never {
 
       // We got here without being cancelled, so deallocate the storage for
       // the flag word and continuation.
-      wordPtr.deallocate()
+      token.deallocate()
 
       // If we got cancelled before we even started, through the cancellation
       // error now.