1 files changed, 40 insertions, 4 deletions

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;
 }