[XRay][compiler-rt] Coalesce calls to mprotect to reduce patching overhead

Summary:
Before this change, XRay would conservatively patch sections of the code
one sled at a time. Upon testing/profiling, this turns out to take an
inordinate amount of time and cycles. For an instrumented clang binary,
the cycles spent both in the patching/unpatching routine constituted 4%
of the cycles -- this didn't count the time spent in the kernel while
performing the mprotect calls in quick succession.

With this change, we're coalescing the number of calls to mprotect from
being linear to the number of instrumentation points, to now being a
lower constant when patching all the sleds through `__xray_patch()` or
`__xray_unpatch()`. In the case of calling `__xray_patch_function()` or
`__xray_unpatch_function()` we're now doing an mprotect call once for
all the sleds for that function (reduction of at least 2x calls to
mprotect).

Reviewers: kpw, eizan

Subscribers: llvm-commits

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

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

4 files changed, 195 insertions, 137 deletions

diff --git a/lib/xray/xray_flags.h b/lib/xray/xray_flags.h
index f4e30283b..3ed5b8844 100644
--- a/lib/xray/xray_flags.h
+++ b/lib/xray/xray_flags.h
@@ -16,6 +16,7 @@
 #define XRAY_FLAGS_H
 
 #include "sanitizer_common/sanitizer_flag_parser.h"
