[Runtime] Use objc_addLoadImageFunc when available to find out about newly loaded images.

rdar://problem/49742015

Author: mikeash

stdlib/public/runtime/ImageInspectionMachO.cpp

@@ -23,6 +23,7 @@
 #include "ImageInspection.h"
 #include <mach-o/dyld.h>
 #include <mach-o/getsect.h>
+#include <objc/runtime.h>
 #include <assert.h>
 #include <dlfcn.h>
 
@@ -55,7 +56,7 @@ extern "C" void *_NSGetMachExecuteHeader();
 
 template <const char *SEGMENT_NAME, const char *SECTION_NAME,
          void CONSUME_BLOCK(const void *start, uintptr_t size)>
-void addImageCallback(const mach_header *mh, intptr_t vmaddr_slide) {
+void addImageCallback(const mach_header *mh) {
 #if __POINTER_WIDTH__ == 64
   assert(mh->magic == MH_MAGIC_64 && "loaded non-64-bit image?!");
 #endif
@@ -72,12 +73,17 @@ void addImageCallback(const mach_header *mh, intptr_t vmaddr_slide) {
 
   CONSUME_BLOCK(section, size);
 }
+template <const char *SEGMENT_NAME, const char *SECTION_NAME,
+         void CONSUME_BLOCK(const void *start, uintptr_t size)>
+void addImageCallback(const mach_header *mh, intptr_t vmaddr_slide) {
+  addImageCallback<SEGMENT_NAME, SECTION_NAME, CONSUME_BLOCK>(mh);
+}
 
 template <const char *SEGMENT_NAME, const char *SECTION_NAME,
           const char *SEGMENT_NAME2, const char *SECTION_NAME2,
           void CONSUME_BLOCK(const void *start, uintptr_t size,
                              const void *start2, uintptr_t size2)>
-void addImageCallback2Sections(const mach_header *mh, intptr_t vmaddr_slide) {
+void addImageCallback2Sections(const mach_header *mh) {
 #if __POINTER_WIDTH__ == 64
   assert(mh->magic == MH_MAGIC_64 && "loaded non-64-bit image?!");
 #endif
@@ -103,27 +109,48 @@ void addImageCallback2Sections(const mach_header *mh, intptr_t vmaddr_slide) {
 
   CONSUME_BLOCK(section, size, section2, size2);
 }
+template <const char *SEGMENT_NAME, const char *SECTION_NAME,
+          const char *SEGMENT_NAME2, const char *SECTION_NAME2,
+          void CONSUME_BLOCK(const void *start, uintptr_t size,
+                             const void *start2, uintptr_t size2)>
+void addImageCallback2Sections(const mach_header *mh, intptr_t vmaddr_slide) {
+  addImageCallback2Sections<SEGMENT_NAME, SECTION_NAME,
+                            SEGMENT_NAME2, SECTION_NAME2,
+                            CONSUME_BLOCK>(mh);
+}
+
 } // end anonymous namespace
 
+#if OBJC_ADDLOADIMAGEFUNC_DEFINED
+#define REGISTER_FUNC(...)                                               \
+  if (__builtin_available(macOS 10.15, iOS 13, tvOS 13, watchOS 6, *)) { \
+    objc_addLoadImageFunc(__VA_ARGS__);                                  \
+  } else {                                                               \
+    _dyld_register_func_for_add_image(__VA_ARGS__);                      \
+  }
+#else
+#define REGISTER_FUNC(...) _dyld_register_func_for_add_image(__VA_ARGS__)
+#endif
+
 void swift::initializeProtocolLookup() {
-  _dyld_register_func_for_add_image(
+  REGISTER_FUNC(
     addImageCallback<TextSegment, ProtocolsSection,
                      addImageProtocolsBlockCallback>);
 }
 
 void swift::initializeProtocolConformanceLookup() {
-  _dyld_register_func_for_add_image(
+  REGISTER_FUNC(
     addImageCallback<TextSegment, ProtocolConformancesSection,
                      addImageProtocolConformanceBlockCallback>);
 }
 void swift::initializeTypeMetadataRecordLookup() {
-  _dyld_register_func_for_add_image(
+  REGISTER_FUNC(
     addImageCallback<TextSegment, TypeMetadataRecordSection,
                      addImageTypeMetadataRecordBlockCallback>);
 }
 
 void swift::initializeDynamicReplacementLookup() {
-  _dyld_register_func_for_add_image(
+  REGISTER_FUNC(
       addImageCallback2Sections<TextSegment, DynamicReplacementSection,
                                 TextSegment, DynamicReplacementSomeSection,
                                 addImageDynamicReplacementBlockCallback>);

test/stdlib/dlopen_race.swift

@@ -30,6 +30,11 @@ func register_func_for_add_image(_ f: add_image_callback) {
 // for them to be used. rdar://problem/49742015
 var add_image_count = 0
 DlopenRaceTests.test("race") {
+  // This test is expected to fail unless the ObjC notification is supported.
+  let RTLD_DEFAULT = UnsafeMutableRawPointer(bitPattern: -2)
+  let objc_addLoadImageFunc = dlsym(RTLD_DEFAULT, "objc_addLoadImageFunc");
+  if objc_addLoadImageFunc == nil { return }
+  
   register_func_for_add_image({ header, slide in
     // The protocol conformance check in the print call is enough to trigger
     // ObjC class initialization in the newly opened image if Swift has