[XRay] [compiler-rt] [DSO] introduce a context for patching/unpatching

[pelikan]

Summary

For the upcoming DSO support, we want to reuse code between the "static"
version of XRay, already within the main binary and possibly initialized
at startup, and the "dynamic" part bundled with .so libraries which will
need to communicate with it. Wherever patching/unpatching happens needs
to be clearly defined for the patching code to recognize which memory is
it actually operating on.

This diff introduces an "always visible" pointer (*DSOContext), which is
local for each DSO, only visible within its own "copy" of XRay. Whoever
initializes that context (xray_init.cc in the main binary or xray_dso.cc
in a DSO during dlopen(3)) will populate that pointer. The advantage of
having a pointer instead of scattered functions and variables is:

• we have something to identify a DSO (it will have a numeric ID soon)
• it is easier to add per-DSO behaviour later, and control it globally
• it is clearer for the reader which C++ code is intended to be shared

When this is in, it is easy to take xray_interface{.cc,_internal.h} and
build a DSO version of the XRay runtime with it.

[dberris]

@pelikan and I are discussing this offline, pulling in @eugenis and @kcc -- we're prepping to get an RFC onto the mailing list(s) on how we might proceed with our options before continuing on this.

lib/xray/xray_init.cc

 extern "C" {
 void __xray_init();
 extern const XRaySledEntry __start_xray_instr_map[] __attribute__((weak));
 extern const XRaySledEntry __stop_xray_instr_map[] __attribute__((weak));
 extern const XRayFunctionSledIndex __start_xray_fn_idx[] __attribute__((weak));
 extern const XRayFunctionSledIndex __stop_xray_fn_idx[] __attribute__((weak));
 }
 
+namespace __xray {
+XRayDSOContext *DSOContext;
+}
+
 using namespace __xray;
 
 // When set to 'true' this means the XRay runtime has been initialised. We use
 // the weak symbols defined above (__start_xray_inst_map and
 // __stop_xray_instr_map) to initialise the instrumentation map that XRay uses
 // for runtime patching/unpatching of instrumentation points.
 //
 // FIXME: Support DSO instrumentation maps too. The current solution only works
 // for statically linked executables.
 __sanitizer::atomic_uint8_t XRayInitialized{0};
 
 // This should always be updated before XRayInitialized is updated.
-__sanitizer::SpinMutex XRayInstrMapMutex;
-XRaySledMap XRayInstrMap;
+XRayDSOContext MainBinaryContext;
 
 // Global flag to determine whether the flags have been initialized.
 __sanitizer::atomic_uint8_t XRayFlagsInitialized{0};
 
 // A mutex to allow only one thread to initialize the XRay data structures.
 __sanitizer::SpinMutex XRayInitMutex;
 
 // __xray_init() will do the actual loading of the current process' memory map
void __xray_init() XRAY_NEVER_INSTRUMENT {
 
   if (__start_xray_instr_map == nullptr) {
     if (Verbosity())
       Report("XRay instrumentation map missing. Not initializing XRay.\n");
     return;
   }
 
   {
-    __sanitizer::SpinMutexLock Guard(&XRayInstrMapMutex);
-    XRayInstrMap.Sleds = __start_xray_instr_map;
-    XRayInstrMap.Entries = __stop_xray_instr_map - __start_xray_instr_map;
-    XRayInstrMap.SledsIndex = __start_xray_fn_idx;
-    XRayInstrMap.Functions = __stop_xray_fn_idx - __start_xray_fn_idx;
+    __sanitizer::SpinMutexLock Guard(&MainBinaryContext.Lock);
+    XRaySledMap &Map = MainBinaryContext.Map;
+    Map.Sleds = __start_xray_instr_map;
+    Map.Entries = __stop_xray_instr_map - __start_xray_instr_map;
+    Map.SledsIndex = __start_xray_fn_idx;
+    Map.Functions = __stop_xray_fn_idx - __start_xray_fn_idx;
   }
+  DSOContext = &MainBinaryContext;
   __sanitizer::atomic_store(&XRayInitialized, true,
                             __sanitizer::memory_order_release);
 
 #ifndef XRAY_NO_PREINIT
   if (flags()->patch_premain)
     __xray_patch();
 #endif
 }
 
 #ifndef XRAY_NO_PREINIT
 __attribute__((section(".preinit_array"),
                used)) void (*__local_xray_preinit)(void) = __xray_init;
 #endif

lib/xray/xray_interface.cc

~MProtectHelper() XRAY_NEVER_INSTRUMENT {
     if (MustCleanup) {
       mprotect(PageAlignedAddr, MProtectLen, PROT_READ | PROT_EXEC);
     }
   }
 };
 
 } // namespace __xray
 
-extern __sanitizer::SpinMutex XRayInstrMapMutex;
 extern __sanitizer::atomic_uint8_t XRayInitialized;
