Add test that exercises "too large removal buffer" code

ConflictSet.cpp

@@ -5347,6 +5347,11 @@ struct __attribute__((visibility("hidden"))) ConflictSet::Impl {
     gc_iterations_total.add(set_oldest_iterations_accum);
     if (n == nullptr) {
       removalKey = {};
+      if (removalBufferSize > kMaxRemovalBufferSize) {
+        safe_free(removalBuffer, removalBufferSize);
+        removalBufferSize = kMinRemovalBufferSize;
+        removalBuffer = (uint8_t *)safe_malloc(removalBufferSize);
+      }
       oldestExtantVersion = oldestVersionAtGcBegin;
       oldest_extant_version.set(oldestExtantVersion);
       oldestVersionAtGcBegin = oldestVersionFullPrecision;
@@ -5357,12 +5362,47 @@ struct __attribute__((visibility("hidden"))) ConflictSet::Impl {
               oldestExtantVersion, oldestVersionAtGcBegin);
 #endif
     } else {
-      removalKeyArena = Arena();
-      removalKey = getSearchPath(removalKeyArena, n);
+      // Store the current search path to resume the scan later
+      saveRemovalKey(n);
     }
     return fuel;
   }
 
+  void saveRemovalKey(Node *n) {
+    uint8_t *cursor = removalBuffer + removalBufferSize;
+    int size = 0;
+    auto reserve = [&](int delta) {
+      if (size + delta > removalBufferSize) [[unlikely]] {
+        int newBufSize = std::max(removalBufferSize * 2, size + delta);
+        uint8_t *newBuf = (uint8_t *)safe_malloc(newBufSize);
+        memcpy(newBuf + newBufSize - size, cursor, size);
+        safe_free(removalBuffer, removalBufferSize);
+        removalBuffer = newBuf;
+        removalBufferSize = newBufSize;
+        cursor = newBuf + newBufSize - size;
+      }
+    };
+    for (;;) {
+      auto partialKey = TrivialSpan{n->partialKey(), n->partialKeyLen};
+      reserve(partialKey.size());
+      size += partialKey.size();
+      cursor -= partialKey.size();
+      memcpy(cursor, partialKey.data(), partialKey.size());
+
+      if (n->parent == nullptr) {
+        break;
+      }
+
+      reserve(1);
+      ++size;
+      --cursor;
+      *cursor = n->parentsIndex;
+
+      n = n->parent;
+    }
+    removalKey = {cursor, size};
+  }
+
   void setOldestVersion(int64_t newOldestVersion) {
     assert(newOldestVersion >= 0);
     assert(newOldestVersion <= newestVersionFullPrecision);
@@ -5413,7 +5453,7 @@ struct __attribute__((visibility("hidden"))) ConflictSet::Impl {
     writeContext.~WriteContext();
     new (&writeContext) WriteContext();
 
-    removalKeyArena = Arena{};
+    // Leave removalBuffer as is
     removalKey = {};
     keyUpdates = 10;
 
@@ -5445,11 +5485,16 @@ struct __attribute__((visibility("hidden"))) ConflictSet::Impl {
   ~Impl() {
     eraseTree(root, &writeContext);
     safe_free(metrics, metricsCount * sizeof(metrics[0]));
+    safe_free(removalBuffer, removalBufferSize);
   }
 
   WriteContext writeContext;
 
-  Arena removalKeyArena;
+  static constexpr int kMinRemovalBufferSize = 1 << 10;
+  // Eventually downsize if larger than this value
+  static constexpr int kMaxRemovalBufferSize = 1 << 16;
+  uint8_t *removalBuffer = (uint8_t *)safe_malloc(kMinRemovalBufferSize);
+  int removalBufferSize = kMinRemovalBufferSize;
   TrivialSpan removalKey;
   int64_t keyUpdates;
 

test_conflict_set.py

@@ -1,3 +1,4 @@
+import struct
 from conflict_set import *
 
 
@@ -164,6 +165,16 @@ def test_fixup_256():
         cs.check(read(0, bytes([1]), bytes([2])))
 
 
+def test_large_removal_buffer():
+    with DebugConflictSet() as cs:
+        for i in range(1000):
+            # create extra gc work
+            for j in range(100):
+                cs.addWrites(1000 + i)
+            cs.addWrites(1000 + i, write(struct.pack(">l", i) + bytes([0] * 100000)))
+            cs.setOldestVersion(i)
+
+
 if __name__ == "__main__":
     # budget "pytest" for ctest integration without pulling in a dependency. You can of course still use pytest in local development.
     import argparse