[XRay][compiler-rt] FDR Mode Controller

Summary:
This change implements a controller for abstracting away the details of
what happens when tracing with FDR mode. This controller type allows us
to test in isolation the various cases where we're encountering function
entry, exit, and other kinds of events we are handling when FDR mode is
enabled.

This change introduces a number of testing facilities we've needed to
better support expressing the conditions we need for the unit tests. We
leave some TODOs for moving those utilities into the LLVM project,
sitting in the `Testing` library, to make matching conditions on XRay
`Trace` instances through googlemock more manageable and declarative.

We don't wire in the controller right away, to allow us to incrementally
update the implementation(s) as we increase testing coverage of the
controller type. There's a need to re-think the way we're managing
buffers in a multi-threaded environment, which is more invasive than
this implementation.

This step in the process allows us to encode our assumptions in the
implementation of the controller, and then evolve the buffer queue
implementation to support generational buffer management to ensure we
can continue to support the cases we're already supporting with the
controller.

Reviewers: mboerger, eizan

Subscribers: mgorny, llvm-commits, jfb

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

git-svn-id: https://llvm.org/svn/llvm-project/compiler-rt/trunk@344488 91177308-0d34-0410-b5e6-96231b3b80d8

8 files changed, 758 insertions, 29 deletions

diff --git a/lib/xray/CMakeLists.txt b/lib/xray/CMakeLists.txt
index 4bb25d454..541e181af 100644
--- a/lib/xray/CMakeLists.txt
+++ b/lib/xray/CMakeLists.txt
@@ -67,9 +67,11 @@ set(XRAY_IMPL_HEADERS
   xray_basic_logging.h
   xray_buffer_queue.h
   xray_defs.h
+  xray_fdr_controller.h
   xray_fdr_flags.h
   xray_fdr_flags.inc
   xray_fdr_log_records.h
+  xray_fdr_log_writer.h
   xray_fdr_logging.h
   xray_flags.h
   xray_flags.inc
diff --git a/lib/xray/tests/unit/CMakeLists.txt b/lib/xray/tests/unit/CMakeLists.txt
index d0ead947d..42ea43750 100644
--- a/lib/xray/tests/unit/CMakeLists.txt
+++ b/lib/xray/tests/unit/CMakeLists.txt
@@ -1,8 +1,10 @@
 add_xray_unittest(XRayTest SOURCES
-  buffer_queue_test.cc
   allocator_test.cc
-  segmented_array_test.cc
+  buffer_queue_test.cc
+  fdr_controller_test.cc
+  fdr_log_writer_test.cc
   function_call_trie_test.cc
   profile_collector_test.cc
-  fdr_log_writer_test.cc
+  segmented_array_test.cc
+  test_helpers.cc
   xray_unit_test_main.cc)
diff --git a/lib/xray/tests/unit/fdr_controller_test.cc b/lib/xray/tests/unit/fdr_controller_test.cc
new file mode 100644
index 000000000..73666a76b
--- /dev/null
+++ b/lib/xray/tests/unit/fdr_controller_test.cc
@@ -0,0 +1,242 @@
+//===-- fdr_controller_test.cc --------------------------------------------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+// This file is a part of XRay, a function call tracing system.
+//
+//===----------------------------------------------------------------------===//
+#include <algorithm>
+#include <memory>
+#include <time.h>
+
+#include "test_helpers.h"
+#include "xray/xray_records.h"
+#include "xray_buffer_queue.h"
+#include "xray_fdr_controller.h"
+#include "xray_fdr_log_writer.h"
+#include "llvm/Support/DataExtractor.h"
+#include "llvm/Testing/Support/Error.h"
+#include "llvm/XRay/Trace.h"
+#include "llvm/XRay/XRayRecord.h"
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+
+namespace __xray {
+namespace {
+
+using ::llvm::HasValue;
+using ::llvm::xray::testing::FuncId;
+using ::llvm::xray::testing::HasArg;
+using ::llvm::xray::testing::RecordType;
+using ::testing::AllOf;
+using ::testing::ElementsAre;
+using ::testing::Eq;
+using ::testing::Field;
+using ::testing::IsEmpty;
+using ::testing::SizeIs;
+
+class FunctionSequenceTest : public ::testing::Test {
+protected:
+  BufferQueue::Buffer B{};
+  std::unique_ptr<BufferQueue> BQ;
+  std::unique_ptr<FDRLogWriter> W;
+  std::unique_ptr<FDRController<>> C;
+
+public:
+  void SetUp() override {
+    bool Success;
+    BQ = llvm::make_unique<BufferQueue>(4096, 1, Success);
+    ASSERT_TRUE(Success);
+    ASSERT_EQ(BQ->getBuffer(B), BufferQueue::ErrorCode::Ok);
+    W = llvm::make_unique<FDRLogWriter>(B);
+    C = llvm::make_unique<FDRController<>>(BQ.get(), B, *W, clock_gettime, 0);
+  }
+};
+
+TEST_F(FunctionSequenceTest, DefaultInitFinalizeFlush) {
+  ASSERT_TRUE(C->functionEnter(1, 2, 3));
+  ASSERT_TRUE(C->functionExit(1, 2, 3));
+  ASSERT_TRUE(C->flush());
+  ASSERT_EQ(BQ->finalize(), BufferQueue::ErrorCode::Ok);
+
+  // Serialize the buffers then test to see we find the expected records.
+  std::string Serialized = serialize(*BQ, 3);
+  llvm::DataExtractor DE(Serialized, true, 8);
+  auto TraceOrErr = llvm::xray::loadTrace(DE);
+  EXPECT_THAT_EXPECTED(
+      TraceOrErr,
+      HasValue(ElementsAre(
+          AllOf(FuncId(1), RecordType(llvm::xray::RecordTypes::ENTER)),
+          AllOf(FuncId(1), RecordType(llvm::xray::RecordTypes::EXIT)))));
+}
+
+TEST_F(FunctionSequenceTest, ThresholdsAreEnforced) {
+  C = llvm::make_unique<FDRController<>>(BQ.get(), B, *W, clock_gettime, 1000);
+  ASSERT_TRUE(C->functionEnter(1, 2, 3));
+  ASSERT_TRUE(C->functionExit(1, 2, 3));
+  ASSERT_TRUE(C->flush());
+  ASSERT_EQ(BQ->finalize(), BufferQueue::ErrorCode::Ok);
+
+  // Serialize the buffers then test to see we find the *no* records, because
+  // the function entry-exit comes under the cycle threshold.
+  std::string Serialized = serialize(*BQ, 3);
+  llvm::DataExtractor DE(Serialized, true, 8);
+  auto TraceOrErr = llvm::xray::loadTrace(DE);
+  EXPECT_THAT_EXPECTED(TraceOrErr, HasValue(IsEmpty()));
+}
+
+TEST_F(FunctionSequenceTest, ArgsAreHandledAndKept) {
+  C = llvm::make_unique<FDRController<>>(BQ.get(), B, *W, clock_gettime, 1000);
+  ASSERT_TRUE(C->functionEnterArg(1, 2, 3, 4));
+  ASSERT_TRUE(C->functionExit(1, 2, 3));
+  ASSERT_TRUE(C->flush());
+  ASSERT_EQ(BQ->finalize(), BufferQueue::ErrorCode::Ok);
+
+  // Serialize the buffers then test to see we find the function enter arg
+  // record with the specified argument.
+  std::string Serialized = serialize(*BQ, 3);
+  llvm::DataExtractor DE(Serialized, true, 8);
+  auto TraceOrErr = llvm::xray::loadTrace(DE);
+  EXPECT_THAT_EXPECTED(
+      TraceOrErr,
+      HasValue(ElementsAre(
+          AllOf(FuncId(1), RecordType(llvm::xray::RecordTypes::ENTER_ARG),
+                HasArg(4)),
+          AllOf(FuncId(1), RecordType(llvm::xray::RecordTypes::EXIT)))));
+}
+
+TEST_F(FunctionSequenceTest, RewindingMultipleCalls) {
+  C = llvm::make_unique<FDRController<>>(BQ.get(), B, *W, clock_gettime, 1000);
+
+  // First we construct an arbitrarily deep function enter/call stack.
+  // We also ensure that we are in the same CPU.
+  uint64_t TSC = 1;
+  uint16_t CPU = 1;
+  ASSERT_TRUE(C->functionEnter(1, TSC++, CPU));
+  ASSERT_TRUE(C->functionEnter(2, TSC++, CPU));
+  ASSERT_TRUE(C->functionEnter(3, TSC++, CPU));
+
+  // Then we exit them one at a time, in reverse order of entry.
+  ASSERT_TRUE(C->functionExit(3, TSC++, CPU));
+  ASSERT_TRUE(C->functionExit(2, TSC++, CPU));
+  ASSERT_TRUE(C->functionExit(1, TSC++, CPU));
+
+  ASSERT_TRUE(C->flush());
+  ASSERT_EQ(BQ->finalize(), BufferQueue::ErrorCode::Ok);
+
+  // Serialize the buffers then test to see we find that all the calls have been
+  // unwound because all of them are under the cycle counter threshold.
+  std::string Serialized = serialize(*BQ, 3);
+  llvm::DataExtractor DE(Serialized, true, 8);
+  auto TraceOrErr = llvm::xray::loadTrace(DE);
+  EXPECT_THAT_EXPECTED(TraceOrErr, HasValue(IsEmpty()));
+}
+
+TEST_F(FunctionSequenceTest, RewindingIntermediaryTailExits) {
+  C = llvm::make_unique<FDRController<>>(BQ.get(), B, *W, clock_gettime, 1000);
+
+  // First we construct an arbitrarily deep function enter/call stack.
+  // We also ensure that we are in the same CPU.
+  uint64_t TSC = 1;
+  uint16_t CPU = 1;
+  ASSERT_TRUE(C->functionEnter(1, TSC++, CPU));
+  ASSERT_TRUE(C->functionEnter(2, TSC++, CPU));
+  ASSERT_TRUE(C->functionEnter(3, TSC++, CPU));
+
+  // Next we tail-exit into a new function multiple times.
+  ASSERT_TRUE(C->functionTailExit(3, TSC++, CPU));
+  ASSERT_TRUE(C->functionEnter(4, TSC++, CPU));
+  ASSERT_TRUE(C->functionTailExit(4, TSC++, CPU));
+  ASSERT_TRUE(C->functionEnter(5, TSC++, CPU));
+  ASSERT_TRUE(C->functionTailExit(5, TSC++, CPU));
+  ASSERT_TRUE(C->functionEnter(6, TSC++, CPU));
+
+  // Then we exit them one at a time, in reverse order of entry.
+  ASSERT_TRUE(C->functionExit(6, TSC++, CPU));
+  ASSERT_TRUE(C->functionExit(2, TSC++, CPU));
+  ASSERT_TRUE(C->functionExit(1, TSC++, CPU));
+  ASSERT_TRUE(C->flush());
+  ASSERT_EQ(BQ->finalize(), BufferQueue::ErrorCode::Ok);
+
+  // Serialize the buffers then test to see we find that all the calls have been
+  // unwound because all of them are under the cycle counter threshold.
+  std::string Serialized = serialize(*BQ, 3);
+  llvm::DataExtractor DE(Serialized, true, 8);
+  auto TraceOrErr = llvm::xray::loadTrace(DE);
+  EXPECT_THAT_EXPECTED(TraceOrErr, HasValue(IsEmpty()));
+}
+
+class BufferManagementTest : public ::testing::Test {
+protected:
+  BufferQueue::Buffer B{};
+  std::unique_ptr<BufferQueue> BQ;
+  std::unique_ptr<FDRLogWriter> W;
+  std::unique_ptr<FDRController<>> C;
+
+  static constexpr size_t kBuffers = 10;
+
+public:
+  void SetUp() override {
+    bool Success;
+    BQ = llvm::make_unique<BufferQueue>(sizeof(MetadataRecord) * 4 +
+                                            sizeof(FunctionRecord) * 2,
+                                        kBuffers, Success);
+    ASSERT_TRUE(Success);
+    ASSERT_EQ(BQ->getBuffer(B), BufferQueue::ErrorCode::Ok);
+    W = llvm::make_unique<FDRLogWriter>(B);
+    C = llvm::make_unique<FDRController<>>(BQ.get(), B, *W, clock_gettime, 0);
+  }
+};
+
+constexpr size_t BufferManagementTest::kBuffers;
+
+TEST_F(BufferManagementTest, HandlesOverflow) {
+  uint64_t TSC = 1;
+  uint16_t CPU = 1;
+  for (size_t I = 0; I < kBuffers; ++I) {
+    ASSERT_TRUE(C->functionEnter(1, TSC++, CPU));
+    ASSERT_TRUE(C->functionExit(1, TSC++, CPU));
+  }
+  C->flush();
+  ASSERT_EQ(BQ->finalize(), BufferQueue::ErrorCode::Ok);
+
+  std::string Serialized = serialize(*BQ, 3);
+  llvm::DataExtractor DE(Serialized, true, 8);
+  auto TraceOrErr = llvm::xray::loadTrace(DE);
+  EXPECT_THAT_EXPECTED(TraceOrErr, HasValue(SizeIs(kBuffers * 2)));
+}
+
+TEST_F(BufferManagementTest, HandlesFinalizedBufferQueue) {
+  uint64_t TSC = 1;
+  uint16_t CPU = 1;
+
+  // First write one function entry.
+  ASSERT_TRUE(C->functionEnter(1, TSC++, CPU));
+
+  // Then we finalize the buffer queue, simulating the case where the logging
+  // has been finalized.
+  ASSERT_EQ(BQ->finalize(), BufferQueue::ErrorCode::Ok);
+
+  // At this point further calls to the controller must fail.
+  ASSERT_FALSE(C->functionExit(1, TSC++, CPU));
+
+  // But flushing should succeed.
+  ASSERT_TRUE(C->flush());
+
+  // We expect that we'll only be able to find the function enter event, but not
+  // the function exit event.
+  std::string Serialized = serialize(*BQ, 3);
+  llvm::DataExtractor DE(Serialized, true, 8);
+  auto TraceOrErr = llvm::xray::loadTrace(DE);
+  EXPECT_THAT_EXPECTED(
+      TraceOrErr, HasValue(ElementsAre(AllOf(
+                      FuncId(1), RecordType(llvm::xray::RecordTypes::ENTER)))));
+}
+
+} // namespace
+} // namespace __xray
diff --git a/lib/xray/tests/unit/fdr_log_writer_test.cc b/lib/xray/tests/unit/fdr_log_writer_test.cc
index 3a2138cd8..1aeaf099d 100644
--- a/lib/xray/tests/unit/fdr_log_writer_test.cc
+++ b/lib/xray/tests/unit/fdr_log_writer_test.cc
@@ -12,6 +12,7 @@
 //===----------------------------------------------------------------------===//
 #include <time.h>
 
+#include "test_helpers.h"
 #include "xray/xray_records.h"
 #include "xray_fdr_log_writer.h"
 #include "llvm/Support/DataExtractor.h"
@@ -26,8 +27,12 @@ namespace {
 static constexpr size_t kSize = 4096;
 
 using ::llvm::HasValue;
+using ::llvm::xray::testing::FuncId;
+using ::llvm::xray::testing::RecordType;
 using ::testing::Eq;
-using ::testing::SizeIs;
+using ::testing::AllOf;
+using ::testing::IsEmpty;
+using ::testing::ElementsAre;
 
 // Exercise the common code path where we initialize a buffer and are able to
 // write some records successfully.
@@ -58,34 +63,47 @@ TEST(FdrLogWriterTest, WriteSomeRecords) {
   // We then need to go through each element of the Buffers, and re-create a
   // flat buffer that we would see if they were laid out in a file. This also
   // means we need to write out the header manually.
-  // TODO: Isolate the file header writing.
-  std::string Serialized;
-  std::aligned_storage<sizeof(XRayFileHeader), alignof(XRayFileHeader)>::type
-      HeaderStorage;
-  auto *Header = reinterpret_cast<XRayFileHeader *>(&HeaderStorage);
-  new (Header) XRayFileHeader();
-  Header->Version = 3;
-  Header->Type = FileTypes::FDR_LOG;
-  Header->CycleFrequency = 3e9;
-  Header->ConstantTSC = 1;
-  Header->NonstopTSC = 1;
-  Serialized.append(reinterpret_cast<const char *>(&HeaderStorage),
-                    sizeof(XRayFileHeader));
-  size_t BufferCount = 0;
-  Buffers.apply([&](const BufferQueue::Buffer &B) {
-    ++BufferCount;
-    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);
-  });
-  ASSERT_EQ(BufferCount, 1u);
+  std::string Serialized = serialize(Buffers, 3);
+  llvm::DataExtractor DE(Serialized, true, 8);
+  auto TraceOrErr = llvm::xray::loadTrace(DE);
+  EXPECT_THAT_EXPECTED(
+      TraceOrErr,
+      HasValue(ElementsAre(
+          AllOf(FuncId(1), RecordType(llvm::xray::RecordTypes::ENTER)),
+          AllOf(FuncId(1), RecordType(llvm::xray::RecordTypes::EXIT)))));
+}
+
+TEST(FdrLogWriterTest, UnwriteRecords) {
+  bool Success = false;
+  BufferQueue Buffers(kSize, 1, Success);
+  BufferQueue::Buffer B;
+  ASSERT_EQ(Buffers.getBuffer(B), BufferQueue::ErrorCode::Ok);
+
+  FDRLogWriter Writer(B);
+  MetadataRecord Preamble[] = {
+      createMetadataRecord<MetadataRecord::RecordKinds::NewBuffer>(int32_t{1}),
+      createMetadataRecord<MetadataRecord::RecordKinds::WalltimeMarker>(
+          int64_t{1}, int32_t{2}),
+      createMetadataRecord<MetadataRecord::RecordKinds::Pid>(int32_t{1}),
+  };
+  ASSERT_THAT(Writer.writeMetadataRecords(Preamble),
+              Eq(sizeof(MetadataRecord) * 3));
+  ASSERT_TRUE(Writer.writeMetadata<MetadataRecord::RecordKinds::NewCPUId>(1));
+  ASSERT_TRUE(
+      Writer.writeFunction(FDRLogWriter::FunctionRecordKind::Enter, 1, 1));
+  ASSERT_TRUE(
+      Writer.writeFunction(FDRLogWriter::FunctionRecordKind::Exit, 1, 1));
+  Writer.undoWrites(sizeof(FunctionRecord) * 2);
+  ASSERT_EQ(Buffers.releaseBuffer(B), BufferQueue::ErrorCode::Ok);
+  ASSERT_EQ(B.Data, nullptr);
+  ASSERT_EQ(Buffers.finalize(), BufferQueue::ErrorCode::Ok);
 
+  // We've un-done the two function records we've written, and now we expect
+  // that we don't have any function records in the trace.
+  std::string Serialized = serialize(Buffers, 3);
   llvm::DataExtractor DE(Serialized, true, 8);
   auto TraceOrErr = llvm::xray::loadTrace(DE);
-  EXPECT_THAT_EXPECTED(TraceOrErr, HasValue(SizeIs(2)));
+  EXPECT_THAT_EXPECTED(TraceOrErr, HasValue(IsEmpty()));
 }
 
 } // namespace
