[Concurrency] Relax task-local mis-use prevention in task groups

Author: ktoso

include/swift/ABI/TaskLocal.h

@@ -44,6 +44,10 @@ class TaskLocal {
     /// lookups by skipping empty parent tasks during get(), and explained
     /// in depth in `createParentLink`.
     IsParent = 0b01,
+    /// The task local binding was created inside the body of a `withTaskGroup`,
+    /// and therefore must either copy it, or crash when a child task is created
+    /// using 'group.addTask' and it would refer to this task local.
+    IsNextCreatedInTaskGroupBody = 0b10,
   };
 
   class Item {
@@ -101,10 +105,20 @@ class TaskLocal {
     /// the Item linked list into the appropriate parent.
     static Item *createParentLink(AsyncTask *task, AsyncTask *parent);
 
+    static Item *createLink(AsyncTask *task,
+                            const HeapObject *key,
+                            const Metadata *valueType,
+                            bool inTaskGroupBody);
+
     static Item *createLink(AsyncTask *task,
                             const HeapObject *key,
                             const Metadata *valueType);
 
+    static Item *createLinkInTaskGroup(
+        AsyncTask *task,
+        const HeapObject *key,
+        const Metadata *valueType);
+
     void destroy(AsyncTask *task);
 
     Item *getNext() {
@@ -115,6 +129,16 @@ class TaskLocal {
       return static_cast<NextLinkType>(next & statusMask);
     }
 
+    bool isNextLinkPointer() const {
+      return static_cast<NextLinkType>(next & statusMask) ==
+             NextLinkType::IsNext;
+    }
+
+    bool IsNextCreatedInTaskGroupBody() const {
+      return static_cast<NextLinkType>(next & statusMask) ==
+             NextLinkType::IsNextCreatedInTaskGroupBody;
+    }
+
     /// Item does not contain any actual value, and is only used to point at
     /// a specific parent item.
     bool isEmpty() const {
@@ -136,9 +160,9 @@ class TaskLocal {
       if (valueType) {
         size_t alignment = valueType->vw_alignment();
         return (offset + alignment - 1) & ~(alignment - 1);
-      } else {
-        return offset;
       }
+
+      return offset;
     }
 
     /// Determine the size of the item given a particular value type.
@@ -200,6 +224,8 @@ class TaskLocal {
     /// can be safely disposed of.
     bool popValue(AsyncTask *task);
 
+    NextLinkType peekHeadLinkType() const;
+
     /// Copy all task-local bindings to the target task.
     ///
     /// The new bindings allocate their own items and can out-live the current task.

stdlib/public/Concurrency/Task.cpp

@@ -729,19 +729,31 @@ swift_task_create_commonImpl(size_t rawTaskCreateFlags,
     assert(initialContextSize >= sizeof(FutureAsyncContext));
   }
 
-  // Add to the task group, if requested.
-  if (taskCreateFlags.addPendingGroupTaskUnconditionally()) {
-    assert(group && "Missing group");
-    swift_taskGroup_addPending(group, /*unconditionally=*/true);
-  }
-
   AsyncTask *parent = nullptr;
   AsyncTask *currentTask = swift_task_getCurrent();
   if (jobFlags.task_isChildTask()) {
     parent = currentTask;
     assert(parent != nullptr && "creating a child task with no active task");
   }
 
+  if (group) {
+    assert(parent && "a task created in a group must be a child task");
+
+    // Prevent task-local misuse;
+    // We must not allow an addTask {} wrapped immediately with a withValue {}
+    auto taskLocalHeadLinkType = parent->_private().Local.peekHeadLinkType();
+    if (taskLocalHeadLinkType == swift::TaskLocal::NextLinkType::IsNextCreatedInTaskGroupBody) {
+      swift_task_reportIllegalTaskLocalBindingWithinWithTaskGroup(nullptr, 0, true, 0);
+      abort();
+    }
+
+    // Add to the task group, if requested.
+    if (taskCreateFlags.addPendingGroupTaskUnconditionally()) {
+      assert(group && "Missing group");
+      swift_taskGroup_addPending(group, /*unconditionally=*/true);
+    }
+  }
+
   // Start with user specified priority at creation time (if any)
   JobPriority basePriority = (taskCreateFlags.getRequestedPriority());
 

stdlib/public/Concurrency/TaskLocal.cpp

@@ -171,6 +171,7 @@ TaskLocal::Item::createParentLink(AsyncTask *task, AsyncTask *parent) {
                        static_cast<uintptr_t>(NextLinkType::IsParent);
           break;
         case NextLinkType::IsNext:
+        case NextLinkType::IsNextCreatedInTaskGroupBody:
           if (parentHead->getNext()) {
             assert(false && "empty taskValue head in parent task, yet parent's 'head' is `IsNext`, "
                             "this should not happen, as it implies the parent must have stored some value.");
@@ -194,7 +195,8 @@ TaskLocal::Item::createParentLink(AsyncTask *task, AsyncTask *parent) {
 TaskLocal::Item*
 TaskLocal::Item::createLink(AsyncTask *task,
                             const HeapObject *key,
-                            const Metadata *valueType) {
+                            const Metadata *valueType,
+                            bool inTaskGroupBody) {
   size_t amountToAllocate = Item::itemSize(valueType);
   void *allocation = task ? _swift_task_alloc_specific(task, amountToAllocate)
                           : malloc(amountToAllocate);
@@ -203,11 +205,28 @@ TaskLocal::Item::createLink(AsyncTask *task,
   auto next = task ? task->_private().Local.head
                    : FallbackTaskLocalStorage::get()->head;
   item->next = reinterpret_cast<uintptr_t>(next) |
-      static_cast<uintptr_t>(NextLinkType::IsNext);
+               static_cast<uintptr_t>(
+                   inTaskGroupBody ? NextLinkType::IsNextCreatedInTaskGroupBody
+                                   : NextLinkType::IsNext);
 
   return item;
 }
 
+TaskLocal::Item*
+TaskLocal::Item::createLink(AsyncTask *task,
+                            const HeapObject *key,
+                            const Metadata *valueType) {
+  return createLink(task, key, valueType, /*=inTaskGroupBody=*/false);
+}
+
+
+TaskLocal::Item*
+TaskLocal::Item::createLinkInTaskGroup(AsyncTask *task,
+                                       const HeapObject *key,
+                                       const Metadata *valueType) {
+  return createLink(task, key, valueType, /*=inTaskGroupBody=*/true);
+}
+
 
 void TaskLocal::Item::copyTo(AsyncTask *target) {
   assert(target && "TaskLocal item attempt to copy to null target task!");
@@ -232,13 +251,11 @@ void TaskLocal::Item::copyTo(AsyncTask *target) {
 
 SWIFT_CC(swift)
 static void swift_task_reportIllegalTaskLocalBindingWithinWithTaskGroupImpl(
-    const unsigned char *file, uintptr_t fileLength,
-    bool fileIsASCII, uintptr_t line) {
+    const unsigned char *_unused_file, uintptr_t _unused_fileLength,
+    bool _unused_fileIsASCII, uintptr_t _unused_line) {
 
-  char *message;
-  swift_asprintf(
-      &message,
-      "error: task-local: detected illegal task-local value binding at %.*s:%d.\n"
+  char *message =
+      "error: task-local: detected illegal task-local value binding.\n"
       "Task-local values must only be set in a structured-context, such as: "
       "around any (synchronous or asynchronous function invocation), "
       "around an 'async let' declaration, or around a 'with(Throwing)TaskGroup(...){ ... }' "
@@ -272,9 +289,7 @@ static void swift_task_reportIllegalTaskLocalBindingWithinWithTaskGroupImpl(
       "        }\n"
       "\n"
       "        group.addTask { ... }\n"
-      "    }\n",
-      (int)fileLength, file,
-      (int)line);
+      "    }\n";
 
   if (_swift_shouldReportFatalErrorsToDebugger()) {
     RuntimeErrorDetails details = {
@@ -329,6 +344,7 @@ void TaskLocal::Storage::destroy(AsyncTask *task) {
     auto linkType = item->getNextLinkType();
     switch (linkType) {
     case TaskLocal::NextLinkType::IsNext:
+    case TaskLocal::NextLinkType::IsNextCreatedInTaskGroupBody:
         next = item->getNext();
         item->destroy(task);
         item = next;
@@ -351,8 +367,25 @@ void TaskLocal::Storage::pushValue(AsyncTask *task,
                                    /* +1 */ OpaqueValue *value,
                                    const Metadata *valueType) {
   assert(value && "Task local value must not be nil");
+  assert(swift_task_getCurrent() == task &&
+         "must only be pushing task locals onto current task");
+
+  // We're in a task group body.
+  // We specifically need to prevent this pattern:
+  //
+  //    $number.withValue(0xBAADF00D) { // push
+  //      group.addTask { ... }
+  //    } // pop! BOOM!
+  //
+  // because the end of the withValue scope would pop the value,
+  // and thus if the child task didn't copy the value, it'd refer to a bad
+  // memory location at this point.
+
+  TaskLocal::Item* item = Item::createLink(
+      task, key, valueType,
+      /*inTaskGroupBody=*/swift_task_hasTaskGroupStatusRecord());
+
 
-  auto item = Item::createLink(task, key, valueType);
   valueType->vw_initializeWithTake(item->getStoragePtr(), value);
   head = item;
 }
@@ -367,6 +400,11 @@ bool TaskLocal::Storage::popValue(AsyncTask *task) {
   return head != nullptr;
 }
 
+TaskLocal::NextLinkType
+TaskLocal::Storage::peekHeadLinkType() const {
+  return head->getNextLinkType();
+}
+
 OpaqueValue* TaskLocal::Storage::getValue(AsyncTask *task,
                                           const HeapObject *key) {
   assert(key && "TaskLocal key must not be null.");

stdlib/public/Concurrency/TaskLocal.swift

@@ -252,9 +252,6 @@ public final class TaskLocal<Value: Sendable>: Sendable, CustomStringConvertible
                                  operation: () async throws -> R,
                                  isolation: isolated (any Actor)?,
                                  file: String = #fileID, line: UInt = #line) async rethrows -> R {
-    // check if we're not trying to bind a value from an illegal context; this may crash
-    _checkIllegalTaskLocalBindingWithinWithTaskGroup(file: file, line: line)
-
     _taskLocalValuePush(key: key, value: consume valueDuringOperation)
     defer { _taskLocalValuePop() }
 
@@ -269,9 +266,6 @@ public final class TaskLocal<Value: Sendable>: Sendable, CustomStringConvertible
   internal func withValueImpl<R>(_ valueDuringOperation: __owned Value,
                                  operation: () async throws -> R,
                                  file: String = #fileID, line: UInt = #line) async rethrows -> R {
-    // check if we're not trying to bind a value from an illegal context; this may crash
-    _checkIllegalTaskLocalBindingWithinWithTaskGroup(file: file, line: line)
-
     _taskLocalValuePush(key: key, value: consume valueDuringOperation)
     defer { _taskLocalValuePop() }
 
@@ -296,9 +290,6 @@ public final class TaskLocal<Value: Sendable>: Sendable, CustomStringConvertible
   @discardableResult
   public func withValue<R>(_ valueDuringOperation: Value, operation: () throws -> R,
                            file: String = #fileID, line: UInt = #line) rethrows -> R {
-    // check if we're not trying to bind a value from an illegal context; this may crash
-    _checkIllegalTaskLocalBindingWithinWithTaskGroup(file: file, line: line)
-
     _taskLocalValuePush(key: key, value: valueDuringOperation)
     defer { _taskLocalValuePop() }
 

test/Concurrency/Runtime/async_task_locals_prevent_illegal_use_but_specific_case_is_ok.swift

@@ -0,0 +1,39 @@
+// RUN: %target-run-simple-swift( -plugin-path %swift-plugin-dir -Xfrontend -disable-availability-checking -parse-as-library %import-libdispatch) 2>&1 | %FileCheck %s --dump-input=always
+
+// REQUIRES: executable_test
+// REQUIRES: concurrency
+// REQUIRES: libdispatch
+// REQUIRES: concurrency_runtime
+// UNSUPPORTED: back_deployment_runtime
+
+@available(SwiftStdlib 5.1, *)
+enum TL {
+  @TaskLocal
+  static var number: Int = 2
+}
+
+// ==== ------------------------------------------------------------------------
+
+func bindAroundGroupAddTask() async {
+  await TL.$number.withValue(1111) { // ok
+    await withTaskGroup(of: Int.self) { group in
+      // CHECK-NOT: error: task-local: detected illegal
+
+      TL.$number.withValue(2222) { // this is OK, there's no addTask being wrapped
+        print("Survived, inside withValue, value: \(TL.number)") // CHECK: Survived, inside withValue, value: 2222
+      }
+
+      group.addTask {
+        print("Survived, inside addTask, value: \(TL.number)") // CHECK: Survived, inside addTask, value: 1111
+        return TL.number
+      }
+    }
+    print("Survived, done") // CHECK: Survived, done
+  }
+}
+
+@main struct Main {
+  static func main() async {
+    await bindAroundGroupAddTask()
+  }
+}

test/Concurrency/Runtime/async_task_locals_prevent_illegal_use_discarding_taskgroup.swift

@@ -22,11 +22,11 @@ enum TL {
 func bindAroundGroupAddTask() async {
   await TL.$number.withValue(1111) { // ok
     await withTaskGroup(of: Int.self) { group in
-      // CHECK: error: task-local: detected illegal task-local value binding at {{.*}}illegal_use_discarding_taskgroup.swift:[[# @LINE + 1]]
+      // CHECK: error: task-local: detected illegal task-local value binding
       TL.$number.withValue(2222) { // bad!
-        print("Survived, inside withValue!") // CHECK-NOT: Survived, inside withValue!
         group.addTask {
-          0 // don't actually perform the read, it would be unsafe.
+          print("Survived, inside withValue!") // CHECK-NOT: Survived, inside withValue!
+          return 0 // don't actually perform the read, it would be unsafe.
         }
       }
 

test/Concurrency/Runtime/async_task_locals_prevent_illegal_use_taskgroup.swift

@@ -21,11 +21,11 @@ enum TL {
 func bindAroundDiscardingGroupAddTask() async {
   await TL.$number.withValue(1111) { // ok
     await withTaskGroup(of: Int.self) { group in
-      // CHECK: error: task-local: detected illegal task-local value binding at {{.*}}illegal_use_taskgroup.swift:[[# @LINE + 1]]
+      // CHECK: error: task-local: detected illegal task-local value binding
       TL.$number.withValue(2222) { // bad!
-        print("Survived, inside withValue!") // CHECK-NOT: Survived, inside withValue!
         group.addTask {
-          0 // don't actually perform the read, it would be unsafe.
+          print("Survived, inside withValue!") // CHECK-NOT: Survived, inside withValue!
+          return 0 // don't actually perform the read, it would be unsafe.
         }
       }