Implement disabling task-local values in isolated deinit

stdlib/public/Concurrency/Actor.cpp

@@ -2350,7 +2350,8 @@ static void swift_task_deinitOnExecutorImpl(void *object,
   // when running on the main thread without any executor.
   if (swift_task_isCurrentExecutorWithFlags(
           newExecutor, swift_task_is_current_executor_flag::None)) {
-    return work(object); // 'return' forces tail call
+    TaskLocal::StopLookupScope scope;
+    return work(object);
   }
 
 #if SWIFT_CONCURRENCY_ACTORS_AS_LOCKS
@@ -2384,7 +2385,10 @@ static void swift_task_deinitOnExecutorImpl(void *object,
       trackingInfo.enterAndShadow(newExecutor, taskExecutor);
 
       // Run the work.
-      work(object);
+      {
+        TaskLocal::StopLookupScope scope;
+        work(object);
+      }
 
       // `work` is a synchronous function, it cannot call swift_task_switch()
       // If it calls any synchronous API that may change executor inside

stdlib/public/Concurrency/TaskLocal.cpp

@@ -181,7 +181,7 @@ void TaskLocal::Storage::initializeLinkParent(AsyncTask* task,
   if (!item)
     return;
 
-  auto tail = TaskLocal::ParentTaskMarkerItem::create(task);
+  auto tail = TaskLocal::MarkerItem::createParentTaskMarker(task);
   head = tail;
 
   // Set of keys for which we already have copied to the new task.
@@ -234,18 +234,24 @@ void TaskLocal::Storage::initializeLinkParent(AsyncTask* task,
     item = item->getNext();
   }
 
+  if (item && item->getKind() == Item::Kind::StopLookupMarker) {
+    // Stop marker also could have been created inside task group body
+    // But we don't need to copy it. Instead we can break the chain.
+    item = nullptr;
+  }
+
   // The next item is not the "risky one" so we can directly link to it,
   // as we would have within normal child task relationships. E.g. this is
   // a parent or next pointer to a "safe" (withValue { withTaskGroup { ... } })
   // binding, so we re-link our current head to point at this item.
   tail->setNext(item);
 }
 
-TaskLocal::ParentTaskMarkerItem *
-TaskLocal::ParentTaskMarkerItem::create(AsyncTask *task) {
-  size_t amountToAllocate = sizeof(ParentTaskMarkerItem);
+TaskLocal::MarkerItem *TaskLocal::MarkerItem::create(AsyncTask *task,
+                                                     Item *next, Kind kind) {
+  size_t amountToAllocate = sizeof(MarkerItem);
   void *allocation = _swift_task_alloc_specific(task, amountToAllocate);
-  return new (allocation) ParentTaskMarkerItem(nullptr);
+  return new (allocation) MarkerItem(next, kind);
 }
 
 TaskLocal::ValueItem *TaskLocal::ValueItem::create(AsyncTask *task,
@@ -367,11 +373,12 @@ bool TaskLocal::Item::destroy(AsyncTask *task) {
     cast<ValueItem>(this)->~ValueItem();
     break;
   case Kind::ParentTaskMarker:
-    cast<ParentTaskMarkerItem>(this)->~ParentTaskMarkerItem();
-
     // 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.
     stop = true;
+    LLVM_FALLTHROUGH;
+  case Kind::StopLookupMarker:
+    cast<MarkerItem>(this)->~MarkerItem();
     break;
   }
 
@@ -442,6 +449,21 @@ bool TaskLocal::Storage::popValue(AsyncTask *task) {
   return head != nullptr;
 }
 
+void TaskLocal::Storage::pushStopLookup(AsyncTask *task) {
+  head = MarkerItem::createStopLookupMarker(task, head);
+  SWIFT_TASK_LOCAL_DEBUG_LOG(nullptr, "push stop node item:%p", head);
+}
+
+void TaskLocal::Storage::popStopLookup(AsyncTask *task) {
+  assert(head && "attempted to pop stop node off empty task-local stack");
+  assert(head->getKind() == Item::Kind::StopLookupMarker &&
+         "attempted to pop wrong node type");
+  auto old = head;
+  SWIFT_TASK_LOCAL_DEBUG_LOG(nullptr, "pop stop node item:%p", old);
+  head = head->getNext();
+  old->destroy(task);
+}
+
 OpaqueValue* TaskLocal::Storage::getValue(AsyncTask *task,
                                           const HeapObject *key) {
   assert(key && "TaskLocal key must not be null.");
@@ -452,6 +474,8 @@ OpaqueValue* TaskLocal::Storage::getValue(AsyncTask *task,
       if (valueItem->key == key) {
         return valueItem->getStoragePtr();
       }
+    } else if (item->getKind() == Item::Kind::StopLookupMarker) {
+      break;
     }
 
     item = item->getNext();
@@ -487,10 +511,30 @@ void TaskLocal::Storage::copyTo(AsyncTask *target) {
             "skip copy, already copied most recent value, value was [%p]",
             valueItem->getStoragePtr());
       }
+    } else if (item->getKind() == Item::Kind::StopLookupMarker) {
+      break;
     }
     item = item->getNext();
   }
 }
 
+TaskLocal::StopLookupScope::StopLookupScope() {
+  task = swift_task_getCurrent();
+  storage = Storage::getCurrent(task);
+  if (storage && storage->isEmpty()) {
+    storage = nullptr;
+  }
+
+  if (storage) {
+    storage->pushStopLookup(task);
+  }
+}
+
+TaskLocal::StopLookupScope::~StopLookupScope() {
+  if (storage) {
+    storage->popStopLookup(task);
+  }
+}
+
 #define OVERRIDE_TASK_LOCAL COMPATIBILITY_OVERRIDE
 #include "../CompatibilityOverride/CompatibilityOverrideIncludePath.h"

stdlib/public/Concurrency/TaskLocal.h

@@ -59,7 +59,12 @@ class TaskLocal {
       /// does not "contribute" any task local values. This is to speed up
       /// lookups by skipping empty parent tasks during get(), and explained
       /// in depth in `initializeLinkParent()`.
-      ParentTaskMarker = 2
+      ParentTaskMarker = 2,
+
+      /// Marker item that indicates that all earlier items should be ignored.
+      /// Allows to temporary disable all task local values in O(1).
+      /// Is used to disable task-locals for fast path of isolated deinit.
+      StopLookupMarker = 3,
     };
 
   private:
@@ -131,14 +136,22 @@ class TaskLocal {
     }
   };
 
