Merge pull request #36572 from mikeash/async-task-dispatch-integration2

[Concurrency] Make Job objects work as Dispatch objects.

Author: mikeash

include/swift/ABI/Metadata.h

@@ -1338,18 +1338,20 @@ using ClassMetadata = TargetClassMetadata<InProcess>;
 /// dispatch expects to see, with padding to place them at the expected offsets.
 template <typename Runtime>
 struct TargetDispatchClassMetadata : public TargetHeapMetadata<Runtime> {
-  using DummyVTableCall = void (*)(void);
+  using InvokeCall = void (*)(void *, void *, uint32_t);
 
-  TargetDispatchClassMetadata(MetadataKind Kind,
-                              DummyVTableCall DummyVTableEntry)
-      : TargetHeapMetadata<Runtime>(Kind), DummyVTableEntry(DummyVTableEntry) {}
+  TargetDispatchClassMetadata(MetadataKind Kind, unsigned long VTableType,
+                              InvokeCall Invoke)
+      : TargetHeapMetadata<Runtime>(Kind), VTableType(VTableType),
+        VTableInvoke(Invoke) {}
 
   TargetPointer<Runtime, void> Opaque;
 #if SWIFT_OBJC_INTEROP
   TargetPointer<Runtime, void> OpaqueObjC[3];
 #endif
 
-  TargetSignedPointer<Runtime, DummyVTableCall> DummyVTableEntry;
+  unsigned long VTableType;
+  TargetSignedPointer<Runtime, InvokeCall __ptrauth_swift_dispatch_invoke_function> VTableInvoke;
 };
 using DispatchClassMetadata = TargetDispatchClassMetadata<InProcess>;
 

include/swift/ABI/MetadataValues.h

@@ -1227,6 +1227,9 @@ namespace SpecialPointerAuthDiscriminators {
 
   /// Swift async context parameter stored in the extended frame info.
   const uint16_t SwiftAsyncContextExtendedFrameEntry = 0xc31a; // = 49946
+
+  /// Dispatch integration.
+  const uint16_t DispatchInvokeFunction = 0xf493; // = 62611
 }
 
 /// The number of arguments that will be passed directly to a generic

include/swift/ABI/Task.h

@@ -40,14 +40,20 @@ extern FullMetadata<DispatchClassMetadata> jobHeapMetadata;
 
 /// A schedulable job.
 class alignas(2 * alignof(void*)) Job : public HeapObject {
-protected:
+public:
   // Indices into SchedulerPrivate, for use by the runtime.
   enum {
     /// The next waiting task link, an AsyncTask that is waiting on a future.
     NextWaitingTaskIndex = 0,
+
+    /// An opaque field used by Dispatch when enqueueing Jobs directly.
+    DispatchLinkageIndex = 0,
+
+    /// The dispatch queue being used when enqueueing a Job directly with
+    /// Dispatch.
+    DispatchQueueIndex = 1,
   };
 
-public:
   // Reserved for the use of the scheduler.
   void *SchedulerPrivate[2];
 

include/swift/Runtime/Config.h

@@ -247,6 +247,9 @@ extern uintptr_t __COMPATIBILITY_LIBRARIES_CANNOT_CHECK_THE_IS_SWIFT_BIT_DIRECTL
 #define __ptrauth_swift_escalation_notification_function                       \
   __ptrauth(ptrauth_key_function_pointer, 1,                                   \
             SpecialPointerAuthDiscriminators::EscalationNotificationFunction)
+#define __ptrauth_swift_dispatch_invoke_function                               \
+  __ptrauth(ptrauth_key_process_independent_code, 1,                           \
+            SpecialPointerAuthDiscriminators::DispatchInvokeFunction)
 #define swift_ptrauth_sign_opaque_read_resume_function(__fn, __buffer)         \
   ptrauth_auth_and_resign(__fn, ptrauth_key_function_pointer, 0,               \
                           ptrauth_key_process_independent_code,                \
@@ -272,6 +275,7 @@ extern uintptr_t __COMPATIBILITY_LIBRARIES_CANNOT_CHECK_THE_IS_SWIFT_BIT_DIRECTL
 #define __ptrauth_swift_async_context_yield
 #define __ptrauth_swift_cancellation_notification_function
 #define __ptrauth_swift_escalation_notification_function
+#define __ptrauth_swift_dispatch_invoke_function
 #define __ptrauth_swift_runtime_function_entry
 #define __ptrauth_swift_runtime_function_entry_with_key(__key)
 #define __ptrauth_swift_runtime_function_entry_strip(__fn) (__fn)

include/swift/Runtime/EnvironmentVariables.h

@@ -36,6 +36,10 @@ extern OnceToken_t initializeToken;
   }
 #include "../../../stdlib/public/runtime/EnvironmentVariables.def"
 
+// Wrapper around SWIFT_ENABLE_ASYNC_JOB_DISPATCH_INTEGRATION that the
+// Concurrency library can call.
+SWIFT_RUNTIME_STDLIB_SPI bool concurrencyEnableJobDispatchIntegration();
+
 } // end namespace environment
 } // end namespace runtime
 } // end namespace Swift

