[XRay] Move buffer extents back to the heap

[dberris]

Summary

This change addresses an issue which shows up with the synchronised race
between threads writing into a buffer, and another thread reading the
buffer.

In a lot of cases, we cannot guarantee that threads will always see the
signal to finalise their buffers in time despite the grace periods and
state machine maintained through atomic variables. This change addresses
it by ensuring that the same instance being updated to indicate how much
of the buffer is "used" by the writing thread is the same instance being
read by the thread processing the buffer to be written out to disk or
handled through the iterators.

To do this, we ensure that all the "extents" instances live in their own
the backing store, in a different contiguous page from the
buffer-specific backing store. We also take precautions to ensure that
the atomic variables are cache-line-sized to prevent false-sharing from
unnecessarily causing cache contention on unrelated writes/reads.

It's feasible that we may in the future be able to move the storage of
the extents objects into the single backing store, slightly changing the
way to compute the size(s) of the buffers, but in the meantime we'll
settle for the isolation afforded by having a different backing store
for the extents instances.

[mboerger]

compiler-rt/lib/xray/xray_buffer_queue.cc

73

Shouldn't this be cache line size minus size of Extends?

[dberris]

compiler-rt/lib/xray/xray_buffer_queue.cc

73

Nope, the union is meant to ensure that the size of the member is the longer of these two elements.

[mboerger]

LGTM

[dberris]

clang-format with LLVM style

[mboerger]

LGTM

compiler-rt/lib/xray/tests/unit/test_helpers.cc

std::string serialize(BufferQueue &Buffers, int32_t Version) {
   Header->Version = Version;
   Header->Type = FileTypes::FDR_LOG;
   Header->CycleFrequency = 3e9;
   Header->ConstantTSC = 1;
   Header->NonstopTSC = 1;
   Serialized.append(reinterpret_cast<const char *>(&HeaderStorage),
                     sizeof(XRayFileHeader));
   Buffers.apply([&](const BufferQueue::Buffer &B) {
-    auto Size = atomic_load_relaxed(&B.Extents);
+    auto Size = atomic_load_relaxed(B.Extents);
     auto Extents =
         createMetadataRecord<MetadataRecord::RecordKinds::BufferExtents>(Size);
     Serialized.append(reinterpret_cast<const char *>(&Extents),
                       sizeof(Extents));
     Serialized.append(reinterpret_cast<const char *>(B.Data), Size);
   });
   return Serialized;
 }
 
 } // namespace __xray

compiler-rt/lib/xray/xray_buffer_queue.h

 /// BufferQueue implements a circular queue of fixed sized buffers (much like a
 /// 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:
   /// ControlBlock represents the memory layout of how we interpret the backing
