[Reflection] Add API for inspecting async task allocation slabs.

include/swift/Reflection/ReflectionContext.h

@@ -131,6 +131,12 @@ class ReflectionContext
     ChunkKind Kind;
   };
 
+  struct AsyncTaskSlabInfo {
+    StoredPointer NextSlab;
+    StoredSize SlabSize;
+    std::vector<AsyncTaskAllocationChunk> Chunks;
+  };
+
   explicit ReflectionContext(std::shared_ptr<MemoryReader> reader)
     : super(std::move(reader), *this)
   {}
@@ -1346,44 +1352,45 @@ class ReflectionContext
     return llvm::None;
   }
 
-  llvm::Optional<std::string> iterateAsyncTaskAllocations(
-      StoredPointer AsyncTaskPtr,
-      std::function<void(StoredPointer, unsigned, AsyncTaskAllocationChunk[])>
-          Call) {
-    using AsyncTask = AsyncTask<Runtime>;
+  std::pair<llvm::Optional<std::string>, AsyncTaskSlabInfo>
+  asyncTaskSlabAllocations(StoredPointer SlabPtr) {
     using StackAllocator = StackAllocator<Runtime>;
+    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;
+
+    // Total slab size is the slab's capacity plus the slab struct itself.
+    StoredPointer SlabSize = Slab->Capacity + sizeof(*Slab);
+
+    return {llvm::None, {Slab->Next, SlabSize, {Chunk}}};
+  }
+
+  std::pair<llvm::Optional<std::string>, StoredPointer>
+  asyncTaskSlabPtr(StoredPointer AsyncTaskPtr) {
+    using AsyncTask = AsyncTask<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");
+      return {std::string("failure reading async task"), 0};
 
     StoredPointer SlabPtr = AsyncTaskObj->PrivateStorage.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;
+    return {llvm::None, SlabPtr};
   }
 
 private:

include/swift/Reflection/RuntimeInternals.h

@@ -86,6 +86,7 @@ struct StackAllocator {
   bool FirstSlabIsPreallocated;
 
   struct Slab {
+    typename Runtime::StoredPointer Metadata;
     typename Runtime::StoredPointer Next;
     uint32_t Capacity;
     uint32_t CurrentOffset;

include/swift/Remote/MetadataReader.h

@@ -2242,25 +2242,21 @@ class MetadataReader {
       return false;
     };
 
-    bool isTypeContext = false;
     switch (auto contextKind = descriptor->getKind()) {
     case ContextDescriptorKind::Class:
       if (!getContextName())
         return nullptr;
       nodeKind = Demangle::Node::Kind::Class;
-      isTypeContext = true;
       break;
     case ContextDescriptorKind::Struct:
       if (!getContextName())
         return nullptr;
       nodeKind = Demangle::Node::Kind::Structure;
-      isTypeContext = true;
       break;
     case ContextDescriptorKind::Enum:
       if (!getContextName())
         return nullptr;
       nodeKind = Demangle::Node::Kind::Enum;
-      isTypeContext = true;
       break;
     case ContextDescriptorKind::Protocol: {
       if (!getContextName())

include/swift/SwiftRemoteMirror/SwiftRemoteMirror.h

@@ -44,6 +44,9 @@ extern unsigned long long swift_reflection_classIsSwiftMask;
 ///     first attempted fix to use the right AsyncTask layout.
 /// 1 - Indicates that swift_reflection_iterateAsyncTaskAllocations has been
 ///     actually fixed to use the right AsyncTask layout.
+/// 2 - swift_reflection_iterateAsyncTaskAllocations has been replaced by
+///     swift_reflection_asyncTaskSlabPointer and
+///     swift_reflection_asyncTaskSlabAllocations.
 SWIFT_REMOTE_MIRROR_LINKAGE extern uint32_t swift_reflection_libraryVersion;
 
 /// Get the metadata version supported by the Remote Mirror library.
@@ -395,33 +398,40 @@ 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.
+/// Get the first allocation slab for a 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.
+/// It is possible that the async task object hasn't allocated a slab yet, in
+/// which case the slab pointer will be NULL. If non-NULL, the returned slab
+/// pointer may be a separate heap allocation, or it may be interior to some
+/// allocation used by the task.
+SWIFT_REMOTE_MIRROR_LINKAGE
+swift_async_task_slab_return_t
+swift_reflection_asyncTaskSlabPointer(SwiftReflectionContextRef ContextRef,
+                                      swift_reflection_ptr_t AsyncTaskPtr);
+
+/// Iterate over the allocations in the given async task allocator slab.
+/// This allocation must have an "isa" value (scare quotes because it's not a
+/// real object) equal to _swift_concurrency_debug_asyncTaskSlabMetadata.
+///
+/// Calls the passed in Call function for each allocation in the slab. 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.
+/// A slab may be part of a chain of slabs, so the
+/// function may be called more than once.
 ///
 /// 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);
+swift_async_task_slab_allocations_return_t
+swift_reflection_asyncTaskSlabAllocations(SwiftReflectionContextRef ContextRef,
+                                          swift_reflection_ptr_t SlabPtr);
 
 #ifdef __cplusplus
 } // extern "C"

include/swift/SwiftRemoteMirror/SwiftRemoteMirrorTypes.h

@@ -189,12 +189,42 @@ typedef struct swift_metadata_cache_node {
   swift_reflection_ptr_t Right;
 } swift_metadata_cache_node_t;
 
+/// The return value when getting an async task's slab pointer.
+typedef struct swift_async_task_slab_return {
+  /// On failure, a pointer to a string describing the error. On success, NULL.
+  /// This pointer remains valid until the next
+  /// swift_reflection call on the given context.
+  const char *Error;
+
+  /// The task's slab pointer, if no error occurred.
+  swift_reflection_ptr_t SlabPtr;
+} swift_async_task_slab_return_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;
 
+typedef struct swift_async_task_slab_allocations_return {
+  /// On failure, a pointer to a string describing the error. On success, NULL.
+  /// This pointer remains valid until the next
+  /// swift_reflection call on the given context.
+  const char *Error;
+
+  /// The remote pointer to the next slab, or NULL/0 if none.
+  swift_reflection_ptr_t NextSlab;
+
+  /// The size of the entire slab, in bytes.
+  unsigned SlabSize;
+
+  /// The number of chunks pointed to by Chunks.
+  unsigned ChunkCount;
+
+  /// A pointer to the chunks, if no error occurred.
+  swift_async_task_allocation_chunk_t *Chunks;
+} swift_async_task_slab_allocations_return_t;
+
 /// An opaque pointer to a context which maintains state and
 /// caching of reflection structure for heap instances.
 typedef struct SwiftReflectionContext *SwiftReflectionContextRef;

stdlib/public/Concurrency/Debug.h

@@ -32,6 +32,10 @@ const void *const _swift_concurrency_debug_jobMetadata;
 SWIFT_EXPORT_FROM(swift_Concurrency)
 const void *const _swift_concurrency_debug_asyncTaskMetadata;
 
+/// 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;
+
 } // namespace swift
 
 #endif

stdlib/public/Concurrency/Task.cpp

@@ -68,6 +68,10 @@ using namespace swift;
 using FutureFragment = AsyncTask::FutureFragment;
 using TaskGroup = swift::TaskGroup;
 
+Metadata swift::TaskAllocatorSlabMetadata;
+const void *const swift::_swift_concurrency_debug_asyncTaskSlabMetadata =
+    &TaskAllocatorSlabMetadata;
+
 void FutureFragment::destroy() {
   auto queueHead = waitQueue.load(std::memory_order_acquire);
   switch (queueHead.getStatus()) {

stdlib/public/Concurrency/TaskPrivate.h

@@ -270,8 +270,9 @@ class alignas(sizeof(void*) * 2) ActiveTaskStatus {
 
 /// The size of an allocator slab.
 static constexpr size_t SlabCapacity = 1000;
+extern Metadata TaskAllocatorSlabMetadata;
 
-using TaskAllocator = StackAllocator<SlabCapacity>;
+using TaskAllocator = StackAllocator<SlabCapacity, &TaskAllocatorSlabMetadata>;
 
 /// Private storage in an AsyncTask object.
 struct AsyncTask::PrivateStorage {

stdlib/public/SwiftRemoteMirror/SwiftRemoteMirror.cpp

@@ -18,7 +18,7 @@ extern "C" {
 SWIFT_REMOTE_MIRROR_LINKAGE
 unsigned long long swift_reflection_classIsSwiftMask = 2;
 
-SWIFT_REMOTE_MIRROR_LINKAGE uint32_t swift_reflection_libraryVersion = 1;
+SWIFT_REMOTE_MIRROR_LINKAGE uint32_t swift_reflection_libraryVersion = 2;
 }
 
 #include "swift/Demangling/Demangler.h"
@@ -43,6 +43,7 @@ struct SwiftReflectionContext {
   std::vector<std::function<void()>> freeFuncs;
   std::vector<std::tuple<swift_addr_t, swift_addr_t>> dataSegments;
   std::string lastString;
+  std::vector<swift_async_task_allocation_chunk_t> lastChunks;
 
   SwiftReflectionContext(MemoryReaderImpl impl) {
     auto Reader = std::make_shared<CMemoryReader>(impl);
@@ -774,22 +775,46 @@ 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);
+swift_async_task_slab_return_t
+swift_reflection_asyncTaskSlabPointer(SwiftReflectionContextRef ContextRef,
+                                      swift_reflection_ptr_t AsyncTaskPtr) {
+  auto Context = ContextRef->nativeContext;
+  llvm::Optional<std::string> Error;
+  NativeReflectionContext::StoredPointer SlabPtr;
+  std::tie(Error, SlabPtr) = Context->asyncTaskSlabPtr(AsyncTaskPtr);
+
+  swift_async_task_slab_return_t Result = {};
+  Result.Error = returnableCString(ContextRef, Error);
+  Result.SlabPtr = SlabPtr;
+  return Result;
+}
+
+swift_async_task_slab_allocations_return_t
+swift_reflection_asyncTaskSlabAllocations(SwiftReflectionContextRef ContextRef,
+                                          swift_reflection_ptr_t SlabPtr) {
+  auto Context = ContextRef->nativeContext;
+  llvm::Optional<std::string> Error;
+  NativeReflectionContext::AsyncTaskSlabInfo Info;
+  std::tie(Error, Info) = Context->asyncTaskSlabAllocations(SlabPtr);
+
+  swift_async_task_slab_allocations_return_t Result = {};
+  Result.Error = returnableCString(ContextRef, Error);
+
+  Result.NextSlab = Info.NextSlab;
+  Result.SlabSize = Info.SlabSize;
+
+  ContextRef->lastChunks.clear();
+  ContextRef->lastChunks.reserve(Info.Chunks.size());
+  for (auto &Chunk : Info.Chunks) {
+    swift_async_task_allocation_chunk_t ConvertedChunk;
+    ConvertedChunk.Start = Chunk.Start;
+    ConvertedChunk.Length = Chunk.Length;
+    ConvertedChunk.Kind = convertAllocationChunkKind(Chunk.Kind);
+    ContextRef->lastChunks.push_back(ConvertedChunk);
+  }
+
+  Result.ChunkCount = ContextRef->lastChunks.size();
+  Result.Chunks = ContextRef->lastChunks.data();
+
+  return Result;
 }

stdlib/public/runtime/StackAllocator.h

@@ -30,7 +30,7 @@ namespace swift {
 ///
 /// StackAllocator performs fast allocation and deallocation of memory by
 /// implementing a bump-pointer allocation strategy.
-/// 
+///
 /// This isn't strictly a bump-pointer allocator as it uses backing slabs of
 /// memory rather than relying on a boundless contiguous heap. However, it has
 /// bump-pointer semantics in that it is a monotonically growing pool of memory
@@ -45,7 +45,10 @@ namespace swift {
 /// It's possible to place the first slab into pre-allocated memory.
 ///
 /// The SlabCapacity specifies the capacity for newly allocated slabs.
-template <size_t SlabCapacity>
+///
+/// SlabMetadataPtr specifies a fake metadata pointer to place at the beginning
+/// of slab allocations, so analysis tools can identify them.
+template <size_t SlabCapacity, Metadata *SlabMetadataPtr>
 class StackAllocator {
 private:
 
@@ -86,6 +89,10 @@ class StackAllocator {
   /// This struct is actually just the slab header. The slab buffer is tail
   /// allocated after Slab.
   struct Slab {
+    /// A fake metadata pointer that analysis tools can use to identify slab
+    /// allocations.
+    const void *metadata;
+
     /// A single linked list of all allocated slabs.
     Slab *next = nullptr;
 
@@ -95,7 +102,8 @@ class StackAllocator {
 
     // Here starts the tail allocated memory buffer of the slab.
 
-    Slab(size_t newCapacity) : capacity(newCapacity) {
+    Slab(size_t newCapacity)
+        : metadata(SlabMetadataPtr), capacity(newCapacity) {
       assert((size_t)capacity == newCapacity && "capacity overflow");
     }