[WIP] Task Cancellation Shields

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.

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

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.
 ///

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,

include/swift/SIL/AddressWalker.h

@@ -296,6 +296,8 @@ TransitiveAddressWalker<Impl>::walk(SILValue projectedAddress) {
         case BuiltinValueKind::FlowSensitiveSelfIsolation:
         case BuiltinValueKind::FlowSensitiveDistributedSelfIsolation:
         case BuiltinValueKind::TaskLocalValuePush:
+        case BuiltinValueKind::TaskCancellationShieldPush:
+        case BuiltinValueKind::TaskCancellationShieldPop:
           callVisitUse(op);
           continue;
         default:

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!");

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:

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) {

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
 

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

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
 

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<ManagedValue> 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<ManagedValue> args, SGFContext C) {
+//   auto *b =
+//       SGF.B.createBuiltin(loc, BuiltinNames::TaskCancellationShieldPop,
+//                           SILType::getUnsafeRawPointer(SGF.getASTContext()),
+//                           subs, {});
+//   return ManagedValue::forRValueWithoutOwnership(b);
+// }
+
 std::optional<SpecializedEmitter>
 SpecializedEmitter::forDecl(SILGenModule &SGM, SILDeclRef function) {
   // Only consider standalone declarations in the Builtin module.

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;
     }
   }

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::, , )

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(

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<T, E>(
+  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<T, E>(
+  operation: () throws(E) -> T,
+) throws(E) -> T {
+  let record = unsafe Builtin.taskCancellationShieldPush()
+  defer { unsafe Builtin.taskCancellationShieldPop() }
+
+  return try operation()
+}

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