-  class ParentTaskMarkerItem : public Item {
-    ParentTaskMarkerItem(Item *next) : Item(next, Kind::ParentTaskMarker) {}
+  class MarkerItem : public Item {
+    MarkerItem(Item *next, Kind kind) : Item(next, kind) {}
+
+    static MarkerItem *create(AsyncTask *task, Item *next, Kind kind);
 
   public:
-    static ParentTaskMarkerItem *create(AsyncTask *task);
+    static MarkerItem *createParentTaskMarker(AsyncTask *task) {
+      return create(task, nullptr, Kind::ParentTaskMarker);
+    }
+    static MarkerItem *createStopLookupMarker(AsyncTask *task, Item *next) {
+      return create(task, next, Kind::StopLookupMarker);
+    }
 
     static bool classof(const Item *item) {
-      return item->getKind() == Kind::ParentTaskMarker;
+      return item->getKind() == Kind::ParentTaskMarker ||
+             item->getKind() == Kind::StopLookupMarker;
     }
   };
 
@@ -185,6 +198,8 @@ class TaskLocal {
 
     void initializeLinkParent(AsyncTask *task, AsyncTask *parent);
 
+    bool isEmpty() const { return head == nullptr; }
+
     void pushValue(AsyncTask *task,
                    const HeapObject *key,
                    /* +1 */ OpaqueValue *value, const Metadata *valueType);
@@ -196,6 +211,9 @@ class TaskLocal {
     /// can be safely disposed of.
     bool popValue(AsyncTask *task);
 
+    void pushStopLookup(AsyncTask *task);
+    void popStopLookup(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.
@@ -213,6 +231,15 @@ class TaskLocal {
     /// Items owned by a parent task are left untouched, since we do not own them.
     void destroy(AsyncTask *task);
   };
+
+  class StopLookupScope {
+    AsyncTask *task;
+    Storage *storage;
+
+  public:
+    StopLookupScope();
+    ~StopLookupScope();
+  };
 };
 
 } // end namespace swift

test/Concurrency/Runtime/async_task_locals_isolated_deinit.swift

@@ -142,7 +142,7 @@ if #available(SwiftStdlib 5.1, *) {
       // FIXME: isolated deinit should be clearing task locals
       await TL.$number.withValue(42) {
         await AnotherActor.shared.performTesting {
-          preventAllocationOnStack(ClassNoOp(expectedNumber: 42, group: group))
+          preventAllocationOnStack(ClassNoOp(expectedNumber: 0, group: group))
         }
       }
     }
@@ -168,7 +168,7 @@ if #available(SwiftStdlib 5.1, *) {
       TL.$number.withValue(99) {
         // Despite last release happening not on the actor itself,
         // this is still a fast path due to optimisation for deallocating actors.
-        preventAllocationOnStack(ActorNoOp(expectedNumber: 99, group: group))
+        preventAllocationOnStack(ActorNoOp(expectedNumber: 0, group: group))
       }
     }
     group.wait()