[Concurrency] Retain the actor around the CAS in enqueue() when necessary.

It's possible that the job we enqueue holds the last strong reference to the actor. If that job runs on another thread after we enqueue it, then it's possible for `this` to be destroyed while we're still in this function. We need to use `this` after the enqueue when the priorities don't match. When it looks like that will happen, retain `this` before the enqueue to ensure it stays alive until we're done with it.

Introduce a defensive retain helper class that makes it easy to do a single retain under certain conditions even in a loop, and does RAII to balance it with a release when the scope exits.

rdar://135400933

Author: mikeash

include/swift/Runtime/HeapObject.h

@@ -386,6 +386,57 @@ class SwiftRAII {
   HeapObject *operator *() const { return object; }
 };
 
+/// RAII object that wraps a Swift object and optionally performs a single
+/// retain on that object. Multiple requests to retain the object only perform a
+/// single retain, and if that retain has been done then it's automatically
+/// released when leaving the scope. This helps implement a defensive retain
+/// pattern where you may need to retain an object in some circumstances. This
+/// helper makes it easy to retain the object only once even when loops are
+/// involved, and do a release to balance the retain on all paths out of the
+/// scope.
+class SwiftDefensiveRetainRAII {
+  HeapObject *object;
+  bool didRetain;
+
+public:
+  // Noncopyable.
+  SwiftDefensiveRetainRAII(const SwiftDefensiveRetainRAII &) = delete;
+  SwiftDefensiveRetainRAII &operator=(const SwiftDefensiveRetainRAII &) = delete;
+
+  /// Create a new helper with the given object. The object is not retained
+  /// initially.
+  SwiftDefensiveRetainRAII(HeapObject *object)
+      : object(object), didRetain(false) {}
+
+  ~SwiftDefensiveRetainRAII() {
+    if (didRetain)
+      swift_release(object);
+  }
+
+  /// Perform a defensive retain of the object. If a defensive retain has
+  /// already been performed, this is a no-op.
+  void defensiveRetain() {
+    if (!didRetain) {
+      swift_retain(object);
+      didRetain = true;
+    }
+  }
+
+  /// Take the retain from the helper. This is an optimization for code paths
+  /// that want to retain the object long-term, and avoids doing a redundant
+  /// retain/release pair. If a defensive retain has not been done, then this
+  /// will retain the object, so the caller always gets a +1 on the object.
+  void takeRetain() {
+    if (!didRetain)
+      swift_retain(object);
+    didRetain = false;
+  }
+
+  /// Returns true if the object was defensively retained (and takeRetain not
+  /// called). Intended for use in asserts.
+  bool isRetained() { return didRetain; }
+};
+
 /*****************************************************************************/
 /**************************** UNOWNED REFERENCES *****************************/
 /*****************************************************************************/

stdlib/public/Concurrency/Actor.cpp

@@ -1357,6 +1357,7 @@ void DefaultActorImpl::enqueue(Job *job, JobPriority priority) {
   bool distributedActorIsRemote = swift_distributed_actor_is_remote(this);
 
   auto oldState = _status().load(std::memory_order_relaxed);
+  SwiftDefensiveRetainRAII thisRetainHelper{this};
   while (true) {
     auto newState = oldState;
 
@@ -1386,15 +1387,29 @@ void DefaultActorImpl::enqueue(Job *job, JobPriority priority) {
     if (needsScheduling || needsStealer)
       taskExecutor = TaskExecutorRef::fromTaskExecutorPreference(job);
 
+    // In some cases (we aren't scheduling the actor and priorities don't
+    // match) then we need to access `this` after the enqueue. But the enqueue
+    // can cause the job to run and release `this`, so we need to retain `this`
+    // in those cases. The conditional here matches the conditions where we can
+    // get to the code below that uses `this`.
+    bool willSchedule = !oldState.isScheduled() && newState.isScheduled();
+    bool priorityMismatch = oldState.getMaxPriority() != newState.getMaxPriority();
+    if (!willSchedule && priorityMismatch)
+        thisRetainHelper.defensiveRetain();
+
     // This needs to be a store release so that we also publish the contents of
     // the new Job we are adding to the atomic job queue. Pairs with consume
     // in drainOne.
     if (_status().compare_exchange_weak(oldState, newState,
                    /* success */ std::memory_order_release,
                    /* failure */ std::memory_order_relaxed)) {
       // NOTE: `job` is off limits after this point, as another thread might run
-      // and destroy it now that it's enqueued.
+      // and destroy it now that it's enqueued. `this` is only accessible if
+      // `retainedThis` is true.
+      job = nullptr; // Ensure we can't use it accidentally.
 
+      // NOTE: only the pointer value of `this` is used here, so this one
+      // doesn't need a retain.
       traceActorStateTransition(this, oldState, newState, distributedActorIsRemote);
 
       if (!oldState.isScheduled() && newState.isScheduled()) {
@@ -1405,6 +1420,9 @@ void DefaultActorImpl::enqueue(Job *job, JobPriority priority) {
 
 #if SWIFT_CONCURRENCY_ENABLE_PRIORITY_ESCALATION
       if (oldState.getMaxPriority() != newState.getMaxPriority()) {
+        // We still need `this`, assert that we did a defensive retain.
+        assert(thisRetainHelper.isRetained());
+
         if (newState.isRunning()) {
           // Actor is running on a thread, escalate the thread running it
           SWIFT_TASK_DEBUG_LOG("[Override] Escalating actor %p which is running on %#x to %#x priority", this, newState.currentDrainer(), priority);
@@ -1414,11 +1432,12 @@ void DefaultActorImpl::enqueue(Job *job, JobPriority priority) {
         } else {
           // We are scheduling a stealer for an actor due to priority override.
           // This extra processing job has a reference on the actor. See
-          // ownership rule (2).
+          // ownership rule (2). That means that we need to retain `this`, which
+          // we'll take from the retain helper.
+          thisRetainHelper.takeRetain();
           SWIFT_TASK_DEBUG_LOG(
               "[Override] Scheduling a stealer for actor %p at %#x priority",
               this, newState.getMaxPriority());
-          swift_retain(this);
 
           scheduleActorProcessJob(newState.getMaxPriority(), taskExecutor);
         }