-extern __xray::XRaySledMap XRayInstrMap;
 
 int __xray_set_handler(void (*entry)(int32_t,
                                      XRayEntryType)) XRAY_NEVER_INSTRUMENT {
   if (__sanitizer::atomic_load(&XRayInitialized,
                                __sanitizer::memory_order_acquire)) {
 
     __sanitizer::atomic_store(&__xray::XRayPatchedFunction,
                               reinterpret_cast<uintptr_t>(entry),
if (!PatchingSuccess)
       __sanitizer::atomic_store(&XRayPatching, false,
                                 __sanitizer::memory_order_release);
   });
 
   // Step 1: Compute the function id, as a unique identifier per function in the
   // instrumentation map.
   XRaySledMap InstrMap;
   {
-    __sanitizer::SpinMutexLock Guard(&XRayInstrMapMutex);
-    InstrMap = XRayInstrMap;
+    __sanitizer::SpinMutexLock Guard(&__xray::DSOContext->Lock);
+    InstrMap = __xray::DSOContext->Map;
   }
   if (InstrMap.Entries == 0)
     return XRayPatchingStatus::NOT_INITIALIZED;
 
   const uint64_t PageSize = GetPageSizeCached();
   if ((PageSize == 0) || ((PageSize & (PageSize - 1)) != 0)) {
     Report("System page size is not a power of two: %lld\n", PageSize);
     return XRayPatchingStatus::FAILED;
XRayPatchingStatus patchFunction(int32_t FuncId,
   uint8_t NotPatching = false;
   if (!__sanitizer::atomic_compare_exchange_strong(
           &XRayPatching, &NotPatching, true, __sanitizer::memory_order_acq_rel))
     return XRayPatchingStatus::ONGOING; // Already patching.
 
   // Next, we look for the function index.
   XRaySledMap InstrMap;
   {
-    __sanitizer::SpinMutexLock Guard(&XRayInstrMapMutex);
-    InstrMap = XRayInstrMap;
+    __sanitizer::SpinMutexLock Guard(&__xray::DSOContext->Lock);
+    InstrMap = __xray::DSOContext->Map;
   }
 
   // If we don't have an index, we can't patch individual functions.
   if (InstrMap.Functions == 0)
     return XRayPatchingStatus::NOT_INITIALIZED;
 
   // FuncId must be a positive number, less than the number of functions
   // instrumented.
int __xray_set_handler_arg1(void (*entry)(int32_t, XRayEntryType, uint64_t)) {
   __sanitizer::atomic_store(&XRayArgLogger, reinterpret_cast<uint64_t>(entry),
                             __sanitizer::memory_order_release);
   return 1;
 }
 
 int __xray_remove_handler_arg1() { return __xray_set_handler_arg1(nullptr); }
 
 uintptr_t __xray_function_address(int32_t FuncId) XRAY_NEVER_INSTRUMENT {
-  __sanitizer::SpinMutexLock Guard(&XRayInstrMapMutex);
-  if (FuncId <= 0 || static_cast<size_t>(FuncId) > XRayInstrMap.Functions)
+  __sanitizer::SpinMutexLock Guard(&__xray::DSOContext->Lock);
+  XRaySledMap &Map = __xray::DSOContext->Map;
+  if (FuncId <= 0 || static_cast<size_t>(FuncId) > Map.Functions)
     return 0;
-  return XRayInstrMap.SledsIndex[FuncId - 1].Begin->Function
+  return Map.SledsIndex[FuncId - 1].Begin->Function
 // On PPC, function entries are always aligned to 16 bytes. The beginning of a
 // sled might be a local entry, which is always +8 based on the global entry.
 // Always return the global entry.
 #ifdef __PPC__
          & ~0xf
 #endif
       ;
 }
 
 size_t __xray_max_function_id() XRAY_NEVER_INSTRUMENT {
-  __sanitizer::SpinMutexLock Guard(&XRayInstrMapMutex);
-  return XRayInstrMap.Functions;
+  __sanitizer::SpinMutexLock Guard(&__xray::DSOContext->Lock);
+  return __xray::DSOContext->Map.Functions;
 }

lib/xray/xray_interface_internal.h

 // This file is a part of XRay, a dynamic runtime instrumentation system.
 //
 // Implementation of the API functions. See also include/xray/xray_interface.h.
 //
 //===----------------------------------------------------------------------===//
 #ifndef XRAY_INTERFACE_INTERNAL_H
 #define XRAY_INTERFACE_INTERNAL_H
 
+#include "sanitizer_common/sanitizer_mutex.h"
 #include "sanitizer_common/sanitizer_platform.h"
 #include "xray/xray_interface.h"
 #include <cstddef>
 #include <cstdint>
 
 extern "C" {
 
 struct XRaySledEntry {

 
 struct XRaySledMap {
   const XRaySledEntry *Sleds;
   size_t Entries;
   const XRayFunctionSledIndex *SledsIndex;
   size_t Functions;
 };
 
+struct XRayDSOContext {
+  __sanitizer::SpinMutex Lock;
+  XRaySledMap Map;
+};
+
+extern XRayDSOContext *DSOContext;
+
 bool patchFunctionEntry(bool Enable, uint32_t FuncId,
                         const XRaySledEntry &Sled, void (*Trampoline)());
 bool patchFunctionExit(bool Enable, uint32_t FuncId, const XRaySledEntry &Sled);
 bool patchFunctionTailExit(bool Enable, uint32_t FuncId,
                            const XRaySledEntry &Sled);
 bool patchCustomEvent(bool Enable, uint32_t FuncId, const XRaySledEntry &Sled);
 
 } // namespace __xray