[RemoteMirror] Add a call for iterating over the allocations belonging to an AsyncTask.

rdar://72907056

include/swift/Reflection/ReflectionContext.h

@@ -114,6 +114,22 @@ public:
   using typename super::StoredSignedPointer;
   using typename super::StoredSize;
 
+  struct AsyncTaskAllocationChunk {
+    enum class ChunkKind {
+      Unknown,
+      NonPointer,
+      RawPointer,
+      StrongReference,
+      UnownedReference,
+      WeakReference,
+      UnmanagedReference
+    };
+
+    StoredPointer Start;
+    unsigned Length;
+    ChunkKind Kind;
+  };
+
   explicit ReflectionContext(std::shared_ptr<MemoryReader> reader)
     : super(std::move(reader), *this)
   {}
@@ -1165,6 +1181,48 @@ public:
     return llvm::None;
   }
 
+  llvm::Optional<std::string> iterateAsyncTaskAllocations(
+      StoredPointer AsyncTaskPtr,
+      std::function<void(StoredPointer, unsigned, AsyncTaskAllocationChunk[])>
+          Call) {
+    using AsyncTask = AsyncTask<Runtime>;
+    using StackAllocator = StackAllocator<Runtime>;
+
+    auto AsyncTaskBytes =
+        getReader().readBytes(RemoteAddress(AsyncTaskPtr), sizeof(AsyncTask));
+    auto *AsyncTaskObj =
+        reinterpret_cast<const AsyncTask *>(AsyncTaskBytes.get());
+    if (!AsyncTaskObj)
+      return std::string("failure reading async task");
+
+    auto *Allocator = reinterpret_cast<const StackAllocator *>(
+        &AsyncTaskObj->AllocatorPrivate);
+    StoredPointer SlabPtr = Allocator->FirstSlab;
+    while (SlabPtr) {
+      auto SlabBytes = getReader().readBytes(
+          RemoteAddress(SlabPtr), sizeof(typename StackAllocator::Slab));
+      auto Slab = reinterpret_cast<const typename StackAllocator::Slab *>(
+          SlabBytes.get());
+      if (!Slab)
+        return std::string("failure reading slab");
+
+      // For now, we won't try to walk the allocations in the slab, we'll just
+      // provide the whole thing as one big chunk.
+      size_t HeaderSize =
+          llvm::alignTo(sizeof(*Slab), llvm::Align(alignof(std::max_align_t)));
+      AsyncTaskAllocationChunk Chunk;
+
+      Chunk.Start = SlabPtr + HeaderSize;
+      Chunk.Length = Slab->CurrentOffset;
+      Chunk.Kind = AsyncTaskAllocationChunk::ChunkKind::Unknown;
+      Call(SlabPtr, 1, &Chunk);
+
+      SlabPtr = Slab->Next;
+    }
+
+    return llvm::None;
+  }
+
 private:
   const TypeInfo *
   getClosureContextInfo(StoredPointer Context, const ClosureContextInfo &Info,

include/swift/Reflection/RuntimeInternals.h

@@ -60,6 +60,40 @@ template <typename Runtime> struct ConformanceCacheEntry {
   typename Runtime::StoredPointer Witness;
 };
 
+template <typename Runtime>
+struct HeapObject {
+  typename Runtime::StoredPointer Metadata;
+  typename Runtime::StoredSize RefCounts;
+};
+
+template <typename Runtime>
+struct Job {
+  typename Runtime::StoredPointer Opaque[4];
+};
+
+template <typename Runtime>
+struct StackAllocator {
+  typename Runtime::StoredPointer LastAllocation;
+  typename Runtime::StoredPointer FirstSlab;
+  int32_t NumAllocatedSlabs;
+  bool FirstSlabIsPreallocated;
+
+  struct Slab {
+    typename Runtime::StoredPointer Next;
+    uint32_t Capacity;
+    uint32_t CurrentOffset;
+  };
+};
+
+template <typename Runtime>
+struct AsyncTask {
+  HeapObject<Runtime> HeapObject;
+  Job<Runtime> Job;
+  typename Runtime::StoredPointer ResumeContext;
+  typename Runtime::StoredSize Status;
+  typename Runtime::StoredPointer AllocatorPrivate[4];
+};
+
 } // end namespace reflection
 } // end namespace swift
 

include/swift/SwiftRemoteMirror/SwiftRemoteMirror.h

@@ -372,6 +372,34 @@ const char *swift_reflection_iterateMetadataAllocationBacktraces(
     SwiftReflectionContextRef ContextRef,
     swift_metadataAllocationBacktraceIterator Call, void *ContextPtr);
 
