diff --git a/include/swift/ABI/Task.h b/include/swift/ABI/Task.h index a587747658283..3fc89d0299aef 100644 --- a/include/swift/ABI/Task.h +++ b/include/swift/ABI/Task.h @@ -468,7 +468,10 @@ class AsyncTask : public Job { /// Check whether this task has been cancelled. /// Checking this is, of course, inherently race-prone on its own. - bool isCancelled() const; + /// + /// \param ignoreShield if cancellation shield should be ignored. + /// Cancellation shields prevent the observation of the isCancelled flag while active. + bool isCancelled(bool ignoreShield) const; // ==== INITIAL TASK RECORDS ================================================= // A task may have a number of "initial" records set, they MUST be set in the @@ -531,6 +534,12 @@ class AsyncTask : public Job { /// Returns true if storage has still more bindings. bool localValuePop(); + // ==== Cancellation Shields ------------------------------------------------- + + /// Returns true if the shield was installed and should be removed when leaving the shielded scope. + bool cancellationShieldPush(); + void cancellationShieldPop(); + // ==== Child Fragment ------------------------------------------------------- /// A fragment of an async task structure that happens to be a child task. diff --git a/include/swift/AST/Builtins.def b/include/swift/AST/Builtins.def index 826f8305a1392..801ad197f2616 100644 --- a/include/swift/AST/Builtins.def +++ b/include/swift/AST/Builtins.def @@ -1210,6 +1210,16 @@ BUILTIN_MISC_OPERATION(TaskLocalValuePush, "taskLocalValuePush", "", Special) /// Signature: () -> () BUILTIN_MISC_OPERATION(TaskLocalValuePop, "taskLocalValuePop", "", Special) +/// Equivalent to calling swift_task_cancellationShieldPush. +/// +/// Signature: () -> () +BUILTIN_MISC_OPERATION(TaskCancellationShieldPush, "taskCancellationShieldPush", "", Special) + +/// Equivalent to calling swift_task_cancellationShieldPop. +/// +/// Signature: () -> () +BUILTIN_MISC_OPERATION(TaskCancellationShieldPop, "taskCancellationShieldPop", "", Special) + #undef BUILTIN_TYPE_TRAIT_OPERATION #undef BUILTIN_UNARY_OPERATION #undef BUILTIN_BINARY_PREDICATE diff --git a/include/swift/Runtime/Concurrency.h b/include/swift/Runtime/Concurrency.h index 3e091c3976aca..d1c8a66055758 100644 --- a/include/swift/Runtime/Concurrency.h +++ b/include/swift/Runtime/Concurrency.h @@ -721,6 +721,12 @@ void swift_task_localValuePop(); SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift) void swift_task_localsCopyTo(AsyncTask* target); +SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift) +void swift_task_cancellationShieldPush(); + +SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift) +void swift_task_cancellationShieldPop(); + /// Switch the current task to a new executor if we aren't already /// running on a compatible executor. /// diff --git a/include/swift/Runtime/RuntimeFunctions.def b/include/swift/Runtime/RuntimeFunctions.def index 842e4c531a29e..32c93b3ffda3f 100644 --- a/include/swift/Runtime/RuntimeFunctions.def +++ b/include/swift/Runtime/RuntimeFunctions.def @@ -2734,6 +2734,26 @@ FUNCTION(TaskLocalValuePop, EFFECT(RuntimeEffect::Concurrency), UNKNOWN_MEMEFFECTS) +// void swift_task_cancellationShieldPush(); +FUNCTION(TaskCancellationShieldPush, + _Concurrency, swift_task_cancellationShieldPush, SwiftCC, + ConcurrencyAvailability, + RETURNS(), // TODO: return a bool here + ARGS(), + ATTRS(NoUnwind), + EFFECT(RuntimeEffect::Concurrency), + UNKNOWN_MEMEFFECTS) + +// void swift_task_cancellationShieldPop(); +FUNCTION(TaskCancellationShieldPop, + _Concurrency, swift_task_cancellationShieldPop, SwiftCC, + ConcurrencyAvailability, + RETURNS(), + ARGS(), + ATTRS(NoUnwind), + EFFECT(RuntimeEffect::Concurrency), + UNKNOWN_MEMEFFECTS) + // AutoDiffLinearMapContext *swift_autoDiffCreateLinearMapContextWithType(const Metadata *); FUNCTION(AutoDiffCreateLinearMapContextWithType, Swift, swift_autoDiffCreateLinearMapContextWithType, SwiftCC, diff --git a/include/swift/SIL/AddressWalker.h b/include/swift/SIL/AddressWalker.h index e9f7868bdaa6e..f0c3c77880f86 100644 --- a/include/swift/SIL/AddressWalker.h +++ b/include/swift/SIL/AddressWalker.h @@ -296,6 +296,8 @@ TransitiveAddressWalker::walk(SILValue projectedAddress) { case BuiltinValueKind::FlowSensitiveSelfIsolation: case BuiltinValueKind::FlowSensitiveDistributedSelfIsolation: case BuiltinValueKind::TaskLocalValuePush: + case BuiltinValueKind::TaskCancellationShieldPush: + case BuiltinValueKind::TaskCancellationShieldPop: callVisitUse(op); continue; default: diff --git a/lib/AST/Builtins.cpp b/lib/AST/Builtins.cpp index 0fd240a0eac62..94e1f991ea1e9 100644 --- a/lib/AST/Builtins.cpp +++ b/lib/AST/Builtins.cpp @@ -2410,6 +2410,14 @@ static ValueDecl *getTaskLocalValuePop(ASTContext &ctx, Identifier id) { return getBuiltinFunction(ctx, id, _thin, _parameters(), _void); } +static ValueDecl *getTaskCancellationShieldPush(ASTContext &ctx, Identifier id) { + return getBuiltinFunction(ctx, id, _thin, _parameters(), _void); +} + +static ValueDecl *getTaskCancellationShieldPop(ASTContext &ctx, Identifier id) { + return getBuiltinFunction(ctx, id, _thin, _parameters(), _void); +} + /// An array of the overloaded builtin kinds. static const OverloadedBuiltinKind OverloadedBuiltinKinds[] = { OverloadedBuiltinKind::None, @@ -3516,6 +3524,12 @@ ValueDecl *swift::getBuiltinValueDecl(ASTContext &Context, Identifier Id) { case BuiltinValueKind::TaskLocalValuePop: return getTaskLocalValuePop(Context, Id); + + case BuiltinValueKind::TaskCancellationShieldPush: + return getTaskCancellationShieldPush(Context, Id); + + case BuiltinValueKind::TaskCancellationShieldPop: + return getTaskCancellationShieldPop(Context, Id); } llvm_unreachable("bad builtin value!"); diff --git a/lib/IRGen/GenBuiltin.cpp b/lib/IRGen/GenBuiltin.cpp index e278a41911c69..d7e541574ce1b 100644 --- a/lib/IRGen/GenBuiltin.cpp +++ b/lib/IRGen/GenBuiltin.cpp @@ -1587,6 +1587,10 @@ void irgen::emitBuiltinCall(IRGenFunction &IGF, const BuiltinInfo &Builtin, auto *valueMetatype = IGF.emitTypeMetadataRef(argTypes[1].getASTType()); return emitBuiltinTaskLocalValuePush(IGF, key, value, valueMetatype); } + case BuiltinValueKind::TaskCancellationShieldPush: + return emitBuiltinTaskCancellationShieldPush(IGF); + case BuiltinValueKind::TaskCancellationShieldPop: + return emitBuiltinTaskCancellationShieldPop(IGF); // Builtins without IRGen implementations. case BuiltinValueKind::None: diff --git a/lib/IRGen/GenConcurrency.cpp b/lib/IRGen/GenConcurrency.cpp index 3f7f476daa291..3f9065f71e86b 100644 --- a/lib/IRGen/GenConcurrency.cpp +++ b/lib/IRGen/GenConcurrency.cpp @@ -368,6 +368,24 @@ void irgen::emitBuiltinTaskLocalValuePop(IRGenFunction &IGF) { call->setCallingConv(IGF.IGM.SwiftCC); } +void irgen::emitBuiltinTaskCancellationShieldPush(IRGenFunction &IGF) { + auto *call = + IGF.Builder.CreateCall(IGF.IGM.getTaskCancellationShieldPushFunctionPointer(), {}); + call->setDoesNotThrow(); + call->setCallingConv(IGF.IGM.SwiftCC); + // llvm::Value *identity = + // IGF.Builder.CreatePtrToInt(executor, IGF.IGM.ExecutorFirstTy); + + // out.add +} + +void irgen::emitBuiltinTaskCancellationShieldPop(IRGenFunction &IGF) { + auto *call = + IGF.Builder.CreateCall(IGF.IGM.getTaskCancellationShieldPopFunctionPointer(), {}); + call->setDoesNotThrow(); + call->setCallingConv(IGF.IGM.SwiftCC); +} + void irgen::emitFinishAsyncLet(IRGenFunction &IGF, llvm::Value *asyncLet, llvm::Value *resultBuffer) { diff --git a/lib/IRGen/GenConcurrency.h b/lib/IRGen/GenConcurrency.h index 062b4f74cbd9b..2d216c10aaca2 100644 --- a/lib/IRGen/GenConcurrency.h +++ b/lib/IRGen/GenConcurrency.h @@ -138,6 +138,10 @@ void emitBuiltinTaskLocalValuePush(IRGenFunction &IGF, llvm::Value *key, void emitBuiltinTaskLocalValuePop(IRGenFunction &IGF); +void emitBuiltinTaskCancellationShieldPush(IRGenFunction &IGF); + +void emitBuiltinTaskCancellationShieldPop(IRGenFunction &IGF); + } // end namespace irgen } // end namespace swift diff --git a/lib/SIL/IR/OperandOwnership.cpp b/lib/SIL/IR/OperandOwnership.cpp index b97d3de1f9a59..44556f021615a 100644 --- a/lib/SIL/IR/OperandOwnership.cpp +++ b/lib/SIL/IR/OperandOwnership.cpp @@ -1065,6 +1065,9 @@ BUILTIN_OPERAND_OWNERSHIP(TrivialUse, TaskRemovePriorityEscalationHandler) // second is an address to our generic Value. BUILTIN_OPERAND_OWNERSHIP(TrivialUse, TaskLocalValuePush) +BUILTIN_OPERAND_OWNERSHIP(TrivialUse, TaskCancellationShieldPush) +BUILTIN_OPERAND_OWNERSHIP(TrivialUse, TaskCancellationShieldPop) + #undef BUILTIN_OPERAND_OWNERSHIP #define SHOULD_NEVER_VISIT_BUILTIN(ID) \ @@ -1076,6 +1079,8 @@ BUILTIN_OPERAND_OWNERSHIP(TrivialUse, TaskLocalValuePush) SHOULD_NEVER_VISIT_BUILTIN(GetCurrentAsyncTask) SHOULD_NEVER_VISIT_BUILTIN(GetCurrentExecutor) SHOULD_NEVER_VISIT_BUILTIN(TaskLocalValuePop) +// SHOULD_NEVER_VISIT_BUILTIN(TaskCancellationShieldPush) +// SHOULD_NEVER_VISIT_BUILTIN(TaskCancellationShieldPop) #undef SHOULD_NEVER_VISIT_BUILTIN // Builtins that should be lowered to SIL instructions so we should never see diff --git a/lib/SIL/IR/ValueOwnership.cpp b/lib/SIL/IR/ValueOwnership.cpp index 3cb2dcc55f325..f96779385f6e6 100644 --- a/lib/SIL/IR/ValueOwnership.cpp +++ b/lib/SIL/IR/ValueOwnership.cpp @@ -670,6 +670,8 @@ CONSTANT_OWNERSHIP_BUILTIN(None, TaskAddPriorityEscalationHandler) CONSTANT_OWNERSHIP_BUILTIN(None, TaskRemovePriorityEscalationHandler) CONSTANT_OWNERSHIP_BUILTIN(None, TaskLocalValuePush) CONSTANT_OWNERSHIP_BUILTIN(None, TaskLocalValuePop) +CONSTANT_OWNERSHIP_BUILTIN(None, TaskCancellationShieldPush) +CONSTANT_OWNERSHIP_BUILTIN(None, TaskCancellationShieldPop) #undef CONSTANT_OWNERSHIP_BUILTIN diff --git a/lib/SILGen/SILGenBuiltin.cpp b/lib/SILGen/SILGenBuiltin.cpp index e1babae5c518a..98db348c22b10 100644 --- a/lib/SILGen/SILGenBuiltin.cpp +++ b/lib/SILGen/SILGenBuiltin.cpp @@ -2200,6 +2200,26 @@ static ManagedValue emitBuiltinTaskAddPriorityEscalationHandler( return ManagedValue::forRValueWithoutOwnership(b); } +// static ManagedValue emitBuiltinTaskCancellationShieldPush( +// SILGenFunction &SGF, SILLocation loc, SubstitutionMap subs, +// ArrayRef args, SGFContext C) { +// auto *b = +// SGF.B.createBuiltin(loc, BuiltinNames::TaskCancellationShieldPush, +// SILType::getUnsafeRawPointer(SGF.getASTContext()), +// subs, {}); +// return ManagedValue::forRValueWithoutOwnership(b); +// } + +// static ManagedValue emitBuiltinTaskCancellationShieldPop( +// SILGenFunction &SGF, SILLocation loc, SubstitutionMap subs, +// ArrayRef args, SGFContext C) { +// auto *b = +// SGF.B.createBuiltin(loc, BuiltinNames::TaskCancellationShieldPop, +// SILType::getUnsafeRawPointer(SGF.getASTContext()), +// subs, {}); +// return ManagedValue::forRValueWithoutOwnership(b); +// } + std::optional SpecializedEmitter::forDecl(SILGenModule &SGM, SILDeclRef function) { // Only consider standalone declarations in the Builtin module. diff --git a/lib/SILOptimizer/Transforms/AccessEnforcementReleaseSinking.cpp b/lib/SILOptimizer/Transforms/AccessEnforcementReleaseSinking.cpp index 1134d0f0f4b58..41927ef12eed9 100644 --- a/lib/SILOptimizer/Transforms/AccessEnforcementReleaseSinking.cpp +++ b/lib/SILOptimizer/Transforms/AccessEnforcementReleaseSinking.cpp @@ -220,6 +220,8 @@ static bool isBarrier(SILInstruction *inst) { case BuiltinValueKind::TaskRemovePriorityEscalationHandler: case BuiltinValueKind::TaskLocalValuePush: case BuiltinValueKind::TaskLocalValuePop: + case BuiltinValueKind::TaskCancellationShieldPush: + case BuiltinValueKind::TaskCancellationShieldPop: return true; } } diff --git a/stdlib/public/CompatibilityOverride/CompatibilityOverrideConcurrency.def b/stdlib/public/CompatibilityOverride/CompatibilityOverrideConcurrency.def index 9d4b20a72156c..12ae7ae0762d7 100644 --- a/stdlib/public/CompatibilityOverride/CompatibilityOverrideConcurrency.def +++ b/stdlib/public/CompatibilityOverride/CompatibilityOverrideConcurrency.def @@ -381,6 +381,14 @@ OVERRIDE_TASK_LOCAL(task_localsCopyTo, void, (AsyncTask *target), (target)) +OVERRIDE_TASK(task_cancellationShieldPush, void, + SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift), + swift::, , ) + +OVERRIDE_TASK(task_cancellationShieldPop, void, + SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift), + swift::, , ) + OVERRIDE_TASK_STATUS(task_hasTaskGroupStatusRecord, bool, SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift), swift::, , ) diff --git a/stdlib/public/Concurrency/Task.cpp b/stdlib/public/Concurrency/Task.cpp index 8011df6247211..04eb3344665d2 100644 --- a/stdlib/public/Concurrency/Task.cpp +++ b/stdlib/public/Concurrency/Task.cpp @@ -1840,6 +1840,22 @@ static void swift_task_removePriorityEscalationHandlerImpl( swift_task_dealloc(record); } +SWIFT_CC(swift) +static void swift_task_cancellationShieldPushImpl() { + if (AsyncTask *task = swift_task_getCurrent()) { + // return // TODO: impl bool return + auto pushed = task->cancellationShieldPush(); + } + // return false; // TODO: impl bool return +} + +SWIFT_CC(swift) +static void swift_task_cancellationShieldPopImpl() { + if (AsyncTask *task = swift_task_getCurrent()) { + task->cancellationShieldPop(); + } +} + SWIFT_CC(swift) static NullaryContinuationJob* swift_task_createNullaryContinuationJobImpl( diff --git a/stdlib/public/Concurrency/TaskCancellation.swift b/stdlib/public/Concurrency/TaskCancellation.swift index b4d453e72651c..179718e76b99d 100644 --- a/stdlib/public/Concurrency/TaskCancellation.swift +++ b/stdlib/public/Concurrency/TaskCancellation.swift @@ -163,3 +163,25 @@ public struct CancellationError: Error { // no extra information, cancellation is intended to be light-weight public init() {} } + +// ==== Task Cancellation Shielding ------------------------------------------- + +@available(SwiftStdlib 6.3, *) +public nonisolated(nonsending) func withTaskCancellationShield( + operation: nonisolated(nonsending) () async throws(E) -> T, +) async throws(E) -> T { + let record = unsafe Builtin.taskCancellationShieldPush() + defer { unsafe Builtin.taskCancellationShieldPop() } + + return try await operation() +} + +@available(SwiftStdlib 6.3, *) +public func withTaskCancellationShield( + operation: () throws(E) -> T, +) throws(E) -> T { + let record = unsafe Builtin.taskCancellationShieldPush() + defer { unsafe Builtin.taskCancellationShieldPop() } + + return try operation() +} diff --git a/stdlib/public/Concurrency/TaskPrivate.h b/stdlib/public/Concurrency/TaskPrivate.h index 5e2d353c98a2e..e49afa94feaca 100644 --- a/stdlib/public/Concurrency/TaskPrivate.h +++ b/stdlib/public/Concurrency/TaskPrivate.h @@ -486,6 +486,8 @@ class alignas(2 * sizeof(void*)) ActiveTaskStatus { /// use the task executor preference when we'd otherwise be running on /// the generic global pool. HasTaskExecutorPreference = 0x8000, + + HasActiveTaskCancellationShield = 0x10000, }; // Note: this structure is mirrored by ActiveTaskStatusWithEscalation and @@ -536,12 +538,32 @@ class alignas(2 * sizeof(void*)) ActiveTaskStatus { #endif /// Is the task currently cancelled? - bool isCancelled() const { return Flags & IsCancelled; } + /// This does take into account cancellation shields, i.e. while a shield is active this function will always return 'false'. + bool isCancelled(bool ignoreShield = false) const { + return (Flags & IsCancelled) && (ignoreShield || !(Flags & HasActiveTaskCancellationShield)); + } + bool isCancelledIgnoringShield() const { return Flags & IsCancelled; } ActiveTaskStatus withCancelled() const { #if SWIFT_CONCURRENCY_ENABLE_PRIORITY_ESCALATION return ActiveTaskStatus(Record, Flags | IsCancelled, ExecutionLock); #else return ActiveTaskStatus(Record, Flags | IsCancelled); +#endif + } + + bool hasCancellationShield() const { return Flags & HasActiveTaskCancellationShield; } + ActiveTaskStatus withCancellationShield() const { +#if SWIFT_CONCURRENCY_ENABLE_PRIORITY_ESCALATION + return ActiveTaskStatus(Record, Flags | HasActiveTaskCancellationShield, ExecutionLock); +#else + return ActiveTaskStatus(Record, Flags | HasActiveTaskCancellationShield); +#endif + } + ActiveTaskStatus withoutCancellationShield() const { +#if SWIFT_CONCURRENCY_ENABLE_PRIORITY_ESCALATION + return ActiveTaskStatus(Record, Flags & ~HasActiveTaskCancellationShield, ExecutionLock); +#else + return ActiveTaskStatus(Record, Flags & ~HasActiveTaskCancellationShield); #endif } @@ -965,9 +987,9 @@ inline const AsyncTask::PrivateStorage &AsyncTask::_private() const { return Private.get(); } -inline bool AsyncTask::isCancelled() const { +inline bool AsyncTask::isCancelled(bool ignoreShield = false) const { return _private()._status().load(std::memory_order_relaxed) - .isCancelled(); + .isCancelled(ignoreShield); } inline uint32_t AsyncTask::flagAsRunning() { @@ -1296,6 +1318,49 @@ inline bool AsyncTask::localValuePop() { return _private().Local.popValue(this); } +// ==== Cancellation Shields -------------------------------------------------- + +inline bool AsyncTask::cancellationShieldPush() { + while (true) { + auto oldStatus = _private()._status().load(std::memory_order_relaxed); + if (oldStatus.hasCancellationShield()) { + return false; + } + + auto newStatus = oldStatus.withCancellationShield(); + assert(newStatus.hasCancellationShield()); + + // if (oldStatus == newStatus) { + // // we already had a cancellation shield active, no-op + // return false; + // } + + if (_private()._status().compare_exchange_weak(oldStatus, newStatus, + /* success */ std::memory_order_relaxed, + /* failure */ std::memory_order_relaxed)) { + return true; // we did successfully install the shield + } + } +} + +inline void AsyncTask::cancellationShieldPop() { + while (true) { + auto oldStatus = _private()._status().load(std::memory_order_relaxed); + if (!oldStatus.hasCancellationShield()) { + return; + } + + auto newStatus = oldStatus.withoutCancellationShield(); + assert(!newStatus.hasCancellationShield()); + + if (_private()._status().compare_exchange_weak(oldStatus, newStatus, + /* success */ std::memory_order_relaxed, + /* failure */ std::memory_order_relaxed)) { + return; + } + } +} + } // end namespace swift #endif diff --git a/stdlib/public/Concurrency/TaskStatus.cpp b/stdlib/public/Concurrency/TaskStatus.cpp index c4c5d849161c9..25fc5789f14b1 100644 --- a/stdlib/public/Concurrency/TaskStatus.cpp +++ b/stdlib/public/Concurrency/TaskStatus.cpp @@ -799,7 +799,7 @@ void swift::_swift_taskGroup_detachChild(TaskGroup *group, /// The caller must guarantee that this is called while holding the owning /// task's status record lock. void swift::_swift_taskGroup_cancel(TaskGroup *group) { - (void) group->statusCancel(); + (void) group->statusCancel(); // TODO: do prevent the task group from being cancelled? I think probably no, we cancel and "don't observe", same as Task // Because only the owning task of the task group can modify the // child list of a task group status record, and it can only do so @@ -830,7 +830,7 @@ void swift::_swift_taskGroup_cancel_unlocked(TaskGroup *group, /**************************************************************************/ /// Perform any cancellation actions required by the given record. -static void performCancellationAction(TaskStatusRecord *record) { +static void performCancellationAction(ActiveTaskStatus status, TaskStatusRecord *record) { switch (record->getKind()) { // Child tasks need to be recursively cancelled. case TaskStatusRecordKind::ChildTask: { @@ -853,6 +853,10 @@ static void performCancellationAction(TaskStatusRecord *record) { case TaskStatusRecordKind::CancellationNotification: { auto notification = cast(record); + if (status.hasCancellationShield()) { + SWIFT_TASK_DEBUG_LOG("cancellation shielded: skip cancellation handler invocation in task = %p", swift_task_getCurrent()); + return; + } notification->run(); return; } @@ -889,7 +893,9 @@ static void swift_task_cancelImpl(AsyncTask *task) { auto oldStatus = task->_private()._status().load(std::memory_order_relaxed); auto newStatus = oldStatus; while (true) { - if (oldStatus.isCancelled()) { + // Are we already cancelled? + // Even if we have a cancellation shield active, we do want to set the isCancelled flag. + if (oldStatus.isCancelled(/*ignoreShield=*/false)) { return; } @@ -907,8 +913,8 @@ static void swift_task_cancelImpl(AsyncTask *task) { newStatus.traceStatusChanged(task, false, oldStatus.isRunning()); if (newStatus.getInnermostRecord() == nullptr) { - // No records, nothing to propagate - return; + // No records, nothing to propagate + return; } withStatusRecordLock(task, newStatus, [&](ActiveTaskStatus status) { @@ -917,7 +923,9 @@ static void swift_task_cancelImpl(AsyncTask *task) { // modify this list that is being iterated. However, cancellation is // happening from outside of the task so we know that no new records will // be added since that's only possible while on task. - performCancellationAction(cur); + // + // Each action must independently take care of how to deal with cancellation shields. + performCancellationAction(newStatus, cur); } }); } diff --git a/stdlib/toolchain/Compatibility56/CompatibilityOverrideConcurrency.def b/stdlib/toolchain/Compatibility56/CompatibilityOverrideConcurrency.def index 8d53124180875..60ef678235bb2 100644 --- a/stdlib/toolchain/Compatibility56/CompatibilityOverrideConcurrency.def +++ b/stdlib/toolchain/Compatibility56/CompatibilityOverrideConcurrency.def @@ -307,6 +307,14 @@ OVERRIDE_TASK_LOCAL(task_localsCopyTo, void, (AsyncTask *target), (target)) +OVERRIDE_TASK(task_cancellationShieldPush, void, + SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift), + swift::, ,) + +OVERRIDE_TASK(task_cancellationShieldPop, void, + SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift), + swift::, ,) + OVERRIDE_TASK_STATUS(task_addStatusRecord, bool, SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift), swift::, (TaskStatusRecord *newRecord), (newRecord)) diff --git a/test/Concurrency/Runtime/async_task_cancellation_shield.swift b/test/Concurrency/Runtime/async_task_cancellation_shield.swift new file mode 100644 index 0000000000000..c675b41e0e0f9 --- /dev/null +++ b/test/Concurrency/Runtime/async_task_cancellation_shield.swift @@ -0,0 +1,270 @@ +// RUN: %target-run-simple-swift( -Xfrontend -disable-availability-checking %import-libdispatch -parse-as-library) | %FileCheck %s --dump-input=always + +// REQUIRES: executable_test +// REQUIRES: concurrency +// REQUIRES: synchronization + +// rdar://76038845 +// REQUIRES: concurrency_runtime +// UNSUPPORTED: back_deployment_runtime + +import Synchronization + +@available(SwiftStdlib 6.3, *) +func test_task_cancel_shield() async { + print(#function) // CHECK: test_task_cancel_shield + + let t = Task { + print("Inside task, before cancel: isCancelled:\(Task.isCancelled)") + // CHECK: Inside task, before cancel: isCancelled:false + + print("cancel self") + withUnsafeCurrentTask { $0?.cancel() } // cancel myself! + + print("Inside task, after cancel: isCancelled:\(Task.isCancelled)") + // CHECK: Inside task, after cancel: isCancelled:true + + withTaskCancellationShield { + print("Inside task, shielded: isCancelled:\(Task.isCancelled)") + // CHECK: Inside task, shielded: isCancelled:false + } + print("Inside task, after shield: isCancelled:\(Task.isCancelled)") + // CHECK: Inside task, after shield: isCancelled:true + + await withTaskCancellationShield { + await withTaskCancellationHandler { + print("shielded, withTaskCancellationHandler - operation") + // CHECK: shielded, withTaskCancellationHandler - operation + + } onCancel: { + // MUST NOT execute because we're shielded + print("shielded, withTaskCancellationHandler - onCancel") + // CHECK-NOT: shielded, withTaskCancellationHandler - onCancel + } + } + + async let child = { + print("Inside child-task: isCancelled:\(Task.isCancelled)") + // CHECK: Inside child-task: isCancelled:true + withTaskCancellationShield { + print("Inside child-task, shielded: isCancelled:\(Task.isCancelled)") + // CHECK: Inside child-task, shielded: isCancelled:false + } + }() + _ = await child + + let t2 = Task { + let unsafeT2 = withUnsafeCurrentTask { $0! } // escape self reference unsafely + + await withTaskCancellationShield { + await withTaskCancellationHandler { + print("cancel self") + unsafeT2.cancel() // cancel self, but we're shielded, so the handler MUST NOT run anyway + // CHECK-NOT: Task{}.cancel, withTaskCancellationHandler - onCancel + print("Task{}.cancel, shielded, withTaskCancellationHandler - operation") + // CHECK: Task{}.cancel, shielded, withTaskCancellationHandler - operation + } onCancel: { + // MUST NOT execute because we're shielded + print("Task{}.cancel, withTaskCancellationHandler - onCancel") + // CHECK-NOT: Task{}.cancel, withTaskCancellationHandler - onCancel + } + } + } + await t2.value + } + + await t.value +} + +@available(SwiftStdlib 6.3, *) +func test_defer_cancel_shield() async { + let task = Task { + withUnsafeCurrentTask { $0?.cancel() } + + print("Inside task, before cancel: isCancelled:\(Task.isCancelled)") + // CHECK: Inside task, before cancel: isCancelled:true + + defer { + defer { + print("Inside defer, after cancel: isCancelled:\(Task.isCancelled)") + // CHECK: Inside defer, after cancel: isCancelled:true + + withTaskCancellationShield { + print("Inside defer, shielded, after cancel: isCancelled:\(Task.isCancelled)") + // CHECK: Inside defer, shielded, after cancel: isCancelled:false + } + + print("Inside defer, after cancel: isCancelled:\(Task.isCancelled)") + // CHECK: Inside defer, after cancel: isCancelled:true + } + } + } + + await task.value +} + +@available(SwiftStdlib 6.3, *) +func test_nested_shields() async { + print(#function) // CHECK: test_nested_shields + + let task = Task { + withUnsafeCurrentTask { $0?.cancel() } + + withTaskCancellationShield { + print("Inside shield, after cancel: isCancelled:\(Task.isCancelled)") + // CHECK: Inside shield, after cancel: isCancelled:false + + withTaskCancellationShield { + print("Inside shield shield, after cancel: isCancelled:\(Task.isCancelled)") + // CHECK: Inside shield shield, after cancel: isCancelled:false + } + + // make sure we didn't remove the "outer" shield by accident + print("Popped shield, still shielded, after cancel: isCancelled:\(Task.isCancelled)") + // CHECK: Popped shield, still shielded, after cancel: isCancelled:false + } + + print("Inside task, before cancel: isCancelled:\(Task.isCancelled)") + // CHECK: Inside task, before cancel: isCancelled:false + } + + await task.value +} + +@available(SwiftStdlib 6.3, *) +func test_sleep_cancel_shield() async { + print(#function) // CHECK: test_sleep_cancel_shield + + let task = Task { + await withTaskCancellationShield { + try! await Task.sleep(for: .seconds(1)) + } + } + + task.cancel() + await task.value +} + +@available(SwiftStdlib 6.3, *) +func test_async_stream_cancel_shield() async { + print(#function) // CHECK: test_async_stream_cancel_shield + + let (stream, continuation) = AsyncThrowingStream.makeStream() + let task = Task { + try await withTaskCancellationShield { + for try await value in stream { + print("Inside loop") + // CHECK: Inside loop + } + } + } + + task.cancel() + continuation.yield(()) + continuation.finish() + try! await task.value +} + +@available(SwiftStdlib 6.3, *) +func test_child_task_cancel_shield() async { + print(#function) // CHECK: test_child_task_cancel_shield + + let cancellableContinuation = CancellableContinuation() + await withTaskGroup { group in + group.addTask { + await withTaskCancellationShield { + try! await cancellableContinuation.wait() + } + } + + group.cancelAll() + cancellableContinuation.complete() + } +} + +@available(SwiftStdlib 6.3, *) +func test_task_group_cancel_shield() async { + print(#function) // CHECK: test_task_group_cancel_shield + + await withTaskGroup(of: Void.self) { group in + group.cancelAll() + group.addTask { + await withTaskCancellationShield { + print("Inside cancellation shield, isCancelled:\(Task.isCancelled)") + // CHECK: Inside cancellation shield, isCancelled:false + + await withTaskGroup { group in + group.addTask { + print("Inside child task, isCancelled:\(Task.isCancelled)") + // CHECK: Inside child task, isCancelled:false + } + print("Inside task group, isCancelled:\(Task.isCancelled)") + // CHECK: Inside task group, isCancelled:false + } + } + } + } +} + +@available(SwiftStdlib 6.3, *) +func test_add_task_cancel_shield() async { + print(#function) // CHECK: test_add_task_cancel_shield + + await withTaskGroup(of: Void.self) { group in + group.cancelAll() + // The cancellation shield is only protecting the addTask but not the child task + withTaskCancellationShield { + group.addTask { + print("Inside child task, isCancelled:\(Task.isCancelled)") + // CHECK: Inside child task, isCancelled:true + } + } + group.addTask { + print("Inside child task, isCancelled:\(Task.isCancelled)") + // CHECK: Inside child task, isCancelled:true + } + } +} + +@available(SwiftStdlib 6.3, *) +struct CancellableContinuation: ~Copyable { + let state = Mutex<(Bool, (CheckedContinuation?))>((false, nil)) + + func wait() async throws { + try await withTaskCancellationHandler { + try await withCheckedThrowingContinuation { continuation in + self.state.withLock { state -> CheckedContinuation? in + if state.0 { + return continuation + } else { + state.1 = continuation + return nil + } + }?.resume(throwing: CancellationError()) + } + } onCancel: { + self.state.withLock { state in + state.0 = true + return state.1 + }?.resume(throwing: CancellationError()) + } + } + + func complete() { + self.state.withLock { $0.1 }?.resume() + } +} + +@available(SwiftStdlib 6.3, *) +@main struct Main { + static func main() async { + await test_task_cancel_shield() + await test_defer_cancel_shield() + await test_sleep_cancel_shield() + await test_nested_shields() + // await test_async_stream_cancel_shield() + // await test_child_task_cancel_shield() + // await test_task_group_cancel_shield() + // await test_add_task_cancel_shield() + } +} diff --git a/test/Distributed/distributed_package_import.swift b/test/Distributed/distributed_package_import.swift new file mode 100644 index 0000000000000..7b3cb4ef6ab92 --- /dev/null +++ b/test/Distributed/distributed_package_import.swift @@ -0,0 +1,14 @@ +// RUN: %empty-directory(%t) +// RUN: %target-swift-frontend-emit-module -emit-module-path %t/FakeDistributedActorSystems.swiftmodule -module-name FakeDistributedActorSystems -target %target-swift-5.7-abi-triple %S/Inputs/FakeDistributedActorSystems.swift +// RUN: %target-swift-frontend -typecheck -abi -module MyModule -verify -target %target-swift-5.7-abi-triple -I %t 2>&1 %s +// REQUIRES: concurrency +// REQUIRES: distributed + +// Should NOT produce a warning, we do use the module in a public declaration, see MyPublicDistributedActor +package import Distributed +import FakeDistributedActorSystems + + +package distributed actor MyPublicDistributedActor { + package typealias ActorSystem = FakeActorSystem +} diff --git a/test/IRGen/builtins.swift b/test/IRGen/builtins.swift index e478f0a86be87..6e169e89cbe9e 100644 --- a/test/IRGen/builtins.swift +++ b/test/IRGen/builtins.swift @@ -985,6 +985,20 @@ nonisolated(nonsending) func testTaskLocalValuePush(_ key: Builtin.RawPoi // CHECK: call swiftcc {} @swift_task_localValuePop() nonisolated(nonsending) func testTaskLocalValuePop() async { Builtin.taskLocalValuePop() +} + +// CHECK-LABEL: define {{.*}}void @"$s8builtins22testTaskLocalValuePushyyBp_xntYalF"(ptr swiftasync %0, i64 %1, i64 %2, ptr %3, ptr noalias %4, ptr %Value) +// CHECK: [[TASK_VAR:%.*]] = call swiftcc ptr @swift_task_alloc({{.*}} +// CHECK: call ptr %InitializeWithTake(ptr noalias [[TASK_VAR]], ptr noalias {{%.*}}, ptr %Value) +// CHECK: call swiftcc {} @swift_task_localValuePush(ptr %3, ptr [[TASK_VAR]], ptr %Value) +nonisolated(nonsending) func testTaskCancellationShieldPush(_ key: Builtin.RawPointer, _ value: consuming Value) async { + Builtin.taskCancellationShieldPush() +} + +// CHECK-LABEL: define {{.*}}void @"$s8builtins21testTaskLocalValuePopyyYaF"(ptr swiftasync %0, i64 %1, i64 %2) +// CHECK: call swiftcc {} @swift_task_localValuePop() +nonisolated(nonsending) func taskCancellationShieldPop() async { + Builtin.taskCancellationShieldPop() } // CHECK: ![[R]] = !{i64 0, i64 9223372036854775807} diff --git a/test/SILGen/builtins.swift b/test/SILGen/builtins.swift index cf9fe9f2dbdbf..99b5c5497154f 100644 --- a/test/SILGen/builtins.swift +++ b/test/SILGen/builtins.swift @@ -979,3 +979,17 @@ func testTaskLocalValuePush(_ key: Builtin.RawPointer, _ value: consuming func testTaskLocalValuePop() async { Builtin.taskLocalValuePop() } + +// CHECK-LABEL: sil hidden [ossa] @$s8builtins21taskCancellationShieldPushyyYaF : $@convention(thin) @async () -> () { +// CHECK: builtin "taskCancellationShieldPush"() : $() +// CHECK: } // end sil function '$s8builtins21taskCancellationShieldPushyyYaF' +func taskCancellationShieldPush() async { + Builtin.taskCancellationShieldPush() +} + +// CHECK-LABEL: sil hidden [ossa] @$s8builtins21taskCancellationShieldPopyyYaF : $@convention(thin) @async () -> () { +// CHECK: builtin "taskCancellationShieldPop"() : $() +// CHECK: } // end sil function '$s8builtins21taskCancellationShieldPopyyYaF' +func taskCancellationShieldPop() async { + Builtin.taskCancellationShieldPop() +} diff --git a/test/abi/macOS/arm64/concurrency.swift b/test/abi/macOS/arm64/concurrency.swift index f847cef841826..3706dc446491e 100644 --- a/test/abi/macOS/arm64/concurrency.swift +++ b/test/abi/macOS/arm64/concurrency.swift @@ -400,6 +400,10 @@ Added: _$sScfsE25isIsolatingCurrentContextSbSgyF Added: _$sScf25isIsolatingCurrentContextSbSgyFTj Added: _$sScf25isIsolatingCurrentContextSbSgyFTq +// Cancellation shields +Added: _swift_task_cancellationShieldPush +Added: _swift_task_cancellationShieldPop + // CoroutineAccessors Added: _swift_task_dealloc_through diff --git a/test/abi/macOS/x86_64/concurrency.swift b/test/abi/macOS/x86_64/concurrency.swift index b826dec8c0dbe..5a03c42161ba0 100644 --- a/test/abi/macOS/x86_64/concurrency.swift +++ b/test/abi/macOS/x86_64/concurrency.swift @@ -400,6 +400,10 @@ Added: _$sScfsE25isIsolatingCurrentContextSbSgyF Added: _$sScf25isIsolatingCurrentContextSbSgyFTj Added: _$sScf25isIsolatingCurrentContextSbSgyFTq +// Cancellation shields +Added: _swift_task_cancellationShieldPush +Added: _swift_task_cancellationShieldPop + // CoroutineAccessors Added: _swift_task_dealloc_through