Initial Task Executor implementation Task(on:), addTask(on:) etc. (#68793)

Co-authored-by: John McCall <rjmccall@gmail.com>

Author: ktoso

include/swift/ABI/Executor.h

@@ -208,6 +208,12 @@ class TaskExecutorRef {
 
 public:
 
+  // Only public for CompatibilityOverrideConcurrency stubs.
+  // Prefer `TaskExecutorRef::undefined` instead.
+  explicit TaskExecutorRef() : Identity(nullptr), Implementation(0) {
+    assert(isUndefined());
+  }
+
   constexpr static TaskExecutorRef undefined() {
     return TaskExecutorRef(nullptr, 0);
   }
@@ -232,6 +238,8 @@ class TaskExecutorRef {
     return Identity == 0;
   }
 
+  bool isDefined() const { return !isUndefined(); }
+
   TaskExecutorKind getExecutorKind() const {
     return static_cast<TaskExecutorKind>(Implementation & ~WitnessTableMask);
   }

include/swift/ABI/MetadataValues.h

@@ -2437,6 +2437,7 @@ class TaskCreateFlags : public FlagSet<size_t> {
 /// Flags for schedulable jobs.
 class JobFlags : public FlagSet<uint32_t> {
 public:
+  // clang-format off
   enum {
     Kind           = 0,
     Kind_width     = 8,
@@ -2448,12 +2449,14 @@ class JobFlags : public FlagSet<uint32_t> {
 
     // Kind-specific flags.
 
-    Task_IsChildTask           = 24,
-    Task_IsFuture              = 25,
-    Task_IsGroupChildTask      = 26,
+    Task_IsChildTask                      = 24,
+    Task_IsFuture                         = 25,
+    Task_IsGroupChildTask                 = 26,
     // 27 is currently unused
-    Task_IsAsyncLetTask        = 28,
+    Task_IsAsyncLetTask                   = 28,
+    Task_HasInitialTaskExecutorPreference = 29,
   };
+  // clang-format on
 
   explicit JobFlags(uint32_t bits) : FlagSet(bits) {}
   JobFlags(JobKind kind) { setKind(kind); }
@@ -2485,6 +2488,9 @@ class JobFlags : public FlagSet<uint32_t> {
   FLAGSET_DEFINE_FLAG_ACCESSORS(Task_IsAsyncLetTask,
                                 task_isAsyncLetTask,
                                 task_setIsAsyncLetTask)
+  FLAGSET_DEFINE_FLAG_ACCESSORS(Task_HasInitialTaskExecutorPreference,
+                                task_hasInitialTaskExecutorPreference,
+                                task_setHasInitialTaskExecutorPreference)
 };
 
 /// Kinds of task status record.
@@ -2510,6 +2516,10 @@ enum class TaskStatusRecordKind : uint8_t {
   /// escalated.
   EscalationNotification = 4,
 
+  /// A task executor preference, which may impact what executor a task will be
+  /// enqueued on.
+  TaskExecutorPreference = 5,
+
   // Kinds >= 192 are private to the implementation.
   First_Reserved = 192,
   Private_RecordLock = 192
@@ -2518,12 +2528,12 @@ enum class TaskStatusRecordKind : uint8_t {
 /// Kinds of option records that can be passed to creating asynchronous tasks.
 enum class TaskOptionRecordKind : uint8_t {
   /// Request a task to be kicked off, or resumed, on a specific executor.
-  Executor  = 0,
+  InitialTaskExecutor = 0,
   /// Request a child task to be part of a specific task group.
   TaskGroup = 1,
   /// DEPRECATED. AsyncLetWithBuffer is used instead.
   /// Request a child task for an 'async let'.
-  AsyncLet  = 2,
+  AsyncLet = 2,
   /// Request a child task for an 'async let'.
   AsyncLetWithBuffer = 3,
   /// Information about the result type of the task, used in embedded Swift.

include/swift/ABI/Task.h

@@ -38,6 +38,7 @@ struct OpaqueValue;
 struct SwiftError;
 class TaskStatusRecord;
 class TaskDependencyStatusRecord;
+class TaskExecutorPreferenceStatusRecord;
 class TaskOptionRecord;
 class TaskGroup;
 class ContinuationAsyncContext;
@@ -399,7 +400,7 @@ class AsyncTask : public Job {
 
 public:
   /// Flag that the task is to be enqueued on the provided executor and actually
-  /// enqueue it
+  /// enqueue it.
   void flagAsAndEnqueueOnExecutor(SerialExecutorRef newExecutor);
 
   /// Flag that this task is now completed. This normally does not do anything
@@ -410,6 +411,30 @@ class AsyncTask : public Job {
   /// Checking this is, of course, inherently race-prone on its own.
   bool isCancelled() const;
 
+  // ==== Task Executor Preference ---------------------------------------------
+
+  /// Get the preferred task executor reference if there is one set for this
+  /// task.
+  TaskExecutorRef getPreferredTaskExecutor();
+
+  /// WARNING: Only to be used during task creation, in other situations prefer
+  /// to use `swift_task_pushTaskExecutorPreference` and
+  /// `swift_task_popTaskExecutorPreference`.
+  void pushInitialTaskExecutorPreference(TaskExecutorRef preferred);
+
+  /// WARNING: Only to be used during task completion (destroy).
+  ///
+  /// This is because between task creation and its destory, we cannot carry the
+  /// exact record to `pop(record)`, and instead assume that there will be
+  /// exactly one record remaining -- the "initial" record (added during
+  /// creating the task), and it must be that record that is removed by this
+  /// api.
+  ///
+  /// All other situations from user code should be using the
+  /// `swift_task_pushTaskExecutorPreference`, and
+  /// `swift_task_popTaskExecutorPreference(record)` method pair.
+  void dropInitialTaskExecutorPreferenceRecord();
+
   // ==== Task Local Values ----------------------------------------------------
 
   void localValuePush(const HeapObject *key,
@@ -513,6 +538,26 @@ class AsyncTask : public Job {
     return reinterpret_cast<GroupChildFragment *>(offset);
   }
 
+  // ==== Task Executor Preference --------------------------------------------
+
+  /// Returns true if the task has a task executor preference set,
+  /// specifically at creation time of the task. This may be from
+  /// inheriting the preference from a parent task, or by explicitly
+  /// setting it during creation (`Task(_on:...)`).
+  ///
+  /// This means that during task tear down the record should be deallocated
+  /// because it will not be taken care of by a paired "pop" as the normal
+  /// user-land "push / pop" pair of setting a task executor preference would
+  /// have been.
+  bool hasInitialTaskExecutorPreferenceRecord() const {
+    return Flags.task_hasInitialTaskExecutorPreference();
+  }
+
+  /// Returns true if the current task has any task preference record,
+  /// including if it has an initial task preference record or onces
+  /// set during the lifetime of the task.
+  bool hasTaskExecutorPreferenceRecord() const;
+
   // ==== Future ---------------------------------------------------------------
 
   class FutureFragment {

include/swift/ABI/TaskOptions.h

@@ -75,26 +75,25 @@ class TaskGroupTaskOptionRecord : public TaskOptionRecord {
   }
 };
 
-
 /// Task option to specify on what executor the task should be executed.
 ///
-/// Not passing this option implies that a "best guess" or good default
-/// executor should be used instead, most often this may mean the global
-/// concurrent executor, or the enclosing actor's executor.
-class ExecutorTaskOptionRecord : public TaskOptionRecord {
-  const SerialExecutorRef Executor;
+/// Not passing this option implies that an inferred (e.g. surrounding actor
+/// when we inherit execution context) or the default executor should be used.
+///
+/// Lack of this option usually means that the global concurrent executor, or
+/// the executor of the enclosing actor will be used.
+class InitialTaskExecutorPreferenceTaskOptionRecord : public TaskOptionRecord {
+  const TaskExecutorRef Executor;
 
 public:
-  ExecutorTaskOptionRecord(SerialExecutorRef executor)
-    : TaskOptionRecord(TaskOptionRecordKind::Executor),
-      Executor(executor) {}
+  InitialTaskExecutorPreferenceTaskOptionRecord(TaskExecutorRef executor)
+      : TaskOptionRecord(TaskOptionRecordKind::InitialTaskExecutor),
+        Executor(executor) {}
 
-  SerialExecutorRef getExecutor() const {
-    return Executor;
-  }
+  TaskExecutorRef getExecutorRef() const { return Executor; }
 
   static bool classof(const TaskOptionRecord *record) {
-    return record->getKind() == TaskOptionRecordKind::Executor;
+    return record->getKind() == TaskOptionRecordKind::InitialTaskExecutor;
   }
 };
 

include/swift/ABI/TaskStatus.h

@@ -276,6 +276,29 @@ class EscalationNotificationStatusRecord : public TaskStatusRecord {
   }
 };
 
+/// This record signifies that the task has an executor preference.
+/// This preference may be added or removed at runtime, e.g. when multiple
+/// `_withTaskExecutor { ... }` blocks are nested, they add more executor
+/// preferences.
+///
+/// Any number of these preferences may be present at runtime, and the
+/// innermost preference takes priority.
+class TaskExecutorPreferenceStatusRecord : public TaskStatusRecord {
+private:
+  const TaskExecutorRef Preferred;
+
+public:
+  TaskExecutorPreferenceStatusRecord(TaskExecutorRef executor)
+      : TaskStatusRecord(TaskStatusRecordKind::TaskExecutorPreference),
+        Preferred(executor) {}
+
+  TaskExecutorRef getPreferredExecutor() { return Preferred; }
+
+  static bool classof(const TaskStatusRecord *record) {
+    return record->getKind() == TaskStatusRecordKind::TaskExecutorPreference;
+  }
+};
+
 // This record is allocated for a task to record what it is dependent on before
 // the task can make progress again.
 class TaskDependencyStatusRecord : public TaskStatusRecord {

include/swift/AST/ASTContext.h

@@ -916,6 +916,9 @@ class ASTContext final {
   /// Get the runtime availability of support for concurrency.
   AvailabilityContext getConcurrencyAvailability();
 
+  /// Get the runtime availability of task executors.
+  AvailabilityContext getTaskExecutorAvailability();
+
   /// Get the runtime availability of the `DiscardingTaskGroup`,
   /// and supporting runtime functions function
   AvailabilityContext getConcurrencyDiscardingTaskGroupAvailability();

include/swift/AST/ASTSynthesis.h

@@ -41,16 +41,16 @@ enum SingletonTypeSynthesizer {
   _any,
   _bridgeObject,
   _error,
-  _executor,
+  _executor, // the 'BuiltinExecutor' type
   _job,
   _nativeObject,
   _never,
   _rawPointer,
   _rawUnsafeContinuation,
   _void,
   _word,
-  _serialExecutor,
-  _taskExecutor,
+  _serialExecutor, // the '_Concurrency.SerialExecutor' protocol
+  _taskExecutor,   // the '_Concurrency._TaskExecutor' protocol
 };
 inline Type synthesizeType(SynthesisContext &SC,
                            SingletonTypeSynthesizer kind) {
@@ -71,7 +71,7 @@ inline Type synthesizeType(SynthesisContext &SC,
     return SC.Context.getProtocol(KnownProtocolKind::SerialExecutor)
       ->getDeclaredInterfaceType();
   case _taskExecutor:
-    return SC.Context.getProtocol(KnownProtocolKind::TaskExecutor)
+    return SC.Context.getProtocol(KnownProtocolKind::_TaskExecutor)
       ->getDeclaredInterfaceType();
   }
 }

include/swift/AST/Builtins.def

@@ -986,6 +986,34 @@ BUILTIN_MISC_OPERATION_WITH_SILGEN(CreateAsyncTask,
 BUILTIN_MISC_OPERATION_WITH_SILGEN(CreateAsyncTaskInGroup,
     "createAsyncTaskInGroup", "", Special)
 
+/// createAsyncTaskWithExecutor(): (
+///     Int,                            // flags
+///     Builtin.Executor,               // executor
+///     @escaping () async throws -> T  // function
+/// ) -> Builtin.NativeObject
+///
+/// Create a new asynchronous task future, given flags, a parent task,
+/// task group and a function to execute.
+BUILTIN_MISC_OPERATION_WITH_SILGEN(CreateAsyncTaskWithExecutor,
+    "createAsyncTaskWithExecutor", "", Special)
+
+/// createAsyncTaskWithExecutor(): (
+///     Int,                            // flags
+///     Builtin.RawPointer,             // group
+///     Builtin.Executor,               // executor
+///     @escaping () async throws -> T  // function
+/// ) -> Builtin.NativeObject
+///
+/// Create a new asynchronous task future, given flags, a parent task,
+/// task group and a function to execute.
+BUILTIN_MISC_OPERATION_WITH_SILGEN(CreateAsyncTaskInGroupWithExecutor,
+    "createAsyncTaskInGroupWithExecutor", "", Special)
+
+/// Build a Builtin.Executor value from an "ordinary" task executor
+/// reference.
+BUILTIN_MISC_OPERATION_WITH_SILGEN(BuildOrdinaryTaskExecutorRef,
+    "buildOrdinaryTaskExecutorRef", "n", Special)
+
 /// globalStringTablePointer has type String -> Builtin.RawPointer.
 /// It returns an immortal, global string table pointer for strings constructed
 /// from string literals. We consider it effects as readnone meaning that it

include/swift/AST/KnownProtocols.def

@@ -86,7 +86,7 @@ PROTOCOL(FixedWidthInteger)
 PROTOCOL(RangeReplaceableCollection)
 PROTOCOL(Executor)
 PROTOCOL(SerialExecutor)
-PROTOCOL(TaskExecutor)
+PROTOCOL(_TaskExecutor)
 PROTOCOL(GlobalActor)
 
 PROTOCOL_(BridgedNSError)

include/swift/AST/KnownSDKTypes.def

@@ -40,11 +40,11 @@ KNOWN_SDK_TYPE_DECL(ObjectiveC, ObjCBool, StructDecl, 0)
 KNOWN_SDK_TYPE_DECL(Concurrency, CheckedContinuation, NominalTypeDecl, 2)
 KNOWN_SDK_TYPE_DECL(Concurrency, UnsafeContinuation, NominalTypeDecl, 2)
 KNOWN_SDK_TYPE_DECL(Concurrency, MainActor, NominalTypeDecl, 0)
-KNOWN_SDK_TYPE_DECL(Concurrency, Job, StructDecl, 0) // TODO: remove in favor of ExecutorJob
+KNOWN_SDK_TYPE_DECL(Concurrency, Job, StructDecl, 0) // legacy type; prefer ExecutorJob
 KNOWN_SDK_TYPE_DECL(Concurrency, ExecutorJob, StructDecl, 0)
 KNOWN_SDK_TYPE_DECL(Concurrency, UnownedJob, StructDecl, 0)
 KNOWN_SDK_TYPE_DECL(Concurrency, Executor, NominalTypeDecl, 0)
-KNOWN_SDK_TYPE_DECL(Concurrency, TaskExecutor, NominalTypeDecl, 0)
+KNOWN_SDK_TYPE_DECL(Concurrency, _TaskExecutor, NominalTypeDecl, 0)
 KNOWN_SDK_TYPE_DECL(Concurrency, SerialExecutor, NominalTypeDecl, 0)
 KNOWN_SDK_TYPE_DECL(Concurrency, UnownedSerialExecutor, NominalTypeDecl, 0)
 

include/swift/Basic/Features.def

@@ -82,10 +82,13 @@ LANGUAGE_FEATURE(BuiltinHopToActor, 0, "Builtin.HopToActor", true)
 LANGUAGE_FEATURE(BuiltinTaskGroupWithArgument, 0, "TaskGroup builtins", true)
 LANGUAGE_FEATURE(InheritActorContext, 0, "@_inheritActorContext attribute", true)
 LANGUAGE_FEATURE(ImplicitSelfCapture, 0, "@_implicitSelfCapture attribute", true)
+LANGUAGE_FEATURE(BuiltinBuildTaskExecutor, 0, "TaskExecutor-building builtins", true)
 LANGUAGE_FEATURE(BuiltinBuildExecutor, 0, "Executor-building builtins", true)
 LANGUAGE_FEATURE(BuiltinBuildComplexEqualityExecutor, 0, "Executor-building for 'complexEquality executor' builtins", true)
 LANGUAGE_FEATURE(BuiltinBuildMainExecutor, 0, "MainActor executor building builtin", true)
-LANGUAGE_FEATURE(BuiltinCreateAsyncTaskInGroup, 0, "MainActor executor building builtin", true)
+LANGUAGE_FEATURE(BuiltinCreateAsyncTaskInGroup, 0, "Task create in task group builtin with extra flags", true)
+LANGUAGE_FEATURE(BuiltinCreateAsyncTaskInGroupWithExecutor, 0, "Task create in task group builtin with extra flags", true)
+LANGUAGE_FEATURE(BuiltinCreateAsyncTaskWithExecutor, 0, "Task create builtin with extra executor preference", true)
 LANGUAGE_FEATURE(BuiltinCopy, 0, "Builtin.copy()", true)
 LANGUAGE_FEATURE(BuiltinStackAlloc, 0, "Builtin.stackAlloc", true)
 LANGUAGE_FEATURE(BuiltinUnprotectedStackAlloc, 0, "Builtin.unprotectedStackAlloc", true)

include/swift/Demangling/StandardTypesMangling.def

@@ -88,7 +88,7 @@ STANDARD_TYPE_CONCURRENCY(Protocol, F, Executor)
 STANDARD_TYPE_CONCURRENCY(Protocol, f, SerialExecutor)
 STANDARD_TYPE_CONCURRENCY(Structure, G, TaskGroup)
 STANDARD_TYPE_CONCURRENCY(Structure, g, ThrowingTaskGroup)
-STANDARD_TYPE_CONCURRENCY(Protocol, h, TaskExecutor)
+STANDARD_TYPE_CONCURRENCY(Protocol, h, _TaskExecutor)
 STANDARD_TYPE_CONCURRENCY(Protocol, I, AsyncIteratorProtocol)
 STANDARD_TYPE_CONCURRENCY(Protocol, i, AsyncSequence)
 STANDARD_TYPE_CONCURRENCY(Structure, J, UnownedJob)