[XRay] Use optimistic logging model for FDR mode

[dberris]

Summary

Before this change, the FDR mode implementation relied on at thread-exit
handling to return buffers back to the (global) buffer queue. This
introduces issues with the initialisation of the thread_local objects
which, even through the use of pthread_setspecific(...) may eventually
call into an allocation function. Similar to previous changes in this
line, we're finding that there is a huge potential for deadlocks when
initialising these thread-locals when the memory allocation
implementation is also xray-instrumented.

In this change, we limit the call to pthread_setspecific(...) to provide
a non-null value to associate to the key created with
pthread_key_create(...). While this doesn't completely eliminate the
potential for the deadlock(s), it does allow us to still clean up at
thread exit when we need to. The change is that we don't need to do more
work when starting and ending a thread's lifetime. We also have a test
to make sure that we actually can safely recycle the buffers in case we
end up re-using the buffer(s) available from the queue on multiple
thread entry/exits.

This change cuts across both LLVM and compiler-rt to allow us to update
both the XRay runtime implementation as well as the library support for
loading these new versions of the FDR mode logging. Version 2 of the FDR
logging implementation makes the following changes:

• Introduction of a new 'BufferExtents' metadata record that's outside of the buffer's contents but are written before the actual buffer. This data is associated to the Buffer handed out by the BufferQueue rather than a record that occupies bytes in the actual buffer.
• Removal of the "end of buffer" records. This is in-line with the changes we described above, to allow for optimistic logging without explicit record writing at thread exit.

The optimistic logging model operates under the following assumptions:

• Threads writing to the buffers will potentially race with the thread attempting to flush the log. To avoid this situation from occuring, we make sure that when we've finalized the logging implementation, that threads will see this finalization state on the next write, and either choose to not write records the thread would have written or write the record(s) in two phases -- first write the record(s), then update the extents metadata.
• We change the buffer queue implementation so that once it's handed out a buffer to a thread, that we assume that buffer is marked "used" to be able to capture partial writes. None of this will be safe to handle if threads are racing to write the extents records and the reader thread is attempting to flush the log. The optimism comes from the finalization routine being required to complete before we attempt to flush the log.

This is a fairly significant semantics change for the FDR
implementation. This is why we've decided to update the version number
for FDR mode logs. The tools, however, still need to be able to support
older versions of the log until we finally deprecate those earlier
versions.

Additional changes:

• Move per-thread key initialisation to when we initialise the logging implementation
• Define a grace period for finalisation (default to 100ms).
• Use aligned allocations for the buffers handed out by the BufferQueue.

[dberris]

• fixup: reintroduce thread exit handler, and test

[dberris]

Re-wording a bit to explain the retention of one call to pthread_setspecific.

[dberris]

• fixup: some changes;

[dberris]

Okay, now it's ready for a look. I'm done tweaking this for now. :)

Thanks in advance!

[dberris]

Actually, I just realised there's an issue with this approach -- I may need to ensure that the size_t being written doesn't straddle a cache line so that partial writes couldn't cause the thread reading the data to see an incorrect value of the extents. My options for fixing this are:

• Attaching additional metadata to the Buffer type, that's cache-line aligned instead of using the metadata record in the buffer that's character-aligned.
• Create a dedicated section in the Buffer that's treated as metadata, and is guaranteed cache-aligned. This will mean that we use an aligning allocator (posix_memalign perhaps) to get a suitably-aligned metadata section (buffer header?) that gets updated with the useful data about the buffer.

I'll think about this a little more, and update with a more correct solution.

[dberris]

Update implementation to be more thread-safe, despite accepting potential races
in the implementation.

[dberris]

Update description to reflect new state of the change.

[dberris]

Fix description.

[dberris]

This is now ready for a look @pelikan @kpw.

Thanks in advance!

[dberris]

• fixup: remove unnecessary 'version' propagation

[dberris]

Okay, I'm going to stop updating this now and wait for review. :)

[pelikan]

partial review only

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

167

I can't see a call to munmap(2). Does this mean each TEST_F() is run in a separate process?

compiler-rt/lib/xray/xray_fdr_logging.cc

79

I'd be inclined to use a semaphore instead but I suspect its overhead would be too high due to MOESI traffic for the memory of its internal count. Please confirm.

142–146

We might benefit from turning this into a writev(2):

struct iov { iov_vec = {{ExtendsRecord, 8}, {B.Buffer, Extents} }, iov_cnt = 2 };

294

I trust it's safe to incrementExtents after the memcpy(3) because this will never run concurrently in multiple threads.

345–350

I think it would improve readability if this indented left a bit, as in:

if (TLD.BQ == nullptr) return;

then the rest.

[pelikan]

still not 100% finished - I'll get to it tomorrow. I am convinced the overall concept will work though.

compiler-rt/lib/xray/xray_fdr_logging.cc

294

It's correct because incrementExtents is acq_rel so before that technically nothing happened :-) Marking as done.

compiler-rt/lib/xray/xray_fdr_logging_impl.h

56

You speak of MemPtr but your words mean nothing.

195

Since we have different types of "records", how about calling the variable Metadata?

316–319

Plus million for this coding style. Much better than wrapping in the middle of the screen.

[dberris]

• fixup: Address some review comments

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

167

Good catch -- it just so happens that the lit test runner does it one test case per run. Fixed to register the map to the ScopedFileCloserAndDeleter.

compiler-rt/lib/xray/xray_fdr_logging.cc

79

We don't want a semaphore because that will cause synchronisation across threads. This explicitly not require synchronisation, and allow threads to observe some other signalling (atomic) mechanism instead, i.e. whether the buffer queue is finalising.

142–146

Good idea. Maybe changing the API of retryingWriteAll to support a number of ranges might be useful to do. Probably not in this change.

compiler-rt/lib/xray/xray_fdr_logging_impl.h

316–319

Thank clang-format -- this only happened because we removed the MemPtr argument.

[pelikan]

compiler-rt/lib/xray/xray_fdr_logging_impl.h

316–319

It didn't do that before, and I *really* tried to make it.

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

 //===----------------------------------------------------------------------===//
 #include "xray_fdr_logging.h"
 #include "gtest/gtest.h"
 
 #include <fcntl.h>
 #include <iostream>
 #include <sys/mman.h>
 #include <sys/stat.h>
+#include <sys/syscall.h>
 #include <sys/types.h>
 #include <system_error>
