[Concurrency] Initial Task Local Values implementation

Author: ktoso

include/swift/ABI/MetadataValues.h

@@ -1936,9 +1936,10 @@ class JobFlags : public FlagSet<size_t> {
 
     // Kind-specific flags.
 
-    Task_IsChildTask  = 24,
-    Task_IsFuture     = 25,
-    Task_IsTaskGroup  = 26,
+    Task_IsChildTask    = 24,
+    Task_IsFuture       = 25,
+    Task_IsTaskGroup    = 26,
+    Task_HasLocalValues = 27
   };
 
   explicit JobFlags(size_t bits) : FlagSet(bits) {}
@@ -1965,11 +1966,12 @@ class JobFlags : public FlagSet<size_t> {
   FLAGSET_DEFINE_FLAG_ACCESSORS(Task_IsFuture,
                                 task_isFuture,
                                 task_setIsFuture)
-
   FLAGSET_DEFINE_FLAG_ACCESSORS(Task_IsTaskGroup,
                                 task_isTaskGroup,
                                 task_setIsTaskGroup)
-
+  FLAGSET_DEFINE_FLAG_ACCESSORS(Task_HasLocalValues,
+                                task_hasLocalValues,
+                                task_setHasLocalValues)
 };
 
 /// Kinds of task status record.

include/swift/ABI/Task.h

