Revert "Future-proof the layout of AsyncTask."

This reverts commit 7eb1052.

Author: ktoso

include/swift/ABI/MetadataValues.h

@@ -47,9 +47,6 @@ 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 SWIFT_POINTER_IS_8_BYTES
+#if defined(__LP64__) || defined(_WIN64)
 static_assert(sizeof(Job) == 6 * sizeof(void*),
               "Job size is wrong");
 #else
@@ -136,6 +136,46 @@ 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:
@@ -172,16 +212,6 @@ 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
@@ -192,57 +222,36 @@ class AsyncTask : public Job {
   /// prevent it from being corrupted in flight.
   AsyncContext * __ptrauth_swift_task_resume_context ResumeContext;
 
-#if SWIFT_POINTER_IS_8_BYTES
-  void *Reserved64;
-#endif
+  /// The currently-active information about cancellation.
+  std::atomic<ActiveTaskStatus> Status;
 
-  struct PrivateStorage;
+  /// Reserved for the use of the task-local stack allocator.
+  void *AllocatorPrivate[4];
 
-  /// Private storage for the use of the runtime.
-  struct alignas(2 * alignof(void*)) OpaquePrivateStorage {
-    void *Storage[8];
+  /// Task local values storage container.
+  TaskLocal::Storage Local;
 
-    /// 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) {
+      ResumeContext(initialContext),
+      Status(ActiveTaskStatus()),
+      Local(TaskLocal::Storage()) {
     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) {
+      ResumeContext(initialContext),
+      Status(ActiveTaskStatus()),
+      Local(TaskLocal::Storage()) {
     assert(flags.isAsyncTask());
     Id = getNextTaskId();
   }
@@ -260,18 +269,24 @@ 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;
+  bool isCancelled() const {
+    return Status.load(std::memory_order_relaxed).isCancelled();
+  }
 
   // ==== Task Local Values ----------------------------------------------------
 
   void localValuePush(const HeapObject *key,
-                      /* +1 */ OpaqueValue *value,
-                      const Metadata *valueType);
+                      /* +1 */ OpaqueValue *value, const Metadata *valueType) {
+    Local.pushValue(this, key, value, valueType);
+  }
 
-  OpaqueValue *localValueGet(const HeapObject *key);
+  OpaqueValue* localValueGet(const HeapObject *key) {
+    return Local.getValue(this, key);
+  }
 
-  /// Returns true if storage has still more bindings.
-  void localValuePop();
+  void localValuePop() {
+    Local.popValue(this);
+  }
 
   // ==== Child Fragment -------------------------------------------------------
 
@@ -513,7 +528,7 @@ class AsyncTask : public Job {
 };
 
 // The compiler will eventually assume these.
-static_assert(sizeof(AsyncTask) == NumWords_AsyncTask * sizeof(void*),
+static_assert(sizeof(AsyncTask) == 14 * sizeof(void*),
               "AsyncTask size is wrong");
 static_assert(alignof(AsyncTask) == 2 * alignof(void*),
               "AsyncTask alignment is wrong");

include/swift/ABI/TaskStatus.h

@@ -82,6 +82,11 @@ 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,13 +128,4 @@
 #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,6 +86,10 @@ static AsyncLetImpl *asImpl(const AsyncLet *alet) {
       const_cast<AsyncLet*>(alet));
 }
 
+static AsyncLet *asAbstract(AsyncLetImpl *alet) {
+  return reinterpret_cast<AsyncLet*>(alet);
+}
+
 // =============================================================================
 // ==== start ------------------------------------------------------------------
 
@@ -130,6 +134,7 @@ 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,7 +193,8 @@ AsyncTask::~AsyncTask() {
     futureFragment()->destroy();
   }
 
-  Private.destroy();
+  // Release any objects potentially held as task local values.
+  Local.destroy(this);
 }
 
 SWIFT_CC(swift)
@@ -282,7 +283,13 @@ static void completeTaskImpl(AsyncTask *task,
       reinterpret_cast<char *>(context) - sizeof(AsyncContextPrefix));
   asyncContextPrefix->errorResult = error;
 
-  task->Private.complete(task);
+  // 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);
 
 #if SWIFT_TASK_PRINTF_DEBUG
   fprintf(stderr, "[%p] task %p completed\n", pthread_self(), task);
@@ -514,19 +521,6 @@ 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.
@@ -536,9 +530,6 @@ 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.
@@ -551,6 +542,26 @@ 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,25 +20,46 @@
 #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 task->Private.get().Allocator;
+    return reinterpret_cast<TaskAllocator &>(task->AllocatorPrivate);
 
   // 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
@@ -47,6 +68,10 @@ 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);
 }

stdlib/public/Concurrency/TaskLocal.cpp

@@ -72,7 +72,7 @@ TaskLocal::Item::createParentLink(AsyncTask *task, AsyncTask *parent) {
   Item *item = new(allocation) Item();
 
   // FIXME: parent pointer must to be the parent STORAGE not just the next item.
-  auto parentHead = parent->_private().Local.head;
+  auto parentHead = parent->Local.head;
   if (parentHead) {
     if (parentHead->isEmpty()) {
       switch (parentHead->getNextLinkType()) {
@@ -189,7 +189,7 @@ TaskLocal::Item::createLink(AsyncTask *task,
   void *allocation = _swift_task_alloc_specific(task, amountToAllocate);
   Item *item = new(allocation) Item(key, valueType);
 
-  auto next = task->_private().Local.head;
+  auto next = task->Local.head;
   item->next = reinterpret_cast<uintptr_t>(next) |
       static_cast<uintptr_t>(NextLinkType::IsNext);