+#include <thread>
 #include <unistd.h>
 
 #include "xray/xray_records.h"
 
 namespace __xray {
 namespace {
 
 constexpr auto kBufferSize = 16384;

 
 TEST(FDRLoggingTest, Simple) {
   FDRLoggingOptions Options;
   Options.ReportErrors = true;
   char TmpFilename[] = "fdr-logging-test.XXXXXX";
   Options.Fd = mkstemp(TmpFilename);
   ASSERT_NE(Options.Fd, -1);
   ASSERT_EQ(fdrLoggingInit(kBufferSize, kBufferMax, &Options,
-                            sizeof(FDRLoggingOptions)),
             XRayLogInitStatus::XRAY_LOG_INITIALIZED);
   fdrLoggingHandleArg0(1, XRayEntryType::ENTRY);
   fdrLoggingHandleArg0(1, XRayEntryType::EXIT);
   ASSERT_EQ(fdrLoggingFinalize(), XRayLogInitStatus::XRAY_LOG_FINALIZED);
   ASSERT_EQ(fdrLoggingFlush(), XRayLogFlushStatus::XRAY_LOG_FLUSHED);
 
   // To do this properly, we have to close the file descriptor then re-open the
   // file for reading this time.
   ASSERT_EQ(close(Options.Fd), 0);
   int Fd = open(TmpFilename, O_RDONLY);
   ASSERT_NE(-1, Fd);
   ScopedFileCloserAndDeleter Guard(Fd, TmpFilename);
   auto Size = lseek(Fd, 0, SEEK_END);
   ASSERT_NE(Size, 0);
   // Map the file contents.
   const char *Contents = static_cast<const char *>(
       mmap(NULL, Size, PROT_READ, MAP_PRIVATE, Fd, 0));
   ASSERT_NE(Contents, nullptr);
 
   XRayFileHeader H;
   memcpy(&H, Contents, sizeof(XRayFileHeader));
-  ASSERT_EQ(H.Version, 1);
+  ASSERT_EQ(H.Version, 2);
   ASSERT_EQ(H.Type, FileTypes::FDR_LOG);
 
-  // We require one buffer at least to have the "start of buffer" metadata
-  // record.
-  MetadataRecord MDR;
-  memcpy(&MDR, Contents + sizeof(XRayFileHeader), sizeof(MetadataRecord));
-  ASSERT_EQ(MDR.RecordKind, uint8_t(MetadataRecord::RecordKinds::NewBuffer));
+  // We require one buffer at least to have the "extents" metadata record,
+  // followed by the NewBuffer record.
+  MetadataRecord MDR0, MDR1;
+  memcpy(&MDR0, Contents + sizeof(XRayFileHeader), sizeof(MetadataRecord));
+  memcpy(&MDR1, Contents + sizeof(XRayFileHeader) + sizeof(MetadataRecord),
+         sizeof(MetadataRecord));
+  ASSERT_EQ(MDR0.RecordKind,
+            uint8_t(MetadataRecord::RecordKinds::BufferExtents));
+  ASSERT_EQ(MDR1.RecordKind, uint8_t(MetadataRecord::RecordKinds::NewBuffer));
 }
 
 TEST(FDRLoggingTest, Multiple) {
   FDRLoggingOptions Options;
   char TmpFilename[] = "fdr-logging-test.XXXXXX";
   Options.Fd = mkstemp(TmpFilename);
   ASSERT_NE(Options.Fd, -1);
   ASSERT_EQ(fdrLoggingInit(kBufferSize, kBufferMax, &Options,
-                            sizeof(FDRLoggingOptions)),
             XRayLogInitStatus::XRAY_LOG_INITIALIZED);
   for (uint64_t I = 0; I < 100; ++I) {
     fdrLoggingHandleArg0(1, XRayEntryType::ENTRY);
     fdrLoggingHandleArg0(1, XRayEntryType::EXIT);
   }
   ASSERT_EQ(fdrLoggingFinalize(), XRayLogInitStatus::XRAY_LOG_FINALIZED);
   ASSERT_EQ(fdrLoggingFlush(), XRayLogFlushStatus::XRAY_LOG_FLUSHED);
 
   // To do this properly, we have to close the file descriptor then re-open the
   // file for reading this time.
   ASSERT_EQ(close(Options.Fd), 0);
   int Fd = open(TmpFilename, O_RDONLY);
   ASSERT_NE(-1, Fd);
   ScopedFileCloserAndDeleter Guard(Fd, TmpFilename);
   auto Size = lseek(Fd, 0, SEEK_END);
   ASSERT_NE(Size, 0);
   // Map the file contents.
   const char *Contents = static_cast<const char *>(
       mmap(NULL, Size, PROT_READ, MAP_PRIVATE, Fd, 0));
   ASSERT_NE(Contents, nullptr);
 
   XRayFileHeader H;
   memcpy(&H, Contents, sizeof(XRayFileHeader));
-  ASSERT_EQ(H.Version, 1);
+  ASSERT_EQ(H.Version, 2);
+  ASSERT_EQ(H.Type, FileTypes::FDR_LOG);
+
+  MetadataRecord MDR0, MDR1;
+  memcpy(&MDR0, Contents + sizeof(XRayFileHeader), sizeof(MetadataRecord));
+  memcpy(&MDR1, Contents + sizeof(XRayFileHeader) + sizeof(MetadataRecord),
+         sizeof(MetadataRecord));
+  ASSERT_EQ(MDR0.RecordKind,
+            uint8_t(MetadataRecord::RecordKinds::BufferExtents));
+  ASSERT_EQ(MDR1.RecordKind, uint8_t(MetadataRecord::RecordKinds::NewBuffer));
+}
+
+TEST(FDRLoggingTest, MultiThreadedCycling) {
+  FDRLoggingOptions Options;
+  char TmpFilename[] = "fdr-logging-test.XXXXXX";
+  Options.Fd = mkstemp(TmpFilename);
+  ASSERT_NE(Options.Fd, -1);
+  ASSERT_EQ(fdrLoggingInit(kBufferSize, 1, &Options, sizeof(FDRLoggingOptions)),
+            XRayLogInitStatus::XRAY_LOG_INITIALIZED);
+
+  // Now we want to create one thread, do some logging, then create another one,
+  // in succession and making sure that we're able to get thread records from
+  // the latest thread (effectively being able to recycle buffers).
+  std::array<pid_t, 2> Threads;
+  for (uint64_t I = 0; I < 2; ++I) {
+    std::thread t{[I, &Threads] {
+      fdrLoggingHandleArg0(I + 1, XRayEntryType::ENTRY);
+      fdrLoggingHandleArg0(I + 1, XRayEntryType::EXIT);
+      Threads[I] = syscall(SYS_gettid);
+    }};
+    t.join();
+  }
+  ASSERT_EQ(fdrLoggingFinalize(), XRayLogInitStatus::XRAY_LOG_FINALIZED);
+  ASSERT_EQ(fdrLoggingFlush(), XRayLogFlushStatus::XRAY_LOG_FLUSHED);
+
+  // To do this properly, we have to close the file descriptor then re-open the
+  // file for reading this time.
+  ASSERT_EQ(close(Options.Fd), 0);
+  int Fd = open(TmpFilename, O_RDONLY);
+  ASSERT_NE(-1, Fd);
+  ScopedFileCloserAndDeleter Guard(Fd, TmpFilename);
+  auto Size = lseek(Fd, 0, SEEK_END);
+  ASSERT_NE(Size, 0);
+  // Map the file contents.
+  const char *Contents = static_cast<const char *>(
+      mmap(NULL, Size, PROT_READ, MAP_PRIVATE, Fd, 0));
+  ASSERT_NE(Contents, nullptr);
+
+  XRayFileHeader H;
+  memcpy(&H, Contents, sizeof(XRayFileHeader));
+  ASSERT_EQ(H.Version, 2);
   ASSERT_EQ(H.Type, FileTypes::FDR_LOG);
 
-  MetadataRecord MDR0;
+  MetadataRecord MDR0, MDR1;
   memcpy(&MDR0, Contents + sizeof(XRayFileHeader), sizeof(MetadataRecord));
-  ASSERT_EQ(MDR0.RecordKind, uint8_t(MetadataRecord::RecordKinds::NewBuffer));
+  memcpy(&MDR1, Contents + sizeof(XRayFileHeader) + sizeof(MetadataRecord),
+         sizeof(MetadataRecord));
+  ASSERT_EQ(MDR0.RecordKind,
+            uint8_t(MetadataRecord::RecordKinds::BufferExtents));
+  ASSERT_EQ(MDR1.RecordKind, uint8_t(MetadataRecord::RecordKinds::NewBuffer));
+  pid_t Latest = 0;
+  memcpy(&Latest, MDR1.Data, sizeof(pid_t));
+  ASSERT_EQ(Latest, Threads[1]);
 }
 
 } // namespace
 } // 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 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
 /// trace collection.
 class BufferQueue {
  public:
+  struct alignas(64) BufferExtents {
+    __sanitizer::atomic_uint64_t Size;
+  };
+
   struct Buffer {
     void *Buffer = nullptr;
     size_t Size = 0;
+    BufferExtents* Extents;
   };
 
  private:
   struct BufferRep {
     // The managed buffer.
     Buffer Buff;
 
     // This is true if the buffer has been returned to the available queue, and

compiler-rt/lib/xray/xray_buffer_queue.cc

 #include "sanitizer_common/sanitizer_allocator_internal.h"
 #include "sanitizer_common/sanitizer_common.h"
 #include "sanitizer_common/sanitizer_libc.h"
 
 using namespace __xray;
 using namespace __sanitizer;
 
 BufferQueue::BufferQueue(size_t B, size_t N, bool &Success)
-    : BufferSize(B),
-      Buffers(new BufferRep[N]()),
-      BufferCount(N),
-      Finalizing{0},
-      OwnedBuffers(new void *[N]()),
-      Next(Buffers),
-      First(Buffers),
+    : BufferSize(B), Buffers(new BufferRep[N]()), BufferCount(N), Finalizing{0},
+      OwnedBuffers(new void *[N]()), Next(Buffers), First(Buffers),
       LiveBuffers(0) {
   for (size_t i = 0; i < N; ++i) {
     auto &T = Buffers[i];
-    void *Tmp = InternalAlloc(BufferSize);
+    void *Tmp = InternalAlloc(BufferSize, nullptr, 64);
     if (Tmp == nullptr) {
       Success = false;
       return;
     }
+    void *Extents = InternalAlloc(sizeof(BufferExtents), nullptr, 64);
+    if (Extents == nullptr) {
+      Success = false;
+      return;
+    }
     auto &Buf = T.Buff;
     Buf.Buffer = Tmp;
     Buf.Size = B;
+    Buf.Extents = reinterpret_cast<BufferExtents *>(Extents);
     OwnedBuffers[i] = Tmp;
   }
   Success = true;
 }
 
 BufferQueue::ErrorCode BufferQueue::getBuffer(Buffer &Buf) {
   if (__sanitizer::atomic_load(&Finalizing, __sanitizer::memory_order_acquire))
     return ErrorCode::QueueFinalizing;
   __sanitizer::SpinMutexLock Guard(&Mutex);
-  if (LiveBuffers == BufferCount) return ErrorCode::NotEnoughMemory;
+  if (LiveBuffers == BufferCount)
+    return ErrorCode::NotEnoughMemory;
 
   auto &T = *Next;
   auto &B = T.Buff;
   Buf = B;
+  T.Used = true;
   ++LiveBuffers;
 
-  if (++Next == (Buffers + BufferCount)) Next = Buffers;
+  if (++Next == (Buffers + BufferCount))
+    Next = Buffers;
 
   return ErrorCode::Ok;
 }
 
 BufferQueue::ErrorCode BufferQueue::releaseBuffer(Buffer &Buf) {
   // Blitz through the buffers array to find the buffer.
   bool Found = false;
   for (auto I = OwnedBuffers, E = OwnedBuffers + BufferCount; I != E; ++I) {
     if (*I == Buf.Buffer) {
       Found = true;
       break;
     }
   }
-  if (!Found) return ErrorCode::UnrecognizedBuffer;
+  if (!Found)
+    return ErrorCode::UnrecognizedBuffer;
 
   __sanitizer::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;
+  if (LiveBuffers == 0)
+    return ErrorCode::NotEnoughMemory;
 
   // Now that the buffer has been released, we mark it as "used".
   First->Buff = Buf;
   First->Used = true;
   Buf.Buffer = nullptr;
   Buf.Size = 0;
   --LiveBuffers;
-  if (++First == (Buffers + BufferCount)) First = Buffers;
+  if (++First == (Buffers + BufferCount))
+    First = Buffers;
 
   return ErrorCode::Ok;
 }
 
 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 I = Buffers, E = Buffers + BufferCount; I != E; ++I) {
     auto &T = *I;
     auto &Buf = T.Buff;
     InternalFree(Buf.Buffer);
+    InternalFree(Buf.Extents);
   }
   delete[] Buffers;
   delete[] OwnedBuffers;
 }

compiler-rt/lib/xray/xray_fdr_log_records.h

struct alignas(16) MetadataRecord {
   enum class RecordKinds : uint8_t {
     NewBuffer,
     EndOfBuffer,
     NewCPUId,
     TSCWrap,
     WalltimeMarker,
     CustomEventMarker,
     CallArgument,
+    BufferExtents,
   };
+
   // Use 7 bits to identify this record type.
   /* RecordKinds */ uint8_t RecordKind : 7;
   char Data[15];
 } __attribute__((packed));
 
 static_assert(sizeof(MetadataRecord) == 16, "Wrong size for MetadataRecord.");
 
 struct alignas(8) FunctionRecord {

compiler-rt/lib/xray/xray_fdr_logging.cc

 // This file is a part of XRay, a dynamic runtime instrumentation system.
 //
 // Here we implement the Flight Data Recorder mode for XRay, where we use
 // compact structures to store records in memory as well as when writing out the
 // data to files.
 //
 //===----------------------------------------------------------------------===//
 #include "xray_fdr_logging.h"
+#include <errno.h>
 #include <sys/syscall.h>
 #include <sys/time.h>
-#include <errno.h>
 #include <time.h>
 #include <unistd.h>
 
 #include "sanitizer_common/sanitizer_atomic.h"
 #include "sanitizer_common/sanitizer_common.h"
 #include "xray/xray_interface.h"
 #include "xray/xray_records.h"
 #include "xray_buffer_queue.h"
 #include "xray_defs.h"
 #include "xray_fdr_logging_impl.h"
 #include "xray_flags.h"
 #include "xray_tsc.h"
 #include "xray_utils.h"
 
 namespace __xray {
 
 // Global BufferQueue.
-// NOTE: This is a pointer to avoid having to do atomic operations at
-// initialization time. This is OK to leak as there will only be one bufferqueue
-// for the runtime, initialized once through the fdrInit(...) sequence.
-std::shared_ptr<BufferQueue> *BQ = nullptr;
+BufferQueue *BQ = nullptr;
 
 __sanitizer::atomic_sint32_t LogFlushStatus = {
     XRayLogFlushStatus::XRAY_LOG_NOT_FLUSHING};
 
 FDRLoggingOptions FDROptions;
 
 __sanitizer::SpinMutex FDROptionsMutex;
 
 // Must finalize before flushing.
 XRayLogFlushStatus fdrLoggingFlush() XRAY_NEVER_INSTRUMENT {
   if (__sanitizer::atomic_load(&LoggingStatus,
                                __sanitizer::memory_order_acquire) !=
-      XRayLogInitStatus::XRAY_LOG_FINALIZED)
+      XRayLogInitStatus::XRAY_LOG_FINALIZED) {
+    if (__sanitizer::Verbosity())
+      Report("Not flushing log, implementation is not finalized.\n");
     return XRayLogFlushStatus::XRAY_LOG_NOT_FLUSHING;
+  }
 
   s32 Result = XRayLogFlushStatus::XRAY_LOG_NOT_FLUSHING;
   if (!__sanitizer::atomic_compare_exchange_strong(
           &LogFlushStatus, &Result, XRayLogFlushStatus::XRAY_LOG_FLUSHING,
-          __sanitizer::memory_order_release))
+          __sanitizer::memory_order_release)) {
+
+    if (__sanitizer::Verbosity())
+      Report("Not flushing log, implementation is still finalizing.\n");
     return static_cast<XRayLogFlushStatus>(Result);
+  }
+
+  if (BQ == nullptr) {
+    if (__sanitizer::Verbosity())
+      Report("Cannot flush when global buffer queue is null.\n");
+    return XRayLogFlushStatus::XRAY_LOG_NOT_FLUSHING;
+  }
 
