[XRay][compiler-rt] Use sanitizer_common's atomic ops

[dberris]

Summary

Instead of std::atomic APIs for atomic operations, we instead use APIs
include with sanitizer_common. This allows us to, at runtime, not have
to depend on potentially dynamically provided implementations of these
atomic operations.

Currently only do this for x86_64.

Fixes http://llvm.org/PR32274.

[dberris]

Retitle.

[pelikan]

Overall, I like it. Except for burning CPU while waiting on someone else's write(2) to (potentially) spinning rust.

lib/xray/xray_buffer_queue.h

44

This lock is being held over all Buffers (which may not be that bad actually). It is also held during the only use of apply(), which calls retryingWriteAll() and therefore will take a lot of time, including passively waiting on system calls. It would therefore make sense to use the blocking mutex here, to avoid burning CPU for potentially milliseconds.

Is there any downside to a blocking lock in this scenario? The other ones look fine as the locks only protect very small regions of memory.

[dberris]

• fixup: use a blocking mutex in the buffer queue

lib/xray/xray_buffer_queue.h

44

For the usage patterns we have, this should not be much of a problem. Let me explain why:

• We expect that getting buffers and returning buffers to happen more frequently than we do calling apply(...). That should be fairly quick.
• The caller to apply usually is a thread that has already ensured that all threads that want to write into the buffers will have returned the buffers already. While this is an implementation detail at a higher level (in the FDR logging implementation), we kind-of know it from here.

In the review thread for this original change, we decided to not hand-write our own spinlock, but it's more apt for the get and release cases. I'm fine using an explicitly blocking mutex here, but that means we're trading off the potential chance of spinning for a while (fairly low based on the way we're using this data structure) on some threads, with the potential that when we're doing FDR mode that we're blocking threads.

It's easy enough to switch anyway, so I've changed it to use a blocking mutex.

[dberris]

Submitted as rL298833.

lib/xray/tests/CMakeLists.txt

foreach(arch ${XRAY_TEST_ARCH})
       if(NOT APPLE)
         add_compiler_rt_test(XRayUnitTests ${testname}-${arch}
           OBJECTS ${TEST_OBJECTS}
           DEPS ${TEST_DEPS}
           LINK_FLAGS ${TARGET_LINK_FLAGS}
           -lstdc++ -lm ${CMAKE_THREAD_LIBS_INIT}
           -lpthread
           -L${COMPILER_RT_LIBRARY_OUTPUT_DIR} -lclang_rt.xray-${arch}
-          -latomic -ldl -lrt)
+          -ldl -lrt)
       endif()
       # FIXME: Figure out how to run even just the unit tests on APPLE.
     endforeach()
   endif()
 endmacro()
 
 if(COMPILER_RT_CAN_EXECUTE_TESTS)
   add_subdirectory(unit)
 endif()

lib/xray/xray_buffer_queue.h

 // This file is a part of XRay, a dynamic runtime instrumentation system.
 //
 // Defines the interface for a buffer queue implementation.
 //
 //===----------------------------------------------------------------------===//
 #ifndef XRAY_BUFFER_QUEUE_H
 #define XRAY_BUFFER_QUEUE_H
 
-#include <atomic>
-#include <cstdint>
+#include "sanitizer_common/sanitizer_atomic.h"
+#include "sanitizer_common/sanitizer_mutex.h"
 #include <deque>
