[compiler-rt][XRay] Address follow-up comments to initial interface and initialisation code

This addresses some comments from D21612, which contains the following changes:

- Update __xray_patch() and __xray_unpatch() API documentation to not imply asynchrony.
- Introduce a scope cleanup mechanism to make sure we can roll-back changes to the XRayPatching global atomic.
- Introduce a few more comments for potential extension points for other platforms (for the implementation details of patching and un-patching).

Reviewers: eugenis, rnk, kcc, echristo, majnemer

Subscribers: llvm-commits, mehdi_amini

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

llvm-svn: 277124

3 files changed, 54 insertions, 24 deletions

diff --git a/compiler-rt/include/xray/xray_interface.h b/compiler-rt/include/xray/xray_interface.h
index b8a0a61920b..22f137d110f 100644
--- a/compiler-rt/include/xray/xray_interface.h
+++ b/compiler-rt/include/xray/xray_interface.h
@@ -41,26 +41,18 @@ extern int __xray_remove_handler();
 
 enum XRayPatchingStatus {
   NOT_INITIALIZED = 0,
-  NOTIFIED = 1,
+  SUCCESS = 1,
   ONGOING = 2,
-  FAILED = 3
+  FAILED = 3,
 };
 
-// This tells XRay to patch the instrumentation points. This is an asynchronous
-// process, and returns the following status in specific cases:
-//
-//   - 0 : XRay is not initialized.
-//   - 1 : We've done the notification.
-//   - 2 : Patching / un-patching is on-going.
+// This tells XRay to patch the instrumentation points. See XRayPatchingStatus
+// for possible result values.
 extern XRayPatchingStatus __xray_patch();
 
-// Reverses the effect of __xray_patch(). This is an asynchronous process, and
-// returns the following status in specific cases.
-//
-//   - 0 : XRay is not initialized.
-//   - 1 : We've done the notification.
-//   - 2 : Patching / un-patching is on-going.
-extern int __xray_unpatch();
+// Reverses the effect of __xray_patch(). See XRayPatchingStatus for possible
+// result values.
+extern XRayPatchingStatus __xray_unpatch();
 }
 
 #endif
diff --git a/compiler-rt/lib/xray/xray_init.cc b/compiler-rt/lib/xray/xray_init.cc
index 8c2a5d380c8..e5dc0d9557d 100644
--- a/compiler-rt/lib/xray/xray_init.cc
+++ b/compiler-rt/lib/xray/xray_init.cc
@@ -22,18 +22,20 @@
 #include "xray_interface_internal.h"
 
 extern "C" {
-extern void __xray_init();
+void __xray_init();
 extern const XRaySledEntry __start_xray_instr_map[] __attribute__((weak));
 extern const XRaySledEntry __stop_xray_instr_map[] __attribute__((weak));
 }
 
 using namespace __xray;
 
-// We initialize some global variables that pertain to specific sections of XRay
-// data structures in the binary. We do this for the current process using
-// /proc/curproc/map and make sure that we're able to get it. We signal failure
-// via a global atomic boolean to indicate whether we've initialized properly.
+// When set to 'true' this means the XRay runtime has been initialised. We use
+// the weak symbols defined above (__start_xray_inst_map and
+// __stop_xray_instr_map) to initialise the instrumentation map that XRay uses
+// for runtime patching/unpatching of instrumentation points.
 //
+// FIXME: Support DSO instrumentation maps too. The current solution only works
+// for statically linked executables.
 std::atomic<bool> XRayInitialized{false};
 
 // This should always be updated before XRayInitialized is updated.
diff --git a/compiler-rt/lib/xray/xray_interface.cc b/compiler-rt/lib/xray/xray_interface.cc
index 474b47fae5f..ccbe7e2f2f7 100644
--- a/compiler-rt/lib/xray/xray_interface.cc
+++ b/compiler-rt/lib/xray/xray_interface.cc
@@ -14,6 +14,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "xray_interface_internal.h"
+
 #include <atomic>
 #include <cstdint>
 #include <cstdio>
@@ -21,6 +22,8 @@
 #include <limits>
 #include <sys/mman.h>
 
