[XRay][compiler-rt] Support patching/unpatching specific functions

Summary:
This change allows us to patch/unpatch specific functions using the
function ID. This is useful in cases where implementations might want to
do coverage-style, or more fine-grained control of which functions to
patch or un-patch at runtime.

Depends on D32693.

Reviewers: dblaikie, echristo, kpw

Subscribers: llvm-commits

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

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

5 files changed, 203 insertions, 33 deletions

diff --git a/include/xray/xray_interface.h b/include/xray/xray_interface.h
index 52a7e1d9e..e2d0a89d6 100644
--- a/include/xray/xray_interface.h
+++ b/include/xray/xray_interface.h
@@ -67,6 +67,14 @@ extern XRayPatchingStatus __xray_patch();
 // result values.
 extern XRayPatchingStatus __xray_unpatch();
 
+// This patches a specific function id. See XRayPatchingStatus for possible
+// result values.
+extern XRayPatchingStatus __xray_patch_function(int32_t FuncId);
+
+// This unpatches a specific function id. See XRayPatchingStatus for possible
+// result values.
+extern XRayPatchingStatus __xray_unpatch_function(int32_t FuncId);
+
 // Use XRay to log the first argument of each (instrumented) function call.
 // When this function exits, all threads will have observed the effect and
 // start logging their subsequent affected function calls (if patched).
diff --git a/lib/xray/xray_init.cc b/lib/xray/xray_init.cc
index 6f558d656..aa660baa9 100644
--- a/lib/xray/xray_init.cc
+++ b/lib/xray/xray_init.cc
@@ -25,6 +25,8 @@ extern "C" {
 void __xray_init();
 extern const XRaySledEntry __start_xray_instr_map[] __attribute__((weak));
 extern const XRaySledEntry __stop_xray_instr_map[] __attribute__((weak));
+extern const XRayFunctionSledIndex __start_xray_fn_idx[] __attribute__((weak));
+extern const XRayFunctionSledIndex __stop_xray_fn_idx[] __attribute__((weak));
 }
 
 using namespace __xray;
@@ -55,6 +57,8 @@ void __xray_init() XRAY_NEVER_INSTRUMENT {
     __sanitizer::SpinMutexLock Guard(&XRayInstrMapMutex);
     XRayInstrMap.Sleds = __start_xray_instr_map;
     XRayInstrMap.Entries = __stop_xray_instr_map - __start_xray_instr_map;
+    XRayInstrMap.SledsIndex = __start_xray_fn_idx;
+    XRayInstrMap.Functions = __stop_xray_fn_idx - __start_xray_fn_idx;
   }
   __sanitizer::atomic_store(&XRayInitialized, true,
                             __sanitizer::memory_order_release);
diff --git a/lib/xray/xray_interface.cc b/lib/xray/xray_interface.cc
index 26ec161fe..73cca80bc 100644
--- a/lib/xray/xray_interface.cc
+++ b/lib/xray/xray_interface.cc
@@ -132,12 +132,48 @@ 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;
+  case XRayEntryType::TAIL:
+    Success = patchFunctionTailExit(Enable, FuncId, Sled);
+    break;
+  case XRayEntryType::LOG_ARGS_ENTRY:
+    Success = patchFunctionEntry(Enable, FuncId, Sled, __xray_ArgLoggerEntry);
+    break;
+  default:
+    Report("Unsupported sled kind '%d' @%04x\n", Sled.Address, int(Sled.Kind));
+    return false;
+  }
+  return Success;
+}
+
 // controlPatching implements the common internals of the patching/unpatching
 // implementation. |Enable| defines whether we're enabling or disabling the
 // runtime XRay instrumentation.
 XRayPatchingStatus controlPatching(bool Enable) XRAY_NEVER_INSTRUMENT {
   if (!__sanitizer::atomic_load(&XRayInitialized,
-                               __sanitizer::memory_order_acquire))
+                                __sanitizer::memory_order_acquire))
     return XRayPatchingStatus::NOT_INITIALIZED; // Not initialized.
 
   uint8_t NotPatching = false;
