Add asyncMainDrainQueue hook

The swift_task_asyncMainDrainQueue function acts as the entrypoint into
driving the main queues, ultimately running the whole program and acting
as the backing driver of the main actor. Making the function hookable
means that custom concurrency runtimes can implement their own async
entrypoints, allowing async top-level code and async-main to "just
work".

Author: etcwilde

include/swift/Runtime/Concurrency.h

@@ -770,6 +770,19 @@ SWIFT_EXPORT_FROM(swift_Concurrency)
 SWIFT_CC(swift) void (*swift_task_enqueueMainExecutor_hook)(
     Job *job, swift_task_enqueueMainExecutor_original original);
 
+/// A hook to override the entrypoint to the main runloop used to drive the
+/// concurrency runtime and drain the main queue. This function must not return.
+/// Note: If the hook is wrapping the original function and the `compatOverride`
+///       is passed in, the `original` function pointer must be passed into the
+///       compatibility override function as the original function.
+typedef SWIFT_CC(swift) void (*swift_task_asyncMainDrainQueue_original)();
+typedef SWIFT_CC(swift) void (*swift_task_asyncMainDrainQueue_override)(
+    swift_task_asyncMainDrainQueue_original original);
+SWIFT_EXPORT_FROM(swift_Concurrency)
+SWIFT_CC(swift) void (*swift_task_asyncMainDrainQueue_hook)(
+    swift_task_asyncMainDrainQueue_original original,
+    swift_task_asyncMainDrainQueue_override compatOverride);
+
 /// Initialize the runtime storage for a default actor.
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
 void swift_defaultActor_initialize(DefaultActor *actor);

stdlib/public/CompatibilityOverride/CompatibilityOverrideConcurrency.def

@@ -77,9 +77,9 @@
 #  define OVERRIDE_TASK_GROUP OVERRIDE
 #  define OVERRIDE_TASK_LOCAL OVERRIDE
 #  define OVERRIDE_TASK_STATUS OVERRIDE
-#ifndef OVERRIDE_TASK_NORETURN
-# define OVERRIDE_TASK_NORETURN(name, attrs, ccAttrs, namespace, typedArgs, \
-                                namedArgs) \
+#ifndef HOOKED_OVERRIDE_TASK_NORETURN
+# define HOOKED_OVERRIDE_TASK_NORETURN(name, attrs, ccAttrs, namespace,       \
+                                      typedArgs, namedArgs)                   \
   OVERRIDE(name, void, attrs, ccAttrs, namespace, typedArgs, namedArgs)
 #endif
 #else
@@ -101,8 +101,8 @@
 #  ifndef OVERRIDE_TASK_STATUS
 #    define OVERRIDE_TASK_STATUS(...)
 #  endif
-#  ifndef OVERRIDE_TASK_NORETURN
-#    define OVERRIDE_TASK_NORETURN(...)
+#  ifndef HOOKED_OVERRIDE_TASK_NORETURN
+#    define HOOKED_OVERRIDE_TASK_NORETURN(...)
 #  endif
 #endif
 
@@ -189,9 +189,10 @@ OVERRIDE_TASK(task_createNullaryContinuationJob, NullaryContinuationJob *,
               (size_t priority,
                AsyncTask *continuation), (priority, continuation))
 
-OVERRIDE_TASK_NORETURN(task_asyncMainDrainQueue,
-              SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift), swift::,
-              , )
+HOOKED_OVERRIDE_TASK_NORETURN(task_asyncMainDrainQueue,
+              SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift),
+              swift::, ,)
+
 
 OVERRIDE_TASK(task_suspend, AsyncTask *,
               SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift),
@@ -388,4 +389,4 @@ OVERRIDE_TASK_STATUS(task_escalate, JobPriority,
 #undef OVERRIDE_TASK_GROUP
 #undef OVERRIDE_TASK_LOCAL
 #undef OVERRIDE_TASK_STATUS
-#undef OVERRIDE_TASK_NORETURN
+#undef HOOKED_OVERRIDE_TASK_NORETURN

stdlib/public/Concurrency/Task.cpp

@@ -1589,20 +1589,31 @@ static void swift_task_asyncMainDrainQueueImpl() {
 #endif
 }
 
+SWIFT_CC(swift)
+void (*swift::swift_task_asyncMainDrainQueue_hook)(
+    swift_task_asyncMainDrainQueue_original original,
+    swift_task_asyncMainDrainQueue_override compatOverride) = nullptr;
+
 #define OVERRIDE_TASK COMPATIBILITY_OVERRIDE
 
 #ifdef SWIFT_STDLIB_SUPPORT_BACK_DEPLOYMENT
 /// The original COMPATIBILITY_OVERRIDE defined in CompatibilityOverride.h
 /// returns the result of the impl function and override function. This results
 /// in a warning emitted for noreturn functions. Overriding the override macro
 /// to not return.
-#define OVERRIDE_TASK_NORETURN(name, attrs, ccAttrs, namespace, typedArgs,     \
-                               namedArgs)                                      \
+#define HOOKED_OVERRIDE_TASK_NORETURN(name, attrs, ccAttrs, namespace,         \
+                                      typedArgs, namedArgs)                    \
   attrs ccAttrs void namespace swift_##name COMPATIBILITY_PAREN(typedArgs) {   \
     static Override_##name Override;                                           \
     static swift_once_t Predicate;                                             \
     swift_once(                                                                \
         &Predicate, [](void *) { Override = getOverride_##name(); }, nullptr); \
+    if (swift_##name##_hook) {                                                 \
+      swift_##name##_hook(COMPATIBILITY_UNPAREN_WITH_COMMA(namedArgs)          \
+                              swift_##name##Impl,                              \
+                          Override);                                           \
+      abort();                                                                 \
+    }                                                                          \
     if (Override != nullptr)                                                   \
       Override(COMPATIBILITY_UNPAREN_WITH_COMMA(namedArgs)                     \
                    swift_##name##Impl);                                        \
@@ -1612,11 +1623,15 @@ static void swift_task_asyncMainDrainQueueImpl() {
 #else // ifndef SWIFT_STDLIB_SUPPORT_BACK_DEPLOYMENT
 // Call directly through to the original implementation when we don't support
 // overrides.
-#define OVERRIDE_TASK_NORETURN(name, attrs, ccAttrs, namespace, typedArgs,     \
-                               namedArgs)                                      \
+#define HOOKED_OVERRIDE_TASK_NORETURN(name, attrs, ccAttrs, namespace,         \
+                                      typedArgs, namedArgs)                    \
   attrs ccAttrs void namespace swift_##name COMPATIBILITY_PAREN(typedArgs) {   \
-      swift_##name##Impl COMPATIBILITY_PAREN(namedArgs);                       \
-    }
+    if (swift_##name##_hook) {                                                 \
+      swift_##name##_hook(swift_##name##Impl, nullptr);                        \
+      abort();                                                                 \
+    }                                                                          \
+    swift_##name##Impl COMPATIBILITY_PAREN(namedArgs);                         \
+  }
 #endif // #else SWIFT_STDLIB_SUPPORT_BACK_DEPLOYMENT
 
 #include COMPATIBILITY_OVERRIDE_INCLUDE_PATH