[XRay][compiler-rt] Migrate basic mode logging to the XRay framework

Summary:
Before this patch, XRay's basic (naive mode) logging would be
initialised and installed in an adhoc manner. This patch ports the
implementation of the basic (naive mode) logging implementation to use
the common XRay framework.

We also make the following changes to reduce the variance between the
usage model of basic mode from FDR (flight data recorder) mode:

  - Allow programmatic control of the size of the buffers dedicated to
    per-thread records. This removes some hard-coded constants and turns
    them into runtime-controllable flags and through an Options
    structure.

  - Default the `xray_naive_log` option to false. For now, the only way
    to start basic mode is to set the environment variable, or set the
    default at build-time compiler options. Because of this change we've
    had to update a couple of tests relying on basic mode being always
    on.

  - Removed the reliance on a non-trivially destructible per-thread
    resource manager. We use a similar trick done in D39526 to use
    pthread_key_create() and pthread_setspecific() to ensure that the
    per-thread cleanup handling is performed at thread-exit time.

We also radically simplify the code structure for basic mode, to move
most of the implementation in the `__xray` namespace.

Reviewers: pelikan, eizan, kpw

Subscribers: llvm-commits

Differential Revision: https://reviews.llvm.org/D40164

git-svn-id: https://llvm.org/svn/llvm-project/compiler-rt/trunk@318734 91177308-0d34-0410-b5e6-96231b3b80d8

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();