@@ -179,38 +215,7 @@ XRayPatchingStatus controlPatching(bool Enable) XRAY_NEVER_INSTRUMENT {
       ++FuncId;
       CurFun = F;
     }
-
-    // While we're here, we should patch the nop sled. To do that we mprotect
-    // the page containing the function to be writeable.
-    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;
-    case XRayEntryType::TAIL:
-      Success = patchFunctionTailExit(Enable, FuncId, Sled);
-      break;
-    case XRayEntryType::LOG_ARGS_ENTRY:
-      Success = patchFunctionEntry(Enable, FuncId, Sled, __xray_ArgLoggerEntry);
-      break;
-    default:
-      Report("Unsupported sled kind: %d\n", int(Sled.Kind));
-      continue;
-    }
-    (void)Success;
+    patchSled(Sled, Enable, FuncId);
   }
   __sanitizer::atomic_store(&XRayPatching, false,
                             __sanitizer::memory_order_release);
@@ -226,6 +231,64 @@ 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.
+  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;
+}
+
+XRayPatchingStatus __xray_patch_function(int32_t FuncId) XRAY_NEVER_INSTRUMENT {
+  return patchFunction(FuncId, true);
+}
+
+XRayPatchingStatus
+__xray_unpatch_function(int32_t FuncId) XRAY_NEVER_INSTRUMENT {
+  return patchFunction(FuncId, false);
+}
+
 int __xray_set_handler_arg1(void (*Handler)(int32_t, XRayEntryType, uint64_t)) {
   if (!__sanitizer::atomic_load(&XRayInitialized,
                                 __sanitizer::memory_order_acquire))
diff --git a/lib/xray/xray_interface_internal.h b/lib/xray/xray_interface_internal.h
index 0e3a251f3..ef0c6b158 100644
--- a/lib/xray/xray_interface_internal.h
+++ b/lib/xray/xray_interface_internal.h
@@ -39,6 +39,11 @@ struct XRaySledEntry {
 #error "Unsupported word size."
 #endif
 };
+
+struct XRayFunctionSledIndex {
+  const XRaySledEntry* Begin;
+  const XRaySledEntry* End;
+};
 }
 
 namespace __xray {
@@ -46,6 +51,8 @@ namespace __xray {
 struct XRaySledMap {
   const XRaySledEntry *Sleds;
   size_t Entries;
+  const XRayFunctionSledIndex *SledsIndex;
+  size_t Functions;
 };
 
 bool patchFunctionEntry(bool Enable, uint32_t FuncId,
diff --git a/test/xray/TestCases/Linux/coverage-sample.cc b/test/xray/TestCases/Linux/coverage-sample.cc
new file mode 100644
index 000000000..623b4e345
--- /dev/null
+++ b/test/xray/TestCases/Linux/coverage-sample.cc
@@ -0,0 +1,88 @@
+// 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
+
+#include "xray/xray_interface.h"
+
+#include <set>
+#include <cstdio>
+
+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;
+  function_ids.insert(fid);
+  __xray_unpatch_function(fid);
+  patching = false;
+}
+
+[[clang::xray_always_instrument]] void baz() {
+  // do nothing!
+}
+
+[[clang::xray_always_instrument]] void bar() {
+  baz();
+}
+
+[[clang::xray_always_instrument]] void foo() {
+  bar();
+}
+
+[[clang::xray_always_instrument]] int main(int argc, char *argv[]) {
+  __xray_set_handler(coverage_handler);
+  __xray_patch();
+  foo();
+  __xray_unpatch();
+
+  // 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);
+
+  // then call them again.
+  foo();
+  __xray_unpatch();
+
+  // 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);
+  foo();
+  __xray_unpatch();
+
+  // 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
+}