[XRay][compiler-rt] Generational Buffer Management

Summary:
This change updates the buffer queue implementation to support using a
generation number to identify the lifetime of buffers. This first part
introduces the notion of the generation number, without changing the way
we handle the buffers yet.

What's missing here is the cleanup of the buffers. Ideally we'll keep
the two most recent generations. We need to ensure that before we do any
writes to the buffers, that we check the generation number(s) first.

Those changes will follow-on from this change.

Depends on D52588.

Reviewers: mboerger, eizan

Subscribers: llvm-commits, jfb

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

llvm-svn: 344670

1 files changed, 95 insertions, 55 deletions

diff --git a/compiler-rt/lib/xray/xray_buffer_queue.cc b/compiler-rt/lib/xray/xray_buffer_queue.cc
index 5a88ecd3399..c17138d9972 100644
--- a/compiler-rt/lib/xray/xray_buffer_queue.cc
+++ b/compiler-rt/lib/xray/xray_buffer_queue.cc
@@ -24,58 +24,85 @@
 using namespace __xray;
 using namespace __sanitizer;
 
-BufferQueue::BufferQueue(size_t B, size_t N,
-                         bool &Success) XRAY_NEVER_INSTRUMENT
-    : BufferSize(B),
-      BufferCount(N),
-      Mutex(),
-      Finalizing{0},
-      BackingStore(allocateBuffer(B *N)),
-      Buffers(initArray<BufferQueue::BufferRep>(N)),
-      Next(Buffers),
-      First(Buffers),
-      LiveBuffers(0) {
-  if (BackingStore == nullptr) {
-    Success = false;
-    return;
-  }
-  if (Buffers == nullptr) {
+BufferQueue::ErrorCode BufferQueue::init(size_t BS, size_t BC) {
+  SpinMutexLock Guard(&Mutex);
+
+  if (!finalizing())
+    return BufferQueue::ErrorCode::AlreadyInitialized;
+
+  bool Success = false;
+  BufferSize = BS;
+  BufferCount = BC;
+  BackingStore = allocateBuffer(BufferSize * BufferCount);
+  if (BackingStore == nullptr)
+    return BufferQueue::ErrorCode::NotEnoughMemory;
+
+  auto CleanupBackingStore = __sanitizer::at_scope_exit([&, this] {
+    if (Success)
+      return;
     deallocateBuffer(BackingStore, BufferSize * BufferCount);
-    Success = false;
-    return;
-  }
+  });
+
+  Buffers = initArray<BufferRep>(BufferCount);
+  if (Buffers == nullptr)
+    return BufferQueue::ErrorCode::NotEnoughMemory;
+
+  // At this point we increment the generation number to associate the buffers
+  // to the new generation.
+  atomic_fetch_add(&Generation, 1, memory_order_acq_rel);
 
-  for (size_t i = 0; i < N; ++i) {
+  Success = true;
+  for (size_t i = 0; i < BufferCount; ++i) {
     auto &T = Buffers[i];
     auto &Buf = T.Buff;
-    Buf.Data = reinterpret_cast<char *>(BackingStore) + (BufferSize * i);
-    Buf.Size = B;
     atomic_store(&Buf.Extents, 0, memory_order_release);
+    Buf.Generation = generation();
+    Buf.Data = reinterpret_cast<char *>(BackingStore) + (BufferSize * i);
+    Buf.Size = BufferSize;
     T.Used = false;
   }
-  Success = true;
+
+  Next = Buffers;
+  First = Buffers;
+  LiveBuffers = 0;
+  atomic_store(&Finalizing, 0, memory_order_release);
+  return BufferQueue::ErrorCode::Ok;
+}
+
+BufferQueue::BufferQueue(size_t B, size_t N,
+                         bool &Success) XRAY_NEVER_INSTRUMENT
+    : BufferSize(B),
+      BufferCount(N),
+      Mutex(),
+      Finalizing{1},
+      BackingStore(nullptr),
+      Buffers(nullptr),
+      Next(Buffers),
+      First(Buffers),
+      LiveBuffers(0),
+      Generation{0} {
+  Success = init(B, N) == BufferQueue::ErrorCode::Ok;
 }
 
 BufferQueue::ErrorCode BufferQueue::getBuffer(Buffer &Buf) {
   if (atomic_load(&Finalizing, memory_order_acquire))
     return ErrorCode::QueueFinalizing;
 
-  SpinMutexLock Guard(&Mutex);
-  if (LiveBuffers == BufferCount)
-    return ErrorCode::NotEnoughMemory;
-
-  auto &T = *Next;
-  auto &B = T.Buff;
-  auto Extents = atomic_load(&B.Extents, memory_order_acquire);
-  atomic_store(&Buf.Extents, Extents, memory_order_release);
-  Buf.Data = B.Data;
-  Buf.Size = B.Size;
-  T.Used = true;
-  ++LiveBuffers;
-
-  if (++Next == (Buffers + BufferCount))
-    Next = Buffers;
+  BufferRep *B = nullptr;
+  {
+    SpinMutexLock Guard(&Mutex);
+    if (LiveBuffers == BufferCount)
+      return ErrorCode::NotEnoughMemory;
+    B = Next++;
+    if (Next == (Buffers + BufferCount))
+      Next = Buffers;
+    ++LiveBuffers;
+  }
 
+  Buf.Data = B->Buff.Data;
+  Buf.Generation = generation();
+  Buf.Size = B->Buff.Size;
+  B->Used = true;
   return ErrorCode::Ok;
 }
 
@@ -84,29 +111,42 @@ BufferQueue::ErrorCode BufferQueue::releaseBuffer(Buffer &Buf) {
   // backing store's range.
   if (Buf.Data < BackingStore ||
       Buf.Data >
-          reinterpret_cast<char *>(BackingStore) + (BufferCount * BufferSize))
-    return ErrorCode::UnrecognizedBuffer;
+          reinterpret_cast<char *>(BackingStore) + (BufferCount * BufferSize)) {
+    if (Buf.Generation != generation()) {
+      Buf.Data = nullptr;
+      Buf.Size = 0;
+      Buf.Generation = 0;
+      return BufferQueue::ErrorCode::Ok;
+    }
+    return BufferQueue::ErrorCode::UnrecognizedBuffer;
+  }
 
-  SpinMutexLock Guard(&Mutex);
+  BufferRep *B = nullptr;
+  {
+    SpinMutexLock Guard(&Mutex);
+
+    // This points to a semantic bug, we really ought to not be releasing more
+    // buffers than we actually get.
+    if (LiveBuffers == 0)
+      return ErrorCode::NotEnoughMemory;
 
-  // This points to a semantic bug, we really ought to not be releasing more
-  // buffers than we actually get.
-  if (LiveBuffers == 0)
-    return ErrorCode::NotEnoughMemory;
+    --LiveBuffers;
+    B = First++;
+    if (First == (Buffers + BufferCount))
+      First = Buffers;
+  }
 
   // Now that the buffer has been released, we mark it as "used".
-  auto Extents = atomic_load(&Buf.Extents, memory_order_acquire);
-  atomic_store(&First->Buff.Extents, Extents, memory_order_release);
-  First->Buff.Data = Buf.Data;
-  First->Buff.Size = Buf.Size;
-  First->Used = true;
+  B->Buff.Data = Buf.Data;
+  B->Buff.Size = Buf.Size;
+  B->Buff.Generation = Buf.Generation;
+  B->Used = true;
+  atomic_store(&B->Buff.Extents,
+               atomic_load(&Buf.Extents, memory_order_acquire),
+               memory_order_release);
   Buf.Data = nullptr;
   Buf.Size = 0;
-  atomic_store(&Buf.Extents, 0, memory_order_release);
-  --LiveBuffers;
-  if (++First == (Buffers + BufferCount))
-    First = Buffers;
-
+  Buf.Generation = 0;
   return ErrorCode::Ok;
 }