1 files changed, 68 insertions, 24 deletions

diff --git a/lib/xray/xray_buffer_queue.h b/lib/xray/xray_buffer_queue.h
index bf0b7af9df4d..e051695a297b 100644
--- a/lib/xray/xray_buffer_queue.h
+++ b/lib/xray/xray_buffer_queue.h
@@ -15,12 +15,11 @@
 #ifndef XRAY_BUFFER_QUEUE_H
 #define XRAY_BUFFER_QUEUE_H
 
-#include <atomic>
-#include <cstdint>
+#include "sanitizer_common/sanitizer_atomic.h"
+#include "sanitizer_common/sanitizer_mutex.h"
 #include <deque>
-#include <mutex>
-#include <system_error>
 #include <unordered_set>
+#include <utility>
 
 namespace __xray {
 
@@ -33,19 +32,47 @@ class BufferQueue {
 public:
   struct Buffer {
     void *Buffer = nullptr;
-    std::size_t Size = 0;
+    size_t Size = 0;
   };
 
 private:
-  std::size_t BufferSize;
-  std::deque<Buffer> Buffers;
-  std::mutex Mutex;
+  size_t BufferSize;
+
+  // We use a bool to indicate whether the Buffer has been used in this
+  // freelist implementation.
+  std::deque<std::tuple<Buffer, bool>> Buffers;
+  __sanitizer::BlockingMutex Mutex;
   std::unordered_set<void *> OwnedBuffers;
-  std::atomic<bool> Finalizing;
+  __sanitizer::atomic_uint8_t Finalizing;
 
 public:
-  /// Initialise a queue of size |N| with buffers of size |B|.
-  BufferQueue(std::size_t B, std::size_t N);
+  enum class ErrorCode : unsigned {
+    Ok,
+    NotEnoughMemory,
+    QueueFinalizing,
+    UnrecognizedBuffer,
+    AlreadyFinalized,
+  };
+
+  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";
+    }
+    return "unknown error";
+  }
+
+  /// Initialise a queue of size |N| with buffers of size |B|. We report success
+  /// through |Success|.
+  BufferQueue(size_t B, size_t N, bool &Success);
 
   /// Updates |Buf| to contain the pointer to an appropriate buffer. Returns an
   /// error in case there are no available buffers to return when we will run
@@ -58,24 +85,41 @@ public:
   ///   - std::errc::not_enough_memory on exceeding MaxSize.
   ///   - no error when we find a Buffer.
   ///   - std::errc::state_not_recoverable on finalising BufferQueue.
-  std::error_code getBuffer(Buffer &Buf);
+  ErrorCode getBuffer(Buffer &Buf);
 
   /// Updates |Buf| to point to nullptr, with size 0.
   ///
   /// Returns:
   ///   - ...
-  std::error_code releaseBuffer(Buffer &Buf);
-
-  bool finalizing() const { return Finalizing.load(std::memory_order_acquire); }
-
-  // Sets the state of the BufferQueue to finalizing, which ensures that:
-  //
-  //   - All subsequent attempts to retrieve a Buffer will fail.
-  //   - All releaseBuffer operations will not fail.
-  //
-  // After a call to finalize succeeds, all subsequent calls to finalize will
-  // fail with std::errc::state_not_recoverable.
-  std::error_code finalize();
+  ErrorCode releaseBuffer(Buffer &Buf);
+
+  bool finalizing() const {
+    return __sanitizer::atomic_load(&Finalizing,
+                                    __sanitizer::memory_order_acquire);
+  }
+
+  /// Returns the configured size of the buffers in the buffer queue.
+  size_t ConfiguredBufferSize() const { return BufferSize; }
+
+  /// Sets the state of the BufferQueue to finalizing, which ensures that:
+  ///
+  ///   - All subsequent attempts to retrieve a Buffer will fail.
+  ///   - All releaseBuffer operations will not fail.
+  ///
+  /// After a call to finalize succeeds, all subsequent calls to finalize will
+  /// fail with std::errc::state_not_recoverable.
+  ErrorCode finalize();
+
+  /// 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) {
+    __sanitizer::BlockingMutexLock G(&Mutex);
+    for (const auto &T : Buffers) {
+      if (std::get<1>(T))
+        Fn(std::get<0>(T));
+    }
+  }
 
   // Cleans up allocated buffers.
   ~BufferQueue();