1 files changed, 60 insertions, 26 deletions

diff --git a/lib/xray/xray_interface.cc b/lib/xray/xray_interface.cc
index 20a2b66c4401c..26ec161fe8600 100644
--- a/lib/xray/xray_interface.cc
+++ b/lib/xray/xray_interface.cc
@@ -15,7 +15,6 @@
 
 #include "xray_interface_internal.h"
 
-#include <atomic>
 #include <cstdint>
 #include <cstdio>
 #include <errno.h>
@@ -35,12 +34,21 @@ static const int16_t cSledLength = 12;
 static const int16_t cSledLength = 32;
 #elif defined(__arm__)
 static const int16_t cSledLength = 28;
+#elif SANITIZER_MIPS32
+static const int16_t cSledLength = 48;
+#elif SANITIZER_MIPS64
+static const int16_t cSledLength = 64;
+#elif defined(__powerpc64__)
+static const int16_t cSledLength = 8;
 #else
 #error "Unsupported CPU Architecture"
 #endif /* CPU architecture */
 
 // This is the function to call when we encounter the entry or exit sleds.
-std::atomic<void (*)(int32_t, XRayEntryType)> XRayPatchedFunction{nullptr};
+__sanitizer::atomic_uintptr_t XRayPatchedFunction{0};
+
+// This is the function to call from the arg1-enabled sleds/trampolines.
+__sanitizer::atomic_uintptr_t XRayArgLogger{0};
 
 // MProtectHelper is an RAII wrapper for calls to mprotect(...) that will undo
 // any successful mprotect(...) changes. This is used to make a page writeable
@@ -79,13 +87,18 @@ public:
 
 } // namespace __xray
 
-extern std::atomic<bool> XRayInitialized;
-extern std::atomic<__xray::XRaySledMap> XRayInstrMap;
+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 (XRayInitialized.load(std::memory_order_acquire)) {
-    __xray::XRayPatchedFunction.store(entry, std::memory_order_release);
+  if (__sanitizer::atomic_load(&XRayInitialized,
+                               __sanitizer::memory_order_acquire)) {
+
+    __sanitizer::atomic_store(&__xray::XRayPatchedFunction,
+                              reinterpret_cast<uint64_t>(entry),
+                              __sanitizer::memory_order_release);
     return 1;
   }
   return 0;
@@ -95,7 +108,7 @@ int __xray_remove_handler() XRAY_NEVER_INSTRUMENT {
   return __xray_set_handler(nullptr);
 }
 
-std::atomic<bool> XRayPatching{false};
+__sanitizer::atomic_uint8_t XRayPatching{0};
 
 using namespace __xray;
 
@@ -115,34 +128,37 @@ public:
 };
 
 template <class Function>
-CleanupInvoker<Function> ScopeCleanup(Function Fn) XRAY_NEVER_INSTRUMENT {
+CleanupInvoker<Function> scopeCleanup(Function Fn) XRAY_NEVER_INSTRUMENT {
   return CleanupInvoker<Function>{Fn};
 }
 
-// ControlPatching implements the common internals of the patching/unpatching
+// controlPatching implements the common internals of the patching/unpatching
 // implementation. |Enable| defines whether we're enabling or disabling the
 // runtime XRay instrumentation.
-XRayPatchingStatus ControlPatching(bool Enable) XRAY_NEVER_INSTRUMENT {
-  if (!XRayInitialized.load(std::memory_order_acquire))
+XRayPatchingStatus controlPatching(bool Enable) XRAY_NEVER_INSTRUMENT {
+  if (!__sanitizer::atomic_load(&XRayInitialized,
+                               __sanitizer::memory_order_acquire))
     return XRayPatchingStatus::NOT_INITIALIZED; // Not initialized.
 
-  static bool NotPatching = false;
-  if (!XRayPatching.compare_exchange_strong(NotPatching, true,
-                                            std::memory_order_acq_rel,
-                                            std::memory_order_acquire)) {
+  uint8_t NotPatching = false;
+  if (!__sanitizer::atomic_compare_exchange_strong(
+          &XRayPatching, &NotPatching, true, __sanitizer::memory_order_acq_rel))
     return XRayPatchingStatus::ONGOING; // Already patching.
-  }
 
-  bool PatchingSuccess = false;
-  auto XRayPatchingStatusResetter = ScopeCleanup([&PatchingSuccess] {
-    if (!PatchingSuccess) {
-      XRayPatching.store(false, std::memory_order_release);
-    }
+  uint8_t PatchingSuccess = false;
+  auto XRayPatchingStatusResetter = scopeCleanup([&PatchingSuccess] {
+    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 = XRayInstrMap.load(std::memory_order_acquire);
+  XRaySledMap InstrMap;
+  {
+    __sanitizer::SpinMutexLock Guard(&XRayInstrMapMutex);
+    InstrMap = XRayInstrMap;
+  }
   if (InstrMap.Entries == 0)
     return XRayPatchingStatus::NOT_INITIALIZED;
 
@@ -179,7 +195,7 @@ XRayPatchingStatus ControlPatching(bool Enable) XRAY_NEVER_INSTRUMENT {
     bool Success = false;
     switch (Sled.Kind) {
     case XRayEntryType::ENTRY:
-      Success = patchFunctionEntry(Enable, FuncId, Sled);
+      Success = patchFunctionEntry(Enable, FuncId, Sled, __xray_FunctionEntry);
       break;
     case XRayEntryType::EXIT:
       Success = patchFunctionExit(Enable, FuncId, Sled);
@@ -187,21 +203,39 @@ XRayPatchingStatus ControlPatching(bool Enable) XRAY_NEVER_INSTRUMENT {
     case XRayEntryType::TAIL:
       Success = patchFunctionTailExit(Enable, FuncId, Sled);
       break;
+    case XRayEntryType::LOG_ARGS_ENTRY:
+      Success = patchFunctionEntry(Enable, FuncId, Sled, __xray_ArgLoggerEntry);
+      break;
     default:
       Report("Unsupported sled kind: %d\n", int(Sled.Kind));
       continue;
     }
     (void)Success;
   }
-  XRayPatching.store(false, std::memory_order_release);
+  __sanitizer::atomic_store(&XRayPatching, false,
+                            __sanitizer::memory_order_release);
   PatchingSuccess = true;
   return XRayPatchingStatus::SUCCESS;
 }
 
 XRayPatchingStatus __xray_patch() XRAY_NEVER_INSTRUMENT {
-  return ControlPatching(true);
+  return controlPatching(true);
 }
 
 XRayPatchingStatus __xray_unpatch() XRAY_NEVER_INSTRUMENT {
-  return ControlPatching(false);
+  return controlPatching(false);
+}
+
+int __xray_set_handler_arg1(void (*Handler)(int32_t, XRayEntryType, uint64_t)) {
+  if (!__sanitizer::atomic_load(&XRayInitialized,
+                                __sanitizer::memory_order_acquire))
+    return 0;
+
+  // A relaxed write might not be visible even if the current thread gets
+  // scheduled on a different CPU/NUMA node.  We need to wait for everyone to
+  // have this handler installed for consistency of collected data across CPUs.
+  __sanitizer::atomic_store(&XRayArgLogger, reinterpret_cast<uint64_t>(Handler),
+                            __sanitizer::memory_order_release);
+  return 1;
 }
+int __xray_remove_handler_arg1() { return __xray_set_handler_arg1(nullptr); }