[TaskLocal] Use the task-local stack discipline allocator

Author: ktoso

include/swift/ABI/Task.h

@@ -201,10 +201,6 @@ class AsyncTask : public HeapObject, public Job {
 
   // ==== Task Local Values ----------------------------------------------------
 
-  void localValueInitializeLinkParent(AsyncTask *parent) {
-    Local.initializeLinkParent(this, parent);
-  }
-
   void localValuePush(const Metadata *keyType,
                       /* +1 */ OpaqueValue *value, const Metadata *valueType) {
     Local.pushValue(this, keyType, value, valueType);

include/swift/ABI/TaskLocal.h

@@ -93,6 +93,11 @@ class TaskLocal {
     // uninitialized or contain an instance of \c valueType.
 
   private:
+    explicit Item()
+      : next(0),
+        keyType(nullptr),
+        valueType(nullptr) {}
+
     explicit Item(const Metadata *keyType, const Metadata *valueType)
       : next(0),
         keyType(keyType),
@@ -114,11 +119,7 @@ class TaskLocal {
                             const Metadata *keyType,
                             const Metadata *valueType);
 
-    void destroy(AsyncTask *task) {
-      if (valueType) {
-        valueType->vw_destroy(getStoragePtr());
-      }
-    }
+    void destroy(AsyncTask *task);
 
     Item *getNext() {
       return reinterpret_cast<Item *>(next & ~statusMask);
@@ -210,12 +211,13 @@ class TaskLocal {
 
     void popValue(AsyncTask *task);
 
+    /// Destroy and deallocate all items stored by this specific task.
+    ///
+    /// Items owned by a parent task are left untouched, since we do not own them.
     void destroy(AsyncTask *task);
   };
 };
 
-TaskLocal::Storage* swift_task_localValueStorage(AsyncTask *task);
-
 } // end namespace swift
 
 #endif

stdlib/public/Concurrency/Task.cpp

@@ -190,6 +190,10 @@ static void completeTask(AsyncTask *task, ExecutorRef executor,
   // Set that there's no longer a running task in the current thread.
   _swift_task_clearCurrent();
 
+  // Destroy and deallocate any remaining task local items.
+  // We need to do this before we destroy the task local deallocator.
+  task->Local.destroy(task);
+
   // Tear down the task-local allocator immediately;
   // there's no need to wait for the object to be destroyed.
   _swift_task_alloc_destroy(task);
@@ -288,9 +292,6 @@ AsyncTaskAndContext swift::swift_task_create_group_future_f(
     new(allocation) AsyncTask(&taskHeapMetadata, flags,
                               function, initialContext);
 
-  // Initialize task locals fragment.
-  task->Local.initializeLinkParent(task, parent);
-
   // Initialize the child fragment if applicable.
   if (parent) {
     auto childFragment = task->childFragment();
@@ -343,6 +344,13 @@ AsyncTaskAndContext swift::swift_task_create_group_future_f(
   // TODO: consider providing an initial pre-allocated first slab to the allocator.
   _swift_task_alloc_initialize(task);
 
+  // TODO: if the allocator would be prepared earlier we could do this in some
+  //       other existing if-parent if rather than adding another one here.
+  if (parent) {
+    // Initialize task locals with a link to the parent task.
+    task->Local.initializeLinkParent(task, parent);
+  }
+
   return {task, initialContext};
 }
 

stdlib/public/Concurrency/TaskLocal.cpp

@@ -20,10 +20,6 @@ using namespace swift;
 // =============================================================================
 // ==== ABI --------------------------------------------------------------------
 
-TaskLocal::Storage* swift::swift_task_localValueStorage(AsyncTask *task) {
-  return &task->Local;
-}
-
 void swift::swift_task_localValuePush(AsyncTask *task,
                                       const Metadata *keyType,
                                       /* +1 */ OpaqueValue *value,
@@ -46,25 +42,19 @@ void swift::swift_task_localValuePop(AsyncTask *task) {
 
 void TaskLocal::Storage::initializeLinkParent(AsyncTask* task,
                                               AsyncTask* parent) {
-  assert(!head && "task local storage was already initialized with parent");
-  if (parent) {
-    head = TaskLocal::Item::createParentLink(task, parent);
-  }
+  assert(!head && "initial task local storage was already initialized");
+  assert(parent && "parent must be provided to link to it");
+  head = TaskLocal::Item::createParentLink(task, parent);
 }
 
 TaskLocal::Item*
 TaskLocal::Item::createParentLink(AsyncTask *task, AsyncTask *parent) {
-  assert(parent);
   size_t amountToAllocate = Item::itemSize(/*valueType*/nullptr);
   // assert(amountToAllocate % MaximumAlignment == 0); // TODO: do we need this?
-  void *allocation = malloc(amountToAllocate); // TODO: use task-local allocator
+  void *allocation = swift_task_alloc(task, amountToAllocate);
+  Item *item = new(allocation) Item();
 
-  Item *item =
-    new(allocation) Item(nullptr, nullptr);
-
-   auto parentHead = parent->Local.head;
-//  auto parentLocalStorage = swift_task_localValueStorage(parent);
-//  auto parentHead = parentLocalStorage->head;
+  auto parentHead = parent->Local.head;
   if (parentHead) {
     if (parentHead->isEmpty()) {
       switch (parentHead->getNextLinkType()) {
@@ -105,11 +95,9 @@ TaskLocal::Item::createLink(AsyncTask *task,
   assert(task);
   size_t amountToAllocate = Item::itemSize(valueType);
   // assert(amountToAllocate % MaximumAlignment == 0); // TODO: do we need this?
-  void *allocation = malloc(amountToAllocate); // TODO: use task-local allocator rdar://74218679
-  Item *item =
-    new(allocation) Item(keyType, valueType);
+  void *allocation = swift_task_alloc(task, amountToAllocate);
+  Item *item = new(allocation) Item(keyType, valueType);
 
-  auto nextStorage = swift_task_localValueStorage(task);
   auto next = task->Local.head;
   auto nextLinkType = next ? NextLinkType::IsNext
                            : NextLinkType::IsTerminal;
@@ -122,6 +110,14 @@ TaskLocal::Item::createLink(AsyncTask *task,
 // =============================================================================
 // ==== destroy ----------------------------------------------------------------
 
+void TaskLocal::Item::destroy(AsyncTask *task) {
+  if (valueType) {
+    valueType->vw_destroy(getStoragePtr());
+  }
+
+  swift_task_dealloc(task, this);
+}
+
 void TaskLocal::Storage::destroy(AsyncTask *task) {
   auto item = head;
   head = nullptr;
@@ -131,13 +127,14 @@ void TaskLocal::Storage::destroy(AsyncTask *task) {
       case TaskLocal::NextLinkType::IsNext:
         next = item->getNext();
         item->destroy(task);
-        free(item);
         item = next;
         break;
 
       case TaskLocal::NextLinkType::IsParent:
       case TaskLocal::NextLinkType::IsTerminal:
-        // we're done here, we must not destroy values owned by the parent task.
+        // we're done here; as we must not proceed into the parent owned values.
+        // we do have to destroy the item pointing at the parent/edge itself though.
+        item->destroy(task);
         return;
     }
   }
@@ -159,8 +156,9 @@ void TaskLocal::Storage::pushValue(AsyncTask *task,
 
 void TaskLocal::Storage::popValue(AsyncTask *task) {
   assert(head && "attempted to pop value off empty task-local stack");
-  head->destroy(task);
+  auto old = head;
   head = head->getNext();
+  old->destroy(task);
 }
 
 OpaqueValue* TaskLocal::Storage::getValue(AsyncTask *task,

stdlib/public/Concurrency/TaskLocal.swift

@@ -71,13 +71,12 @@ extension Task {
   ///            bound in the current (or any parent) tasks.
   public static func local<Key>(_ keyPath: KeyPath<TaskLocalValues, Key>)
     -> Key.Value where Key: TaskLocalKey {
-    guard let _task = Task.unsafeCurrent?._task else {
-      fatalError("No async task!")
+    guard let unsafeTask = Task.unsafeCurrent else {
       return Key.defaultValue
     }
 
     let value = _taskLocalValueGet(
-      _task, keyType: Key.self, inheritance: Key.inherit.rawValue)
+      unsafeTask._task, keyType: Key.self, inheritance: Key.inherit.rawValue)
     guard let rawValue = value else {
       return Key.defaultValue
     }

stdlib/public/Concurrency/TaskPrivate.h

@@ -57,7 +57,6 @@ void donateThreadToGlobalExecutorUntil(bool (*condition)(void*),
 #endif
 
 // ==== ------------------------------------------------------------------------
-// TODO: this was moved from Task.cpp to share it with TaskGroup.cpp, where else better place to put it?
 
 namespace {
 

test/Concurrency/Runtime/async_task_locals_async_let.swift

@@ -46,11 +46,11 @@ func async_let_nested() async {
     func test() async {
       printTaskLocal(\.number) // CHECK: NumberKey: 2 {{.*}}
       async let x31 = printTaskLocal(\.number) // CHECK: NumberKey: 2 {{.*}}
-      _ = try! await x31
+      _ = await x31
     }
     async let x3: () = test()
 
-    _ = try! await x2
+    _ = await x2
     await x3
   }
 
@@ -64,8 +64,9 @@ func async_let_nested_skip_optimization() async {
       async let x3: Int? = { () async -> Int? in
         async let x4: Int? = { () async -> Int? in
           async let x5: Int? = { () async -> Int? in
+            assert(Task.local(\.number) == 2)
             async let xx = printTaskLocal(\.number) // CHECK: NumberKey: 2 {{.*}}
-            return try! await xx
+            return await xx
           }()
           return await x5
         }()

test/Concurrency/Runtime/async_task_locals_basic.swift

@@ -53,7 +53,7 @@ func printTaskLocal<Key>(
   _ expected: Key.Value? = nil,
   file: String = #file, line: UInt = #line, function: String = #function
 ) where Key: TaskLocalKey {
-  let value = Task.local(key, file: file, line: line, function: function)
+  let value = Task.local(key)
   print("\(Key.self): \(value) at \(file):\(line)")
   if let expected = expected {
     assert("\(expected)" == "\(value)",

test/Concurrency/Runtime/async_task_locals_groups.swift

@@ -25,7 +25,7 @@ func printTaskLocal<Key>(
   _ key: KeyPath<TaskLocalValues, Key>,
   _ expected: Key.Value? = nil,
   file: String = #file, line: UInt = #line
-) async throws where Key: TaskLocalKey {
+) where Key: TaskLocalKey {
   let value = Task.local(key)
   print("\(Key.self): \(value) at \(file):\(line)")
   if let expected = expected {
@@ -68,8 +68,8 @@ func groups() async {
       await group.add {
         printTaskLocal(\.number) // CHECK: NumberKey: 2 {{.*}}
 
-        async let childInsideGroupChild = printTaskLocal(\.number)
-        try! await childInsideGroupChild // CHECK: NumberKey: 2 {{.*}}
+        async let childInsideGroupChild: () = printTaskLocal(\.number)
+        await childInsideGroupChild // CHECK: NumberKey: 2 {{.*}}
 
         return Task.local(\.number)
       }

test/Concurrency/Runtime/async_task_locals_inherit_never.swift

@@ -17,7 +17,7 @@ func printTaskLocal<Key>(
   _ key: KeyPath<TaskLocalValues, Key>,
   _ expected: Key.Value? = nil,
   file: String = #file, line: UInt = #line
-) async where Key: TaskLocalKey {
+) where Key: TaskLocalKey {
   let value = Task.local(key)
   print("\(Key.self): \(value) at \(file):\(line)")
   if let expected = expected {