-  /// store for all buffers managed by a BufferQueue instance. The ControlBlock
-  /// has the reference count as the first member, sized according to
-  /// platform-specific cache-line size. We never use the Buffer member of the
-  /// union, which is only there for compiler-supported alignment and sizing.
+  /// store for all buffers and extents managed by a BufferQueue instance. The
+  /// ControlBlock has the reference count as the first member, sized according
+  /// to platform-specific cache-line size. We never use the Buffer member of
+  /// the union, which is only there for compiler-supported alignment and
+  /// sizing.
   ///
   /// This ensures that the `Data` member will be placed at least kCacheLineSize
   /// bytes from the beginning of the structure.
   struct ControlBlock {
     union {
       atomic_uint64_t RefCount;
       char Buffer[kCacheLineSize];
     };
 
     /// We need to make this size 1, to conform to the C++ rules for array data
     /// members. Typically, we want to subtract this 1 byte for sizing
     /// information.
     char Data[1];
   };
 
   struct Buffer {
-    atomic_uint64_t Extents{0};
+    atomic_uint64_t *Extents{0};
     uint64_t Generation{0};
     void *Data = nullptr;
     size_t Size = 0;
 
   private:
     friend class BufferQueue;
     ControlBlock *BackingStore = nullptr;
+    ControlBlock *ExtentsBackingStore = 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
private:
   size_t BufferCount;
 
   SpinMutex Mutex;
   atomic_uint8_t Finalizing;
 
   // The collocated ControlBlock and buffer storage.
   ControlBlock *BackingStore;
 
+  // The collocated ControlBlock and extents storage.
+  ControlBlock *ExtentsBackingStore;
+
   // A dynamically allocated array of BufferRep instances.
   BufferRep *Buffers;
 
   // Pointer to the next buffer to be handed out.
   BufferRep *Next;
 
   // Pointer to the entry in the array where the next released buffer will be
   // placed.

compiler-rt/lib/xray/xray_buffer_queue.cc

 #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 {
 
 BufferQueue::ControlBlock *allocControlBlock(size_t Size, size_t Count) {
   auto B =
       allocateBuffer((sizeof(BufferQueue::ControlBlock) - 1) + (Size * Count));
   return B == nullptr ? nullptr
                       : reinterpret_cast<BufferQueue::ControlBlock *>(B);

 }
 
 void incRefCount(BufferQueue::ControlBlock *C) {
   if (C == nullptr)
     return;
   atomic_fetch_add(&C->RefCount, 1, memory_order_acq_rel);
 }
 
+// We use a struct to ensure that we are allocating one atomic_uint64_t per
+// cache line. This allows us to not worry about false-sharing among atomic
+// objects being updated (constantly) by different threads.
+struct ExtentsPadded {
+  union {
+    atomic_uint64_t Extents;
+    unsigned char Storage[kCacheLineSize];
+  };
+};
+
+constexpr size_t kExtentsSize = sizeof(ExtentsPadded);
+
 } // 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 = allocControlBlock(BufferSize, BufferCount);
   if (BackingStore == nullptr)
     return BufferQueue::ErrorCode::NotEnoughMemory;
 
