Properly save and restore the current task in the runtime so that tasks

can be reentrantly executed.

I don't think doing this is *actually a good idea*, but corrupting the
runtime is an even worse idea, and the overhead here is very low.

Author: rjmccall

include/swift/Threading/ThreadLocalStorage.h

@@ -115,6 +115,15 @@ class ThreadLocal {
   T get() { return value; }
 
   void set(T newValue) { value = newValue; }
+
+  T swap(T newValue) {
+    // There's an implicit optimization here because we implicitly
+    // materialize the address of the thread-local once instead of
+    // separately for two calls to get and set.
+    auto curValue = get();
+    set(newValue);
+    return curValue;
+  }
 };
 #else
 // A wrapper around a TLS key that is lazily initialized using swift::once.
@@ -168,6 +177,12 @@ class ThreadLocal {
     memcpy(&storedValue, &newValue, sizeof(T));
     tls_set(key.getKey(), storedValue);
   }
+
+  T swap(T newValue) {
+    auto curValue = get();
+    set(newValue);
+    return curValue;
+  }
 };
 #endif
 

stdlib/public/Concurrency/Actor.cpp

@@ -196,6 +196,9 @@ class ActiveTask {
 public:
   static void set(AsyncTask *task) { Value.set(task); }
   static AsyncTask *get() { return Value.get(); }
+  static AsyncTask *swap(AsyncTask *newTask) {
+    return Value.swap(newTask);
+  }
 };
 
 /// Define the thread-locals.
@@ -217,7 +220,7 @@ void swift::runJobInEstablishedExecutorContext(Job *job) {
 
   if (auto task = dyn_cast<AsyncTask>(job)) {
     // Update the active task in the current thread.
-    ActiveTask::set(task);
+    auto oldTask = ActiveTask::swap(task);
 
     // Update the task status to say that it's running on the
     // current thread.  If the task suspends somewhere, it should
@@ -231,6 +234,7 @@ void swift::runJobInEstablishedExecutorContext(Job *job) {
 
     assert(ActiveTask::get() == nullptr &&
            "active task wasn't cleared before suspending?");
+    if (oldTask) ActiveTask::set(oldTask);
   } else {
     // There's no extra bookkeeping to do for simple jobs besides swapping in
     // the voucher.
@@ -259,15 +263,11 @@ AsyncTask *swift::swift_task_getCurrent() {
 }
 
 AsyncTask *swift::_swift_task_clearCurrent() {
-  auto task = ActiveTask::get();
-  ActiveTask::set(nullptr);
-  return task;
+  return ActiveTask::swap(nullptr);
 }
 
 AsyncTask *swift::_swift_task_setCurrent(AsyncTask *new_task) {
-  auto task = ActiveTask::get();
-  ActiveTask::set(new_task);
-  return task;
+  return ActiveTask::swap(new_task);
 }
 
 SWIFT_CC(swift)