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

[dberris]

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

[dberris]

• fixup: only update the patching atomic when the patching is actually done

[kpw]

Suggestion for future changes:

1. We could decouple mprotect and __xray_patch_function so that callers that want to mprotect once and then patch individual functions have their cake and eat it too.
2. Should xray make an effort to restore the mprotect status of pages after it's done? Probably not worth the complexity.
3. If these regions never lose write permissions, why even call mprotect for unpatching?

compiler-rt/lib/xray/xray_flags.inc

22–23

It's a bit unfortunate that there is now an undocumented requirement that code #including this file must have already imported sanitizer_internal_defs.h.

Perhaps you could add an #ifdef check for SANITIZER_DEFS_H or the uptr typedef itself with an #error assertion.

Is uptr even the right type? Page size doesn't represent a pointer, it's more like a size_t (although we might have to use it as an offset and add it to an address).

compiler-rt/lib/xray/xray_interface.cc

211–213

What evidence or reasoning leads us to conclude that this assumption holds (ignoring DSO)?

221

const auto &Sled?

228

Usage suggests defining the flag as size_t.

232

Consider listing this requirement in the flag description.

320

This function name is very non-descriptive.

__mprotect_and_patch_function()?

[dberris]

• fixup: Address review comments, refactor things a bit

[dberris]

In D41153#954307, @kpw wrote:

Suggestion for future changes:

1. We could decouple mprotect and __xray_patch_function so that callers that want to mprotect once and then patch individual functions have their cake and eat it too.

Yeah, I'd like to think about that a little more. I suspect the call to mprotect is an implementation detail of patching/unpatching.

2. Should xray make an effort to restore the mprotect status of pages after it's done? Probably not worth the complexity.

I think we do that already, no? MProtectHelper already does that in the destructor, to make the pages read+execute only as opposed to read+write+execute.

3. If these regions never lose write permissions, why even call mprotect for unpatching?

They do, MProtectHelper does RAII on the mprotect status of pages. It does make the assumption that the pages being mprotected are r+x to begin with.

compiler-rt/lib/xray/xray_flags.inc

22–23

Yeah, this file could only really be included from xray_flags.h and xray_flags.cc which should include the types. Fixed it from that side instead.

Unfortunately the flag parser doesn't support an overload for size_t reliably, and of the options we have there, uptr seems to fit the best.

compiler-rt/lib/xray/xray_interface.cc

211–213

This is the convention used by ELF and thus the Linux kernel to map the .text segment as contiguous virtual memory. If this is somehow not true, then our calls to mprotect will fail (i.e. if .text was somehow further segmented or non-contiguous). Is there a way to somehow improve this comment to make that clearer?

228

Unfortunately the flag parser may not know how to parse size_t but does know how to parse uptr.

320

Good point. I've moved this into the __xray namespace, put it in an anonymous namespace in there, and used a more descriptive name mprotectAndPatchFunction.

[kpw]

Thanks for pointing out the scope cleanup from MProtectHelper and making some improvements/clarifications.

compiler-rt/lib/xray/xray_flags.inc

 #error "Define XRAY_FLAG prior to including this file!"
 #endif
 
 XRAY_FLAG(bool, patch_premain, false,
           "Whether to patch instrumentation points before main.")
 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.")
 
 // Basic (Naive) Mode logging options.
 XRAY_FLAG(bool, xray_naive_log, false,
           "DEPRECATED: Use xray_mode=xray-basic instead.")
 XRAY_FLAG(int, xray_naive_log_func_duration_threshold_us, 5,
           "Naive logging will try to skip functions that execute for fewer "
           "microseconds than this threshold.")
 XRAY_FLAG(int, xray_naive_log_max_stack_depth, 64,

compiler-rt/lib/xray/xray_interface.cc

 #include <cstdint>
 #include <cstdio>
 #include <errno.h>
 #include <limits>
 #include <sys/mman.h>
 
 #include "sanitizer_common/sanitizer_common.h"
 #include "xray_defs.h"
+#include "xray_flags.h"
 
 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;
 #elif defined(__arm__)
 static const int16_t cSledLength = 28;
 #elif SANITIZER_MIPS32
 static const int16_t cSledLength = 48;
 #elif SANITIZER_MIPS64
if (__sanitizer::atomic_load(&XRayInitialized,
     __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);
 }
 

 
 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;
-  }
-
   bool Success = false;
   switch (Sled.Kind) {
   case XRayEntryType::ENTRY:
     Success = patchFunctionEntry(Enable, FuncId, Sled, __xray_FunctionEntry);
     break;
   case XRayEntryType::EXIT:
     Success = patchFunctionExit(Enable, FuncId, Sled);
     break;
XRayPatchingStatus controlPatching(bool Enable) XRAY_NEVER_INSTRUMENT {
 
   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.
   XRaySledMap InstrMap;
   {
     __sanitizer::SpinMutexLock Guard(&XRayInstrMapMutex);
     InstrMap = XRayInstrMap;
   }
   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++) {
+    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;
     if (F != CurFun) {
       ++FuncId;
       CurFun = F;
     }
     patchSled(Sled, Enable, FuncId);
   }
     __sanitizer::atomic_store(&XRayPatching, false,
                               __sanitizer::memory_order_release);
   PatchingSuccess = true;
   return XRayPatchingStatus::SUCCESS;
 }
 
 XRayPatchingStatus __xray_patch() XRAY_NEVER_INSTRUMENT {
   return controlPatching(true);
 }
 