-  auto CleanupBackingStore = __sanitizer::at_scope_exit([&, this] {
+  auto CleanupBackingStore = at_scope_exit([&, this] {
     if (Success)
       return;
     deallocControlBlock(BackingStore, BufferSize, BufferCount);
     BackingStore = nullptr;
   });
 
+  // Initialize enough atomic_uint64_t instances, each
+  ExtentsBackingStore = allocControlBlock(kExtentsSize, BufferCount);
+  if (ExtentsBackingStore == nullptr)
+    return BufferQueue::ErrorCode::NotEnoughMemory;
+
+  auto CleanupExtentsBackingStore = at_scope_exit([&, this] {
+    if (Success)
+      return;
+    deallocControlBlock(ExtentsBackingStore, kExtentsSize, BufferCount);
+    ExtentsBackingStore = 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);
 
   // First, we initialize the refcount in the ControlBlock, which we treat as
   // being at the start of the BackingStore pointer.
   atomic_store(&BackingStore->RefCount, 1, memory_order_release);
+  atomic_store(&ExtentsBackingStore->RefCount, 1, memory_order_release);
 
   // Then we initialise the individual buffers that sub-divide the whole backing
   // store. Each buffer will start at the `Data` member of the ControlBlock, and
   // will be offsets from these locations.
   for (size_t i = 0; i < BufferCount; ++i) {
     auto &T = Buffers[i];
     auto &Buf = T.Buff;
-    atomic_store(&Buf.Extents, 0, memory_order_release);
+    auto *E = reinterpret_cast<ExtentsPadded *>(&ExtentsBackingStore->Data +
+                                                (kExtentsSize * i));
+    Buf.Extents = &E->Extents;
+    atomic_store(Buf.Extents, 0, memory_order_release);
     Buf.Generation = generation();
     Buf.Data = &BackingStore->Data + (BufferSize * i);
     Buf.Size = BufferSize;
     Buf.BackingStore = BackingStore;
+    Buf.ExtentsBackingStore = ExtentsBackingStore;
     Buf.Count = BufferCount;
     T.Used = false;
   }
 
   Next = Buffers;
   First = Buffers;
   LiveBuffers = 0;
   atomic_store(&Finalizing, 0, memory_order_release);
   Success = true;
   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),
+      ExtentsBackingStore(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;
 
   BufferRep *B = nullptr;
   {
     SpinMutexLock Guard(&Mutex);
     if (LiveBuffers == BufferCount)
       return ErrorCode::NotEnoughMemory;
     B = Next++;
     if (Next == (Buffers + BufferCount))
       Next = Buffers;
     ++LiveBuffers;
   }
 
   incRefCount(BackingStore);
+  incRefCount(ExtentsBackingStore);
   Buf = B->Buff;
   Buf.Generation = generation();
   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.Generation != generation() || LiveBuffers == 0) {
       Buf = {};
       decRefCount(Buf.BackingStore, Buf.Size, Buf.Count);
+      decRefCount(Buf.ExtentsBackingStore, kExtentsSize, Buf.Count);
       return BufferQueue::ErrorCode::Ok;
     }
 
     if (Buf.Data < &BackingStore->Data ||
         Buf.Data > &BackingStore->Data + (BufferCount * BufferSize))
       return BufferQueue::ErrorCode::UnrecognizedBuffer;
 
     --LiveBuffers;
     B = First++;
     if (First == (Buffers + BufferCount))
       First = Buffers;
   }
 
   // Now that the buffer has been released, we mark it as "used".
   B->Buff = Buf;
   B->Used = true;
   decRefCount(Buf.BackingStore, Buf.Size, Buf.Count);
-  atomic_store(&B->Buff.Extents,
-               atomic_load(&Buf.Extents, memory_order_acquire),
+  decRefCount(Buf.ExtentsBackingStore, kExtentsSize, Buf.Count);
+  atomic_store(B->Buff.Extents, atomic_load(Buf.Extents, memory_order_acquire),
                memory_order_release);
   Buf = {};
   return ErrorCode::Ok;
 }
 
 BufferQueue::ErrorCode BufferQueue::finalize() {
   if (atomic_exchange(&Finalizing, 1, memory_order_acq_rel))
     return ErrorCode::QueueFinalizing;
   return ErrorCode::Ok;
 }
 
 void BufferQueue::cleanupBuffers() {
   for (auto B = Buffers, E = Buffers + BufferCount; B != E; ++B)
     B->~BufferRep();
   deallocateBuffer(Buffers, BufferCount);
   decRefCount(BackingStore, BufferSize, BufferCount);
+  decRefCount(ExtentsBackingStore, kExtentsSize, BufferCount);
   BackingStore = nullptr;
+  ExtentsBackingStore = nullptr;
   Buffers = nullptr;
   BufferCount = 0;
   BufferSize = 0;
 }
 
 BufferQueue::~BufferQueue() { cleanupBuffers(); }

compiler-rt/lib/xray/xray_fdr_controller.h

bool getNewBuffer() XRAY_NEVER_INSTRUMENT {
 
     W.resetRecord();
     DCHECK_EQ(W.getNextRecord(), B.Data);
     LatestTSC = 0;
     LatestCPU = 0;
     First = true;
     UndoableFunctionEnters = 0;
     UndoableTailExits = 0;
-    atomic_store(&B.Extents, 0, memory_order_release);
+    atomic_store(B.Extents, 0, memory_order_release);
     return true;
   }
 
   bool setupNewBuffer() XRAY_NEVER_INSTRUMENT {
     if (finalized())
       return false;
 
     DCHECK(hasSpace(sizeof(MetadataRecord) * 3));
if (UNLIKELY(!hasSpace(S))) {
         return false;
       if (!setupNewBuffer())
         return false;
     }
 
     if (First) {
       First = false;
       W.resetRecord();
-      atomic_store(&B.Extents, 0, memory_order_release);
+      atomic_store(B.Extents, 0, memory_order_release);
       return setupNewBuffer();
     }
 
     return true;
   }
 
   bool returnBuffer() XRAY_NEVER_INSTRUMENT {
     if (BQ == nullptr)

compiler-rt/lib/xray/xray_fdr_log_writer.h

class FDRLogWriter {
 
   template <class T> void writeRecord(const T &R) {
     internal_memcpy(NextRecord, reinterpret_cast<const char *>(&R), sizeof(T));
     NextRecord += sizeof(T);
     // We need this atomic fence here to ensure that other threads attempting to
     // read the bytes in the buffer will see the writes committed before the
     // extents are updated.
     atomic_thread_fence(memory_order_release);
-    atomic_fetch_add(&Buffer.Extents, sizeof(T), memory_order_acq_rel);
+    atomic_fetch_add(Buffer.Extents, sizeof(T), memory_order_acq_rel);
   }
 
 public:
   explicit FDRLogWriter(BufferQueue::Buffer &B, char *P)
       : Buffer(B), NextRecord(P) {
     DCHECK_NE(Buffer.Data, nullptr);
     DCHECK_NE(NextRecord, nullptr);
   }
public:
   template <size_t N> size_t writeMetadataRecords(MetadataRecord (&Recs)[N]) {
     constexpr auto Size = sizeof(MetadataRecord) * N;
     internal_memcpy(NextRecord, reinterpret_cast<const char *>(Recs), Size);
     NextRecord += Size;
     // We need this atomic fence here to ensure that other threads attempting to
     // read the bytes in the buffer will see the writes committed before the
     // extents are updated.
     atomic_thread_fence(memory_order_release);
-    atomic_fetch_add(&Buffer.Extents, Size, memory_order_acq_rel);
+    atomic_fetch_add(Buffer.Extents, Size, memory_order_acq_rel);
     return Size;
   }
 
   enum class FunctionRecordKind : uint8_t {
     Enter = 0x00,
     Exit = 0x01,
     TailExit = 0x02,
     EnterArg = 0x03,
NextRecord = reinterpret_cast<char *>(internal_memcpy(
                  sizeof(R);
     NextRecord = reinterpret_cast<char *>(internal_memcpy(
                      NextRecord, reinterpret_cast<char *>(&A), sizeof(A))) +
                  sizeof(A);
     // We need this atomic fence here to ensure that other threads attempting to
     // read the bytes in the buffer will see the writes committed before the
     // extents are updated.
     atomic_thread_fence(memory_order_release);
-    atomic_fetch_add(&Buffer.Extents, sizeof(R) + sizeof(A),
+    atomic_fetch_add(Buffer.Extents, sizeof(R) + sizeof(A),
                      memory_order_acq_rel);
     return true;
   }
 
   bool writeCustomEvent(int32_t Delta, const void *Event, int32_t EventSize) {
     // We write the metadata record and the custom event data into the buffer
     // first, before we atomically update the extents for the buffer. This
     // allows us to ensure that any threads reading the extents of the buffer
     // will only ever see the full metadata and custom event payload accounted
     // (no partial writes accounted).
     MetadataRecord R =
         createMetadataRecord<MetadataRecord::RecordKinds::CustomEventMarker>(
             EventSize, Delta);
     NextRecord = reinterpret_cast<char *>(internal_memcpy(
                      NextRecord, reinterpret_cast<char *>(&R), sizeof(R))) +
                  sizeof(R);
     NextRecord = reinterpret_cast<char *>(
                      internal_memcpy(NextRecord, Event, EventSize)) +
                  EventSize;
 
     // We need this atomic fence here to ensure that other threads attempting to
     // read the bytes in the buffer will see the writes committed before the
     // extents are updated.
     atomic_thread_fence(memory_order_release);
-    atomic_fetch_add(&Buffer.Extents, sizeof(R) + EventSize,
+    atomic_fetch_add(Buffer.Extents, sizeof(R) + EventSize,
                      memory_order_acq_rel);
     return true;
   }
 
   bool writeTypedEvent(int32_t Delta, uint16_t EventType, const void *Event,
                        int32_t EventSize) {
     // We do something similar when writing out typed events, see
     // writeCustomEvent(...) above for details.
     MetadataRecord R =
         createMetadataRecord<MetadataRecord::RecordKinds::TypedEventMarker>(
             EventSize, Delta, EventType);
     NextRecord = reinterpret_cast<char *>(internal_memcpy(
                      NextRecord, reinterpret_cast<char *>(&R), sizeof(R))) +
                  sizeof(R);
     NextRecord = reinterpret_cast<char *>(
                      internal_memcpy(NextRecord, Event, EventSize)) +
                  EventSize;
 
     // We need this atomic fence here to ensure that other threads attempting to
     // read the bytes in the buffer will see the writes committed before the
     // extents are updated.
     atomic_thread_fence(memory_order_release);
-    atomic_fetch_add(&Buffer.Extents, EventSize, memory_order_acq_rel);
+    atomic_fetch_add(Buffer.Extents, EventSize, memory_order_acq_rel);
     return true;
   }
 
   char *getNextRecord() const { return NextRecord; }
 
   void resetRecord() {
     NextRecord = reinterpret_cast<char *>(Buffer.Data);
-    atomic_store(&Buffer.Extents, 0, memory_order_release);
+    atomic_store(Buffer.Extents, 0, memory_order_release);
   }
 
   void undoWrites(size_t B) {
     DCHECK_GE(NextRecord - B, reinterpret_cast<char *>(Buffer.Data));
     NextRecord -= B;
-    atomic_fetch_sub(&Buffer.Extents, B, memory_order_acq_rel);
+    atomic_fetch_sub(Buffer.Extents, B, memory_order_acq_rel);
   }
 
 }; // namespace __xray
 
 } // namespace __xray
 
 #endif // COMPILER-RT_LIB_XRAY_XRAY_FDR_LOG_WRITER_H_

compiler-rt/lib/xray/xray_fdr_logging.cc

XRayBuffer fdrIterator(const XRayBuffer B) {
   // handlers see the empty buffers in the queue.
   //
   // We need this atomic fence here to ensure that writes happening to the
   // buffer have been committed before we load the extents atomically. Because
   // the buffer is not explicitly synchronised across threads, we rely on the
   // fence ordering to ensure that writes we expect to have been completed
   // before the fence are fully committed before we read the extents.
   atomic_thread_fence(memory_order_acquire);
-  auto BufferSize = atomic_load(&It->Extents, memory_order_acquire);
+  auto BufferSize = atomic_load(It->Extents, memory_order_acquire);
   SerializedBufferSize = BufferSize + sizeof(MetadataRecord);
   CurrentBuffer = allocateBuffer(SerializedBufferSize);
   if (CurrentBuffer == nullptr)
     return {nullptr, 0};
 
   // Write out the extents as a Metadata Record into the CurrentBuffer.
   MetadataRecord ExtentsRecord;
   ExtentsRecord.Type = uint8_t(RecordType::Metadata);
XRayLogFlushStatus fdrLoggingFlush() XRAY_NEVER_INSTRUMENT {
 
   BQ->apply([&](const BufferQueue::Buffer &B) {
     // Starting at version 2 of the FDR logging implementation, we only write
     // the records identified by the extents of the buffer. We use the Extents
     // from the Buffer and write that out as the first record in the buffer.  We
     // still use a Metadata record, but fill in the extents instead for the
     // data.
     MetadataRecord ExtentsRecord;
-    auto BufferExtents = atomic_load(&B.Extents, memory_order_acquire);
+    auto BufferExtents = atomic_load(B.Extents, memory_order_acquire);
     DCHECK(BufferExtents <= B.Size);
     ExtentsRecord.Type = uint8_t(RecordType::Metadata);
     ExtentsRecord.RecordKind =
         uint8_t(MetadataRecord::RecordKinds::BufferExtents);
     internal_memcpy(ExtentsRecord.Data, &BufferExtents, sizeof(BufferExtents));
     if (BufferExtents > 0) {
       LW->WriteAll(reinterpret_cast<char *>(&ExtentsRecord),
                    reinterpret_cast<char *>(&ExtentsRecord) +