@@ -137,11 +137,12 @@ class ActiveTaskStatus {
 /// ### Fragments
 /// An AsyncTask may have the following fragments:
 ///
-///    +------------------+
-///    | childFragment?   |
-///    | groupFragment?   |
-///    | futureFragment?  |*
-///    +------------------+
+///    +--------------------------+
+///    | childFragment?           |
+///    | taskLocalValuesFragment  |
+///    | groupFragment?           |
+///    | futureFragment?          |*
+///    +--------------------------+
 ///
 /// The future fragment is dynamic in size, based on the future result type
 /// it can hold, and thus must be the *last* fragment.
@@ -175,13 +176,15 @@ class AsyncTask : public HeapObject, public Job {
   void run(ExecutorRef currentExecutor) {
     ResumeTask(this, currentExecutor, ResumeContext);
   }
-
+  
   /// Check whether this task has been cancelled.
   /// Checking this is, of course, inherently race-prone on its own.
   bool isCancelled() const {
     return Status.load(std::memory_order_relaxed).isCancelled();
   }
 
+  // ==== Child Fragment -------------------------------------------------------
+
   /// A fragment of an async task structure that happens to be a child task.
   class ChildFragment {
     /// The parent task of this task.
@@ -205,14 +208,212 @@ class AsyncTask : public HeapObject, public Job {
     }
   };
 
-  // TODO: rename? all other functions are `is...` rather than `has...Fragment`
-  bool hasChildFragment() const { return Flags.task_isChildTask(); }
+  bool hasChildFragment() const {
+    return Flags.task_isChildTask();
+  }
 
   ChildFragment *childFragment() {
     assert(hasChildFragment());
     return reinterpret_cast<ChildFragment*>(this + 1);
   }
 
+  // ==== Task Locals Values-- -------------------------------------------------
+
+  class TaskLocalValuesFragment {
+  public:
+    /// Type of the pointed at `next` task local item.
+    enum class NextLinkType : uintptr_t {
+      /// This task is known to be a "terminal" node in the lookup of task locals.
+      /// In other words, even if it had a parent, the parent (and its parents)
+      /// are known to not contain any any more task locals, and thus any further
+      /// search beyond this task.
+      IsTerminal = 0b00,
+      /// The storage pointer points at the next TaskLocalChainItem in this task.
+      IsNext     = 0b01,
+      /// The storage pointer points at a parent AsyncTask,
+      /// in which we should continue the lookup.
+      IsParent   = 0b11
+    };
+
+    class TaskLocalItem {
+    private:
+      /// Mask used for the low status bits in a task local chain item.
+      static const uintptr_t statusMask = 0x03;
+
+      /// Pointer to the next task local item; be it in this task or in a parent.
+      /// Low bits encode `NextLinkType`.
+      /// TaskLocalItem *next = nullptr;
+      uintptr_t next;
+
+    public:
+      /// The type of the key with which this value is associated.
+      const Metadata *keyType;
+      /// The type of the value stored by this item.
+      const Metadata *valueType;
+
+      // Trailing storage for the value itself. The storage will be
+      // uninitialized or contain an instance of \c valueType.
+
+    private:
+      explicit TaskLocalItem(const Metadata *keyType, const Metadata *valueType)
+          : keyType(keyType),
+            valueType(valueType),
+            next(0) { }
+
+    public:
+      /// TaskLocalItem which does not by itself store any value, but only points
+      /// to the nearest task-local-value containing parent's first task item.
+      ///
+      /// This item type is used to link to the appropriate parent task's item,
+      /// when the current task itself does not have any task local values itself.
+      ///
+      /// When a task actually has its own task locals, it should rather point
+      /// to the parent's *first* task-local item in its *last* item, extending
+      /// the TaskLocalItem linked list into the appropriate parent.
+      static TaskLocalItem* createParentLink(AsyncTask *task, AsyncTask *parent) {
+        assert(parent);
+        assert(parent->hasTaskLocalValues());
+        assert(task->hasTaskLocalValues());
+        size_t amountToAllocate = TaskLocalItem::itemSize(/*valueType*/nullptr);
+        // assert(amountToAllocate % MaximumAlignment == 0); // TODO: do we need this?
+        void *allocation = malloc(amountToAllocate); // TODO: use task-local allocator
+        fprintf(stderr, "MALLOC parent link item: %d\n", allocation);
+
+        TaskLocalItem *item =
+            new(allocation) TaskLocalItem(nullptr, nullptr);
+
+        auto next = parent->localValuesFragment()->head;
+        auto nextLinkType = next ? NextLinkType::IsParent : NextLinkType::IsTerminal;
+        item->next = reinterpret_cast<uintptr_t>(next) |
+                          static_cast<uintptr_t>(nextLinkType);
+
+        fprintf(stderr, "error: %s [%s:%d] created parent item: task=%d -> parentTask=%d  :: item=%d -> item->getNext()=%d\n", __FUNCTION__, __FILE_NAME__, __LINE__,
+                task, parent, item, item->getNext());
+
+        return item;
+      }
+
+      static TaskLocalItem* createLink(AsyncTask *task,
+                                       const Metadata *keyType,
+                                       const Metadata *valueType) {
+        assert(task);
+        assert(task->hasTaskLocalValues());
+        size_t amountToAllocate = TaskLocalItem::itemSize(valueType);
+        // assert(amountToAllocate % MaximumAlignment == 0); // TODO: do we need this?
+        void *allocation = malloc(amountToAllocate); // TODO: use task-local allocator
+        fprintf(stderr, "MALLOC link item: %d\n", allocation);
+        TaskLocalItem *item =
+            new(allocation) TaskLocalItem(keyType, valueType);
+
+        auto next = task->localValuesFragment()->head;
+        auto nextLinkType = next ? NextLinkType::IsNext : NextLinkType::IsTerminal;
+        item->next = reinterpret_cast<uintptr_t>(next) |
+            static_cast<uintptr_t>(nextLinkType);
+
+        return item;
+      }
+
+      void destroy() {
+        if (valueType) {
+          valueType->vw_destroy(getStoragePtr());
+        }
+      }
+
+      TaskLocalItem *getNext() {
+        return reinterpret_cast<TaskLocalItem *>(next & ~statusMask);
+      }
+
+      NextLinkType getNextLinkType() {
+        return static_cast<NextLinkType>(next & statusMask);
+      }
+
+      /// Retrieve a pointer to the storage of the value.
+      OpaqueValue *getStoragePtr() {
+//        assert(valueType && "valueType must be set before accessing storage pointer.");
+        return reinterpret_cast<OpaqueValue *>(
+            reinterpret_cast<char *>(this) + storageOffset(valueType));
+      }
+
+      /// Compute the offset of the storage from the base of the item.
+      static size_t storageOffset(const Metadata *valueType) {
+        size_t offset = sizeof(TaskLocalItem);
+        if (valueType) {
+          size_t alignment = valueType->vw_alignment();
+          return (offset + alignment - 1) & ~(alignment - 1);
+        } else {
+          return offset;
+        }
+      }
+
+      /// Determine the size of the item given a particular value type.
+      static size_t itemSize(const Metadata *valueType) {
+        size_t offset = storageOffset(valueType);
+        if (valueType) {
+          offset += valueType->vw_size();
+        }
+        return offset;
+      }
+    };
+
+  private:
+    /// Single-linked list of task local values.
+    /// Once task local values within this task are traversed, the list continues
+    /// to the "next parent that contributes task local values," or if no such
+    /// parent exists it terminates with null.
+    ///
+    /// If the TaskLocalValuesFragment was allocated, it is expected that this
+    /// value should be NOT null; it either has own values, or at least one
+    /// parent that has values. If this task does not have any values, the head
+    /// pointer MAY immediately point at this task's parent task which has values.
+    ///
+    /// NOTE: Check the higher bits to know if this is a self or parent value.
+    TaskLocalItem *head = nullptr;
+
+  public:
+      TaskLocalValuesFragment() {}
+
+    void destroy();
+
+    /// If the parent task has task local values defined, point to in
+    /// the task local values chain.
+    void initializeLinkParent(AsyncTask* task, AsyncTask* parent);
+
+    void pushValue(AsyncTask *task, const Metadata *keyType,
+                   /* +1 */ OpaqueValue *value, const Metadata *valueType);
+
+    void popValue(AsyncTask *task);
+
+    OpaqueValue* get(const Metadata *keyType);
+  };
+
+  bool hasTaskLocalValues() const {
+    return Flags.task_hasLocalValues();
+  }
+
+  TaskLocalValuesFragment *localValuesFragment() {
+    assert(hasTaskLocalValues());
+
+    auto offset = reinterpret_cast<char*>(this);
+    offset += sizeof(AsyncTask);
+
+    if (hasChildFragment()) {
+      offset += sizeof(ChildFragment);
+    }
+
+    return reinterpret_cast<TaskLocalValuesFragment*>(offset);
+  }
+
+  OpaqueValue* localValueGet(const Metadata *keyType) {
+    if (hasTaskLocalValues()) {
+      return localValuesFragment()->get(keyType);
+    } else {
+      // We are guaranteed to have a task-local fragment even if this task has
+      // no bindings, but its parent tasks do. Thus, if no fragment, we can
+      // immediately return null.
+      return nullptr;
+    }
+  }
+
   // ==== TaskGroup ------------------------------------------------------------
 
   class GroupFragment {
@@ -516,12 +717,18 @@ class AsyncTask : public HeapObject, public Job {
   GroupFragment *groupFragment() {
     assert(isTaskGroup());
 
+    auto offset = reinterpret_cast<char*>(this);
+    offset += sizeof(AsyncTask);
+
     if (hasChildFragment()) {
-      return reinterpret_cast<GroupFragment *>(
-          reinterpret_cast<ChildFragment*>(this + 1) + 1);
+      offset += sizeof(ChildFragment);
     }
 
-    return reinterpret_cast<GroupFragment *>(this + 1);
+    if (hasTaskLocalValues()) {
+      offset += sizeof(TaskLocalValuesFragment);
+    }
+
+    return reinterpret_cast<GroupFragment *>(offset);
   }
 
   /// Offer result of a task into this channel.
@@ -647,13 +854,17 @@ class AsyncTask : public HeapObject, public Job {
   FutureFragment *futureFragment() {
     assert(isFuture());
 
-    auto offset = reinterpret_cast<uintptr_t>(this); // TODO: char* instead?
+    auto offset = reinterpret_cast<char*>(this);
     offset += sizeof(AsyncTask);
 
     if (hasChildFragment()) {
       offset += sizeof(ChildFragment);
     }
 
+    if (hasTaskLocalValues()) {
+      offset += sizeof(TaskLocalValuesFragment);
+    }
+
     if (isTaskGroup()) {
       offset += sizeof(GroupFragment);
     }

include/swift/Runtime/Concurrency.h

@@ -18,6 +18,7 @@
 #define SWIFT_RUNTIME_CONCURRENCY_H
 
 #include "swift/ABI/TaskStatus.h"
+#include "swift/Runtime/ExistentialContainer.h"
 
 namespace swift {
 class DefaultActor;
@@ -225,6 +226,63 @@ size_t swift_task_getJobFlags(AsyncTask* task);
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
 bool swift_task_isCancelled(AsyncTask* task);
 
+/// Get a task local value from the passed in task. 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 Metadata *keyType);
+
+/// Add a task local value to the passed in task.
+///
+/// This must be only invoked by the task itself to avoid concurrent writes.
+///
+/// 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 Metadata *keyType,
+                         /* +1 */ OpaqueValue *value,
+                         const Metadata *valueType);
+
+/// Remove task `count` local bindings from the task local binding stack.
+/// Crashes if `count` is greater if the number of task locals stored in the task.
+///
+/// This must be only invoked by the task itself to avoid concurrent writes.
+///
+/// Its Swift signature is
+///
+/// \code
+///  public func _taskLocalValuePop(
+///    _ task: Builtin.NativeObject,
+///    count: Int
+///  )
+/// \endcode
+SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
+void swift_task_localValuePop(AsyncTask* task, int count);
+
+/// Checks if task (or any of its parent tasks) has task local values.
+///
+/// \code
+/// func swift_task_hasTaskLocalValues<Key>(
+///   _ task: Builtin.NativeObject,
+/// ) -> Bool
+/// \endcode
+SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
+bool swift_task_hasTaskLocalValues(AsyncTask* task);
+
 /// This should have the same representation as an enum like this:
 ///    enum NearestTaskDeadline {
 ///      case none