[XRay] [compiler-rt] Allow logging the first argument of a function call.

[pelikan]

Summary

Functions with the LOG_ARGS_ENTRY sled kind at their beginning will be handled
in a way to (optionally) pass their first call argument to your logging handler.

For practical and performance reasons, only the first argument is supported, and
only up to 64 bits.

[dberris]

include/xray/xray_interface.h

26–27 ↗

(On Diff #87590)

Use 3 instead (we don't have anything using 3 yet)?

lib/xray/xray_AArch64.cc

124–126 ↗

(On Diff #87590)

I don't think this will work -- this will be mangled differently from an extern "C" { ... } function, which this may have to be to make the symbol consistent. This means you'd have to define that symbol outside the __xray namespace, and mark it exter "C" {...}. As in:

} // namespace __xray

extern "C" {
void __xray_ArgLoggerEntry() XRAY_NEVER_INSTRUMENT {
  // FIXME: ...
}
}

lib/xray/xray_arm.cc

160–162 ↗

(On Diff #87590)

Same here.

lib/xray/xray_interface.cc

217 ↗

(On Diff #87590)

I think it would be better to make this store a release, for a stronger guarantee.

[pelikan]

• fix number and extern "C" ARM's empty stubs
• fix test

[pelikan]

lib/xray/xray_interface.cc

217 ↗

(On Diff #87590)

Stronger guarantee of what? There's nothing we're synchronizing *with*, we only care about atomicity (which this provides). When I see a release/consume I automatically look for an acquire, and in this case there isn't one.

[dberris]

lib/xray/xray_interface.cc

217 ↗

(On Diff #87590)

In x86-64 the acquire is in the assembly (spelled movq) because of the stronger memory model. That would change in weaker memory models, and since this is in a cross-platform part of the code, we have to be a bit more careful about this.

[dberris]

LGTM

lib/xray/xray_fdr_logging.cc

497 ↗

(On Diff #90435)

Potentially provide a comment to suggest that we ought to support this explicitly later (something like // FIXME: Handle actual argument logging later).

[pelikan]

I think this should land ASAP so we can get some testing. Things can be polished later.

lib/xray/xray_interface.cc

217 ↗

(On Diff #87590)

Why do we care which order do the store and load happen? I can't think of a scenario where I'm enabling arg1 logging on one CPU and worrying about what happens on other CPUs at the same time with respect to me just enabling it. How is eventual consistency not enough?

Similarly, due to the strong memory model, you can look at the MOV as a relaxed operation that just got unnecessarily expensive because of Intel ;-)

[dberris]

lib/xray/xray_interface.cc

217 ↗

(On Diff #87590)

Because on non-x86_64, these operations will really be relaxed. And because this isn't x86-specific code, we need to be careful about it.

For instance, the load in arm and ppc will need to be stronger versions (as it already is) to avoid having to deal with the issue of some threads not seeing the pointer update for potentially longer than when logging is on -- or they can stay relaxed but the write has to be strong so the appropriate caches are invalidated. A relaxed store here will mean that there's no strong guarantee that the data will show up as soon as the logging handler is installed, and threads can keep executing without seeing the handlers updated even if the sleds have been patched appropriately. We don't want this to happen because it affects data quality.

Does that make more sense?

[pelikan]

• strenghten and clarify how we treat the set_handler() operation w.r.t. the memory model

[pelikan]

lib/xray/xray_interface.cc

217 ↗

(On Diff #87590)

None of this answers the question "why do we care". The test behaviour depending on this, however, does. I've added a comment about why maintaining this is important.

[pelikan]

• unify the API and clarify comments

[pelikan]

• remove handler needs to return and int too, for consistency

[dberris]

LGTM -- will land soon enough. Thanks @pelikan!

compiler-rt/trunk/include/xray/xray_interface.h

 //===----------------------------------------------------------------------===//
 #ifndef XRAY_XRAY_INTERFACE_H
 #define XRAY_XRAY_INTERFACE_H
 
 #include <cstdint>
 
 extern "C" {
 
-enum XRayEntryType { ENTRY = 0, EXIT = 1, TAIL = 2 };
+// Synchronize this with AsmPrinter::SledKind in LLVM.
+enum XRayEntryType {
+  ENTRY = 0,
+  EXIT = 1,
+  TAIL = 2,
+  LOG_ARGS_ENTRY = 3,
+};
 
 // Provide a function to invoke for when instrumentation points are hit. This is
 // a user-visible control surface that overrides the default implementation. The
 // function provided should take the following arguments:
 //
 //   - function id: an identifier that indicates the id of a function; this id
 //                  is generated by xray; the mapping between the function id
 //                  and the actual function pointer is available through

 
 // 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();
+
+// 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

compiler-rt/trunk/lib/xray/xray_AArch64.cc

std::atomic_store_explicit(
         uint32_t(PatchOpcodes::PO_B32), std::memory_order_release);
   }
   __clear_cache(reinterpret_cast<char*>(FirstAddress),
       reinterpret_cast<char*>(CurAddress));
   return true;
 }
 
 bool patchFunctionEntry(const bool Enable, const uint32_t FuncId,
-                        const XRaySledEntry &Sled) XRAY_NEVER_INSTRUMENT {
-  return patchSled(Enable, FuncId, Sled, __xray_FunctionEntry);
+                        const XRaySledEntry &Sled,
+                        void (*Trampoline)()) XRAY_NEVER_INSTRUMENT {
+  return patchSled(Enable, FuncId, Sled, Trampoline);
 }
 
 bool patchFunctionExit(const bool Enable, const uint32_t FuncId,
                        const XRaySledEntry &Sled) XRAY_NEVER_INSTRUMENT {
   return patchSled(Enable, FuncId, Sled, __xray_FunctionExit);
 }
 
 bool patchFunctionTailExit(const bool Enable, const uint32_t FuncId,
                            const XRaySledEntry &Sled) XRAY_NEVER_INSTRUMENT {
   return patchSled(Enable, FuncId, Sled, __xray_FunctionTailExit);
 }
 
 // FIXME: Maybe implement this better?
 bool probeRequiredCPUFeatures() XRAY_NEVER_INSTRUMENT { return true; }
 
 } // namespace __xray
+
+extern "C" void __xray_ArgLoggerEntry() XRAY_NEVER_INSTRUMENT {
+  // FIXME: this will have to be implemented in the trampoline assembly file
+}

compiler-rt/trunk/lib/xray/xray_arm.cc

std::atomic_store_explicit(
         uint32_t(PatchOpcodes::PO_B20), std::memory_order_release);
   }
   __clear_cache(reinterpret_cast<char*>(FirstAddress),
       reinterpret_cast<char*>(CurAddress));
   return true;
 }
 
 bool patchFunctionEntry(const bool Enable, const uint32_t FuncId,
-                        const XRaySledEntry &Sled) XRAY_NEVER_INSTRUMENT {
-  return patchSled(Enable, FuncId, Sled, __xray_FunctionEntry);
+                        const XRaySledEntry &Sled,
+                        void (*Trampoline)()) XRAY_NEVER_INSTRUMENT {
+  return patchSled(Enable, FuncId, Sled, Trampoline);
 }
 
 bool patchFunctionExit(const bool Enable, const uint32_t FuncId,
                        const XRaySledEntry &Sled) XRAY_NEVER_INSTRUMENT {
   return patchSled(Enable, FuncId, Sled, __xray_FunctionExit);
 }
 
 bool patchFunctionTailExit(const bool Enable, const uint32_t FuncId,
                            const XRaySledEntry &Sled) XRAY_NEVER_INSTRUMENT {
   return patchSled(Enable, FuncId, Sled, __xray_FunctionTailExit);
 }
 
 // FIXME: Maybe implement this better?
 bool probeRequiredCPUFeatures() XRAY_NEVER_INSTRUMENT { return true; }
 
 } // namespace __xray
+
+extern "C" void __xray_ArgLoggerEntry() XRAY_NEVER_INSTRUMENT {
+  // FIXME: this will have to be implemented in the trampoline assembly file
+}

compiler-rt/trunk/lib/xray/xray_fdr_logging.cc

void fdrLoggingHandleArg0(int32_t FuncId,
 
   // We then update our "LastTSC" and "CurrentCPU" thread-local variables to aid
   // us in future computations of this TSC delta value.
   LastTSC = TSC;
   CurrentCPU = CPU;
 
   switch (Entry) {
   case XRayEntryType::ENTRY:
+  case XRayEntryType::LOG_ARGS_ENTRY:
     FuncRecord.RecordKind = uint8_t(FunctionRecord::RecordKinds::FunctionEnter);
     break;
   case XRayEntryType::EXIT:
     FuncRecord.RecordKind = uint8_t(FunctionRecord::RecordKinds::FunctionExit);
     break;
   case XRayEntryType::TAIL:
     FuncRecord.RecordKind =
         uint8_t(FunctionRecord::RecordKinds::FunctionTailExit);

compiler-rt/trunk/lib/xray/xray_interface.cc

 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.
 std::atomic<void (*)(int32_t, XRayEntryType)> XRayPatchedFunction{nullptr};
 
+// This is the function to call from the arg1-enabled sleds/trampolines.
+std::atomic<void (*)(int32_t, XRayEntryType, uint64_t)> XRayArgLogger{nullptr};
+
 // 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.
for (std::size_t I = 0; I < InstrMap.Entries; I++) {
     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);
+      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;
   }
   XRayPatching.store(false, std::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);
 }
+
+int __xray_set_handler_arg1(void (*Handler)(int32_t, XRayEntryType, uint64_t))
+{
+  if (!XRayInitialized.load(std::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.
+  XRayArgLogger.store(Handler, std::memory_order_release);
+  return 1;
+}
+int __xray_remove_handler_arg1() { return __xray_set_handler_arg1(nullptr); }

compiler-rt/trunk/lib/xray/xray_interface_internal.h

 namespace __xray {
 
 struct XRaySledMap {
   const XRaySledEntry *Sleds;
   size_t Entries;
 };
 
 bool patchFunctionEntry(bool Enable, uint32_t FuncId,
-                        const XRaySledEntry &Sled);
+                        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);
 
 } // namespace __xray
 
 extern "C" {
 // The following functions have to be defined in assembler, on a per-platform
 // basis. See xray_trampoline_*.S files for implementations.
 extern void __xray_FunctionEntry();
 extern void __xray_FunctionExit();
 extern void __xray_FunctionTailExit();
+extern void __xray_ArgLoggerEntry();
 }
 
 #endif

compiler-rt/trunk/lib/xray/xray_trampoline_x86_64.S

.macro RESTORE_REGISTERS
   movq  32(%rsp), %rcx
   movq  24(%rsp), %r8
   movq  16(%rsp), %r9
   addq  $200, %rsp
 .endm
 
   .text
   .file "xray_trampoline_x86.S"
+
+//===----------------------------------------------------------------------===//
+
   .globl __xray_FunctionEntry
   .align 16, 0x90
   .type __xray_FunctionEntry,@function
 
 __xray_FunctionEntry:
   .cfi_startproc
   pushq %rbp
   .cfi_def_cfa_offset 16

 .Ltmp0:
   RESTORE_REGISTERS
   popq  %rbp
   retq
 .Ltmp1:
   .size __xray_FunctionEntry, .Ltmp1-__xray_FunctionEntry
   .cfi_endproc
 
+//===----------------------------------------------------------------------===//
+
   .globl __xray_FunctionExit
   .align 16, 0x90
   .type __xray_FunctionExit,@function
 __xray_FunctionExit:
   .cfi_startproc
   // Save the important registers first. Since we're assuming that this
   // function is only jumped into, we only preserve the registers for
   // returning.
.Ltmp2:
   movq  8(%rsp), %rdx
   addq  $56, %rsp
   popq  %rbp
   retq
 .Ltmp3:
   .size __xray_FunctionExit, .Ltmp3-__xray_FunctionExit
   .cfi_endproc
 
+//===----------------------------------------------------------------------===//
+
   .global __xray_FunctionTailExit
   .align 16, 0x90
   .type __xray_FunctionTailExit,@function
 __xray_FunctionTailExit:
   .cfi_startproc
   // Save the important registers as in the entry trampoline, but indicate that
   // this is an exit. In the future, we will introduce a new entry type that
   // differentiates between a normal exit and a tail exit, but we'd have to do

 
 .Ltmp4:
   RESTORE_REGISTERS
   popq  %rbp
   retq
 .Ltmp5:
   .size __xray_FunctionTailExit, .Ltmp5-__xray_FunctionTailExit
   .cfi_endproc
+
+//===----------------------------------------------------------------------===//
+
+  .globl __xray_ArgLoggerEntry
+  .align 16, 0x90
+  .type __xray_ArgLoggerEntry,@function
+__xray_ArgLoggerEntry:
+  .cfi_startproc
+  pushq %rbp
+  .cfi_def_cfa_offset 16
+  SAVE_REGISTERS
+
+  // Again, these function pointer loads must be atomic; MOV is fine.
+  movq  _ZN6__xray13XRayArgLoggerE(%rip), %rax
+  testq %rax, %rax
+  jne .Larg1entryLog
+
+  // If [arg1 logging handler] not set, defer to no-arg logging.
+  movq  _ZN6__xray19XRayPatchedFunctionE(%rip), %rax
+  testq %rax, %rax
+  je  .Larg1entryFail
+
+.Larg1entryLog:
+  movq  %rdi, %rdx  // first argument will become the third
+  xorq  %rsi, %rsi  // XRayEntryType::ENTRY into the second
+  movl  %r10d, %edi // 32-bit function ID becomes the first
+  callq *%rax
+
+.Larg1entryFail:
+  RESTORE_REGISTERS
+  popq  %rbp
+  retq
+
+.Larg1entryEnd:
+  .size __xray_ArgLoggerEntry, .Larg1entryEnd-__xray_ArgLoggerEntry
+  .cfi_endproc

compiler-rt/trunk/lib/xray/xray_x86_64.cc

 static constexpr uint16_t Jmp9Seq = 0x09eb;
 static constexpr uint8_t JmpOpCode = 0xe9;
 static constexpr uint8_t RetOpCode = 0xc3;
 
 static constexpr int64_t MinOffset{std::numeric_limits<int32_t>::min()};
 static constexpr int64_t MaxOffset{std::numeric_limits<int32_t>::max()};
 
 bool patchFunctionEntry(const bool Enable, const uint32_t FuncId,
-                        const XRaySledEntry &Sled) XRAY_NEVER_INSTRUMENT {
+                        const XRaySledEntry &Sled,
+                        void (*Trampoline)()) XRAY_NEVER_INSTRUMENT {
   // Here we do the dance of replacing the following sled:
   //
   // xray_sled_n:
   //   jmp +9
   //   <9 byte nop>
   //
   // With the following:
   //
   //   mov r10d, <function id>
   //   call <relative 32bit offset to entry trampoline>
   //
   // We need to do this in the following order:
   //
   // 1. Put the function id first, 2 bytes from the start of the sled (just
   // after the 2-byte jmp instruction).
   // 2. Put the call opcode 6 bytes from the start of the sled.
   // 3. Put the relative offset 7 bytes from the start of the sled.
   // 4. Do an atomic write over the jmp instruction for the "mov r10d"
   // opcode and first operand.
   //
-  // Prerequisite is to compute the relative offset to the
-  // __xray_FunctionEntry function's address.
-  int64_t TrampolineOffset = reinterpret_cast<int64_t>(__xray_FunctionEntry) -
+  // Prerequisite is to compute the relative offset to the trampoline's address.
+  int64_t TrampolineOffset = reinterpret_cast<int64_t>(Trampoline) -
                              (static_cast<int64_t>(Sled.Address) + 11);
   if (TrampolineOffset < MinOffset || TrampolineOffset > MaxOffset) {
     Report("XRay Entry trampoline (%p) too far from sled (%p)\n",
-           __xray_FunctionEntry, reinterpret_cast<void *>(Sled.Address));
+           Trampoline, reinterpret_cast<void *>(Sled.Address));
     return false;
   }
   if (Enable) {
     *reinterpret_cast<uint32_t *>(Sled.Address + 2) = FuncId;
     *reinterpret_cast<uint8_t *>(Sled.Address + 6) = CallOpCode;
     *reinterpret_cast<uint32_t *>(Sled.Address + 7) = TrampolineOffset;
     std::atomic_store_explicit(
         reinterpret_cast<std::atomic<uint16_t> *>(Sled.Address), MovR10Seq,

compiler-rt/trunk/test/xray/TestCases/Linux/arg1-logger.cc

+// Check that we can get the first function argument logged
+// using a custom logging function.
+//
+// RUN: %clangxx_xray -std=c++11 %s -o %t
+// RUN: XRAY_OPTIONS="patch_premain=true verbosity=1 xray_logfile_base=arg1-logger-" %run %t 2>&1 | FileCheck %s
+//
+// After all that, clean up the XRay log file.
+//
+// RUN: rm arg1-logger-*
+//
+// At the time of writing, the ARM trampolines weren't written yet.
+// XFAIL: arm || aarch64
+
+#include "xray/xray_interface.h"
+
+#include <cinttypes>
+#include <cstdio>
+
+void arg1logger(int32_t fn, XRayEntryType t, uint64_t a1) {
+  printf("Arg1: %" PRIx64 ", XRayEntryType %u\n", a1, t);
+}
+
+[[clang::xray_always_instrument, clang::xray_log_args(1)]] void foo(void *) {}
+
+int main() {
+  // CHECK: XRay: Log file in 'arg1-logger-{{.*}}'
+
+  __xray_set_handler_arg1(arg1logger);
+  foo(nullptr);
+  // CHECK: Arg1: 0, XRayEntryType 0
+
+  __xray_remove_handler_arg1();
+  foo((void *) 0xBADC0DE);
+  // nothing expected to see here
+
+  __xray_set_handler_arg1(arg1logger);
+  foo((void *) 0xDEADBEEFCAFE);
+  // CHECK-NEXT: Arg1: deadbeefcafe, XRayEntryType 0
+  foo((void *) -1);
+  // CHECK-NEXT: Arg1: ffffffffffffffff, XRayEntryType 0
+}