Merge pull request #34381 from rjmccall/task-creation-api

Refine task APIs and introduce an API for task creation

Author: rjmccall

include/swift/ABI/MetadataKind.def

@@ -85,6 +85,10 @@ METADATAKIND(HeapGenericLocalVariable,
 METADATAKIND(ErrorObject,
              1 | MetadataKindIsNonType | MetadataKindIsRuntimePrivate)
 
+/// A heap-allocated simple task.
+METADATAKIND(SimpleTask,
+             2 | MetadataKindIsNonType | MetadataKindIsRuntimePrivate)
+
 // getEnumeratedMetadataKind assumes that all the enumerated values here
 // will be <= LastEnumeratedMetadataKind.
 

include/swift/ABI/MetadataValues.h

@@ -1914,9 +1914,8 @@ class JobFlags : public FlagSet<size_t> {
 
     // Kind-specific flags.
 
-    Task_IsHeapObject = 24,
-    Task_IsChildTask  = 25,
-    Task_IsFuture     = 26
+    Task_IsChildTask  = 24,
+    Task_IsFuture     = 25
   };
 
   explicit JobFlags(size_t bits) : FlagSet(bits) {}
@@ -1933,9 +1932,6 @@ class JobFlags : public FlagSet<size_t> {
     return getKind() == JobKind::Task;
   }
 
-  FLAGSET_DEFINE_FLAG_ACCESSORS(Task_IsHeapObject,
-                                task_isHeapObject,
-                                task_setIsHeapObject)
   FLAGSET_DEFINE_FLAG_ACCESSORS(Task_IsChildTask,
                                 task_isChildTask,
                                 task_setIsChildTask)
@@ -2034,8 +2030,8 @@ class AsyncContextFlags : public FlagSet<uint32_t> {
   /// allocated as part of the caller's context, or if the callee will
   /// be called multiple times.
   FLAGSET_DEFINE_FLAG_ACCESSORS(ShouldNotDeallocate,
-                                shouldNotDeallocateInCaller,
-                                setShouldNotDeallocateInCaller)
+                                shouldNotDeallocateInCallee,
+                                setShouldNotDeallocateInCallee)
 };
 
 } // end namespace swift

include/swift/ABI/Task.h

@@ -17,6 +17,7 @@
 #ifndef SWIFT_ABI_TASK_H
 #define SWIFT_ABI_TASK_H
 
+#include "swift/Basic/RelativePointer.h"
 #include "swift/ABI/HeapObject.h"
 #include "swift/ABI/MetadataValues.h"
 #include "swift/Runtime/Config.h"
@@ -32,7 +33,21 @@ class TaskStatusRecord;
 
 /// An ExecutorRef isn't necessarily just a pointer to an executor
 /// object; it may have other bits set.
-using ExecutorRef = Executor *;
+struct ExecutorRef {
+  Executor *Pointer;
+
+  /// Get an executor ref that represents a lack of preference about
+  /// where execution resumes.  This is only valid in continuations,
+  /// return contexts, and so on; it is not generally passed to
+  /// executing functions.
+  static ExecutorRef noPreference() {
+    return { nullptr };
+  }
+
+  bool operator==(ExecutorRef other) const {
+    return Pointer == other.Pointer;
+  }
+};
 
 using JobInvokeFunction =
   SWIFT_CC(swift)
@@ -42,6 +57,33 @@ using TaskContinuationFunction =
   SWIFT_CC(swift)
   void (AsyncTask *, ExecutorRef, AsyncContext *);
 
+template <class Fn>
+struct AsyncFunctionTypeImpl;
+template <class Result, class... Params>
+struct AsyncFunctionTypeImpl<Result(Params...)> {
+  // TODO: expand and include the arguments in the parameters.
+  using type = TaskContinuationFunction;
+};
+
+template <class Fn>
+using AsyncFunctionType = typename AsyncFunctionTypeImpl<Fn>::type;
+
+/// A "function pointer" for an async function.
+///
+/// Eventually, this will always be signed with the data key
+/// using a type-specific discriminator.
+template <class FnType>
+class AsyncFunctionPointer {
+public:
+  /// The function to run.
+  RelativeDirectPointer<AsyncFunctionType<FnType>,
+                        /*nullable*/ false,
+                        int32_t> Function;
+
+  /// The expected size of the context.
+  uint32_t ExpectedContextSize;
+};
+
 /// A schedulable job.
 class alignas(2 * alignof(void*)) Job {
 public:
@@ -76,7 +118,7 @@ class alignas(2 * alignof(void*)) Job {
   }
 
   /// Run this job.
-  void run(Executor *currentExecutor);
+  void run(ExecutorRef currentExecutor);
 };
 
 // The compiler will eventually assume these.
