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

We remove the existing `swift_reflection_iterateAsyncTaskAllocations` API that attempts to provide all necessary information about a tasks's allocations starting from the task. Instead, we split it into two pieces: `swift_reflection_asyncTaskSlabPointer` to get the first slab for a task, and `+swift_reflection_asyncTaskSlabAllocations` to get the allocations in a slab, and a pointer to the next slab.

We also add a dummy metadata pointer to the beginning of each slab. This allows tools to identify slab allocations on the heap without needing to locate every single async task object. They can then use `swift_reflection_asyncTaskSlabAllocations` on such allocations to find out about the contents.

rdar://82549631

Author: mikeash

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");
     }