1 files changed, 175 insertions, 84 deletions

diff --git a/lib/xray/xray_inmemory_log.cc b/lib/xray/xray_inmemory_log.cc
index 188bae642..43cb50a54 100644
--- a/lib/xray/xray_inmemory_log.cc
+++ b/lib/xray/xray_inmemory_log.cc
@@ -19,64 +19,53 @@
 #include <cstring>
 #include <errno.h>
 #include <fcntl.h>
+#include <pthread.h>
 #include <sys/stat.h>
 #include <sys/syscall.h>
 #include <sys/types.h>
 #include <time.h>
 #include <unistd.h>
 
+#include "sanitizer_common/sanitizer_allocator_internal.h"
 #include "sanitizer_common/sanitizer_libc.h"
 #include "xray/xray_records.h"
 #include "xray_defs.h"
 #include "xray_flags.h"
+#include "xray_inmemory_log.h"
 #include "xray_interface_internal.h"
 #include "xray_tsc.h"
 #include "xray_utils.h"
 
-// __xray_InMemoryRawLog will use a thread-local aligned buffer capped to a
-// certain size (32kb by default) and use it as if it were a circular buffer for
-// events. We store simple fixed-sized entries in the log for external analysis.
-
-extern "C" {
-void __xray_InMemoryRawLog(int32_t FuncId,
-                           XRayEntryType Type) XRAY_NEVER_INSTRUMENT;
-}
-
 namespace __xray {
 
 __sanitizer::SpinMutex LogMutex;
 
-class ThreadExitFlusher {
-  int Fd;
-  XRayRecord *Start;
-  size_t &Offset;
-
-public:
-  explicit ThreadExitFlusher(int Fd, XRayRecord *Start,
-                             size_t &Offset) XRAY_NEVER_INSTRUMENT
-      : Fd(Fd),
-        Start(Start),
-        Offset(Offset) {}
+// We use elements of this type to record the entry TSC of every function ID we
+// see as we're tracing a particular thread's execution.
+struct StackEntry {
+  int32_t FuncId;
+  uint64_t TSC;
+};
 
-  ~ThreadExitFlusher() XRAY_NEVER_INSTRUMENT {
-    __sanitizer::SpinMutexLock L(&LogMutex);
-    if (Fd > 0 && Start != nullptr) {
-      retryingWriteAll(Fd, reinterpret_cast<char *>(Start),
-                       reinterpret_cast<char *>(Start + Offset));
-      // Because this thread's exit could be the last one trying to write to the
-      // file and that we're not able to close out the file properly, we sync
-      // instead and hope that the pending writes are flushed as the thread
-      // exits.
-      fsync(Fd);
-    }
-  }
+struct alignas(64) ThreadLocalData {
+  XRayRecord *InMemoryBuffer = nullptr;
+  size_t BufferSize = 0;
+  size_t BufferOffset = 0;
+  StackEntry *ShadowStack = nullptr;
+  size_t StackSize = 0;
+  size_t StackEntries = 0;
+  int Fd = -1;
 };
 
-} // namespace __xray
+static pthread_key_t PThreadKey;
 
-using namespace __xray;
+static __sanitizer::atomic_uint8_t BasicInitialized{0};
 
-static int __xray_OpenLogFile() XRAY_NEVER_INSTRUMENT {
+BasicLoggingOptions GlobalOptions;
+
+thread_local volatile bool RecusionGuard = false;
+
+static int openLogFile() XRAY_NEVER_INSTRUMENT {
   int F = getLogFD();
   if (F == -1)
     return -1;
@@ -103,37 +92,49 @@ static int __xray_OpenLogFile() XRAY_NEVER_INSTRUMENT {
   return F;
 }
 
-using Buffer =
-    std::aligned_storage<sizeof(XRayRecord), alignof(XRayRecord)>::type;
-
-static constexpr size_t BuffLen = 1024;
-thread_local size_t Offset = 0;
-
-Buffer (&getThreadLocalBuffer())[BuffLen] XRAY_NEVER_INSTRUMENT {
-  thread_local static Buffer InMemoryBuffer[BuffLen] = {};
-  return InMemoryBuffer;
-}
-
 pid_t getTId() XRAY_NEVER_INSTRUMENT {
   thread_local pid_t TId = syscall(SYS_gettid);
   return TId;
 }
 
 int getGlobalFd() XRAY_NEVER_INSTRUMENT {
-  static int Fd = __xray_OpenLogFile();
+  static int Fd = openLogFile();
   return Fd;
 }
 
-thread_local volatile bool RecusionGuard = false;
+ThreadLocalData &getThreadLocalData() XRAY_NEVER_INSTRUMENT {
+  thread_local ThreadLocalData TLD;
+  thread_local bool UNUSED TOnce = [] {
+    if (GlobalOptions.ThreadBufferSize == 0)
+      return false;
+    pthread_setspecific(PThreadKey, &TLD);
+    TLD.Fd = getGlobalFd();
+    TLD.InMemoryBuffer = reinterpret_cast<XRayRecord *>(
+        InternalAlloc(sizeof(XRayRecord) * GlobalOptions.ThreadBufferSize,
+                      nullptr, alignof(XRayRecord)));
+    TLD.BufferSize = GlobalOptions.ThreadBufferSize;
+    TLD.BufferOffset = 0;
+    if (GlobalOptions.MaxStackDepth == 0)
+      return false;
+    TLD.ShadowStack = reinterpret_cast<StackEntry *>(
+        InternalAlloc(sizeof(StackEntry) * GlobalOptions.MaxStackDepth, nullptr,
+                      alignof(StackEntry)));
+    TLD.StackSize = GlobalOptions.MaxStackDepth;
+    TLD.StackEntries = 0;
+    return false;
+  }();
+  return TLD;
+}
+
 template <class RDTSC>
-void __xray_InMemoryRawLog(int32_t FuncId, XRayEntryType Type,
-                           RDTSC ReadTSC) XRAY_NEVER_INSTRUMENT {
-  auto &InMemoryBuffer = getThreadLocalBuffer();
+void InMemoryRawLog(int32_t FuncId, XRayEntryType Type,
+                    RDTSC ReadTSC) XRAY_NEVER_INSTRUMENT {
+  auto &TLD = getThreadLocalData();
+  auto &InMemoryBuffer = TLD.InMemoryBuffer;
+  auto &Offset = TLD.BufferOffset;
   int Fd = getGlobalFd();
   if (Fd == -1)
     return;
-  thread_local __xray::ThreadExitFlusher Flusher(
-      Fd, reinterpret_cast<__xray::XRayRecord *>(InMemoryBuffer), Offset);
 
   // Use a simple recursion guard, to handle cases where we're already logging
   // and for one reason or another, this function gets called again in the same
@@ -150,8 +151,7 @@ void __xray_InMemoryRawLog(int32_t FuncId, XRayEntryType Type,
   R.TId = getTId();
   R.Type = Type;
   R.FuncId = FuncId;
-  ++Offset;
-  if (Offset == BuffLen) {
+  if (++Offset == TLD.BufferSize) {
     __sanitizer::SpinMutexLock L(&LogMutex);
     auto RecordBuffer = reinterpret_cast<__xray::XRayRecord *>(InMemoryBuffer);
     retryingWriteAll(Fd, reinterpret_cast<char *>(RecordBuffer),
@@ -163,10 +163,12 @@ void __xray_InMemoryRawLog(int32_t FuncId, XRayEntryType Type,
 }
 
 template <class RDTSC>
-void __xray_InMemoryRawLogWithArg(int32_t FuncId, XRayEntryType Type,
-                                  uint64_t Arg1,
-                                  RDTSC ReadTSC) XRAY_NEVER_INSTRUMENT {
-  auto &InMemoryBuffer = getThreadLocalBuffer();
+void InMemoryRawLogWithArg(int32_t FuncId, XRayEntryType Type, uint64_t Arg1,
+                           RDTSC ReadTSC) XRAY_NEVER_INSTRUMENT {
+  auto &TLD = getThreadLocalData();
+  auto &InMemoryBuffer = TLD.InMemoryBuffer;
+  auto &Offset = TLD.BufferOffset;
+  const auto &BuffLen = TLD.BufferSize;
   int Fd = getGlobalFd();
   if (Fd == -1)
     return;
@@ -183,7 +185,7 @@ void __xray_InMemoryRawLogWithArg(int32_t FuncId, XRayEntryType Type,
   }
 
   // Then we write the "we have an argument" record.
-  __xray_InMemoryRawLog(FuncId, Type, ReadTSC);
+  InMemoryRawLog(FuncId, Type, ReadTSC);
 
   if (RecusionGuard)
     return;
@@ -199,8 +201,7 @@ void __xray_InMemoryRawLogWithArg(int32_t FuncId, XRayEntryType Type,
   auto EntryPtr =
       &reinterpret_cast<__xray::XRayArgPayload *>(&InMemoryBuffer)[Offset];
   std::memcpy(EntryPtr, &R, sizeof(R));
-  ++Offset;
-  if (Offset == BuffLen) {
+  if (++Offset == BuffLen) {
     __sanitizer::SpinMutexLock L(&LogMutex);
     auto RecordBuffer = reinterpret_cast<__xray::XRayRecord *>(InMemoryBuffer);
     retryingWriteAll(Fd, reinterpret_cast<char *>(RecordBuffer),
@@ -211,14 +212,14 @@ void __xray_InMemoryRawLogWithArg(int32_t FuncId, XRayEntryType Type,
   RecusionGuard = false;
 }
 
-void __xray_InMemoryRawLogRealTSC(int32_t FuncId,
-                                  XRayEntryType Type) XRAY_NEVER_INSTRUMENT {
-  __xray_InMemoryRawLog(FuncId, Type, __xray::readTSC);
+void basicLoggingHandleArg0RealTSC(int32_t FuncId,
+                                   XRayEntryType Type) XRAY_NEVER_INSTRUMENT {
+  InMemoryRawLog(FuncId, Type, __xray::readTSC);
 }
 
-void __xray_InMemoryEmulateTSC(int32_t FuncId,
-                               XRayEntryType Type) XRAY_NEVER_INSTRUMENT {
-  __xray_InMemoryRawLog(FuncId, Type, [](uint8_t &CPU) XRAY_NEVER_INSTRUMENT {
+void basicLoggingHandleArg0EmulateTSC(int32_t FuncId, XRayEntryType Type)
+    XRAY_NEVER_INSTRUMENT {
+  InMemoryRawLog(FuncId, Type, [](uint8_t &CPU) XRAY_NEVER_INSTRUMENT {
     timespec TS;
     int result = clock_gettime(CLOCK_REALTIME, &TS);
     if (result != 0) {
@@ -230,14 +231,14 @@ void __xray_InMemoryEmulateTSC(int32_t FuncId,
   });
 }
 
-void __xray_InMemoryRawLogWithArgRealTSC(int32_t FuncId, XRayEntryType Type,
-                                         uint64_t Arg1) XRAY_NEVER_INSTRUMENT {
-  __xray_InMemoryRawLogWithArg(FuncId, Type, Arg1, __xray::readTSC);
+void basicLoggingHandleArg1RealTSC(int32_t FuncId, XRayEntryType Type,
+                                   uint64_t Arg1) XRAY_NEVER_INSTRUMENT {
+  InMemoryRawLogWithArg(FuncId, Type, Arg1, __xray::readTSC);
 }
 
-void __xray_InMemoryRawLogWithArgEmulateTSC(
-    int32_t FuncId, XRayEntryType Type, uint64_t Arg1) XRAY_NEVER_INSTRUMENT {
-  __xray_InMemoryRawLogWithArg(
+void basicLoggingHandleArg1EmulateTSC(int32_t FuncId, XRayEntryType Type,
+                                      uint64_t Arg1) XRAY_NEVER_INSTRUMENT {
+  InMemoryRawLogWithArg(
       FuncId, Type, Arg1, [](uint8_t &CPU) XRAY_NEVER_INSTRUMENT {
         timespec TS;
         int result = clock_gettime(CLOCK_REALTIME, &TS);
@@ -250,18 +251,108 @@ void __xray_InMemoryRawLogWithArgEmulateTSC(
       });
 }
 
-static auto UNUSED Unused = [] {
-  auto UseRealTSC = probeRequiredCPUFeatures();
-  if (!UseRealTSC)
+XRayLogInitStatus basicLoggingInit(size_t BufferSize, size_t BufferMax,
+                                   void *Options,
+                                   size_t OptionsSize) XRAY_NEVER_INSTRUMENT {
+  static bool UNUSED Once = [] {
+    pthread_key_create(&PThreadKey, +[](void *P) {
+      ThreadLocalData &TLD = *reinterpret_cast<ThreadLocalData *>(P);
+      if (TLD.Fd == -1 || TLD.BufferOffset == 0)
+        return;
+
+      {
+        __sanitizer::SpinMutexLock L(&LogMutex);
+        retryingWriteAll(
+            TLD.Fd, reinterpret_cast<char *>(TLD.InMemoryBuffer),
+            reinterpret_cast<char *>(TLD.InMemoryBuffer + TLD.BufferOffset));
+      }
+
+      // Because this thread's exit could be the last one trying to write to
+      // the file and that we're not able to close out the file properly, we
+      // sync instead and hope that the pending writes are flushed as the
+      // thread exits.
+      fsync(TLD.Fd);
+
+      // Clean up dynamic resources.
+      if (TLD.InMemoryBuffer)
+        InternalFree(TLD.InMemoryBuffer);
+      if (TLD.ShadowStack)
+        InternalFree(TLD.ShadowStack);
+    });
+    return false;
+  }();
+
+  uint8_t Expected = 0;
+  if (!__sanitizer::atomic_compare_exchange_strong(
+          &BasicInitialized, &Expected, 1, __sanitizer::memory_order_acq_rel)) {
+    if (__sanitizer::Verbosity())
+      Report("Basic logging already initialized.\n");
+    return XRayLogInitStatus::XRAY_LOG_INITIALIZED;
+  }
+
+  if (OptionsSize != sizeof(BasicLoggingOptions)) {
+    Report("Invalid options size, potential ABI mismatch; expected %d got %d",
+           sizeof(BasicLoggingOptions), OptionsSize);
+    return XRayLogInitStatus::XRAY_LOG_UNINITIALIZED;
+  }
+
+  static auto UseRealTSC = probeRequiredCPUFeatures();
+  if (!UseRealTSC && __sanitizer::Verbosity())
     Report("WARNING: Required CPU features missing for XRay instrumentation, "
            "using emulation instead.\n");
+
+  GlobalOptions = *reinterpret_cast<BasicLoggingOptions *>(Options);
+  __xray_set_handler_arg1(UseRealTSC ? basicLoggingHandleArg1RealTSC
+                                     : basicLoggingHandleArg1EmulateTSC);
+  __xray_set_handler(UseRealTSC ? basicLoggingHandleArg0RealTSC
+                                : basicLoggingHandleArg0EmulateTSC);
+  __xray_remove_customevent_handler();
+  return XRayLogInitStatus::XRAY_LOG_INITIALIZED;
+}
+
+XRayLogInitStatus basicLoggingFinalize() XRAY_NEVER_INSTRUMENT {
+  uint8_t Expected = 0;
+  if (!__sanitizer::atomic_compare_exchange_strong(
+          &BasicInitialized, &Expected, 0, __sanitizer::memory_order_acq_rel) &&
+      __sanitizer::Verbosity())
+    Report("Basic logging already finalized.\n");
+
+  // Nothing really to do aside from marking state of the global to be
+  // uninitialized.
+
+  return XRayLogInitStatus::XRAY_LOG_FINALIZED;
+}
+
+XRayLogFlushStatus basicLoggingFlush() XRAY_NEVER_INSTRUMENT {
+  // This really does nothing, since flushing the logs happen at the end of a
+  // thread's lifetime, or when the buffers are full.
+  return XRayLogFlushStatus::XRAY_LOG_FLUSHED;
+}
+
+// This is a handler that, effectively, does nothing.
+void basicLoggingHandleArg0Empty(int32_t, XRayEntryType) XRAY_NEVER_INSTRUMENT {
+}
+
+bool basicLogDynamicInitializer() XRAY_NEVER_INSTRUMENT {
   if (flags()->xray_naive_log) {
-    __xray_set_handler_arg1(UseRealTSC
-                                ? __xray_InMemoryRawLogWithArgRealTSC
-                                : __xray_InMemoryRawLogWithArgEmulateTSC);
-    __xray_set_handler(UseRealTSC ? __xray_InMemoryRawLogRealTSC
-                                  : __xray_InMemoryEmulateTSC);
+    XRayLogImpl Impl{
+        basicLoggingInit,
+        basicLoggingFinalize,
+        basicLoggingHandleArg0Empty,
+        basicLoggingFlush,
+    };
+    __xray_set_log_impl(Impl);
+    BasicLoggingOptions Options;
+    Options.DurationFilterMicros =
+        flags()->xray_naive_log_func_duration_threshold_us;
+    Options.MaxStackDepth = flags()->xray_naive_log_max_stack_depth;
+    Options.ThreadBufferSize = flags()->xray_naive_log_thread_buffer_size;
+    __xray_log_init(flags()->xray_naive_log_thread_buffer_size, 0, &Options,
+                    sizeof(BasicLoggingOptions));
   }
-
   return true;
-}();
+}
+
+} // namespace __xray
+
+static auto UNUSED Unused = __xray::basicLogDynamicInitializer();