+#include "sanitizer_common/sanitizer_internal_defs.h"
 
 namespace __xray {
 
diff --git a/lib/xray/xray_flags.inc b/lib/xray/xray_flags.inc
index 934fb2986..29f1fce7d 100644
--- a/lib/xray/xray_flags.inc
+++ b/lib/xray/xray_flags.inc
@@ -19,7 +19,9 @@ XRAY_FLAG(bool, patch_premain, false,
 XRAY_FLAG(const char *, xray_logfile_base, "xray-log.",
           "Filename base for the xray logfile.")
 XRAY_FLAG(const char *, xray_mode, "", "Mode to install by default.")
-
+XRAY_FLAG(uptr, xray_page_size_override, 0,
+          "Override the default page size for the system, in bytes. The size "
+          "should be a power-of-two.")
 
 // Basic (Naive) Mode logging options.
 XRAY_FLAG(bool, xray_naive_log, false,
diff --git a/lib/xray/xray_interface.cc b/lib/xray/xray_interface.cc
index 7ad6a9b84..766313e85 100644
--- a/lib/xray/xray_interface.cc
+++ b/lib/xray/xray_interface.cc
@@ -23,12 +23,15 @@
 
 #include "sanitizer_common/sanitizer_common.h"
 #include "xray_defs.h"
+#include "xray_flags.h"
+
+extern __sanitizer::SpinMutex XRayInstrMapMutex;
+extern __sanitizer::atomic_uint8_t XRayInitialized;
+extern __xray::XRaySledMap XRayInstrMap;
 
 namespace __xray {
 
 #if defined(__x86_64__)
-// FIXME: The actual length is 11 bytes. Why was length 12 passed to mprotect()
-// ?
 static const int16_t cSledLength = 12;
 #elif defined(__aarch64__)
 static const int16_t cSledLength = 32;
@@ -53,6 +56,10 @@ __sanitizer::atomic_uintptr_t XRayArgLogger{0};
 // This is the function to call when we encounter a custom event log call.
 __sanitizer::atomic_uintptr_t XRayPatchedCustomEvent{0};
 
+// This is the global status to determine whether we are currently
+// patching/unpatching.
+__sanitizer::atomic_uint8_t XRayPatching{0};
+
 // 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
@@ -88,85 +95,10 @@ public:
   }
 };
 
-} // namespace __xray
-
-extern __sanitizer::SpinMutex XRayInstrMapMutex;
-extern __sanitizer::atomic_uint8_t XRayInitialized;
-extern __xray::XRaySledMap XRayInstrMap;
-
-int __xray_set_handler(void (*entry)(int32_t,
-                                     XRayEntryType)) XRAY_NEVER_INSTRUMENT {
-  if (__sanitizer::atomic_load(&XRayInitialized,
-                               __sanitizer::memory_order_acquire)) {
-
-    __sanitizer::atomic_store(&__xray::XRayPatchedFunction,
-                              reinterpret_cast<uintptr_t>(entry),
-                              __sanitizer::memory_order_release);
-    return 1;
-  }
-  return 0;
-}
-
-int __xray_set_customevent_handler(void (*entry)(void *, size_t))
-    XRAY_NEVER_INSTRUMENT {
-  if (__sanitizer::atomic_load(&XRayInitialized,
-                               __sanitizer::memory_order_acquire)) {
-    __sanitizer::atomic_store(&__xray::XRayPatchedCustomEvent,
-                              reinterpret_cast<uintptr_t>(entry),
-                              __sanitizer::memory_order_release);
-    return 1;
-  }
-  return 0;
-}
-
-
-int __xray_remove_handler() XRAY_NEVER_INSTRUMENT {
-  return __xray_set_handler(nullptr);
-}
-
-int __xray_remove_customevent_handler() XRAY_NEVER_INSTRUMENT {
-  return __xray_set_customevent_handler(nullptr);
-}
-
-__sanitizer::atomic_uint8_t XRayPatching{0};
-
-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) XRAY_NEVER_INSTRUMENT : Fn(Fn) {}
-  CleanupInvoker(const CleanupInvoker &) XRAY_NEVER_INSTRUMENT = default;
-  CleanupInvoker(CleanupInvoker &&) XRAY_NEVER_INSTRUMENT = default;
-  CleanupInvoker &
-  operator=(const CleanupInvoker &) XRAY_NEVER_INSTRUMENT = delete;
-  CleanupInvoker &operator=(CleanupInvoker &&) XRAY_NEVER_INSTRUMENT = delete;
-  ~CleanupInvoker() XRAY_NEVER_INSTRUMENT { Fn(); }
-};
-
-template <class Function>
-CleanupInvoker<Function> scopeCleanup(Function Fn) XRAY_NEVER_INSTRUMENT {
-  return CleanupInvoker<Function>{Fn};
-}
-
-inline bool patchSled(const XRaySledEntry &Sled, bool Enable,
-                      int32_t FuncId) XRAY_NEVER_INSTRUMENT {
-  // While we're here, we should patch the nop sled. To do that we mprotect
-  // the page containing the function to be writeable.
-  const uint64_t PageSize = GetPageSizeCached();
-  void *PageAlignedAddr =
-      reinterpret_cast<void *>(Sled.Address & ~(PageSize - 1));
-  std::size_t MProtectLen = (Sled.Address + cSledLength) -
-                            reinterpret_cast<uint64_t>(PageAlignedAddr);
-  MProtectHelper Protector(PageAlignedAddr, MProtectLen);
-  if (Protector.MakeWriteable() == -1) {
-    printf("Failed mprotect: %d\n", errno);
-    return XRayPatchingStatus::FAILED;
-  }
+namespace {
 
+bool patchSled(const XRaySledEntry &Sled, bool Enable,
+               int32_t FuncId) XRAY_NEVER_INSTRUMENT {
   bool Success = false;
   switch (Sled.Kind) {
   case XRayEntryType::ENTRY:
@@ -191,6 +123,55 @@ inline bool patchSled(const XRaySledEntry &Sled, bool Enable,
   return Success;
 }
 
+XRayPatchingStatus patchFunction(int32_t FuncId,
+                                 bool Enable) XRAY_NEVER_INSTRUMENT {
+  if (!__sanitizer::atomic_load(&XRayInitialized,
+                                __sanitizer::memory_order_acquire))
+    return XRayPatchingStatus::NOT_INITIALIZED; // Not initialized.
+
+  uint8_t NotPatching = false;
+  if (!__sanitizer::atomic_compare_exchange_strong(
+          &XRayPatching, &NotPatching, true, __sanitizer::memory_order_acq_rel))
+    return XRayPatchingStatus::ONGOING; // Already patching.
+
+  // Next, we look for the function index.
+  XRaySledMap InstrMap;
+  {
+    __sanitizer::SpinMutexLock Guard(&XRayInstrMapMutex);
+    InstrMap = XRayInstrMap;
+  }
+
+  // If we don't have an index, we can't patch individual functions.
+  if (InstrMap.Functions == 0)
+    return XRayPatchingStatus::NOT_INITIALIZED;
+
+  // FuncId must be a positive number, less than the number of functions
+  // instrumented.
+  if (FuncId <= 0 || static_cast<size_t>(FuncId) > InstrMap.Functions) {
+    Report("Invalid function id provided: %d\n", FuncId);
+    return XRayPatchingStatus::FAILED;
+  }
+
+  // Now we patch ths sleds for this specific function.
+  auto SledRange = InstrMap.SledsIndex[FuncId - 1];
+  auto *f = SledRange.Begin;
+  auto *e = SledRange.End;
+
+  bool SucceedOnce = false;
+  while (f != e)
+    SucceedOnce |= patchSled(*f++, Enable, FuncId);
+
+  __sanitizer::atomic_store(&XRayPatching, false,
+                            __sanitizer::memory_order_release);
+
+  if (!SucceedOnce) {
+    Report("Failed patching any sled for function '%d'.", FuncId);
+    return XRayPatchingStatus::FAILED;
+  }
+
+  return XRayPatchingStatus::SUCCESS;
+}
+
 // controlPatching implements the common internals of the patching/unpatching
 // implementation. |Enable| defines whether we're enabling or disabling the
 // runtime XRay instrumentation.
@@ -205,14 +186,13 @@ XRayPatchingStatus controlPatching(bool Enable) XRAY_NEVER_INSTRUMENT {
     return XRayPatchingStatus::ONGOING; // Already patching.
 
   uint8_t PatchingSuccess = false;
-  auto XRayPatchingStatusResetter = scopeCleanup([&PatchingSuccess] {
-    if (!PatchingSuccess)
-      __sanitizer::atomic_store(&XRayPatching, false,
-                                __sanitizer::memory_order_release);
-  });
-
-  // Step 1: Compute the function id, as a unique identifier per function in the
-  // instrumentation map.
+  auto XRayPatchingStatusResetter =
+      __sanitizer::at_scope_exit([&PatchingSuccess] {
+        if (!PatchingSuccess)
+          __sanitizer::atomic_store(&XRayPatching, false,
+                                    __sanitizer::memory_order_release);
+      });
+
   XRaySledMap InstrMap;
   {
     __sanitizer::SpinMutexLock Guard(&XRayInstrMapMutex);
@@ -221,16 +201,47 @@ XRayPatchingStatus controlPatching(bool Enable) XRAY_NEVER_INSTRUMENT {
   if (InstrMap.Entries == 0)
     return XRayPatchingStatus::NOT_INITIALIZED;
 
-  const uint64_t PageSize = GetPageSizeCached();
+  uint32_t FuncId = 1;
+  uint64_t CurFun = 0;
+
+  // First we want to find the bounds for which we have instrumentation points,
+  // and try to get as few calls to mprotect(...) as possible. We're assuming
+  // that all the sleds for the instrumentation map are contiguous as a single
+  // set of pages. When we do support dynamic shared object instrumentation,
+  // we'll need to do this for each set of page load offsets per DSO loaded. For
+  // now we're assuming we can mprotect the whole section of text between the
+  // minimum sled address and the maximum sled address (+ the largest sled
+  // size).
+  auto MinSled = InstrMap.Sleds[0];
+  auto MaxSled = InstrMap.Sleds[InstrMap.Entries - 1];
+  for (std::size_t I = 0; I < InstrMap.Entries; I++) {
+    const auto &Sled = InstrMap.Sleds[I];
+    if (Sled.Address < MinSled.Address)
+      MinSled = Sled;
+    if (Sled.Address > MaxSled.Address)
+      MaxSled = Sled;
+  }
+
+  const size_t PageSize = flags()->xray_page_size_override > 0
+                              ? flags()->xray_page_size_override
+                              : GetPageSizeCached();
   if ((PageSize == 0) || ((PageSize & (PageSize - 1)) != 0)) {
     Report("System page size is not a power of two: %lld\n", PageSize);
     return XRayPatchingStatus::FAILED;
   }
 
-  uint32_t FuncId = 1;
-  uint64_t CurFun = 0;
-  for (std::size_t I = 0; I < InstrMap.Entries; I++) {
-    auto Sled = InstrMap.Sleds[I];
+  void *PageAlignedAddr =
+      reinterpret_cast<void *>(MinSled.Address & ~(PageSize - 1));
+  size_t MProtectLen =
+      (MaxSled.Address - reinterpret_cast<uptr>(PageAlignedAddr)) + cSledLength;
+  MProtectHelper Protector(PageAlignedAddr, MProtectLen);
+  if (Protector.MakeWriteable() == -1) {
+    Report("Failed mprotect: %d\n", errno);
+    return XRayPatchingStatus::FAILED;
+  }
+
+  for (std::size_t I = 0; I < InstrMap.Entries; ++I) {
+    auto &Sled = InstrMap.Sleds[I];
     auto F = Sled.Function;
     if (CurFun == 0)
       CurFun = F;
@@ -246,36 +257,14 @@ XRayPatchingStatus controlPatching(bool Enable) XRAY_NEVER_INSTRUMENT {
   return XRayPatchingStatus::SUCCESS;
 }
 
-XRayPatchingStatus __xray_patch() XRAY_NEVER_INSTRUMENT {
-  return controlPatching(true);
-}
-
-XRayPatchingStatus __xray_unpatch() XRAY_NEVER_INSTRUMENT {
-  return controlPatching(false);
-}
-
-XRayPatchingStatus patchFunction(int32_t FuncId,
-                                 bool Enable) XRAY_NEVER_INSTRUMENT {
-  if (!__sanitizer::atomic_load(&XRayInitialized,
-                                __sanitizer::memory_order_acquire))
-    return XRayPatchingStatus::NOT_INITIALIZED; // Not initialized.
-
-  uint8_t NotPatching = false;
-  if (!__sanitizer::atomic_compare_exchange_strong(
-          &XRayPatching, &NotPatching, true, __sanitizer::memory_order_acq_rel))
-    return XRayPatchingStatus::ONGOING; // Already patching.
-
-  // Next, we look for the function index.
+XRayPatchingStatus mprotectAndPatchFunction(int32_t FuncId,
+                                            bool Enable) XRAY_NEVER_INSTRUMENT {
   XRaySledMap InstrMap;
   {
     __sanitizer::SpinMutexLock Guard(&XRayInstrMapMutex);
     InstrMap = XRayInstrMap;
   }
 
-  // If we don't have an index, we can't patch individual functions.
-  if (InstrMap.Functions == 0)
-    return XRayPatchingStatus::NOT_INITIALIZED;
-
   // FuncId must be a positive number, less than the number of functions
   // instrumented.
   if (FuncId <= 0 || static_cast<size_t>(FuncId) > InstrMap.Functions) {
@@ -283,33 +272,98 @@ XRayPatchingStatus patchFunction(int32_t FuncId,
     return XRayPatchingStatus::FAILED;
   }
 
-  // Now we patch ths sleds for this specific function.
+  const size_t PageSize = flags()->xray_page_size_override > 0
+                              ? flags()->xray_page_size_override
+                              : GetPageSizeCached();
+  if ((PageSize == 0) || ((PageSize & (PageSize - 1)) != 0)) {
+    Report("Provided page size is not a power of two: %lld\n", PageSize);
+    return XRayPatchingStatus::FAILED;
+  }
+
+  // Here we compute the minumum sled and maximum sled associated with a
+  // particular function ID.
   auto SledRange = InstrMap.SledsIndex[FuncId - 1];
   auto *f = SledRange.Begin;
   auto *e = SledRange.End;
+  auto MinSled = *f;
+  auto MaxSled = *(SledRange.End - 1);
+  while (f != e) {
+    if (f->Address < MinSled.Address)
+      MinSled = *f;
+    if (f->Address > MaxSled.Address)
+      MaxSled = *f;
+    ++f;
+  }
 
-  bool SucceedOnce = false;
-  while (f != e)
-    SucceedOnce |= patchSled(*f++, Enable, FuncId);
+  void *PageAlignedAddr =
+      reinterpret_cast<void *>(MinSled.Address & ~(PageSize - 1));
+  size_t MProtectLen =
+      (MaxSled.Address - reinterpret_cast<uptr>(PageAlignedAddr)) + cSledLength;
+  MProtectHelper Protector(PageAlignedAddr, MProtectLen);
+  if (Protector.MakeWriteable() == -1) {
+    Report("Failed mprotect: %d\n", errno);
+    return XRayPatchingStatus::FAILED;
+  }
+  return patchFunction(FuncId, Enable);
+}
 
-  __sanitizer::atomic_store(&XRayPatching, false,
-                            __sanitizer::memory_order_release);
+} // namespace
 
-  if (!SucceedOnce) {
-    Report("Failed patching any sled for function '%d'.", FuncId);
-    return XRayPatchingStatus::FAILED;
+} // namespace __xray
+
+using namespace __xray;
+
+// The following functions are declared `extern "C" {...}` in the header, hence
+// they're defined in the global namespace.
+
+int __xray_set_handler(void (*entry)(int32_t,
+                                     XRayEntryType)) XRAY_NEVER_INSTRUMENT {
+  if (__sanitizer::atomic_load(&XRayInitialized,
+                               __sanitizer::memory_order_acquire)) {
+
+    __sanitizer::atomic_store(&__xray::XRayPatchedFunction,
+                              reinterpret_cast<uintptr_t>(entry),
+                              __sanitizer::memory_order_release);
+    return 1;
   }
+  return 0;
+}
 
-  return XRayPatchingStatus::SUCCESS;
+int __xray_set_customevent_handler(void (*entry)(void *, size_t))
+    XRAY_NEVER_INSTRUMENT {
+  if (__sanitizer::atomic_load(&XRayInitialized,
+                               __sanitizer::memory_order_acquire)) {
+    __sanitizer::atomic_store(&__xray::XRayPatchedCustomEvent,
+                              reinterpret_cast<uintptr_t>(entry),
+                              __sanitizer::memory_order_release);
+    return 1;
+  }
+  return 0;
+}
+
+int __xray_remove_handler() XRAY_NEVER_INSTRUMENT {
+  return __xray_set_handler(nullptr);
+}
+
+int __xray_remove_customevent_handler() XRAY_NEVER_INSTRUMENT {
+  return __xray_set_customevent_handler(nullptr);
+}
+
+XRayPatchingStatus __xray_patch() XRAY_NEVER_INSTRUMENT {
+  return controlPatching(true);
+}
+
+XRayPatchingStatus __xray_unpatch() XRAY_NEVER_INSTRUMENT {
+  return controlPatching(false);
 }
 
 XRayPatchingStatus __xray_patch_function(int32_t FuncId) XRAY_NEVER_INSTRUMENT {
-  return patchFunction(FuncId, true);
+  return mprotectAndPatchFunction(FuncId, true);
 }
 
 XRayPatchingStatus
 __xray_unpatch_function(int32_t FuncId) XRAY_NEVER_INSTRUMENT {
-  return patchFunction(FuncId, false);
+  return mprotectAndPatchFunction(FuncId, false);
 }
 
 int __xray_set_handler_arg1(void (*entry)(int32_t, XRayEntryType, uint64_t)) {
diff --git a/test/xray/TestCases/Linux/coverage-sample.cc b/test/xray/TestCases/Linux/coverage-sample.cc
index 9ec22dcd6..62c13ba3d 100644
--- a/test/xray/TestCases/Linux/coverage-sample.cc
+++ b/test/xray/TestCases/Linux/coverage-sample.cc
@@ -9,6 +9,7 @@
 
 #include <set>
 #include <cstdio>
+#include <cassert>
 
 std::set<int32_t> function_ids;
 
@@ -36,9 +37,9 @@ std::set<int32_t> function_ids;
 
 [[clang::xray_always_instrument]] int main(int argc, char *argv[]) {
   __xray_set_handler(coverage_handler);
-  __xray_patch();
+  assert(__xray_patch() == XRayPatchingStatus::SUCCESS);
   foo();
-  __xray_unpatch();
+  assert(__xray_unpatch() == XRayPatchingStatus::SUCCESS);
 
   // print out the function_ids.
   printf("first pass.\n");
@@ -58,11 +59,11 @@ std::set<int32_t> function_ids;
 
   // patch the functions we've called before.
   for (const auto id : called_fns)
-    __xray_patch_function(id);
+    assert(__xray_patch_function(id) == XRayPatchingStatus::SUCCESS);
 
   // then call them again.
   foo();
-  __xray_unpatch();
+  assert(__xray_unpatch() == XRayPatchingStatus::SUCCESS);
 
   // confirm that we've seen the same functions again.
   printf("second pass.\n");
@@ -76,10 +77,10 @@ std::set<int32_t> function_ids;
   // Now we want to make sure that if we unpatch one, that we're only going to
   // see two calls of the coverage_handler.
   function_ids.clear();
-  __xray_patch();
-  __xray_unpatch_function(1);
+  assert(__xray_patch() == XRayPatchingStatus::SUCCESS);
+  assert(__xray_unpatch_function(1) == XRayPatchingStatus::SUCCESS);
   foo();
-  __xray_unpatch();
+  assert(__xray_unpatch() == XRayPatchingStatus::SUCCESS);
 
   // confirm that we don't see function id one called anymore.
   printf("missing 1.\n");