diff --git a/lib/xray/tests/unit/test_helpers.cc b/lib/xray/tests/unit/test_helpers.cc
new file mode 100644
index 000000000..ebc40f88c
--- /dev/null
+++ b/lib/xray/tests/unit/test_helpers.cc
@@ -0,0 +1,91 @@
+//===-- test_helpers.cc ---------------------------------------------------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+// This file is a part of XRay, a function call tracing system.
+//
+//===----------------------------------------------------------------------===//
+#include "test_helpers.h"
+#include "xray/xray_records.h"
+#include "xray_buffer_queue.h"
+#include "xray_fdr_log_writer.h"
+#include <type_traits>
+
+// TODO: Move these to llvm/include/Testing/XRay/...
+namespace llvm {
+namespace xray {
+
+std::string RecordTypeAsString(RecordTypes T) {
+  switch (T) {
+  case RecordTypes::ENTER:
+    return "llvm::xray::RecordTypes::ENTER";
+  case RecordTypes::EXIT:
+    return "llvm::xray::RecordTypes::EXIT";
+  case RecordTypes::TAIL_EXIT:
+    return "llvm::xray::RecordTypes::TAIL_EXIT";
+  case RecordTypes::ENTER_ARG:
+    return "llvm::xray::RecordTypes::ENTER_ARG";
+  }
+  return "<UNKNOWN>";
+}
+
+void PrintTo(RecordTypes T, std::ostream *OS) {
+  *OS << RecordTypeAsString(T);
+}
+
+void PrintTo(const XRayRecord &R, std::ostream *OS) {
+  *OS << "XRayRecord { CPU = " << R.CPU
+      << "; Type = " << RecordTypeAsString(R.Type) << "; FuncId = " << R.FuncId
+      << "; TSC = " << R.TSC << "; TId = " << R.TId << "; PId = " << R.PId
+      << " Args = " << ::testing::PrintToString(R.CallArgs) << " }";
+}
+
+void PrintTo(const Trace &T, std::ostream *OS) {
+  const auto &H = T.getFileHeader();
+  *OS << "XRay Trace:\nHeader: { Version = " << H.Version
+      << "; Type = " << H.Type
+      << "; ConstantTSC = " << ::testing::PrintToString(H.ConstantTSC)
+      << "; NonstopTSC = " << ::testing::PrintToString(H.NonstopTSC)
+      << "; CycleFrequency = " << H.CycleFrequency << "; FreeFormData = '"
+      << ::testing::PrintToString(H.FreeFormData) << "' }\n";
+  for (const auto &R : T) {
+    PrintTo(R, OS);
+    *OS << "\n";
+  }
+}
+
+} // namespace xray
+} // namespace llvm
+
+namespace __xray {
+
+std::string serialize(BufferQueue &Buffers, int32_t Version) {
+  std::string Serialized;
+  std::aligned_storage<sizeof(XRayFileHeader), alignof(XRayFileHeader)>::type
+      HeaderStorage;
+  auto *Header = reinterpret_cast<XRayFileHeader *>(&HeaderStorage);
+  new (Header) XRayFileHeader();
+  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 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
diff --git a/lib/xray/tests/unit/test_helpers.h b/lib/xray/tests/unit/test_helpers.h
new file mode 100644
index 000000000..6eca7795e
--- /dev/null
+++ b/lib/xray/tests/unit/test_helpers.h
@@ -0,0 +1,59 @@
+//===-- test_helpers.h ----------------------------------------------------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+// This file is a part of XRay, a function call tracing system.
+//
+//===----------------------------------------------------------------------===//
+#ifndef COMPILER_RT_LIB_XRAY_TESTS_TEST_HELPERS_H_
+#define COMPILER_RT_LIB_XRAY_TESTS_TEST_HELPERS_H_
+
+#include "xray_buffer_queue.h"
+#include "llvm/XRay/XRayRecord.h"
+#include "llvm/XRay/Trace.h"
+#include "gmock/gmock.h"
+
+// TODO: Move these to llvm/include/Testing/XRay/...
+namespace llvm {
+namespace xray {
+
+std::string RecordTypeAsString(RecordTypes T);
+void PrintTo(RecordTypes T, std::ostream *OS);
+void PrintTo(const XRayRecord &R, std::ostream *OS);
+void PrintTo(const Trace &T, std::ostream *OS);
+
+namespace testing {
+
+MATCHER_P(FuncId, F, "") {
+  *result_listener << "where the function id is " << F;
+  return arg.FuncId == F;
+}
+
+MATCHER_P(RecordType, T, "") {
+  *result_listener << "where the record type is " << RecordTypeAsString(T);
+  return arg.Type == T;
+}
+
+MATCHER_P(HasArg, A, "") {
+  *result_listener << "where args contains " << A;
+  return !arg.CallArgs.empty() &&
+         std::any_of(arg.CallArgs.begin(), arg.CallArgs.end(),
+                     [this](decltype(A) V) { return V == A; });
+}
+
+} // namespace testing
+} // namespace xray
+} // namespace llvm
+
+namespace __xray {
+
+std::string serialize(BufferQueue &Buffers, int32_t Version);
+
+} // namespace __xray
+
+#endif // COMPILER_RT_LIB_XRAY_TESTS_TEST_HELPERS_H_
diff --git a/lib/xray/xray_fdr_controller.h b/lib/xray/xray_fdr_controller.h
new file mode 100644
index 000000000..210a32dc2
--- /dev/null
+++ b/lib/xray/xray_fdr_controller.h
@@ -0,0 +1,304 @@
+//===-- xray_fdr_controller.h ---------------------------------------------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+// This file is a part of XRay, a function call tracing system.
+//
+//===----------------------------------------------------------------------===//
+#ifndef COMPILER_RT_LIB_XRAY_XRAY_FDR_CONTROLLER_H_
+#define COMPILER_RT_LIB_XRAY_XRAY_FDR_CONTROLLER_H_
+
+#include <time.h>
+
+#include "xray/xray_interface.h"
+#include "xray/xray_records.h"
+#include "xray_buffer_queue.h"
+#include "xray_fdr_log_writer.h"
+
+namespace __xray {
+
+template <size_t Version = 3> class FDRController {
+  BufferQueue *BQ;
+  BufferQueue::Buffer &B;
+  FDRLogWriter &W;
+  int (*WallClockReader)(clockid_t, struct timespec *) = 0;
+  uint64_t CycleThreshold = 0;
+
+  uint64_t LastFunctionEntryTSC = 0;
+  uint64_t LatestTSC = 0;
+  uint16_t LatestCPU = 0;
+  tid_t TId = 0;
+  pid_t PId = 0;
+  bool First = true;
+
+  uint32_t UndoableFunctionEnters = 0;
+  uint32_t UndoableTailExits = 0;
+
+  bool finalized() const { return BQ == nullptr || BQ->finalizing(); }
+
+  bool hasSpace(size_t S) {
+    return B.Data != nullptr &&
+           W.getNextRecord() + S <= reinterpret_cast<char *>(B.Data) + B.Size;
+  }
+
+  constexpr int32_t mask(int32_t FuncId) const {
+    return FuncId & ((1 << 29) - 1);
+  }
+
+  bool getNewBuffer() {
+    if (!returnBuffer())
+      return false;
+    if (BQ->getBuffer(B) != BufferQueue::ErrorCode::Ok)
+      return false;
+
+    W.resetRecord();
+    DCHECK_EQ(W.getNextRecord(), B.Data);
+    LatestTSC = 0;
+    LatestCPU = 0;
+    atomic_store(&B.Extents, 0, memory_order_release);
+    return true;
+  }
+
+  bool setupNewBuffer() {
+    if (finalized())
+      return false;
+
+    DCHECK(hasSpace(sizeof(MetadataRecord) * 3));
+    TId = GetTid();
+    PId = internal_getpid();
+    struct timespec TS {
+      0, 0
+    };
+    WallClockReader(CLOCK_MONOTONIC, &TS);
+
+    MetadataRecord Metadata[] = {
+        // 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 (tid_t, cast to 4 bytes type due to Darwin being 8
+        //   bytes)
+        createMetadataRecord<MetadataRecord::RecordKinds::NewBuffer>(
+            static_cast<int32_t>(TId)),
+
+        // Also write the WalltimeMarker record. 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.
+        createMetadataRecord<MetadataRecord::RecordKinds::WalltimeMarker>(
+            static_cast<int64_t>(TS.tv_sec),
+            static_cast<int32_t>(TS.tv_nsec / 1000)),
+
+        // Also write the Pid record.
+        createMetadataRecord<MetadataRecord::RecordKinds::Pid>(
+            static_cast<int32_t>(PId)),
+    };
+
+    if (finalized())
+      return false;
+    return W.writeMetadataRecords(Metadata);
+  }
+
+  bool prepareBuffer(size_t S) {
+    if (finalized())
+      return returnBuffer();
+
+    if (UNLIKELY(!hasSpace(S))) {
+      if (!getNewBuffer())
+        return false;
+      if (!setupNewBuffer())
+        return false;
+    }
+
+    if (First) {
+      First = false;
+      W.resetRecord();
+      atomic_store(&B.Extents, 0, memory_order_release);
+      return setupNewBuffer();
+    }
+
+    return true;
+  }
+
+  bool returnBuffer() {
+    if (BQ == nullptr)
+      return false;
+
+    if (finalized()) {
+      BQ->releaseBuffer(B); // ignore result.
+      return false;
+    }
+
+    First = true;
+    if (BQ->releaseBuffer(B) != BufferQueue::ErrorCode::Ok)
+      return false;
+    return true;
+  }
+
+  enum class PreambleResult { NoChange, WroteMetadata };
+  PreambleResult functionPreamble(uint64_t TSC, uint16_t CPU) {
+    if (UNLIKELY(LatestCPU != CPU || LatestTSC == 0)) {
+      // We update our internal tracking state for the Latest TSC and CPU we've
+      // seen, then write out the appropriate metadata and function records.
+      LatestTSC = TSC;
+      LatestCPU = CPU;
+      W.writeMetadata<MetadataRecord::RecordKinds::NewCPUId>(CPU, TSC);
+      return PreambleResult::WroteMetadata;
+    }
+
+    if (UNLIKELY(LatestCPU == LatestCPU && LatestTSC > TSC)) {
+      // The TSC has wrapped around, from the last TSC we've seen.
+      LatestTSC = TSC;
+      W.writeMetadata<MetadataRecord::RecordKinds::TSCWrap>(TSC);
+      return PreambleResult::WroteMetadata;
+    }
+
+    return PreambleResult::NoChange;
+  }
+
+  void rewindRecords(int32_t FuncId, uint64_t TSC, uint16_t CPU) {
+    // Undo one enter record, because at this point we are either at the state
+    // of:
+    // - We are exiting a function that we recently entered.
+    // - We are exiting a function that was the result of a sequence of tail
+    //   exits, and we can check whether the tail exits can be re-wound.
+    //
+    FunctionRecord F;
+    W.undoWrites(sizeof(FunctionRecord));
+    internal_memcpy(&F, W.getNextRecord(), sizeof(FunctionRecord));
+
+    DCHECK(F.RecordKind ==
+               uint8_t(FunctionRecord::RecordKinds::FunctionEnter) &&
+           "Expected to find function entry recording when rewinding.");
+    DCHECK_EQ(F.FuncId, FuncId & ~(0x0F << 28));
+
+    LatestTSC -= F.TSCDelta;
+    if (--UndoableFunctionEnters != 0) {
+      LastFunctionEntryTSC -= F.TSCDelta;
+      return;
+    }
+
+    LastFunctionEntryTSC = 0;
+    auto RewindingTSC = LatestTSC;
+    auto RewindingRecordPtr = W.getNextRecord() - sizeof(FunctionRecord);
+    while (UndoableTailExits) {
+      internal_memcpy(&F, RewindingRecordPtr, sizeof(FunctionRecord));
+      DCHECK_EQ(F.RecordKind,
+                uint8_t(FunctionRecord::RecordKinds::FunctionTailExit));
+      RewindingTSC -= F.TSCDelta;
+      RewindingRecordPtr -= sizeof(FunctionRecord);
+      internal_memcpy(&F, RewindingRecordPtr, sizeof(FunctionRecord));
+
+      // This tail call exceeded the threshold duration. It will not be erased.
+      if ((TSC - RewindingTSC) >= CycleThreshold) {
+        UndoableTailExits = 0;
+        return;
+      }
+
+      --UndoableTailExits;
+      W.undoWrites(sizeof(FunctionRecord) * 2);
+      LatestTSC = RewindingTSC;
+    }
+  }
+
+public:
+  template <class WallClockFunc>
+  FDRController(BufferQueue *BQ, BufferQueue::Buffer &B, FDRLogWriter &W,
+                WallClockFunc R, uint64_t C)
+      : BQ(BQ), B(B), W(W), WallClockReader(R), CycleThreshold(C) {}
+
+  bool functionEnter(int32_t FuncId, uint64_t TSC, uint16_t CPU) {
+    if (finalized())
+      return returnBuffer();
+
+    if (!prepareBuffer(sizeof(MetadataRecord) + sizeof(FunctionRecord)))
+      return returnBuffer();
+
+    if (functionPreamble(TSC, CPU) == PreambleResult::WroteMetadata) {
+      UndoableFunctionEnters = 1;
+    } else {
+      ++UndoableFunctionEnters;
+    }
+
+    LastFunctionEntryTSC = TSC;
+    LatestTSC = TSC;
+    return W.writeFunction(FDRLogWriter::FunctionRecordKind::Enter,
+                           mask(FuncId), TSC - LatestTSC);
+  }
+
+  bool functionTailExit(int32_t FuncId, uint64_t TSC, uint16_t CPU) {
+    if (finalized())
+      return returnBuffer();
+
+    if (!prepareBuffer(sizeof(MetadataRecord) + sizeof(FunctionRecord)))
+      return returnBuffer();
+
+    if (functionPreamble(TSC, CPU) == PreambleResult::NoChange &&
+        UndoableFunctionEnters != 0 &&
+        TSC - LastFunctionEntryTSC < CycleThreshold) {
+      rewindRecords(FuncId, TSC, CPU);
+      return true;
+    }
+
+    UndoableTailExits = UndoableFunctionEnters ? UndoableTailExits + 1 : 0;
+    UndoableFunctionEnters = 0;
+    LatestTSC = TSC;
+    return W.writeFunction(FDRLogWriter::FunctionRecordKind::TailExit,
+                           mask(FuncId), TSC - LatestTSC);
+  }
+
+  bool functionEnterArg(int32_t FuncId, uint64_t TSC, uint16_t CPU,
+                        uint64_t Arg) {
+    if (finalized())
+      return returnBuffer();
+
+    if (!prepareBuffer((2 * sizeof(MetadataRecord)) + sizeof(FunctionRecord)))
+      return returnBuffer();
+
+    // Ignore the result of writing out the preamble.
+    functionPreamble(TSC, CPU);
+
+    LatestTSC = TSC;
+    LastFunctionEntryTSC = 0;
+    UndoableFunctionEnters = 0;
+    UndoableTailExits = 0;
+
+    W.writeFunction(FDRLogWriter::FunctionRecordKind::EnterArg, mask(FuncId),
+                    TSC - LatestTSC);
+    return W.writeMetadata<MetadataRecord::RecordKinds::CallArgument>(Arg);
+  }
+
+  bool functionExit(int32_t FuncId, uint64_t TSC, uint16_t CPU) {
+    if (finalized())
+      return returnBuffer();
+
+    if (functionPreamble(TSC, CPU) == PreambleResult::NoChange &&
+        UndoableFunctionEnters != 0 &&
+        TSC - LastFunctionEntryTSC < CycleThreshold) {
+      rewindRecords(FuncId, TSC, CPU);
+      return true;
+    }
+
+    LatestTSC = TSC;
+    UndoableFunctionEnters = 0;
+    UndoableTailExits = 0;
+    return W.writeFunction(FDRLogWriter::FunctionRecordKind::Exit, mask(FuncId),
+                           TSC - LatestTSC);
+  }
+
+  bool flush() {
+    if (finalized()) {
+      returnBuffer(); // ignore result.
+      return true;
+    }
+    return returnBuffer();
+  }
+};
+
+} // namespace __xray
+
+#endif // COMPILER-RT_LIB_XRAY_XRAY_FDR_CONTROLLER_H_
diff --git a/lib/xray/xray_fdr_log_writer.h b/lib/xray/xray_fdr_log_writer.h
index 28af356e7..06b7239ec 100644
--- a/lib/xray/xray_fdr_log_writer.h
+++ b/lib/xray/xray_fdr_log_writer.h
@@ -112,6 +112,17 @@ public:
 
   char *getNextRecord() const { return NextRecord; }
 
+  void resetRecord() {
+    NextRecord = reinterpret_cast<char *>(Buffer.Data);
+    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);
+  }
+
 }; // namespace __xray
 
 } // namespace __xray