[Concurrency] Propagate Darwin vouchers across async tasks.

Darwin OSes support vouchers, which are key/value sets that can be adopted on a thread to influence its execution, or sent to another process. APIs like Dispatch propagate vouchers to worker threads when running async code. This change makes Swift Concurrency do the same.

The change consists of a few different parts:

1. A set of shims (in VoucherShims.h) which provides declarations for the necessary calls when they're not available from the SDK, and stub implementations for non-Darwin platforms.
2. One of Job's reserved fields is now used to store the voucher associated with a job.
3. Jobs grab the current thread's voucher when they're created.
4. A VoucherManager class manages adoption of vouchers when running a Job, and replacing vouchers in suspended tasks.
5. A VoucherManager instance is maintained in ExecutionTrackingInfo, and is updated as necessary throughout a Job/Task's lifecycle.

rdar://76080222

Author: mikeash

include/swift/ABI/Task.h

@@ -23,6 +23,7 @@
 #include "swift/ABI/Metadata.h"
 #include "swift/ABI/MetadataValues.h"
 #include "swift/Runtime/Config.h"
+#include "swift/Runtime/VoucherShims.h"
 #include "swift/Basic/STLExtras.h"
 #include "bitset"
 #include "queue" // TODO: remove and replace with our own mpsc
@@ -72,7 +73,13 @@ class alignas(2 * alignof(void*)) Job :
   // Derived classes can use this to store a Job Id.
   uint32_t Id = 0;
 
-  void *Reserved[2] = {};
+  /// The voucher associated with the job. Note: this is currently unused on
+  /// non-Darwin platforms, with stub implementations of the functions for
+  /// consistency.
+  voucher_t Voucher = nullptr;
+
+  /// Reserved for future use.
+  void *Reserved = nullptr;
 
   // We use this union to avoid having to do a second indirect branch
   // when resuming an asynchronous task, which we expect will be the
@@ -88,23 +95,32 @@ class alignas(2 * alignof(void*)) Job :
   Job(JobFlags flags, JobInvokeFunction *invoke,
       const HeapMetadata *metadata = &jobHeapMetadata)
       : HeapObject(metadata), Flags(flags), RunJob(invoke) {
+    Voucher = voucher_copy();
     assert(!isAsyncTask() && "wrong constructor for a task");
   }
 
   Job(JobFlags flags, TaskContinuationFunction *invoke,
-      const HeapMetadata *metadata = &jobHeapMetadata)
+      const HeapMetadata *metadata = &jobHeapMetadata,
+      bool captureCurrentVoucher = true)
       : HeapObject(metadata), Flags(flags), ResumeTask(invoke) {
+    if (captureCurrentVoucher)
+      Voucher = voucher_copy();
     assert(isAsyncTask() && "wrong constructor for a non-task job");
   }
 
   /// Create a job with "immortal" reference counts.
   /// Used for async let tasks.
   Job(JobFlags flags, TaskContinuationFunction *invoke,
-      const HeapMetadata *metadata, InlineRefCounts::Immortal_t immortal)
+      const HeapMetadata *metadata, InlineRefCounts::Immortal_t immortal,
+      bool captureCurrentVoucher = true)
       : HeapObject(metadata, immortal), Flags(flags), ResumeTask(invoke) {
+    if (captureCurrentVoucher)
+      Voucher = voucher_copy();
     assert(isAsyncTask() && "wrong constructor for a non-task job");
   }
 
+  ~Job() { swift_voucher_release(Voucher); }
+
   bool isAsyncTask() const {
     return Flags.isAsyncTask();
   }
