Merge pull request #37340 from ktoso/wip-locals-sync

[TaskLocals] Enable sync functions to bind task-locals; Keep Storage in TLS

Author: ktoso

include/swift/ABI/Task.h

@@ -271,8 +271,9 @@ class AsyncTask : public Job {
     return Local.getValue(this, key);
   }
 
-  void localValuePop() {
-    Local.popValue(this);
+  /// Returns true if storage has still more bindings.
+  bool localValuePop() {
+    return Local.popValue(this);
   }
 
   // ==== Child Fragment -------------------------------------------------------

include/swift/ABI/TaskLocal.h

@@ -118,6 +118,8 @@ class TaskLocal {
         reinterpret_cast<char *>(this) + storageOffset(valueType));
     }
 
+    void copyTo(AsyncTask *task);
+
     /// Compute the offset of the storage from the base of the item.
     static size_t storageOffset(const Metadata *valueType) {
       size_t offset = sizeof(Item);
@@ -184,7 +186,22 @@ class TaskLocal {
 
     OpaqueValue* getValue(AsyncTask *task, const HeapObject *key);
 
-    void popValue(AsyncTask *task);
+    /// Returns `true` of more bindings remain in this storage,
+    /// and `false` if the just popped value was the last one and the storage
+    /// can be safely disposed of.
+    bool popValue(AsyncTask *task);
+
+    /// Copy all task-local bindings to the target task.
+    ///
+    /// The new bindings allocate their own items and can out-live the current task.
+    ///
+    /// ### Optimizations
+    /// Only the most recent binding of a value is copied over, i.e. given
+    /// a key bound to `A` and then `B`, only the `B` binding will be copied.
+    /// This is safe and correct because the new task would never have a chance
+    /// to observe the `A` value, because it semantically will never observe a
+    /// "pop" of the `B` value - it was spawned from a scope where only B was observable.
+    void copyTo(AsyncTask *target);
 
     /// Destroy and deallocate all items stored by this specific task.
     ///

include/swift/Runtime/Concurrency.h

@@ -433,50 +433,62 @@ void swift_task_reportIllegalTaskLocalBindingWithinWithTaskGroup(
     const unsigned char *file, uintptr_t fileLength,
     bool fileIsASCII, uintptr_t line);
 
-/// Get a task local value from the passed in task. Its Swift signature is
+/// Get a task local value from either the current task, or fallback task-local
+/// storage.
+///
+/// Its Swift signature is
 ///
 /// \code
 /// func _taskLocalValueGet<Key>(
-///   _ task: Builtin.NativeObject,
 ///   keyType: Any.Type /*Key.Type*/
 /// ) -> UnsafeMutableRawPointer? where Key: TaskLocalKey
 /// \endcode
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
 OpaqueValue*
-swift_task_localValueGet(AsyncTask* task, const HeapObject *key);
+swift_task_localValueGet(const HeapObject *key);
 
-/// Add a task local value to the passed in task.
-///
-/// This must be only invoked by the task itself to avoid concurrent writes.
+/// Bind a task local key to a value in the context of either the current
+/// AsyncTask if present, or in the thread-local fallback context if no task
+/// available.
 ///
 /// Its Swift signature is
 ///
 /// \code
 ///  public func _taskLocalValuePush<Value>(
-///    _ task: Builtin.NativeObject,
 ///    keyType: Any.Type/*Key.Type*/,
 ///    value: __owned Value
 ///  )
 /// \endcode
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
-void swift_task_localValuePush(AsyncTask* task,
-                         const HeapObject *key,
-                         /* +1 */ OpaqueValue *value,
-                         const Metadata *valueType);
+void swift_task_localValuePush(const HeapObject *key,
+                                   /* +1 */ OpaqueValue *value,
+                                   const Metadata *valueType);
 
-/// Remove task a local binding from the task local values stack.
+/// Pop a single task local binding from the binding stack of the current task,
+/// or the fallback thread-local storage if no task is available.
 ///
-/// This must be only invoked by the task itself to avoid concurrent writes.
+/// This operation must be paired up with a preceding "push" operation, as otherwise
+/// it may attempt to "pop" off an empty value stuck which will lead to a crash.
+///
+/// The Swift surface API ensures proper pairing of push and pop operations.
 ///
 /// Its Swift signature is
 ///
 /// \code
-///  public func _taskLocalValuePop(
-///    _ task: Builtin.NativeObject
-///  )
+///  public func _taskLocalValuePop()
+/// \endcode
+SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
+void swift_task_localValuePop();
+
+/// Copy all task locals from the current context to the target task.
+///
+/// Its Swift signature is
+///
+/// \code
+/// func _taskLocalValueGet<Key>(AsyncTask* task)
 /// \endcode
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
-void swift_task_localValuePop(AsyncTask* task);
+void swift_task_localsCopyTo(AsyncTask* target);
 
 /// This should have the same representation as an enum like this:
 ///    enum NearestTaskDeadline {

include/swift/Runtime/ThreadLocal.h

@@ -65,7 +65,7 @@ namespace swift {
 //   itself thread-local, and no internal support is required.
 //
 //   Note that this includes platforms that set
-//   SWIFT_STDLIB_SINGLE_THREADED_RUNTIME, for whhch
+//   SWIFT_STDLIB_SINGLE_THREADED_RUNTIME, for which
 //   SWIFT_RUNTIME_ATTRIBUTE_THREAD_LOCAL is empty;
 //   thread-local declarations then create an ordinary global.
 //

stdlib/public/CompatibilityOverride/CompatibilityOverrideConcurrency.def

@@ -233,19 +233,25 @@ OVERRIDE_TASK_LOCAL(task_reportIllegalTaskLocalBindingWithinWithTaskGroup, void,
 OVERRIDE_TASK_LOCAL(task_localValuePush, void,
                     SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift),
                     swift::,
-                    (AsyncTask *task, const HeapObject *key,
-                     OpaqueValue *value, const Metadata *valueType),
-                    (task, key, value, valueType))
+                    (const HeapObject *key, OpaqueValue *value,
+                     const Metadata *valueType),
+                    (key, value, valueType))
 
 OVERRIDE_TASK_LOCAL(task_localValueGet, OpaqueValue *,
                     SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift),
                     swift::,
-                    (AsyncTask *task, const HeapObject *key),
-                    (task, key))
+                    (const HeapObject *key),
+                    (key))
 
 OVERRIDE_TASK_LOCAL(task_localValuePop, void,
                     SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift),
