diff --git a/stdlib/public/Concurrency/Actor.cpp b/stdlib/public/Concurrency/Actor.cpp index 8673b72be8df6..8f82f5c0f33fa 100644 --- a/stdlib/public/Concurrency/Actor.cpp +++ b/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 diff --git a/stdlib/public/Concurrency/TaskLocal.cpp b/stdlib/public/Concurrency/TaskLocal.cpp index de7dfa00067b0..1e5e363de1d10 100644 --- a/stdlib/public/Concurrency/TaskLocal.cpp +++ b/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,6 +234,12 @@ 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 { ... } }) @@ -241,11 +247,11 @@ void TaskLocal::Storage::initializeLinkParent(AsyncTask* task, 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(this)->~ValueItem(); break; case Kind::ParentTaskMarker: - cast(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(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" diff --git a/stdlib/public/Concurrency/TaskLocal.h b/stdlib/public/Concurrency/TaskLocal.h index d6eb3e622172a..dc861250fc336 100644 --- a/stdlib/public/Concurrency/TaskLocal.h +++ b/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 diff --git a/test/Concurrency/Runtime/async_task_locals_isolated_deinit.swift b/test/Concurrency/Runtime/async_task_locals_isolated_deinit.swift index 5c0e1b5728f93..8b8d25c8413af 100644 --- a/test/Concurrency/Runtime/async_task_locals_isolated_deinit.swift +++ b/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()