[Concurrency] Add a unique Task ID to AsyncTask

This commit changes JobFlags storage to be 32bits, but leaves the runtime
API expressed in terms of size_t. This allows us to pack an Id in the
32bits we freed up.

The offset of this Id in the AsyncTask is an ABI constant. This way
introspection tools can extract the currently running task identifier
without any need for special APIs.

Author: fredriss

include/swift/ABI/MetadataValues.h

@@ -2000,7 +2000,7 @@ enum class JobPriority : size_t {
 };
 
 /// Flags for schedulable jobs.
-class JobFlags : public FlagSet<size_t> {
+class JobFlags : public FlagSet<uint32_t> {
 public:
   enum {
     Kind           = 0,
@@ -2019,7 +2019,7 @@ class JobFlags : public FlagSet<size_t> {
     Task_IsContinuingAsyncTask      = 27,
   };
 
-  explicit JobFlags(size_t bits) : FlagSet(bits) {}
+  explicit JobFlags(uint32_t bits) : FlagSet(bits) {}
   JobFlags(JobKind kind) { setKind(kind); }
   JobFlags(JobKind kind, JobPriority priority) {
     setKind(kind);

include/swift/ABI/Task.h

@@ -68,6 +68,9 @@ class alignas(2 * alignof(void*)) Job :
 
   JobFlags Flags;
 
+  // Derived classes can use this to store a Job Id.
+  uint32_t Id = 0;
+
   // We use this union to avoid having to do a second indirect branch
   // when resuming an asynchronous task, which we expect will be the
   // common case.
@@ -121,8 +124,13 @@ class alignas(2 * alignof(void*)) Job :
 };
 
 // The compiler will eventually assume these.
+#if defined(__LP64__) || defined(_WIN64)
 static_assert(sizeof(Job) == 6 * sizeof(void*),
               "Job size is wrong");
+#else
+static_assert(sizeof(Job) == 8 * sizeof(void*),
+              "Job size is wrong");
+#endif
 static_assert(alignof(Job) == 2 * alignof(void*),
               "Job alignment is wrong");
 
@@ -229,6 +237,7 @@ class AsyncTask : public Job {
       Status(ActiveTaskStatus()),
       Local(TaskLocal::Storage()) {
     assert(flags.isAsyncTask());
+    Id = getNextTaskId();
   }
 
   /// Create a task with "immortal" reference counts.
@@ -242,6 +251,7 @@ class AsyncTask : public Job {
       Status(ActiveTaskStatus()),
       Local(TaskLocal::Storage()) {
     assert(flags.isAsyncTask());
+    Id = getNextTaskId();
   }
 
   ~AsyncTask();
@@ -505,13 +515,24 @@ class AsyncTask : public Job {
     return reinterpret_cast<AsyncTask *&>(
         SchedulerPrivate[NextWaitingTaskIndex]);
   }
+
+  /// Get the next non-zero Task ID.
+  uint32_t getNextTaskId() {
+    static std::atomic<uint32_t> Id(1);
+    uint32_t Next = Id.fetch_add(1, std::memory_order_relaxed);
+    if (Next == 0) Next = Id.fetch_add(1, std::memory_order_relaxed);
+    return Next;
+  }
 };
 
 // The compiler will eventually assume these.
 static_assert(sizeof(AsyncTask) == 14 * sizeof(void*),
               "AsyncTask size is wrong");
 static_assert(alignof(AsyncTask) == 2 * alignof(void*),
               "AsyncTask alignment is wrong");
+// Libc hardcodes this offset to extract the TaskID
+static_assert(offsetof(AsyncTask, Id) == 4 * sizeof(void *) + 4,
+              "AsyncTask::Id offset is wrong");
 
 inline void Job::runInFullyEstablishedContext() {
   if (auto task = dyn_cast<AsyncTask>(this))

include/swift/Runtime/Concurrency.h

@@ -38,7 +38,7 @@ struct AsyncTaskAndContext {
 /// 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,
+AsyncTaskAndContext swift_task_create_f(size_t flags,
                              ThinNullaryAsyncSignature::FunctionType *function,
                                         size_t initialContextSize);
 
@@ -51,7 +51,7 @@ using FutureAsyncSignature =
 /// closure.
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
 AsyncTaskAndContext swift_task_create_future(
-    JobFlags flags,
+    size_t flags,
     const Metadata *futureResultType,
     void *closureEntryPoint,
     HeapObject * /* +1 */ closureContext);
@@ -60,7 +60,7 @@ AsyncTaskAndContext swift_task_create_future(
 /// function.
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
 AsyncTaskAndContext swift_task_create_future_f(
-    JobFlags flags,
+    size_t flags,
     const Metadata *futureResultType,
     FutureAsyncSignature::FunctionType *function,
     size_t initialContextSize);
@@ -69,7 +69,7 @@ AsyncTaskAndContext swift_task_create_future_f(
 /// closure, and offer its result to the task group
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
 AsyncTaskAndContext swift_task_create_group_future(
-    JobFlags flags, TaskGroup *group,
+    size_t flags, TaskGroup *group,
     const Metadata *futureResultType,
     void *closureEntryPoint,
     HeapObject * /* +1 */ closureContext);
@@ -78,7 +78,7 @@ AsyncTaskAndContext swift_task_create_group_future(
 /// function, and offer its result to the task group
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
 AsyncTaskAndContext swift_task_create_group_future_f(
-    JobFlags flags,
+    size_t flags,
     TaskGroup *group,
     const Metadata *futureResultType,
     FutureAsyncSignature::FunctionType *function,

lib/IRGen/IRGenFunction.cpp

@@ -527,8 +527,8 @@ static llvm::Value *unsafeContinuationFromTask(IRGenFunction &IGF,
 static llvm::Value *emitLoadOfResumeContextFromTask(IRGenFunction &IGF,
                                                     llvm::Value *task) {
   auto addr = Address(task, IGF.IGM.getPointerAlignment());
-  auto resumeContextAddr =
-    IGF.Builder.CreateStructGEP(addr, 5, 6 * IGF.IGM.getPointerSize());
+  auto resumeContextAddr = IGF.Builder.CreateStructGEP(
+    addr, 6, (5 * IGF.IGM.getPointerSize()) + Size(8));
   llvm::Value *resumeContext = IGF.Builder.CreateLoad(resumeContextAddr);
   if (auto &schema = IGF.getOptions().PointerAuth.TaskResumeContext) {
     auto info = PointerAuthInfo::emit(IGF, schema,

lib/IRGen/IRGenModule.cpp

@@ -610,7 +610,8 @@ IRGenModule::IRGenModule(IRGenerator &irgen,
   SwiftTaskTy = createStructType(*this, "swift.task", {
     RefCountedStructTy,   // object header
     Int8PtrTy, Int8PtrTy, // Job.SchedulerPrivate
-    SizeTy,               // Job.Flags
+    Int32Ty,              // Job.Flags
+    Int32Ty,              // Job.ID
     FunctionPtrTy,        // Job.RunJob/Job.ResumeTask
     SwiftContextPtrTy,    // Task.ResumeContext
     IntPtrTy              // Task.Status
@@ -629,7 +630,8 @@ IRGenModule::IRGenModule(IRGenerator &irgen,
   SwiftJobTy = createStructType(*this, "swift.job", {
     RefCountedStructTy,   // object header
     Int8PtrTy, Int8PtrTy, // SchedulerPrivate
-    SizeTy,               // flags
+    Int32Ty,              // flags
+    Int32Ty,              // ID
     FunctionPtrTy,        // RunJob/ResumeTask
   });
   SwiftJobPtrTy = SwiftJobTy->getPointerTo(DefaultAS);

stdlib/public/CompatibilityOverride/CompatibilityOverrideConcurrency.def

@@ -97,7 +97,7 @@ OVERRIDE_ACTOR(task_switch, void,
                (resumeToContext, resumeFunction, newExecutor))
 
 OVERRIDE_TASK(task_create_group_future_common, AsyncTaskAndContext, , , ,
-              (JobFlags flags, TaskGroup *group,
+              (size_t flags, TaskGroup *group,
                const Metadata *futureResultType,
                FutureAsyncSignature::FunctionType *function,
                void *closureContext, bool isAsyncLetTask,

stdlib/public/Concurrency/AsyncLet.cpp

@@ -106,7 +106,7 @@ static void swift_asyncLet_startImpl(AsyncLet *alet,
   flags.task_setIsChildTask(true);
 
   auto childTaskAndContext = swift_task_create_async_let_future(
-      flags,
+      flags.getOpaqueValue(),
       futureResultType,
       closureEntryPoint,
       closureContext);

stdlib/public/Concurrency/Task.cpp

@@ -386,11 +386,12 @@ static void task_wait_throwing_resume_adapter(SWIFT_ASYNC_CONTEXT AsyncContext *
 /// Also, async-let tasks are not heap allcoated, but allcoated with the parent
 /// task's stack allocator.
 static AsyncTaskAndContext swift_task_create_group_future_commonImpl(
-    JobFlags flags, TaskGroup *group,
+    size_t rawFlags, TaskGroup *group,
     const Metadata *futureResultType,
     FutureAsyncSignature::FunctionType *function,
     void *closureContext, bool isAsyncLetTask,
     size_t initialContextSize) {
+  JobFlags flags(rawFlags);
   assert((futureResultType != nullptr) == flags.task_isFuture());
   assert(!flags.task_isFuture() ||
          initialContextSize >= sizeof(FutureAsyncContext));
@@ -568,32 +569,32 @@ static AsyncTaskAndContext swift_task_create_group_future_commonImpl(
 }
 
 static AsyncTaskAndContext swift_task_create_group_future_common(
-    JobFlags flags, TaskGroup *group, const Metadata *futureResultType,
+    size_t flags, TaskGroup *group, const Metadata *futureResultType,
     FutureAsyncSignature::FunctionType *function,
     void *closureContext, bool isAsyncLetTask,
     size_t initialContextSize);
 
 AsyncTaskAndContext
-swift::swift_task_create_f(JobFlags flags,
+swift::swift_task_create_f(size_t flags,
                 ThinNullaryAsyncSignature::FunctionType *function,
                            size_t initialContextSize) {
   return swift_task_create_future_f(
       flags, nullptr, function, initialContextSize);
 }
 
 AsyncTaskAndContext swift::swift_task_create_future_f(
-    JobFlags flags,
+    size_t flags,
     const Metadata *futureResultType,
     FutureAsyncSignature::FunctionType *function, size_t initialContextSize) {
-  assert(!flags.task_isGroupChildTask() &&
+  assert(!JobFlags(flags).task_isGroupChildTask() &&
   "use swift_task_create_group_future_f to initialize task group child tasks");
   return swift_task_create_group_future_f(
       flags, /*group=*/nullptr, futureResultType,
       function, initialContextSize);
 }
 
 AsyncTaskAndContext swift::swift_task_create_group_future_f(
-    JobFlags flags, TaskGroup *group,
+    size_t flags, TaskGroup *group,
     const Metadata *futureResultType,
     FutureAsyncSignature::FunctionType *function,
     size_t initialContextSize) {
@@ -621,7 +622,7 @@ getAsyncClosureEntryPointAndContextSize(void *function,
           fnPtr->ExpectedContextSize};
 }
 
-AsyncTaskAndContext swift::swift_task_create_future(JobFlags flags,
+AsyncTaskAndContext swift::swift_task_create_future(size_t flags,
                      const Metadata *futureResultType,
                      void *closureEntry,
                      HeapObject * /* +1 */ closureContext) {
@@ -640,7 +641,7 @@ AsyncTaskAndContext swift::swift_task_create_future(JobFlags flags,
       initialContextSize);
 }
 
-AsyncTaskAndContext swift::swift_task_create_async_let_future(JobFlags flags,
+AsyncTaskAndContext swift::swift_task_create_async_let_future(size_t flags,
                      const Metadata *futureResultType,
                      void *closureEntry,
                      void *closureContext) {
@@ -661,7 +662,7 @@ AsyncTaskAndContext swift::swift_task_create_async_let_future(JobFlags flags,
 
 AsyncTaskAndContext
 swift::swift_task_create_group_future(
-                        JobFlags flags, TaskGroup *group,
+                        size_t flags, TaskGroup *group,
                         const Metadata *futureResultType,
                         void *closureEntry,
                         HeapObject * /*+1*/closureContext) {
@@ -853,7 +854,7 @@ void swift::swift_task_runAndBlockThread(const void *function,
   // Set up a task that runs the runAndBlock async function above.
   auto flags = JobFlags(JobKind::Task, JobPriority::Default);
   auto pair = swift_task_create_f(
-      flags,
+      flags.getOpaqueValue(),
       reinterpret_cast<ThinNullaryAsyncSignature::FunctionType *>(
           &runAndBlock_start),
       sizeof(RunAndBlockContext));

stdlib/public/Concurrency/Task.swift

@@ -253,12 +253,12 @@ extension Task {
   /// This is a port of the C++ FlagSet.
   struct JobFlags {
     /// Kinds of schedulable jobs.
-    enum Kind: Int {
+    enum Kind: Int32 {
       case task = 0
     }
 
     /// The actual bit representation of these flags.
-    var bits: Int = 0
+    var bits: Int32 = 0
 
     /// The kind of job described by these flags.
     var kind: Kind {
@@ -277,11 +277,11 @@ extension Task {
     /// The priority given to the job.
     var priority: Priority? {
       get {
-        Priority(rawValue: (bits & 0xFF00) >> 8)
+        Priority(rawValue: (Int(bits) & 0xFF00) >> 8)
       }
 
       set {
-        bits = (bits & ~0xFF00) | ((newValue?.rawValue ?? 0) << 8)
+        bits = (bits & ~0xFF00) | Int32((newValue?.rawValue ?? 0) << 8)
       }
     }
 
@@ -410,7 +410,7 @@ public func detach<T>(
   flags.isFuture = true
 
   // Create the asynchronous task future.
-  let (task, _) = Builtin.createAsyncTaskFuture(flags.bits, operation)
+  let (task, _) = Builtin.createAsyncTaskFuture(Int(flags.bits), operation)
 
   // Enqueue the resulting job.
   _enqueueJobGlobal(Builtin.convertTaskToJob(task))
@@ -463,7 +463,7 @@ public func detach<T>(
   flags.isFuture = true
 
   // Create the asynchronous task future.
-  let (task, _) = Builtin.createAsyncTaskFuture(flags.bits, operation)
+  let (task, _) = Builtin.createAsyncTaskFuture(Int(flags.bits), operation)
 
   // Enqueue the resulting job.
   _enqueueJobGlobal(Builtin.convertTaskToJob(task))
@@ -533,7 +533,7 @@ public func async<T>(
   flags.isContinuingAsyncTask = true
 
   // Create the asynchronous task future.
-  let (task, _) = Builtin.createAsyncTaskFuture(flags.bits, operation)
+  let (task, _) = Builtin.createAsyncTaskFuture(Int(flags.bits), operation)
 
   // Copy all task locals to the newly created task.
   // We must copy them rather than point to the current task since the new task
@@ -578,7 +578,7 @@ public func async<T>(
   flags.isContinuingAsyncTask = true
 
   // Create the asynchronous task future.
-  let (task, _) = Builtin.createAsyncTaskFuture(flags.bits, operation)
+  let (task, _) = Builtin.createAsyncTaskFuture(Int(flags.bits), operation)
 
   // Enqueue the resulting job.
   _enqueueJobGlobal(Builtin.convertTaskToJob(task))

stdlib/public/Concurrency/TaskGroup.swift

@@ -269,7 +269,7 @@ public struct TaskGroup<ChildTaskResult> {
 
     // Create the asynchronous task future.
     let (childTask, _) = Builtin.createAsyncTaskGroupFuture(
-      flags.bits, _group, operation)
+      Int(flags.bits), _group, operation)
 
     // Attach it to the group's task record in the current task.
     _taskGroupAttachChild(group: _group, child: childTask)
@@ -332,7 +332,7 @@ public struct TaskGroup<ChildTaskResult> {
 
     // Create the asynchronous task future.
     let (childTask, _) = Builtin.createAsyncTaskGroupFuture(
-      flags.bits, _group, operation)
+      Int(flags.bits), _group, operation)
 
     // Attach it to the group's task record in the current task.
     _taskGroupAttachChild(group: _group, child: childTask)
@@ -558,7 +558,7 @@ public struct ThrowingTaskGroup<ChildTaskResult, Failure: Error> {
 
     // Create the asynchronous task future.
     let (childTask, _) = Builtin.createAsyncTaskGroupFuture(
-      flags.bits, _group, operation)
+      Int(flags.bits), _group, operation)
 
     // Attach it to the group's task record in the current task.
     _taskGroupAttachChild(group: _group, child: childTask)
@@ -621,7 +621,7 @@ public struct ThrowingTaskGroup<ChildTaskResult, Failure: Error> {
 
     // Create the asynchronous task future.
     let (childTask, _) = Builtin.createAsyncTaskGroupFuture(
-      flags.bits, _group, operation)
+      Int(flags.bits), _group, operation)
 
     // Attach it to the group's task record in the current task.
     _taskGroupAttachChild(group: _group, child: childTask)