[6.0][Concurrency] TaskExecutor ownership fixes (#74174)

Author: ktoso

include/swift/ABI/Executor.h

@@ -260,12 +260,6 @@ class TaskExecutorRef {
     return reinterpret_cast<const TaskExecutorWitnessTable*>(table);
   }
 
-//  /// Do we have to do any work to start running as the requested
-//  /// executor?
-//  bool mustSwitchToRun(TaskExecutorRef newExecutor) const {
-//    return Identity != newExecutor.Identity;
-//  }
-
   /// Get the raw value of the Implementation field, for tracing.
   uintptr_t getRawImplementation() const {
     return Implementation & WitnessTableMask;

include/swift/ABI/MetadataValues.h

@@ -2697,7 +2697,8 @@ enum class TaskOptionRecordKind : uint8_t {
   /// Information about the result type of the task, used in embedded Swift.
   ResultTypeInfo = 4,
   /// Set the initial task executor preference of the task.
-  InitialTaskExecutor = 5,
+  InitialTaskExecutorUnowned = 5,
+  InitialTaskExecutorOwned = 6,
   /// Request a child task for swift_task_run_inline.
   RunInline = UINT8_MAX,
 };

include/swift/ABI/Task.h

@@ -419,12 +419,16 @@ class AsyncTask : public Job {
 
   /// Get the preferred task executor reference if there is one set for this
   /// task.
-  TaskExecutorRef getPreferredTaskExecutor();
+  TaskExecutorRef getPreferredTaskExecutor(bool assumeHasRecord = false);
 
   /// 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);
+  ///
+  /// The `owned` parameter indicates if the executor is owned by the task,
+  /// and must be released when the task completes.
+  void pushInitialTaskExecutorPreference(
+      TaskExecutorRef preferred, bool owned);
 
   /// WARNING: Only to be used during task completion (destroy).
   ///

include/swift/ABI/TaskOptions.h

@@ -77,23 +77,58 @@ class TaskGroupTaskOptionRecord : public TaskOptionRecord {
 
 /// Task option to specify on what executor the task should be executed.
 ///
-/// Not passing this option implies that an inferred (e.g. surrounding actor
-/// when we inherit execution context) or the default executor should be used.
+/// Not passing this option (or it's alternative "owned" version) 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 {
+class InitialTaskExecutorRefPreferenceTaskOptionRecord : public TaskOptionRecord {
   const TaskExecutorRef Executor;
 
 public:
-  InitialTaskExecutorPreferenceTaskOptionRecord(TaskExecutorRef executor)
-      : TaskOptionRecord(TaskOptionRecordKind::InitialTaskExecutor),
+  InitialTaskExecutorRefPreferenceTaskOptionRecord(TaskExecutorRef executor)
+      : TaskOptionRecord(TaskOptionRecordKind::InitialTaskExecutorUnowned),
         Executor(executor) {}
 
   TaskExecutorRef getExecutorRef() const { return Executor; }
 
   static bool classof(const TaskOptionRecord *record) {
-    return record->getKind() == TaskOptionRecordKind::InitialTaskExecutor;
+    return record->getKind() == TaskOptionRecordKind::InitialTaskExecutorUnowned;
+  }
+};
+
+/// This is quite similar to `InitialTaskExecutorRefPreferenceTaskOptionRecord`
+/// however it takes a "raw" TaskExecutor existential in the form of an Identity
+/// and WitnessTable - rather than the specific UnownedTaskExecutor which already
+/// may have specific "flags" set on it.
+///
+/// In order to use the executor in the runtime, we need to call into the type's
+/// `asUnownedTaskExecutor` which is done by
+/// `getExecutorRefFromUnownedTaskExecutor`.
+class InitialTaskExecutorOwnedPreferenceTaskOptionRecord
+    : public TaskOptionRecord {
+
+  // These look similar to TaskExecutorRef but are NOT the same!
+  // A TaskExecutorRef is obtained through calling user defined
+  // `asUnownedTaskExecutor` which is what we need to do on these to get a real executor ref.
+  HeapObject *Identity;
+  const TaskExecutorWitnessTable *WitnessTable;
+
+public:
+  InitialTaskExecutorOwnedPreferenceTaskOptionRecord(
+      HeapObject *executor, uintptr_t witnessTable)
+      : TaskOptionRecord(TaskOptionRecordKind::InitialTaskExecutorOwned),
+        Identity(executor) {
+    WitnessTable = reinterpret_cast<const TaskExecutorWitnessTable*>(witnessTable);
+  }
+
+  /// Invokes Swift implemented `asUnownedTaskExecutor` in order to obtain an
+  /// `TaskExecutorRef` which is properly populated with any flags it might need.
+  TaskExecutorRef getExecutorRefFromUnownedTaskExecutor() const;
+
+  static bool classof(const TaskOptionRecord *record) {
+    return record->getKind() == TaskOptionRecordKind::InitialTaskExecutorOwned;
   }
 };
 

include/swift/ABI/TaskStatus.h

@@ -20,6 +20,7 @@
 #ifndef SWIFT_ABI_TASKSTATUS_H
 #define SWIFT_ABI_TASKSTATUS_H
 
+#include "swift/Basic/OptionSet.h"
 #include "swift/ABI/MetadataValues.h"
 #include "swift/ABI/Task.h"
 #include "swift/ABI/Executor.h"
@@ -285,15 +286,35 @@ class EscalationNotificationStatusRecord : public TaskStatusRecord {
 /// innermost preference takes priority.
 class TaskExecutorPreferenceStatusRecord : public TaskStatusRecord {
 private:
+  enum class Flags : uint8_t {
+    /// The executor was retained during this task's creation,
+    /// and therefore must be released when this task completes.
+    ///
+    /// The only tasks which need to manually retain/release the task executor
+    /// are those which cannot structurally guarantee its lifetime. E.g. an async
+    /// let does not need to do so, because it structurally always will end
+    /// before/// we leave the scope in which it was defined -- and such scope
+    /// must have been keeping alive the executor.
+    HasRetainedExecutor = 1 << 0
+  };
+  OptionSet<Flags> flags;
   const TaskExecutorRef Preferred;
 
 public:
-  TaskExecutorPreferenceStatusRecord(TaskExecutorRef executor)
+  TaskExecutorPreferenceStatusRecord(TaskExecutorRef executor, bool retainedExecutor)
       : TaskStatusRecord(TaskStatusRecordKind::TaskExecutorPreference),
-        Preferred(executor) {}
+        Preferred(executor) {
+    if (retainedExecutor) {
+      flags = Flags::HasRetainedExecutor;
+    }
+  }
 
   TaskExecutorRef getPreferredExecutor() { return Preferred; }
 
+  bool hasRetainedExecutor() const {
+    return flags.contains(Flags::HasRetainedExecutor);
+  }
+
   static bool classof(const TaskStatusRecord *record) {
     return record->getKind() == TaskStatusRecordKind::TaskExecutorPreference;
   }

include/swift/AST/ASTSynthesis.h

@@ -78,8 +78,11 @@ inline Type synthesizeType(SynthesisContext &SC,
     return SC.Context.getProtocol(KnownProtocolKind::SerialExecutor)
       ->getDeclaredInterfaceType();
   case _taskExecutor:
-    return SC.Context.getProtocol(KnownProtocolKind::TaskExecutor)
-      ->getDeclaredInterfaceType();
+    if (auto ty = SC.Context.getProtocol(KnownProtocolKind::TaskExecutor)) {
+      return ty->getDeclaredInterfaceType();
+    } else {
+      return nullptr;
+    }
   case _actor:
     return SC.Context.getProtocol(KnownProtocolKind::Actor)
       ->getDeclaredInterfaceType();
@@ -175,6 +178,28 @@ Type synthesizeType(SynthesisContext &SC,
   return ExistentialType::get(synthesizeType(SC, M.Sub));
 }
 
+/// A synthesizer which generates an existential type from a requirement type.
+template <class S, class FallbackS>
+struct BincompatIfTypeAvailableTypeSynthesizer {
+  bool condition;
+  S Sub;
+  FallbackS FallbackSub;
+};
+template <class S, class FallbackS>
+constexpr BincompatIfTypeAvailableTypeSynthesizer<S, FallbackS> _bincompatType(
+    bool bincompatCondition, S sub, FallbackS fallbackSub) {
+  return {bincompatCondition, sub, {fallbackSub}};
+}
+template <class S, class FallbackS>
+Type synthesizeType(SynthesisContext &SC,
+                    const BincompatIfTypeAvailableTypeSynthesizer<S, FallbackS> &M) {
+  if (M.condition) {
+    return synthesizeType(SC, M.Sub);
+  } else {
+    return synthesizeType(SC, M.FallbackSub);
+  }
+}
+
 MetatypeRepresentation
 inline synthesizeMetatypeRepresentation(RepresentationSynthesizer rep) {
   switch (rep) {
@@ -332,6 +357,10 @@ constexpr SpecifiedParamSynthesizer<G> _owned(G sub) {
   return {ParamSpecifier::LegacyOwned, sub};
 }
 template <class G>
+constexpr SpecifiedParamSynthesizer<G> _consuming(G sub) {
+  return {ParamSpecifier::Consuming, sub};
+}
+template <class G>
 constexpr SpecifiedParamSynthesizer<G> _inout(G sub) {
   return {ParamSpecifier::InOut, sub};
 }

lib/AST/Builtins.cpp

@@ -1534,31 +1534,50 @@ Type swift::getAsyncTaskAndContextType(ASTContext &ctx) {
 }
 
 static ValueDecl *getCreateTask(ASTContext &ctx, Identifier id) {
+  auto taskExecutorIsAvailable =
+      ctx.getProtocol(swift::KnownProtocolKind::TaskExecutor) != nullptr;
+
   return getBuiltinFunction(
       ctx, id, _thin, _generics(_unrestricted, _conformsToDefaults(0)),
       _parameters(
         _label("flags", _swiftInt),
         _label("initialSerialExecutor", _defaulted(_optional(_executor), _nil)),
         _label("taskGroup", _defaulted(_optional(_rawPointer), _nil)),
         _label("initialTaskExecutor", _defaulted(_optional(_executor), _nil)),
-        _label("operation", _function(_async(_throws(_sendable(_thick))),
+        _label("initialTaskExecutorConsuming",
+                 _defaulted(_consuming(_optional(_bincompatType(
+                                /*if*/ taskExecutorIsAvailable,
+                                _existential(_taskExecutor),
+                                /*else*/ _executor))),
+                            _nil)),
+          _label("operation", _function(_async(_throws(_sendable(_thick))),
                                       _typeparam(0), _parameters()))),
       _tuple(_nativeObject, _rawPointer));
 }
 
 static ValueDecl *getCreateDiscardingTask(ASTContext &ctx, Identifier id) {
+  auto taskExecutorIsAvailable =
+      ctx.getProtocol(swift::KnownProtocolKind::TaskExecutor) != nullptr;
+
   return getBuiltinFunction(
       ctx, id, _thin,
       _parameters(
         _label("flags", _swiftInt),
         _label("initialSerialExecutor", _defaulted(_optional(_executor), _nil)),
         _label("taskGroup", _defaulted(_optional(_rawPointer), _nil)),
         _label("initialTaskExecutor", _defaulted(_optional(_executor), _nil)),
+        _label("initialTaskExecutorConsuming",
+               _defaulted(_consuming(_optional(_bincompatType(
+                              /*if*/ taskExecutorIsAvailable,
+                              _existential(_taskExecutor),
+                              /*else*/ _executor))),
+                          _nil)),
         _label("operation", _function(_async(_throws(_sendable(_thick))),
                                       _void, _parameters()))),
       _tuple(_nativeObject, _rawPointer));
 }
 
+// Legacy entry point, prefer `createAsyncTask`
 static ValueDecl *getCreateAsyncTask(ASTContext &ctx, Identifier id,
                                      bool inGroup, bool withTaskExecutor,
                                      bool isDiscarding) {
@@ -1568,6 +1587,7 @@ static ValueDecl *getCreateAsyncTask(ASTContext &ctx, Identifier id,
   if (inGroup) {
     builder.addParameter(makeConcrete(ctx.TheRawPointerType)); // group
   }
+
   if (withTaskExecutor) {
     builder.addParameter(makeConcrete(ctx.TheExecutorType)); // executor
   }
@@ -3003,7 +3023,7 @@ ValueDecl *swift::getBuiltinValueDecl(ASTContext &Context, Identifier Id) {
 
   case BuiltinValueKind::FixLifetime:
     return getFixLifetimeOperation(Context, Id);
-      
+
   case BuiltinValueKind::CanBeObjCClass:
     return getCanBeObjCClassOperation(Context, Id);
 

lib/IRGen/GenConcurrency.cpp

@@ -472,7 +472,7 @@ static llvm::Value *addOptionRecord(IRGenFunction &IGF,
 }
 
 /// Add a task option record to the options list if the given value
-/// is presernt.
+/// is present.
 template <class RecordTraits>
 static llvm::Value *maybeAddOptionRecord(IRGenFunction &IGF,
                                          llvm::Value *curRecordPointer,
@@ -645,15 +645,15 @@ struct TaskGroupRecordTraits {
   }
 };
 
-struct InitialTaskExecutorRecordTraits {
+struct InitialTaskExecutorUnownedRecordTraits {
   static StringRef getLabel() {
-    return "task_executor";
+    return "task_executor_unowned";
   }
   static llvm::StructType *getRecordType(IRGenModule &IGM) {
-    return IGM.SwiftInitialTaskExecutorPreferenceTaskOptionRecordTy;
+    return IGM.SwiftInitialTaskExecutorUnownedPreferenceTaskOptionRecordTy;
   }
   static TaskOptionRecordFlags getRecordFlags() {
-    return TaskOptionRecordFlags(TaskOptionRecordKind::InitialTaskExecutor);
+    return TaskOptionRecordFlags(TaskOptionRecordKind::InitialTaskExecutorUnowned);
   }
   static CanType getValueType(ASTContext &ctx) {
     return ctx.TheExecutorType;
@@ -670,6 +670,36 @@ struct InitialTaskExecutorRecordTraits {
   }
 };
 
+struct InitialTaskExecutorOwnedRecordTraits {
+  static StringRef getLabel() {
+    return "task_executor_owned";
+  }
+  static llvm::StructType *getRecordType(IRGenModule &IGM) {
+    return IGM.SwiftInitialTaskExecutorOwnedPreferenceTaskOptionRecordTy;
+  }
+  static TaskOptionRecordFlags getRecordFlags() {
+    return TaskOptionRecordFlags(TaskOptionRecordKind::InitialTaskExecutorOwned);
+  }
+  static CanType getValueType(ASTContext &ctx) {
+    return OptionalType::get(ctx.getProtocol(KnownProtocolKind::TaskExecutor)
+                                 ->getDeclaredInterfaceType())
+        ->getCanonicalType();
+  }
+
+  void initialize(IRGenFunction &IGF, Address recordAddr,
+                  Explosion &taskExecutor) const {
+    auto executorRecord =
+      IGF.Builder.CreateStructGEP(recordAddr, 1, 2 * IGF.IGM.getPointerSize());
+
+    // This relies on the fact that the HeapObject is directly followed by a
+    // pointer to the witness table.
+    IGF.Builder.CreateStore(taskExecutor.claimNext(),
+        IGF.Builder.CreateStructGEP(executorRecord, 0, Size()));
+    IGF.Builder.CreateStore(taskExecutor.claimNext(),
+        IGF.Builder.CreateStructGEP(executorRecord, 1, Size()));
+  }
+};
+
 } // end anonymous namespace
 
 static llvm::Value *
@@ -693,15 +723,25 @@ maybeAddInitialTaskExecutorOptionRecord(IRGenFunction &IGF,
                                         llvm::Value *prevOptions,
                                         OptionalExplosion &taskExecutor) {
   return maybeAddOptionRecord(IGF, prevOptions,
-                              InitialTaskExecutorRecordTraits(),
+                              InitialTaskExecutorUnownedRecordTraits(),
                               taskExecutor);
 }
 
+static llvm::Value *
+maybeAddInitialTaskExecutorOwnedOptionRecord(IRGenFunction &IGF,
+                                        llvm::Value *prevOptions,
+                                        OptionalExplosion &taskExecutorExistential) {
+  return maybeAddOptionRecord(IGF, prevOptions,
+                              InitialTaskExecutorOwnedRecordTraits(),
+                              taskExecutorExistential);
+}
+
 std::pair<llvm::Value *, llvm::Value *>
 irgen::emitTaskCreate(IRGenFunction &IGF, llvm::Value *flags,
                       OptionalExplosion &serialExecutor,
                       OptionalExplosion &taskGroup,
-                      OptionalExplosion &taskExecutor,
+                      OptionalExplosion &taskExecutorUnowned,
+                      OptionalExplosion &taskExecutorExistential,
                       Explosion &taskFunction,
                       SubstitutionMap subs) {
   llvm::Value *taskOptions =
@@ -729,8 +769,14 @@ irgen::emitTaskCreate(IRGenFunction &IGF, llvm::Value *flags,
   taskOptions = maybeAddTaskGroupOptionRecord(IGF, taskOptions, taskGroup);
 
   // Add an option record for the initial task executor, if present.
-  taskOptions =
-    maybeAddInitialTaskExecutorOptionRecord(IGF, taskOptions, taskExecutor);
+  {
+    // Deprecated: This is the UnownedTaskExecutor? which is NOT consuming
+    taskOptions = maybeAddInitialTaskExecutorOptionRecord(
+        IGF, taskOptions, taskExecutorUnowned);
+    // Take an (any TaskExecutor)? which we retain until task has completed
+    taskOptions = maybeAddInitialTaskExecutorOwnedOptionRecord(
+        IGF, taskOptions, taskExecutorExistential);
+  }
 
   // In embedded Swift, create and pass result type info.
   taskOptions = maybeAddEmbeddedSwiftResultTypeInfo(IGF, taskOptions, resultType);

lib/IRGen/GenConcurrency.h

@@ -104,7 +104,8 @@ std::pair<llvm::Value *, llvm::Value *>
 emitTaskCreate(IRGenFunction &IGF, llvm::Value *flags,
                OptionalExplosion &initialExecutor,
                OptionalExplosion &taskGroup,
-               OptionalExplosion &taskExecutor,
+               OptionalExplosion &taskExecutorUnowned,
+               OptionalExplosion &taskExecutorExistential,
                Explosion &taskFunction,
                SubstitutionMap subs);