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

[dberris]

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.

[dberris]

Remember to use the mutex/atomic that guards the patching from clobbering each other.

[dblaikie]

Looks good - maybe the test case could be improved. Provided some ideas.

lib/xray/xray_interface.cc

251

'unpatch' or should this be 'patch' (since the function is called 'patchSpecificFunction')?

264–266

any chance the InstrMap data structure could be typed in this way to begin with, to avoid the need for casts here? (also - should const be added here? I guess maybe not, since these pointers point to the sleds which will be modified)

268–269

Does the GCC implementation support partial patching failures/success/etc?

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

56–58

I'd consider changing the test to patch a specific function - as it stands it patches them all, which doesn't look all that different to calling the __xray_patch() function.

Also - what's the purpose of unpatching each function as its called in the handler? Would the handler behave differently/fail the test, etc, if the __xray_unpatch_function wasn't called there.

Maybe start with them all patched so you can record the function IDs, then unpatch one, test that it wasn't visited. Unpatch all+repatch one, test that only that one was visited?

[dberris]

• fixup: address comments, make better tests/semantics
• fixup: in 32-bit platforms, we need to pad the struct to be 16-bytes

[dberris]

lib/xray/xray_interface.cc

264–266

It seems like we could just use a pointer directly, good point. Also yes, made it a pointer to a const XRaySledEntry.

268–269

Actually now that I think about it, we only report errors and keep going anyway (this happens when we encounter a sled we don't know how to patch). It might happen for example when we add a new sled type and some object files have this new sled type but they're linked against an old version of the XRay library.

Updated it now to only fail if we didn't get to patch any sled for the function. Also added some diagnostics for debug-ability.

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

56–58

Good idea -- added a third case where we explicitly un-patch function id 1.

include/xray/xray_interface.h

 // This tells XRay to patch the instrumentation points. See XRayPatchingStatus
 // for possible result values.
 extern XRayPatchingStatus __xray_patch();
 
 // Reverses the effect of __xray_patch(). See XRayPatchingStatus for possible
 // 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).
 //
 // Returns 1 on success, 0 on error.
 extern int __xray_set_handler_arg1(void (*)(int32_t, XRayEntryType, uint64_t));
 
 // Disables the XRay handler used to log first arguments of function calls.
 // Returns 1 on success, 0 on error.
 extern int __xray_remove_handler_arg1();
 }
 
 #endif

lib/xray/xray_init.cc

 #include "xray_defs.h"
 #include "xray_flags.h"
 #include "xray_interface_internal.h"
 
 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;
 
 // 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.
if (__start_xray_instr_map == nullptr) {
     Report("XRay instrumentation map missing. Not initializing XRay.\n");
     return;
   }
 
   {
     __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);
 
   if (flags()->patch_premain)
     __xray_patch();
 }
 
 __attribute__((section(".preinit_array"),
                used)) void (*__local_xray_preinit)(void) = __xray_init;

lib/xray/xray_interface.cc

public:
   ~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;
+  }
+
+  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));
+    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))
     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.
 
   uint8_t PatchingSuccess = false;
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;
     }
-
-    // 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);
   PatchingSuccess = true;
   return XRayPatchingStatus::SUCCESS;
 }
 
 XRayPatchingStatus __xray_patch() XRAY_NEVER_INSTRUMENT {
   return controlPatching(true);
 }
 
 XRayPatchingStatus __xray_unpatch() XRAY_NEVER_INSTRUMENT {
   return controlPatching(false);
 }
 
+XRayPatchingStatus patchSpecificFunction(int32_t FuncId, bool Enable) XRAY_NEVER_INSTRUMENT {
+  if (!__sanitizer::atomic_load(&XRayInitialized,
+                                __sanitizer::memory_order_acquire))
+    return XRayPatchingStatus::NOT_INITIALIZED; // Not initialized.
+
+  // 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 unpatch individual functions.
+  if (InstrMap.Functions == 0)
+    return XRayPatchingStatus::NOT_INITIALIZED;
+
+  // FuncId may not be zero or negative.
+  if (FuncId <= 0)
+    return XRayPatchingStatus::FAILED;
+
+  // Check whether the function id is out of bounds.
+  if (static_cast<size_t>(FuncId) >= InstrMap.Functions)
+    return XRayPatchingStatus::FAILED;
+
+  // Now we patch ths sleds for this specific function.
+  auto SledRange = InstrMap.SledsIndex[FuncId - 1];
+  auto *f = reinterpret_cast<XRaySledEntry *>(SledRange.Begin);
+  auto *e = reinterpret_cast<XRaySledEntry *>(SledRange.End);
+
+  // TODO: Figure out how we can undo partial patching, or signal partial
+  // success.
+  while (f != e) {
+    if (!patchSled(*f, Enable, FuncId))
+      return XRayPatchingStatus::FAILED;
+    ++f;
+  }
+
+  return XRayPatchingStatus::SUCCESS;
+}
+
+XRayPatchingStatus __xray_patch_function(int32_t FuncId) XRAY_NEVER_INSTRUMENT {
+  return patchSpecificFunction(FuncId, true);
+}
+
+XRayPatchingStatus __xray_unpatch_function(int32_t FuncId) XRAY_NEVER_INSTRUMENT {
+  return patchSpecificFunction(FuncId, false);
+}
+
 int __xray_set_handler_arg1(void (*Handler)(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
   // scheduled on a different CPU/NUMA node.  We need to wait for everyone to
   // have this handler installed for consistency of collected data across CPUs.
   __sanitizer::atomic_store(&XRayArgLogger, reinterpret_cast<uint64_t>(Handler),
                             __sanitizer::memory_order_release);
   return 1;
 }
 int __xray_remove_handler_arg1() { return __xray_set_handler_arg1(nullptr); }

lib/xray/xray_interface_internal.h

#elif SANITIZER_WORDSIZE == 32
   uint32_t Function;
   unsigned char Kind;
   unsigned char AlwaysInstrument;
   unsigned char Padding[6]; // Need 16 bytes
 #else
 #error "Unsupported word size."
 #endif
 };
+
+struct XRayFunctionSledIndex {
+#if SANITIZER_WORDSIZE == 64
+  uint64_t Begin;
+  uint64_t End;
+#elif SANITIZER_WORDSIZE == 32
+  uint32_t Begin;
+  uint32_t End;
+#else
+#error "Unsupported word size."
+#endif
+};
 }
 
 namespace __xray {
 
 struct XRaySledMap {
   const XRaySledEntry *Sleds;
   size_t Entries;
+  const XRayFunctionSledIndex *SledsIndex;
+  size_t Functions;
 };
 
 bool patchFunctionEntry(bool Enable, uint32_t FuncId,
                         const XRaySledEntry &Sled, void (*Trampoline)());
 bool patchFunctionExit(bool Enable, uint32_t FuncId, const XRaySledEntry &Sled);
 bool patchFunctionTailExit(bool Enable, uint32_t FuncId,
                            const XRaySledEntry &Sled);
 

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
+
+#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) {
+  static 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.
+  for (const auto id : function_ids)
+    printf("patched: %d\n", id);
+
+  // 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.
+  for (const auto id : function_ids)
+    printf("patched: %d\n", id);
+  // CHECK-DAG: patched: [[F1]]
+  // CHECK-DAG: patched: [[F2]]
+  // CHECK-DAG: patched: [[F3]]
+
+  return function_ids == called_fns ? 0 : 1;
+}