[XRay][compiler-rt] Refactor recursion guard for Basic and FDR Mode

Summary:
This change extracts the recursion guard implementation from FDR Mode
and updates it to do the following:

- Do the atomic operation correctly to be signal-handler safe.

- Make it usable in both FDR and Basic Modes.

Before this change, the recursion guard relied on an unsynchronised read
and write on a volatile thread-local. A signal handler could then run in
between the read and the write, and then be able to run instrumented
code as part of the signal handling. Using an atomic exchange instead
fixes that by doing a proper mutual exclusion even in the presence of
signal handling.

Reviewers: kpw, eizan, jfb

Reviewed By: eizan

Subscribers: llvm-commits

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

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

4 files changed, 67 insertions, 36 deletions

diff --git a/lib/xray/tests/CMakeLists.txt b/lib/xray/tests/CMakeLists.txt
index 2f41e6ebc..2a65e118b 100644
--- a/lib/xray/tests/CMakeLists.txt
+++ b/lib/xray/tests/CMakeLists.txt
@@ -35,6 +35,7 @@ set(XRAY_IMPL_FILES
   ../../xray_profile_collector.h
   ../../xray_profiler_flags.cc
   ../../xray_profiler_flags.h
+  ../../xray_recursion_guard.h
   ../../xray_segmented_array.h
   ../../xray_trampoline_powerpc64.cc
   ../../xray_tsc.h
diff --git a/lib/xray/xray_basic_logging.cc b/lib/xray/xray_basic_logging.cc
index 525a77b42..9a6172beb 100644
--- a/lib/xray/xray_basic_logging.cc
+++ b/lib/xray/xray_basic_logging.cc
@@ -29,6 +29,7 @@
 #include "sanitizer_common/sanitizer_allocator_internal.h"
 #include "sanitizer_common/sanitizer_libc.h"
 #include "xray/xray_records.h"
+#include "xray_recursion_guard.h"
 #include "xray_basic_flags.h"
 #include "xray_basic_logging.h"
 #include "xray_defs.h"
@@ -70,7 +71,7 @@ static atomic_uint8_t BasicInitialized{0};
 
 BasicLoggingOptions GlobalOptions;
 
-thread_local volatile bool RecursionGuard = false;
+thread_local atomic_uint8_t Guard{0};
 
 static uint64_t thresholdTicks() XRAY_NEVER_INSTRUMENT {
   static uint64_t TicksPerSec = probeRequiredCPUFeatures()
@@ -165,10 +166,9 @@ void InMemoryRawLog(int32_t FuncId, XRayEntryType Type,
   // 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
   // thread.
-  if (RecursionGuard)
+  RecursionGuard G(Guard);
+  if (!G)
     return;
-  RecursionGuard = true;
-  auto ExitGuard = at_scope_exit([] { RecursionGuard = false; });
 
   uint8_t CPU = 0;
   uint64_t TSC = ReadTSC(CPU);
@@ -276,10 +276,9 @@ void InMemoryRawLogWithArg(int32_t FuncId, XRayEntryType Type, uint64_t Arg1,
   // Then we write the "we have an argument" record.
   InMemoryRawLog(FuncId, Type, ReadTSC);
 
-  if (RecursionGuard)
+  RecursionGuard G(Guard);
+  if (!G)
     return;
-  RecursionGuard = true;
-  auto ExitGuard = at_scope_exit([] { RecursionGuard = false; });
 
   // And from here on write the arg payload.
   XRayArgPayload R;
diff --git a/lib/xray/xray_fdr_logging.cc b/lib/xray/xray_fdr_logging.cc
index 610689ffc..cec083950 100644
--- a/lib/xray/xray_fdr_logging.cc
+++ b/lib/xray/xray_fdr_logging.cc
@@ -28,6 +28,7 @@
 #include "sanitizer_common/sanitizer_common.h"
 #include "xray/xray_interface.h"
 #include "xray/xray_records.h"
+#include "xray_recursion_guard.h"
 #include "xray_buffer_queue.h"
 #include "xray_defs.h"
 #include "xray_fdr_flags.h"
@@ -122,33 +123,6 @@ static ThreadLocalData &getThreadLocalData() {
   return TLD;
 }
 
-namespace {
-
-class RecursionGuard {
-  volatile bool &Running;
-  const bool Valid;
-
-public:
-  explicit RecursionGuard(volatile 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;
-  }
-};
-
-} // namespace
-
 static void writeNewBufferPreamble(tid_t Tid,
                                    timespec TS) XRAY_NEVER_INSTRUMENT {
   static constexpr int InitRecordsCount = 2;
@@ -553,7 +527,7 @@ static void endBufferIfFull() XRAY_NEVER_INSTRUMENT {
   }
 }
 
-thread_local volatile bool Running = false;
+thread_local atomic_uint8_t Running{0};
 
 /// 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
@@ -574,7 +548,7 @@ static void processFunctionHook(int32_t FuncId, XRayEntryType Entry,
   // handleArg0 to happen at any given time.
   RecursionGuard Guard{Running};
   if (!Guard) {
-    DCHECK(Running == true && "RecursionGuard is buggy!");
+    DCHECK(atomic_load_relaxed(&Running) && "RecursionGuard is buggy!");
     return;
   }
 
diff --git a/lib/xray/xray_recursion_guard.h b/lib/xray/xray_recursion_guard.h
new file mode 100644
index 000000000..6edadea56
--- /dev/null
+++ b/lib/xray/xray_recursion_guard.h
@@ -0,0 +1,57 @@
+//===-- xray_recursion_guard.h ---------------------------------*- 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.
+//
+//===----------------------------------------------------------------------===//
+#ifndef XRAY_XRAY_RECURSION_GUARD_H
+#define XRAY_XRAY_RECURSION_GUARD_H
+
+#include "sanitizer_common/sanitizer_atomic.h"
+
+namespace __xray {
+
+/// The RecursionGuard is useful for guarding against signal handlers which are
+/// also potentially calling XRay-instrumented functions. To use the
+/// RecursionGuard, you'll typically need a thread_local atomic_uint8_t:
+///
+///   thread_local atomic_uint8_t Guard{0};
+///
+///   // In a handler function:
+///   void handleArg0(int32_t F, XRayEntryType T) {
+///     RecursionGuard G(Guard);
+///     if (!G)
+///       return;  // Failed to acquire the guard.
+///     ...
+///   }
+///
+class RecursionGuard {
+  atomic_uint8_t &Running;
+  const bool Valid;
+
+public:
+  explicit inline RecursionGuard(atomic_uint8_t &R)
+      : Running(R), Valid(!atomic_exchange(&R, 1, memory_order_acq_rel)) {}
+
+  inline RecursionGuard(const RecursionGuard &) = delete;
+  inline RecursionGuard(RecursionGuard &&) = delete;
+  inline RecursionGuard &operator=(const RecursionGuard &) = delete;
+  inline RecursionGuard &operator=(RecursionGuard &&) = delete;
+
+  explicit inline operator bool() const { return Valid; }
+
+  inline ~RecursionGuard() noexcept {
+    if (Valid)
+      atomic_store(&Running, 0, memory_order_release);
+  }
+};
+
+} // namespace __xray
+
+#endif // XRAY_XRAY_RECURSION_GUARD_H