[RemoteMirror] Fix AsyncTask child iteration.

Iterating child tasks depends on knowing the size of AsyncTask, and changing the size of the task broke it. Instead of relying on mirroring the full structure in our out-of-process definitions, add a debug variable to libswift_Concurrency that contains the size of AsyncTask.

While we're there, add some more validation to child task enumeration. Check each child task's metadata pointer to make sure that it actually points to the AsyncTask metadata, and have the inner loop also increment and check ChildTaskLoopCount to stop runaway iteration in that loop.

Author: mikeash

include/swift/RemoteInspection/ReflectionContext.h

@@ -134,11 +134,13 @@ class ReflectionContext
   std::vector<std::tuple<RemoteAddress, RemoteAddress>> dataRanges;
 
   bool setupTargetPointers = false;
+  typename super::StoredPointer target_asyncTaskMetadata = 0;
   typename super::StoredPointer target_non_future_adapter = 0;
   typename super::StoredPointer target_future_adapter = 0;
   typename super::StoredPointer target_task_wait_throwing_resume_adapter = 0;
   typename super::StoredPointer target_task_future_wait_resume_adapter = 0;
   bool supportsPriorityEscalation = false;
+  typename super::StoredSize asyncTaskSize = 0;
 
 public:
   using super::getBuilder;
@@ -1816,16 +1818,39 @@ class ReflectionContext
           ChildTask = RecordObj->FirstChild;
       }
 
