[Concurrency] TaskPool, close cousin to TaskGroup

include/swift/ABI/MetadataValues.h

@@ -55,6 +55,9 @@ enum {
   /// The number of words in a task group.
   NumWords_TaskGroup = 32,
 
+  /// The number of words in a task pool.
+  NumWords_TaskPool = 32,
+
   /// The number of words in an AsyncLet (flags + child task context & allocation)
   NumWords_AsyncLet = 80, // 640 bytes ought to be enough for anyone
 
@@ -145,6 +148,9 @@ const size_t Alignment_TaskGroup = MaximumAlignment;
 /// The alignment of an AsyncLet.
 const size_t Alignment_AsyncLet = MaximumAlignment;
 
+/// The alignment of a TaskPool.
+const size_t Alignment_TaskPool = MaximumAlignment;
+
 /// Flags stored in the value-witness table.
 template <typename int_type>
 class TargetValueWitnessFlags {
@@ -2249,17 +2255,17 @@ class TaskCreateFlags : public FlagSet<size_t> {
 public:
   enum {
     // Priority that user specified while creating the task
-    RequestedPriority = 0,
-    RequestedPriority_width = 8,
-
-    Task_IsChildTask                              = 8,
-    // Should only be set in task-to-thread model where Task.runInline is
-    // available
-    Task_IsInlineTask                             = 9,
-    Task_CopyTaskLocals                           = 10,
-    Task_InheritContext                           = 11,
-    Task_EnqueueJob                               = 12,
-    Task_AddPendingGroupTaskUnconditionally       = 13,
+    RequestedPriority                       = 0,
+    RequestedPriority_width                 = 8,
+
+    Task_IsChildTask                        = 8,
+    // Should only be set in task-to-thread model
+    // where Task.runInline is available
+    Task_IsInlineTask                       = 9,
+    Task_CopyTaskLocals                     = 10,
+    Task_InheritContext                     = 11,
+    Task_EnqueueJob                         = 12,
+    Task_AddPendingGroupTaskUnconditionally = 13, // used for: TaskGroup, TaskPool
   };
 
   explicit constexpr TaskCreateFlags(size_t bits) : FlagSet(bits) {}
@@ -2284,8 +2290,8 @@ class TaskCreateFlags : public FlagSet<size_t> {
                                 enqueueJob,
                                 setEnqueueJob)
   FLAGSET_DEFINE_FLAG_ACCESSORS(Task_AddPendingGroupTaskUnconditionally,
-                                addPendingGroupTaskUnconditionally,
-                                setAddPendingGroupTaskUnconditionally)
+                                addPendingTaskUnconditionally,
+                                setaddPendingTaskUnconditionally)
 };
 
 /// Flags for schedulable jobs.
@@ -2305,7 +2311,7 @@ class JobFlags : public FlagSet<uint32_t> {
     Task_IsChildTask           = 24,
     Task_IsFuture              = 25,
     Task_IsGroupChildTask      = 26,
-    // 27 is currently unused
+    Task_IsPoolChildTask       = 27,
     Task_IsAsyncLetTask        = 28,
   };
 
@@ -2339,6 +2345,9 @@ class JobFlags : public FlagSet<uint32_t> {
   FLAGSET_DEFINE_FLAG_ACCESSORS(Task_IsAsyncLetTask,
                                 task_isAsyncLetTask,
                                 task_setIsAsyncLetTask)
+  FLAGSET_DEFINE_FLAG_ACCESSORS(Task_IsPoolChildTask,
+                                task_isPoolChildTask,
+                                task_setIsPoolChildTask)
 };
 
 /// Kinds of task status record.
@@ -2363,6 +2372,9 @@ enum class TaskStatusRecordKind : uint8_t {
   /// escalated.
   EscalationNotification = 4,
 
+  /// TaskPool
+  TaskPool = 5,
+
   // Kinds >= 192 are private to the implementation.
   First_Reserved = 192,
   Private_RecordLock = 192
@@ -2379,6 +2391,8 @@ enum class TaskOptionRecordKind : uint8_t {
   AsyncLet  = 2,
   /// Request a child task for an 'async let'.
   AsyncLetWithBuffer = 3,
+  /// Request a child task to be part of a specific task pool.
+  TaskPool = 4,
   /// Request a child task for swift_task_run_inline.
   RunInline = UINT8_MAX,
 };

include/swift/ABI/Task.h

@@ -39,6 +39,7 @@ struct SwiftError;
 class TaskStatusRecord;
 class TaskOptionRecord;
 class TaskGroup;
+class TaskPool;
 
 extern FullMetadata<DispatchClassMetadata> jobHeapMetadata;
 
@@ -183,11 +184,11 @@ class NullaryContinuationJob : public Job {
 /// ### Fragments
 /// An AsyncTask may have the following fragments:
 ///
-///    +--------------------------+
-///    | childFragment?           |
-///    | groupChildFragment?      |
-///    | futureFragment?          |*
-///    +--------------------------+
+///    +------------------------------------------------+
+///    | childFragment?                                 |
+///    | (groupChildFragment | poolChildFragment)?      |
+///    | futureFragment?                                |*
+///    +------------------------------------------------+
 ///
 /// * The future fragment is dynamic in size, based on the future result type
 ///   it can hold, and thus must be the *last* fragment.
@@ -425,6 +426,7 @@ class AsyncTask : public Job {
 
   GroupChildFragment *groupChildFragment() {
     assert(hasGroupChildFragment());
+    assert(!hasPoolChildFragment()); // pool and group are mutually exclusive
 
     auto offset = reinterpret_cast<char*>(this);
     offset += sizeof(AsyncTask);
@@ -434,6 +436,45 @@ class AsyncTask : public Job {
     return reinterpret_cast<GroupChildFragment *>(offset);
   }
 
+  // ==== TaskPool Child ------------------------------------------------------
+
+  /// A child task created by `pool.addTask` is called a "task pool child."
+  /// Upon completion, in addition to the usual future notifying all its waiters,
+  /// it must also `pool->offer` itself to the pool.
+  class PoolChildFragment {
+  private:
+    TaskPool* Pool;
+
+    friend class AsyncTask;
+    friend class TaskPool;
+
+  public:
+    explicit PoolChildFragment(TaskPool *pool)
+        : Pool(pool) {}
+
+    /// Return the group this task should offer into when it completes.
+    TaskPool* getPool() {
+      return Pool;
+    }
+  };
+
+  // Checks if task is a child of a TaskPool task.
+  //
+  // A child task that is a group child knows that it's parent is a group
+  // and therefore may `groupOffer` to it upon completion.
+  bool hasPoolChildFragment() const { return Flags.task_isPoolChildTask(); }
+
+  PoolChildFragment *poolChildFragment() {
+    assert(hasPoolChildFragment());
+
+    auto offset = reinterpret_cast<char*>(this);
+    offset += sizeof(AsyncTask);
+    if (hasChildFragment())
+      offset += sizeof(ChildFragment);
+
+    return reinterpret_cast<PoolChildFragment *>(offset);
+  }
+
   // ==== Future ---------------------------------------------------------------
 
   class FutureFragment {
@@ -549,12 +590,15 @@ class AsyncTask : public Job {
 
   FutureFragment *futureFragment() {
     assert(isFuture());
-    auto offset = reinterpret_cast<char*>(this);
+    auto offset = reinterpret_cast<char *>(this);
     offset += sizeof(AsyncTask);
     if (hasChildFragment())
       offset += sizeof(ChildFragment);
-    if (hasGroupChildFragment())
+    if (hasGroupChildFragment()) {
       offset += sizeof(GroupChildFragment);
+    } else if (hasPoolChildFragment()) {
+      offset += sizeof(PoolChildFragment);
+    }
 
     return reinterpret_cast<FutureFragment *>(offset);
   }

include/swift/ABI/TaskLocal.h

@@ -27,6 +27,7 @@ struct OpaqueValue;
 struct SwiftError;
 class TaskStatusRecord;
 class TaskGroup;
+class TaskPool;
 
 // ==== Task Locals Values ---------------------------------------------------
 

include/swift/ABI/TaskOptions.h

@@ -75,6 +75,23 @@ class TaskGroupTaskOptionRecord : public TaskOptionRecord {
   }
 };
 
+class TaskPoolTaskOptionRecord : public TaskOptionRecord {
+  TaskPool * const Pool;
+
+  public:
+    TaskPoolTaskOptionRecord(TaskPool *pool)
+    : TaskOptionRecord(TaskOptionRecordKind::TaskPool),
+      Pool(pool) {}
+
+  TaskPool *getPool() const {
+    return Pool;
+  }
+
+  static bool classof(const TaskOptionRecord *record) {
+    return record->getKind() == TaskOptionRecordKind::TaskPool;
+  }
+};
+
 
 /// Task option to specify on what executor the task should be executed.
 ///

include/swift/ABI/TaskPool.h

@@ -0,0 +1,61 @@
+//===--- TaskPool.h - ABI structures for task pools  -00--------*- C++ -*-===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2022 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+//
+// Swift ABI describing task pools.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SWIFT_ABI_TASK_POOL_H
+#define SWIFT_ABI_TASK_POOL_H
+
+#include "swift/ABI/Task.h"
+#include "swift/ABI/TaskStatus.h"
+#include "swift/ABI/HeapObject.h"
+#include "swift/Runtime/Concurrency.h"
+#include "swift/Runtime/Config.h"
+#include "swift/Basic/RelativePointer.h"
+#include "swift/Basic/STLExtras.h"
+
+namespace swift {
+
+/// The task pool is responsible for maintaining dynamically created child tasks.
+class alignas(Alignment_TaskPool) TaskPool {
+public:
+  // These constructors do not initialize the group instance, and the
+  // destructor does not destroy the group instance; you must call
+  // swift_taskGroup_{initialize,destroy} yourself.
+  constexpr TaskPool()
+    : PrivateData{} {}
+
+  void *PrivateData[NumWords_TaskPool];
+
+  /// Upon a future task's completion, offer it to the task pool it belongs to.
+  void offer(AsyncTask *completed, AsyncContext *context);
+
+  /// Checks the cancellation status of the group.
+  bool isCancelled();
+
+  // Add a child task to the task pool. Always called while holding the
+  // status record lock of the task pool's owning task.
+  void addChildTask(AsyncTask *task);
+
+  // Remove a child task from the task pool. Always called while holding
+  // the status record lock of the task pool's owning task.
+  void removeChildTask(AsyncTask *task);
+
+  // Provide accessor for task pool's status record
+  TaskPoolTaskStatusRecord *getTaskRecord();
+};
+
+} // end namespace swift
+
+#endif // SWIFT_ABI_TASK_POOL_H

include/swift/ABI/TaskStatus.h

@@ -260,6 +260,95 @@ class TaskGroupTaskStatusRecord : public TaskStatusRecord {
   }
 };
 
+class TaskPoolTaskStatusRecord : public TaskStatusRecord {
+  AsyncTask *FirstChild;
+  AsyncTask *LastChild;
+
+public:
+  TaskPoolTaskStatusRecord()
+      : TaskStatusRecord(TaskStatusRecordKind::TaskPool),
+        FirstChild(nullptr),
+        LastChild(nullptr) {
+  }
+
+  TaskPoolTaskStatusRecord(AsyncTask *child)
+      : TaskStatusRecord(TaskStatusRecordKind::TaskPool),
+        FirstChild(child),
+        LastChild(child) {
+    assert(!LastChild || !LastChild->childFragment()->getNextChild());
+  }
+
+  TaskPool *getPool() { return reinterpret_cast<TaskPool *>(this); }
+
+  /// Return the first child linked by this record.  This may be null;
+  /// if not, it (and all of its successors) are guaranteed to satisfy
+  /// `isChildTask()`.
+  AsyncTask *getFirstChild() const { return FirstChild; }
+
+  /// Attach the passed in `child` task to this group.
+  void attachChild(AsyncTask *child) {
+    assert(child->hasPoolChildFragment());
+    assert(child->poolChildFragment()->getPool() == getPool());
+
+    auto oldLastChild = LastChild;
+    LastChild = child;
+
+    if (!FirstChild) {
+      // This is the first child we ever attach, so store it as FirstChild.
+      FirstChild = child;
+      return;
+    }
+
+    oldLastChild->childFragment()->setNextChild(child);
+  }
+
+  void detachChild(AsyncTask *child) {
+    assert(child && "cannot remove a null child from group");
+    if (FirstChild == child) {
+      FirstChild = getNextChildTask(child);
+      if (FirstChild == nullptr) {
+        LastChild = nullptr;
+      }
+      return;
+    }
+
+    AsyncTask *prev = FirstChild;
+    // Remove the child from the linked list, i.e.:
+    //     prev -> afterPrev -> afterChild
+    //                 ==
+    //               child   -> afterChild
+    // Becomes:
+    //     prev --------------> afterChild
+    while (prev) {
+      auto afterPrev = getNextChildTask(prev);
+
+      if (afterPrev == child) {
+        auto afterChild = getNextChildTask(child);
+        prev->childFragment()->setNextChild(afterChild);
+        if (child == LastChild) {
+          LastChild = prev;
+        }
+        return;
+      }
+
+      prev = afterPrev;
+    }
+  }
+
+  static AsyncTask *getNextChildTask(AsyncTask *task) {
+    return task->childFragment()->getNextChild();
+  }
+
+  using child_iterator = LinkedListIterator<AsyncTask, getNextChildTask>;
+  llvm::iterator_range<child_iterator> children() const {
+    return child_iterator::rangeBeginning(getFirstChild());
+  }
+
+  static bool classof(const TaskStatusRecord *record) {
+    return record->getKind() == TaskStatusRecordKind::TaskGroup;
+  }
+};
+
 /// A cancellation record which states that a task has an arbitrary
 /// function that needs to be called if the task is cancelled.
 ///