3 files changed, 92 insertions, 3 deletions
diff --git a/include/xray/xray_interface.h b/include/xray/xray_interface.h
index 564613417069..d08039a67fa2 100644
--- a/include/xray/xray_interface.h
+++ b/include/xray/xray_interface.h
@@ -106,6 +106,14 @@ extern uintptr_t __xray_function_address(int32_t FuncId);
 /// encounter errors (when there are no instrumented functions, etc.).
 extern size_t __xray_max_function_id();
 
+/// Initialize the required XRay data structures. This is useful in cases where
+/// users want to control precisely when the XRay instrumentation data
+/// structures are initialized, for example when the XRay library is built with
+/// the XRAY_NO_PREINIT preprocessor definition.
+///
+/// Calling __xray_init() more than once is safe across multiple threads.
+extern void __xray_init();
+
 } // end extern "C"
 
 #endif // XRAY_XRAY_INTERFACE_H
diff --git a/include/xray/xray_log_interface.h b/include/xray/xray_log_interface.h
index cdb20094dc26..6c53cdaebfe2 100644
--- a/include/xray/xray_log_interface.h
+++ b/include/xray/xray_log_interface.h
@@ -128,6 +128,16 @@ enum XRayLogFlushStatus {
   XRAY_LOG_FLUSHED = 2,
 };
 
+/// This enum indicates the installation state of a logging implementation, when
+/// associating a mode to a particular logging implementation through
+/// `__xray_log_register_impl(...)` or through `__xray_log_select_mode(...`.
+enum XRayLogRegisterStatus {
+  XRAY_REGISTRATION_OK = 0,
+  XRAY_DUPLICATE_MODE = 1,
+  XRAY_MODE_NOT_FOUND = 2,
+  XRAY_INCOMPLETE_IMPL = 3,
+};
+
 /// A valid XRay logging implementation MUST provide all of the function
 /// pointers in XRayLogImpl when being installed through `__xray_set_log_impl`.
 /// To be precise, ALL the functions pointers MUST NOT be nullptr.
@@ -159,6 +169,9 @@ struct XRayLogImpl {
   /// always have a handler for function entry and exit events. In case the
   /// implementation wants to support arg1 (or other future extensions to XRay
   /// logging) those MUST be installed by the installed 'log_init' handler.
+  ///
+  /// Because we didn't want to change the ABI of this struct, the arg1 handler
+  /// may be silently overwritten during initialization as well.
   void (*handle_arg0)(int32_t, XRayEntryType);
 
   /// The log implementation provided routine for when __xray_log_flushLog() is
@@ -186,6 +199,34 @@ struct XRayLogImpl {
 /// called while in any other states.
 void __xray_set_log_impl(XRayLogImpl Impl);
 
+/// This function registers a logging implementation against a "mode"
+/// identifier. This allows multiple modes to be registered, and chosen at
+/// runtime using the same mode identifier through
+/// `__xray_log_select_mode(...)`.
+///
+/// We treat the Mode identifier as a null-terminated byte string, as the
+/// identifier used when retrieving the log impl.
+///
+/// Returns:
+///   - XRAY_REGISTRATION_OK on success.
+///   - XRAY_DUPLICATE_MODE when an implementation is already associated with
+///     the provided Mode; does not update the already-registered
+///     implementation.
+XRayLogRegisterStatus __xray_log_register_mode(const char *Mode,
+                                               XRayLogImpl Impl);
+
+/// This function selects the implementation associated with Mode that has been
+/// registered through __xray_log_register_mode(...) and installs that
+/// implementation (as if through calling __xray_set_log_impl(...)). The same
+/// caveats apply to __xray_log_select_mode(...) as with
+/// __xray_log_set_log_impl(...).
+///
+/// Returns:
+///   - XRAY_REGISTRATION_OK on success.
+///   - XRAY_MODE_NOT_FOUND if there is no implementation associated with Mode;
+///     does not update the currently installed implementation.
+XRayLogRegisterStatus __xray_log_select_mode(const char *Mode);
+
 /// This function removes the currently installed implementation. It will also
 /// uninstall any handlers that have been previously installed. It does NOT
 /// unpatch the instrumentation sleds.
@@ -220,12 +261,19 @@ XRayLogFlushStatus __xray_log_flushLog();
 
 namespace __xray {
 
-// Options used by the LLVM XRay FDR implementation.
+/// Options used by the LLVM XRay FDR logging implementation.
 struct FDRLoggingOptions {
   bool ReportErrors = false;
   int Fd = -1;
 };
 
+/// Options used by the LLVM XRay Basic (Naive) logging implementation.
+struct BasicLoggingOptions {
+  int DurationFilterMicros = 0;
+  size_t MaxStackDepth = 0;
+  size_t ThreadBufferSize = 0;
+};
+
 } // namespace __xray
 
 #endif // XRAY_XRAY_LOG_INTERFACE_H
diff --git a/include/xray/xray_records.h b/include/xray/xray_records.h
index feb8d228b2fd..d4b7b4c31a9a 100644
--- a/include/xray/xray_records.h
+++ b/include/xray/xray_records.h
@@ -17,6 +17,8 @@
 #ifndef XRAY_XRAY_RECORDS_H
 #define XRAY_XRAY_RECORDS_H
 
+#include <cstdint>
+
 namespace __xray {
 
 enum FileTypes {
@@ -65,18 +67,23 @@ static_assert(sizeof(XRayFileHeader) == 32, "XRayFileHeader != 32 bytes");
 
 enum RecordTypes {
   NORMAL = 0,
+  ARG_PAYLOAD = 1,
 };
 
 struct alignas(32) XRayRecord {
   // This is the type of the record being written. We use 16 bits to allow us to
   // treat this as a discriminant, and so that the first 4 bytes get packed
   // properly. See RecordTypes for more supported types.
-  uint16_t RecordType = 0;
+  uint16_t RecordType = RecordTypes::NORMAL;
 
   // The CPU where the thread is running. We assume number of CPUs <= 256.
   uint8_t CPU = 0;
 
-  // The type of the event. Usually either ENTER = 0 or EXIT = 1.
+  // The type of the event. One of the following:
+  //   ENTER = 0
+  //   EXIT = 1
+  //   TAIL_EXIT = 2
+  //   ENTER_ARG = 3
   uint8_t Type = 0;
 
   // The function ID for the record.
@@ -94,6 +101,32 @@ struct alignas(32) XRayRecord {
 
 static_assert(sizeof(XRayRecord) == 32, "XRayRecord != 32 bytes");
 
+struct alignas(32) XRayArgPayload {
+  // We use the same 16 bits as a discriminant for the records in the log here
+  // too, and so that the first 4 bytes are packed properly.
+  uint16_t RecordType = RecordTypes::ARG_PAYLOAD;
+
+  // Add a few bytes to pad.
+  uint8_t Padding[2] = {};
+
+  // The function ID for the record.
+  int32_t FuncId = 0;
+
+  // The thread ID for the currently running thread.
+  uint32_t TId = 0;
+
+  // Add more padding.
+  uint8_t Padding2[4] = {};
+
+  // The argument payload.
+  uint64_t Arg = 0;
+
+  // The rest of this record ought to be left as padding.
+  uint8_t TailPadding[8] = {};
+} __attribute__((packed));
+
+static_assert(sizeof(XRayArgPayload) == 32, "XRayArgPayload != 32 bytes");
+
 } // namespace __xray
 
 #endif // XRAY_XRAY_RECORDS_H