[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.

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,32 @@ private:
           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};
+
         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 +1945,7 @@ private:
     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 +1954,27 @@ private:
         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/private/SwiftReflectionTest/SwiftReflectionTest.swift

@@ -209,6 +209,7 @@ public enum InstanceKind: UInt8 {
   case Enum
   case EnumValue
   case AsyncTask
+  case LogString
 }
 
 /// Represents a section in a loaded image in this process.
@@ -642,6 +643,15 @@ public func reflect(asyncTask: UInt) {
   reflect(instanceAddress: asyncTask, kind: .AsyncTask)
 }
 
+/// Log a string to the test's output. Use instead of print, which gets
+/// captured by the parent and read as commands.
+public func reflectionLog(str: String) {
+  str.withCString {
+    let addr = UInt(bitPattern: $0)
+    reflect(instanceAddress: addr, kind: .LogString);
+  }
+}
+
 /// Call this function to indicate to the parent that there are
 /// no more instances to look at.
 public func doneReflecting() {

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/messages.h

@@ -29,5 +29,6 @@ typedef enum InstanceKind {
   Closure,
   Enum,
   EnumValue,
-  AsyncTask
+  AsyncTask,
+  LogString,
 } InstanceKind;

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,67 @@ 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 logString(SwiftReflectionContextRef RC, const PipeMemoryReader *Reader) {
+#pragma clang diagnostic push
+#pragma clang diagnostic ignored "-Wcast-qual"
+  void *Context = (void *)Reader;
+#pragma clang diagnostic pop
+
+  swift_addr_t StringPointer = PipeMemoryReader_receiveInstanceAddress(Context);
+  uint64_t StringLength =
+      PipeMemoryReader_getStringLength(Context, StringPointer);
+
+  void *FreeContext;
+  // Read length+1 bytes to get the NUL terminator too.
+  const void *String = PipeMemoryReader_readBytes(
+      Context, StringPointer, StringLength + 1, &FreeContext);
+
+  printf("%s\n", (const char *)String);
+  PipeMemoryReader_freeBytes(Context, String, FreeContext);
+
+  PipeMemoryReader_sendDoneMessage(Context);
+  return 1;
+}
 
 int doDumpHeapInstance(const char *BinaryFilename, PipeMemoryReader *Reader) {
 #if defined(_WIN32)
@@ -926,6 +1005,11 @@ int doDumpHeapInstance(const char *BinaryFilename, PipeMemoryReader *Reader) {
             return EXIT_SUCCESS;
           break;
         }
+        case LogString: {
+          if (!logString(RC, Reader))
+            return EXIT_SUCCESS;
+          break;
+        }
         case None:
           swift_reflection_destroyReflectionContext(RC);
           printf("Done.\n");

test/Concurrency/Reflection/reflect_task.swift

@@ -28,28 +28,72 @@ func _getCurrentTaskShim() -> UInt
 func add(_ a: UInt, _ b: UInt) async -> UInt {
   if b == 0 {
     reflect(asyncTask: _getCurrentTaskShim())
-    // CHECK: Reflecting an async task.
-    // CHECK: Async task {{0x[0-9a-fA-F]*}}
-
-    // The actual number of chunks we'll get depends on internal implementation
-    // details that we don't want this test to depend on. We'll just make sure
-    // we get at least two, and ignore the details.
-    // CHECK:   Slab pointer {{0x[0-9a-fA-F]*}}
-    // CHECK:     Slab size {{[0-9]{2,}()}}
-    // CHECK:     Chunk at {{0x[0-9a-fA-F]*}} length {{[1-9][0-9]*}} kind {{[0-9]*}}
-    // CHECK:   Slab pointer {{0x[0-9a-fA-F]*}}
-    // CHECK:     Slab size {{[0-9]{2,}()}}
-    // CHECK:     Chunk at {{0x[0-9a-fA-F]*}} length {{[1-9[[0-9]*}} kind {{[0-9]*}}
     return a
   } else {
     return await add(a, b - 1) + 1
   }
 }
 
+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.")
+  }
+}
+
+func testNestedCallsTask() async {
+  reflectionLog(str: "testNestedCallsTask")
+  // CHECK: testNestedCallsTask
+
+  let n = await add(100, 100)
+  reflect(any: n)
+  // CHECK: Reflecting an async task.
+  // CHECK: Async task {{0x[0-9a-fA-F]*}}
+
+  // The actual number of chunks we'll get depends on internal implementation
+  // details that we don't want this test to depend on. We'll just make sure
+  // we get at least two, and ignore the details.
+  // CHECK:   Slab pointer {{0x[0-9a-fA-F]*}}
+  // CHECK:     Slab size {{[0-9]{2,}()}}
+  // CHECK:     Chunk at {{0x[0-9a-fA-F]*}} length {{[1-9][0-9]*}} kind {{[0-9]*}}
+  // CHECK:   Slab pointer {{0x[0-9a-fA-F]*}}
+  // CHECK:     Slab size {{[0-9]{2,}()}}
+  // CHECK:     Chunk at {{0x[0-9a-fA-F]*}} length {{[1-9[[0-9]*}} kind {{[0-9]*}}
+}
+
+func testOneAsyncLet() async {
+  reflectionLog(str: "testOneAsyncLet")
+  // CHECK: testOneAsyncLet
+
+  async let alet = sleepForever()
+  reflect(asyncTask: _getCurrentTaskShim())
+  // CHECK: Async task {{0x[0-9a-fA-F]*}}
+  // CHECK: children = {
+  // CHECK:   Async task {{0x[0-9a-fA-F]*}}
+  // CHECK: }
+}
+
+func testMultipleAsyncLet() async {
+  reflectionLog(str: "testMultipleAsyncLet")
+  // CHECK: testMultipleAsyncLet
+
+  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: }
+}
+
 @main struct Main {
   static func main() async {
-    let n = await add(100, 100)
-    reflect(any: n)
+    await testNestedCallsTask()
+    await testOneAsyncLet()
+    await testMultipleAsyncLet()
 
     doneReflecting()
   }

test/abi/Inputs/macOS/arm64/concurrency/baseline

@@ -914,6 +914,7 @@ _$ss9TaskLocalCyxGs23CustomStringConvertiblesMc
 _$ss9_contains_5whereSbx_Sb7ElementQzYaKXEtYaKSciRzlF
 _$ss9_contains_5whereSbx_Sb7ElementQzYaKXEtYaKSciRzlFTu
 __swift_concurrency_debug_asyncTaskMetadata
+__swift_concurrency_debug_asyncTaskSize
 __swift_concurrency_debug_asyncTaskSlabMetadata
 __swift_concurrency_debug_future_adapter
 __swift_concurrency_debug_jobMetadata

test/abi/Inputs/macOS/arm64/concurrency/baseline-asserts

@@ -1160,6 +1160,7 @@ _$ss9TaskLocalCyxGs23CustomStringConvertiblesMc
 _$ss9_contains_5whereSbx_Sb7ElementQzYaKXEtYaKSciRzlF
 _$ss9_contains_5whereSbx_Sb7ElementQzYaKXEtYaKSciRzlFTu
 __swift_concurrency_debug_asyncTaskMetadata
+__swift_concurrency_debug_asyncTaskSize
 __swift_concurrency_debug_asyncTaskSlabMetadata
 __swift_concurrency_debug_future_adapter
 __swift_concurrency_debug_jobMetadata

test/abi/Inputs/macOS/x86_64/concurrency/baseline

@@ -914,6 +914,7 @@ _$ss9TaskLocalCyxGs23CustomStringConvertiblesMc
 _$ss9_contains_5whereSbx_Sb7ElementQzYaKXEtYaKSciRzlF
 _$ss9_contains_5whereSbx_Sb7ElementQzYaKXEtYaKSciRzlFTu
 __swift_concurrency_debug_asyncTaskMetadata
+__swift_concurrency_debug_asyncTaskSize
 __swift_concurrency_debug_asyncTaskSlabMetadata
 __swift_concurrency_debug_future_adapter
 __swift_concurrency_debug_jobMetadata

test/abi/Inputs/macOS/x86_64/concurrency/baseline-asserts

@@ -1160,6 +1160,7 @@ _$ss9TaskLocalCyxGs23CustomStringConvertiblesMc
 _$ss9_contains_5whereSbx_Sb7ElementQzYaKXEtYaKSciRzlF
 _$ss9_contains_5whereSbx_Sb7ElementQzYaKXEtYaKSciRzlFTu
 __swift_concurrency_debug_asyncTaskMetadata
+__swift_concurrency_debug_asyncTaskSize
 __swift_concurrency_debug_asyncTaskSlabMetadata
 __swift_concurrency_debug_future_adapter
 __swift_concurrency_debug_jobMetadata