-  // Make a copy of the BufferQueue pointer to prevent other threads that may be
-  // resetting it from blowing away the queue prematurely while we're dealing
-  // with it.
-  auto LocalBQ = *BQ;
+  // We wait a number of microseconds to allow threads to see that we've
+  // finalised before attempting to flush the log.
+  struct timespec TS;
+  TS.tv_sec = flags()->xray_fdr_log_grace_period_us / 1000000;
+  TS.tv_nsec = (flags()->xray_fdr_log_grace_period_us % 1000000) * 1000;
+  struct timespec Rem;
+  while (clock_nanosleep(CLOCK_REALTIME, 0, &TS, &Rem) &&
+         (Rem.tv_sec != 0 || Rem.tv_nsec != 0))
+    TS = Rem;
 
   // We write out the file in the following format:
   //
   //   1) We write down the XRay file header with version 1, type FDR_LOG.
   //   2) Then we use the 'apply' member of the BufferQueue that's live, to
   //      ensure that at this point in time we write down the buffers that have
   //      been released (and marked "used") -- we dump the full buffer for now
   //      (fixed-sized) and let the tools reading the buffers deal with the data
XRayLogFlushStatus fdrLoggingFlush() XRAY_NEVER_INSTRUMENT {
   }
 
   // Test for required CPU features and cache the cycle frequency
   static bool TSCSupported = probeRequiredCPUFeatures();
   static uint64_t CycleFrequency =
       TSCSupported ? getTSCFrequency() : __xray::NanosecondsPerSecond;
 
   XRayFileHeader Header;
-  Header.Version = 1;
+
+  // Version 2 of the log writes the extents of the buffer, instead of relying
+  // on an end-of-buffer record.
+  Header.Version = 2;
   Header.Type = FileTypes::FDR_LOG;
   Header.CycleFrequency = CycleFrequency;
+
   // FIXME: Actually check whether we have 'constant_tsc' and 'nonstop_tsc'
   // before setting the values in the header.
   Header.ConstantTSC = 1;
   Header.NonstopTSC = 1;
-  Header.FdrData = FdrAdditionalHeaderData{LocalBQ->ConfiguredBufferSize()};
+  Header.FdrData = FdrAdditionalHeaderData{BQ->ConfiguredBufferSize()};
   retryingWriteAll(Fd, reinterpret_cast<char *>(&Header),
                    reinterpret_cast<char *>(&Header) + sizeof(Header));
 
-  LocalBQ->apply([&](const BufferQueue::Buffer &B) {
-    uint64_t BufferSize = B.Size;
-    if (BufferSize > 0) {
+  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 = __sanitizer::atomic_load(
+        &B.Extents->Size, __sanitizer::memory_order_acquire);
+    assert(BufferExtents <= B.Size);
+    ExtentsRecord.Type = uint8_t(RecordType::Metadata);
+    ExtentsRecord.RecordKind =
+        uint8_t(MetadataRecord::RecordKinds::BufferExtents);
+    std::memcpy(ExtentsRecord.Data, &BufferExtents, sizeof(BufferExtents));
+    if (BufferExtents > 0) {
+      retryingWriteAll(Fd, reinterpret_cast<char *>(&ExtentsRecord),
+                       reinterpret_cast<char *>(&ExtentsRecord) +
+                           sizeof(MetadataRecord));
       retryingWriteAll(Fd, reinterpret_cast<char *>(B.Buffer),
-                       reinterpret_cast<char *>(B.Buffer) + B.Size);
+                       reinterpret_cast<char *>(B.Buffer) + BufferExtents);
     }
   });
 
-  // The buffer for this particular thread would have been finalised after
-  // we've written everything to disk, and we'd lose the thread's trace.
-  auto &TLD = __xray::__xray_fdr_internal::getThreadLocalData();
-  if (TLD.Buffer.Buffer != nullptr) {
-    __xray::__xray_fdr_internal::writeEOBMetadata();
-    auto Start = reinterpret_cast<char *>(TLD.Buffer.Buffer);
-    retryingWriteAll(Fd, Start, Start + TLD.Buffer.Size);
-  }
-
   __sanitizer::atomic_store(&LogFlushStatus,
                             XRayLogFlushStatus::XRAY_LOG_FLUSHED,
                             __sanitizer::memory_order_release);
   return XRayLogFlushStatus::XRAY_LOG_FLUSHED;
 }
 
 XRayLogInitStatus fdrLoggingFinalize() XRAY_NEVER_INSTRUMENT {
   s32 CurrentStatus = XRayLogInitStatus::XRAY_LOG_INITIALIZED;
   if (!__sanitizer::atomic_compare_exchange_strong(
           &LoggingStatus, &CurrentStatus,
           XRayLogInitStatus::XRAY_LOG_FINALIZING,
-          __sanitizer::memory_order_release))
+          __sanitizer::memory_order_release)) {
+    if (__sanitizer::Verbosity())
+      Report("Cannot finalize log, implementation not initialized.\n");
     return static_cast<XRayLogInitStatus>(CurrentStatus);
+  }
 
   // Do special things to make the log finalize itself, and not allow any more
   // operations to be performed until re-initialized.
-  (*BQ)->finalize();
+  BQ->finalize();
 
   __sanitizer::atomic_store(&LoggingStatus,
                             XRayLogInitStatus::XRAY_LOG_FINALIZED,
                             __sanitizer::memory_order_release);
   return XRayLogInitStatus::XRAY_LOG_FINALIZED;
 }
 
 XRayLogInitStatus fdrLoggingReset() XRAY_NEVER_INSTRUMENT {
   s32 CurrentStatus = XRayLogInitStatus::XRAY_LOG_FINALIZED;
   if (__sanitizer::atomic_compare_exchange_strong(
           &LoggingStatus, &CurrentStatus,
           XRayLogInitStatus::XRAY_LOG_INITIALIZED,
           __sanitizer::memory_order_release))
     return static_cast<XRayLogInitStatus>(CurrentStatus);
 
   // Release the in-memory buffer queue.
-  BQ->reset();
+  delete BQ;
+  BQ = nullptr;
 
   // Spin until the flushing status is flushed.
   s32 CurrentFlushingStatus = XRayLogFlushStatus::XRAY_LOG_FLUSHED;
   while (__sanitizer::atomic_compare_exchange_weak(
       &LogFlushStatus, &CurrentFlushingStatus,
       XRayLogFlushStatus::XRAY_LOG_NOT_FLUSHING,
       __sanitizer::memory_order_release)) {
     if (CurrentFlushingStatus == XRayLogFlushStatus::XRAY_LOG_NOT_FLUSHING)
       break;
     CurrentFlushingStatus = XRayLogFlushStatus::XRAY_LOG_FLUSHED;
   }
 
   // At this point, we know that the status is flushed, and that we can assume
   return XRayLogInitStatus::XRAY_LOG_UNINITIALIZED;
 }
 
 struct TSCAndCPU {
-  uint64_t TSC;
-  unsigned char CPU;
+  uint64_t TSC = 0;
+  unsigned char CPU = 0;
 };
 
 static TSCAndCPU getTimestamp() XRAY_NEVER_INSTRUMENT {
   // We want to get the TSC as early as possible, so that we can check whether
   // we've seen this CPU before. We also do it before we load anything else, to
   // allow for forward progress with the scheduling.
   TSCAndCPU Result;
 
if (TSCSupported) {
     Result.TSC = TS.tv_sec * __xray::NanosecondsPerSecond + TS.tv_nsec;
   }
   return Result;
 }
 
 void fdrLoggingHandleArg0(int32_t FuncId,
                           XRayEntryType Entry) XRAY_NEVER_INSTRUMENT {
   auto TC = getTimestamp();
-  __xray_fdr_internal::processFunctionHook(FuncId, Entry, TC.TSC,
-                                           TC.CPU, 0, clock_gettime, *BQ);
+  __xray_fdr_internal::processFunctionHook(FuncId, Entry, TC.TSC, TC.CPU, 0,
+                                           clock_gettime, BQ);
 }
 
 void fdrLoggingHandleArg1(int32_t FuncId, XRayEntryType Entry,
                           uint64_t Arg) XRAY_NEVER_INSTRUMENT {
   auto TC = getTimestamp();
-  __xray_fdr_internal::processFunctionHook(
-      FuncId, Entry, TC.TSC, TC.CPU, Arg, clock_gettime, *BQ);
+  __xray_fdr_internal::processFunctionHook(FuncId, Entry, TC.TSC, TC.CPU, Arg,
+                                           clock_gettime, BQ);
 }
 
 void fdrLoggingHandleCustomEvent(void *Event,
                                  std::size_t EventSize) XRAY_NEVER_INSTRUMENT {
   using namespace __xray_fdr_internal;
   auto TC = getTimestamp();
   auto &TSC = TC.TSC;
   auto &CPU = TC.CPU;
   RecursionGuard Guard{Running};
-  if (!Guard) {
-    assert(Running && "RecursionGuard is buggy!");
+  if (!Guard)
     return;
-  }
   if (EventSize > std::numeric_limits<int32_t>::max()) {
     using Empty = struct {};
     static Empty Once = [&] {
       Report("Event size too large = %zu ; > max = %d\n", EventSize,
              std::numeric_limits<int32_t>::max());
       return Empty();
     }();
     (void)Once;
   }
   int32_t ReducedEventSize = static_cast<int32_t>(EventSize);
   auto &TLD = getThreadLocalData();
-  if (!isLogInitializedAndReady(TLD.LocalBQ, TSC, CPU, clock_gettime))
+  if (!isLogInitializedAndReady(TLD.BQ, TSC, CPU, clock_gettime))
     return;
 
   // Here we need to prepare the log to handle:
   //   - The metadata record we're going to write. (16 bytes)
   //   - The additional data we're going to write. Currently, that's the size of
   //   the event we're going to dump into the log as free-form bytes.
   if (!prepareBuffer(TSC, CPU, clock_gettime, MetadataRecSize + EventSize)) {
-    TLD.LocalBQ = nullptr;
+    TLD.BQ = nullptr;
     return;
   }
 
   // Write the custom event metadata record, which consists of the following
   // information:
   //   - 8 bytes (64-bits) for the full TSC when the event started.
   //   - 4 bytes (32-bits) for the length of the data.
   MetadataRecord CustomEvent;
   CustomEvent.Type = uint8_t(RecordType::Metadata);
   CustomEvent.RecordKind =
       uint8_t(MetadataRecord::RecordKinds::CustomEventMarker);
   constexpr auto TSCSize = sizeof(TC.TSC);
   std::memcpy(&CustomEvent.Data, &ReducedEventSize, sizeof(int32_t));
   std::memcpy(&CustomEvent.Data[sizeof(int32_t)], &TSC, TSCSize);
   std::memcpy(TLD.RecordPtr, &CustomEvent, sizeof(CustomEvent));
   TLD.RecordPtr += sizeof(CustomEvent);
   std::memcpy(TLD.RecordPtr, Event, ReducedEventSize);
+  incrementExtents(MetadataRecSize + EventSize);
   endBufferIfFull();
 }
 
 XRayLogInitStatus fdrLoggingInit(std::size_t BufferSize, std::size_t BufferMax,
                                  void *Options,
                                  size_t OptionsSize) XRAY_NEVER_INSTRUMENT {
-  if (OptionsSize != sizeof(FDRLoggingOptions))
+  if (OptionsSize != sizeof(FDRLoggingOptions)) {
+    if (__sanitizer::Verbosity())
+      Report("Cannot initialize FDR logging; wrong size for options: %d\n",
+             OptionsSize);
     return static_cast<XRayLogInitStatus>(__sanitizer::atomic_load(
         &LoggingStatus, __sanitizer::memory_order_acquire));
+  }
   s32 CurrentStatus = XRayLogInitStatus::XRAY_LOG_UNINITIALIZED;
   if (!__sanitizer::atomic_compare_exchange_strong(
           &LoggingStatus, &CurrentStatus,
           XRayLogInitStatus::XRAY_LOG_INITIALIZING,
-          __sanitizer::memory_order_release))
+          __sanitizer::memory_order_release)) {
+    if (__sanitizer::Verbosity())
+      Report("Cannot initialize already initialized implementation.\n");
     return static_cast<XRayLogInitStatus>(CurrentStatus);
+  }
 
   {
     __sanitizer::SpinMutexLock Guard(&FDROptionsMutex);
     memcpy(&FDROptions, Options, OptionsSize);
   }
 
   bool Success = false;