-                    swift::, (AsyncTask *task), (task))
+                    swift::, ,)
+
+OVERRIDE_TASK_LOCAL(task_localsCopyTo, void,
+                    SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift),
+                    swift::,
+                    (AsyncTask *target),
+                    (target))
 
 OVERRIDE_TASK_STATUS(task_addStatusRecord, bool,
                      SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift),

stdlib/public/Concurrency/Task.cpp

@@ -195,6 +195,10 @@ AsyncTask::~AsyncTask() {
   }
 
   // Release any objects potentially held as task local values.
+  //
+  // This must be called last when destroying a task - to keep stack discipline of the allocator.
+  // because it may have created some task-locals immediately upon creation,
+  // e.g. if the task is spawned with async{} and inherited some task-locals.
   Local.destroy(this);
 }
 

stdlib/public/Concurrency/Task.swift

@@ -535,6 +535,15 @@ public func async<T>(
   // Create the asynchronous task future.
   let (task, _) = Builtin.createAsyncTaskFuture(flags.bits, operation)
 
+  // Copy all task locals to the newly created task.
+  // We must copy them rather than point to the current task since the new task
+  // is not structured and may out-live the current task.
+  //
+  // WARNING: This MUST be done BEFORE we enqueue the task,
+  // because it acts as-if it was running inside the task and thus does not
+  // take any extra steps to synchronize the task-local operations.
+  _taskLocalsCopy(to: task)
+
   // Enqueue the resulting job.
   _enqueueJobGlobal(Builtin.convertTaskToJob(task))