Implement trampoline and handler for typed xray event tracing.

Summary:
Compiler-rt support first before defining the __xray_typedevent() lowering in
llvm. I'm looking for some early feedback before I touch much more code.

Reviewers: dberris

Subscribers: delcypher, llvm-commits, #sanitizers

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

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

9 files changed, 208 insertions, 12 deletions

diff --git a/include/xray/xray_interface.h b/include/xray/xray_interface.h
index d08039a67..ba4c0e812 100644
--- a/include/xray/xray_interface.h
+++ b/include/xray/xray_interface.h
@@ -27,6 +27,7 @@ enum XRayEntryType {
   TAIL = 2,
   LOG_ARGS_ENTRY = 3,
   CUSTOM_EVENT = 4,
+  TYPED_EVENT = 5,
 };
 
 /// Provide a function to invoke for when instrumentation points are hit. This
@@ -68,12 +69,23 @@ extern int __xray_set_handler_arg1(void (*entry)(int32_t, XRayEntryType,
 extern int __xray_remove_handler_arg1();
 
 /// Provide a function to invoke when XRay encounters a custom event.
-extern int __xray_set_customevent_handler(void (*entry)(void*, std::size_t));
+extern int __xray_set_customevent_handler(void (*entry)(void *, std::size_t));
 
 /// This removes whatever the currently provided custom event handler is.
 /// Returns 1 on success, 0 on error.
 extern int __xray_remove_customevent_handler();
 
+/// Set a handler for xray typed event logging. The first parameter is a type
+/// identifier, the second is a payload, and the third is the payload size.
+extern int __xray_set_typedevent_handler(void (*entry)(uint16_t, const void *,
+                                                       std::size_t));
+
+/// Removes the currently set typed event handler.
+/// Returns 1 on success, 0 on error.
+extern int __xray_remove_typedevent_handler();
+
+extern uint16_t __xray_register_event_type(const char *event_type);
+
 enum XRayPatchingStatus {
   NOT_INITIALIZED = 0,
   SUCCESS = 1,
diff --git a/lib/xray/xray_fdr_log_records.h b/lib/xray/xray_fdr_log_records.h
index 324208db8..1a5518d7a 100644
--- a/lib/xray/xray_fdr_log_records.h
+++ b/lib/xray/xray_fdr_log_records.h
@@ -32,6 +32,7 @@ struct alignas(16) MetadataRecord {
     CustomEventMarker,
     CallArgument,
     BufferExtents,
+    TypedEventMarker,
   };
 
   // Use 7 bits to identify this record type.
diff --git a/lib/xray/xray_fdr_logging.cc b/lib/xray/xray_fdr_logging.cc
index 72776e2c1..fcb33fc87 100644
--- a/lib/xray/xray_fdr_logging.cc
+++ b/lib/xray/xray_fdr_logging.cc
@@ -289,6 +289,63 @@ void fdrLoggingHandleCustomEvent(void *Event,
   endBufferIfFull();
 }
 
+void fdrLoggingHandleTypedEvent(
+    uint16_t EventType, const void *Event,
+    std::size_t EventSize) noexcept XRAY_NEVER_INSTRUMENT {
+  using namespace __xray_fdr_internal;
+  auto TC = getTimestamp();
+  auto &TSC = TC.TSC;
+  auto &CPU = TC.CPU;
+  RecursionGuard Guard{Running};
+  if (!Guard)
+    return;
+  if (EventSize > std::numeric_limits<int32_t>::max()) {
+    using Empty = struct {};
+    static Empty Once = [&] {
+      Report("Event size too large = %zu ; > max = %d\n", EventSize,
+             std::numeric_limits<int32_t>::max());
+      return Empty();
+    }();
+    (void)Once;
+  }
+  int32_t ReducedEventSize = static_cast<int32_t>(EventSize);
+  auto &TLD = getThreadLocalData();
+  if (!isLogInitializedAndReady(TLD.BQ, TSC, CPU, clock_gettime))
+    return;
+
+  // Here we need to prepare the log to handle:
+  //   - The metadata record we're going to write. (16 bytes)
+  //   - The additional data we're going to write. Currently, that's the size of
+  //   the event we're going to dump into the log as free-form bytes.
+  if (!prepareBuffer(TSC, CPU, clock_gettime, MetadataRecSize + EventSize)) {
+    TLD.BQ = nullptr;
+    return;
+  }
+  // Write the custom event metadata record, which consists of the following
+  // information:
+  //   - 8 bytes (64-bits) for the full TSC when the event started.
+  //   - 4 bytes (32-bits) for the length of the data.
+  //   - 2 bytes (16-bits) for the event type. 3 bytes remain since one of the
+  //       bytes has the record type (Metadata Record) and kind (TypedEvent).
+  //       We'll log the error if the event type is greater than 2 bytes.
+  //       Event types are generated sequentially, so 2^16 is enough.
+  MetadataRecord TypedEvent;
+  TypedEvent.Type = uint8_t(RecordType::Metadata);
+  TypedEvent.RecordKind =
+      uint8_t(MetadataRecord::RecordKinds::TypedEventMarker);
+  constexpr auto TSCSize = sizeof(TC.TSC);
+  std::memcpy(&TypedEvent.Data, &ReducedEventSize, sizeof(int32_t));
+  std::memcpy(&TypedEvent.Data[sizeof(int32_t)], &TSC, TSCSize);
+  std::memcpy(&TypedEvent.Data[sizeof(int32_t) + TSCSize], &EventType,
+              sizeof(EventType));
+  std::memcpy(TLD.RecordPtr, &TypedEvent, sizeof(TypedEvent));
+
+  TLD.RecordPtr += sizeof(TypedEvent);
+  std::memcpy(TLD.RecordPtr, Event, ReducedEventSize);
+  incrementExtents(MetadataRecSize + EventSize);
+  endBufferIfFull();
+}
+
 XRayLogInitStatus fdrLoggingInit(std::size_t BufferSize, std::size_t BufferMax,
                                  void *Options,
                                  size_t OptionsSize) XRAY_NEVER_INSTRUMENT {
@@ -352,6 +409,7 @@ XRayLogInitStatus fdrLoggingInit(std::size_t BufferSize, std::size_t BufferMax,
   // Install the actual handleArg0 handler after initialising the buffers.
   __xray_set_handler(fdrLoggingHandleArg0);
   __xray_set_customevent_handler(fdrLoggingHandleCustomEvent);
+  __xray_set_typedevent_handler(fdrLoggingHandleTypedEvent);
 
   __sanitizer::atomic_store(&LoggingStatus,
                             XRayLogInitStatus::XRAY_LOG_INITIALIZED,
diff --git a/lib/xray/xray_fdr_logging_impl.h b/lib/xray/xray_fdr_logging_impl.h
index ed0b5ca90..3de310e9c 100644
--- a/lib/xray/xray_fdr_logging_impl.h
+++ b/lib/xray/xray_fdr_logging_impl.h
@@ -364,6 +364,16 @@ writeFunctionRecord(int FuncId, uint32_t TSCDelta,
     (void)Once;
     return;
   }
+  case XRayEntryType::TYPED_EVENT: {
+    static bool Once = [&] {
+      Report("Internal error: patched an XRay typed event call as a function; "
+             "func id = %d\n",
+             FuncId);
+      return true;
+    }();
+    (void)Once;
+    return;
+  }
   }
 
   std::memcpy(TLD.RecordPtr, &FuncRecord, sizeof(FunctionRecord));
@@ -689,6 +699,16 @@ inline void processFunctionHook(int32_t FuncId, XRayEntryType Entry,
     (void)Once;
     return;
   }
+  case XRayEntryType::TYPED_EVENT: {
+    static bool Once = [&] {
+      Report("Internal error: patched an XRay typed event call as a function; "
+             "func id = %d\n",
+             FuncId);
+      return true;
+    }();
+    (void)Once;
+    return;
+  }
   }
 
   writeFunctionRecord(FuncId, RecordTSCDelta, Entry);
diff --git a/lib/xray/xray_inmemory_log.cc b/lib/xray/xray_inmemory_log.cc
index ae437766f..4d1189695 100644
--- a/lib/xray/xray_inmemory_log.cc
+++ b/lib/xray/xray_inmemory_log.cc
@@ -401,6 +401,7 @@ XRayLogInitStatus basicLoggingInit(size_t BufferSize, size_t BufferMax,
   __xray_set_handler(UseRealTSC ? basicLoggingHandleArg0RealTSC
                                 : basicLoggingHandleArg0EmulateTSC);
   __xray_remove_customevent_handler();
+  __xray_remove_typedevent_handler();
 
   return XRayLogInitStatus::XRAY_LOG_INITIALIZED;
 }
diff --git a/lib/xray/xray_interface.cc b/lib/xray/xray_interface.cc
index 766313e85..edfb59c6c 100644
--- a/lib/xray/xray_interface.cc
+++ b/lib/xray/xray_interface.cc
@@ -19,9 +19,12 @@
 #include <cstdio>
 #include <errno.h>
 #include <limits>
+#include <string.h>
 #include <sys/mman.h>
 
+#include "sanitizer_common/sanitizer_addrhashmap.h"
 #include "sanitizer_common/sanitizer_common.h"
+
 #include "xray_defs.h"
 #include "xray_flags.h"
 
@@ -56,18 +59,32 @@ __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 function to call when we encounter a typed event log call.
+__sanitizer::atomic_uintptr_t XRayPatchedTypedEvent{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.
+struct TypeDescription {
+  uint32_t type_id;
+  std::size_t description_string_length;
+};
+
+using TypeDescriptorMapType = __sanitizer::AddrHashMap<TypeDescription, 11>;
+// An address map from immutable descriptors to type ids.
+TypeDescriptorMapType TypeDescriptorAddressMap{};
+
+__sanitizer::atomic_uint32_t TypeEventDescriptorCounter{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.
+// instrumentation points. This assumes that the executable pages are
+// originally read-and-execute only.
 class MProtectHelper {
   void *PageAlignedAddr;
   std::size_t MProtectLen;
@@ -116,6 +133,9 @@ bool patchSled(const XRaySledEntry &Sled, bool Enable,
   case XRayEntryType::CUSTOM_EVENT:
     Success = patchCustomEvent(Enable, FuncId, Sled);
     break;
+  case XRayEntryType::TYPED_EVENT:
+    Success = patchTypedEvent(Enable, FuncId, Sled);
+    break;
   default:
     Report("Unsupported sled kind '%d' @%04x\n", Sled.Address, int(Sled.Kind));
     return false;
@@ -341,6 +361,18 @@ int __xray_set_customevent_handler(void (*entry)(void *, size_t))
   return 0;
 }
 
+int __xray_set_typedevent_handler(void (*entry)(
+    uint16_t, const void *, size_t)) noexcept XRAY_NEVER_INSTRUMENT {
+  if (__sanitizer::atomic_load(&XRayInitialized,
+                               __sanitizer::memory_order_acquire)) {
+    __sanitizer::atomic_store(&__xray::XRayPatchedTypedEvent,
+                              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);
 }
@@ -349,6 +381,21 @@ int __xray_remove_customevent_handler() XRAY_NEVER_INSTRUMENT {
   return __xray_set_customevent_handler(nullptr);
 }
 
+int __xray_remove_typedevent_handler() noexcept XRAY_NEVER_INSTRUMENT {
+  return __xray_set_typedevent_handler(nullptr);
+}
+
+uint16_t __xray_register_event_type(
+    const char *const event_type) noexcept XRAY_NEVER_INSTRUMENT {
+  TypeDescriptorMapType::Handle h(&TypeDescriptorAddressMap, (uptr)event_type);
+  if (h.created()) {
+    h->type_id = __sanitizer::atomic_fetch_add(
+        &TypeEventDescriptorCounter, 1, __sanitizer::memory_order_acq_rel);
+    h->description_string_length = strnlen(event_type, 1024);
+  }
+  return h->type_id;
+}
+
 XRayPatchingStatus __xray_patch() XRAY_NEVER_INSTRUMENT {
   return controlPatching(true);
 }
diff --git a/lib/xray/xray_interface_internal.h b/lib/xray/xray_interface_internal.h
index 5811e2b73..8ca874574 100644
--- a/lib/xray/xray_interface_internal.h
+++ b/lib/xray/xray_interface_internal.h
@@ -43,8 +43,8 @@ struct XRaySledEntry {
 };
 
 struct XRayFunctionSledIndex {
-  const XRaySledEntry* Begin;
-  const XRaySledEntry* End;
+  const XRaySledEntry *Begin;
+  const XRaySledEntry *End;
 };
 }
 
@@ -57,12 +57,13 @@ struct XRaySledMap {
   size_t Functions;
 };
 
-bool patchFunctionEntry(bool Enable, uint32_t FuncId,
-                        const XRaySledEntry &Sled, void (*Trampoline)());
+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);
 bool patchCustomEvent(bool Enable, uint32_t FuncId, const XRaySledEntry &Sled);
+bool patchTypedEvent(bool Enable, uint32_t FuncId, const XRaySledEntry &Sled);
 
 } // namespace __xray
 
@@ -74,6 +75,7 @@ extern void __xray_FunctionExit();
 extern void __xray_FunctionTailExit();
 extern void __xray_ArgLoggerEntry();
 extern void __xray_CustomEvent();
+extern void __xray_TypedEvent();
 }
 
 #endif
diff --git a/lib/xray/xray_trampoline_x86_64.S b/lib/xray/xray_trampoline_x86_64.S
index 370a561e1..83b83e81b 100644
--- a/lib/xray/xray_trampoline_x86_64.S
+++ b/lib/xray/xray_trampoline_x86_64.S
@@ -242,4 +242,29 @@ ASM_SYMBOL(__xray_CustomEvent):
 	ASM_SIZE(__xray_CustomEvent)
 	CFI_ENDPROC
 
+//===----------------------------------------------------------------------===//
+
+	.global ASM_SYMBOL(__xray_TypedEvent)
+	.align 16, 0x90
+	ASM_TYPE_FUNCTION(__xray_TypedEvent)
+ASM_SYMBOL(__xray_TypedEvent):
+	CFI_STARTPROC
+	SAVE_REGISTERS
+
+	// We pass three arguments to this trampoline, which should be in rdi, rsi
+	// and rdx without our intervention.
+	movq ASM_SYMBOL(_ZN6__xray21XRayPatchedTypedEventE)(%rip), %rax
+	testq %rax,%rax
+	je .LtypedEventCleanup
+
+	ALIGNED_CALL_RAX
+
+.LtypedEventCleanup:
+	RESTORE_REGISTERS
+	retq
+	ASM_SIZE(__xray_TypedEvent)
+	CFI_ENDPROC
+
+//===----------------------------------------------------------------------===//
+
 NO_EXEC_STACK_DIRECTIVE
diff --git a/lib/xray/xray_x86_64.cc b/lib/xray/xray_x86_64.cc
index 52c5535c4..2021632f0 100644
--- a/lib/xray/xray_x86_64.cc
+++ b/lib/xray/xray_x86_64.cc
@@ -99,7 +99,6 @@ uint64_t getTSCFrequency() XRAY_NEVER_INSTRUMENT {
     }
 
     return 0;
-    
 }
 #else
 uint64_t getTSCFrequency() XRAY_NEVER_INSTRUMENT {
@@ -287,6 +286,37 @@ bool patchCustomEvent(const bool Enable, const uint32_t FuncId,
   return false;
 }
 
+bool patchTypedEvent(const bool Enable, const uint32_t FuncId,
+                      const XRaySledEntry &Sled) XRAY_NEVER_INSTRUMENT {
+  // Here we do the dance of replacing the following sled:
+  //
+  // xray_sled_n:
+  //   jmp +20          // 2 byte instruction
+  //   ...
+  //
+  // With the following:
+  //
+  //   nopw             // 2 bytes
+  //   ...
+  //
+  //
+  // The "unpatch" should just turn the 'nopw' back to a 'jmp +20'.
+  // The 20 byte sled stashes three argument registers, calls the trampoline,
+  // unstashes the registers and returns. If the arguments are already in
+  // the correct registers, the stashing and unstashing become equivalently
+  // sized nops.
+  if (Enable) {
+    std::atomic_store_explicit(
+        reinterpret_cast<std::atomic<uint16_t> *>(Sled.Address), NopwSeq,
+        std::memory_order_release);
+  } else {
+      std::atomic_store_explicit(
+          reinterpret_cast<std::atomic<uint16_t> *>(Sled.Address), Jmp20Seq,
+          std::memory_order_release);
+  }
+  return false;
+}
+
 // We determine whether the CPU we're running on has the correct features we
 // need. In x86_64 this will be rdtscp support.
 bool probeRequiredCPUFeatures() XRAY_NEVER_INSTRUMENT {