+/// Allocation iterator passed to swift_reflection_iterateAsyncTaskAllocations
+typedef void (*swift_asyncTaskAllocationIterator)(
+    swift_reflection_ptr_t AllocationPtr, unsigned Count,
+    swift_async_task_allocation_chunk_t Chunks[], void *ContextPtr);
+
+/// Iterate over the allocations associated with the given async task object.
+/// This object must have an isa value equal to
+/// _swift_concurrency_debug_asyncTaskMetadata.
+///
+/// Calls the passed in Call function for each allocation associated with the
+/// async task object. The function is passed the allocation pointer and an
+/// array of chunks. Each chunk consists of a start, length, and kind for that
+/// chunk of the allocated memory. Any regions of the allocation that are not
+/// covered by a chunk are unallocated or garbage. The chunk array is valid only
+/// for the duration of the call.
+///
+/// An async task may have more than one allocation associated with it, so the
+/// function may be called more than once. It may also have no allocations, in
+/// which case the function is not called.
+///
+/// Returns NULL on success. On error, returns a pointer to a C string
+/// describing the error. This pointer remains valid until the next
+/// swift_reflection call on the given context.
+SWIFT_REMOTE_MIRROR_LINKAGE
+const char *swift_reflection_iterateAsyncTaskAllocations(
+    SwiftReflectionContextRef ContextRef, swift_reflection_ptr_t AsyncTaskPtr,
+    swift_asyncTaskAllocationIterator Call, void *ContextPtr);
+
 #ifdef __cplusplus
 } // extern "C"
 #endif

include/swift/SwiftRemoteMirror/SwiftRemoteMirrorTypes.h

@@ -189,6 +189,12 @@ typedef struct swift_metadata_cache_node {
   swift_reflection_ptr_t Right;
 } swift_metadata_cache_node_t;
 
+typedef struct swift_async_task_allocation_chunk {
+  swift_reflection_ptr_t Start;
+  unsigned Length;
+  swift_layout_kind_t Kind;
+} swift_async_task_allocation_chunk_t;
+
 /// An opaque pointer to a context which maintains state and
 /// caching of reflection structure for heap instances.
 typedef struct SwiftReflectionContext *SwiftReflectionContextRef;

stdlib/private/SwiftReflectionTest/SwiftReflectionTest.swift

@@ -179,6 +179,7 @@ public enum InstanceKind: UInt8 {
   case Closure
   case Enum
   case EnumValue
+  case AsyncTask
 }
 
 /// Represents a section in a loaded image in this process.
@@ -579,6 +580,12 @@ public func reflect(function: @escaping (Int, String, AnyObject?) -> Void) {
   fn.deallocate()
 }
 
+
+/// Reflect an AsyncTask.
+public func reflect(asyncTask: UInt) {
+  reflect(instanceAddress: asyncTask, kind: .AsyncTask)
+}
+
 /// Call this function to indicate to the parent that there are
 /// no more instances to look at.
 public func doneReflecting() {

stdlib/public/SwiftRemoteMirror/SwiftRemoteMirror.cpp

@@ -463,6 +463,30 @@ static swift_childinfo_t convertChild(const TypeInfo *TI, unsigned Index) {
   };
 }
 
+static swift_layout_kind_t convertAllocationChunkKind(
+    NativeReflectionContext::AsyncTaskAllocationChunk::ChunkKind Kind) {
+  switch (Kind) {
+  case NativeReflectionContext::AsyncTaskAllocationChunk::ChunkKind::Unknown:
+    return SWIFT_UNKNOWN;
+  case NativeReflectionContext::AsyncTaskAllocationChunk::ChunkKind::NonPointer:
+    return SWIFT_BUILTIN;
+  case NativeReflectionContext::AsyncTaskAllocationChunk::ChunkKind::RawPointer:
+    return SWIFT_RAW_POINTER;
+  case NativeReflectionContext::AsyncTaskAllocationChunk::ChunkKind::
+      StrongReference:
+    return SWIFT_STRONG_REFERENCE;
+  case NativeReflectionContext::AsyncTaskAllocationChunk::ChunkKind::
+      UnownedReference:
+    return SWIFT_UNOWNED_REFERENCE;
+  case NativeReflectionContext::AsyncTaskAllocationChunk::ChunkKind::
+      WeakReference:
+    return SWIFT_WEAK_REFERENCE;
+  case NativeReflectionContext::AsyncTaskAllocationChunk::ChunkKind::
+      UnmanagedReference:
+    return SWIFT_UNMANAGED_REFERENCE;
+  }
+}
+
 static const char *returnableCString(SwiftReflectionContextRef ContextRef,
                                       llvm::Optional<std::string> String) {
   if (String) {
@@ -718,3 +742,23 @@ const char *swift_reflection_iterateMetadataAllocationBacktraces(
       });
   return returnableCString(ContextRef, Error);
 }