@@ -128,7 +170,7 @@ class ActiveTaskStatus {
 /// An asynchronous task.  Tasks are the analogue of threads for
 /// asynchronous functions: that is, they are a persistent identity
 /// for the overall async computation.
-class AsyncTask : public Job {
+class AsyncTask : public HeapObject, public Job {
 public:
   /// The context for resuming the job.  When a task is scheduled
   /// as a job, the next continuation should be installed as the
@@ -146,9 +188,10 @@ class AsyncTask : public Job {
   /// Reserved for the use of the task-local stack allocator.
   void *AllocatorPrivate[4];
 
-  AsyncTask(JobFlags flags, TaskContinuationFunction *run,
+  AsyncTask(const HeapMetadata *metadata, JobFlags flags,
+            TaskContinuationFunction *run,
             AsyncContext *initialContext)
-    : Job(flags, run),
+    : HeapObject(metadata), Job(flags, run),
       ResumeContext(initialContext),
       Status(ActiveTaskStatus()) {
     assert(flags.isAsyncTask());
@@ -164,12 +207,6 @@ class AsyncTask : public Job {
     return Status.load(std::memory_order_relaxed).isCancelled();
   }
 
-  bool isHeapObject() const { return Flags.task_isHeapObject(); }
-  HeapObject *heapObjectHeader() {
-    assert(isHeapObject());
-    return reinterpret_cast<HeapObject*>(this) - 1;
-  }
-
   /// A fragment of an async task structure that happens to be a child task.
   class ChildFragment {
     /// The parent task of this task.
@@ -182,6 +219,8 @@ class AsyncTask : public Job {
     AsyncTask *NextChild = nullptr;
 
   public:
+    ChildFragment(AsyncTask *parent) : Parent(parent) {}
+
     AsyncTask *getParent() const {
       return Parent;
     }
@@ -191,6 +230,8 @@ class AsyncTask : public Job {
     }
   };
 
+  bool isFuture() const { return Flags.task_isFuture(); }
+
   bool hasChildFragment() const { return Flags.task_isChildTask(); }
   ChildFragment *childFragment() {
     assert(hasChildFragment());
@@ -205,7 +246,7 @@ class AsyncTask : public Job {
 };
 
 // The compiler will eventually assume these.
-static_assert(sizeof(AsyncTask) == 10 * sizeof(void*),
+static_assert(sizeof(AsyncTask) == 12 * sizeof(void*),
               "AsyncTask size is wrong");
 static_assert(alignof(AsyncTask) == 2 * alignof(void*),
               "AsyncTask alignment is wrong");
@@ -237,6 +278,9 @@ class alignas(MaximumAlignment) AsyncContext {
   TaskContinuationFunction * __ptrauth_swift_async_context_resume
     ResumeParent;
 
+  /// The executor that the parent needs to be resumed on.
+  ExecutorRef ResumeParentExecutor;
+
   /// Flags describing this context.
   ///
   /// Note that this field is only 32 bits; any alignment padding
@@ -245,31 +289,44 @@ class alignas(MaximumAlignment) AsyncContext {
   /// is of course interrupted by the YieldToParent field.
   AsyncContextFlags Flags;
 
-  // Fields following this point may not be valid in all instances
-  // of AsyncContext.
-
-  /// The function to call to temporarily resume running in the
-  /// parent context temporarily.  Generally this means a semantic
-  /// yield.  Requires Flags.hasYieldFunction().
-  TaskContinuationFunction * __ptrauth_swift_async_context_yield
-    YieldToParent;
-
-  AsyncContext(AsyncContextFlags flags,
-               TaskContinuationFunction *resumeParent,
-               AsyncContext *parent)
-    : Parent(parent), ResumeParent(resumeParent), Flags(flags) {}
-
   AsyncContext(AsyncContextFlags flags,
                TaskContinuationFunction *resumeParent,
-               TaskContinuationFunction *yieldToParent,
+               ExecutorRef resumeParentExecutor,
                AsyncContext *parent)
-    : Parent(parent), ResumeParent(resumeParent), Flags(flags),
-      YieldToParent(yieldToParent) {}
+    : Parent(parent), ResumeParent(resumeParent),
+      ResumeParentExecutor(resumeParentExecutor),
+      Flags(flags) {}
 
   AsyncContext(const AsyncContext &) = delete;
   AsyncContext &operator=(const AsyncContext &) = delete;
 };
 
+/// An async context that supports yielding.
+class YieldingAsyncContext : public AsyncContext {
+public:
+  /// The function to call to temporarily resume running in the
+  /// parent context.  Generally this means a semantic yield.
+  TaskContinuationFunction * __ptrauth_swift_async_context_yield
+    YieldToParent;
+
+  /// The executor that the parent context needs to be yielded to on.
+  ExecutorRef YieldToParentExecutor;
+
+  YieldingAsyncContext(AsyncContextFlags flags,
+                       TaskContinuationFunction *resumeParent,
+                       ExecutorRef resumeParentExecutor,
+                       TaskContinuationFunction *yieldToParent,
+                       ExecutorRef yieldToParentExecutor,
+                       AsyncContext *parent)
+    : AsyncContext(flags, resumeParent, resumeParentExecutor, parent),
+      YieldToParent(yieldToParent),
+      YieldToParentExecutor(yieldToParentExecutor) {}
+
+  static bool classof(const AsyncContext *context) {
+    return context->Flags.getKind() == AsyncContextKind::Yielding;
+  }
+};
+
 } // end namespace swift
 
 #endif

include/swift/ABI/TaskStatus.h

@@ -82,7 +82,7 @@ class TaskStatusRecord {
   }
 };
 
-TaskStatusRecord *
+inline TaskStatusRecord *
 ActiveTaskStatus::getStatusRecordParent(TaskStatusRecord *ptr) {
   return ptr->getParent();
 }

include/swift/Runtime/Concurrency.h

@@ -21,6 +21,50 @@
 
 namespace swift {
 
+struct AsyncTaskAndContext {
+  AsyncTask *Task;
+  AsyncContext *InitialContext;
+};
+
+/// Create a task object with no future which will run the given
+/// function.
+///
+/// The task is not yet scheduled.
+///
+/// If a parent task is provided, flags.task_hasChildFragment() must
+/// be true, and this must be called synchronously with the parent.
+/// The parent is responsible for creating a ChildTaskStatusRecord.
+/// TODO: should we have a single runtime function for creating a task
+/// and doing this child task status record management?
+SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
+AsyncTaskAndContext swift_task_create(JobFlags flags,
+                                      AsyncTask *parent,
+                                const AsyncFunctionPointer<void()> *function);
+
+/// Create a task object with no future which will run the given
+/// function.
+SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
+AsyncTaskAndContext swift_task_create_f(JobFlags flags,
+                                        AsyncTask *parent,
+                                        AsyncFunctionType<void()> *function,
+                                        size_t initialContextSize);
+
+/// Allocate memory in a task.
+///
+/// This must be called synchronously with the task.
+///
+/// All allocations will be rounded to a multiple of MAX_ALIGNMENT.
+SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
+void *swift_task_alloc(AsyncTask *task, size_t size);
+
+/// Deallocate memory in a task.
+///
+/// The pointer provided must be the last pointer allocated on
+/// this task that has not yet been deallocated; that is, memory
+/// must be allocated and deallocated in a strict stack discipline.
+SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
+void swift_task_dealloc(AsyncTask *task, void *ptr);
+
 /// Cancel a task and all of its child tasks.
 ///
 /// This can be called from any thread.

include/swift/Runtime/RuntimeFunctions.def

@@ -1461,17 +1461,17 @@ FUNCTION(GetTypeByMangledNameInContextInMetadataState,
               Int8PtrPtrTy),
          ATTRS(NoUnwind, ArgMemOnly))
 
-// void *swift_taskAlloc(SwiftTask *task, size_t size);
+// void *swift_task_alloc(AsyncTask *task, size_t size);
 FUNCTION(TaskAlloc,
-         swift_taskAlloc, SwiftCC,
+         swift_task_alloc, SwiftCC,
          ConcurrencyAvailability,
          RETURNS(Int8PtrTy),
          ARGS(SwiftTaskPtrTy, SizeTy),
          ATTRS(NoUnwind, ArgMemOnly))
 
-// void swift_taskDealloc(SwiftTask *task, void *ptr);
+// void swift_task_dealloc(AsyncTask *task, void *ptr);
 FUNCTION(TaskDealloc,
-         swift_taskDealloc, SwiftCC,
+         swift_task_dealloc, SwiftCC,
          ConcurrencyAvailability,
          RETURNS(VoidTy),
          ARGS(SwiftTaskPtrTy, Int8PtrTy),

stdlib/public/Concurrency/CMakeLists.txt

@@ -13,6 +13,8 @@
 add_swift_target_library(swift_Concurrency ${SWIFT_STDLIB_LIBRARY_BUILD_TYPES} IS_STDLIB
   Actor.swift
   PartialAsyncTask.swift
+  Task.cpp
+  TaskAlloc.cpp
   TaskStatus.cpp
   Mutex.cpp