-      while (ChildTask) {
+      while (ChildTask && ChildTaskLoopCount++ < ChildTaskLimit) {
+        // Read the child task.
+        auto ChildTaskObj = readObj<AsyncTaskType>(ChildTask);
+        if (!ChildTaskObj)
+          return {std::string("found unreadable child task pointer"), Info};
+
+        // Make sure the task's metadata pointer is what we expect.
+        if (stripSignedPointer(ChildTaskObj->HeapObject.Metadata) !=
+            target_asyncTaskMetadata)
+          return {
+              std::string("found child task pointer with incorrect metadata"),
+              Info};
+
         Info.ChildTasks.push_back(ChildTask);
 
-        StoredPointer ChildFragmentAddr = ChildTask + sizeof(*AsyncTaskObj);
-        auto ChildFragmentObj =
-            readObj<ChildFragment<Runtime>>(ChildFragmentAddr);
-        if (ChildFragmentObj)
-          ChildTask = ChildFragmentObj->NextChild;
-        else
+        swift::JobFlags ChildJobFlags(AsyncTaskObj->Flags);
+        if (ChildJobFlags.task_isChildTask()) {
+          if (asyncTaskSize == 0)
+            return {std::string("target async task size unknown, unable to "
+                                "iterate child tasks"),
+                    Info};
+
+          StoredPointer ChildFragmentAddr = ChildTask + asyncTaskSize;
+          auto ChildFragmentObj =
+              readObj<ChildFragment<Runtime>>(ChildFragmentAddr);
+          if (ChildFragmentObj)
+            ChildTask = ChildFragmentObj->NextChild;
+          else
+            ChildTask = 0;
+        } else {
+          // No child fragment, so we're done iterating.
           ChildTask = 0;
+        }
       }
 
       RecordPtr = RecordObj->Parent;
@@ -1927,7 +1952,7 @@ class ReflectionContext
     if (setupTargetPointers)
       return;
 
-    auto getFunc = [&](const std::string &name) -> StoredPointer {
+    auto getPointer = [&](const std::string &name) -> StoredPointer {
       auto Symbol = getReader().getSymbolAddress(name);
       if (!Symbol)
         return 0;
@@ -1936,19 +1961,27 @@ class ReflectionContext
         return 0;
       return Pointer->getResolvedAddress().getAddressData();
     };
+    target_asyncTaskMetadata =
+        getPointer("_swift_concurrency_debug_asyncTaskMetadata");
     target_non_future_adapter =
-        getFunc("_swift_concurrency_debug_non_future_adapter");
-    target_future_adapter = getFunc("_swift_concurrency_debug_future_adapter");
-    target_task_wait_throwing_resume_adapter =
-        getFunc("_swift_concurrency_debug_task_wait_throwing_resume_adapter");
+        getPointer("_swift_concurrency_debug_non_future_adapter");
+    target_future_adapter =
+        getPointer("_swift_concurrency_debug_future_adapter");
+    target_task_wait_throwing_resume_adapter = getPointer(
+        "_swift_concurrency_debug_task_wait_throwing_resume_adapter");
     target_task_future_wait_resume_adapter =
-        getFunc("_swift_concurrency_debug_task_future_wait_resume_adapter");
+        getPointer("_swift_concurrency_debug_task_future_wait_resume_adapter");
     auto supportsPriorityEscalationAddr = getReader().getSymbolAddress(
         "_swift_concurrency_debug_supportsPriorityEscalation");
     if (supportsPriorityEscalationAddr) {
       getReader().readInteger(supportsPriorityEscalationAddr,
                               &supportsPriorityEscalation);
     }
+    auto asyncTaskSizeAddr =
+        getReader().getSymbolAddress("_swift_concurrency_debug_asyncTaskSize");
+    if (asyncTaskSizeAddr) {
+      getReader().readInteger(asyncTaskSizeAddr, &asyncTaskSize);
+    }
 
     setupTargetPointers = true;
   }

include/swift/RemoteInspection/RuntimeInternals.h

@@ -64,7 +64,7 @@ template <typename Runtime> struct ConformanceCacheEntry {
 
 template <typename Runtime>
 struct HeapObject {
-  typename Runtime::StoredPointer Metadata;
+  typename Runtime::StoredSignedPointer Metadata;
   typename Runtime::StoredSize RefCounts;
 };
 
@@ -131,10 +131,7 @@ struct AsyncTask: Job<Runtime> {
   // On 64-bit, there's a Reserved64 after ResumeContext.
   typename Runtime::StoredPointer ResumeContextAndReserved[
     sizeof(typename Runtime::StoredPointer) == 8 ? 2 : 1];
-  union {
-    AsyncTaskPrivateStorage<Runtime, ActiveTaskStatus> PrivateStorage;
-    typename Runtime::StoredPointer PrivateStorageRaw[14];
-  };
+  AsyncTaskPrivateStorage<Runtime, ActiveTaskStatus> PrivateStorage;
 };
 
 template <typename Runtime>

stdlib/public/Concurrency/Debug.h

@@ -34,6 +34,10 @@ const void *const _swift_concurrency_debug_jobMetadata;
 SWIFT_EXPORT_FROM(swift_Concurrency)
 const void *const _swift_concurrency_debug_asyncTaskMetadata;
 
+/// The size of an AsyncTask, in bytes.
+SWIFT_EXPORT_FROM(swift_Concurrency)
+const size_t _swift_concurrency_debug_asyncTaskSize;
+
 /// A fake metadata pointer placed at the start of async task slab allocations.
 SWIFT_EXPORT_FROM(swift_Concurrency)
 const void *const _swift_concurrency_debug_asyncTaskSlabMetadata;

stdlib/public/Concurrency/Task.cpp

@@ -441,6 +441,8 @@ const void *const swift::_swift_concurrency_debug_jobMetadata =
 const void *const swift::_swift_concurrency_debug_asyncTaskMetadata =
     static_cast<Metadata *>(&taskHeapMetadata);
 
+const size_t swift::_swift_concurrency_debug_asyncTaskSize = sizeof(AsyncTask);
+
 const HeapMetadata *swift::jobHeapMetadataPtr =
     static_cast<HeapMetadata *>(&jobHeapMetadata);
 const HeapMetadata *swift::taskHeapMetadataPtr =

stdlib/tools/swift-reflection-test/swift-reflection-test.c

@@ -23,6 +23,7 @@
 #include <assert.h>
 #include <errno.h>
 #include <inttypes.h>
+#include <stdarg.h>
 #include <stdint.h>
 #include <stdio.h>
 #include <stdlib.h>
@@ -84,6 +85,20 @@ static void errnoAndExit(const char *message) {
 #define DEBUG_LOG(fmt, ...) (void)0
 #endif
 
+#ifdef __clang__
+__attribute((__format__(__printf__, 2, 3)))
+#endif
+static void
+indented_printf(unsigned indentLevel, const char *fmt, ...) {
+  for (unsigned i = 0; i < indentLevel; i++)
+    fputs("  ", stdout);
+
+  va_list args;
+  va_start(args, fmt);
+  vprintf(fmt, args);
+  va_end(args);
+}
+
 static const size_t ReadEnd = 0;
 static const size_t WriteEnd = 1;
 
@@ -774,10 +789,40 @@ int reflectEnumValue(SwiftReflectionContextRef RC,
 
 }
 
-int reflectAsyncTask(SwiftReflectionContextRef RC,
-                     const PipeMemoryReader *Reader) {
-  uintptr_t AsyncTaskInstance = PipeMemoryReader_receiveInstanceAddress(Reader);
-  printf("Async task %#" PRIx64 "\n", (uint64_t)AsyncTaskInstance);
+static int reflectAsyncTaskInstance(SwiftReflectionContextRef RC,
+                                    uintptr_t AsyncTaskInstance,
+                                    const PipeMemoryReader *Reader,
+                                    unsigned indentLevel) {
+  indented_printf(indentLevel, "Async task %#" PRIx64 "\n",
+                  (uint64_t)AsyncTaskInstance);
+
+  swift_async_task_info_t TaskInfo =
+      swift_reflection_asyncTaskInfo(RC, AsyncTaskInstance);
+  if (TaskInfo.Error) {
+    printf("swift_reflection_asyncTaskInfo failed: %s\n", TaskInfo.Error);
+  } else {
+    indented_printf(indentLevel, "id %" PRIu64 "\n", TaskInfo.Id);
+    indented_printf(indentLevel, "enqueuePriority %u\n",
+                    TaskInfo.EnqueuePriority);
+    if (TaskInfo.ChildTaskCount > 0) {
+      indented_printf(indentLevel, "children = {\n");
+
+      // The memory for ChildTasks is only valid until the next Remote Mirror
+      // call, so we need to copy it.
+      swift_reflection_ptr_t *ChildTasks =
+          calloc(TaskInfo.ChildTaskCount, sizeof(swift_reflection_ptr_t));
+      memcpy(ChildTasks, TaskInfo.ChildTasks,
+             TaskInfo.ChildTaskCount * sizeof(swift_reflection_ptr_t));
+
+      for (unsigned i = 0; i < TaskInfo.ChildTaskCount; i++)
+        reflectAsyncTaskInstance(RC, ChildTasks[i], Reader, indentLevel + 1);
+
+      free(ChildTasks);
+      indented_printf(indentLevel, "}\n");
+    } else {
+      indented_printf(indentLevel, "children = {}\n");
+    }
+  }
 
   swift_async_task_slab_return_t SlabPtrResult =
       swift_reflection_asyncTaskSlabPointer(RC, AsyncTaskInstance);
@@ -787,33 +832,45 @@ int reflectAsyncTask(SwiftReflectionContextRef RC,
   } else {
     swift_reflection_ptr_t SlabPtr = SlabPtrResult.SlabPtr;
     while (SlabPtr) {
-      printf("  Slab pointer %#" PRIx64 "\n", (uint64_t)SlabPtr);
+      indented_printf(indentLevel, "  Slab pointer %#" PRIx64 "\n",
+                      (uint64_t)SlabPtr);
       swift_async_task_slab_allocations_return_t AllocationsResult =
           swift_reflection_asyncTaskSlabAllocations(RC, SlabPtr);
       if (AllocationsResult.Error) {
-        printf("swift_reflection_asyncTaskSlabAllocations failed: %s\n",
-               AllocationsResult.Error);
+        indented_printf(
+            indentLevel,
+            "swift_reflection_asyncTaskSlabAllocations failed: %s\n",
+            AllocationsResult.Error);
         SlabPtr = 0;
       } else {
-        printf("    Slab size %" PRIu64 "\n",
-               (uint64_t)AllocationsResult.SlabSize);
+        indented_printf(indentLevel, "    Slab size %" PRIu64 "\n",
+                        (uint64_t)AllocationsResult.SlabSize);
         for (unsigned i = 0; i < AllocationsResult.ChunkCount; i++) {
           swift_async_task_allocation_chunk_t Chunk =
               AllocationsResult.Chunks[i];
-          printf("    Chunk at %#" PRIx64 " length %u kind %u\n",
-                 (uint64_t)Chunk.Start, Chunk.Length, Chunk.Kind);
+          indented_printf(indentLevel,
+                          "    Chunk at %#" PRIx64 " length %u kind %u\n",
+                          (uint64_t)Chunk.Start, Chunk.Length, Chunk.Kind);
         }
         SlabPtr = AllocationsResult.NextSlab;
       }
     }
   }
 
-  printf("\n\n");
-  PipeMemoryReader_sendDoneMessage(Reader);
+  if (indentLevel == 0) {
+    printf("\n\n");
+  }
   fflush(stdout);
   return 1;
 }
 
+int reflectAsyncTask(SwiftReflectionContextRef RC,
+                     const PipeMemoryReader *Reader) {
+  uintptr_t AsyncTaskInstance = PipeMemoryReader_receiveInstanceAddress(Reader);
+  int result = reflectAsyncTaskInstance(RC, AsyncTaskInstance, Reader, 0);
+  PipeMemoryReader_sendDoneMessage(Reader);
+  return result;
+}
 
 int doDumpHeapInstance(const char *BinaryFilename, PipeMemoryReader *Reader) {
 #if defined(_WIN32)

test/Concurrency/Reflection/reflect_task.swift

@@ -46,11 +46,28 @@ func add(_ a: UInt, _ b: UInt) async -> UInt {
   }
 }
 
+func sleepForever() async -> Int {
+  if #available(SwiftStdlib 5.7, *) {
+    try? await Task.sleep(for: .seconds(1_000_000_000))
+    return 42
+  } else {
+    fatalError("This test shouldn't be running against old stdlibs.")
+  }
+}
+
 @main struct Main {
   static func main() async {
     let n = await add(100, 100)
     reflect(any: n)
 
+    async let alet1 = sleepForever()
+    async let alet2 = sleepForever()
+    reflect(asyncTask: _getCurrentTaskShim())
+    // CHECK: Async task {{0x[0-9a-fA-F]*}}
+    // CHECK: children = {
+    // CHECK:   Async task {{0x[0-9a-fA-F]*}}
+    // CHECK:   Async task {{0x[0-9a-fA-F]*}}
+    // CHECK: }
     doneReflecting()
   }
 }