[XRay] Refcount backing store for buffers

[dberris]

Summary

This change implements the ref-counting for backing stores associated
with generational buffer management. We do this as an implementation
detail of the buffer queue, instead of exposing this to the interface.

This change allows us to keep the buffer queue interface and usage model
the same.

Depends on D53551.

[dberris]

Update description.

[dberris]

Remove changes from D53551.

[dberris]

Update description for real this time.

[mboerger]

LGTM

compiler-rt/lib/xray/xray_buffer_queue.h

 #include "sanitizer_common/sanitizer_mutex.h"
 #include "xray_defs.h"
 #include <cstddef>
 #include <cstdint>
 
 namespace __xray {
 
 /// BufferQueue implements a circular queue of fixed sized buffers (much like a
-/// freelist) but is concerned mostly with making it really quick to initialise,
-/// finalise, and get/return buffers to the queue. This is one key component of
-/// the "flight data recorder" (FDR) mode to support ongoing XRay function call
+/// freelist) but is concerned with making it quick to initialise, finalise, and
+/// get from or return buffers to the queue. This is one key component of the
+/// "flight data recorder" (FDR) mode to support ongoing XRay function call
 /// trace collection.
 class BufferQueue {
 public:
   struct Buffer {
     atomic_uint64_t Extents{0};
     uint64_t Generation{0};
     void *Data = nullptr;
     size_t Size = 0;
+
+  private:
+    friend class BufferQueue;
+    unsigned char *BackingStore = nullptr;
+    size_t Count = 0;
   };
 
   struct BufferRep {
     // The managed buffer.
     Buffer Buff;
 
     // This is true if the buffer has been returned to the available queue, and
     // is considered "used" by another thread.
private:
 
   // Count of buffers that have been handed out through 'getBuffer'.
   size_t LiveBuffers;
 
   // We use a generation number to identify buffers and which generation they're
   // associated with.
   atomic_uint64_t Generation;
 
+  /// Releases references to the buffers backed by the current buffer queue.
+  void cleanupBuffers();
+
 public:
   enum class ErrorCode : unsigned {
     Ok,
     NotEnoughMemory,
     QueueFinalizing,
     UnrecognizedBuffer,
     AlreadyFinalized,
     AlreadyInitialized,

compiler-rt/lib/xray/xray_buffer_queue.cc

 //===-- xray_buffer_queue.cc -----------------------------------*- C++ -*-===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 //
 // This file is a part of XRay, a dynamic runtime instruementation system.
 //
 // Defines the interface for a buffer queue implementation.
 //
 //===----------------------------------------------------------------------===//
 #include "xray_buffer_queue.h"
+#include "sanitizer_common/sanitizer_atomic.h"
 #include "sanitizer_common/sanitizer_common.h"
 #include "sanitizer_common/sanitizer_libc.h"
 #include "sanitizer_common/sanitizer_posix.h"
 #include "xray_allocator.h"
 #include "xray_defs.h"
 #include <memory>
 #include <sys/mman.h>
 
 using namespace __xray;
 using namespace __sanitizer;
 
+namespace {
+
+void decRefCount(unsigned char *ControlBlock, size_t Size, size_t Count) {
+  if (ControlBlock == nullptr)
+    return;
+  auto *RefCount = reinterpret_cast<atomic_uint64_t *>(ControlBlock);
+  if (atomic_fetch_sub(RefCount, 1, memory_order_acq_rel) == 1)
+    deallocateBuffer(ControlBlock, (Size * Count) + kCacheLineSize);
+}
+
+void incRefCount(unsigned char *ControlBlock) {
+  if (ControlBlock == nullptr)
+    return;
+  auto *RefCount = reinterpret_cast<atomic_uint64_t *>(ControlBlock);
+  atomic_fetch_add(RefCount, 1, memory_order_acq_rel);
+}
+
+} // namespace
+
 BufferQueue::ErrorCode BufferQueue::init(size_t BS, size_t BC) {
   SpinMutexLock Guard(&Mutex);
 
   if (!finalizing())
     return BufferQueue::ErrorCode::AlreadyInitialized;
 
+  cleanupBuffers();
+
   bool Success = false;
   BufferSize = BS;
   BufferCount = BC;
-  BackingStore = allocateBuffer(BufferSize * BufferCount);
+  BackingStore = allocateBuffer((BufferSize * BufferCount) + kCacheLineSize);
   if (BackingStore == nullptr)
     return BufferQueue::ErrorCode::NotEnoughMemory;
 
   auto CleanupBackingStore = __sanitizer::at_scope_exit([&, this] {
     if (Success)
       return;
-    deallocateBuffer(BackingStore, BufferSize * BufferCount);
+    deallocateBuffer(BackingStore, (BufferSize * BufferCount) + kCacheLineSize);
+    BackingStore = nullptr;
   });
 
   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);
 
   Success = true;
+
+  // First, we initialize the refcount in the RefCountedBackingStore, which we
+  // treat as being at the start of the BackingStore pointer.
+  auto ControlBlock = reinterpret_cast<atomic_uint64_t *>(BackingStore);
+  atomic_store(ControlBlock, 1, memory_order_release);
+
   for (size_t i = 0; i < BufferCount; ++i) {
     auto &T = Buffers[i];
     auto &Buf = T.Buff;
     atomic_store(&Buf.Extents, 0, memory_order_release);
     Buf.Generation = generation();
-    Buf.Data = reinterpret_cast<char *>(BackingStore) + (BufferSize * i);
+    Buf.Data = BackingStore + kCacheLineSize + (BufferSize * i);
     Buf.Size = BufferSize;
+    Buf.BackingStore = BackingStore;
+    Buf.Count = BufferCount;
     T.Used = false;
   }
 
   Next = Buffers;
   First = Buffers;
   LiveBuffers = 0;
   atomic_store(&Finalizing, 0, memory_order_release);
   return BufferQueue::ErrorCode::Ok;
BufferRep *B = nullptr;
     if (LiveBuffers == BufferCount)
       return ErrorCode::NotEnoughMemory;
     B = Next++;
     if (Next == (Buffers + BufferCount))
       Next = Buffers;
     ++LiveBuffers;
   }
 
+  incRefCount(BackingStore);
   Buf.Data = B->Buff.Data;
   Buf.Generation = generation();
   Buf.Size = B->Buff.Size;
+  Buf.BackingStore = BackingStore;
+  Buf.Count = BufferCount;
   B->Used = true;
   return ErrorCode::Ok;
 }
 
 BufferQueue::ErrorCode BufferQueue::releaseBuffer(Buffer &Buf) {
   // Check whether the buffer being referred to is within the bounds of the
   // backing store's range.
   BufferRep *B = nullptr;
   {
     SpinMutexLock Guard(&Mutex);
     if (Buf.Data < BackingStore ||
         Buf.Data > reinterpret_cast<char *>(BackingStore) +
                        (BufferCount * BufferSize)) {
       if (Buf.Generation != generation()) {
+        decRefCount(Buf.BackingStore, Buf.Size, Buf.Count);
         Buf.Data = nullptr;
         Buf.Size = 0;
         Buf.Generation = 0;
+        Buf.Count = 0;
+        Buf.BackingStore = nullptr;
         return BufferQueue::ErrorCode::Ok;
       }
       return BufferQueue::ErrorCode::UnrecognizedBuffer;
     }
 
     if (LiveBuffers == 0) {
+      decRefCount(Buf.BackingStore, Buf.Size, Buf.Count);
       Buf.Data = nullptr;
       Buf.Size = Buf.Size;
       Buf.Generation = 0;
+      Buf.BackingStore = nullptr;
+      Buf.Count = 0;
       return ErrorCode::Ok;
     }
 
     --LiveBuffers;
     B = First++;
     if (First == (Buffers + BufferCount))
       First = Buffers;
   }
 
   // Now that the buffer has been released, we mark it as "used".
   B->Buff.Data = Buf.Data;
   B->Buff.Size = Buf.Size;
   B->Buff.Generation = Buf.Generation;
+  B->Buff.BackingStore = Buf.BackingStore;
+  B->Buff.Count = Buf.Count;
   B->Used = true;
+  decRefCount(Buf.BackingStore, Buf.Size, Buf.Count);
   atomic_store(&B->Buff.Extents,
                atomic_load(&Buf.Extents, memory_order_acquire),
                memory_order_release);
   Buf.Data = nullptr;
   Buf.Size = 0;
   Buf.Generation = 0;
+  Buf.BackingStore = nullptr;
+  Buf.Count = 0;
   return ErrorCode::Ok;
 }
 
 BufferQueue::ErrorCode BufferQueue::finalize() {
   if (atomic_exchange(&Finalizing, 1, memory_order_acq_rel))
     return ErrorCode::QueueFinalizing;
   return ErrorCode::Ok;
 }
 
-BufferQueue::~BufferQueue() {
+void BufferQueue::cleanupBuffers() {
   for (auto B = Buffers, E = Buffers + BufferCount; B != E; ++B)
     B->~BufferRep();
   deallocateBuffer(Buffers, BufferCount);
-  deallocateBuffer(BackingStore, BufferSize * BufferCount);
+  decRefCount(BackingStore, BufferSize, BufferCount);
+  BackingStore = nullptr;
+  Buffers = nullptr;
+  BufferCount = 0;
+  BufferSize = 0;
 }
+
+BufferQueue::~BufferQueue() { cleanupBuffers(); }