+#include "sanitizer_common/sanitizer_common.h"
+
 namespace __xray {
 
 // This is the function to call when we encounter the entry or exit sleds.
@@ -83,8 +86,25 @@ std::atomic<bool> XRayPatching{false};
 
 using namespace __xray;
 
+// FIXME: Figure out whether we can move this class to sanitizer_common instead
+// as a generic "scope guard".
+template <class Function> class CleanupInvoker {
+  Function Fn;
+
+public:
+  explicit CleanupInvoker(Function Fn) : Fn(Fn) {}
+  CleanupInvoker(const CleanupInvoker &) = default;
+  CleanupInvoker(CleanupInvoker &&) = default;
+  CleanupInvoker &operator=(const CleanupInvoker &) = delete;
+  CleanupInvoker &operator=(CleanupInvoker &&) = delete;
+  ~CleanupInvoker() { Fn(); }
+};
+
+template <class Function> CleanupInvoker<Function> ScopeCleanup(Function Fn) {
+  return CleanupInvoker<Function>{Fn};
+}
+
 XRayPatchingStatus __xray_patch() {
-  // FIXME: Make this happen asynchronously. For now just do this sequentially.
   if (!XRayInitialized.load(std::memory_order_acquire))
     return XRayPatchingStatus::NOT_INITIALIZED; // Not initialized.
 
@@ -95,6 +115,13 @@ XRayPatchingStatus __xray_patch() {
     return XRayPatchingStatus::ONGOING; // Already patching.
   }
 
+  bool PatchingSuccess = false;
+  auto XRayPatchingStatusResetter = ScopeCleanup([&PatchingSuccess] {
+    if (!PatchingSuccess) {
+      XRayPatching.store(false, std::memory_order_release);
+    }
+  });
+
   // Step 1: Compute the function id, as a unique identifier per function in the
   // instrumentation map.
   XRaySledMap InstrMap = XRayInstrMap.load(std::memory_order_acquire);
@@ -131,6 +158,7 @@ XRayPatchingStatus __xray_patch() {
     static constexpr int64_t MinOffset{std::numeric_limits<int32_t>::min()};
     static constexpr int64_t MaxOffset{std::numeric_limits<int32_t>::max()};
     if (Sled.Kind == XRayEntryType::ENTRY) {
+      // FIXME: Implement this in a more extensible manner, per-platform.
       // Here we do the dance of replacing the following sled:
       //
       // xray_sled_n:
@@ -157,7 +185,10 @@ XRayPatchingStatus __xray_patch() {
           reinterpret_cast<int64_t>(__xray_FunctionEntry) -
           (static_cast<int64_t>(Sled.Address) + 11);
       if (TrampolineOffset < MinOffset || TrampolineOffset > MaxOffset) {
-        // FIXME: Print out an error here.
+        Report("XRay Entry trampoline (%p) too far from sled (%p); distance = "
+               "%ld\n",
+               __xray_FunctionEntry, reinterpret_cast<void *>(Sled.Address),
+               TrampolineOffset);
         continue;
       }
       *reinterpret_cast<uint32_t *>(Sled.Address + 2) = FuncId;
@@ -169,6 +200,7 @@ XRayPatchingStatus __xray_patch() {
     }
 
     if (Sled.Kind == XRayEntryType::EXIT) {
+      // FIXME: Implement this in a more extensible manner, per-platform.
       // Here we do the dance of replacing the following sled:
       //
       // xray_sled_n:
@@ -193,7 +225,10 @@ XRayPatchingStatus __xray_patch() {
           reinterpret_cast<int64_t>(__xray_FunctionExit) -
           (static_cast<int64_t>(Sled.Address) + 11);
       if (TrampolineOffset < MinOffset || TrampolineOffset > MaxOffset) {
-        // FIXME: Print out an error here.
+        Report("XRay Exit trampoline (%p) too far from sled (%p); distance = "
+               "%ld\n",
+               __xray_FunctionExit, reinterpret_cast<void *>(Sled.Address),
+               TrampolineOffset);
         continue;
       }
       *reinterpret_cast<uint32_t *>(Sled.Address + 2) = FuncId;
@@ -205,5 +240,6 @@ XRayPatchingStatus __xray_patch() {
     }
   }
   XRayPatching.store(false, std::memory_order_release);
-  return XRayPatchingStatus::NOTIFIED;
+  PatchingSuccess = true;
+  return XRayPatchingStatus::SUCCESS;
 }