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

[dberris]

Summary

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).

[eugenis]

lib/xray/xray_init.cc

39

This only works for the main executable, right? Is this a temporary limitation? I imagine you could emit constructors to all DSOs that would add per-DSO sled maps to the global list.

[dberris]

• Add FIXME for supporting instrumentation maps in DSOs

lib/xray/xray_init.cc

39

Currently, yes. I haven't quite thought out a solution for DSO's as it hasn't come up yet.

But as you mention, yes we can do something special for DSO's that would be treated especially by the runtime. I've added a FIXME to indicate that.

include/xray/xray_interface.h

 extern int __xray_set_handler(void (*entry)(int32_t, XRayEntryType));
 
 // This removes whatever the currently provided handler is. Returns 1 on
 // success, 0 on error.
 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

lib/xray/xray_init.cc

 #include <strings.h>
 #include <unistd.h>
 
 #include "sanitizer_common/sanitizer_common.h"
 #include "xray_flags.h"
 #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.
 std::atomic<__xray::XRaySledMap> XRayInstrMap{};
 
 // __xray_init() will do the actual loading of the current process' memory map
 // and then proceed to look for the .xray_instr_map section/segment.
 void __xray_init() {
   InitializeFlags();

lib/xray/xray_interface.cc

 //===-- xray_interface.cpp --------------------------------------*- 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.
 //
 // Implementation of the API functions.
 //
 //===----------------------------------------------------------------------===//
 
 #include "xray_interface_internal.h"
+
 #include <atomic>
 #include <cstdint>
 #include <cstdio>
 #include <errno.h>
 #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.
 std::atomic<void (*)(int32_t, XRayEntryType)> XRayPatchedFunction{nullptr};
 
 // MProtectHelper is an RAII wrapper for calls to mprotect(...) that will undo
 // any successful mprotect(...) changes. This is used to make a page writeable
 // and executable, and upon destruction if it was successful in doing so returns
int __xray_set_handler(void (*entry)(int32_t, XRayEntryType)) {
   }
   return 0;
 }
 
 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.
 
   static bool NotPatching = false;
   if (!XRayPatching.compare_exchange_strong(NotPatching, true,
                                             std::memory_order_acq_rel,
                                             std::memory_order_acquire)) {
     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);
   if (InstrMap.Entries == 0)
     return XRayPatchingStatus::NOT_INITIALIZED;
 
   int32_t FuncId = 1;
   static constexpr uint8_t CallOpCode = 0xe8;
for (std::size_t I = 0; I < InstrMap.Entries; I++) {
     if (Protector.MakeWriteable() == -1) {
       printf("Failed mprotect: %d\n", errno);
       return XRayPatchingStatus::FAILED;
     }
 
     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:
       //   jmp +9
       //   <9 byte nop>
       //
       // With the following:
       //
if (Sled.Kind == XRayEntryType::ENTRY) {
       // opcode and first operand.
       //
       // Prerequisite is to compute the relative offset to the
       // __xray_FunctionEntry function's address.
       int64_t TrampolineOffset =
           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;
       *reinterpret_cast<uint8_t *>(Sled.Address + 6) = CallOpCode;
       *reinterpret_cast<uint32_t *>(Sled.Address + 7) = TrampolineOffset;
       std::atomic_store_explicit(
           reinterpret_cast<std::atomic<uint16_t> *>(Sled.Address), MovR10Seq,
           std::memory_order_release);
     }
 
     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:
       //   ret
       //   <10 byte nop>
       //
       // With the following:
       //
       //   mov r10d, <function id>
       //   jmp <relative 32bit offset to exit trampoline>
       //
       // 1. Put the function id first, 2 bytes from the start of the sled (just
       // after the 1-byte ret instruction).
       // 2. Put the jmp opcode 6 bytes from the start of the sled.
       // 3. Put the relative offset 7 bytes from the start of the sled.
       // 4. Do an atomic write over the jmp instruction for the "mov r10d"
       // opcode and first operand.
       //
       // Prerequisite is to compute the relative offset fo the
       // __xray_FunctionExit function's address.
       int64_t TrampolineOffset =
           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;
       *reinterpret_cast<uint8_t *>(Sled.Address + 6) = JmpOpCode;
       *reinterpret_cast<uint32_t *>(Sled.Address + 7) = TrampolineOffset;
       std::atomic_store_explicit(
           reinterpret_cast<std::atomic<uint16_t> *>(Sled.Address), MovR10Seq,
           std::memory_order_release);
     }
   }
   XRayPatching.store(false, std::memory_order_release);
-  return XRayPatchingStatus::NOTIFIED;
+  PatchingSuccess = true;
+  return XRayPatchingStatus::SUCCESS;
 }