Add exported symbol on non-Apple platforms to enumerate MetadataSection structures at runtime.

Author: grynspan

include/swift/Runtime/Concurrent.h

@@ -520,10 +520,14 @@ template <class ElemTy> struct ConcurrentReadableArray {
     // These are marked as ref-qualified (the &) to make sure they can't be
     // called on temporaries, since the temporary would be destroyed before the
     // return value can be used, making it invalid.
-    const ElemTy *begin() & { return Start; }
-    const ElemTy *end() & { return Start + Count; }
+    const ElemTy *begin() const& { return Start; }
+    const ElemTy *end() const& { return Start + Count; }
+    const ElemTy& operator [](size_t index) const& {
+      assert(index < count() && "out-of-bounds access to snapshot element");
+      return Start[index];
+    }
 
-    size_t count() { return Count; }
+    size_t count() const { return Count; }
   };
 
   // This type cannot be safely copied or moved.

stdlib/public/SwiftShims/MetadataSections.h

@@ -21,6 +21,10 @@
 #ifndef SWIFT_STDLIB_SHIMS_METADATASECTIONS_H
 #define SWIFT_STDLIB_SHIMS_METADATASECTIONS_H
 
+#if defined(__cplusplus) && !defined(__swift__)
+#include <atomic>
+#endif
+
 #include "SwiftStddef.h"
 #include "SwiftStdint.h"
 
@@ -36,7 +40,12 @@ typedef struct MetadataSectionRange {
 } MetadataSectionRange;
 
 
