Clear the active task immediately when suspending.

Author: rjmccall

stdlib/public/Concurrency/Actor.cpp

@@ -239,8 +239,8 @@ void swift::runJobInEstablishedExecutorContext(Job *job) {
 
     task->runInFullyEstablishedContext();
 
-    // Clear the active task.
-    ActiveTask::set(nullptr);
+    assert(ActiveTask::get() == nullptr &&
+           "active task wasn't cleared before susspending?");
   } else {
     // There's no extra bookkeeping to do for simple jobs.
     job->runSimpleInFullyEstablishedContext();
@@ -1937,6 +1937,7 @@ static void swift_task_switchImpl(SWIFT_ASYNC_CONTEXT AsyncContext *resumeContex
           _swift_get_thread_id(), task, newExecutor.getIdentity());
 #endif
   task->flagAsSuspended();
+  _swift_task_clearCurrent();
   swift_task_enqueue(task, newExecutor);
 }
 

stdlib/public/Concurrency/Task.cpp

@@ -124,6 +124,7 @@ FutureFragment::Status AsyncTask::waitFuture(AsyncTask *waitingTask,
       // Escalate the priority of this task based on the priority
       // of the waiting task.
       swift_task_escalate(this, waitingTask->Flags.getPriority());
+      _swift_task_clearCurrent();
       return FutureFragment::Status::Executing;
     }
   }
@@ -937,7 +938,7 @@ size_t swift::swift_task_getJobFlags(AsyncTask *task) {
 
 SWIFT_CC(swift)
 static AsyncTask *swift_task_suspendImpl() {
-  auto task = swift_task_getCurrent();
+  auto task = _swift_task_clearCurrent();
   task->flagAsSuspended();
   return task;
 }
@@ -963,16 +964,21 @@ static AsyncTask *swift_continuation_initImpl(ContinuationAsyncContext *context,
                                         : ContinuationStatus::Pending, 
                                       std::memory_order_relaxed);
 
-  auto task = swift_task_getCurrent();
-  assert(task && "initializing a continuation with no current task");
-  task->ResumeContext = context;
-  task->ResumeTask = context->ResumeParent;
+  AsyncTask *task;
 
   // A preawait immediately suspends the task.
   if (flags.isPreawaited()) {
+    task = _swift_task_clearCurrent();
+    assert(task && "initializing a continuation with no current task");
     task->flagAsSuspended();
+  } else {
+    task = swift_task_getCurrent();
+    assert(task && "initializing a continuation with no current task");
   }
 
+  task->ResumeContext = context;
+  task->ResumeTask = context->ResumeParent;
+
   return task;
 }
 
@@ -1016,7 +1022,10 @@ static void swift_continuation_awaitImpl(ContinuationAsyncContext *context) {
                                  /*failure*/ std::memory_order_acquire);
 
   // If that succeeded, we have nothing to do.
-  if (success) return;
+  if (success) {
+    _swift_task_clearCurrent();
+    return;
+  }
 
   // If it failed, it should be because someone concurrently resumed
   // (note that the compare-exchange above is strong).

stdlib/public/Concurrency/TaskGroup.cpp

@@ -810,10 +810,10 @@ PollResult TaskGroupImpl::poll(AsyncTask *waitingTask) {
       mutex.unlock(); // TODO: remove fragment lock, and use status for synchronization
       // no ready tasks, so we must wait.
       result.status = PollStatus::MustWait;
+      _swift_task_clearCurrent();
       return result;
     } // else, try again
   }
-  assert(false && "must successfully compare exchange the waiting task.");
 }
 
 // =============================================================================