+
+  if (BQ != nullptr) {
+    delete BQ;
+    BQ = nullptr;
+  }
+
   if (BQ == nullptr)
-    BQ = new std::shared_ptr<BufferQueue>();
+    BQ = new BufferQueue(BufferSize, BufferMax, Success);
 
-  *BQ = std::make_shared<BufferQueue>(BufferSize, BufferMax, Success);
   if (!Success) {
     Report("BufferQueue init failed.\n");
+    if (BQ != nullptr) {
+      delete BQ;
+      BQ = nullptr;
+    }
     return XRayLogInitStatus::XRAY_LOG_UNINITIALIZED;
   }
 
+  static bool UNUSED Once = [] {
+    pthread_key_create(&__xray_fdr_internal::Key, +[](void *) {
+      auto &TLD = __xray_fdr_internal::getThreadLocalData();
+      if (TLD.BQ != nullptr) {
+        auto EC = TLD.BQ->releaseBuffer(TLD.Buffer);
+        if (EC != BufferQueue::ErrorCode::Ok)
+          Report("At thread exit, failed to release buffer at %p; error=%s\n",
+                 TLD.Buffer.Buffer, BufferQueue::getErrorString(EC));
+      }
+    });
+    return false;
+  }();
+
   // Arg1 handler should go in first to avoid concurrent code accidentally
   // falling back to arg0 when it should have ran arg1.
   __xray_set_handler_arg1(fdrLoggingHandleArg1);
   // Install the actual handleArg0 handler after initialising the buffers.
   __xray_set_handler(fdrLoggingHandleArg0);
   __xray_set_customevent_handler(fdrLoggingHandleCustomEvent);
 
   __sanitizer::atomic_store(&LoggingStatus,
                             XRayLogInitStatus::XRAY_LOG_INITIALIZED,
                             __sanitizer::memory_order_release);
