[XRay][compiler-rt] Support in-memory processing of FDR mode logs

Summary:
This change allows for handling the in-memory data associated with the
FDR mode implementation through the new `__xray_log_process_buffers`
API. With this change, we can now allow users to process the data
in-memory of the process instead of through writing files.

This for example allows users to stream the data of the FDR logging
implementation through network sockets, or through other mechanisms
instead of saving them to local files.

We introduce an FDR-specific flag, for "no_file_flush" which lets the
flushing logic skip opening/writing to files.

This option can be defaulted to `true` when building the compiler-rt
XRay runtime through the `XRAY_FDR_OPTIONS` preprocessor macro.

Reviewers: kpw, echristo, pelikan, eizan

Subscribers: llvm-commits

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

llvm-svn: 332208

1 files changed, 91 insertions, 23 deletions

diff --git a/compiler-rt/lib/xray/xray_buffer_queue.h b/compiler-rt/lib/xray/xray_buffer_queue.h
index 32f852d40fa..3f61263a8ed 100644
--- a/compiler-rt/lib/xray/xray_buffer_queue.h
+++ b/compiler-rt/lib/xray/xray_buffer_queue.h
@@ -15,9 +15,10 @@
 #ifndef XRAY_BUFFER_QUEUE_H
 #define XRAY_BUFFER_QUEUE_H
 
-#include <cstddef>
 #include "sanitizer_common/sanitizer_atomic.h"
+#include "sanitizer_common/sanitizer_common.h"
 #include "sanitizer_common/sanitizer_mutex.h"
+#include <cstddef>
 
 namespace __xray {
 
@@ -27,7 +28,7 @@ namespace __xray {
 /// the "flight data recorder" (FDR) mode to support ongoing XRay function call
 /// trace collection.
 class BufferQueue {
- public:
+public:
   struct alignas(64) BufferExtents {
     __sanitizer::atomic_uint64_t Size;
   };
@@ -35,10 +36,10 @@ class BufferQueue {
   struct Buffer {
     void *Data = nullptr;
     size_t Size = 0;
-    BufferExtents* Extents;
+    BufferExtents *Extents;
   };
 
- private:
+private:
   struct BufferRep {
     // The managed buffer.
     Buffer Buff;
@@ -48,10 +49,67 @@ class BufferQueue {
     bool Used = false;
   };
 
+  // This models a ForwardIterator. |T| Must be either a `Buffer` or `const
+  // Buffer`. Note that we only advance to the "used" buffers, when
+  // incrementing, so that at dereference we're always at a valid point.
+  template <class T> class Iterator {
+  public:
+    BufferRep *Buffers = nullptr;
+    size_t Offset = 0;
+    size_t Max = 0;
+
+    Iterator &operator++() {
+      DCHECK_NE(Offset, Max);
+      do {
+        ++Offset;
+      } while (!Buffers[Offset].Used && Offset != Max);
+      return *this;
+    }
+
+    Iterator operator++(int) {
+      Iterator C = *this;
+      ++(*this);
+      return C;
+    }
+
+    T &operator*() const { return Buffers[Offset].Buff; }
+
+    T *operator->() const { return &(Buffers[Offset].Buff); }
+
+    Iterator(BufferRep *Root, size_t O, size_t M)
+        : Buffers(Root), Offset(O), Max(M) {
+      // We want to advance to the first Offset where the 'Used' property is
+      // true, or to the end of the list/queue.
+      while (!Buffers[Offset].Used && Offset != Max) {
+        ++Offset;
+      }
+    }
+
+    Iterator() = default;
+    Iterator(const Iterator &) = default;
+    Iterator(Iterator &&) = default;
+    Iterator &operator=(const Iterator &) = default;
+    Iterator &operator=(Iterator &&) = default;
+    ~Iterator() = default;
+
+    template <class V>
+    friend bool operator==(const Iterator &L, const Iterator<V> &R) {
+      DCHECK_EQ(L.Max, R.Max);
+      return L.Buffers == R.Buffers && L.Offset == R.Offset;
+    }
+
+    template <class V>
+    friend bool operator!=(const Iterator &L, const Iterator<V> &R) {
+      return !(L == R);
+    }
+  };
+
   // Size of each individual Buffer.
   size_t BufferSize;
 
   BufferRep *Buffers;
+
+  // Amount of pre-allocated buffers.
   size_t BufferCount;
 
   __sanitizer::SpinMutex Mutex;
@@ -70,7 +128,7 @@ class BufferQueue {
   // Count of buffers that have been handed out through 'getBuffer'.
   size_t LiveBuffers;
 
- public:
+public:
   enum class ErrorCode : unsigned {
     Ok,
     NotEnoughMemory,
@@ -81,16 +139,16 @@ class BufferQueue {
 
   static const char *getErrorString(ErrorCode E) {
     switch (E) {
-      case ErrorCode::Ok:
-        return "(none)";
-      case ErrorCode::NotEnoughMemory:
-        return "no available buffers in the queue";
-      case ErrorCode::QueueFinalizing:
-        return "queue already finalizing";
-      case ErrorCode::UnrecognizedBuffer:
-        return "buffer being returned not owned by buffer queue";
-      case ErrorCode::AlreadyFinalized:
-        return "queue already finalized";
+    case ErrorCode::Ok:
+      return "(none)";
+    case ErrorCode::NotEnoughMemory:
+      return "no available buffers in the queue";
+    case ErrorCode::QueueFinalizing:
+      return "queue already finalizing";
+    case ErrorCode::UnrecognizedBuffer:
+      return "buffer being returned not owned by buffer queue";
+    case ErrorCode::AlreadyFinalized:
+      return "queue already finalized";
     }
     return "unknown error";
   }
@@ -141,19 +199,29 @@ class BufferQueue {
   /// Applies the provided function F to each Buffer in the queue, only if the
   /// Buffer is marked 'used' (i.e. has been the result of getBuffer(...) and a
   /// releaseBuffer(...) operation).
-  template <class F>
-  void apply(F Fn) {
+  template <class F> void apply(F Fn) {
     __sanitizer::SpinMutexLock G(&Mutex);
-    for (auto I = Buffers, E = Buffers + BufferCount; I != E; ++I) {
-      const auto &T = *I;
-      if (T.Used) Fn(T.Buff);
-    }
+    for (auto I = begin(), E = end(); I != E; ++I)
+      Fn(*I);
+  }
+
+  using const_iterator = Iterator<const Buffer>;
+  using iterator = Iterator<Buffer>;
+
+  /// Provides iterator access to the raw Buffer instances.
+  iterator begin() const { return iterator(Buffers, 0, BufferCount); }
+  const_iterator cbegin() const {
+    return const_iterator(Buffers, 0, BufferCount);
+  }
+  iterator end() const { return iterator(Buffers, BufferCount, BufferCount); }
+  const_iterator cend() const {
+    return const_iterator(Buffers, BufferCount, BufferCount);
   }
 
   // Cleans up allocated buffers.
   ~BufferQueue();
 };
 
-}  // namespace __xray
+} // namespace __xray
 
-#endif  // XRAY_BUFFER_QUEUE_H
+#endif // XRAY_BUFFER_QUEUE_H