Don't destroy the task private storage until task is destroyed since

other threads/tasks which have a reference to the task might need to
access parts of it even if the task itself is completed.

Radar-Id: rdar://problem/88093007

Author: rokhinip

include/swift/ABI/Task.h

@@ -325,7 +325,7 @@ class AsyncTask : public Job {
 
   /// Flag that this task is now completed. This normally does not do anything
   /// but can be used to locally insert logging.
-  void flagAsCompleted();
+  void flagAsDestroyed();
 
   /// Check whether this task has been cancelled.
   /// Checking this is, of course, inherently race-prone on its own.

stdlib/public/Concurrency/Task.cpp

@@ -330,7 +330,7 @@ static void destroyJob(SWIFT_CONTEXT HeapObject *obj) {
 }
 
 AsyncTask::~AsyncTask() {
-  flagAsCompleted();
+  flagAsDestroyed();
 
   // For a future, destroy the result.
   if (isFuture()) {
@@ -372,7 +372,7 @@ static void destroyTask(SWIFT_CONTEXT HeapObject *obj) {
   // the task-local allocator.  There's actually nothing else to clean up
   // here.
 
-  SWIFT_TASK_DEBUG_LOG("destroy task %p", task);
+  SWIFT_TASK_DEBUG_LOG("Destroyed task %p", task);
   free(task);
 }
 

stdlib/public/Concurrency/TaskPrivate.h

@@ -287,11 +287,9 @@ class alignas(2 * sizeof(void*)) ActiveTaskStatus {
     /// or is completed.
     IsEnqueued = 0x1000,
 
-#ifndef NDEBUG
     /// Task has been completed.  This is purely used to enable an assertion
     /// that the task is completed when we destroy it.
     IsComplete = 0x2000,
-#endif
   };
 
   // Note: this structure is mirrored by ActiveTaskStatusWithEscalation and
@@ -409,7 +407,6 @@ class alignas(2 * sizeof(void*)) ActiveTaskStatus {
 #endif
   }
 
-#ifndef NDEBUG
   bool isComplete() const {
     return Flags & IsComplete;
   }
@@ -421,7 +418,6 @@ class alignas(2 * sizeof(void*)) ActiveTaskStatus {
     return ActiveTaskStatus(Record, Flags | IsComplete);
 #endif
   }
-#endif
 
   /// Is there a lock on the linked list of status records?
   bool isStatusRecordLocked() const { return Flags & IsStatusRecordLocked; }
@@ -596,9 +592,7 @@ struct AsyncTask::PrivateStorage {
       auto newStatus = oldStatus.withRunning(false);
       newStatus = newStatus.withoutStoredPriorityEscalation();
       newStatus = newStatus.withoutEnqueued();
-#ifndef NDEBUG
       newStatus = newStatus.withComplete();
-#endif
 
       // This can fail since the task can still get concurrently cancelled or
       // escalated.
@@ -615,10 +609,22 @@ struct AsyncTask::PrivateStorage {
       }
     }
 
-    // Destroy and deallocate any remaining task local items.
-    // We need to do this before we destroy the task local deallocator.
+    // Destroy and deallocate any remaining task local items since the task is
+    // completed. We need to do this before we destroy the task local
+    // deallocator.
     Local.destroy(task);
 
+    // Don't destroy the task private storage as a whole since others which have
+    // a reference to the task might still need to access the ActiveTaskStatus.
+    // It is destroyed when task is destroyed
+  }
+
+  // Destroy the opaque storage of the task
+  void destroy() {
+#if NDEBUG
+    auto oldStatus = task->_private()._status().load(std::memory_order_relaxed);
+    assert(oldStatus.isComplete());
+#endif
     this->~PrivateStorage();
   }
 
@@ -653,11 +659,13 @@ inline void AsyncTask::OpaquePrivateStorage::initializeWithSlab(
     JobPriority basePri, void *slab, size_t slabCapacity) {
   ::new (this) PrivateStorage(basePri, slab, slabCapacity);
 }
+
 inline void AsyncTask::OpaquePrivateStorage::complete(AsyncTask *task) {
   get().complete(task);
 }
+
 inline void AsyncTask::OpaquePrivateStorage::destroy() {
-  // nothing else to do
+  get().destroy();
 }
 
 inline AsyncTask::PrivateStorage &AsyncTask::_private() {
@@ -822,10 +830,10 @@ inline void AsyncTask::flagAsSuspended() {
 }
 
 // READ ME: This is not a dead function! Do not remove it! This is a function
-// that can be used when debugging locally to instrument when a task actually is
-// dealloced.
-inline void AsyncTask::flagAsCompleted() {
-  SWIFT_TASK_DEBUG_LOG("task completed %p", this);
+// that can be used when debugging locally to instrument when a task
+// actually is dealloced.
+inline void AsyncTask::flagAsDestroyed() {
+  SWIFT_TASK_DEBUG_LOG("task destroyed %p", this);
 }
 
 inline void AsyncTask::localValuePush(const HeapObject *key,