+
+  if (__sanitizer::Verbosity())
     Report("XRay FDR init successful.\n");
   return XRayLogInitStatus::XRAY_LOG_INITIALIZED;
 }
 
 } // namespace __xray
 
 static auto UNUSED Unused = [] {
   using namespace __xray;
   if (flags()->xray_fdr_log) {

compiler-rt/lib/xray/xray_fdr_logging_impl.h

 
 #include <cassert>
 #include <cstddef>
 #include <cstring>
 #include <limits>
 #include <pthread.h>
 #include <sys/syscall.h>
 #include <time.h>
+#include <type_traits>
 #include <unistd.h>
 
-// FIXME: Implement analogues to std::shared_ptr and std::weak_ptr
-#include <memory>
-
 #include "sanitizer_common/sanitizer_common.h"
 #include "xray/xray_log_interface.h"
 #include "xray_buffer_queue.h"
 #include "xray_defs.h"
 #include "xray_fdr_log_records.h"
 #include "xray_flags.h"
 #include "xray_tsc.h"
 

 /// design, callers that wish to test events occuring on different threads
 /// will actually have to run them on different threads.
 ///
 /// This also means that it is possible to break invariants maintained by
 /// cooperation with xray_fdr_logging class, so be careful and think twice.
 namespace __xray_fdr_internal {
 
 /// Writes the new buffer record and wallclock time that begin a buffer for a
 /// thread to MemPtr and increments MemPtr. Bypasses the thread local state
 /// machine and writes directly to memory without checks.
-static void writeNewBufferPreamble(pid_t Tid, timespec TS, char *&MemPtr);
-
-/// Write a metadata record to switch to a new CPU to MemPtr and increments
-/// MemPtr. Bypasses the thread local state machine and writes directly to
-/// memory without checks.
-static void writeNewCPUIdMetadata(uint16_t CPU, uint64_t TSC, char *&MemPtr);
-
-/// Writes an EOB metadata record to MemPtr and increments MemPtr. Bypasses the
-/// thread local state machine and writes directly to memory without checks.
-static void writeEOBMetadata(char *&MemPtr);
-
-/// Writes a TSC Wrap metadata record to MemPtr and increments MemPtr. Bypasses
-/// the thread local state machine and directly writes to memory without checks.
-static void writeTSCWrapMetadata(uint64_t TSC, char *&MemPtr);
+static void writeNewBufferPreamble(pid_t Tid, timespec TS);
 
 /// Writes a Function Record to MemPtr and increments MemPtr. Bypasses the
 /// thread local state machine and writes the function record directly to
 /// memory.
 static void writeFunctionRecord(int FuncId, uint32_t TSCDelta,
-                                XRayEntryType EntryType, char *&MemPtr);
+                                XRayEntryType EntryType);
 
 /// Sets up a new buffer in thread_local storage and writes a preamble. The
 /// wall_clock_reader function is used to populate the WallTimeRecord entry.
 static void setupNewBuffer(int (*wall_clock_reader)(clockid_t,
                                                     struct timespec *));
 
-/// Called to record CPU time for a new CPU within the current thread.
-static void writeNewCPUIdMetadata(uint16_t CPU, uint64_t TSC);
-
-/// Called to close the buffer when the thread exhausts the buffer or when the
-/// thread exits (via a thread local variable destructor).
-static void writeEOBMetadata();
-
 /// TSC Wrap records are written when a TSC delta encoding scheme overflows.
 static void writeTSCWrapMetadata(uint64_t TSC);
 
 // Group together thread-local-data in a struct, then hide it behind a function
 // call so that it can be initialized on first use instead of as a global. We
 // force the alignment to 64-bytes for x86 cache line alignment, as this
 // structure is used in the hot path of implementation.
 struct alignas(64) ThreadLocalData {
struct alignas(64) ThreadLocalData {
   // We assume that we'll support only 65536 CPUs for x86_64.
   uint16_t CurrentCPU = std::numeric_limits<uint16_t>::max();
   uint64_t LastTSC = 0;
   uint64_t LastFunctionEntryTSC = 0;
 
   // Make sure a thread that's ever called handleArg0 has a thread-local
   // live reference to the buffer queue for this particular instance of
   // FDRLogging, and that we're going to clean it up when the thread exits.
-  std::shared_ptr<BufferQueue> LocalBQ = nullptr;
+  BufferQueue *BQ = nullptr;
 };
 
-// Forward-declare, defined later.
-static ThreadLocalData &getThreadLocalData();
+static_assert(std::is_trivially_destructible<ThreadLocalData>::value,
+              "ThreadLocalData must be trivially destructible");
 
 static constexpr auto MetadataRecSize = sizeof(MetadataRecord);
 static constexpr auto FunctionRecSize = sizeof(FunctionRecord);
 
+// Use a global pthread key to identify thread-local data for logging.
+static pthread_key_t Key;
+
 // This function will initialize the thread-local data structure used by the FDR
 // logging implementation and return a reference to it. The implementation
 // details require a bit of care to maintain.
 //
 // First, some requirements on the implementation in general:
 //
 //   - XRay handlers should not call any memory allocation routines that may
 //     delegate to an instrumented implementation. This means functions like
 //     malloc() and free() should not be called while instrumenting.
 //
 //   - We would like to use some thread-local data initialized on first-use of
 //     the XRay instrumentation. These allow us to implement unsynchronized
 //     routines that access resources associated with the thread.
 //
 // The implementation here uses a few mechanisms that allow us to provide both
 // the requirements listed above. We do this by:
 //
 //   1. Using a thread-local aligned storage buffer for representing the
 //      ThreadLocalData struct. This data will be uninitialized memory by
 //      design.
 //
-//   2. Using pthread_once(...) to initialize the thread-local data structures
-//      on first use, for every thread. We don't use std::call_once so we don't
-//      have a reliance on the C++ runtime library.
-//
-//   3. Registering a cleanup function that gets run at the end of a thread's
-//      lifetime through pthread_create_key(...). The cleanup function would
-//      allow us to release the thread-local resources in a manner that would
-//      let the rest of the XRay runtime implementation handle the records
-//      written for this thread's active buffer.
+//   2. Not requiring a thread exit handler/implementation, keeping the
+//      thread-local as purely a collection of references/data that do not
+//      require cleanup.
 //
 // We're doing this to avoid using a `thread_local` object that has a
 // non-trivial destructor, because the C++ runtime might call std::malloc(...)
 // to register calls to destructors. Deadlocks may arise when, for example, an
 // externally provided malloc implementation is XRay instrumented, and
 // initializing the thread-locals involves calling into malloc. A malloc
 // implementation that does global synchronization might be holding a lock for a
 // critical section, calling a function that might be XRay instrumented (and
 // thus in turn calling into malloc by virtue of registration of the
 // thread_local's destructor).
-//
-// With the approach taken where, we attempt to avoid the potential for
-// deadlocks by relying instead on pthread's memory management routines.
 static ThreadLocalData &getThreadLocalData() {
-  thread_local pthread_key_t key;
-
-  // We need aligned, uninitialized storage for the TLS object which is
-  // trivially destructible. We're going to use this as raw storage and
-  // placement-new the ThreadLocalData object into it later.
-  alignas(alignof(ThreadLocalData)) thread_local unsigned char
-      TLSBuffer[sizeof(ThreadLocalData)];
-
-  // Ensure that we only actually ever do the pthread initialization once.
-  thread_local bool UNUSED Unused = [] {
-    new (&TLSBuffer) ThreadLocalData();
-    auto result = pthread_key_create(&key, +[](void *) {
-      auto &TLD = *reinterpret_cast<ThreadLocalData *>(&TLSBuffer);
-      auto &RecordPtr = TLD.RecordPtr;
-      auto &Buffers = TLD.LocalBQ;
-      auto &Buffer = TLD.Buffer;
-      if (RecordPtr == nullptr)
-        return;
-
-      // We make sure that upon exit, a thread will write out the EOB
-      // MetadataRecord in the thread-local log, and also release the buffer
-      // to the queue.
-      assert((RecordPtr + MetadataRecSize) -
-                 static_cast<char *>(Buffer.Buffer) >=
-             static_cast<ptrdiff_t>(MetadataRecSize));
-      if (Buffers) {
-        writeEOBMetadata();
-        auto EC = Buffers->releaseBuffer(Buffer);
-        if (EC != BufferQueue::ErrorCode::Ok)
-          Report("Failed to release buffer at %p; error=%s\n", Buffer.Buffer,
-                 BufferQueue::getErrorString(EC));
-        Buffers = nullptr;
-        return;
-      }
-    });
-    if (result != 0) {
-      Report("Failed to allocate thread-local data through pthread; error=%d",
-             result);
-      return false;
-    }
-    pthread_setspecific(key, &TLSBuffer);
-    return true;
+  static_assert(alignof(ThreadLocalData) >= 64,
+                "ThreadLocalData must be cache line aligned.");
+  thread_local ThreadLocalData TLD;
+  thread_local bool UNUSED ThreadOnce = [] {
+    pthread_setspecific(Key, &TLD);
   }();
-
-  return *reinterpret_cast<ThreadLocalData *>(TLSBuffer);
+  return TLD;
 }
 
 //-----------------------------------------------------------------------------|
 // The rest of the file is implementation.                                     |
 //-----------------------------------------------------------------------------|
 // Functions are implemented in the header for inlining since we don't want    |
 // to grow the stack when we've hijacked the binary for logging.               |
 //-----------------------------------------------------------------------------|
public:
   ~RecursionGuard() noexcept {
     if (Valid)
       Running = false;
   }
 };
 
 } // namespace
 
-inline void writeNewBufferPreamble(pid_t Tid, timespec TS,
-                                   char *&MemPtr) XRAY_NEVER_INSTRUMENT {
+static void writeNewBufferPreamble(pid_t Tid,
+                                   timespec TS) XRAY_NEVER_INSTRUMENT {
   static constexpr int InitRecordsCount = 2;
-  alignas(alignof(MetadataRecord)) unsigned char
-      Records[InitRecordsCount * MetadataRecSize];
+  auto &TLD = getThreadLocalData();
+  MetadataRecord Records[InitRecordsCount];
   {
     // Write out a MetadataRecord to signify that this is the start of a new
     // buffer, associated with a particular thread, with a new CPU.  For the
     // data, we have 15 bytes to squeeze as much information as we can.  At this
     // point we only write down the following bytes:
     //   - Thread ID (pid_t, 4 bytes)
-    auto &NewBuffer = *reinterpret_cast<MetadataRecord *>(Records);
+    auto &NewBuffer = Records[0];
     NewBuffer.Type = uint8_t(RecordType::Metadata);
     NewBuffer.RecordKind = uint8_t(MetadataRecord::RecordKinds::NewBuffer);
     std::memcpy(&NewBuffer.Data, &Tid, sizeof(pid_t));
   }
+
   // Also write the WalltimeMarker record.
   {
     static_assert(sizeof(time_t) <= 8, "time_t needs to be at most 8 bytes");
-    auto &WalltimeMarker =
-        *reinterpret_cast<MetadataRecord *>(Records + MetadataRecSize);
+    auto &WalltimeMarker = Records[1];
     WalltimeMarker.Type = uint8_t(RecordType::Metadata);
     WalltimeMarker.RecordKind =
         uint8_t(MetadataRecord::RecordKinds::WalltimeMarker);
 
     // We only really need microsecond precision here, and enforce across
     // platforms that we need 64-bit seconds and 32-bit microseconds encoded in
     // the Metadata record.
     int32_t Micros = TS.tv_nsec / 1000;
     int64_t Seconds = TS.tv_sec;
     std::memcpy(WalltimeMarker.Data, &Seconds, sizeof(Seconds));
     std::memcpy(WalltimeMarker.Data + sizeof(Seconds), &Micros, sizeof(Micros));
   }
-  std::memcpy(MemPtr, Records, sizeof(MetadataRecord) * InitRecordsCount);
-  MemPtr += sizeof(MetadataRecord) * InitRecordsCount;
-  auto &TLD = getThreadLocalData();
+
   TLD.NumConsecutiveFnEnters = 0;
   TLD.NumTailCalls = 0;
+  if (TLD.BQ && !TLD.BQ->finalizing()) {
+    std::memcpy(TLD.RecordPtr, Records,
+                sizeof(MetadataRecord) * InitRecordsCount);
+    TLD.RecordPtr += sizeof(MetadataRecord) * InitRecordsCount;
+    // Since we write out the extents as the first metadata record of the
+    // buffer, we need to write out the extents including the extents record.
+    __sanitizer::atomic_store(&TLD.Buffer.Extents->Size,
+                              sizeof(MetadataRecord) * InitRecordsCount,
+                              __sanitizer::memory_order_release);
+  }
 }
 
 inline void setupNewBuffer(int (*wall_clock_reader)(
     clockid_t, struct timespec *)) XRAY_NEVER_INSTRUMENT {
   auto &TLD = getThreadLocalData();
-  auto &Buffer = TLD.Buffer;
-  auto &RecordPtr = TLD.RecordPtr;
-  RecordPtr = static_cast<char *>(Buffer.Buffer);
+  auto &B = TLD.Buffer;
+  TLD.RecordPtr = static_cast<char *>(B.Buffer);
   pid_t Tid = syscall(SYS_gettid);
   timespec TS{0, 0};
   // This is typically clock_gettime, but callers have injection ability.
   wall_clock_reader(CLOCK_MONOTONIC, &TS);
-  writeNewBufferPreamble(Tid, TS, RecordPtr);
+  writeNewBufferPreamble(Tid, TS);
   TLD.NumConsecutiveFnEnters = 0;
   TLD.NumTailCalls = 0;
 }
 
-inline void writeNewCPUIdMetadata(uint16_t CPU, uint64_t TSC,
-                                  char *&MemPtr) XRAY_NEVER_INSTRUMENT {
+static void incrementExtents(size_t Add) {
+  auto &TLD = getThreadLocalData();
+  __sanitizer::atomic_fetch_add(&TLD.Buffer.Extents->Size, Add,
+                                __sanitizer::memory_order_acq_rel);
+}
+
+static void decrementExtents(size_t Subtract) {
+  auto &TLD = getThreadLocalData();
+  __sanitizer::atomic_fetch_sub(&TLD.Buffer.Extents->Size, Subtract,
+                                __sanitizer::memory_order_acq_rel);
+}
+
+inline void writeNewCPUIdMetadata(uint16_t CPU,
+                                  uint64_t TSC) XRAY_NEVER_INSTRUMENT {
   auto &TLD = getThreadLocalData();
   MetadataRecord NewCPUId;
   NewCPUId.Type = uint8_t(RecordType::Metadata);
   NewCPUId.RecordKind = uint8_t(MetadataRecord::RecordKinds::NewCPUId);
 
   // The data for the New CPU will contain the following bytes:
   //   - CPU ID (uint16_t, 2 bytes)
   //   - Full TSC (uint64_t, 8 bytes)
   // Total = 10 bytes.
   std::memcpy(&NewCPUId.Data, &CPU, sizeof(CPU));
   std::memcpy(&NewCPUId.Data[sizeof(CPU)], &TSC, sizeof(TSC));
-  std::memcpy(MemPtr, &NewCPUId, sizeof(MetadataRecord));
-  MemPtr += sizeof(MetadataRecord);
-  TLD.NumConsecutiveFnEnters = 0;
-  TLD.NumTailCalls = 0;
-}
-
-inline void writeNewCPUIdMetadata(uint16_t CPU,
-                                  uint64_t TSC) XRAY_NEVER_INSTRUMENT {
-  writeNewCPUIdMetadata(CPU, TSC, getThreadLocalData().RecordPtr);
-}
-
-inline void writeEOBMetadata(char *&MemPtr) XRAY_NEVER_INSTRUMENT {
-  auto &TLD = getThreadLocalData();
-  MetadataRecord EOBMeta;
-  EOBMeta.Type = uint8_t(RecordType::Metadata);
-  EOBMeta.RecordKind = uint8_t(MetadataRecord::RecordKinds::EndOfBuffer);
-  // For now we don't write any bytes into the Data field.
-  std::memcpy(MemPtr, &EOBMeta, sizeof(MetadataRecord));
-  MemPtr += sizeof(MetadataRecord);
+  std::memcpy(TLD.RecordPtr, &NewCPUId, sizeof(MetadataRecord));
+  TLD.RecordPtr += sizeof(MetadataRecord);
   TLD.NumConsecutiveFnEnters = 0;
   TLD.NumTailCalls = 0;
+  incrementExtents(sizeof(MetadataRecord));
 }
 
-inline void writeEOBMetadata() XRAY_NEVER_INSTRUMENT {
-  writeEOBMetadata(getThreadLocalData().RecordPtr);
-}
-
-inline void writeTSCWrapMetadata(uint64_t TSC,
-                                 char *&MemPtr) XRAY_NEVER_INSTRUMENT {
+inline void writeTSCWrapMetadata(uint64_t TSC) XRAY_NEVER_INSTRUMENT {
   auto &TLD = getThreadLocalData();
   MetadataRecord TSCWrap;
   TSCWrap.Type = uint8_t(RecordType::Metadata);
   TSCWrap.RecordKind = uint8_t(MetadataRecord::RecordKinds::TSCWrap);
 
   // The data for the TSCWrap record contains the following bytes:
   //   - Full TSC (uint64_t, 8 bytes)
   // Total = 8 bytes.
   std::memcpy(&TSCWrap.Data, &TSC, sizeof(TSC));
-  std::memcpy(MemPtr, &TSCWrap, sizeof(MetadataRecord));
-  MemPtr += sizeof(MetadataRecord);
+  std::memcpy(TLD.RecordPtr, &TSCWrap, sizeof(MetadataRecord));
+  TLD.RecordPtr += sizeof(MetadataRecord);
   TLD.NumConsecutiveFnEnters = 0;
   TLD.NumTailCalls = 0;
-}
-
-inline void writeTSCWrapMetadata(uint64_t TSC) XRAY_NEVER_INSTRUMENT {
-  writeTSCWrapMetadata(TSC, getThreadLocalData().RecordPtr);
+  incrementExtents(sizeof(MetadataRecord));
 }
 
 // Call Argument metadata records store the arguments to a function in the
 // order of their appearance; holes are not supported by the buffer format.
 static inline void writeCallArgumentMetadata(uint64_t A) XRAY_NEVER_INSTRUMENT {
   auto &TLD = getThreadLocalData();
   MetadataRecord CallArg;
   CallArg.Type = uint8_t(RecordType::Metadata);
   CallArg.RecordKind = uint8_t(MetadataRecord::RecordKinds::CallArgument);
 
   std::memcpy(CallArg.Data, &A, sizeof(A));
   std::memcpy(TLD.RecordPtr, &CallArg, sizeof(MetadataRecord));
   TLD.RecordPtr += sizeof(MetadataRecord);
+  incrementExtents(sizeof(MetadataRecord));
 }
 
-static inline void writeFunctionRecord(int FuncId, uint32_t TSCDelta,
-                                       XRayEntryType EntryType,
-                                       char *&MemPtr) XRAY_NEVER_INSTRUMENT {
+static inline void
+writeFunctionRecord(int FuncId, uint32_t TSCDelta,
+                    XRayEntryType EntryType) XRAY_NEVER_INSTRUMENT {
   FunctionRecord FuncRecord;
   FuncRecord.Type = uint8_t(RecordType::Function);
   // Only take 28 bits of the function id.
   FuncRecord.FuncId = FuncId & ~(0x0F << 28);
   FuncRecord.TSCDelta = TSCDelta;
 
   auto &TLD = getThreadLocalData();
   switch (EntryType) {
   case XRayEntryType::ENTRY:
static bool Once = [&] {
              FuncId);
       return true;
     }();
     (void)Once;
     return;
   }
   }
 
-  std::memcpy(MemPtr, &FuncRecord, sizeof(FunctionRecord));
-  MemPtr += sizeof(FunctionRecord);
+  std::memcpy(TLD.RecordPtr, &FuncRecord, sizeof(FunctionRecord));
+  TLD.RecordPtr += sizeof(FunctionRecord);
+  incrementExtents(sizeof(FunctionRecord));
 }
 
 static uint64_t thresholdTicks() {
   static uint64_t TicksPerSec = probeRequiredCPUFeatures()
                                     ? getTSCFrequency()
                                     : __xray::NanosecondsPerSecond;
   static const uint64_t ThresholdTicks =
       TicksPerSec * flags()->xray_fdr_log_func_duration_threshold_us / 1000000;
   return ThresholdTicks;
 }
 
 // Re-point the thread local pointer into this thread's Buffer before the recent
 // "Function Entry" record and any "Tail Call Exit" records after that.
 static void rewindRecentCall(uint64_t TSC, uint64_t &LastTSC,
                              uint64_t &LastFunctionEntryTSC, int32_t FuncId) {
   auto &TLD = getThreadLocalData();
   TLD.RecordPtr -= FunctionRecSize;
+  decrementExtents(FunctionRecSize);
   FunctionRecord FuncRecord;
   std::memcpy(&FuncRecord, TLD.RecordPtr, FunctionRecSize);
   assert(FuncRecord.RecordKind ==
              uint8_t(FunctionRecord::RecordKinds::FunctionEnter) &&
          "Expected to find function entry recording when rewinding.");
   assert(FuncRecord.FuncId == (FuncId & ~(0x0F << 28)) &&
          "Expected matching function id when rewinding Exit");
   --TLD.NumConsecutiveFnEnters;
while (TLD.NumTailCalls > 0) {
 
     // We can erase a tail exit pair that we're exiting through since
     // its duration is under threshold.
     --TLD.NumTailCalls;
     RewindingRecordPtr -= FunctionRecSize;
     RewindingTSC -= ExpectedFunctionEntry.TSCDelta;
     TLD.RecordPtr -= 2 * FunctionRecSize;
     LastTSC = RewindingTSC;
+    decrementExtents(2 * FunctionRecSize);
   }
 }
 
 inline bool releaseThreadLocalBuffer(BufferQueue &BQArg) {
   auto &TLD = getThreadLocalData();
   auto EC = BQArg.releaseBuffer(TLD.Buffer);
   if (EC != BufferQueue::ErrorCode::Ok) {
     Report("Failed to release buffer at %p; error=%s\n", TLD.Buffer.Buffer,
            BufferQueue::getErrorString(EC));
     return false;
   }
   return true;
 }
 
 inline bool prepareBuffer(uint64_t TSC, unsigned char CPU,
                           int (*wall_clock_reader)(clockid_t,
                                                    struct timespec *),
                           size_t MaxSize) XRAY_NEVER_INSTRUMENT {
   auto &TLD = getThreadLocalData();
   char *BufferStart = static_cast<char *>(TLD.Buffer.Buffer);
-  if ((TLD.RecordPtr + MaxSize) >
-      (BufferStart + TLD.Buffer.Size - MetadataRecSize)) {
-    writeEOBMetadata();
-    if (!releaseThreadLocalBuffer(*TLD.LocalBQ))
+  if ((TLD.RecordPtr + MaxSize) > (BufferStart + TLD.Buffer.Size)) {
+    if (!releaseThreadLocalBuffer(*TLD.BQ))
       return false;
-    auto EC = TLD.LocalBQ->getBuffer(TLD.Buffer);
+    auto EC = TLD.BQ->getBuffer(TLD.Buffer);
     if (EC != BufferQueue::ErrorCode::Ok) {
       Report("Failed to acquire a buffer; error=%s\n",
              BufferQueue::getErrorString(EC));
       return false;
     }
     setupNewBuffer(wall_clock_reader);
 
     // Always write the CPU metadata as the first record in the buffer.
     writeNewCPUIdMetadata(CPU, TSC);
   }
   return true;
 }
 
-inline bool isLogInitializedAndReady(
-    std::shared_ptr<BufferQueue> &LBQ, uint64_t TSC, unsigned char CPU,
-    int (*wall_clock_reader)(clockid_t,
-                             struct timespec *)) XRAY_NEVER_INSTRUMENT {
+inline bool
+isLogInitializedAndReady(BufferQueue *LBQ, uint64_t TSC, unsigned char CPU,
+                         int (*wall_clock_reader)(clockid_t, struct timespec *))
+    XRAY_NEVER_INSTRUMENT {
   // Bail out right away if logging is not initialized yet.
   // We should take the opportunity to release the buffer though.
   auto Status = __sanitizer::atomic_load(&LoggingStatus,
                                          __sanitizer::memory_order_acquire);
   auto &TLD = getThreadLocalData();
   if (Status != XRayLogInitStatus::XRAY_LOG_INITIALIZED) {
     if (TLD.RecordPtr != nullptr &&
         (Status == XRayLogInitStatus::XRAY_LOG_FINALIZING ||
          Status == XRayLogInitStatus::XRAY_LOG_FINALIZED)) {
-      writeEOBMetadata();
       if (!releaseThreadLocalBuffer(*LBQ))
         return false;
       TLD.RecordPtr = nullptr;
-      LBQ = nullptr;
       return false;
     }
     return false;
   }
 
   if (__sanitizer::atomic_load(&LoggingStatus,
                                __sanitizer::memory_order_acquire) !=
           XRayLogInitStatus::XRAY_LOG_INITIALIZED ||
       LBQ->finalizing()) {
-    writeEOBMetadata();
     if (!releaseThreadLocalBuffer(*LBQ))
       return false;
     TLD.RecordPtr = nullptr;
   }
 
   if (TLD.Buffer.Buffer == nullptr) {
     auto EC = LBQ->getBuffer(TLD.Buffer);
     if (EC != BufferQueue::ErrorCode::Ok) {
inline uint32_t writeCurrentCPUTSC(ThreadLocalData &TLD, uint64_t TSC,
 
   writeTSCWrapMetadata(TSC);
   return 0;
 }
 
 inline void endBufferIfFull() XRAY_NEVER_INSTRUMENT {
   auto &TLD = getThreadLocalData();
   auto BufferStart = static_cast<char *>(TLD.Buffer.Buffer);
-  if ((TLD.RecordPtr + MetadataRecSize) - BufferStart == MetadataRecSize) {
-    writeEOBMetadata();
-    if (!releaseThreadLocalBuffer(*TLD.LocalBQ))
+  if ((TLD.RecordPtr + MetadataRecSize) - BufferStart <=
+      ptrdiff_t{MetadataRecSize}) {
+    if (!releaseThreadLocalBuffer(*TLD.BQ))
       return;
     TLD.RecordPtr = nullptr;
   }
 }
 
 thread_local volatile bool Running = false;
 
 /// Here's where the meat of the processing happens. The writer captures
 /// function entry, exit and tail exit points with a time and will create
 /// TSCWrap, NewCPUId and Function records as necessary. The writer might
 /// walk backward through its buffer and erase trivial functions to avoid
 /// polluting the log and may use the buffer queue to obtain or release a
 /// buffer.
-inline void processFunctionHook(
-    int32_t FuncId, XRayEntryType Entry, uint64_t TSC, unsigned char CPU,
-    uint64_t Arg1, int (*wall_clock_reader)(clockid_t, struct timespec *),
-    const std::shared_ptr<BufferQueue> &BQ) XRAY_NEVER_INSTRUMENT {
+inline void processFunctionHook(int32_t FuncId, XRayEntryType Entry,
+                                uint64_t TSC, unsigned char CPU, uint64_t Arg1,
+                                int (*wall_clock_reader)(clockid_t,
+                                                         struct timespec *),
+                                BufferQueue *BQ) XRAY_NEVER_INSTRUMENT {
   // Prevent signal handler recursion, so in case we're already in a log writing
   // mode and the signal handler comes in (and is also instrumented) then we
   // don't want to be clobbering potentially partial writes already happening in
   // the thread. We use a simple thread_local latch to only allow one on-going
   // handleArg0 to happen at any given time.
   RecursionGuard Guard{Running};
   if (!Guard) {
     assert(Running == true && "RecursionGuard is buggy!");
     return;
   }
 
   auto &TLD = getThreadLocalData();
 
   // In case the reference has been cleaned up before, we make sure we
   // initialize it to the provided BufferQueue.
-  if (TLD.LocalBQ == nullptr)
-    TLD.LocalBQ = BQ;
+  if (TLD.BQ == nullptr)
+    TLD.BQ = BQ;
 
-  if (!isLogInitializedAndReady(TLD.LocalBQ, TSC, CPU, wall_clock_reader))
+  if (!isLogInitializedAndReady(TLD.BQ, TSC, CPU, wall_clock_reader))
     return;
 
   // Before we go setting up writing new function entries, we need to be really
   // careful about the pointer math we're doing. This means we need to ensure
   // that the record we are about to write is going to fit into the buffer,
   // without overflowing the buffer.
   //
   // To do this properly, we use the following assumptions:
inline void processFunctionHook(int32_t FuncId, XRayEntryType Entry,
   //       2. When we learn that we've moved CPUs, we need to write a
   //          MetadataRecord to indicate a "cpu change", and thus write out the
   //          current TSC for that CPU before writing out the actual
   //          FunctionRecord.
   //       3. When we learn about a new CPU ID, we need to write down a "new cpu
   //          id" MetadataRecord before writing out the actual FunctionRecord.
   //       4. The second MetadataRecord is the optional function call argument.
   //
-  //   - An End-of-Buffer (EOB) MetadataRecord is 16 bytes.
-  //
-  // So the math we need to do is to determine whether writing 24 bytes past the
-  // current pointer leaves us with enough bytes to write the EOB
-  // MetadataRecord. If we don't have enough space after writing as much as 24
-  // bytes in the end of the buffer, we need to write out the EOB, get a new
-  // Buffer, set it up properly before doing any further writing.
+  // So the math we need to do is to determine whether writing 40 bytes past the
+  // current pointer exceeds the buffer's maximum size. If we don't have enough
+  // space to write 40 bytes in the buffer, we need get a new Buffer, set it up
+  // properly before doing any further writing.
   size_t MaxSize = FunctionRecSize + 2 * MetadataRecSize;
   if (!prepareBuffer(TSC, CPU, wall_clock_reader, MaxSize)) {
-    TLD.LocalBQ = nullptr;
+    TLD.BQ = nullptr;
     return;
   }
 
   // By this point, we are now ready to write up to 40 bytes (explained above).
   assert((TLD.RecordPtr + MaxSize) - static_cast<char *>(TLD.Buffer.Buffer) >=
              static_cast<ptrdiff_t>(MetadataRecSize) &&
          "Misconfigured BufferQueue provided; Buffer size not large enough.");
 
static bool Once = [&] {
              FuncId);
       return true;
     }();
     (void)Once;
     return;
   }
   }
 
-  writeFunctionRecord(FuncId, RecordTSCDelta, Entry, TLD.RecordPtr);
+  writeFunctionRecord(FuncId, RecordTSCDelta, Entry);
   if (Entry == XRayEntryType::LOG_ARGS_ENTRY)
     writeCallArgumentMetadata(Arg1);
 
   // If we've exhausted the buffer by this time, we then release the buffer to
   // make sure that other threads may start using this buffer.
   endBufferIfFull();
 }
 
 } // namespace __xray_fdr_internal
 } // namespace __xray
 
 #endif // XRAY_XRAY_FDR_LOGGING_IMPL_H

compiler-rt/lib/xray/xray_flags.inc

XRAY_FLAG(bool, xray_naive_log, true,
           "Whether to install the naive log implementation.")
 XRAY_FLAG(const char *, xray_logfile_base, "xray-log.",
           "Filename base for the xray logfile.")
 XRAY_FLAG(bool, xray_fdr_log, false,
           "Whether to install the flight data recorder logging implementation.")
 XRAY_FLAG(int, xray_fdr_log_func_duration_threshold_us, 5,
           "FDR logging will try to skip functions that execute for fewer "
           "microseconds than this threshold.")
+XRAY_FLAG(int, xray_fdr_log_grace_period_us, 100000,
+          "FDR logging will wait this much time in microseconds before "
+          "actually flushing the log; this gives a chance for threads to "
+          "notice that the log has been finalized and clean up.")

llvm/lib/XRay/Trace.cpp

struct FDRState {
   enum class Token {
     NEW_BUFFER_RECORD_OR_EOF,
     WALLCLOCK_RECORD,
     NEW_CPU_ID_RECORD,
     FUNCTION_SEQUENCE,
     SCAN_TO_END_OF_THREAD_BUF,
     CUSTOM_EVENT_DATA,
     CALL_ARGUMENT,
+    BUFFER_EXTENTS,
   };
   Token Expects;
 
   // Each threads buffer may have trailing garbage to scan over, so we track our
   // progress.
   uint64_t CurrentBufferSize;
   uint64_t CurrentBufferConsumed;
 };
const char *fdrStateToTwine(const FDRState::Token &state) {
   case FDRState::Token::FUNCTION_SEQUENCE:
     return "FUNCTION_SEQUENCE";
   case FDRState::Token::SCAN_TO_END_OF_THREAD_BUF:
     return "SCAN_TO_END_OF_THREAD_BUF";
   case FDRState::Token::CUSTOM_EVENT_DATA:
     return "CUSTOM_EVENT_DATA";
   case FDRState::Token::CALL_ARGUMENT:
     return "CALL_ARGUMENT";
+  case FDRState::Token::BUFFER_EXTENTS:
+    return "BUFFER_EXTENTS";
   }
   return "UNKNOWN";
 }
 
 /// State transition when a NewBufferRecord is encountered.
 Error processFDRNewBufferRecord(FDRState &State, uint8_t RecordFirstByte,
                                 DataExtractor &RecordExtractor) {
 
   if (State.Expects != FDRState::Token::NEW_BUFFER_RECORD_OR_EOF)
     return make_error<StringError>(
-        "Malformed log. Read New Buffer record kind out of sequence",
+        Twine("Malformed log. Read New Buffer record kind out of sequence; "
+              "expected: ") +
+            fdrStateToTwine(State.Expects),
         std::make_error_code(std::errc::executable_format_error));
   uint32_t OffsetPtr = 1; // 1 byte into record.
   State.ThreadId = RecordExtractor.getU16(&OffsetPtr);
   State.Expects = FDRState::Token::WALLCLOCK_RECORD;
   return Error::success();
 }
 
 /// State transition when an EndOfBufferRecord is encountered.
 Error processFDREndOfBufferRecord(FDRState &State, uint8_t RecordFirstByte,
                                   DataExtractor &RecordExtractor) {
   if (State.Expects == FDRState::Token::NEW_BUFFER_RECORD_OR_EOF)
     return make_error<StringError>(
-        "Malformed log. Received EOB message without current buffer.",
+        Twine("Malformed log. Received EOB message without current buffer; "
+              "expected: ") +
+            fdrStateToTwine(State.Expects),
         std::make_error_code(std::errc::executable_format_error));
   State.Expects = FDRState::Token::SCAN_TO_END_OF_THREAD_BUF;
   return Error::success();
 }
 
 /// State transition when a NewCPUIdRecord is encountered.
 Error processFDRNewCPUIdRecord(FDRState &State, uint8_t RecordFirstByte,
                                DataExtractor &RecordExtractor) {
   if (State.Expects != FDRState::Token::FUNCTION_SEQUENCE &&
       State.Expects != FDRState::Token::NEW_CPU_ID_RECORD)
     return make_error<StringError>(
-        "Malformed log. Read NewCPUId record kind out of sequence",
+        Twine("Malformed log. Read NewCPUId record kind out of sequence; "
+              "expected: ") +
+            fdrStateToTwine(State.Expects),
         std::make_error_code(std::errc::executable_format_error));
   uint32_t OffsetPtr = 1; // Read starting after the first byte.
   State.CPUId = RecordExtractor.getU16(&OffsetPtr);
   State.BaseTSC = RecordExtractor.getU64(&OffsetPtr);
   State.Expects = FDRState::Token::FUNCTION_SEQUENCE;
   return Error::success();
 }
 
 /// State transition when a TSCWrapRecord (overflow detection) is encountered.
 Error processFDRTSCWrapRecord(FDRState &State, uint8_t RecordFirstByte,
                               DataExtractor &RecordExtractor) {
   if (State.Expects != FDRState::Token::FUNCTION_SEQUENCE)
     return make_error<StringError>(
-        "Malformed log. Read TSCWrap record kind out of sequence",
+        Twine("Malformed log. Read TSCWrap record kind out of sequence; "
+              "expecting: ") +
+            fdrStateToTwine(State.Expects),
         std::make_error_code(std::errc::executable_format_error));
   uint32_t OffsetPtr = 1; // Read starting after the first byte.
   State.BaseTSC = RecordExtractor.getU64(&OffsetPtr);
   return Error::success();
 }
 
 /// State transition when a WallTimeMarkerRecord is encountered.
 Error processFDRWallTimeRecord(FDRState &State, uint8_t RecordFirstByte,
-                               DataExtractor &RecordExtractor) {
+                               DataExtractor &RecordExtractor,
+                               uint16_t Version) {
   if (State.Expects != FDRState::Token::WALLCLOCK_RECORD)
     return make_error<StringError>(
-        "Malformed log. Read Wallclock record kind out of sequence",
+        Twine("Malformed log. Read Wallclock record kind out of sequence; "
+              "expecting: ") +
+            fdrStateToTwine(State.Expects),
         std::make_error_code(std::errc::executable_format_error));
-  // We don't encode the wall time into any of the records.
-  // XRayRecords are concerned with the TSC instead.
+
+  // TODO: Someday, reconcile the TSC ticks to wall clock time for presentation
+  // purposes. For now, we're ignoring these records.
   State.Expects = FDRState::Token::NEW_CPU_ID_RECORD;
   return Error::success();
 }
 
 /// State transition when a CustomEventMarker is encountered.
 Error processCustomEventMarker(FDRState &State, uint8_t RecordFirstByte,
                                DataExtractor &RecordExtractor,
                                size_t &RecordSize) {
   // We can encounter a CustomEventMarker anywhere in the log, so we can handle
   // it regardless of the expectation. However, we do set the expectation to
   // read a set number of fixed bytes, as described in the metadata.
   uint32_t OffsetPtr = 1; // Read after the first byte.
   uint32_t DataSize = RecordExtractor.getU32(&OffsetPtr);
   uint64_t TSC = RecordExtractor.getU64(&OffsetPtr);
 
   // FIXME: Actually represent the record through the API. For now we only
   // skip through the data.
   (void)TSC;
   RecordSize = 16 + DataSize;
   return Error::success();
 }
 
+/// State transition when an BufferExtents record is encountered.
+Error processBufferExtents(FDRState &State, uint8_t RecordFirstByte,
+                           DataExtractor &RecordExtractor) {
+  if (State.Expects != FDRState::Token::BUFFER_EXTENTS)
+    return make_error<StringError>(
+        Twine("Malformed log. Buffer Extents unexpected; expected: ") +
+            fdrStateToTwine(State.Expects),
+        std::make_error_code(std::errc::executable_format_error));
+  uint32_t OffsetPtr = 1; // Read after the first byte.
+  State.CurrentBufferSize = RecordExtractor.getU64(&OffsetPtr);
+  State.Expects = FDRState::Token::NEW_BUFFER_RECORD_OR_EOF;
+  return Error::success();
+}
+
 /// State transition when a CallArgumentRecord is encountered.
 Error processFDRCallArgumentRecord(FDRState &State, uint8_t RecordFirstByte,
                                    DataExtractor &RecordExtractor,
                                    std::vector<XRayRecord> &Records) {
   uint32_t OffsetPtr = 1; // Read starting after the first byte.
   auto &Enter = Records.back();
 
   if (Enter.Type != RecordTypes::ENTER)
     return make_error<StringError>(
         "CallArgument needs to be right after a function entry",
         std::make_error_code(std::errc::executable_format_error));
   Enter.Type = RecordTypes::ENTER_ARG;
   Enter.CallArgs.emplace_back(RecordExtractor.getU64(&OffsetPtr));
   return Error::success();
 }
 
 /// Advances the state machine for reading the FDR record type by reading one
 /// Metadata Record and updating the State appropriately based on the kind of
 /// record encountered. The RecordKind is encoded in the first byte of the
 /// Record, which the caller should pass in because they have already read it
 /// to determine that this is a metadata record as opposed to a function record.
+///
+/// Beginning with Version 2 of the FDR log, we do not depend on the size of the
+/// buffer, but rather use the extents to determine how far to read in the log
+/// for this particular buffer.
 Error processFDRMetadataRecord(FDRState &State, uint8_t RecordFirstByte,
                                DataExtractor &RecordExtractor,
                                size_t &RecordSize,
-                               std::vector<XRayRecord> &Records) {
+                               std::vector<XRayRecord> &Records,
+                               uint16_t Version, size_t &BufferSize) {
   // The remaining 7 bits are the RecordKind enum.
   uint8_t RecordKind = RecordFirstByte >> 1;
   switch (RecordKind) {
   case 0: // NewBuffer
     if (auto E =
             processFDRNewBufferRecord(State, RecordFirstByte, RecordExtractor))
       return E;
     break;
   case 1: // EndOfBuffer
+    if (Version >= 2)
+      return make_error<StringError>(
+          "Since Version 2 of FDR logging, we no longer support EOB records.",
+          std::make_error_code(std::errc::executable_format_error));
     if (auto E = processFDREndOfBufferRecord(State, RecordFirstByte,
                                              RecordExtractor))
       return E;
     break;
   case 2: // NewCPUId
     if (auto E =
             processFDRNewCPUIdRecord(State, RecordFirstByte, RecordExtractor))
       return E;
     break;
   case 3: // TSCWrap
     if (auto E =
             processFDRTSCWrapRecord(State, RecordFirstByte, RecordExtractor))
       return E;
     break;
   case 4: // WallTimeMarker
-    if (auto E =
-            processFDRWallTimeRecord(State, RecordFirstByte, RecordExtractor))
+    if (auto E = processFDRWallTimeRecord(State, RecordFirstByte,
+                                          RecordExtractor, Version))
       return E;
     break;
   case 5: // CustomEventMarker
     if (auto E = processCustomEventMarker(State, RecordFirstByte,
                                           RecordExtractor, RecordSize))
       return E;
     break;
   case 6: // CallArgument
     if (auto E = processFDRCallArgumentRecord(State, RecordFirstByte,
                                               RecordExtractor, Records))
       return E;
     break;
+  case 7: // BufferExtents
+    if (auto E = processBufferExtents(State, RecordFirstByte, RecordExtractor))
+      return E;
+    break;
   default:
     // Widen the record type to uint16_t to prevent conversion to char.
     return make_error<StringError>(
         Twine("Illegal metadata record type: ")
             .concat(Twine(static_cast<unsigned>(RecordKind))),
         std::make_error_code(std::errc::executable_format_error));
   }
   return Error::success();

 /// The following is an attempt to document the grammar of the format, which is
 /// parsed by this function for little-endian machines. Since the format makes
 /// use of BitFields, when we support big-endian architectures, we will need to
 /// adjust not only the endianness parameter to llvm's RecordExtractor, but also
 /// the bit twiddling logic, which is consistent with the little-endian
 /// convention that BitFields within a struct will first be packed into the
 /// least significant bits the address they belong to.
 ///
-/// We expect a format complying with the grammar in the following pseudo-EBNF.
+/// We expect a format complying with the grammar in the following pseudo-EBNF
+/// in Version 1 of the FDR log.
 ///
 /// FDRLog: XRayFileHeader ThreadBuffer*
 /// XRayFileHeader: 32 bytes to identify the log as FDR with machine metadata.
 ///     Includes BufferSize
 /// ThreadBuffer: NewBuffer WallClockTime NewCPUId FunctionSequence EOB
 /// BufSize: 8 byte unsigned integer indicating how large the buffer is.
 /// NewBuffer: 16 byte metadata record with Thread Id.
 /// WallClockTime: 16 byte metadata record with human readable time.
 /// NewCPUId: 16 byte metadata record with CPUId and a 64 bit TSC reading.
 /// EOB: 16 byte record in a thread buffer plus mem garbage to fill BufSize.
 /// FunctionSequence: NewCPUId | TSCWrap | FunctionRecord
 /// TSCWrap: 16 byte metadata record with a full 64 bit TSC reading.
 /// FunctionRecord: 8 byte record with FunctionId, entry/exit, and TSC delta.
+///
+/// In Version 2, we make the following changes:
+///
+/// ThreadBuffer: BufferExtents NewBuffer WallClockTime NewCPUId
+///               FunctionSequence
+/// BufferExtents: 16 byte metdata record describing how many usable bytes are
+///                in the buffer. This is measured from the start of the buffer
+///                and must always be at least 48 (bytes).
+/// EOB: *deprecated*
 Error loadFDRLog(StringRef Data, XRayFileHeader &FileHeader,
                  std::vector<XRayRecord> &Records) {
   if (Data.size() < 32)
     return make_error<StringError>(
         "Not enough bytes for an XRay log.",
         std::make_error_code(std::errc::invalid_argument));
 
   // For an FDR log, there are records sized 16 and 8 bytes.
Error loadFDRLog(StringRef Data, XRayFileHeader &FileHeader,
 
   uint64_t BufferSize = 0;
   {
     StringRef ExtraDataRef(FileHeader.FreeFormData, 16);
     DataExtractor ExtraDataExtractor(ExtraDataRef, true, 8);
     uint32_t ExtraDataOffset = 0;
     BufferSize = ExtraDataExtractor.getU64(&ExtraDataOffset);
   }
-  FDRState State{0,          0, 0, FDRState::Token::NEW_BUFFER_RECORD_OR_EOF,
-                 BufferSize, 0};
+
+  FDRState::Token InitialExpectation;
+  switch (FileHeader.Version) {
+  case 1:
+    InitialExpectation = FDRState::Token::NEW_BUFFER_RECORD_OR_EOF;
+    break;
+  case 2:
+    InitialExpectation = FDRState::Token::BUFFER_EXTENTS;
+    break;
+  default:
+    return make_error<StringError>(
+        Twine("Unsupported version '") + Twine(FileHeader.Version) + "'",
+        std::make_error_code(std::errc::executable_format_error));
+  }
+  FDRState State{0, 0, 0, InitialExpectation, BufferSize, 0};
+
   // RecordSize will tell the loop how far to seek ahead based on the record
   // type that we have just read.
   size_t RecordSize = 0;
   for (auto S = Data.drop_front(32); !S.empty(); S = S.drop_front(RecordSize)) {
     DataExtractor RecordExtractor(S, true, 8);
     uint32_t OffsetPtr = 0;
     if (State.Expects == FDRState::Token::SCAN_TO_END_OF_THREAD_BUF) {
       RecordSize = State.CurrentBufferSize - State.CurrentBufferConsumed;
       if (S.size() < RecordSize) {
         return make_error<StringError>(
             Twine("Incomplete thread buffer. Expected at least ") +
                 Twine(RecordSize) + " bytes but found " + Twine(S.size()),
             make_error_code(std::errc::invalid_argument));
       }
       State.CurrentBufferConsumed = 0;
       State.Expects = FDRState::Token::NEW_BUFFER_RECORD_OR_EOF;
       continue;
     }
     uint8_t BitField = RecordExtractor.getU8(&OffsetPtr);
     bool isMetadataRecord = BitField & 0x01uL;
+    bool isBufferExtents =
+        (BitField >> 1) == 7; // BufferExtents record kind == 7
     if (isMetadataRecord) {
       RecordSize = 16;
-      if (auto E = processFDRMetadataRecord(State, BitField, RecordExtractor,
-                                            RecordSize, Records))
+      if (auto E = processFDRMetadataRecord(
+              State, BitField, RecordExtractor, RecordSize, Records,
+              FileHeader.Version, State.CurrentBufferSize))
         return E;
     } else { // Process Function Record
       RecordSize = 8;
       if (auto E = processFDRFunctionRecord(State, BitField, RecordExtractor,
                                             Records))
         return E;
     }
+
+    // The BufferExtents record is technically not part of the buffer, so we
+    // don't count the size of that record against the buffer's actual size.
+    if (!isBufferExtents)
       State.CurrentBufferConsumed += RecordSize;
+    assert(State.CurrentBufferConsumed <= State.CurrentBufferSize);
+    if (FileHeader.Version == 2 &&
+        State.CurrentBufferSize == State.CurrentBufferConsumed) {
+      // In Version 2 of the log, we don't need to scan to the end of the thread
+      // buffer if we've already consumed all the bytes we need to.
+      State.Expects = FDRState::Token::BUFFER_EXTENTS;
+      State.CurrentBufferSize = BufferSize;
+      State.CurrentBufferConsumed = 0;
+    }
   }
 
   // Having iterated over everything we've been given, we've either consumed
   // everything and ended up in the end state, or were told to skip the rest.
   bool Finished = State.Expects == FDRState::Token::SCAN_TO_END_OF_THREAD_BUF &&
                   State.CurrentBufferSize == State.CurrentBufferConsumed;
-  if (State.Expects != FDRState::Token::NEW_BUFFER_RECORD_OR_EOF && !Finished)
+  if ((State.Expects != FDRState::Token::NEW_BUFFER_RECORD_OR_EOF &&
+       State.Expects != FDRState::Token::BUFFER_EXTENTS) &&
+      !Finished)
     return make_error<StringError>(
         Twine("Encountered EOF with unexpected state expectation ") +
             fdrStateToTwine(State.Expects) +
             ". Remaining expected bytes in thread buffer total " +
             Twine(State.CurrentBufferSize - State.CurrentBufferConsumed),
         std::make_error_code(std::errc::executable_format_error));
 
   return Error::success();
Expected<Trace> llvm::xray::loadTraceFile(StringRef Filename, bool Sort) {
 
   // Attempt to detect the file type using file magic. We have a slight bias
   // towards the binary format, and we do this by making sure that the first 4
   // bytes of the binary file is some combination of the following byte
   // patterns: (observe the code loading them assumes they're little endian)
   //
   //   0x01 0x00 0x00 0x00 - version 1, "naive" format
   //   0x01 0x00 0x01 0x00 - version 1, "flight data recorder" format
+  //   0x02 0x00 0x01 0x00 - version 2, "flight data recorder" format
   //
   // YAML files don't typically have those first four bytes as valid text so we
   // try loading assuming YAML if we don't find these bytes.
   //
   // Only if we can't load either the binary or the YAML format will we yield an
   // error.
   StringRef Magic(MappedFile.data(), 4);
   DataExtractor HeaderExtractor(Magic, true, 8);
   uint32_t OffsetPtr = 0;
   uint16_t Version = HeaderExtractor.getU16(&OffsetPtr);
   uint16_t Type = HeaderExtractor.getU16(&OffsetPtr);
 
   enum BinaryFormatType { NAIVE_FORMAT = 0, FLIGHT_DATA_RECORDER_FORMAT = 1 };
 
   Trace T;
-  if (Type == NAIVE_FORMAT && (Version == 1 || Version == 2)) {
+  switch (Type) {
+  case NAIVE_FORMAT:
+    if (Version == 1 || Version == 2) {
       if (auto E = loadNaiveFormatLog(Data, T.FileHeader, T.Records))
         return std::move(E);
-  } else if (Version == 1 && Type == FLIGHT_DATA_RECORDER_FORMAT) {
+    } else {
+      return make_error<StringError>(
+          Twine("Unsupported version for Basic/Naive Mode logging: ") +
+              Twine(Version),
+          std::make_error_code(std::errc::executable_format_error));
+    }
+    break;
+  case FLIGHT_DATA_RECORDER_FORMAT:
+    if (Version == 1 || Version == 2) {
       if (auto E = loadFDRLog(Data, T.FileHeader, T.Records))
         return std::move(E);
     } else {
+      return make_error<StringError>(
+          Twine("Unsupported version for FDR Mode logging: ") + Twine(Version),
+          std::make_error_code(std::errc::executable_format_error));
+    }
+    break;
+  default:
     if (auto E = loadYAMLLog(Data, T.FileHeader, T.Records))
       return std::move(E);
   }
 
   if (Sort)
     std::sort(T.Records.begin(), T.Records.end(),
               [&](const XRayRecord &L, const XRayRecord &R) {
                 return L.TSC < R.TSC;
               });
 
   return std::move(T);
 }