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

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.

Reviewers: javed.absar, dberris

Reviewed By: dberris

Subscribers: llvm-commits

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

llvm-svn: 297000

7 files changed, 85 insertions, 11 deletions

diff --git a/compiler-rt/lib/xray/xray_AArch64.cc b/compiler-rt/lib/xray/xray_AArch64.cc
index 43c097d11a6..8d1c7c5d807 100644
--- a/compiler-rt/lib/xray/xray_AArch64.cc
+++ b/compiler-rt/lib/xray/xray_AArch64.cc
@@ -92,8 +92,9 @@ inline static bool patchSled(const bool Enable, const uint32_t FuncId,
 }
 
 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,
@@ -110,3 +111,7 @@ bool patchFunctionTailExit(const bool Enable, const uint32_t FuncId,
 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
+}
diff --git a/compiler-rt/lib/xray/xray_arm.cc b/compiler-rt/lib/xray/xray_arm.cc
index fb78eb6e850..26d673ec23a 100644
--- a/compiler-rt/lib/xray/xray_arm.cc
+++ b/compiler-rt/lib/xray/xray_arm.cc
@@ -128,8 +128,9 @@ inline static bool patchSled(const bool Enable, const uint32_t FuncId,
 }
 
 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,
@@ -146,3 +147,7 @@ bool patchFunctionTailExit(const bool Enable, const uint32_t FuncId,
 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
+}
diff --git a/compiler-rt/lib/xray/xray_fdr_logging.cc b/compiler-rt/lib/xray/xray_fdr_logging.cc
index 6db384519e6..bae7d4c4d07 100644
--- a/compiler-rt/lib/xray/xray_fdr_logging.cc
+++ b/compiler-rt/lib/xray/xray_fdr_logging.cc
@@ -494,6 +494,7 @@ void fdrLoggingHandleArg0(int32_t FuncId,
 
   switch (Entry) {
   case XRayEntryType::ENTRY:
+  case XRayEntryType::LOG_ARGS_ENTRY:
     FuncRecord.RecordKind = uint8_t(FunctionRecord::RecordKinds::FunctionEnter);
     break;
   case XRayEntryType::EXIT:
diff --git a/compiler-rt/lib/xray/xray_interface.cc b/compiler-rt/lib/xray/xray_interface.cc
index 306dbcae720..39cf8efb340 100644
--- a/compiler-rt/lib/xray/xray_interface.cc
+++ b/compiler-rt/lib/xray/xray_interface.cc
@@ -48,6 +48,9 @@ static const int16_t cSledLength = 8;
 // 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
@@ -185,7 +188,7 @@ XRayPatchingStatus controlPatching(bool Enable) XRAY_NEVER_INSTRUMENT {
     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);
@@ -193,6 +196,9 @@ XRayPatchingStatus controlPatching(bool Enable) XRAY_NEVER_INSTRUMENT {
     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;
@@ -211,3 +217,16 @@ XRayPatchingStatus __xray_patch() XRAY_NEVER_INSTRUMENT {
 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); }
diff --git a/compiler-rt/lib/xray/xray_interface_internal.h b/compiler-rt/lib/xray/xray_interface_internal.h
index c4f72433d17..0e3a251f3ad 100644
--- a/compiler-rt/lib/xray/xray_interface_internal.h
+++ b/compiler-rt/lib/xray/xray_interface_internal.h
@@ -49,7 +49,7 @@ struct XRaySledMap {
 };
 
 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);
@@ -62,6 +62,7 @@ extern "C" {
 extern void __xray_FunctionEntry();
 extern void __xray_FunctionExit();
 extern void __xray_FunctionTailExit();
+extern void __xray_ArgLoggerEntry();
 }
 
 #endif
diff --git a/compiler-rt/lib/xray/xray_trampoline_x86_64.S b/compiler-rt/lib/xray/xray_trampoline_x86_64.S
index d90c30cd98e..b9fef6dad9a 100644
--- a/compiler-rt/lib/xray/xray_trampoline_x86_64.S
+++ b/compiler-rt/lib/xray/xray_trampoline_x86_64.S
@@ -53,6 +53,9 @@
 
 	.text
 	.file "xray_trampoline_x86.S"
+
+//===----------------------------------------------------------------------===//
+
 	.globl __xray_FunctionEntry
 	.align 16, 0x90
 	.type __xray_FunctionEntry,@function
@@ -81,6 +84,8 @@ __xray_FunctionEntry:
 	.size __xray_FunctionEntry, .Ltmp1-__xray_FunctionEntry
 	.cfi_endproc
 
+//===----------------------------------------------------------------------===//
+
 	.globl __xray_FunctionExit
 	.align 16, 0x90
 	.type __xray_FunctionExit,@function
@@ -117,6 +122,8 @@ __xray_FunctionExit:
 	.size __xray_FunctionExit, .Ltmp3-__xray_FunctionExit
 	.cfi_endproc
 
+//===----------------------------------------------------------------------===//
+
 	.global __xray_FunctionTailExit
 	.align 16, 0x90
 	.type __xray_FunctionTailExit,@function
@@ -145,3 +152,39 @@ __xray_FunctionTailExit:
 .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
diff --git a/compiler-rt/lib/xray/xray_x86_64.cc b/compiler-rt/lib/xray/xray_x86_64.cc
index dd56e421437..1b9131316db 100644
--- a/compiler-rt/lib/xray/xray_x86_64.cc
+++ b/compiler-rt/lib/xray/xray_x86_64.cc
@@ -82,7 +82,8 @@ 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:
@@ -103,13 +104,12 @@ bool patchFunctionEntry(const bool Enable, const uint32_t FuncId,
   // 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) {