stdlib/public/Concurrency/GlobalExecutor.cpp

@@ -55,10 +55,15 @@
 
 #include "../CompatibilityOverride/CompatibilityOverride.h"
 #include "swift/Runtime/Concurrency.h"
+#include "swift/Runtime/EnvironmentVariables.h"
 #include "TaskPrivate.h"
 
 #include <dispatch/dispatch.h>
 
+#if !defined(_WIN32)
+#include <dlfcn.h>
+#endif
+
 using namespace swift;
 
 SWIFT_CC(swift)
@@ -180,18 +185,78 @@ void swift::donateThreadToGlobalExecutorUntil(bool (*condition)(void *),
 
 #else
 
+// Ensure that Job's layout is compatible with what Dispatch expects.
+// Note: MinimalDispatchObjectHeader just has the fields we care about, it is
+// not complete and should not be used for anything other than these asserts.
+struct MinimalDispatchObjectHeader {
+  const void *VTable;
+  void *Opaque;
+  void *Linkage;
+};
+static_assert(
+    offsetof(Job, metadata) == offsetof(MinimalDispatchObjectHeader, VTable),
+    "Job Metadata field must match location of Dispatch VTable field.");
+static_assert(offsetof(Job, SchedulerPrivate[Job::DispatchLinkageIndex]) ==
+                  offsetof(MinimalDispatchObjectHeader, Linkage),
+              "Dispatch Linkage field must match Job "
+              "SchedulerPrivate[DispatchLinkageIndex].");
+
 /// The function passed to dispatch_async_f to execute a job.
 static void __swift_run_job(void *_job) {
   Job *job = (Job*) _job;
-  swift_job_run(job, ExecutorRef::generic());
+  auto metadata =
+      reinterpret_cast<const DispatchClassMetadata *>(job->metadata);
+  metadata->VTableInvoke(job, nullptr, 0);
 }
 
-/// A specialized version of __swift_run_job to execute the job on the main
-/// executor.
-/// FIXME: only exists for the quick-and-dirty MainActor implementation.
-static void __swift_run_job_main_executor(void *_job) {
-  Job *job = (Job*) _job;
-  swift_job_run(job, ExecutorRef::mainExecutor());
+/// The type of a function pointer for enqueueing a Job object onto a dispatch
+/// queue.
+typedef void (*dispatchEnqueueFuncType)(dispatch_queue_t queue, void *obj,
+                                        dispatch_qos_class_t qos);
+
+/// Initialize dispatchEnqueueFunc and then call through to the proper
+/// implementation.
+static void initializeDispatchEnqueueFunc(dispatch_queue_t queue, void *obj,
+                                          dispatch_qos_class_t qos);
+
+/// A function pointer to the function used to enqueue a Job onto a dispatch
+/// queue. Initially set to initializeDispatchEnqueueFunc, so that the first
+/// call will initialize it. initializeDispatchEnqueueFunc sets it to point
+/// either to dispatch_async_swift_job when it's available, otherwise to
+/// dispatchEnqueueDispatchAsync.
+static std::atomic<dispatchEnqueueFuncType> dispatchEnqueueFunc{
+    initializeDispatchEnqueueFunc};
+
+/// A small adapter that dispatches a Job onto a queue using dispatch_async_f.
+static void dispatchEnqueueDispatchAsync(dispatch_queue_t queue, void *obj,
+                                         dispatch_qos_class_t qos) {
+  dispatch_async_f(queue, obj, __swift_run_job);
+}
+
+static void initializeDispatchEnqueueFunc(dispatch_queue_t queue, void *obj,
+                                          dispatch_qos_class_t qos) {
+  dispatchEnqueueFuncType func = nullptr;
+
+  // Always fall back to plain dispatch_async_f on Windows for now.
+#if !defined(_WIN32)
+  if (runtime::environment::concurrencyEnableJobDispatchIntegration())
+    func = reinterpret_cast<dispatchEnqueueFuncType>(
+        dlsym(RTLD_NEXT, "dispatch_async_swift_job"));
+#endif
+
+  if (!func)
+    func = dispatchEnqueueDispatchAsync;
+
+  dispatchEnqueueFunc.store(func, std::memory_order_relaxed);
+
+  func(queue, obj, qos);
+}
+
+/// Enqueue a Job onto a dispatch queue using dispatchEnqueueFunc.
+static void dispatchEnqueue(dispatch_queue_t queue, Job *job,
+                            dispatch_qos_class_t qos, void *executorQueue) {
+  job->SchedulerPrivate[Job::DispatchQueueIndex] = executorQueue;
+  dispatchEnqueueFunc.load(std::memory_order_relaxed)(queue, job, qos);
 }
 
 static constexpr size_t globalQueueCacheCount =
@@ -259,14 +324,12 @@ static void swift_task_enqueueGlobalImpl(Job *job) {
   // the priorities of work added to this queue using Dispatch's public
   // API, but as discussed above, that is less important than avoiding
   // performance problems.
-  dispatch_function_t dispatchFunction = &__swift_run_job;
-  void *dispatchContext = job;
-
   JobPriority priority = job->getPriority();
 
   auto queue = getGlobalQueue(priority);
 
-  dispatch_async_f(queue, dispatchContext, dispatchFunction);
+  dispatchEnqueue(queue, job, (dispatch_qos_class_t)priority,
+                  DISPATCH_QUEUE_GLOBAL_EXECUTOR);
 #endif
 }
 
@@ -293,6 +356,9 @@ static void swift_task_enqueueGlobalWithDelayImpl(unsigned long long delay,
 
   auto queue = getGlobalQueue(priority);
 
+  job->SchedulerPrivate[Job::DispatchQueueIndex] =
+      DISPATCH_QUEUE_GLOBAL_EXECUTOR;
+
   dispatch_time_t when = dispatch_time(DISPATCH_TIME_NOW, delay);
   dispatch_after_f(when, queue, dispatchContext, dispatchFunction);
 #endif
@@ -317,13 +383,13 @@ static void swift_task_enqueueMainExecutorImpl(Job *job) {
   insertIntoJobQueue(job);
 #else
 
-  dispatch_function_t dispatchFunction = &__swift_run_job_main_executor;
-  void *dispatchContext = job;
+  JobPriority priority = job->getPriority();
 
   // This is an inline function that compiles down to a pointer to a global.
   auto mainQueue = dispatch_get_main_queue();
 
-  dispatch_async_f(mainQueue, dispatchContext, dispatchFunction);
+  dispatchEnqueue(mainQueue, job, (dispatch_qos_class_t)priority,
+                  DISPATCH_QUEUE_MAIN_EXECUTOR);
 
 #endif
 }

stdlib/public/Concurrency/Task.cpp

@@ -102,7 +102,7 @@ void AsyncTask::completeFuture(AsyncContext *context) {
       reinterpret_cast<char *>(context) - sizeof(FutureAsyncContextPrefix));
   bool hadErrorResult = false;
   auto errorObject = asyncContextPrefix->errorResult;
-  printf("asyncTask::completeFuture errorObject: %p\n", errorObject);
+//   printf("asyncTask::completeFuture errorObject: %p\n", errorObject);
   fragment->getError() = errorObject;
   if (errorObject) {
     hadErrorResult = true;
@@ -180,10 +180,26 @@ static void destroyTask(SWIFT_CONTEXT HeapObject *obj) {
   free(task);
 }
 
-static void dummyVTableFunction(void) {
-  abort();
+static ExecutorRef executorForEnqueuedJob(Job *job) {
+  void *jobQueue = job->SchedulerPrivate[Job::DispatchQueueIndex];
+  if (jobQueue == DISPATCH_QUEUE_GLOBAL_EXECUTOR)
+    return ExecutorRef::generic();
+  else if (jobQueue == DISPATCH_QUEUE_MAIN_EXECUTOR)
+    return ExecutorRef::mainExecutor();
+  else
+    swift_unreachable("jobQueue was not a known value.");
 }
 
+static void jobInvoke(void *obj, void *unused, uint32_t flags) {
+  (void)unused;
+  Job *job = reinterpret_cast<Job *>(obj);
+
+  swift_job_run(job, executorForEnqueuedJob(job));
+}
+
+// Magic constant to identify Swift Job vtables to Dispatch.
+static const unsigned long dispatchSwiftObjectType = 1;
+
 FullMetadata<DispatchClassMetadata> swift::jobHeapMetadata = {
   {
     {
@@ -195,7 +211,8 @@ FullMetadata<DispatchClassMetadata> swift::jobHeapMetadata = {
   },
   {
     MetadataKind::Job,
-    dummyVTableFunction
+    dispatchSwiftObjectType,
+    jobInvoke
   }
 };
 
@@ -211,7 +228,8 @@ static FullMetadata<DispatchClassMetadata> taskHeapMetadata = {
   },
   {
     MetadataKind::Task,
-    dummyVTableFunction
+    dispatchSwiftObjectType,
+    jobInvoke
   }
 };
 

stdlib/public/Concurrency/TaskPrivate.h

@@ -71,6 +71,11 @@ void donateThreadToGlobalExecutorUntil(bool (*condition)(void*),
 void _swift_tsan_acquire(void *addr);
 void _swift_tsan_release(void *addr);
 
+/// Special values used with DispatchQueueIndex to indicate the global and main
+/// executors.
+#define DISPATCH_QUEUE_GLOBAL_EXECUTOR (void *)1
+#define DISPATCH_QUEUE_MAIN_EXECUTOR (void *)2
+
 // ==== ------------------------------------------------------------------------
 
 namespace {

stdlib/public/runtime/EnvironmentVariables.cpp

@@ -194,3 +194,7 @@ bool swift_COWChecksEnabled() {
   return runtime::environment::SWIFT_DEBUG_ENABLE_COW_CHECKS();
 }
 
+SWIFT_RUNTIME_STDLIB_SPI bool concurrencyEnableJobDispatchIntegration() {
+  return runtime::environment::
+      SWIFT_ENABLE_ASYNC_JOB_DISPATCH_INTEGRATION();
+}

stdlib/public/runtime/EnvironmentVariables.def

@@ -47,6 +47,9 @@ VARIABLE(SWIFT_DEBUG_ENABLE_MALLOC_SCRIBBLE, bool, false,
 VARIABLE(SWIFT_DEBUG_ENABLE_COW_CHECKS, bool, false,
          "Enable internal checks for copy-on-write operations.")
 
+VARIABLE(SWIFT_ENABLE_ASYNC_JOB_DISPATCH_INTEGRATION, bool, true,
+         "Enable use of dispatch_async_swift_job when available.")
+
 #if defined(__APPLE__) && defined(__MACH__)
 
 VARIABLE(SWIFT_DEBUG_VALIDATE_SHARED_CACHE_PROTOCOL_CONFORMANCES, bool, false,