+
+const char *swift_reflection_iterateAsyncTaskAllocations(
+    SwiftReflectionContextRef ContextRef, swift_reflection_ptr_t AsyncTaskPtr,
+    swift_asyncTaskAllocationIterator Call, void *ContextPtr) {
+  auto Context = ContextRef->nativeContext;
+  auto Error = Context->iterateAsyncTaskAllocations(
+      AsyncTaskPtr, [&](auto AllocationPtr, auto Count, auto Chunks) {
+        std::vector<swift_async_task_allocation_chunk_t> ConvertedChunks;
+        ConvertedChunks.reserve(Count);
+        for (unsigned i = 0; i < Count; i++) {
+          swift_async_task_allocation_chunk_t Chunk;
+          Chunk.Start = Chunks[i].Start;
+          Chunk.Length = Chunks[i].Length;
+          Chunk.Kind = convertAllocationChunkKind(Chunks[i].Kind);
+          ConvertedChunks.push_back(Chunk);
+        }
+        Call(AllocationPtr, Count, ConvertedChunks.data(), ContextPtr);
+      });
+  return returnableCString(ContextRef, Error);
+}

stdlib/tools/swift-reflection-test/messages.h

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

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

@@ -24,6 +24,7 @@
 
 #include <assert.h>
 #include <errno.h>
+#include <inttypes.h>
 #include <stdint.h>
 #include <stdio.h>
 #include <stdlib.h>
@@ -70,7 +71,6 @@ static void errnoAndExit(const char *message) {
 }
 
 #if 0
-#include <inttypes.h>
 #define DEBUG_LOG(fmt, ...) fprintf(stderr, "%s: " fmt "\n",\
                                     __func__, __VA_ARGS__)
 #else
@@ -281,7 +281,7 @@ PipeMemoryReader_receiveImages(SwiftReflectionContextRef RC,
   size_t NumReflectionInfos;
   PipeMemoryReader_collectBytesFromPipe(Reader, &NumReflectionInfos,
                                         sizeof(NumReflectionInfos));
-  DEBUG_LOG("Receiving %z images from child", NumReflectionInfos);
+  DEBUG_LOG("Receiving %zu images from child", NumReflectionInfos);
 
   if (NumReflectionInfos == 0)
     return;
@@ -628,6 +628,29 @@ int reflectEnumValue(SwiftReflectionContextRef RC,
 
 }
 
+static void
+asyncTaskIterationCallback(swift_reflection_ptr_t AllocationPtr, unsigned Count,
+                           swift_async_task_allocation_chunk_t Chunks[],
+                           void *ContextPtr) {
+  printf("  Allocation block %#" PRIx64 "\n", (uint64_t)AllocationPtr);
+  for (unsigned i = 0; i < Count; i++)
+    printf("    Chunk at %#" PRIx64 " length %u kind %u\n",
+           (uint64_t)Chunks[i].Start, Chunks[i].Length, Chunks[i].Kind);
+}
+
+int reflectAsyncTask(SwiftReflectionContextRef RC,
+                     const PipeMemoryReader Pipe) {
+  uintptr_t AsyncTaskInstance = PipeMemoryReader_receiveInstanceAddress(&Pipe);
+  printf("Async task %#" PRIx64 "\n", (uint64_t)AsyncTaskInstance);
+  swift_reflection_iterateAsyncTaskAllocations(
+      RC, AsyncTaskInstance, asyncTaskIterationCallback, NULL);
+
+  printf("\n\n");
+  PipeMemoryReader_sendDoneMessage(&Pipe);
+  fflush(stdout);
+  return 1;
+}
+
 
 int doDumpHeapInstance(const char *BinaryFilename) {
   PipeMemoryReader Pipe = createPipeMemoryReader();
@@ -719,6 +742,12 @@ int doDumpHeapInstance(const char *BinaryFilename) {
             return EXIT_SUCCESS;
           break;
         }
+        case AsyncTask: {
+          printf("Reflecting an async task.\n");
+          if (!reflectAsyncTask(RC, Pipe))
+            return EXIT_SUCCESS;
+          break;
+        }
         case None:
           swift_reflection_destroyReflectionContext(RC);
           printf("Done.\n");