[XRay] Support generational buffers in FDR controller

[dberris]

Summary

This is an intermediary step in the full support for generational buffer
management in the FDR runtime. This change makes the FDR controller
aware of the new generation number in the buffers handed out by the
BufferQueue type.

In the process of making this change, we've realised that the cleanest
way of ensuring that the backing store per generation is live while all
the threads that need access to it will need reference counting to tie
the backing store to the lifetime of all threads that have a handle on
buffers associated with the memory.

We also learn that we're missing the edge-case in the function exit
handler's implementation where the first record being written into the
buffer is a function exit, which is caught/fixed by the test for
generational buffer management.

We still haven't wired the controller into the FDR mode runtime, which
will need the reference counting on the backing store implemented to
ensure that we're being conservatively thread-safe with this approach.

Depends on D52974.

[dberris]

Re-title, adding subscribers and reviewers.

[mboerger]

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

253 ↗

(On Diff #170564)

s/we'll/we will/

[dberris]

Update comment in test.

[dberris]

Add dependency on D53551 in description.

[dberris]

Revert inadvertent inclusion of dependenent change.

[mboerger]

LGTM

compiler-rt/trunk/lib/xray/tests/unit/fdr_controller_test.cc

protected:
   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 +
+    BQ = llvm::make_unique<BufferQueue>(sizeof(MetadataRecord) * 5 +
                                             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) {
+  for (size_t I = 0; I < kBuffers + 1; ++I) {
     ASSERT_TRUE(C->functionEnter(1, TSC++, CPU));
     ASSERT_TRUE(C->functionExit(1, TSC++, CPU));
   }
-  C->flush();
-  ASSERT_EQ(BQ->finalize(), BufferQueue::ErrorCode::Ok);
+  ASSERT_TRUE(C->flush());
+  ASSERT_THAT(BQ->finalize(), Eq(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) {
TEST_F(BufferManagementTest, HandlesFinalizedBufferQueue) {
   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)))));
 }
 
+TEST_F(BufferManagementTest, HandlesGenerationalBufferQueue) {
+  uint64_t TSC = 1;
+  uint16_t CPU = 1;
+
+  ASSERT_TRUE(C->functionEnter(1, TSC++, CPU));
+  ASSERT_THAT(BQ->finalize(), Eq(BufferQueue::ErrorCode::Ok));
+  ASSERT_THAT(BQ->init(sizeof(MetadataRecord) * 4 + sizeof(FunctionRecord) * 2,
+                       kBuffers),
+              Eq(BufferQueue::ErrorCode::Ok));
+  EXPECT_TRUE(C->functionExit(1, TSC++, CPU));
+  ASSERT_TRUE(C->flush());
+
+  // We expect that we will only be able to find the function exit event, but
+  // not the function enter event, since we only have information about the new
+  // generation of the buffers.
+  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::EXIT)))));
+}
+
 } // namespace
 } // namespace __xray

compiler-rt/trunk/lib/xray/xray_fdr_controller.h

template <size_t Version = 3> class FDRController {
   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 &&
+    return B.Data != nullptr && B.Generation == BQ->generation() &&
            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;
+    First = true;
+    UndoableFunctionEnters = 0;
+    UndoableTailExits = 0;
     atomic_store(&B.Extents, 0, memory_order_release);
     return true;
   }
 
   bool setupNewBuffer() {
     if (finalized())
       return false;
 
bool setupNewBuffer() {
     return W.writeMetadataRecords(Metadata);
   }
 
   bool prepareBuffer(size_t S) {
     if (finalized())
       return returnBuffer();
 
     if (UNLIKELY(!hasSpace(S))) {
+      if (!returnBuffer())
+        return false;
       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;
 
+    First = true;
     if (finalized()) {
       BQ->releaseBuffer(B); // ignore result.
       return false;
     }
 
-    First = true;
-    if (BQ->releaseBuffer(B) != BufferQueue::ErrorCode::Ok)
-      return false;
-    return true;
+    return BQ->releaseBuffer(B) == BufferQueue::ErrorCode::Ok;
   }
 
-  enum class PreambleResult { NoChange, WroteMetadata };
+  enum class PreambleResult { NoChange, WroteMetadata, InvalidBuffer };
   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;
+
+      if (B.Generation != BQ->generation())
+        return PreambleResult::InvalidBuffer;
+
       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;
+
+      if (B.Generation != BQ->generation())
+        return PreambleResult::InvalidBuffer;
+
       W.writeMetadata<MetadataRecord::RecordKinds::TSCWrap>(TSC);
       return PreambleResult::WroteMetadata;
     }
 
     return PreambleResult::NoChange;
   }
 
