[Runtime] Only enable metadata allocation tracking when requested with an environment variable.

Expose _swift_debug variables for Remote Mirror to use for tracking metadata allocation, and adopt them in Remote Mirror.

rdar://problem/55481578

Author: mikeash

include/swift/Reflection/ReflectionContext.h

@@ -867,10 +867,30 @@ class ReflectionContext
 
   llvm::Optional<std::string>
   iterateMetadataAllocations(std::function<void (MetadataAllocation)> Call) {
-    std::string AllocationPoolName = "__ZL14AllocationPool";
-    auto AllocationPoolAddr = getReader().getSymbolAddress(AllocationPoolName);
+    std::string IterationEnabledName =
+      "__swift_debug_metadataAllocationIterationEnabled";
+    std::string AllocationPoolPointerName =
+      "__swift_debug_allocationPoolPointer";
+
+    auto IterationEnabledAddr =
+      getReader().getSymbolAddress(IterationEnabledName);
+    if (!IterationEnabledAddr)
+      return "unable to look up debug variable " + IterationEnabledName;
+    char IterationEnabled;
+    if (!getReader().readInteger(IterationEnabledAddr, &IterationEnabled))
+      return "failed to read value of " + IterationEnabledName;
+    if (!IterationEnabled)
+      return std::string("remote process does not have metadata allocation "
+                         "iteration enabled");
+
+    auto AllocationPoolAddrAddr =
+      getReader().getSymbolAddress(AllocationPoolPointerName);
+    if (!AllocationPoolAddrAddr)
+      return "unable to look up debug variable " + AllocationPoolPointerName;
+    auto AllocationPoolAddr =
+      getReader().readPointer(AllocationPoolAddrAddr, sizeof(StoredPointer));
     if (!AllocationPoolAddr)
-      return "unable to look up allocation pool symbol " + AllocationPoolName;
+      return "failed to read value of " + AllocationPoolPointerName;
 
     struct PoolRange {
       StoredPointer Begin;
@@ -886,7 +906,7 @@ class ReflectionContext
     };
 
     auto PoolBytes = getReader()
-      .readBytes(RemoteAddress(AllocationPoolAddr), sizeof(PoolRange));
+      .readBytes(AllocationPoolAddr->getResolvedAddress(), sizeof(PoolRange));
     auto Pool = reinterpret_cast<const PoolRange *>(PoolBytes.get());
     if (!Pool)
       return std::string("failure reading allocation pool contents");

include/swift/Runtime/Debug.h

@@ -231,6 +231,11 @@ bool _swift_reportFatalErrorsToDebugger;
 SWIFT_RUNTIME_STDLIB_SPI
 bool _swift_shouldReportFatalErrorsToDebugger();
 
+SWIFT_RUNTIME_STDLIB_SPI
+bool _swift_debug_metadataAllocationIterationEnabled;
+
+SWIFT_RUNTIME_STDLIB_SPI
+const void * const _swift_debug_allocationPoolPointer;
 
 SWIFT_RUNTIME_ATTRIBUTE_ALWAYS_INLINE
 inline static int swift_asprintf(char **strp, const char *fmt, ...) {

stdlib/public/runtime/Metadata.cpp

@@ -5494,18 +5494,31 @@ static std::atomic<PoolRange>
 AllocationPool{PoolRange{InitialAllocationPool.Pool,
                          sizeof(InitialAllocationPool.Pool)}};
 
+bool swift::_swift_debug_metadataAllocationIterationEnabled = false;
+const void * const swift::_swift_debug_allocationPoolPointer = &AllocationPool;
+
 void *MetadataAllocator::Allocate(size_t size, size_t alignment) {
-  if (Tag == 0) abort();
+  assert(Tag != 0);
   assert(alignment <= alignof(void*));
   assert(size % alignof(void*) == 0);
 
+  static OnceToken_t getenvToken;
+  SWIFT_ONCE_F(getenvToken, [](void *) {
+    const char *value =
+      getenv("SWIFT_DEBUG_ENABLE_METADATA_ALLOCATION_ITERATION");
+    if (value && (value[0] == '1' || value[0] == 'y' || value[0] == 'Y'))
+      _swift_debug_metadataAllocationIterationEnabled = true;
+  }, nullptr);
+
   // If the size is larger than the maximum, just use malloc.
   if (size > PoolRange::MaxPoolAllocationSize)
     return malloc(size);
 
   // Allocate out of the pool.
+  auto sizeWithHeader = size;
+  if (SWIFT_UNLIKELY(_swift_debug_metadataAllocationIterationEnabled))
+    sizeWithHeader += sizeof(AllocationHeader);
   PoolRange curState = AllocationPool.load(std::memory_order_relaxed);
-  auto sizeWithHeader = size + sizeof(AllocationHeader);
   while (true) {
     char *allocation;
     PoolRange newState;
@@ -5518,18 +5531,21 @@ void *MetadataAllocator::Allocate(size_t size, size_t alignment) {
       newState = PoolRange{curState.Begin + sizeWithHeader,
                            curState.Remaining - sizeWithHeader};
     } else {
-      auto poolSize = PoolRange::PageSize - sizeof(PoolTrailer);
+      auto poolSize = PoolRange::PageSize;
+      if (SWIFT_UNLIKELY(_swift_debug_metadataAllocationIterationEnabled))
+        poolSize -= sizeof(PoolTrailer);
       allocatedNewPage = true;
       allocation = new char[PoolRange::PageSize];
 
-      char *prevTrailer = curState.Begin + curState.Remaining;
-      PoolTrailer *newTrailer = (PoolTrailer *)(allocation + poolSize);
-      newTrailer->PrevTrailer = prevTrailer;
-      newTrailer->PoolSize = poolSize;
-
+      if (SWIFT_UNLIKELY(_swift_debug_metadataAllocationIterationEnabled)) {
+        PoolTrailer *newTrailer = (PoolTrailer *)(allocation + poolSize);
+        char *prevTrailer = curState.Begin + curState.Remaining;
+        newTrailer->PrevTrailer = prevTrailer;
+        newTrailer->PoolSize = poolSize;
+      }
       newState = PoolRange{allocation + sizeWithHeader,
                            poolSize - sizeWithHeader};
-      __asan_poison_memory_region(allocation, PoolRange::PageSize);
+      __asan_poison_memory_region(allocation, newState.Remaining);
     }
 
     // Swap in the new state.
@@ -5541,11 +5557,15 @@ void *MetadataAllocator::Allocate(size_t size, size_t alignment) {
       __msan_allocated_memory(allocation, sizeWithHeader);
       __asan_unpoison_memory_region(allocation, sizeWithHeader);
 
-      AllocationHeader *header = (AllocationHeader *)allocation;
-      header->Size = size;
-      header->Tag = Tag;
+      if (SWIFT_UNLIKELY(_swift_debug_metadataAllocationIterationEnabled)) {
+        AllocationHeader *header = (AllocationHeader *)allocation;
+        header->Size = size;
+        header->Tag = Tag;
 
-      return allocation + sizeof(AllocationHeader);
+        return allocation + sizeof(AllocationHeader);
+      } else {
+        return allocation;
+      }
     }
 
     // If it failed, go back to a neutral state and try again.