1 files changed, 42 insertions, 26 deletions

diff --git a/lib/xray/xray_buffer_queue.cc b/lib/xray/xray_buffer_queue.cc
index 7e5462fb8e11..7ba755ac3069 100644
--- a/lib/xray/xray_buffer_queue.cc
+++ b/lib/xray/xray_buffer_queue.cc
@@ -13,53 +13,69 @@
 //
 //===----------------------------------------------------------------------===//
 #include "xray_buffer_queue.h"
-#include <cassert>
+#include "sanitizer_common/sanitizer_common.h"
+#include "sanitizer_common/sanitizer_libc.h"
+
 #include <cstdlib>
+#include <tuple>
 
 using namespace __xray;
+using namespace __sanitizer;
 
-BufferQueue::BufferQueue(std::size_t B, std::size_t N)
-    : BufferSize(B), Buffers(N), Mutex(), OwnedBuffers(), Finalizing(false) {
-  for (auto &Buf : Buffers) {
+BufferQueue::BufferQueue(std::size_t B, std::size_t N, bool &Success)
+    : BufferSize(B), Buffers(N), Mutex(), OwnedBuffers(), Finalizing{0} {
+  for (auto &T : Buffers) {
     void *Tmp = malloc(BufferSize);
+    if (Tmp == nullptr) {
+      Success = false;
+      return;
+    }
+
+    auto &Buf = std::get<0>(T);
     Buf.Buffer = Tmp;
     Buf.Size = B;
-    if (Tmp != 0)
-      OwnedBuffers.insert(Tmp);
+    OwnedBuffers.emplace(Tmp);
   }
+  Success = true;
 }
 
-std::error_code BufferQueue::getBuffer(Buffer &Buf) {
-  if (Finalizing.load(std::memory_order_acquire))
-    return std::make_error_code(std::errc::state_not_recoverable);
-  std::lock_guard<std::mutex> Guard(Mutex);
+BufferQueue::ErrorCode BufferQueue::getBuffer(Buffer &Buf) {
+  if (__sanitizer::atomic_load(&Finalizing, __sanitizer::memory_order_acquire))
+    return ErrorCode::QueueFinalizing;
+  __sanitizer::BlockingMutexLock Guard(&Mutex);
   if (Buffers.empty())
-    return std::make_error_code(std::errc::not_enough_memory);
-  Buf = Buffers.front();
+    return ErrorCode::NotEnoughMemory;
+  auto &T = Buffers.front();
+  auto &B = std::get<0>(T);
+  Buf = B;
+  B.Buffer = nullptr;
+  B.Size = 0;
   Buffers.pop_front();
-  return {};
+  return ErrorCode::Ok;
 }
 
-std::error_code BufferQueue::releaseBuffer(Buffer &Buf) {
+BufferQueue::ErrorCode BufferQueue::releaseBuffer(Buffer &Buf) {
   if (OwnedBuffers.count(Buf.Buffer) == 0)
-    return std::make_error_code(std::errc::argument_out_of_domain);
-  std::lock_guard<std::mutex> Guard(Mutex);
-  Buffers.push_back(Buf);
+    return ErrorCode::UnrecognizedBuffer;
+  __sanitizer::BlockingMutexLock Guard(&Mutex);
+
+  // Now that the buffer has been released, we mark it as "used".
+  Buffers.emplace(Buffers.end(), Buf, true /* used */);
   Buf.Buffer = nullptr;
-  Buf.Size = BufferSize;
-  return {};
+  Buf.Size = 0;
+  return ErrorCode::Ok;
 }
 
-std::error_code BufferQueue::finalize() {
-  if (Finalizing.exchange(true, std::memory_order_acq_rel))
-    return std::make_error_code(std::errc::state_not_recoverable);
-  return {};
+BufferQueue::ErrorCode BufferQueue::finalize() {
+  if (__sanitizer::atomic_exchange(&Finalizing, 1,
+                                   __sanitizer::memory_order_acq_rel))
+    return ErrorCode::QueueFinalizing;
+  return ErrorCode::Ok;
 }
 
 BufferQueue::~BufferQueue() {
-  for (auto &Buf : Buffers) {
+  for (auto &T : Buffers) {
+    auto &Buf = std::get<0>(T);
     free(Buf.Buffer);
-    Buf.Buffer = nullptr;
-    Buf.Size = 0;
   }
 }