[Concurrency] Propagate Darwin vouchers across async tasks.

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