-#include <mutex>
 #include <unordered_set>
 #include <utility>
 
 namespace __xray {
 
 /// BufferQueue implements a circular queue of fixed sized buffers (much like a
 /// freelist) but is concerned mostly with making it really quick to initialise,
 /// finalise, and get/return buffers to the queue. This is one key component of
 /// the "flight data recorder" (FDR) mode to support ongoing XRay function call
 /// trace collection.
 class BufferQueue {
 public:
   struct Buffer {
     void *Buffer = nullptr;
     size_t Size = 0;
   };
 
 private:
   size_t BufferSize;
 
   // We use a bool to indicate whether the Buffer has been used in this
   // freelist implementation.
   std::deque<std::tuple<Buffer, bool>> Buffers;
-  std::mutex Mutex;
+  __sanitizer::BlockingMutex Mutex;
   std::unordered_set<void *> OwnedBuffers;
-  std::atomic<bool> Finalizing;
+  __sanitizer::atomic_uint8_t Finalizing;
 
 public:
   enum class ErrorCode : unsigned {
     Ok,
     NotEnoughMemory,
     QueueFinalizing,
     UnrecognizedBuffer,
     AlreadyFinalized,
public:
   ErrorCode getBuffer(Buffer &Buf);
 
   /// Updates |Buf| to point to nullptr, with size 0.
   ///
   /// Returns:
   ///   - ...
   ErrorCode releaseBuffer(Buffer &Buf);
 
-  bool finalizing() const { return Finalizing.load(std::memory_order_acquire); }
+  bool finalizing() const {
+    return __sanitizer::atomic_load(&Finalizing,
+                                    __sanitizer::memory_order_acquire);
+  }
 
   /// Sets the state of the BufferQueue to finalizing, which ensures that:
   ///
   ///   - All subsequent attempts to retrieve a Buffer will fail.
   ///   - All releaseBuffer operations will not fail.
   ///
   /// After a call to finalize succeeds, all subsequent calls to finalize will
   /// fail with std::errc::state_not_recoverable.
   ErrorCode finalize();
 
   /// Applies the provided function F to each Buffer in the queue, only if the
   /// Buffer is marked 'used' (i.e. has been the result of getBuffer(...) and a
   /// releaseBuffer(...) operation.
   template <class F> void apply(F Fn) {
-    std::lock_guard<std::mutex> G(Mutex);
+    __sanitizer::BlockingMutexLock G(&Mutex);
     for (const auto &T : Buffers) {
       if (std::get<1>(T))
         Fn(std::get<0>(T));
     }
   }
 
   // Cleans up allocated buffers.
   ~BufferQueue();
 };
 
 } // namespace __xray
 
 #endif // XRAY_BUFFER_QUEUE_H

lib/xray/xray_buffer_queue.cc

 
 #include <cstdlib>
 #include <tuple>
 
 using namespace __xray;
 using namespace __sanitizer;
 
 BufferQueue::BufferQueue(std::size_t B, std::size_t N, bool &Success)
-    : BufferSize(B), Buffers(N), Mutex(), OwnedBuffers(), Finalizing(false) {
+    : BufferSize(B), Buffers(N), Mutex(), OwnedBuffers(), Finalizing{0} {
   for (auto &T : Buffers) {
     void *Tmp = malloc(BufferSize);
     if (Tmp == nullptr) {
       Success = false;
       return;
     }
 
     auto &Buf = std::get<0>(T);
     Buf.Buffer = Tmp;
     Buf.Size = B;
     OwnedBuffers.emplace(Tmp);
   }
   Success = true;
 }
 
 BufferQueue::ErrorCode BufferQueue::getBuffer(Buffer &Buf) {
-  if (Finalizing.load(std::memory_order_acquire))
+  if (__sanitizer::atomic_load(&Finalizing, __sanitizer::memory_order_acquire))
     return ErrorCode::QueueFinalizing;
-  std::lock_guard<std::mutex> Guard(Mutex);
+  __sanitizer::BlockingMutexLock Guard(&Mutex);
   if (Buffers.empty())
     return ErrorCode::NotEnoughMemory;
   auto &T = Buffers.front();
   auto &B = std::get<0>(T);
   Buf = B;
   B.Buffer = nullptr;
   B.Size = 0;
   Buffers.pop_front();
   return ErrorCode::Ok;
 }
 
 BufferQueue::ErrorCode BufferQueue::releaseBuffer(Buffer &Buf) {
   if (OwnedBuffers.count(Buf.Buffer) == 0)
     return ErrorCode::UnrecognizedBuffer;
-  std::lock_guard<std::mutex> Guard(Mutex);
+  __sanitizer::BlockingMutexLock Guard(&Mutex);
 
   // Now that the buffer has been released, we mark it as "used".
   Buffers.emplace(Buffers.end(), Buf, true /* used */);
   Buf.Buffer = nullptr;
   Buf.Size = 0;
   return ErrorCode::Ok;
 }
 
 BufferQueue::ErrorCode BufferQueue::finalize() {
-  if (Finalizing.exchange(true, std::memory_order_acq_rel))
+  if (__sanitizer::atomic_exchange(&Finalizing, 1,
+                                   __sanitizer::memory_order_acq_rel))
     return ErrorCode::QueueFinalizing;
   return ErrorCode::Ok;
 }
 
 BufferQueue::~BufferQueue() {
   for (auto &T : Buffers) {
     auto &Buf = std::get<0>(T);
     free(Buf.Buffer);
   }
 }

lib/xray/xray_fdr_logging.cc

 // Here we implement the Flight Data Recorder mode for XRay, where we use
 // compact structures to store records in memory as well as when writing out the
 // data to files.
 //
 //===----------------------------------------------------------------------===//
 #include "xray_fdr_logging.h"
 #include <algorithm>
 #include <bitset>
-#include <cassert>
 #include <cstring>
-#include <memory>
 #include <sys/syscall.h>
 #include <sys/time.h>
 #include <time.h>
 #include <unistd.h>
 #include <unordered_map>
 
+#include "sanitizer_common/sanitizer_atomic.h"
 #include "sanitizer_common/sanitizer_common.h"
 #include "xray/xray_interface.h"
 #include "xray/xray_records.h"
 #include "xray_buffer_queue.h"
 #include "xray_defs.h"
 #include "xray_fdr_logging_impl.h"
 #include "xray_flags.h"
 #include "xray_tsc.h"
 #include "xray_utils.h"
 
 namespace __xray {
 
 // Global BufferQueue.
 std::shared_ptr<BufferQueue> BQ;
 
-std::atomic<XRayLogInitStatus> LoggingStatus{
+__sanitizer::atomic_sint32_t LoggingStatus = {
     XRayLogInitStatus::XRAY_LOG_UNINITIALIZED};
 
-std::atomic<XRayLogFlushStatus> LogFlushStatus{
+__sanitizer::atomic_sint32_t LogFlushStatus = {
     XRayLogFlushStatus::XRAY_LOG_NOT_FLUSHING};
 
 std::unique_ptr<FDRLoggingOptions> FDROptions;
 
 XRayLogInitStatus fdrLoggingInit(std::size_t BufferSize, std::size_t BufferMax,
                                  void *Options,
                                  size_t OptionsSize) XRAY_NEVER_INSTRUMENT {
   assert(OptionsSize == sizeof(FDRLoggingOptions));
-  XRayLogInitStatus CurrentStatus = XRayLogInitStatus::XRAY_LOG_UNINITIALIZED;
-  if (!LoggingStatus.compare_exchange_strong(
-          CurrentStatus, XRayLogInitStatus::XRAY_LOG_INITIALIZING,
-          std::memory_order_release, std::memory_order_relaxed))
-    return CurrentStatus;
+  s32 CurrentStatus = XRayLogInitStatus::XRAY_LOG_UNINITIALIZED;
+  if (__sanitizer::atomic_compare_exchange_strong(
+          &LoggingStatus, &CurrentStatus,
+          XRayLogInitStatus::XRAY_LOG_INITIALIZING,
+          __sanitizer::memory_order_release))
+    return static_cast<XRayLogInitStatus>(CurrentStatus);
 
   FDROptions.reset(new FDRLoggingOptions());
   *FDROptions = *reinterpret_cast<FDRLoggingOptions *>(Options);
   if (FDROptions->ReportErrors)
     SetPrintfAndReportCallback(printToStdErr);
 
   bool Success = false;
   BQ = std::make_shared<BufferQueue>(BufferSize, BufferMax, Success);
   if (!Success) {
     Report("BufferQueue init failed.\n");
     return XRayLogInitStatus::XRAY_LOG_UNINITIALIZED;
   }
 
   // Install the actual handleArg0 handler after initialising the buffers.
   __xray_set_handler(fdrLoggingHandleArg0);
 
-  LoggingStatus.store(XRayLogInitStatus::XRAY_LOG_INITIALIZED,
-                      std::memory_order_release);
+  __sanitizer::atomic_store(&LoggingStatus,
+                            XRayLogInitStatus::XRAY_LOG_INITIALIZED,
+                            __sanitizer::memory_order_release);
   return XRayLogInitStatus::XRAY_LOG_INITIALIZED;
 }
 
 // Must finalize before flushing.
 XRayLogFlushStatus fdrLoggingFlush() XRAY_NEVER_INSTRUMENT {
-  if (LoggingStatus.load(std::memory_order_acquire) !=
+  if (__sanitizer::atomic_load(&LoggingStatus,
+                               __sanitizer::memory_order_acquire) !=
       XRayLogInitStatus::XRAY_LOG_FINALIZED)
     return XRayLogFlushStatus::XRAY_LOG_NOT_FLUSHING;
 
-  XRayLogFlushStatus Result = XRayLogFlushStatus::XRAY_LOG_NOT_FLUSHING;
-  if (!LogFlushStatus.compare_exchange_strong(
-          Result, XRayLogFlushStatus::XRAY_LOG_FLUSHING,
-          std::memory_order_release, std::memory_order_relaxed))
-    return Result;
+  s32 Result = XRayLogFlushStatus::XRAY_LOG_NOT_FLUSHING;
+  if (__sanitizer::atomic_compare_exchange_strong(
+          &LogFlushStatus, &Result, XRayLogFlushStatus::XRAY_LOG_FLUSHING,
+          __sanitizer::memory_order_release))
+    return static_cast<XRayLogFlushStatus>(Result);
 
   // Make a copy of the BufferQueue pointer to prevent other threads that may be
   // resetting it from blowing away the queue prematurely while we're dealing
   // with it.
   auto LocalBQ = BQ;
 
   // We write out the file in the following format:
   //
   //   1) We write down the XRay file header with version 1, type FDR_LOG.
   //   2) Then we use the 'apply' member of the BufferQueue that's live, to
   //      ensure that at this point in time we write down the buffers that have
   //      been released (and marked "used") -- we dump the full buffer for now
   //      (fixed-sized) and let the tools reading the buffers deal with the data
   //      afterwards.
   //
   int Fd = FDROptions->Fd;
   if (Fd == -1)
     Fd = getLogFD();
   if (Fd == -1) {
     auto Result = XRayLogFlushStatus::XRAY_LOG_NOT_FLUSHING;
-    LogFlushStatus.store(Result, std::memory_order_release);
+    __sanitizer::atomic_store(&LogFlushStatus, Result,
+                              __sanitizer::memory_order_release);
     return Result;
   }
 
   XRayFileHeader Header;
   Header.Version = 1;
   Header.Type = FileTypes::FDR_LOG;
   Header.CycleFrequency = getTSCFrequency();
   // FIXME: Actually check whether we have 'constant_tsc' and 'nonstop_tsc'
   // before setting the values in the header.
   Header.ConstantTSC = 1;
   Header.NonstopTSC = 1;
   clock_gettime(CLOCK_REALTIME, &Header.TS);
   retryingWriteAll(Fd, reinterpret_cast<char *>(&Header),
                    reinterpret_cast<char *>(&Header) + sizeof(Header));
   LocalBQ->apply([&](const BufferQueue::Buffer &B) {
     retryingWriteAll(Fd, reinterpret_cast<char *>(B.Buffer),
                      reinterpret_cast<char *>(B.Buffer) + B.Size);
   });
-  LogFlushStatus.store(XRayLogFlushStatus::XRAY_LOG_FLUSHED,
-                       std::memory_order_release);
+  __sanitizer::atomic_store(&LogFlushStatus,
+                            XRayLogFlushStatus::XRAY_LOG_FLUSHED,
+                            __sanitizer::memory_order_release);
   return XRayLogFlushStatus::XRAY_LOG_FLUSHED;
 }
 
 XRayLogInitStatus fdrLoggingFinalize() XRAY_NEVER_INSTRUMENT {
-  XRayLogInitStatus CurrentStatus = XRayLogInitStatus::XRAY_LOG_INITIALIZED;
-  if (!LoggingStatus.compare_exchange_strong(
-          CurrentStatus, XRayLogInitStatus::XRAY_LOG_FINALIZING,
-          std::memory_order_release, std::memory_order_relaxed))
-    return CurrentStatus;
+  s32 CurrentStatus = XRayLogInitStatus::XRAY_LOG_INITIALIZED;
+  if (__sanitizer::atomic_compare_exchange_strong(
+          &LoggingStatus, &CurrentStatus,
+          XRayLogInitStatus::XRAY_LOG_FINALIZING,
+          __sanitizer::memory_order_release))
+    return static_cast<XRayLogInitStatus>(CurrentStatus);
 
   // Do special things to make the log finalize itself, and not allow any more
   // operations to be performed until re-initialized.
   BQ->finalize();
 
-  LoggingStatus.store(XRayLogInitStatus::XRAY_LOG_FINALIZED,
-                      std::memory_order_release);
+  __sanitizer::atomic_store(&LoggingStatus,
+                            XRayLogInitStatus::XRAY_LOG_FINALIZED,
+                            __sanitizer::memory_order_release);
   return XRayLogInitStatus::XRAY_LOG_FINALIZED;
 }
 
 XRayLogInitStatus fdrLoggingReset() XRAY_NEVER_INSTRUMENT {
-  XRayLogInitStatus CurrentStatus = XRayLogInitStatus::XRAY_LOG_FINALIZED;
-  if (!LoggingStatus.compare_exchange_strong(
-          CurrentStatus, XRayLogInitStatus::XRAY_LOG_UNINITIALIZED,
-          std::memory_order_release, std::memory_order_relaxed))
-    return CurrentStatus;
+  s32 CurrentStatus = XRayLogInitStatus::XRAY_LOG_FINALIZED;
+  if (__sanitizer::atomic_compare_exchange_strong(
+          &LoggingStatus, &CurrentStatus,
+          XRayLogInitStatus::XRAY_LOG_INITIALIZED,
+          __sanitizer::memory_order_release))
+    return static_cast<XRayLogInitStatus>(CurrentStatus);
 
   // Release the in-memory buffer queue.
   BQ.reset();
 
   // Spin until the flushing status is flushed.
-  XRayLogFlushStatus CurrentFlushingStatus =
+  s32 CurrentFlushingStatus =
       XRayLogFlushStatus::XRAY_LOG_FLUSHED;
-  while (!LogFlushStatus.compare_exchange_weak(
-      CurrentFlushingStatus, XRayLogFlushStatus::XRAY_LOG_NOT_FLUSHING,
-      std::memory_order_release, std::memory_order_relaxed)) {
+  while (__sanitizer::atomic_compare_exchange_weak(
+      &LogFlushStatus, &CurrentFlushingStatus,
+      XRayLogFlushStatus::XRAY_LOG_NOT_FLUSHING,
+      __sanitizer::memory_order_release)) {
     if (CurrentFlushingStatus == XRayLogFlushStatus::XRAY_LOG_NOT_FLUSHING)
       break;
     CurrentFlushingStatus = XRayLogFlushStatus::XRAY_LOG_FLUSHED;
   }
 
   // At this point, we know that the status is flushed, and that we can assume
   return XRayLogInitStatus::XRAY_LOG_UNINITIALIZED;
 }

lib/xray/xray_fdr_logging_impl.h

 //
 //===----------------------------------------------------------------------===//
 #ifndef XRAY_XRAY_FDR_LOGGING_IMPL_H
 #define XRAY_XRAY_FDR_LOGGING_IMPL_H
 
 #include <cassert>
 #include <cstdint>
 #include <cstring>
+#include <limits>
 #include <memory>
 #include <string>
 #include <sys/syscall.h>
 #include <unistd.h>
 
 #include "sanitizer_common/sanitizer_common.h"
 #include "xray/xray_log_interface.h"
 #include "xray_buffer_queue.h"

 static void writeTSCWrapMetadata(uint64_t TSC);
 
 /// Here's where the meat of the processing happens. The writer captures
 /// function entry, exit and tail exit points with a time and will create
 /// TSCWrap, NewCPUId and Function records as necessary. The writer might
 /// walk backward through its buffer and erase trivial functions to avoid
 /// polluting the log and may use the buffer queue to obtain or release a
 /// buffer.
-static void
-processFunctionHook(int32_t FuncId, XRayEntryType Entry, uint64_t TSC,
-                    unsigned char CPU,
-                    int (*wall_clock_reader)(clockid_t, struct timespec *),
-                    const std::atomic<XRayLogInitStatus> &LoggingStatus,
+static void processFunctionHook(int32_t FuncId, XRayEntryType Entry,
+                                uint64_t TSC, unsigned char CPU,
+                                int (*wall_clock_reader)(clockid_t,
+                                                         struct timespec *),
+                                __sanitizer::atomic_sint32_t &LoggingStatus,
                                 const std::shared_ptr<BufferQueue> &BQ);
 
 //-----------------------------------------------------------------------------|
 // The rest of the file is implementation.                                     |
 //-----------------------------------------------------------------------------|
 // Functions are implemented in the header for inlining since we don't want    |
 // to grow the stack when we've hijacked the binary for logging.               |
 //-----------------------------------------------------------------------------|
 
public:
 
   ~RecursionGuard() noexcept {
     if (Valid)
       Running = false;
   }
 };
 
 static inline bool loggingInitialized(
-    const std::atomic<XRayLogInitStatus> &LoggingStatus) XRAY_NEVER_INSTRUMENT {
-  return LoggingStatus.load(std::memory_order_acquire) ==
+    const __sanitizer::atomic_sint32_t &LoggingStatus) XRAY_NEVER_INSTRUMENT {
+  return __sanitizer::atomic_load(&LoggingStatus,
+                                  __sanitizer::memory_order_acquire) ==
          XRayLogInitStatus::XRAY_LOG_INITIALIZED;
 }
 
 } // namespace
 
 static inline void writeNewBufferPreamble(pid_t Tid, timespec TS,
                                           char *&MemPtr) XRAY_NEVER_INSTRUMENT {
   static constexpr int InitRecordsCount = 2;
static inline void writeFunctionRecord(int FuncId, uint32_t TSCDelta,
 
   std::memcpy(MemPtr, &AlignedFuncRecordBuffer, sizeof(FunctionRecord));
   MemPtr += sizeof(FunctionRecord);
 }
 
 static inline void processFunctionHook(
     int32_t FuncId, XRayEntryType Entry, uint64_t TSC, unsigned char CPU,
     int (*wall_clock_reader)(clockid_t, struct timespec *),
-    const std::atomic<XRayLogInitStatus> &LoggingStatus,
+    __sanitizer::atomic_sint32_t &LoggingStatus,
     const std::shared_ptr<BufferQueue> &BQ) XRAY_NEVER_INSTRUMENT {
   // Bail out right away if logging is not initialized yet.
-  if (LoggingStatus.load(std::memory_order_acquire) !=
+  if (__sanitizer::atomic_load(&LoggingStatus,
+                               __sanitizer::memory_order_acquire) !=
       XRayLogInitStatus::XRAY_LOG_INITIALIZED)
     return;
 
   // We use a thread_local variable to keep track of which CPUs we've already
   // run, and the TSC times for these CPUs. This allows us to stop repeating the
   // CPU field in the function records.
   //
   // We assume that we'll support only 65536 CPUs for x86_64.
if (EC != BufferQueue::ErrorCode::Ok) {
       return;
     }
     RecordPtr = nullptr;
   }
 
   if (Buffer.Buffer == nullptr) {
     auto EC = LocalBQ->getBuffer(Buffer);
     if (EC != BufferQueue::ErrorCode::Ok) {
-      auto LS = LoggingStatus.load(std::memory_order_acquire);
+      auto LS = __sanitizer::atomic_load(&LoggingStatus,
+                                         __sanitizer::memory_order_acquire);
       if (LS != XRayLogInitStatus::XRAY_LOG_FINALIZING &&
           LS != XRayLogInitStatus::XRAY_LOG_FINALIZED)
         Report("Failed to acquire a buffer; error=%s\n",
                BufferQueue::getErrorString(EC));
       return;
     }
 
     setupNewBuffer(Buffer, wall_clock_reader);

lib/xray/xray_init.cc

 //===-- xray_init.cc --------------------------------------------*- C++ -*-===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 //
 // This file is a part of XRay, a dynamic runtime instrumentation system.
 //
 // XRay initialisation logic.
 //===----------------------------------------------------------------------===//
 
-#include <atomic>
 #include <fcntl.h>
 #include <strings.h>
 #include <unistd.h>
 
 #include "sanitizer_common/sanitizer_common.h"
 #include "xray_defs.h"
 #include "xray_flags.h"
 #include "xray_interface_internal.h"
 
 extern "C" {
 void __xray_init();
 extern const XRaySledEntry __start_xray_instr_map[] __attribute__((weak));
 extern const XRaySledEntry __stop_xray_instr_map[] __attribute__((weak));
 }
 
-using namespace __sanitizer;
 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.
-std::atomic<bool> XRayInitialized{false};
+__sanitizer::atomic_uint8_t XRayInitialized{0};
 
 // This should always be updated before XRayInitialized is updated.
-std::atomic<__xray::XRaySledMap> XRayInstrMap{};
+__sanitizer::SpinMutex XRayInstrMapMutex;
+XRaySledMap XRayInstrMap;
 
 // __xray_init() will do the actual loading of the current process' memory map
 // and then proceed to look for the .xray_instr_map section/segment.
 void __xray_init() XRAY_NEVER_INSTRUMENT {
   initializeFlags();
   if (__start_xray_instr_map == nullptr) {
     Report("XRay instrumentation map missing. Not initializing XRay.\n");
     return;
   }
 
-  // Now initialize the XRayInstrMap global struct with the address of the
-  // entries, reinterpreted as an array of XRaySledEntry objects. We use the
-  // virtual pointer we have from the section to provide us the correct
-  // information.
-  __xray::XRaySledMap SledMap{};
-  SledMap.Sleds = __start_xray_instr_map;
-  SledMap.Entries = __stop_xray_instr_map - __start_xray_instr_map;
-  XRayInstrMap.store(SledMap, std::memory_order_release);
-  XRayInitialized.store(true, std::memory_order_release);
+  {
+    __sanitizer::SpinMutexLock Guard(&XRayInstrMapMutex);
+    XRayInstrMap.Sleds = __start_xray_instr_map;
+    XRayInstrMap.Entries = __stop_xray_instr_map - __start_xray_instr_map;
+  }
+  __sanitizer::atomic_store(&XRayInitialized, true,
+                            __sanitizer::memory_order_release);
 
   if (flags()->patch_premain)
     __xray_patch();
 }
 
 __attribute__((section(".preinit_array"),
                used)) void (*__local_xray_preinit)(void) = __xray_init;

lib/xray/xray_interface.cc

 // This file is a part of XRay, a dynamic runtime instrumentation system.
 //
 // Implementation of the API functions.
 //
 //===----------------------------------------------------------------------===//
 
 #include "xray_interface_internal.h"
 
-#include <atomic>
 #include <cstdint>
 #include <cstdio>
 #include <errno.h>
 #include <limits>
 #include <sys/mman.h>
 
 #include "sanitizer_common/sanitizer_common.h"
 #include "xray_defs.h"

 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.
-std::atomic<void (*)(int32_t, XRayEntryType, uint64_t)> XRayArgLogger{nullptr};
+__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
 // and executable, and upon destruction if it was successful in doing so returns
 // the page into a read-only and executable page.
 //
 // This is only used specifically for runtime-patching of the XRay
 // instrumentation points. This assumes that the executable pages are originally
~MProtectHelper() XRAY_NEVER_INSTRUMENT {
     if (MustCleanup) {
       mprotect(PageAlignedAddr, MProtectLen, PROT_READ | PROT_EXEC);
     }
   }
 };
 
 } // 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;
 }
 
 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;
 
 // FIXME: Figure out whether we can move this class to sanitizer_common instead
 // as a generic "scope guard".
 template <class Function> class CleanupInvoker {
   Function Fn;
 

 CleanupInvoker<Function> scopeCleanup(Function Fn) XRAY_NEVER_INSTRUMENT {
   return CleanupInvoker<Function>{Fn};
 }
 
 // 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))
+  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;
+  uint8_t PatchingSuccess = false;
   auto XRayPatchingStatusResetter = scopeCleanup([&PatchingSuccess] {
-    if (!PatchingSuccess) {
-      XRayPatching.store(false, std::memory_order_release);
-    }
+    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;
 
   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;
   }
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);
 }
 
 XRayPatchingStatus __xray_unpatch() XRAY_NEVER_INSTRUMENT {
   return controlPatching(false);
 }
 
-int __xray_set_handler_arg1(void (*Handler)(int32_t, XRayEntryType, uint64_t))
-{
-  if (!XRayInitialized.load(std::memory_order_acquire)) {
+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.
-  XRayArgLogger.store(Handler, std::memory_order_release);
+  __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); }

lib/xray/xray_log_interface.cc

 //===-- xray_log_interface.cc ---------------------------------------------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 //
 // This file is a part of XRay, a function call tracing system.
 //
 //===----------------------------------------------------------------------===//
 #include "xray/xray_log_interface.h"
 
+#include "sanitizer_common/sanitizer_atomic.h"
+#include "sanitizer_common/sanitizer_mutex.h"
 #include "xray/xray_interface.h"
 #include "xray_defs.h"
 
 #include <memory>
-#include <mutex>
 
-std::mutex XRayImplMutex;
+__sanitizer::SpinMutex XRayImplMutex;
 std::unique_ptr<XRayLogImpl> GlobalXRayImpl;
 
 void __xray_set_log_impl(XRayLogImpl Impl) XRAY_NEVER_INSTRUMENT {
   if (Impl.log_init == nullptr || Impl.log_finalize == nullptr ||
       Impl.handle_arg0 == nullptr || Impl.flush_log == nullptr) {
-    std::lock_guard<std::mutex> Guard(XRayImplMutex);
+    __sanitizer::SpinMutexLock Guard(&XRayImplMutex);
     GlobalXRayImpl.reset();
     return;
   }
 
-  std::lock_guard<std::mutex> Guard(XRayImplMutex);
+  __sanitizer::SpinMutexLock Guard(&XRayImplMutex);
   GlobalXRayImpl.reset(new XRayLogImpl);
   *GlobalXRayImpl = Impl;
 }
 
 XRayLogInitStatus __xray_init(size_t BufferSize, size_t MaxBuffers, void *Args,
                               size_t ArgsSize) XRAY_NEVER_INSTRUMENT {
-  std::lock_guard<std::mutex> Guard(XRayImplMutex);
+  __sanitizer::SpinMutexLock Guard(&XRayImplMutex);
   if (!GlobalXRayImpl)
     return XRayLogInitStatus::XRAY_LOG_UNINITIALIZED;
   return GlobalXRayImpl->log_init(BufferSize, MaxBuffers, Args, ArgsSize);
 }
 
 XRayLogInitStatus __xray_log_finalize() XRAY_NEVER_INSTRUMENT {
-  std::lock_guard<std::mutex> Guard(XRayImplMutex);
+  __sanitizer::SpinMutexLock Guard(&XRayImplMutex);
   if (!GlobalXRayImpl)
     return XRayLogInitStatus::XRAY_LOG_UNINITIALIZED;
   return GlobalXRayImpl->log_finalize();
 }
 
 XRayLogFlushStatus __xray_log_flushLog() XRAY_NEVER_INSTRUMENT {
-  std::lock_guard<std::mutex> Guard(XRayImplMutex);
+  __sanitizer::SpinMutexLock Guard(&XRayImplMutex);
   if (!GlobalXRayImpl)
     return XRayLogFlushStatus::XRAY_LOG_NOT_FLUSHING;
   return GlobalXRayImpl->flush_log();
 }