Lazy zero segment allocation strategy

c++/src/capnp/arena.c++

@@ -62,6 +62,23 @@ void SegmentBuilder::throwNotWritable() {
       "referenced data, only Readers, because that data is const.");
 }
 
+void SegmentBuilder::doLazyZeroSegment(word* start, size_t words, schema::Type::Which type) {
+  // Get the current arena and lazyZeroSegmentAlloc options.
+  const BuilderArena* arena = getArena();
+  if (!arena) return;
+
+  const auto& lazyZero = arena->getLazyZeroSegmentAlloc();
+
+  // Skip if lazy zero segment alloc is not enabled.
+  if (! lazyZero.enableLazyZero) return;
+
+  // Skip zeroing for types or fields that are configured to be skipped.
+  if (lazyZero.skipLazyZeroTypes.find(type) != lazyZero.skipLazyZeroTypes.end()) return;
+
+  // Perform memset for the remaining memory that requires zeroing.
+  if (words > 0 && start) memset(start, 0, words * sizeof(word));
+}
+
 // =======================================================================================
 
 static SegmentWordCount verifySegmentSize(size_t size) {
@@ -235,7 +252,9 @@ BuilderArena::AllocateResult BuilderArena::allocate(SegmentWordCount amount) {
     kj::ctor(segment0, this, SegmentId(0), ptr.begin(), actualSize, &this->dummyLimiter);
 
     segmentWithSpace = &segment0;
-    return AllocateResult { &segment0, segment0.allocate(amount) };
+    word* wordPtr = segment0.allocate(amount);
+    segment0.doLazyZeroSegment(wordPtr, static_cast<size_t>(POINTER_SIZE_IN_WORDS));
+    return AllocateResult { &segment0, wordPtr };
   } else {
     if (segmentWithSpace != nullptr) {
       // Check if there is space in an existing segment.
@@ -257,8 +276,11 @@ BuilderArena::AllocateResult BuilderArena::allocate(SegmentWordCount amount) {
     // Check this new segment first the next time we need to allocate.
     segmentWithSpace = result;
 
+    word* wordPtr = result->allocate(amount);
+    result->doLazyZeroSegment(wordPtr, static_cast<size_t>(POINTER_SIZE_IN_WORDS));
+
     // Allocating from the new segment is guaranteed to succeed since we made it big enough.
-    return AllocateResult { result, result->allocate(amount) };
+    return AllocateResult { result, wordPtr };
   }
 }
 

c++/src/capnp/arena.h

@@ -33,6 +33,7 @@
 #include "common.h"
 #include "message.h"
 #include "layout.h"
+#include "schema.h"
 #include <kj/map.h>
 
 #if !CAPNP_LITE
@@ -217,6 +218,11 @@ class SegmentBuilder: public SegmentReader {
   // boundaries, then move the end up to `to` and return true. Otherwise, do nothing and return
   // false.
 
+  void doLazyZeroSegment(word* start, size_t words, schema::Type::Which type = schema::Type::ANY_POINTER);
+  // Ensures that the word range [start, start + words) is zeroed lazily.
+  // The zeroing is skipped if LazyZeroSegmentAlloc is not enabled or `type` is listed in skipLazyZeroTypes.
+  // Otherwise, the range is guaranteed to be zeroed.
+
 private:
   word* pos;
   // Pointer to a pointer to the current end point of the segment, i.e. the location where the
@@ -343,6 +349,8 @@ class BuilderArena final: public Arena {
   SegmentReader* tryGetSegment(SegmentId id) override;
   void reportReadLimitReached() override;
 
+  inline const BuilderOptions::LazyZeroSegmentAlloc& getLazyZeroSegmentAlloc() const;
+
 private:
   MessageBuilder* message;
   ReadLimiter dummyLimiter;
@@ -517,6 +525,12 @@ inline bool SegmentBuilder::tryExtend(word* from, word* to) {
   }
 }
 
+static const BuilderOptions::LazyZeroSegmentAlloc defaultAlloc;
+inline const BuilderOptions::LazyZeroSegmentAlloc& BuilderArena::getLazyZeroSegmentAlloc() const {
+  if (message == nullptr) return defaultAlloc;
+  return message->getOptions().lazyZeroSegmentAlloc;
+}
+
 }  // namespace _ (private)
 }  // namespace capnp
 

c++/src/capnp/layout.c++

@@ -458,7 +458,8 @@ struct WireHelpers {
 
   static KJ_ALWAYS_INLINE(word* allocate(
       WirePointer*& ref, SegmentBuilder*& segment, CapTableBuilder* capTable,
-      SegmentWordCount amount, WirePointer::Kind kind, BuilderArena* orphanArena)) {
+      SegmentWordCount amount, WirePointer::Kind kind, BuilderArena* orphanArena,
+      schema::Type::Which type = schema::Type::ANY_POINTER)) {
     // Allocate space in the message for a new object, creating far pointers if necessary. The
     // space is guaranteed to be zero'd (because MessageBuilder implementations are required to
     // return zero'd memory).
@@ -480,6 +481,8 @@ struct WireHelpers {
     //   In this case, `segment` starts out null; the allocation takes place in an arbitrary
     //   segment belonging to the arena.  `ref` will be initialized as a non-far pointer, but its
     //   target offset will be set to zero.
+    // * `type` is an optional hint used for lazily zeroing memory. It is used together with
+    //  LazyZeroSegmentAlloc to skip or do lazy zero for certain schema types.
 
     if (orphanArena == nullptr) {
       if (!ref->isNull()) zeroObject(segment, capTable, ref);
@@ -513,11 +516,21 @@ struct WireHelpers {
 
         // Initialize the landing pad to indicate that the data immediately follows the pad.
         ref = reinterpret_cast<WirePointer*>(ptr);
+
+        // Help lazy zero out segment space for the newly allocated segment if LazyZeroSegmentAlloc is enabled.
+        // Zero the landing pad for a far pointer.
+        segment->doLazyZeroSegment(ptr, static_cast<size_t>(POINTER_SIZE_IN_WORDS));
+        // Zero the actual object segment space following the landing pad (may skip based on configuration).
+        segment->doLazyZeroSegment(ptr + POINTER_SIZE_IN_WORDS, static_cast<size_t>(amount), /*type=*/type);
+
         ref->setKindAndTarget(kind, ptr + POINTER_SIZE_IN_WORDS, segment);
 
         // Allocated space follows new pointer.
         return ptr + POINTER_SIZE_IN_WORDS;
       } else {
+        // Help lazy zero the object's segment space if LazyZeroSegmentAlloc is enabled, ensures the
+        // object memory is clean before initialization (may skip based on configuration).
+        segment->doLazyZeroSegment(ptr, static_cast<size_t>(amount), /*type=*/type);
         ref->setKindAndTarget(kind, ptr, segment);
         return ptr;
       }
@@ -1708,7 +1721,7 @@ struct WireHelpers {
       BuilderArena* orphanArena = nullptr)) {
     // Allocate the space.
     word* ptr = allocate(ref, segment, capTable, roundBytesUpToWords(size),
-                         WirePointer::LIST, orphanArena);
+                         WirePointer::LIST, orphanArena, /*type=*/schema::Type::DATA);
 
     // Initialize the pointer.
     ref->listRef.set(ElementSize::BYTE, size * (ONE * ELEMENTS / BYTES));

c++/src/capnp/message-test.c++

@@ -19,6 +19,8 @@
 // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 // THE SOFTWARE.
 
+#include <cstdlib>
+#include <ctime>
 #include "message.h"
 #include "test-util.h"
 #include <kj/array.h>
@@ -213,6 +215,246 @@ KJ_TEST("MessageBuilder::sizeInWords()") {
   KJ_EXPECT(reader.sizeInWords() == expected);
 }
 
+class MyCustomMessageBuilder : public MessageBuilder {
+public:
+  MyCustomMessageBuilder(BuilderOptions options): MessageBuilder(kj::mv(options)) {
+    std::srand(static_cast<unsigned int>(std::time(nullptr)));
+  }
+
+  kj::ArrayPtr<word> allocateSegment(uint minimumSize) override {
+    auto array = kj::heapArray<word>(minimumSize);
+    auto bytes = array.asBytes();
+    // mock dirty memory
+    for (size_t i = 0; i < bytes.size(); ++i) {
+      bytes[i] = static_cast<uint8_t>(std::rand() % 256);
+    }
+    allocations.add(kj::mv(array));
+    return allocations.back();
+  }
+
+  kj::Vector<kj::Array<word>> allocations;
+};
+
+TEST(Message, LazyZeroCustomBuilder_DataDirty_OthersZero) {
+  // Enable lazy-zero and skip zeroing for DATA type.
+  BuilderOptions options;
+  options.lazyZeroSegmentAlloc.enableLazyZero = true;
+  options.lazyZeroSegmentAlloc.skipLazyZeroTypes.insert(schema::Type::DATA);
+
+  // Use the custom allocator that returns "dirty" memory.
+  MyCustomMessageBuilder builder(options);
+
+  // Init root
+  auto root = builder.initRoot<TestAllTypes>();
+
+  // Allocate a DATA field (size 64) — because we skipped zeroing for DATA,
+  // the returned buffer should contain the allocator's random bytes.
+  auto dataBuf = root.initDataField(64);
+
+  // Check DATA is not all zero (i.e., "dirty")
+  bool dataAllZero = true;
+  for (size_t i = 0; i < dataBuf.size(); ++i) {
+    if (dataBuf[i] != 0) { dataAllZero = false; break; }
+  }
+  EXPECT_FALSE(dataAllZero);
+
+  // Check other primitive/pointer fields are zero / not present by default.
+  // (These names follow the TestAllTypes generated accessors used elsewhere in tests.)
+  EXPECT_EQ(0u, root.getUInt32Field());
+  EXPECT_EQ(0, root.getInt64Field());
+  EXPECT_EQ(0.0, root.getFloat64Field());
+  EXPECT_FALSE(root.hasTextField());
+  EXPECT_FALSE(root.hasStructField()); // if there is a nested struct field, it should be null
+
+  // Also sanity-check that writing/reading doesn't crash: getSegmentsForOutput() is callable.
+  auto segs = builder.getSegmentsForOutput();
+  EXPECT_GE(segs.size(), 1u);
+}
+
+TEST(Message, LazyZeroCustomBuilder_DataWriteAndReadback_Persists) {
+  // Setup builder with lazy-zero skip for DATA.
+  BuilderOptions options;
+  options.lazyZeroSegmentAlloc.enableLazyZero = true;
+  options.lazyZeroSegmentAlloc.skipLazyZeroTypes.insert(schema::Type::DATA);
+
+  MyCustomMessageBuilder builder(options);
+  auto root = builder.initRoot<TestAllTypes>();
+
+  // Fill DATA with a recognizable pattern.
+  const size_t N = 64;
+  auto data = root.initDataField(N);
+  for (size_t i = 0; i < N; ++i) data[i] = static_cast<capnp::byte>(i & 0xFF);
+
+  // Read back by creating a SegmentArrayMessageReader from the builder segments.
+  auto segs = builder.getSegmentsForOutput();
+  capnp::SegmentArrayMessageReader readerFromSegments(segs);
+  auto readBack = readerFromSegments.getRoot<TestAllTypes>();
+  auto readData = readBack.getDataField();
+
+  ASSERT_EQ(readData.size(), N);
+  for (size_t i = 0; i < N; ++i) {
+    EXPECT_EQ(readData[i], static_cast<capnp::byte>(i & 0xFF));
+  }
+}
+
+TEST(Message, LazyZeroCustomBuilder_ClonePreservesDirtyData_ViaSegments) {
+  // Use lazy-zero skipping for DATA to keep allocator's dirty bytes.
+  BuilderOptions options;
+  options.lazyZeroSegmentAlloc.enableLazyZero = true;
+  options.lazyZeroSegmentAlloc.skipLazyZeroTypes.insert(schema::Type::DATA);
+
+  MyCustomMessageBuilder builder(options);
+  auto root = builder.initRoot<TestAllTypes>();
+
+  // Allocate DATA and capture original bytes.
+  const size_t SIZE = 100;
+  auto data = root.initDataField(SIZE);
+  kj::Vector<capnp::byte> original;
+  original.reserve(SIZE);
+  for (size_t i = 0; i < SIZE; ++i) {
+    original.add(data[i]);
+  }
+
+  // Export segments and re-read using SegmentArrayMessageReader to simulate clone/readback.
+  auto segs = builder.getSegmentsForOutput();
+  capnp::SegmentArrayMessageReader sar(segs);
+  auto readBack = sar.getRoot<TestAllTypes>();
+  auto copiedData = readBack.getDataField();
+
+  ASSERT_EQ(copiedData.size(), SIZE);
+  for (size_t i = 0; i < SIZE; ++i) {
+    EXPECT_EQ(copiedData[i], original[i]);
+  }
+}
+
+TEST(Message, LazyZeroCustomBuilder_ManySmallDataAllocations_Stress) {
+  // Setup builder and lazy-zero skip for DATA.
+  BuilderOptions options;
+  options.lazyZeroSegmentAlloc.enableLazyZero = true;
+  options.lazyZeroSegmentAlloc.skipLazyZeroTypes.insert(schema::Type::DATA);
+
+  MyCustomMessageBuilder builder(options);
+  auto root = builder.initRoot<TestAllTypes>();
+
+  // Many small DATA allocations to stress allocation paths.
+  const int COUNT = 256;
+  kj::Vector< kj::ArrayPtr<capnp::byte> > allocated;
+  allocated.reserve(COUNT);
+
+  for (int i = 0; i < COUNT; ++i) {
+    auto d = root.initDataField(8 + (i % 16));
+    allocated.add(d);
+
+    // Quick check: each data allocation should have at least one non-zero byte.
+    bool allZero = true;
+    for (auto b : d) { if (b != 0) { allZero = false; break; } }
+    EXPECT_FALSE(allZero);
+  }
+
+  // Ensure at least one segment exists.
+  auto segs = builder.getSegmentsForOutput();
+  EXPECT_GE(segs.size(), 1u);
+}
+
+TEST(Message, LazyZeroCustomBuilder_PartialOverwriteLeavesRestDirtyForData) {
+  // Setup lazy-zero skip for DATA.
+  BuilderOptions options;
+  options.lazyZeroSegmentAlloc.enableLazyZero = true;
+  options.lazyZeroSegmentAlloc.skipLazyZeroTypes.insert(schema::Type::DATA);
+
+  MyCustomMessageBuilder builder(options);
+  auto root = builder.initRoot<TestAllTypes>();
+
+  const size_t SIZE = 64;
+  auto data = root.initDataField(SIZE);
+
+  // Overwrite only the first half with zeros, leave the second half untouched.
+  for (size_t i = 0; i < SIZE / 2; ++i) data[i] = 0;
+
+  // The second half should remain dirty (not all zero).
+  bool secondHalfAllZero = true;
+  for (size_t i = SIZE / 2; i < SIZE; ++i) {
+    if (data[i] != 0) { secondHalfAllZero = false; break; }
+  }
+  EXPECT_FALSE(secondHalfAllZero);
+
+  // Other fields remain default.
+  EXPECT_EQ(0u, root.getUInt32Field());
+  EXPECT_FALSE(root.hasTextField());
+}
+
+TEST(Message, LazyZeroCustomBuilder_NoSkipTypes_DataZero) {
+  // enableLazyZero = true, but skipLazyZeroTypes is empty -> DATA should be zeroed.
+  BuilderOptions options;
+  options.lazyZeroSegmentAlloc.enableLazyZero = true;
+  // skipLazyZeroTypes left empty
+
+  MyCustomMessageBuilder builder(options);
+  auto root = builder.initRoot<TestAllTypes>();
+
+  const size_t N = 32;
+  auto data = root.initDataField(N);
+
+  bool allZero = true;
+  for (size_t i = 0; i < data.size(); ++i) {
+    if (data[i] != 0) { allZero = false; break; }
+  }
+  EXPECT_TRUE(allZero);
+}
+
+TEST(Message, LazyZeroCustomBuilder_SkipText_ConstructionThrows) {
+  // enableLazyZero = true but skipLazyZeroTypes contains TEXT.
+  // Expect builder construction to fail (throw).
+  BuilderOptions options;
+  options.lazyZeroSegmentAlloc.enableLazyZero = true;
+  options.lazyZeroSegmentAlloc.skipLazyZeroTypes.insert(schema::Type::TEXT);
+
+  // Use any-throw expectation to avoid depending on the exact exception type.
+  EXPECT_ANY_THROW({
+    MyCustomMessageBuilder builder(options);
+  });
+}
+
+TEST(Message, LazyZeroCustomBuilder_SkipData_ModifyDataAndOtherFields_Persist) {
+  // skip DATA for lazy-zero; modify DATA and other fields and verify persistence.
+  BuilderOptions options;
+  options.lazyZeroSegmentAlloc.enableLazyZero = true;
+  options.lazyZeroSegmentAlloc.skipLazyZeroTypes.insert(schema::Type::DATA);
+
+  MyCustomMessageBuilder builder(options);
+  auto root = builder.initRoot<TestAllTypes>();
+
+  // Write pattern into DATA.
+  const size_t N = 24;
+  auto data = root.initDataField(N);
+  for (size_t i = 0; i < N; ++i) data[i] = static_cast<capnp::byte>((i * 7) & 0xFF);
+
+  // Modify other primitive/pointer fields.
+  root.setUInt32Field(0xDEADBEEF);
+  root.setInt64Field(-123456789);
+  root.setFloat64Field(3.14159);
+  root.setTextField("hello-world");
+
+  // Export segments and read back via reader to simulate serialization/clone.
+  auto segs = builder.getSegmentsForOutput();
+  capnp::SegmentArrayMessageReader reader(segs);
+  auto readBack = reader.getRoot<TestAllTypes>();
+
+  // Verify DATA persisted.
+  auto readData = readBack.getDataField();
+  ASSERT_EQ(readData.size(), N);
+  for (size_t i = 0; i < N; ++i) {
+    EXPECT_EQ(readData[i], static_cast<capnp::byte>((i * 7) & 0xFF));
+  }
+
+  // Verify other fields persisted.
+  EXPECT_EQ(readBack.getUInt32Field(), 0xDEADBEEF);
+  EXPECT_EQ(readBack.getInt64Field(), -123456789);
+  EXPECT_DOUBLE_EQ(readBack.getFloat64Field(), 3.14159);
+  EXPECT_TRUE(readBack.hasTextField());
+  EXPECT_EQ(readBack.getTextField(), "hello-world");
+}
+
 // TODO(test):  More tests.
 
 }  // namespace