From 02306bf80eb0d2a62b19c4bb7f435fdadd226537 Mon Sep 17 00:00:00 2001
From: Jon Haslam <jonhaslam@meta.com>
Date: Wed, 15 Nov 2023 16:13:54 +0000
Subject: [PATCH] Make pointer tracking array facility fail gracefully when it
 is full (#411)

---
 oi/FuncGen.cpp     |  8 ++++++++
 oi/OIDebugger.cpp  |  8 ++++++++
 oi/OIDebugger.h    |  2 ++
 oi/OITraceCode.cpp | 37 ++++++++++++++++++++++++++++++-------
 4 files changed, 48 insertions(+), 7 deletions(-)

diff --git a/oi/FuncGen.cpp b/oi/FuncGen.cpp
index 7a1dd09..14eeb61 100644
--- a/oi/FuncGen.cpp
+++ b/oi/FuncGen.cpp
@@ -327,6 +327,8 @@ void FuncGen::DefineTopLevelGetSizeRef(std::string& testCode,
       uintptr_t& writtenSize = data[dataSegOffset++];
       writtenSize = 0;
       uintptr_t& timeTakenNs = data[dataSegOffset++];
+      size_t& pointersSize = data[dataSegOffset++];
+      size_t& pointersCapacity = data[dataSegOffset++];
 
       dataSegOffset *= sizeof(uintptr_t);
       JLOG("%1% @");
@@ -336,6 +338,8 @@ void FuncGen::DefineTopLevelGetSizeRef(std::string& testCode,
       OIInternal::StoreData((uintptr_t)123456789, dataSegOffset);
       writtenSize = dataSegOffset;
       dataBase += dataSegOffset;
+      pointersSize = pointers.size();
+      pointersCapacity = pointers.capacity();
     )";
   if (features[Feature::JitTiming]) {
     func += R"(
@@ -419,6 +423,8 @@ void FuncGen::DefineTopLevelGetSizeSmartPtr(std::string& testCode,
       uintptr_t& writtenSize = data[dataSegOffset++];
       writtenSize = 0;
       uintptr_t& timeTakenNs = data[dataSegOffset++];
+      size_t& pointersSize = data[dataSegOffset++];
+      size_t& pointersCapacity = data[dataSegOffset++];
 
       dataSegOffset *= sizeof(uintptr_t);
 
@@ -427,6 +433,8 @@ void FuncGen::DefineTopLevelGetSizeSmartPtr(std::string& testCode,
       OIInternal::StoreData((uintptr_t)123456789, dataSegOffset);
       writtenSize = dataSegOffset;
       dataBase += dataSegOffset;
+      pointersSize = pointers.size();
+      pointersCapacity = pointers.capacity();
     )";
   if (features[Feature::JitTiming]) {
     func += R"(
diff --git a/oi/OIDebugger.cpp b/oi/OIDebugger.cpp
index 9c98fc6..8edc609 100644
--- a/oi/OIDebugger.cpp
+++ b/oi/OIDebugger.cpp
@@ -2746,6 +2746,14 @@ bool OIDebugger::decodeTargetData(const DataHeader& dataHeader,
     LOG(INFO) << "JIT Timing: " << dataHeader.timeTakenNs << "ns";
   }
 
+  VLOG(1) << "Pointer tracking stats: " << dataHeader.pointersCapacity << "/"
+          << dataHeader.pointersSize;
+
+  if (dataHeader.pointersCapacity == dataHeader.pointersSize) {
+    VLOG(1) << "Pointer tracking array is exhausted! Results may be"
+               " partial.";
+  }
+
   /*
    * Currently  we use MAX_INT to indicate two things:
    *  - a single MAX_INT indicates the end of results for  the current object
diff --git a/oi/OIDebugger.h b/oi/OIDebugger.h
index 804d2cd..3997616 100644
--- a/oi/OIDebugger.h
+++ b/oi/OIDebugger.h
@@ -276,6 +276,8 @@ class OIDebugger {
     uintptr_t cookie;
     uintptr_t size;
     uintptr_t timeTakenNs;
+    size_t pointersSize;
+    size_t pointersCapacity;
 
     /*
      * Flexible Array Member are not standard in C++, but this is
diff --git a/oi/OITraceCode.cpp b/oi/OITraceCode.cpp
index fa0beda..5cc443d 100644
--- a/oi/OITraceCode.cpp
+++ b/oi/OITraceCode.cpp
@@ -44,9 +44,12 @@ class {
  private:
   // 1 MiB of pointers
   std::array<uintptr_t, (1 << 20) / sizeof(uintptr_t)> data;
+  size_t numEntries;
 
-  // twang_mix64 hash function, taken from Folly where it is used
-  // as the default hash function for 64-bit integers
+  /*
+   * twang_mix64 hash function, taken from Folly where it is used as the
+   * default hash function for 64-bit integers.
+   */
   constexpr static uint64_t twang_mix64(uint64_t key) noexcept {
     key = (~key) + (key << 21);  // key *= (1 << 21) - 1; key -= 1;
     key = key ^ (key >> 24);
@@ -61,26 +64,46 @@ class {
  public:
   void initialize() noexcept {
     data.fill(0);
+    numEntries = 0;
   }
 
-  // Adds the pointer to the set.
-  // Returns `true` if the value was newly added,
-  // or `false` if the value was already present.
+  /*
+   * Adds the pointer to the set.
+   * Returns `true` if the value was newly added. `false` may be returned if
+   * the value was already present, a null pointer was passed or if there are
+   * no entries left in the array.
+   */
   bool add(uintptr_t pointer) noexcept {
-    __builtin_assume(pointer > 0);
+    if (pointer == 0) {
+      return false;
+    }
+
     uint64_t index = twang_mix64(pointer) % data.size();
     while (true) {
       uintptr_t entry = data[index];
+
       if (entry == 0) {
         data[index] = pointer;
+        ++numEntries;
         return true;
       }
-      if (entry == pointer) {
+
+      if (entry == pointer || numEntries >= data.size()) {
         return false;
       }
+
       index = (index + 1) % data.size();
     }
   }
+
+  size_t size(void) {
+    return numEntries;
+  }
+
+  size_t capacity(void) {
+    return data.size();
+  }
+
   bool add(const auto* p) {
     return add((uintptr_t)p);
   }