Future-proof the layout of AsyncTask.

Author: rjmccall

include/swift/ABI/MetadataValues.h

@@ -47,6 +47,9 @@ enum {
   /// in a default actor.
   NumWords_DefaultActor = 10,
 
+  /// The number of words in a task.
+  NumWords_AsyncTask = 16,
+
   /// The number of words in a task group.
   NumWords_TaskGroup = 32,
 

include/swift/ABI/Task.h

@@ -126,7 +126,7 @@ class alignas(2 * alignof(void*)) Job :
 };
 
 // The compiler will eventually assume these.
-#if defined(__LP64__) || defined(_WIN64)
+#if SWIFT_POINTER_IS_8_BYTES
 static_assert(sizeof(Job) == 6 * sizeof(void*),
               "Job size is wrong");
 #else
@@ -136,46 +136,6 @@ static_assert(sizeof(Job) == 8 * sizeof(void*),
 static_assert(alignof(Job) == 2 * alignof(void*),
               "Job alignment is wrong");
 
-/// The current state of a task's status records.
-class ActiveTaskStatus {
-  enum : uintptr_t {
-    IsCancelled = 0x1,
-    IsLocked = 0x2,
-    RecordMask = ~uintptr_t(IsCancelled | IsLocked)
-  };
-
-  uintptr_t Value;
-
-public:
-  constexpr ActiveTaskStatus() : Value(0) {}
-  ActiveTaskStatus(TaskStatusRecord *innermostRecord,
-                   bool cancelled, bool locked)
-    : Value(reinterpret_cast<uintptr_t>(innermostRecord)
-                + (locked ? IsLocked : 0)
-                + (cancelled ? IsCancelled : 0)) {}
-
-  /// Is the task currently cancelled?
-  bool isCancelled() const { return Value & IsCancelled; }
-
-  /// Is there an active lock on the cancellation information?
-  bool isLocked() const { return Value & IsLocked; }
-
-  /// Return the innermost cancellation record.  Code running
-  /// asynchronously with this task should not access this record
-  /// without having first locked it; see swift_taskCancel.
-  TaskStatusRecord *getInnermostRecord() const {
-    return reinterpret_cast<TaskStatusRecord*>(Value & RecordMask);
-  }
-
-  static TaskStatusRecord *getStatusRecordParent(TaskStatusRecord *ptr);
-
-  using record_iterator =
-    LinkedListIterator<TaskStatusRecord, getStatusRecordParent>;
-  llvm::iterator_range<record_iterator> records() const {
-    return record_iterator::rangeBeginning(getInnermostRecord());
-  }
-};
-
 class NullaryContinuationJob : public Job {
 
 private:
@@ -212,6 +172,16 @@ class NullaryContinuationJob : public Job {
 ///   it can hold, and thus must be the *last* fragment.
 class AsyncTask : public Job {
 public:
+  // On 32-bit targets, there is a word of tail padding remaining
+  // in Job, and ResumeContext will fit into that, at offset 28.
+  // Private then has offset 32.
+  // On 64-bit targets, there is no tail padding in Job, and so
+  // ResumeContext has offset 48.  There is therefore another word
+  // of reserved storage prior to Private (which needs to have
+  // double-word alignment), which has offset 64.
+  // We therefore converge and end up with 16 words of storage on
+  // all platforms.
+
   /// The context for resuming the job.  When a task is scheduled
   /// as a job, the next continuation should be installed as the
   /// ResumeTask pointer in the job header, with this serving as
@@ -222,36 +192,57 @@ class AsyncTask : public Job {
   /// prevent it from being corrupted in flight.
   AsyncContext * __ptrauth_swift_task_resume_context ResumeContext;
 
-  /// The currently-active information about cancellation.
-  std::atomic<ActiveTaskStatus> Status;
+#if SWIFT_POINTER_IS_8_BYTES
+  void *Reserved64;
+#endif
 
-  /// Reserved for the use of the task-local stack allocator.
-  void *AllocatorPrivate[4];
+  struct PrivateStorage;
 
-  /// Task local values storage container.
-  TaskLocal::Storage Local;
+  /// Private storage for the use of the runtime.
+  struct alignas(2 * alignof(void*)) OpaquePrivateStorage {
+    void *Storage[8];
 
+    /// Initialize this storage during the creation of a task.
+    void initialize(AsyncTask *task);
+    void initializeWithSlab(AsyncTask *task,
+                            void *slab, size_t slabCapacity);
+
+    /// React to the completion of the enclosing task's execution.
+    void complete(AsyncTask *task);
+
+    /// React to the final destruction of the enclosing task.
+    void destroy();
+
+    PrivateStorage &get();
+    const PrivateStorage &get() const;
+  };
+  PrivateStorage &_private();
+  const PrivateStorage &_private() const;
+
+  OpaquePrivateStorage Private;
+
+  /// Create a task.
+  /// This does not initialize Private; callers must call
+  /// Private.initialize separately.
   AsyncTask(const HeapMetadata *metadata, JobFlags flags,
             TaskContinuationFunction *run,
             AsyncContext *initialContext)
     : Job(flags, run, metadata),
-      ResumeContext(initialContext),
-      Status(ActiveTaskStatus()),
-      Local(TaskLocal::Storage()) {
+      ResumeContext(initialContext) {
     assert(flags.isAsyncTask());
     Id = getNextTaskId();
   }
 
   /// Create a task with "immortal" reference counts.
   /// Used for async let tasks.
+  /// This does not initialize Private; callers must call
+  /// Private.initialize separately.
   AsyncTask(const HeapMetadata *metadata, InlineRefCounts::Immortal_t immortal,
             JobFlags flags,
             TaskContinuationFunction *run,
             AsyncContext *initialContext)
     : Job(flags, run, metadata, immortal),
-      ResumeContext(initialContext),
-      Status(ActiveTaskStatus()),
-      Local(TaskLocal::Storage()) {
+      ResumeContext(initialContext) {
     assert(flags.isAsyncTask());
     Id = getNextTaskId();
   }
@@ -269,25 +260,18 @@ class AsyncTask : public Job {
 
   /// Check whether this task has been cancelled.
   /// Checking this is, of course, inherently race-prone on its own.
-  bool isCancelled() const {
-    return Status.load(std::memory_order_relaxed).isCancelled();
-  }
+  bool isCancelled() const;
 
   // ==== Task Local Values ----------------------------------------------------
 
   void localValuePush(const HeapObject *key,
-                      /* +1 */ OpaqueValue *value, const Metadata *valueType) {
-    Local.pushValue(this, key, value, valueType);
-  }
+                      /* +1 */ OpaqueValue *value,
+                      const Metadata *valueType);
 
-  OpaqueValue* localValueGet(const HeapObject *key) {
-    return Local.getValue(this, key);
-  }
+  OpaqueValue *localValueGet(const HeapObject *key);
 
   /// Returns true if storage has still more bindings.
-  bool localValuePop() {
-    return Local.popValue(this);
-  }
+  bool localValuePop();
 
   // ==== Child Fragment -------------------------------------------------------
 
@@ -529,7 +513,7 @@ class AsyncTask : public Job {
 };
 
 // The compiler will eventually assume these.
-static_assert(sizeof(AsyncTask) == 14 * sizeof(void*),
+static_assert(sizeof(AsyncTask) == NumWords_AsyncTask * sizeof(void*),
               "AsyncTask size is wrong");
 static_assert(alignof(AsyncTask) == 2 * alignof(void*),
               "AsyncTask alignment is wrong");

include/swift/ABI/TaskStatus.h

@@ -82,11 +82,6 @@ class TaskStatusRecord {
   }
 };
 
-inline TaskStatusRecord *
-ActiveTaskStatus::getStatusRecordParent(TaskStatusRecord *ptr) {
-  return ptr->getParent();
-}
-
 /// A deadline for the task.  If this is reached, the task will be
 /// automatically cancelled.  The deadline can also be queried and used
 /// in other ways.

include/swift/Basic/Compiler.h

@@ -128,4 +128,13 @@
 #define SWIFT_ASSERT_ONLY(...) do { __VA_ARGS__; } while (false)
 #endif
 
+#if defined(__LP64__) || defined(_WIN64)
+#define SWIFT_POINTER_IS_8_BYTES 1
+#define SWIFT_POINTER_IS_4_BYTES 0
+#else
+// TODO: consider supporting 16-bit targets
+#define SWIFT_POINTER_IS_8_BYTES 0
+#define SWIFT_POINTER_IS_4_BYTES 1
+#endif
+
 #endif // SWIFT_BASIC_COMPILER_H

stdlib/public/Concurrency/AsyncLet.cpp

@@ -86,10 +86,6 @@ static AsyncLetImpl *asImpl(const AsyncLet *alet) {
       const_cast<AsyncLet*>(alet));
 }
 
-static AsyncLet *asAbstract(AsyncLetImpl *alet) {
-  return reinterpret_cast<AsyncLet*>(alet);
-}
-
 // =============================================================================
 // ==== start ------------------------------------------------------------------
 
@@ -134,7 +130,6 @@ SWIFT_CC(swiftasync)
 static void swift_asyncLet_waitImpl(
     OpaqueValue *result, SWIFT_ASYNC_CONTEXT AsyncContext *rawContext,
     AsyncLet *alet, Metadata *T) {
-  auto waitingTask = swift_task_getCurrent();
   auto task = alet->getTask();
   swift_task_future_wait(result, rawContext, task, T);
 }

stdlib/public/Concurrency/Task.cpp

@@ -193,12 +193,7 @@ AsyncTask::~AsyncTask() {
     futureFragment()->destroy();
   }
 
-  // Release any objects potentially held as task local values.
-  //
-  // This must be called last when destroying a task - to keep stack discipline of the allocator.
-  // because it may have created some task-locals immediately upon creation,
-  // e.g. if the task is spawned with async{} and inherited some task-locals.
-  Local.destroy(this);
+  Private.destroy();
 }
 
 SWIFT_CC(swift)
@@ -287,13 +282,7 @@ static void completeTaskImpl(AsyncTask *task,
       reinterpret_cast<char *>(context) - sizeof(AsyncContextPrefix));
   asyncContextPrefix->errorResult = error;
 
-  // Destroy and deallocate any remaining task local items.
-  // We need to do this before we destroy the task local deallocator.
-  task->Local.destroy(task);
-
-  // Tear down the task-local allocator immediately;
-  // there's no need to wait for the object to be destroyed.
-  _swift_task_alloc_destroy(task);
+  task->Private.complete(task);
 
 #if SWIFT_TASK_PRINTF_DEBUG
   fprintf(stderr, "[%p] task %p completed\n", pthread_self(), task);
@@ -525,6 +514,19 @@ static AsyncTaskAndContext swift_task_create_group_future_commonImpl(
   fprintf(stderr, "[%p] creating task %p with parent %p\n", pthread_self(), task, parent);
 #endif
 
+  // Initialize the task-local allocator.
+  if (isAsyncLetTask) {
+    initialContext->ResumeParent = reinterpret_cast<TaskContinuationFunction *>(
+                                                   &completeTask);
+    assert(parent);
+    void *initialSlab = (char*)allocation + amountToAllocate;
+    task->Private.initializeWithSlab(task, initialSlab, initialSlabSize);
+  } else {
+    initialContext->ResumeParent = reinterpret_cast<TaskContinuationFunction *>(
+        closureContext ? &completeTaskWithClosure : &completeTaskAndRelease);
+    task->Private.initialize(task);
+  }
+
   // Perform additional linking between parent and child task.
   if (parent) {
     // If the parent was already cancelled, we carry this flag forward to the child.
@@ -534,6 +536,9 @@ static AsyncTaskAndContext swift_task_create_group_future_commonImpl(
     // so they may have been spawned in any case still.
     if (swift_task_isCancelled(parent))
       swift_task_cancel(task);
+
+    // Initialize task locals with a link to the parent task.
+    task->_private().Local.initializeLinkParent(task, parent);
   }
 
   // Configure the initial context.
@@ -546,26 +551,6 @@ static AsyncTaskAndContext swift_task_create_group_future_commonImpl(
   initialContext->Flags = AsyncContextKind::Ordinary;
   initialContext->Flags.setShouldNotDeallocateInCallee(true);
 
-  // Initialize the task-local allocator.
-  if (isAsyncLetTask) {
-    initialContext->ResumeParent = reinterpret_cast<TaskContinuationFunction *>(
-                                                   &completeTask);
-    assert(parent);
-    void *initialSlab = (char*)allocation + amountToAllocate;
-    _swift_task_alloc_initialize_with_slab(task, initialSlab, initialSlabSize);
-  } else {
-    initialContext->ResumeParent = reinterpret_cast<TaskContinuationFunction *>(
-        closureContext ? &completeTaskWithClosure : &completeTaskAndRelease);
-    _swift_task_alloc_initialize(task);
-  }
-
-  // TODO: if the allocator would be prepared earlier we could do this in some
-  //       other existing if-parent if rather than adding another one here.
-  if (parent) {
-    // Initialize task locals with a link to the parent task.
-    task->Local.initializeLinkParent(task, parent);
-  }
-
   return {task, initialContext};
 }
 

stdlib/public/Concurrency/TaskAlloc.cpp

@@ -20,46 +20,25 @@
 #include "swift/Runtime/Concurrency.h"
 #include "swift/ABI/Task.h"
 #include "TaskPrivate.h"
-#include "../runtime/StackAllocator.h"
+
 #include <stdlib.h>
 
 using namespace swift;
 
 namespace {
 
-/// The size of an allocator slab.
-///
-/// TODO: find the optimal value by experiment.
-static constexpr size_t SlabCapacity = 1024;
-
-using TaskAllocator = StackAllocator<SlabCapacity>;
-
 struct GlobalAllocator {
   TaskAllocator allocator;
   void *spaceForFirstSlab[64];
 
   GlobalAllocator() : allocator(spaceForFirstSlab, sizeof(spaceForFirstSlab)) {}
 };
 
-static_assert(alignof(TaskAllocator) <= alignof(decltype(AsyncTask::AllocatorPrivate)),
-              "task allocator must not be more aligned than "
-              "allocator-private slot");
-
 } // end anonymous namespace
 
-void swift::_swift_task_alloc_initialize(AsyncTask *task) {
-  new (task->AllocatorPrivate) TaskAllocator();
-}
-
-void swift::_swift_task_alloc_initialize_with_slab(AsyncTask *task,
-                                                   void *firstSlabBuffer,
-                                                   size_t bufferCapacity) {
-  new (task->AllocatorPrivate) TaskAllocator(firstSlabBuffer, bufferCapacity);
-}
-
 static TaskAllocator &allocator(AsyncTask *task) {
   if (task)
-    return reinterpret_cast<TaskAllocator &>(task->AllocatorPrivate);
+    return task->Private.get().Allocator;
 
   // FIXME: this fall-back shouldn't be necessary, but it's useful
   // for now, since the current execution tests aren't setting up a task
@@ -68,10 +47,6 @@ static TaskAllocator &allocator(AsyncTask *task) {
   return global.allocator;
 }
 
-void swift::_swift_task_alloc_destroy(AsyncTask *task) {
-  allocator(task).~TaskAllocator();
-}
-
 void *swift::swift_task_alloc(size_t size) {
   return allocator(swift_task_getCurrent()).alloc(size);
 }