@@ -229,8 +245,9 @@ class AsyncTask : public Job {
   /// Private.initialize separately.
   AsyncTask(const HeapMetadata *metadata, JobFlags flags,
             TaskContinuationFunction *run,
-            AsyncContext *initialContext)
-    : Job(flags, run, metadata),
+            AsyncContext *initialContext,
+            bool captureCurrentVoucher)
+    : Job(flags, run, metadata, captureCurrentVoucher),
       ResumeContext(initialContext) {
     assert(flags.isAsyncTask());
     Id = getNextTaskId();
@@ -243,8 +260,9 @@ class AsyncTask : public Job {
   AsyncTask(const HeapMetadata *metadata, InlineRefCounts::Immortal_t immortal,
             JobFlags flags,
             TaskContinuationFunction *run,
-            AsyncContext *initialContext)
-    : Job(flags, run, metadata, immortal),
+            AsyncContext *initialContext,
+            bool captureCurrentVoucher)
+    : Job(flags, run, metadata, immortal, captureCurrentVoucher),
       ResumeContext(initialContext) {
     assert(flags.isAsyncTask());
     Id = getNextTaskId();

include/swift/Runtime/VoucherShims.h

@@ -0,0 +1,87 @@
+//===--- VoucherShims.h - Shims for OS vouchers --------------------*- C++ -*-//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2021 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+//
+// Shims for interfacing with OS voucher calls.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SWIFT_CONCURRENCY_VOUCHERSHIMS_H
+#define SWIFT_CONCURRENCY_VOUCHERSHIMS_H
+
+#include "Config.h"
+
+// swift-corelibs-libdispatch has os/voucher_private.h but it doesn't work for
+// us yet, so only look for it on Apple platforms.
+#if __APPLE__ && __has_include(<os/voucher_private.h>)
+#define SWIFT_HAS_VOUCHER_HEADER 1
+#include <os/voucher_private.h>
+#endif
+
+// A "dead" voucher pointer, indicating that a voucher has been removed from
+// a Job, distinct from a NULL voucher that could just mean no voucher was
+// present. This allows us to catch problems like adopting a voucher from the
+// same Job twice without restoring it.
+#define SWIFT_DEAD_VOUCHER ((voucher_t)-1)
+
+// The OS has voucher support if it has the header or if it has ObjC interop.
+#if SWIFT_HAS_VOUCHER_HEADER || SWIFT_OBJC_INTEROP
+#define SWIFT_HAS_VOUCHERS 1
+#endif
+
+#if SWIFT_HAS_VOUCHERS
+
+// If the header isn't available, declare the necessary calls here.
+#if !SWIFT_HAS_VOUCHER_HEADER
+
+#include <os/object.h>
+
+#pragma clang diagnostic push
+#pragma clang diagnostic ignored "-Wgnu-zero-variadic-macro-arguments"
+OS_OBJECT_DECL_CLASS(voucher);
+#pragma clang diagnostic pop
+
+extern "C" voucher_t _Nullable voucher_copy(void);
+
+// Consumes argument, returns retained value.
+extern "C" voucher_t _Nullable voucher_adopt(voucher_t _Nullable voucher);
+
+static inline bool voucher_needs_adopt(voucher_t _Nullable voucher) {
+  return true;
+}
+
+#endif // __has_include(<os/voucher_private.h>)
+
+static inline void swift_voucher_release(voucher_t _Nullable voucher) {
+  // This NULL check isn't necessary, but NULL vouchers will be common, so
+  // optimize for that.
+  if (!voucher)
+    return;
+  if (voucher == SWIFT_DEAD_VOUCHER)
+    return;
+  os_release(voucher);
+}
+
+#else  // __APPLE__
+
+// Declare some do-nothing stubs for OSes without voucher support.
+typedef void *voucher_t;
+static inline voucher_t _Nullable voucher_copy(void) { return nullptr; }
+static inline voucher_t _Nullable voucher_adopt(voucher_t _Nullable voucher) {
+  return nullptr;
+}
+static inline bool voucher_needs_adopt(voucher_t _Nullable voucher) {
+  return true;
+}
+static inline void swift_voucher_release(voucher_t _Nullable voucher) {}
+#endif // __APPLE__
+
+#endif

stdlib/public/Concurrency/Actor.cpp

@@ -37,6 +37,7 @@
 #include "llvm/Config/config.h"
 #include "llvm/ADT/PointerIntPair.h"
 #include "TaskPrivate.h"
+#include "VoucherSupport.h"
 
 #if !SWIFT_CONCURRENCY_COOPERATIVE_GLOBAL_EXECUTOR
 #include <dispatch/dispatch.h>
@@ -132,6 +133,8 @@ class ExecutorTrackingInfo {
   /// unless the passed-in executor is generic.
   bool AllowsSwitching = true;
 
+  VoucherManager voucherManager;
+
   /// The tracking info that was active when this one was entered.
   ExecutorTrackingInfo *SavedInfo;
 
@@ -150,23 +153,9 @@ class ExecutorTrackingInfo {
     ActiveInfoInThread.set(this);
   }
 
-  /// Initialize a tracking state on the current thread if there
-  /// isn't one already, or else update the current tracking state.
-  ///
-  /// Returns a pointer to the active tracking info.  If this is the
-  /// same as the object on which this was called, leave() must be
-  /// called before the object goes out of scope.
-  ExecutorTrackingInfo *enterOrUpdate(ExecutorRef currentExecutor) {
-    if (auto activeInfo = ActiveInfoInThread.get()) {
-      activeInfo->ActiveExecutor = currentExecutor;
-      return activeInfo;
-    }
+  void swapToJob(Job *job) { voucherManager.swapToJob(job); }
 
-    ActiveExecutor = currentExecutor;
-    SavedInfo = nullptr;
-    ActiveInfoInThread.set(this);
-    return this;
-  }
+  void restoreVoucher(AsyncTask *task) { voucherManager.restoreVoucher(task); }
 
   ExecutorRef getActiveExecutor() const {
     return ActiveExecutor;
@@ -192,6 +181,7 @@ class ExecutorTrackingInfo {
   }
 
   void leave() {
+    voucherManager.leave();
     ActiveInfoInThread.set(SavedInfo);
   }
 };
@@ -242,7 +232,9 @@ void swift::runJobInEstablishedExecutorContext(Job *job) {
     assert(ActiveTask::get() == nullptr &&
            "active task wasn't cleared before susspending?");
   } else {
-    // There's no extra bookkeeping to do for simple jobs.
+    // There's no extra bookkeeping to do for simple jobs besides swapping in
+    // the voucher.
+    ExecutorTrackingInfo::current()->swapToJob(job);
     job->runSimpleInFullyEstablishedContext();
   }
 
@@ -253,6 +245,14 @@ void swift::runJobInEstablishedExecutorContext(Job *job) {
   _swift_tsan_release(job);
 }
 
+void swift::adoptTaskVoucher(AsyncTask *task) {
+  ExecutorTrackingInfo::current()->swapToJob(task);
+}
+
+void swift::restoreTaskVoucher(AsyncTask *task) {
+  ExecutorTrackingInfo::current()->restoreVoucher(task);
+}
+
 SWIFT_CC(swift)
 AsyncTask *swift::swift_task_getCurrent() {
   return ActiveTask::get();
@@ -676,6 +676,12 @@ class DefaultActorImpl : public HeapObject {
 
   friend class ProcessInlineJob;
   union {
+    // When the ProcessInlineJob storage is initialized, its metadata pointer
+    // will point to Job's metadata. When it isn't, the metadata pointer is
+    // NULL. Use HeapObject to initialize the metadata pointer to NULL and allow
+    // it to be checked without fully initializing the ProcessInlineJob.
+    HeapObject JobStorageHeapObject{nullptr};
+
     ProcessInlineJob JobStorage;
   };
 
@@ -685,6 +691,7 @@ class DefaultActorImpl : public HeapObject {
     auto flags = Flags();
     flags.setIsDistributedRemote(isDistributedRemote);
     new (&CurrentState) std::atomic<State>(State{JobRef(), flags});
+    JobStorageHeapObject.metadata = nullptr;
   }
 
   /// Properly destruct an actor, except for the heap header.
@@ -812,6 +819,8 @@ void DefaultActorImpl::deallocate() {
 }
 
 void DefaultActorImpl::deallocateUnconditional() {
+  if (JobStorageHeapObject.metadata != nullptr)
+    JobStorage.~ProcessInlineJob();
   auto metadata = cast<ClassMetadata>(this->metadata);
   swift_deallocObject(this, metadata->getInstanceSize(),
                       metadata->getInstanceAlignMask());
@@ -848,6 +857,8 @@ void DefaultActorImpl::scheduleNonOverrideProcessJob(JobPriority priority,
   if (hasActiveInlineJob) {
     job = new ProcessOutOfLineJob(this, priority);
   } else {
+    if (JobStorageHeapObject.metadata != nullptr)
+      JobStorage.~ProcessInlineJob();
     job = new (&JobStorage) ProcessInlineJob(priority);
   }
   swift_task_enqueueGlobal(job);
@@ -1455,6 +1466,7 @@ static void swift_job_runImpl(Job *job, ExecutorRef executor) {
 
   trackingInfo.enterAndShadow(executor);
 
+  SWIFT_TASK_DEBUG_LOG("%s(%p)", __func__, job);
   runJobInEstablishedExecutorContext(job);
 
   trackingInfo.leave();

stdlib/public/Concurrency/Task.cpp

@@ -586,15 +586,18 @@ static AsyncTaskAndContext swift_task_create_commonImpl(
   // Initialize the task so that resuming it will run the given
   // function on the initial context.
   AsyncTask *task = nullptr;
+  bool captureCurrentVoucher = taskCreateFlags.copyTaskLocals() || jobFlags.task_isChildTask();
   if (asyncLet) {
     // Initialize the refcount bits to "immortal", so that
     // ARC operations don't have any effect on the task.
     task = new(allocation) AsyncTask(&taskHeapMetadata,
                              InlineRefCounts::Immortal, jobFlags,
-                             function, initialContext);
+                             function, initialContext,
+                             captureCurrentVoucher);
   } else {
     task = new(allocation) AsyncTask(&taskHeapMetadata, jobFlags,
-                                    function, initialContext);
+                                    function, initialContext,
+                                    captureCurrentVoucher);
   }
 
   // Initialize the child fragment if applicable.

stdlib/public/Concurrency/TaskPrivate.h

@@ -82,6 +82,14 @@ void _swift_task_dealloc_specific(AsyncTask *task, void *ptr);
 /// related to the active task.
 void runJobInEstablishedExecutorContext(Job *job);
 
+/// Adopt the voucher stored in `task`. This removes the voucher from the task
+/// and adopts it on the current thread.
+void adoptTaskVoucher(AsyncTask *task);
+
+/// Restore the voucher for `task`. This un-adopts the current thread's voucher
+/// and stores it back into the task again.
+void restoreTaskVoucher(AsyncTask *task);
+
 /// Initialize the async let storage for the given async-let child task.
 void asyncLet_addImpl(AsyncTask *task, AsyncLet *asyncLet,
                       bool didAllocateInParentTask);
@@ -354,11 +362,13 @@ inline bool AsyncTask::isCancelled() const {
 }
 
 inline void AsyncTask::flagAsRunning() {
+  SWIFT_TASK_DEBUG_LOG("%p->flagAsRunning()", this);
   auto oldStatus = _private().Status.load(std::memory_order_relaxed);
   while (true) {
     assert(!oldStatus.isRunning());
     if (oldStatus.isLocked()) {
       flagAsRunning_slow();
+      adoptTaskVoucher(this);
       swift_task_enterThreadLocalContext(
           (char *)&_private().ExclusivityAccessSet[0]);
       return;
@@ -373,6 +383,7 @@ inline void AsyncTask::flagAsRunning() {
     if (_private().Status.compare_exchange_weak(oldStatus, newStatus,
                                                 std::memory_order_relaxed,
                                                 std::memory_order_relaxed)) {
+      adoptTaskVoucher(this);
       swift_task_enterThreadLocalContext(
           (char *)&_private().ExclusivityAccessSet[0]);
       return;
@@ -381,13 +392,15 @@ inline void AsyncTask::flagAsRunning() {
 }
 
 inline void AsyncTask::flagAsSuspended() {
+  SWIFT_TASK_DEBUG_LOG("%p->flagAsSuspended()", this);
   auto oldStatus = _private().Status.load(std::memory_order_relaxed);
   while (true) {
     assert(oldStatus.isRunning());
     if (oldStatus.isLocked()) {
       flagAsSuspended_slow();
       swift_task_exitThreadLocalContext(
           (char *)&_private().ExclusivityAccessSet[0]);
+      restoreTaskVoucher(this);
       return;
     }
 
@@ -402,6 +415,7 @@ inline void AsyncTask::flagAsSuspended() {
                                                 std::memory_order_relaxed)) {
       swift_task_exitThreadLocalContext(
           (char *)&_private().ExclusivityAccessSet[0]);
+      restoreTaskVoucher(this);
       return;
     }
   }