[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 ↗

(On Diff #126673)

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 ↗

(On Diff #126673)

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

221 ↗

(On Diff #126673)

const auto &Sled?

228 ↗

(On Diff #126673)

Usage suggests defining the flag as size_t.

232 ↗

(On Diff #126673)

Consider listing this requirement in the flag description.

320 ↗

(On Diff #126673)

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 ↗

(On Diff #126673)

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 ↗

(On Diff #126673)

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 ↗

(On Diff #126673)

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

320 ↗

(On Diff #126673)

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/trunk/lib/xray/xray_flags.h

 //
 // XRay runtime flags.
 //===----------------------------------------------------------------------===//
 
 #ifndef XRAY_FLAGS_H
 #define XRAY_FLAGS_H
 
 #include "sanitizer_common/sanitizer_flag_parser.h"
+#include "sanitizer_common/sanitizer_internal_defs.h"
 
 namespace __xray {
 
 struct Flags {
 #define XRAY_FLAG(Type, Name, DefaultValue, Description) Type Name;
 #include "xray_flags.inc"
 #undef XRAY_FLAG
 

compiler-rt/trunk/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. The size "
+          "should be a power-of-two.")
 
 // 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/trunk/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"
+
+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;
 #elif defined(__arm__)
 static const int16_t cSledLength = 28;
 #elif SANITIZER_MIPS32
 static const int16_t cSledLength = 48;
 #elif SANITIZER_MIPS64
 static const int16_t cSledLength = 64;
 #elif defined(__powerpc64__)
 static const int16_t cSledLength = 8;
 #else
 #error "Unsupported CPU Architecture"
 #endif /* CPU architecture */
 
 // This is the function to call when we encounter the entry or exit sleds.
 __sanitizer::atomic_uintptr_t XRayPatchedFunction{0};
 
 // This is the function to call from the arg1-enabled sleds/trampolines.
 __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
 // the page into a read-only and executable page.
 //
 // This is only used specifically for runtime-patching of the XRay
 // instrumentation points. This assumes that the executable pages are originally
 // read-and-execute only.
public:
 
   ~MProtectHelper() XRAY_NEVER_INSTRUMENT {
     if (MustCleanup) {
       mprotect(PageAlignedAddr, MProtectLen, PROT_READ | PROT_EXEC);
     }
   }
 };
 
-} // 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(); }
-};
+namespace {
 
-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 patchSled(const XRaySledEntry &Sled, bool Enable,
   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;
   case XRayEntryType::CUSTOM_EVENT:
     Success = patchCustomEvent(Enable, FuncId, Sled);
     break;
   default:
     Report("Unsupported sled kind '%d' @%04x\n", Sled.Address, int(Sled.Kind));
     return false;
   }
   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.
 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;
-  auto XRayPatchingStatusResetter = scopeCleanup([&PatchingSuccess] {
+  auto XRayPatchingStatusResetter =
+      __sanitizer::at_scope_exit([&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++) {
+    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;
     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 __xray_unpatch() XRAY_NEVER_INSTRUMENT {
-  return controlPatching(false);
-}
-
-XRayPatchingStatus patchFunction(int32_t FuncId,
+XRayPatchingStatus mprotectAndPatchFunction(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.
+  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,
+} // namespace
+
+} // 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;
+}
 
-  if (!SucceedOnce) {
-    Report("Failed patching any sled for function '%d'.", FuncId);
-    return XRayPatchingStatus::FAILED;
-  }
+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;
 
-  return XRayPatchingStatus::SUCCESS;
+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)) {
   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/trunk/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
 }