XRayPatchingStatus patchFunction(int32_t FuncId,
   if (!SucceedOnce) {
     Report("Failed patching any sled for function '%d'.", FuncId);
     return XRayPatchingStatus::FAILED;
   }
 
   return XRayPatchingStatus::SUCCESS;
 }
 
+static XRayPatchingStatus __xray_patch_helper(int32_t FuncId, bool Enable) {
+  XRaySledMap InstrMap;
+  {
+    __sanitizer::SpinMutexLock Guard(&XRayInstrMapMutex);
+    InstrMap = XRayInstrMap;
+  }
+
+  // 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;
+  }
+
+  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;
+  }
+
+  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);
+}
+
 XRayPatchingStatus __xray_patch_function(int32_t FuncId) XRAY_NEVER_INSTRUMENT {
-  return patchFunction(FuncId, true);
+  return __xray_patch_helper(FuncId, true);
 }
 
 XRayPatchingStatus
 __xray_unpatch_function(int32_t FuncId) XRAY_NEVER_INSTRUMENT {
-  return patchFunction(FuncId, false);
+  return __xray_patch_helper(FuncId, false);
 }
 
 int __xray_set_handler_arg1(void (*entry)(int32_t, XRayEntryType, uint64_t)) {
   if (!__sanitizer::atomic_load(&XRayInitialized,
                                 __sanitizer::memory_order_acquire))
     return 0;
 
   // A relaxed write might not be visible even if the current thread gets

compiler-rt/test/xray/TestCases/Linux/coverage-sample.cc

 // Check that we can patch and unpatch specific function ids.
 //
 // RUN: %clangxx_xray -std=c++11 %s -o %t
 // RUN: XRAY_OPTIONS="patch_premain=false xray_naive_log=false" %run %t | FileCheck %s
 
 // UNSUPPORTED: target-is-mips64,target-is-mips64el
 
 #include "xray/xray_interface.h"
 
 #include <set>
 #include <cstdio>
+#include <cassert>
 
 std::set<int32_t> function_ids;
 
 [[clang::xray_never_instrument]] void coverage_handler(int32_t fid,
                                                        XRayEntryType) {
   thread_local bool patching = false;
   if (patching) return;
   patching = true;

 }
 
 [[clang::xray_always_instrument]] void foo() {
   bar();
 }
 
 [[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");
   for (const auto id : function_ids)
     printf("patched: %d\n", id);
 
   // CHECK-LABEL: first pass.
   // CHECK-DAG: patched: [[F1:.*]]
   // CHECK-DAG: patched: [[F2:.*]]
   // CHECK-DAG: patched: [[F3:.*]]
 
   // make a copy of the function_ids, then patch them later.
   auto called_fns = function_ids;
 
   // clear the function_ids.
   function_ids.clear();
 
   // 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");
   for (const auto id : function_ids)
     printf("patched: %d\n", id);
   // CHECK-LABEL: second pass.
   // CHECK-DAG: patched: [[F1]]
   // CHECK-DAG: patched: [[F2]]
   // CHECK-DAG: patched: [[F3]]
 
   // 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");
   for (const auto id : function_ids)
     printf("patched: %d\n", id);
   // CHECK-LABEL: missing 1.
   // CHECK-NOT: patched: 1
 }