-  void rewindRecords(int32_t FuncId, uint64_t TSC, uint16_t CPU) {
+  bool 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));
+    if (B.Generation != BQ->generation())
+      return false;
     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;
+      return true;
     }
 
     LastFunctionEntryTSC = 0;
     auto RewindingTSC = LatestTSC;
     auto RewindingRecordPtr = W.getNextRecord() - sizeof(FunctionRecord);
     while (UndoableTailExits) {
+      if (B.Generation != BQ->generation())
+        return false;
       internal_memcpy(&F, RewindingRecordPtr, sizeof(FunctionRecord));
       DCHECK_EQ(F.RecordKind,
                 uint8_t(FunctionRecord::RecordKinds::FunctionTailExit));
       RewindingTSC -= F.TSCDelta;
       RewindingRecordPtr -= sizeof(FunctionRecord);
+      if (B.Generation != BQ->generation())
+        return false;
       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;
+        return true;
       }
 
       --UndoableTailExits;
       W.undoWrites(sizeof(FunctionRecord) * 2);
       LatestTSC = RewindingTSC;
     }
+    return true;
   }
 
 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())
+    if (finalized() ||
+        !prepareBuffer(sizeof(MetadataRecord) + sizeof(FunctionRecord)))
       return returnBuffer();
 
-    if (!prepareBuffer(sizeof(MetadataRecord) + sizeof(FunctionRecord)))
+    auto PreambleStatus = functionPreamble(TSC, CPU);
+    if (PreambleStatus == PreambleResult::InvalidBuffer)
       return returnBuffer();
 
-    if (functionPreamble(TSC, CPU) == PreambleResult::WroteMetadata) {
-      UndoableFunctionEnters = 1;
-    } else {
-      ++UndoableFunctionEnters;
-    }
-
+    UndoableFunctionEnters = (PreambleStatus == PreambleResult::WroteMetadata)
+                                 ? 1
+                                 : UndoableFunctionEnters + 1;
     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 &&
+    auto PreambleStatus = functionPreamble(TSC, CPU);
+    if (PreambleStatus == PreambleResult::InvalidBuffer)
+      return returnBuffer();
+
+    if (PreambleStatus == PreambleResult::NoChange &&
         UndoableFunctionEnters != 0 &&
-        TSC - LastFunctionEntryTSC < CycleThreshold) {
-      rewindRecords(FuncId, TSC, CPU);
-      return true;
-    }
+        TSC - LastFunctionEntryTSC < CycleThreshold)
+      return rewindRecords(FuncId, TSC, CPU);
 
     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)))
+    if (finalized() ||
+        !prepareBuffer((2 * sizeof(MetadataRecord)) + sizeof(FunctionRecord)) ||
+        functionPreamble(TSC, CPU) == PreambleResult::InvalidBuffer)
       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())
+    if (finalized() ||
+        !prepareBuffer(sizeof(MetadataRecord) + sizeof(FunctionRecord)))
       return returnBuffer();
 
-    if (functionPreamble(TSC, CPU) == PreambleResult::NoChange &&
+    auto PreambleStatus = functionPreamble(TSC, CPU);
+    if (PreambleStatus == PreambleResult::InvalidBuffer)
+      return returnBuffer();
+
+    if (PreambleStatus == PreambleResult::NoChange &&
         UndoableFunctionEnters != 0 &&
-        TSC - LastFunctionEntryTSC < CycleThreshold) {
-      rewindRecords(FuncId, TSC, CPU);
-      return true;
-    }
+        TSC - LastFunctionEntryTSC < CycleThreshold)
+      return rewindRecords(FuncId, TSC, CPU);
 
     LatestTSC = TSC;
     UndoableFunctionEnters = 0;
     UndoableTailExits = 0;
     return W.writeFunction(FDRLogWriter::FunctionRecordKind::Exit, mask(FuncId),
                            TSC - LatestTSC);
   }