-/// Identifies the address space ranges for the Swift metadata required by the Swift runtime.
+/// Identifies the address space ranges for the Swift metadata required by the
+/// Swift runtime.
+///
+/// \warning If you change the size of this structure by adding fields, it is an
+///   ABI-breaking change on platforms that use it. Make sure to increment
+///   \c CurrentSectionMetadataVersion if you do.
 struct MetadataSections {
   __swift_uintptr_t version;
 
@@ -56,20 +65,23 @@ struct MetadataSections {
   /// loader (i.e. no equivalent of \c __dso_handle or \c __ImageBase), this
   /// field is ignored and should be set to \c nullptr.
   ///
+  /// \bug When imported into Swift, this field is not atomic.
+  ///
   /// \sa swift_addNewDSOImage()
+#if defined(__swift__) || defined(__STDC_NO_ATOMICS__)
   const void *baseAddress;
+#elif defined(__cplusplus)
+  std::atomic<const void *> baseAddress;
+#else
+  _Atomic(const void *) baseAddress;
+#endif
 
-  /// `next` and `prev` are used by the runtime to construct a
-  /// circularly doubly linked list to quickly iterate over the metadata
-  /// from each image loaded into the address space.  These are invasive
-  /// to enable the runtime registration, which occurs at image load time, to
-  /// be allocation-free as it is invoked from an image constructor function
-  /// context where the system may not yet be ready to perform allocations.
-  /// Additionally, avoiding the allocation enables a fast load operation, which
-  /// directly impacts application load time.
-  struct MetadataSections *next;
-  struct MetadataSections *prev;
-
+  /// Unused.
+  ///
+  /// These pointers (or the space they occupy) can be repurposed without
+  /// causing ABI breakage. Set them to \c nullptr.
+  void *unused0;
+  void *unused1;
 
   MetadataSectionRange swift5_protocols;
   MetadataSectionRange swift5_protocol_conformances;

stdlib/public/runtime/ImageInspectionCommon.cpp

@@ -24,68 +24,71 @@
 #include "../SwiftShims/Visibility.h"
 #include "../SwiftShims/MetadataSections.h"
 #include "ImageInspection.h"
+#include "swift/Basic/Lazy.h"
+#include "swift/Runtime/Concurrent.h"
 
+#include <algorithm>
+#include <atomic>
+#include <cstdlib>
 
 namespace swift {
 
-#ifndef NDEBUG
-static swift::MetadataSections *registered = nullptr;
+static Lazy<ConcurrentReadableArray<swift::MetadataSections *>> registered;
 
-static void record(swift::MetadataSections *sections) {
-  if (registered == nullptr) {
-    registered = sections;
-    sections->next = sections->prev = sections;
-  } else {
-    registered->prev->next = sections;
-    sections->next = registered;
-    sections->prev = registered->prev;
-    registered->prev = sections;
-  }
-}
-#endif
-
-static const void *
-getMetadataSectionBaseAddress(swift::MetadataSections *sections) {
-  // If the base address was not set by the caller of swift_addNewDSOImage()
-  // then we can assume that the caller was built against an older version of
-  // the runtime that did not capture a value for this field. Currently nothing
-  // is actively using the image's base address outside of tests that are built
-  // with the runtime/stdlib, so there's no need to try to fix up the value. If
-  // something in the runtime starts using it, we will want to either:
-  // 1. Resolve the address from a known-good address like swift5_protocols when
-  //    the image is first loaded (in this function);
-  // 1. Resolve the address from a known-good address like swift5_protocols when
-  //    the address is first used (and atomically swap the address back so we
-  //    don't incur the cost of lookupSymbol() each time we need it; or
-  // 3. Introduce an ABI-breaking change so that all binaries are rebuilt and
-  //    start supplying a value for this field.
+/// Adjust the \c baseAddress field of a metadata sections structure.
+///
+/// \param sections A pointer to a valid \c swift::MetadataSections structure.
+///
+/// This function should be called at least once before the structure or its
+/// address is passed to code outside this file to ensure that the structure's
+/// \c baseAddress field correctly points to the base address of the image it
+/// is describing.
+static void fixupMetadataSectionBaseAddress(swift::MetadataSections *sections) {
+  bool fixupNeeded = false;
 
-#ifndef NDEBUG
 #if defined(__ELF__)
   // If the base address was set but the image is an ELF image, it is going to
   // be __dso_handle which is not the value we expect (Dl_info::dli_fbase), so
-  // we need to fix it up. Since the base address is currently unused by the
-  // runtime outside tests, we don't normally do this work.
-  if (auto baseAddress = sections->baseAddress) {
-    swift::SymbolInfo symbolInfo;
-    if (lookupSymbol(baseAddress, &symbolInfo) && symbolInfo.baseAddress) {
-      sections->baseAddress = symbolInfo.baseAddress;
-    }
+  // we need to fix it up.
+  fixupNeeded = true;
+#elif !defined(__MACH__)
+  // For non-ELF, non-Apple platforms, if the base address is nullptr, it
+  // implies that this image was built against an older version of the runtime
+  // that did not capture any value for the base address.
+  auto oldBaseAddress = sections->baseAddress.load(std::memory_order_relaxed);
+  if (!oldBaseAddress) {
+    fixupNeeded = true;
   }
-#endif
 #endif
 
-  return sections->baseAddress;
+  if (fixupNeeded) {
+    // We need to fix up the base address. We'll need a known-good address in
+    // the same image: `sections` itself will work nicely.
+    swift::SymbolInfo symbolInfo;
+    if (lookupSymbol(sections, &symbolInfo) && symbolInfo.baseAddress) {
+        sections->baseAddress.store(symbolInfo.baseAddress,
+                                    std::memory_order_relaxed);
+    }
+  }
 }
 }
 
 SWIFT_RUNTIME_EXPORT
 void swift_addNewDSOImage(swift::MetadataSections *sections) {
-#ifndef NDEBUG
-  record(sections);
+#if 0
+  // Ensure the base address of the sections structure is correct.
+  //
+  // Currently disabled because none of the registration functions below
+  // actually do anything with the baseAddress field. Instead,
+  // swift_enumerateAllMetadataSections() is called by other individual
+  // functions, lower in this file, that yield metadata section pointers.
+  //
+  // If one of these registration functions starts needing the baseAddress
+  // field, this call should be enabled and the calls elsewhere in the file can
+  // be removed.
+  swift::fixupMetadataSectionBaseAddress(sections);
 #endif
-
-  auto baseAddress = swift::getMetadataSectionBaseAddress(sections);
+  auto baseAddress = sections->baseAddress.load(std::memory_order_relaxed);
 
   const auto &protocols_section = sections->swift5_protocols;
   const void *protocols = reinterpret_cast<void *>(protocols_section.start);
@@ -125,6 +128,29 @@ void swift_addNewDSOImage(swift::MetadataSections *sections) {
   if (accessible_funcs_section.length)
     swift::addImageAccessibleFunctionsBlockCallback(
         baseAddress, functions, accessible_funcs_section.length);
+
+  // Register this section for future enumeration by clients. This should occur
+  // after this function has done all other relevant work to avoid a race
+  // condition when someone calls swift_enumerateAllMetadataSections() on
+  // another thread.
+  swift::registered->push_back(sections);
+}
+
+SWIFT_RUNTIME_EXPORT
+void swift_enumerateAllMetadataSections(
+  bool (* body)(const swift::MetadataSections *sections, void *context),
+  void *context
+) {
+  auto snapshot = swift::registered->snapshot();
+  for (swift::MetadataSections *sections : snapshot) {
+    // Ensure the base address is fixed up before yielding the pointer.
+    swift::fixupMetadataSectionBaseAddress(sections);
+
+    // Yield the pointer and (if the callback returns false) break the loop.
+    if (!(* body)(sections, context)) {
+      return;
+    }
+  }
 }
 
 void swift::initializeProtocolLookup() {
@@ -146,19 +172,19 @@ void swift::initializeAccessibleFunctionsLookup() {
 
 SWIFT_RUNTIME_EXPORT
 const swift::MetadataSections *swift_getMetadataSection(size_t index) {
-  if (swift::registered == nullptr) {
-    return nullptr;
+  swift::MetadataSections *result = nullptr;
+
+  auto snapshot = swift::registered->snapshot();
+  if (index < snapshot.count()) {
+    result = snapshot[index];
   }
 
-  auto selected = swift::registered;
-  while (index > 0) {
-    selected = selected->next;
-    if (selected == swift::registered) {
-      return nullptr;
-    }
-    --index;
+  if (result) {
+    // Ensure the base address is fixed up before returning it.
+    swift::fixupMetadataSectionBaseAddress(result);
   }
-  return selected;
+
+  return result;
 }
 
 SWIFT_RUNTIME_EXPORT
@@ -184,19 +210,16 @@ void swift_getMetadataSectionBaseAddress(const swift::MetadataSections *section,
     *out_actual = nullptr;
   }
 
-  *out_expected = section->baseAddress;
+  // fixupMetadataSectionBaseAddress() was already called by
+  // swift_getMetadataSection(), presumably on the same thread, so we don't need
+  // to call it again here.
+  *out_expected = section->baseAddress.load(std::memory_order_relaxed);
 }
 
 SWIFT_RUNTIME_EXPORT
 size_t swift_getMetadataSectionCount() {
-  if (swift::registered == nullptr)
-    return 0;
-
-  size_t count = 1;
-  for (const auto *current = swift::registered->next;
-       current != swift::registered; current = current->next, ++count);
-
-  return count;
+  auto snapshot = swift::registered->snapshot();
+  return snapshot.count();
 }
 
 #endif // NDEBUG

stdlib/public/runtime/ImageInspectionCommon.h

@@ -71,6 +71,22 @@ struct SectionInfo {
 SWIFT_RUNTIME_EXPORT
 void swift_addNewDSOImage(struct swift::MetadataSections *sections);
 
+/// Enumerate all metadata sections in the current process that are known to the
+/// Swift runtime.
+///
+/// \param body A function to invoke once per metadata sections structure.
+///   If this function returns \c false, enumeration is stopped.
+/// \param context An additional context pointer to pass to \a body.
+///
+/// On Mach-O-based platforms (i.e. Apple platforms), this function is
+/// unavailable. On those plaforms, use dyld API to enumerate loaded images and
+/// their corresponding metadata sections.
+SWIFT_RUNTIME_EXPORT SWIFT_WEAK_IMPORT
+void swift_enumerateAllMetadataSections(
+  bool (* body)(const swift::MetadataSections *sections, void *context),
+  void *context
+);
+
 #ifndef NDEBUG
 
 SWIFT_RUNTIME_EXPORT

stdlib/public/runtime/SwiftRT-COFF.cpp

@@ -14,6 +14,7 @@
 #include "../SwiftShims/MetadataSections.h"
 
 #include <cstdint>
+#include <new>
 
 extern "C" const char __ImageBase[];
 
@@ -64,9 +65,9 @@ static void swift_image_constructor() {
   { reinterpret_cast<uintptr_t>(&__start_##name) + sizeof(__start_##name),     \
     reinterpret_cast<uintptr_t>(&__stop_##name) - reinterpret_cast<uintptr_t>(&__start_##name) - sizeof(__start_##name) }
 
-  sections = {
+  new (&sections) swift::MetadataSections {
       swift::CurrentSectionMetadataVersion,
-      __ImageBase,
+      { __ImageBase },
 
       nullptr,
       nullptr,

stdlib/public/runtime/SwiftRT-ELF.cpp

@@ -14,6 +14,7 @@
 #include "../SwiftShims/MetadataSections.h"
 
 #include <cstddef>
+#include <new>
 
 extern "C" const char __dso_handle[];
 
@@ -55,9 +56,9 @@ static void swift_image_constructor() {
   { reinterpret_cast<uintptr_t>(&__start_##name),                              \
     static_cast<uintptr_t>(&__stop_##name - &__start_##name) }
 
-  sections = {
+  new (&sections) swift::MetadataSections {
       swift::CurrentSectionMetadataVersion,
-      __dso_handle,
+      { __dso_handle },
 
       nullptr,
       nullptr,

test/Runtime/enumerate_metadata_sections.swift

@@ -0,0 +1,49 @@
+// RUN: %target-run-simple-swift
+// REQUIRES: executable_test
+
+#if os(Linux) || os(Windows)
+import Swift
+import SwiftShims
+import StdlibUnittest
+
+let EnumerateMetadataSections = TestSuite("EnumerateMetadataSections")
+
+@_silgen_name("swift_enumerateAllMetadataSections")
+func swift_enumerateAllMetadataSections(
+  _ body: @convention(c) (
+    _ sections: UnsafePointer<MetadataSections>,
+    _ context: UnsafeMutableRawPointer
+  ) -> Bool,
+  _ context: UnsafeMutableRawPointer
+)
+
+public protocol P { }
+public struct S: P { }
+
+EnumerateMetadataSections.test("swift_enumerateAllMetadataSections works") {
+  var sectionsEnumerated = 0
+  swift_enumerateAllMetadataSections({ sections, context in
+    let sectionsEnumerated = context.bindMemory(to: Int.self, capacity: 1)
+
+    // Confirm that the base address of the metadata sections was loaded.
+    let baseAddress = sections.pointee.baseAddress
+    expectNotNil(baseAddress)
+
+    // Confirm that P and S above have been emitted.
+    if baseAddress == #dsohandle {
+      expectNotNil(sections.pointee.swift5_protocols)
+      expectNotNil(sections.pointee.swift5_protocol_conformances)
+      expectNotNil(sections.pointee.swift5_type_metadata)
+    }
+
+    sectionsEnumerated.pointee += 1
+    return true
+  }, &sectionsEnumerated)
+
+  // Confirm that at least one section was enumerated.
+  expectGT(sectionsEnumerated, 0)
+}
+
+runAllTests()
+
+#endif