[XRay][compiler-rt] Runtime changes to support custom event logging

[dberris]

Summary

This change implements support for the custom event logging sleds and
intrinsics at runtime. For now it only supports handling the sleds in
x86_64, with the implementations for other architectures stubbed out to
do nothing.

Depends on D27503, D30018, and D33032.

[dberris]

Updated to work with most recent changes in D27503. This is now ready for a look.

[timshen]

lib/xray/xray_fdr_logging.cc

508 ↗

(On Diff #91275)

Do you want to assert here?

lib/xray/xray_trampoline_x86_64.S

205 ↗

(On Diff #91275)

Why stash rbp, if we don't clobber it?

[dberris]

lib/xray/xray_fdr_logging.cc

508 ↗

(On Diff #91275)

I thought about it, but then figured it would be bad form to be crashing the process if there was a runtime implementation bug. Is it OK to have assertions in runtime implementations?

I thought about using llvm_unreachable(...) but wasn't sure that was available in compiler-rt.

lib/xray/xray_trampoline_x86_64.S

205 ↗

(On Diff #91275)

I can't fully explain it, but if I don't stash this (seems to be the frame pointer) then I've run into issues in the past. I can try removing it, and see how that goes.

[dberris]

Rebase and add a bit more refactoring to share some core parts of the logic between custom event handling and function entry/exit handling in FDR mode.

[dberris]

Still a work in progress, need to add a bit more testing and ensure that we're able to handle the records in the log.

[dberris]

Ready for a look now.

[timshen]

lib/xray/xray_fdr_logging.cc

202 ↗

(On Diff #98411)

Running &&

206 ↗

(On Diff #98411)

static.

Also, it doesn't have to be a bool. An empty struct or tuple<> will do.

Please also check other places.

235 ↗

(On Diff #98411)

s/sizeof(uint64_t)/64/? By definition. :)

257 ↗

(On Diff #98411)

Does FDROptions really need to be a unique_ptr? How about just using the value type?

lib/xray/xray_fdr_logging_impl.h

464 ↗

(On Diff #98411)

I don't usually see "static inline" around. Usually static defeats the purpose of inline. I suggest to remove static.

lib/xray/xray_trampoline_x86_64.S

205 ↗

(On Diff #91275)

Do you have any clue about the rbp? I believe that rbp is caller-saved, therefore if __xray_CustomEvent doesn't use it, it doesn't have to save it.

[dberris]

Address comments, some cleanup.

lib/xray/xray_fdr_logging.cc

235 ↗

(On Diff #98411)

Actually, sizeof(uint64_t) is 8. However, we're doing this to try and match what the type of the element is. I've made it a local constant instead, to give it a name.

257 ↗

(On Diff #98411)

Good point. Doesn't need to be a unique_ptr. Turned it into a global instead, protected with a spinlock.

lib/xray/xray_trampoline_x86_64.S

205 ↗

(On Diff #91275)

So I looked this up, and it seems that the frame pointer in x86_64 is usually stored in %rbp -- and if I don't put this on the stack, we run into issues locating the frame pointer when unwinding and debugging. If the frame pointer is potentially omitted, we can't know so we're being safe anyway in this case by stashing it regardless.

Does this make sense?

[timshen]

Other than a typo, LGTM. I believe that the atomics are correct as well, especially in the assembly - x86 load-acquire is indeed a no-op.

lib/xray/xray_fdr_logging_impl.h

483 ↗

(On Diff #98579)

s/Redy/Ready/ :)

lib/xray/xray_trampoline_x86_64.S

205 ↗

(On Diff #91275)

Yes! This reminds me of -fomit-frame-pointer.

I think it's good save rbp.

[dberris]

• fixup: typo

compiler-rt/trunk/include/xray/xray_interface.h

 extern "C" {
 
 /// Synchronize this with AsmPrinter::SledKind in LLVM.
 enum XRayEntryType {
   ENTRY = 0,
   EXIT = 1,
   TAIL = 2,
   LOG_ARGS_ENTRY = 3,
+  CUSTOM_EVENT = 4,
 };
 
 /// Provide a function to invoke for when instrumentation points are hit. This
 /// is a user-visible control surface that overrides the default implementation.
 /// The function provided should take the following arguments:
 ///
 ///   - function id: an identifier that indicates the id of a function; this id
 ///                  is generated by xray; the mapping between the function id

 ///
 /// Returns 1 on success, 0 on error.
 extern int __xray_set_handler_arg1(void (*)(int32_t, XRayEntryType, uint64_t));
 
 /// Disables the XRay handler used to log first arguments of function calls.
 /// Returns 1 on success, 0 on error.
 extern int __xray_remove_handler_arg1();
 
+/// Provide a function to invoke when XRay encounters a custom event.
+extern int __xray_set_customevent_handler(void (*entry)(void*, std::size_t));
+
 enum XRayPatchingStatus {
   NOT_INITIALIZED = 0,
   SUCCESS = 1,
   ONGOING = 2,
   FAILED = 3,
 };
 
 /// This tells XRay to patch the instrumentation points. See XRayPatchingStatus

compiler-rt/trunk/lib/xray/xray_AArch64.cc

 //
 //===----------------------------------------------------------------------===//
 #include "sanitizer_common/sanitizer_common.h"
 #include "xray_defs.h"
 #include "xray_interface_internal.h"
 #include <atomic>
 #include <cassert>
 
-
-extern "C" void __clear_cache(void* start, void* end);
+extern "C" void __clear_cache(void *start, void *end);
 
 namespace __xray {
 
 // The machine codes for some instructions used in runtime patching.
 enum class PatchOpcodes : uint32_t {
   PO_StpX0X30SP_m16e = 0xA9BF7BE0, // STP X0, X30, [SP, #-16]!
   PO_LdrW0_12 = 0x18000060,        // LDR W0, #12
   PO_LdrX16_12 = 0x58000070,       // LDR X16, #12
if (Enable) {
     std::atomic_store_explicit(
         reinterpret_cast<std::atomic<uint32_t> *>(FirstAddress),
         uint32_t(PatchOpcodes::PO_StpX0X30SP_m16e), std::memory_order_release);
   } else {
     std::atomic_store_explicit(
         reinterpret_cast<std::atomic<uint32_t> *>(FirstAddress),
         uint32_t(PatchOpcodes::PO_B32), std::memory_order_release);
   }
-  __clear_cache(reinterpret_cast<char*>(FirstAddress),
-      reinterpret_cast<char*>(CurAddress));
+  __clear_cache(reinterpret_cast<char *>(FirstAddress),
+                reinterpret_cast<char *>(CurAddress));
   return true;
 }
 
 bool patchFunctionEntry(const bool Enable, const uint32_t FuncId,
                         const XRaySledEntry &Sled,
                         void (*Trampoline)()) XRAY_NEVER_INSTRUMENT {
   return patchSled(Enable, FuncId, Sled, Trampoline);
 }
 
 bool patchFunctionExit(const bool Enable, const uint32_t FuncId,
                        const XRaySledEntry &Sled) XRAY_NEVER_INSTRUMENT {
   return patchSled(Enable, FuncId, Sled, __xray_FunctionExit);
 }
 
 bool patchFunctionTailExit(const bool Enable, const uint32_t FuncId,
                            const XRaySledEntry &Sled) XRAY_NEVER_INSTRUMENT {
   return patchSled(Enable, FuncId, Sled, __xray_FunctionTailExit);
 }
 
+bool patchCustomEvent(const bool Enable, const uint32_t FuncId,
+                      const XRaySled &Sled)
+    XRAY_NEVER_INSTRUMENT { // FIXME: Implement in aarch64?
+  return false;
+}
+
 // FIXME: Maybe implement this better?
 bool probeRequiredCPUFeatures() XRAY_NEVER_INSTRUMENT { return true; }
 
 } // namespace __xray
 
 extern "C" void __xray_ArgLoggerEntry() XRAY_NEVER_INSTRUMENT {
   // FIXME: this will have to be implemented in the trampoline assembly file
 }

compiler-rt/trunk/lib/xray/xray_arm.cc

 //
 //===----------------------------------------------------------------------===//
 #include "sanitizer_common/sanitizer_common.h"
 #include "xray_defs.h"
 #include "xray_interface_internal.h"
 #include <atomic>
 #include <cassert>
 
-extern "C" void __clear_cache(void* start, void* end);
+extern "C" void __clear_cache(void *start, void *end);
 
 namespace __xray {
 
 // The machine codes for some instructions used in runtime patching.
 enum class PatchOpcodes : uint32_t {
   PO_PushR0Lr = 0xE92D4001, // PUSH {r0, lr}
   PO_BlxIp = 0xE12FFF3C,    // BLX ip
   PO_PopR0Lr = 0xE8BD4001,  // POP {r0, lr}
if (Enable) {
     std::atomic_store_explicit(
         reinterpret_cast<std::atomic<uint32_t> *>(FirstAddress),
         uint32_t(PatchOpcodes::PO_PushR0Lr), std::memory_order_release);
   } else {
     std::atomic_store_explicit(
         reinterpret_cast<std::atomic<uint32_t> *>(FirstAddress),
         uint32_t(PatchOpcodes::PO_B20), std::memory_order_release);
   }
-  __clear_cache(reinterpret_cast<char*>(FirstAddress),
-      reinterpret_cast<char*>(CurAddress));
+  __clear_cache(reinterpret_cast<char *>(FirstAddress),
+                reinterpret_cast<char *>(CurAddress));
   return true;
 }
 
 bool patchFunctionEntry(const bool Enable, const uint32_t FuncId,
                         const XRaySledEntry &Sled,
                         void (*Trampoline)()) XRAY_NEVER_INSTRUMENT {
   return patchSled(Enable, FuncId, Sled, Trampoline);
 }
 
 bool patchFunctionExit(const bool Enable, const uint32_t FuncId,
                        const XRaySledEntry &Sled) XRAY_NEVER_INSTRUMENT {
   return patchSled(Enable, FuncId, Sled, __xray_FunctionExit);
 }
 
 bool patchFunctionTailExit(const bool Enable, const uint32_t FuncId,
                            const XRaySledEntry &Sled) XRAY_NEVER_INSTRUMENT {
   return patchSled(Enable, FuncId, Sled, __xray_FunctionTailExit);
 }
 
+bool patchCustomEvent(const bool Enable, const uint32_t FuncId,
+                      const XRaySled &Sled)
+    XRAY_NEVER_INSTRUMENT { // FIXME: Implement in arm?
+  return false;
+}
+
 // FIXME: Maybe implement this better?
 bool probeRequiredCPUFeatures() XRAY_NEVER_INSTRUMENT { return true; }
 
 } // namespace __xray
 
 extern "C" void __xray_ArgLoggerEntry() XRAY_NEVER_INSTRUMENT {
   // FIXME: this will have to be implemented in the trampoline assembly file
 }

compiler-rt/trunk/lib/xray/xray_fdr_log_records.h

struct alignas(16) MetadataRecord {
   // Each kind of record is represented as a 7-bit value (even though we use an
   // unsigned 8-bit enum class to do so).
   enum class RecordKinds : uint8_t {
     NewBuffer,
     EndOfBuffer,
     NewCPUId,
     TSCWrap,
     WalltimeMarker,
+    CustomEventMarker,
   };
   // Use 7 bits to identify this record type.
   /* RecordKinds */ uint8_t RecordKind : 7;
   char Data[15];
 } __attribute__((packed));
 
 static_assert(sizeof(MetadataRecord) == 16, "Wrong size for MetadataRecord.");
 

compiler-rt/trunk/lib/xray/xray_fdr_logging.cc

 #include "xray_tsc.h"
 #include "xray_utils.h"
 
 namespace __xray {
 
 // Global BufferQueue.
 std::shared_ptr<BufferQueue> BQ;
 
-__sanitizer::atomic_sint32_t LoggingStatus = {
-    XRayLogInitStatus::XRAY_LOG_UNINITIALIZED};
-
 __sanitizer::atomic_sint32_t LogFlushStatus = {
     XRayLogFlushStatus::XRAY_LOG_NOT_FLUSHING};
 
-std::unique_ptr<FDRLoggingOptions> FDROptions;
+FDRLoggingOptions FDROptions;
 
-XRayLogInitStatus fdrLoggingInit(std::size_t BufferSize, std::size_t BufferMax,
-                                 void *Options,
-                                 size_t OptionsSize) XRAY_NEVER_INSTRUMENT {
-  if (OptionsSize != sizeof(FDRLoggingOptions))
-    return static_cast<XRayLogInitStatus>(__sanitizer::atomic_load(
-        &LoggingStatus, __sanitizer::memory_order_acquire));
-  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());
-  memcpy(FDROptions.get(), Options, OptionsSize);
-  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);
-
-  __sanitizer::atomic_store(&LoggingStatus,
-                            XRayLogInitStatus::XRAY_LOG_INITIALIZED,
-                            __sanitizer::memory_order_release);
-  Report("XRay FDR init successful.\n");
-  return XRayLogInitStatus::XRAY_LOG_INITIALIZED;
-}
+__sanitizer::SpinMutex FDROptionsMutex;
 
 // Must finalize before flushing.
 XRayLogFlushStatus fdrLoggingFlush() XRAY_NEVER_INSTRUMENT {
   if (__sanitizer::atomic_load(&LoggingStatus,
                                __sanitizer::memory_order_acquire) !=
       XRayLogInitStatus::XRAY_LOG_FINALIZED)
     return XRayLogFlushStatus::XRAY_LOG_NOT_FLUSHING;
 
XRayLogFlushStatus fdrLoggingFlush() XRAY_NEVER_INSTRUMENT {
   //
   //   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;
+  int Fd = -1;
+  {
+    __sanitizer::SpinMutexLock Guard(&FDROptionsMutex);
+    Fd = FDROptions.Fd;
+  }
   if (Fd == -1)
     Fd = getLogFD();
   if (Fd == -1) {
     auto Result = XRayLogFlushStatus::XRAY_LOG_NOT_FLUSHING;
     __sanitizer::atomic_store(&LogFlushStatus, Result,
                               __sanitizer::memory_order_release);
     return Result;
   }
 
   // Test for required CPU features and cache the cycle frequency
   static bool TSCSupported = probeRequiredCPUFeatures();
-  static uint64_t CycleFrequency = TSCSupported ? getTSCFrequency()
-                                   : __xray::NanosecondsPerSecond;
+  static uint64_t CycleFrequency =
+      TSCSupported ? getTSCFrequency() : __xray::NanosecondsPerSecond;
 
   XRayFileHeader Header;
   Header.Version = 1;
   Header.Type = FileTypes::FDR_LOG;
   Header.CycleFrequency = CycleFrequency;
   // FIXME: Actually check whether we have 'constant_tsc' and 'nonstop_tsc'
   // before setting the values in the header.
   Header.ConstantTSC = 1;
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;
 }
 
-void fdrLoggingHandleArg0(int32_t FuncId,
-                          XRayEntryType Entry) XRAY_NEVER_INSTRUMENT {
+static std::tuple<uint64_t, unsigned char>
+getTimestamp() XRAY_NEVER_INSTRUMENT {
   // We want to get the TSC as early as possible, so that we can check whether
   // we've seen this CPU before. We also do it before we load anything else, to
   // allow for forward progress with the scheduling.
   unsigned char CPU;
   uint64_t TSC;
 
   // Test once for required CPU features
   static bool TSCSupported = probeRequiredCPUFeatures();
 
-  if(TSCSupported) {
+  if (TSCSupported) {
     TSC = __xray::readTSC(CPU);
   } else {
     // FIXME: This code needs refactoring as it appears in multiple locations
     timespec TS;
     int result = clock_gettime(CLOCK_REALTIME, &TS);
     if (result != 0) {
       Report("clock_gettime(2) return %d, errno=%d", result, int(errno));
       TS = {0, 0};
     }
     CPU = 0;
     TSC = TS.tv_sec * __xray::NanosecondsPerSecond + TS.tv_nsec;
   }
+  return std::make_tuple(TSC, CPU);
+}
+
+void fdrLoggingHandleArg0(int32_t FuncId,
+                          XRayEntryType Entry) XRAY_NEVER_INSTRUMENT {
+  auto TSC_CPU = getTimestamp();
+  __xray_fdr_internal::processFunctionHook(FuncId, Entry, std::get<0>(TSC_CPU),
+                                           std::get<1>(TSC_CPU), clock_gettime,
+                                           LoggingStatus, BQ);
+}
+
+void fdrLoggingHandleCustomEvent(void *Event,
+                                 std::size_t EventSize) XRAY_NEVER_INSTRUMENT {
+  using namespace __xray_fdr_internal;
+  auto TSC_CPU = getTimestamp();
+  auto &TSC = std::get<0>(TSC_CPU);
+  auto &CPU = std::get<1>(TSC_CPU);
+  thread_local bool Running = false;
+  RecursionGuard Guard{Running};
+  if (!Guard) {
+    assert(Running && "RecursionGuard is buggy!");
+    return;
+  }
+  if (EventSize > std::numeric_limits<int32_t>::max()) {
+    using Empty = struct {};
+    static Empty Once = [&] {
+      Report("Event size too large = %zu ; > max = %d\n", EventSize,
+             std::numeric_limits<int32_t>::max());
+      return Empty();
+    }();
+    (void)Once;
+  }
+  int32_t ReducedEventSize = static_cast<int32_t>(EventSize);
+  if (!isLogInitializedAndReady(LocalBQ, TSC, CPU, clock_gettime))
+    return;
+
+  // Here we need to prepare the log to handle:
+  //   - The metadata record we're going to write. (16 bytes)
+  //   - The additional data we're going to write. Currently, that's the size of
+  //   the event we're going to dump into the log as free-form bytes.
+  if (!prepareBuffer(clock_gettime, MetadataRecSize + EventSize)) {
+    LocalBQ = nullptr;
+    return;
+  }
+
+  // Write the custom event metadata record, which consists of the following
+  // information:
+  //   - 8 bytes (64-bits) for the full TSC when the event started.
+  //   - 4 bytes (32-bits) for the length of the data.
+  MetadataRecord CustomEvent;
+  CustomEvent.Type = uint8_t(RecordType::Metadata);
+  CustomEvent.RecordKind =
+      uint8_t(MetadataRecord::RecordKinds::CustomEventMarker);
+  constexpr auto TSCSize = sizeof(std::get<0>(TSC_CPU));
+  std::memcpy(&CustomEvent.Data, &ReducedEventSize, sizeof(int32_t));
+  std::memcpy(&CustomEvent.Data[sizeof(int32_t)], &TSC, TSCSize);
+  std::memcpy(RecordPtr, &CustomEvent, sizeof(CustomEvent));
+  RecordPtr += sizeof(CustomEvent);
+  std::memcpy(RecordPtr, Event, ReducedEventSize);
+  endBufferIfFull();
+}
+
+XRayLogInitStatus fdrLoggingInit(std::size_t BufferSize, std::size_t BufferMax,
+                                 void *Options,
+                                 size_t OptionsSize) XRAY_NEVER_INSTRUMENT {
+  if (OptionsSize != sizeof(FDRLoggingOptions))
+    return static_cast<XRayLogInitStatus>(__sanitizer::atomic_load(
+        &LoggingStatus, __sanitizer::memory_order_acquire));
+  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);
 
-  __xray_fdr_internal::processFunctionHook(FuncId, Entry, TSC, CPU,
-                                           clock_gettime, LoggingStatus, BQ);
+  {
+    __sanitizer::SpinMutexLock Guard(&FDROptionsMutex);
+    memcpy(&FDROptions, Options, OptionsSize);
+  }
+
+  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);
+  __xray_set_customevent_handler(fdrLoggingHandleCustomEvent);
+
+  __sanitizer::atomic_store(&LoggingStatus,
+                            XRayLogInitStatus::XRAY_LOG_INITIALIZED,
+                            __sanitizer::memory_order_release);
+  Report("XRay FDR init successful.\n");
+  return XRayLogInitStatus::XRAY_LOG_INITIALIZED;
 }
 
 } // namespace __xray
 
 static auto UNUSED Unused = [] {
   using namespace __xray;
   if (flags()->xray_fdr_log) {
     XRayLogImpl Impl{
         fdrLoggingInit, fdrLoggingFinalize, fdrLoggingHandleArg0,
         fdrLoggingFlush,
     };
     __xray_set_log_impl(Impl);
   }
   return true;
 }();

compiler-rt/trunk/lib/xray/xray_fdr_logging_impl.h

 #include "xray_buffer_queue.h"
 #include "xray_defs.h"
 #include "xray_fdr_log_records.h"
 #include "xray_flags.h"
 #include "xray_tsc.h"
 
 namespace __xray {
 
+__sanitizer::atomic_sint32_t LoggingStatus = {
+    XRayLogInitStatus::XRAY_LOG_UNINITIALIZED};
+
 /// We expose some of the state transitions when FDR logging mode is operating
 /// such that we can simulate a series of log events that may occur without
 /// and test with determinism without worrying about the real CPU time.
 ///
 /// Because the code uses thread_local allocation extensively as part of its
 /// design, callers that wish to test events occuring on different threads
 /// will actually have to run them on different threads.
 ///

 
 // The number of adjacent, consecutive pairs of FunctionEntry, Tail Exit
 // records preceding RecordPtr.
 thread_local uint8_t NumTailCalls = 0;
 
 constexpr auto MetadataRecSize = sizeof(MetadataRecord);
 constexpr auto FunctionRecSize = sizeof(FunctionRecord);
 
+// 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.
+thread_local uint16_t CurrentCPU = std::numeric_limits<uint16_t>::max();
+thread_local uint64_t LastTSC = 0;
+thread_local uint64_t LastFunctionEntryTSC = 0;
+
 class ThreadExitBufferCleanup {
-  std::weak_ptr<BufferQueue> Buffers;
+  std::shared_ptr<BufferQueue> &Buffers;
   BufferQueue::Buffer &Buffer;
 
 public:
-  explicit ThreadExitBufferCleanup(std::weak_ptr<BufferQueue> BQ,
+  explicit ThreadExitBufferCleanup(std::shared_ptr<BufferQueue> &BQ,
                                    BufferQueue::Buffer &Buffer)
       XRAY_NEVER_INSTRUMENT : Buffers(BQ),
                               Buffer(Buffer) {}
 
   ~ThreadExitBufferCleanup() noexcept XRAY_NEVER_INSTRUMENT {
     if (RecordPtr == nullptr)
       return;
 
     // We make sure that upon exit, a thread will write out the EOB
     // MetadataRecord in the thread-local log, and also release the buffer to
     // the queue.
     assert((RecordPtr + MetadataRecSize) - static_cast<char *>(Buffer.Buffer) >=
            static_cast<ptrdiff_t>(MetadataRecSize));
-    if (auto BQ = Buffers.lock()) {
+    if (Buffers) {
       writeEOBMetadata();
-      auto EC = BQ->releaseBuffer(Buffer);
+      auto EC = Buffers->releaseBuffer(Buffer);
       if (EC != BufferQueue::ErrorCode::Ok)
         Report("Failed to release buffer at %p; error=%s\n", Buffer.Buffer,
                BufferQueue::getErrorString(EC));
+      Buffers = nullptr;
       return;
     }
   }
 };
 
+// Make sure a thread that's ever called handleArg0 has a thread-local
+// live reference to the buffer queue for this particular instance of
+// FDRLogging, and that we're going to clean it up when the thread exits.
+thread_local std::shared_ptr<BufferQueue> LocalBQ = nullptr;
+thread_local ThreadExitBufferCleanup Cleanup(LocalBQ, Buffer);
+
 class RecursionGuard {
   bool &Running;
   const bool Valid;
 
 public:
   explicit RecursionGuard(bool &R) : Running(R), Valid(!R) {
     if (Valid)
       Running = true;
   }
 
   RecursionGuard(const RecursionGuard &) = delete;
   RecursionGuard(RecursionGuard &&) = delete;
   RecursionGuard &operator=(const RecursionGuard &) = delete;
   RecursionGuard &operator=(RecursionGuard &&) = delete;
 
   explicit operator bool() const { return Valid; }
 
   ~RecursionGuard() noexcept {
     if (Valid)
       Running = false;
   }
 };
 
-static inline bool loggingInitialized(
+inline bool loggingInitialized(
     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,
+inline void writeNewBufferPreamble(pid_t Tid, timespec TS,
                                           char *&MemPtr) XRAY_NEVER_INSTRUMENT {
   static constexpr int InitRecordsCount = 2;
   std::aligned_storage<sizeof(MetadataRecord)>::type Records[InitRecordsCount];
   {
     // Write out a MetadataRecord to signify that this is the start of a new
     // buffer, associated with a particular thread, with a new CPU.  For the
     // data, we have 15 bytes to squeeze as much information as we can.  At this
     // point we only write down the following bytes:
// Also write the WalltimeMarker record.
     std::memcpy(WalltimeMarker.Data + sizeof(Seconds), &Micros, sizeof(Micros));
   }
   std::memcpy(MemPtr, Records, sizeof(MetadataRecord) * InitRecordsCount);
   MemPtr += sizeof(MetadataRecord) * InitRecordsCount;
   NumConsecutiveFnEnters = 0;
   NumTailCalls = 0;
 }
 
-static inline void setupNewBuffer(int (*wall_clock_reader)(clockid_t,
-                                                           struct timespec *))
-    XRAY_NEVER_INSTRUMENT {
+inline void setupNewBuffer(int (*wall_clock_reader)(
+    clockid_t, struct timespec *)) XRAY_NEVER_INSTRUMENT {
   RecordPtr = static_cast<char *>(Buffer.Buffer);
   pid_t Tid = syscall(SYS_gettid);
   timespec TS{0, 0};
   // This is typically clock_gettime, but callers have injection ability.
   wall_clock_reader(CLOCK_MONOTONIC, &TS);
   writeNewBufferPreamble(Tid, TS, RecordPtr);
   NumConsecutiveFnEnters = 0;
   NumTailCalls = 0;
 }
 
-static inline void writeNewCPUIdMetadata(uint16_t CPU, uint64_t TSC,
+inline void writeNewCPUIdMetadata(uint16_t CPU, uint64_t TSC,
                                          char *&MemPtr) XRAY_NEVER_INSTRUMENT {
   MetadataRecord NewCPUId;
   NewCPUId.Type = uint8_t(RecordType::Metadata);
   NewCPUId.RecordKind = uint8_t(MetadataRecord::RecordKinds::NewCPUId);
 
   // The data for the New CPU will contain the following bytes:
   //   - CPU ID (uint16_t, 2 bytes)
   //   - Full TSC (uint64_t, 8 bytes)
   // Total = 10 bytes.
   std::memcpy(&NewCPUId.Data, &CPU, sizeof(CPU));
   std::memcpy(&NewCPUId.Data[sizeof(CPU)], &TSC, sizeof(TSC));
   std::memcpy(MemPtr, &NewCPUId, sizeof(MetadataRecord));
   MemPtr += sizeof(MetadataRecord);
   NumConsecutiveFnEnters = 0;
   NumTailCalls = 0;
 }
 
-static inline void writeNewCPUIdMetadata(uint16_t CPU,
+inline void writeNewCPUIdMetadata(uint16_t CPU,
                                          uint64_t TSC) XRAY_NEVER_INSTRUMENT {
   writeNewCPUIdMetadata(CPU, TSC, RecordPtr);
 }
 
-static inline void writeEOBMetadata(char *&MemPtr) XRAY_NEVER_INSTRUMENT {
+inline void writeEOBMetadata(char *&MemPtr) XRAY_NEVER_INSTRUMENT {
   MetadataRecord EOBMeta;
   EOBMeta.Type = uint8_t(RecordType::Metadata);
   EOBMeta.RecordKind = uint8_t(MetadataRecord::RecordKinds::EndOfBuffer);
   // For now we don't write any bytes into the Data field.
   std::memcpy(MemPtr, &EOBMeta, sizeof(MetadataRecord));
   MemPtr += sizeof(MetadataRecord);
   NumConsecutiveFnEnters = 0;
   NumTailCalls = 0;
 }
 
-static inline void writeEOBMetadata() XRAY_NEVER_INSTRUMENT {
+inline void writeEOBMetadata() XRAY_NEVER_INSTRUMENT {
   writeEOBMetadata(RecordPtr);
 }
 
-static inline void writeTSCWrapMetadata(uint64_t TSC,
+inline void writeTSCWrapMetadata(uint64_t TSC,
                                         char *&MemPtr) XRAY_NEVER_INSTRUMENT {
   MetadataRecord TSCWrap;
   TSCWrap.Type = uint8_t(RecordType::Metadata);
   TSCWrap.RecordKind = uint8_t(MetadataRecord::RecordKinds::TSCWrap);
 
   // The data for the TSCWrap record contains the following bytes:
   //   - Full TSC (uint64_t, 8 bytes)
   // Total = 8 bytes.
   std::memcpy(&TSCWrap.Data, &TSC, sizeof(TSC));
   std::memcpy(MemPtr, &TSCWrap, sizeof(MetadataRecord));
   MemPtr += sizeof(MetadataRecord);
   NumConsecutiveFnEnters = 0;
   NumTailCalls = 0;
 }
 
-static inline void writeTSCWrapMetadata(uint64_t TSC) XRAY_NEVER_INSTRUMENT {
+inline void writeTSCWrapMetadata(uint64_t TSC) XRAY_NEVER_INSTRUMENT {
   writeTSCWrapMetadata(TSC, RecordPtr);
 }
 
-static inline void writeFunctionRecord(int FuncId, uint32_t TSCDelta,
+inline void writeFunctionRecord(int FuncId, uint32_t TSCDelta,
                                        XRayEntryType EntryType,
                                        char *&MemPtr) XRAY_NEVER_INSTRUMENT {
   std::aligned_storage<sizeof(FunctionRecord), alignof(FunctionRecord)>::type
       AlignedFuncRecordBuffer;
   auto &FuncRecord =
       *reinterpret_cast<FunctionRecord *>(&AlignedFuncRecordBuffer);
   FuncRecord.Type = uint8_t(RecordType::Function);
   // Only take 28 bits of the function id.
if (NumConsecutiveFnEnters > 0) {
       // We will never be able to erase this tail call since we have logged
       // something in between the function entry and tail exit.
       NumTailCalls = 0;
       NumConsecutiveFnEnters = 0;
     }
     FuncRecord.RecordKind =
         uint8_t(FunctionRecord::RecordKinds::FunctionTailExit);
     break;
+  case XRayEntryType::CUSTOM_EVENT: {
+    // This is a bug in patching, so we'll report it once and move on.
+    static bool Once = [&] {
+      Report("Internal error: patched an XRay custom event call as a function; "
+             "func id = %d\n",
+             FuncId);
+      return true;
+    }();
+    (void)Once;
+    return;
+  }
   }
 
   std::memcpy(MemPtr, &AlignedFuncRecordBuffer, sizeof(FunctionRecord));
   MemPtr += sizeof(FunctionRecord);
 }
 
 static uint64_t thresholdTicks() {
-  static uint64_t TicksPerSec = probeRequiredCPUFeatures() ? getTSCFrequency() :
-                                __xray::NanosecondsPerSecond;
+  static uint64_t TicksPerSec = probeRequiredCPUFeatures()
+                                    ? getTSCFrequency()
+                                    : __xray::NanosecondsPerSecond;
   static const uint64_t ThresholdTicks =
       TicksPerSec * flags()->xray_fdr_log_func_duration_threshold_us / 1000000;
   return ThresholdTicks;
 }
 
 // Re-point the thread local pointer into this thread's Buffer before the recent
 // "Function Entry" record and any "Tail Call Exit" records after that.
 static void rewindRecentCall(uint64_t TSC, uint64_t &LastTSC,
const auto &ExpectedTailExit =
             &TailExitRecordBuffer, RewindingRecordPtr, FunctionRecSize));
 
     assert(ExpectedTailExit.RecordKind ==
                uint8_t(FunctionRecord::RecordKinds::FunctionTailExit) &&
            "Expected to find tail exit when rewinding.");
     RewindingRecordPtr -= FunctionRecSize;
     RewindingTSC -= ExpectedTailExit.TSCDelta;
     AlignedFuncStorage FunctionEntryBuffer;
-    const auto &ExpectedFunctionEntry =
-        *reinterpret_cast<FunctionRecord *>(std::memcpy(
-            &FunctionEntryBuffer, RewindingRecordPtr, FunctionRecSize));
+    const auto &ExpectedFunctionEntry = *reinterpret_cast<FunctionRecord *>(
+        std::memcpy(&FunctionEntryBuffer, RewindingRecordPtr, FunctionRecSize));
     assert(ExpectedFunctionEntry.RecordKind ==
                uint8_t(FunctionRecord::RecordKinds::FunctionEnter) &&
            "Expected to find function entry when rewinding tail call.");
     assert(ExpectedFunctionEntry.FuncId == ExpectedTailExit.FuncId &&
            "Expected funcids to match when rewinding tail call.");
 
     // This tail call exceeded the threshold duration. It will not be erased.
     if ((TSC - RewindingTSC) >= thresholdTicks()) {
       NumTailCalls = 0;
       return;
     }
 
     // We can erase a tail exit pair that we're exiting through since
     // its duration is under threshold.
     --NumTailCalls;
     RewindingRecordPtr -= FunctionRecSize;
     RewindingTSC -= ExpectedFunctionEntry.TSCDelta;
     RecordPtr -= 2 * FunctionRecSize;
     LastTSC = RewindingTSC;
   }
 }
 
-static inline bool releaseThreadLocalBuffer(BufferQueue *BQ) {
+inline bool releaseThreadLocalBuffer(BufferQueue *BQ) {
   auto EC = BQ->releaseBuffer(Buffer);
   if (EC != BufferQueue::ErrorCode::Ok) {
     Report("Failed to release buffer at %p; error=%s\n", Buffer.Buffer,
            BufferQueue::getErrorString(EC));
     return false;
   }
   return true;
 }
 
-static inline 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) XRAY_NEVER_INSTRUMENT {
+inline bool prepareBuffer(int (*wall_clock_reader)(clockid_t,
+                                                          struct timespec *),
+                                 size_t MaxSize) XRAY_NEVER_INSTRUMENT {
+  char *BufferStart = static_cast<char *>(Buffer.Buffer);
+  if ((RecordPtr + MaxSize) > (BufferStart + Buffer.Size - MetadataRecSize)) {
+    writeEOBMetadata();
+    if (!releaseThreadLocalBuffer(LocalBQ.get()))
+      return false;
+    auto EC = LocalBQ->getBuffer(Buffer);
+    if (EC != BufferQueue::ErrorCode::Ok) {
+      Report("Failed to acquire a buffer; error=%s\n",
+             BufferQueue::getErrorString(EC));
+      return false;
+    }
+    setupNewBuffer(wall_clock_reader);
+  }
+  return true;
+}
+
+inline bool isLogInitializedAndReady(
+    std::shared_ptr<BufferQueue> &LocalBQ, uint64_t TSC, unsigned char CPU,
+    int (*wall_clock_reader)(clockid_t,
+                             struct timespec *)) XRAY_NEVER_INSTRUMENT {
   // Bail out right away if logging is not initialized yet.
   // We should take the opportunity to release the buffer though.
   auto Status = __sanitizer::atomic_load(&LoggingStatus,
                                          __sanitizer::memory_order_acquire);
   if (Status != XRayLogInitStatus::XRAY_LOG_INITIALIZED) {
     if (RecordPtr != nullptr &&
         (Status == XRayLogInitStatus::XRAY_LOG_FINALIZING ||
          Status == XRayLogInitStatus::XRAY_LOG_FINALIZED)) {
       writeEOBMetadata();
-      if (!releaseThreadLocalBuffer(BQ.get()))
-        return;
+      if (!releaseThreadLocalBuffer(LocalBQ.get()))
+        return false;
       RecordPtr = nullptr;
+      LocalBQ = nullptr;
+      return false;
     }
-    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.
-  thread_local uint16_t CurrentCPU = std::numeric_limits<uint16_t>::max();
-  thread_local uint64_t LastTSC = 0;
-  thread_local uint64_t LastFunctionEntryTSC = 0;
-
-  // Make sure a thread that's ever called handleArg0 has a thread-local
-  // live reference to the buffer queue for this particular instance of
-  // FDRLogging, and that we're going to clean it up when the thread exits.
-  thread_local auto LocalBQ = BQ;
-  thread_local ThreadExitBufferCleanup Cleanup(LocalBQ, Buffer);
-
-  // Prevent signal handler recursion, so in case we're already in a log writing
-  // mode and the signal handler comes in (and is also instrumented) then we
-  // don't want to be clobbering potentially partial writes already happening in
-  // the thread. We use a simple thread_local latch to only allow one on-going
-  // handleArg0 to happen at any given time.
-  thread_local bool Running = false;
-  RecursionGuard Guard{Running};
-  if (!Guard) {
-    assert(Running == true && "RecursionGuard is buggy!");
-    return;
+    return false;
   }
 
   if (!loggingInitialized(LoggingStatus) || LocalBQ->finalizing()) {
     writeEOBMetadata();
-    if (!releaseThreadLocalBuffer(BQ.get()))
-      return;
+    if (!releaseThreadLocalBuffer(LocalBQ.get()))
+      return false;
     RecordPtr = nullptr;
   }
 
   if (Buffer.Buffer == nullptr) {
     auto EC = LocalBQ->getBuffer(Buffer);
     if (EC != BufferQueue::ErrorCode::Ok) {
       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;
+      return false;
     }
 
     setupNewBuffer(wall_clock_reader);
   }
 
   if (CurrentCPU == std::numeric_limits<uint16_t>::max()) {
-    // This means this is the first CPU this thread has ever run on. We set the
-    // current CPU and record this as the first TSC we've seen.
+    // This means this is the first CPU this thread has ever run on. We set
+    // the current CPU and record this as the first TSC we've seen.
     CurrentCPU = CPU;
     writeNewCPUIdMetadata(CPU, TSC);
   }
 
+  return true;
+} // namespace __xray_fdr_internal
+
+inline void endBufferIfFull() XRAY_NEVER_INSTRUMENT {
+  auto BufferStart = static_cast<char *>(Buffer.Buffer);
+  if ((RecordPtr + MetadataRecSize) - BufferStart == MetadataRecSize) {
+    writeEOBMetadata();
+    if (!releaseThreadLocalBuffer(LocalBQ.get()))
+      return;
+    RecordPtr = nullptr;
+  }
+}
+
+inline 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) XRAY_NEVER_INSTRUMENT {
+  // Prevent signal handler recursion, so in case we're already in a log writing
+  // mode and the signal handler comes in (and is also instrumented) then we
+  // don't want to be clobbering potentially partial writes already happening in
+  // the thread. We use a simple thread_local latch to only allow one on-going
+  // handleArg0 to happen at any given time.
+  thread_local bool Running = false;
+  RecursionGuard Guard{Running};
+  if (!Guard) {
+    assert(Running == true && "RecursionGuard is buggy!");
+    return;
+  }
+
+  // In case the reference has been cleaned up before, we make sure we
+  // initialize it to the provided BufferQueue.
+  if (LocalBQ == nullptr)
+    LocalBQ = BQ;
+
+  if (!isLogInitializedAndReady(LocalBQ, TSC, CPU, wall_clock_reader))
+    return;
+
   // Before we go setting up writing new function entries, we need to be really
   // careful about the pointer math we're doing. This means we need to ensure
   // that the record we are about to write is going to fit into the buffer,
   // without overflowing the buffer.
   //
   // To do this properly, we use the following assumptions:
   //
   //   - The least number of bytes we will ever write is 8
inline void processFunctionHook(
   //   - An End-of-Buffer (EOB) MetadataRecord is 16 bytes.
   //
   // So the math we need to do is to determine whether writing 24 bytes past the
   // current pointer leaves us with enough bytes to write the EOB
   // MetadataRecord. If we don't have enough space after writing as much as 24
   // bytes in the end of the buffer, we need to write out the EOB, get a new
   // Buffer, set it up properly before doing any further writing.
   //
-  char *BufferStart = static_cast<char *>(Buffer.Buffer);
-  if ((RecordPtr + (MetadataRecSize + FunctionRecSize)) - BufferStart <
-      static_cast<ptrdiff_t>(MetadataRecSize)) {
-    writeEOBMetadata();
-    if (!releaseThreadLocalBuffer(LocalBQ.get()))
-      return;
-    auto EC = LocalBQ->getBuffer(Buffer);
-    if (EC != BufferQueue::ErrorCode::Ok) {
-      Report("Failed to acquire a buffer; error=%s\n",
-             BufferQueue::getErrorString(EC));
-      return;
-    }
-    setupNewBuffer(wall_clock_reader);
-  }
+  if (!prepareBuffer(wall_clock_reader, FunctionRecSize + MetadataRecSize)) {
+    LocalBQ = nullptr;
 
   // By this point, we are now ready to write at most 24 bytes (one metadata
   // record and one function record).
-  BufferStart = static_cast<char *>(Buffer.Buffer);
+  auto BufferStart = static_cast<char *>(Buffer.Buffer);
   assert((RecordPtr + (MetadataRecSize + FunctionRecSize)) - BufferStart >=
              static_cast<ptrdiff_t>(MetadataRecSize) &&
          "Misconfigured BufferQueue provided; Buffer size not large enough.");
 
   // Here we compute the TSC Delta. There are a few interesting situations we
   // need to account for:
   //
   //   - The thread has migrated to a different CPU. If this is the case, then
   //     we write down the following records:
   //
   //       1. A 'NewCPUId' Metadata record.
   //       2. A FunctionRecord with a 0 for the TSCDelta field.
   //
   //   - The TSC delta is greater than the 32 bits we can store in a
   //     FunctionRecord. In this case we write down the following records:
   //
   //       1. A 'TSCWrap' Metadata record.
   //       2. A FunctionRecord with a 0 for the TSCDelta field.
   //
   //   - The TSC delta is representable within the 32 bits we can store in a
   //     FunctionRecord. In this case we write down just a FunctionRecord with
   //     the correct TSC delta.
   //
-
   uint32_t RecordTSCDelta = 0;
   if (CPU != CurrentCPU) {
     // We've moved to a new CPU.
     writeNewCPUIdMetadata(CPU, TSC);
   } else {
     // If the delta is greater than the range for a uint32_t, then we write out
     // the TSC wrap metadata entry with the full TSC, and the TSC for the
     // function record be 0.
case XRayEntryType::TAIL:
     break;
   case XRayEntryType::EXIT:
     // Break out and write the exit record if we can't erase any functions.
     if (NumConsecutiveFnEnters == 0 ||
         (TSC - LastFunctionEntryTSC) >= thresholdTicks())
       break;
     rewindRecentCall(TSC, LastTSC, LastFunctionEntryTSC, FuncId);
     return; // without writing log.
+  case XRayEntryType::CUSTOM_EVENT: {
+    // This is a bug in patching, so we'll report it once and move on.
+    static bool Once = [&] {
+      Report("Internal error: patched an XRay custom event call as a function; "
+             "func id = %d",
+             FuncId);
+      return true;
+    }();
+    (void)Once;
+    return;
+  }
   }
 
   writeFunctionRecord(FuncId, RecordTSCDelta, Entry, RecordPtr);
 
   // If we've exhausted the buffer by this time, we then release the buffer to
   // make sure that other threads may start using this buffer.
-  if ((RecordPtr + MetadataRecSize) - BufferStart == MetadataRecSize) {
-    writeEOBMetadata();
-    if (!releaseThreadLocalBuffer(LocalBQ.get()))
-      return;
-    RecordPtr = nullptr;
-  }
+  endBufferIfFull();
 }
 
 } // namespace __xray_fdr_internal
-
 } // namespace __xray
+
 #endif // XRAY_XRAY_FDR_LOGGING_IMPL_H

compiler-rt/trunk/lib/xray/xray_interface.cc

 #endif /* CPU architecture */
 
 // This is the function to call when we encounter the entry or exit sleds.
 __sanitizer::atomic_uintptr_t XRayPatchedFunction{0};
 
 // This is the function to call from the arg1-enabled sleds/trampolines.
 __sanitizer::atomic_uintptr_t XRayArgLogger{0};
 
+// This is the function to call when we encounter a custom event log call.
+__sanitizer::atomic_uintptr_t XRayPatchedCustomEvent{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
 // read-and-execute only.

 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<uint64_t>(entry),
+                              reinterpret_cast<uintptr_t>(entry),
+                              __sanitizer::memory_order_release);
+    return 1;
+  }
+  return 0;
+}
+
+int __xray_set_customevent_handler(void (*entry)(void *, size_t))
+    XRAY_NEVER_INSTRUMENT {
+  if (__sanitizer::atomic_load(&XRayInitialized,
+                               __sanitizer::memory_order_acquire)) {
+    __sanitizer::atomic_store(&__xray::XRayPatchedCustomEvent,
+                              reinterpret_cast<uintptr_t>(entry),
                               __sanitizer::memory_order_release);
     return 1;
   }
   return 0;
 }
 
 int __xray_remove_handler() XRAY_NEVER_INSTRUMENT {
   return __xray_set_handler(nullptr);
case XRayEntryType::EXIT:
     Success = patchFunctionExit(Enable, FuncId, Sled);
     break;
   case XRayEntryType::TAIL:
     Success = patchFunctionTailExit(Enable, FuncId, Sled);
     break;
   case XRayEntryType::LOG_ARGS_ENTRY:
     Success = patchFunctionEntry(Enable, FuncId, Sled, __xray_ArgLoggerEntry);
     break;
+  case XRayEntryType::CUSTOM_EVENT:
+    Success = patchCustomEvent(Enable, FuncId, Sled);
+    break;
   default:
     Report("Unsupported sled kind '%d' @%04x\n", Sled.Address, int(Sled.Kind));
     return false;
   }
   return Success;
 }
 
 // controlPatching implements the common internals of the patching/unpatching
int __xray_set_handler_arg1(void (*Handler)(int32_t, XRayEntryType, uint64_t)) {
 
   // 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); }
 
 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)
     return 0;
   return XRayInstrMap.SledsIndex[FuncId - 1].Begin->Address;
 }
 
 size_t __xray_max_function_id() XRAY_NEVER_INSTRUMENT {
   __sanitizer::SpinMutexLock Guard(&XRayInstrMapMutex);
   return XRayInstrMap.Functions;
 }

compiler-rt/trunk/lib/xray/xray_interface_internal.h

struct XRaySledMap {
   size_t Functions;
 };
 
 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
 
 extern "C" {
 // The following functions have to be defined in assembler, on a per-platform
 // basis. See xray_trampoline_*.S files for implementations.
 extern void __xray_FunctionEntry();
 extern void __xray_FunctionExit();
 extern void __xray_FunctionTailExit();
 extern void __xray_ArgLoggerEntry();
+extern void __xray_CustomEvent();
 }
 
 #endif

compiler-rt/trunk/lib/xray/xray_mips.cc

inline static bool patchSled(const bool Enable, const uint32_t FuncId,
   // operation, so that the user code which reaches the sled concurrently
   // either jumps over the whole sled, or executes the whole sled when the
   // latter is ready.
   //
   // When |Enable|==false, we set back the first instruction in the sled to be
   //   B #44
 
   if (Enable) {
-    uint32_t LoTracingHookAddr = reinterpret_cast<int32_t>(TracingHook) & 0xffff;
+    uint32_t LoTracingHookAddr =
+        reinterpret_cast<int32_t>(TracingHook) & 0xffff;
     uint32_t HiTracingHookAddr =
         (reinterpret_cast<int32_t>(TracingHook) >> 16) & 0xffff;
     uint32_t LoFunctionID = FuncId & 0xffff;
     uint32_t HiFunctionID = (FuncId >> 16) & 0xffff;
     *reinterpret_cast<uint32_t *>(Sled.Address + 8) = encodeInstruction(
         PatchOpcodes::PO_SW, RegNum::RN_SP, RegNum::RN_RA, 0x4);
     *reinterpret_cast<uint32_t *>(Sled.Address + 12) = encodeInstruction(
         PatchOpcodes::PO_SW, RegNum::RN_SP, RegNum::RN_T9, 0x0);

 
 bool patchFunctionTailExit(const bool Enable, const uint32_t FuncId,
                            const XRaySledEntry &Sled) XRAY_NEVER_INSTRUMENT {
   // FIXME: In the future we'd need to distinguish between non-tail exits and
   // tail exits for better information preservation.
   return patchSled(Enable, FuncId, Sled, __xray_FunctionExit);
 }
 
+bool patchCustomEvent(const bool Enable, const uint32_t FuncId,
+                      const XRaySledEntry &Sled) XRAY_NEVER_INSTRUMENT {
+  // FIXME: Implement in mips?
+  return false;
+}
+
 } // namespace __xray
 
 extern "C" void __xray_ArgLoggerEntry() XRAY_NEVER_INSTRUMENT {
   // FIXME: this will have to be implemented in the trampoline assembly file
 }

compiler-rt/trunk/lib/xray/xray_mips64.cc

inline static bool patchSled(const bool Enable, const uint32_t FuncId,
   // latter is ready.
   //
   // When |Enable|==false, we set back the first instruction in the sled to be
   //   B #60
 
   if (Enable) {
     uint32_t LoTracingHookAddr =
         reinterpret_cast<int64_t>(TracingHook) & 0xffff;
-    uint32_t HiTracingHookAddr = (reinterpret_cast<int64_t>(TracingHook) >> 16) & 0xffff;
+    uint32_t HiTracingHookAddr =
+        (reinterpret_cast<int64_t>(TracingHook) >> 16) & 0xffff;
     uint32_t HigherTracingHookAddr =
         (reinterpret_cast<int64_t>(TracingHook) >> 32) & 0xffff;
     uint32_t HighestTracingHookAddr =
         (reinterpret_cast<int64_t>(TracingHook) >> 48) & 0xffff;
     uint32_t LoFunctionID = FuncId & 0xffff;
     uint32_t HiFunctionID = (FuncId >> 16) & 0xffff;
     *reinterpret_cast<uint32_t *>(Sled.Address + 8) = encodeInstruction(
         PatchOpcodes::PO_SD, RegNum::RN_SP, RegNum::RN_RA, 0x8);

 
 bool patchFunctionTailExit(const bool Enable, const uint32_t FuncId,
                            const XRaySledEntry &Sled) XRAY_NEVER_INSTRUMENT {
   // FIXME: In the future we'd need to distinguish between non-tail exits and
   // tail exits for better information preservation.
   return patchSled(Enable, FuncId, Sled, __xray_FunctionExit);
 }
 
+bool patchCustomEvent(const bool Enable, const uint32_t FuncId,
+                      const XRaySledEntry &Sled) XRAY_NEVER_INSTRUMENT {
+  // FIXME: Implement in mips64?
+  return false;
+}
 } // namespace __xray
 
 extern "C" void __xray_ArgLoggerEntry() XRAY_NEVER_INSTRUMENT {
   // FIXME: this will have to be implemented in the trampoline assembly file
 }

compiler-rt/trunk/lib/xray/xray_powerpc64.cc

 bool patchFunctionTailExit(const bool Enable, const uint32_t FuncId,
                            const XRaySledEntry &Sled) XRAY_NEVER_INSTRUMENT {
   return patchFunctionExit(Enable, FuncId, Sled);
 }
 
 // FIXME: Maybe implement this better?
 bool probeRequiredCPUFeatures() XRAY_NEVER_INSTRUMENT { return true; }
 
+bool patchCustomEvent(const bool Enable, const uint32_t FuncId,
+                      const XRaySledEntry &Sled) XRAY_NEVER_INSTRUMENT {
+  // FIXME: Implement in powerpc64?
+  return false;
+}
+
 } // namespace __xray
 
 extern "C" void __xray_ArgLoggerEntry() XRAY_NEVER_INSTRUMENT {
   // FIXME: this will have to be implemented in the trampoline assembly file
 }

compiler-rt/trunk/lib/xray/xray_trampoline_x86_64.S

__xray_ArgLoggerEntry:
   jne .Larg1entryLog
 
   // If [arg1 logging handler] not set, defer to no-arg logging.
   movq  _ZN6__xray19XRayPatchedFunctionE(%rip), %rax
   testq %rax, %rax
   je  .Larg1entryFail
 
 .Larg1entryLog:
-  movq  %rdi, %rdx  // first argument will become the third
-  xorq  %rsi, %rsi  // XRayEntryType::ENTRY into the second
-  movl  %r10d, %edi  // 32-bit function ID becomes the first
+
+  // First argument will become the third
+  movq  %rdi, %rdx
+
+  // XRayEntryType::ENTRY into the second
+  xorq  %rsi, %rsi
+
+  // 32-bit function ID becomes the first
+  movl  %r10d, %edi
   callq *%rax
 
 .Larg1entryFail:
   RESTORE_REGISTERS
   retq
 
 .Larg1entryEnd:
   .size __xray_ArgLoggerEntry, .Larg1entryEnd-__xray_ArgLoggerEntry
   .cfi_endproc
 
+//===----------------------------------------------------------------------===//
+
+  .global __xray_CustomEvent
+  .align 16, 0x90
+  .type __xray_CustomEvent,@function
+__xray_CustomEvent:
+  .cfi_startproc
+  subq $16, %rsp
+  .cfi_def_cfa_offset 24
+  movq %rbp, 8(%rsp)
+  movq %rax, 0(%rsp)
+
+  // We take two arguments to this trampoline, which should be in rdi  and rsi
+  // already. We also make sure that we stash %rax because we use that register
+  // to call the logging handler.
+  movq _ZN6__xray22XRayPatchedCustomEventE(%rip), %rax
+  testq %rax,%rax
+  je .LcustomEventCleanup
+
+  // At this point we know that rcx and rdx already has the data, so we just
+  // call the logging handler.
+  callq *%rax
+
+.LcustomEventCleanup:
+  movq 0(%rsp), %rax
+  movq 8(%rsp), %rbp
+  addq $16, %rsp
+  .cfi_def_cfa_offset 8
+  retq
+
+.Ltmp8:
+  .size __xray_CustomEvent, .Ltmp8-__xray_CustomEvent
+  .cfi_endproc
+
 NO_EXEC_STACK_DIRECTIVE

compiler-rt/trunk/lib/xray/xray_x86_64.cc

if (readValueFromFile("/sys/devices/system/cpu/cpu0/tsc_freq_khz",
     Report("Unable to determine CPU frequency for TSC accounting.\n");
   }
   return TSCFrequency == -1 ? 0 : static_cast<uint64_t>(TSCFrequency);
 }
 
 static constexpr uint8_t CallOpCode = 0xe8;
 static constexpr uint16_t MovR10Seq = 0xba41;
 static constexpr uint16_t Jmp9Seq = 0x09eb;
+static constexpr uint16_t Jmp20Seq = 0x14eb;
 static constexpr uint8_t JmpOpCode = 0xe9;
 static constexpr uint8_t RetOpCode = 0xc3;
+static constexpr uint16_t NopwSeq = 0x9066;
 
 static constexpr int64_t MinOffset{std::numeric_limits<int32_t>::min()};
 static constexpr int64_t MaxOffset{std::numeric_limits<int32_t>::max()};
 
 bool patchFunctionEntry(const bool Enable, const uint32_t FuncId,
                         const XRaySledEntry &Sled,
                         void (*Trampoline)()) XRAY_NEVER_INSTRUMENT {
   // Here we do the dance of replacing the following sled:
if (Enable) {
     std::atomic_store_explicit(
         reinterpret_cast<std::atomic<uint16_t> *>(Sled.Address), Jmp9Seq,
         std::memory_order_release);
     // FIXME: Write out the nops still?
   }
   return true;
 }
 
+bool patchCustomEvent(const bool Enable, const uint32_t FuncId,
+                      const XRaySledEntry &Sled) XRAY_NEVER_INSTRUMENT {
+  // Here we do the dance of replacing the following sled:
+  //
+  // xray_sled_n:
+  //   jmp +19          // 2 bytes
+  //   ...
+  //
+  // With the following:
+  //
+  //   nopw             // 2 bytes*
+  //   ...
+  //
+  // We need to do this in the following order:
+  //
+  // 1. Overwrite the 5-byte nop with the call (relative), where (relative) is
+  //    the relative offset to the __xray_CustomEvent trampoline.
+  // 2. Do a two-byte atomic write over the 'jmp +24' to turn it into a 'nopw'.
+  //    This allows us to "enable" this code once the changes have committed.
+  //
+  // The "unpatch" should just turn the 'nopw' back to a 'jmp +24'.
+  //
+  if (Enable) {
+    std::atomic_store_explicit(
+        reinterpret_cast<std::atomic<uint16_t> *>(Sled.Address), NopwSeq,
+        std::memory_order_release);
+  } else {
+    std::atomic_store_explicit(
+        reinterpret_cast<std::atomic<uint16_t> *>(Sled.Address), Jmp20Seq,
+        std::memory_order_release);
+  }
+  return false;
+}
+
 // We determine whether the CPU we're running on has the correct features we
 // need. In x86_64 this will be rdtscp support.
 bool probeRequiredCPUFeatures() XRAY_NEVER_INSTRUMENT {
   unsigned int EAX, EBX, ECX, EDX;
 
   // We check whether rdtscp support is enabled. According to the x86_64 manual,
   // level should be set at 0x80000001, and we should have a look at bit 27 in
   // EDX. That's 0x8000000 (or 1u << 26).

compiler-rt/trunk/test/xray/TestCases/Linux/custom-event-logging.cc

+// Use the clang feature for custom xray event logging.
+//
+// RUN: %clangxx_xray -std=c++11 %s -o %t
+// RUN: XRAY_OPTIONS="patch_premain=false verbosity=1 xray_naive_log=false xray_logfile_base=custom-event-logging.xray-" %run %t 2>&1 | FileCheck %s
+//
+#include <cstdio>
+#include "xray/xray_interface.h"
+
+[[clang::xray_always_instrument]] void foo() {
+  static constexpr char CustomLogged[] = "hello custom logging!";
+  printf("before calling the custom logging...\n");
+  __xray_customevent(CustomLogged, sizeof(CustomLogged));
+  printf("after calling the custom logging...\n");
+}
+
+void myprinter(void* ptr, size_t size) {
+  printf("%.*s\n", static_cast<int>(size), static_cast<const char*>(ptr));
+}
+
+int main() {
+  foo();
+  // CHECK: before calling the custom logging...
+  // CHECK-NEXT: after calling the custom logging...
+  printf("setting up custom event handler...\n");
+  // CHECK-NEXT: setting up custom event handler...
+  __xray_set_customevent_handler(myprinter);
+  __xray_patch();
+  // CHECK-NEXT: before calling the custom logging...
+  foo();
+  // CHECK-NEXT: hello custom logging!
+  // CHECK-NEXT: after calling the custom logging...
+  printf("removing custom event handler...\n");
+  // CHECK-NEXT: removing custom event handler...
+  __xray_set_customevent_handler(nullptr);
+  foo();
+  // CHECK-NEXT: before calling the custom logging...
+  // CHECK-NEXT: after calling the custom logging...
+}