Added fast size estimates to Arena which are O(1) and are usually accurate. When inaccurate, a full size scan will update all estimates in the tree.

flow/Arena.cpp

@@ -106,15 +106,22 @@ void* Arena::allocate4kAlignedBuffer(uint32_t size) {
 	return ArenaBlock::dependOn4kAlignedBuffer(impl, size);
 }
 
-size_t Arena::getSize() const {
+size_t Arena::getSize(bool fastInaccurateEstimate) const {
 	if (impl) {
 		allowAccess(impl.getPtr());
-		auto result = impl->totalSize();
+		size_t result;
+		if (fastInaccurateEstimate) {
+			result = impl->estimatedTotalSize();
+		} else {
+			result = impl->totalSize();
+		}
+
 		disallowAccess(impl.getPtr());
 		return result;
 	}
 	return 0;
 }
+
 bool Arena::hasFree(size_t size, const void* address) {
 	if (impl) {
 		allowAccess(impl.getPtr());
@@ -167,28 +174,38 @@ const void* ArenaBlock::getData() const {
 const void* ArenaBlock::getNextData() const {
 	return (const uint8_t*)getData() + used();
 }
-size_t ArenaBlock::totalSize() {
+
+size_t ArenaBlock::totalSize() const {
 	if (isTiny()) {
 		return size();
 	}
 
-	size_t s = size();
+	// Walk the entire tree to get an accurate size and store it in the estimate for
+	// each block, recursively.
+	totalSizeEstimate = size();
 	int o = nextBlockOffset;
 	while (o) {
 		ArenaBlockRef* r = (ArenaBlockRef*)((char*)getData() + o);
 		makeDefined(r, sizeof(ArenaBlockRef));
 		if (r->aligned4kBufferSize != 0) {
-			s += r->aligned4kBufferSize;
+			totalSizeEstimate += r->aligned4kBufferSize;
 		} else {
 			allowAccess(r->next);
-			s += r->next->totalSize();
+			totalSizeEstimate += r->next->totalSize();
 			disallowAccess(r->next);
 		}
 		o = r->nextBlockOffset;
 		makeNoAccess(r, sizeof(ArenaBlockRef));
 	}
-	return s;
+	return totalSizeEstimate;
 }
+size_t ArenaBlock::estimatedTotalSize() const {
+	if (isTiny()) {
+		return size();
+	}
+	return totalSizeEstimate;
+}
+
 // just for debugging:
 void ArenaBlock::getUniqueBlocks(std::set<ArenaBlock*>& a) {
 	a.insert(this);
@@ -232,6 +249,7 @@ void ArenaBlock::makeReference(ArenaBlock* next) {
 	makeNoAccess(r, sizeof(ArenaBlockRef));
 	nextBlockOffset = bigUsed;
 	bigUsed += sizeof(ArenaBlockRef);
+	totalSizeEstimate += next->estimatedTotalSize();
 }
 
 void* ArenaBlock::make4kAlignedBuffer(uint32_t size) {
@@ -245,6 +263,7 @@ void* ArenaBlock::make4kAlignedBuffer(uint32_t size) {
 	makeNoAccess(r, sizeof(ArenaBlockRef));
 	nextBlockOffset = bigUsed;
 	bigUsed += sizeof(ArenaBlockRef);
+	totalSizeEstimate += size;
 	return result;
 }
 
@@ -341,6 +360,7 @@ ArenaBlock* ArenaBlock::create(int dataSize, Reference<ArenaBlock>& next) {
 				b->bigSize = 8192;
 				INSTRUMENT_ALLOCATE("Arena8192");
 			}
+			b->totalSizeEstimate = b->bigSize;
 			b->tinySize = b->tinyUsed = NOT_TINY;
 			b->bigUsed = sizeof(ArenaBlock);
 		} else {
@@ -350,6 +370,7 @@ ArenaBlock* ArenaBlock::create(int dataSize, Reference<ArenaBlock>& next) {
 			b = (ArenaBlock*)new uint8_t[reqSize];
 			b->tinySize = b->tinyUsed = NOT_TINY;
 			b->bigSize = reqSize;
+			b->totalSizeEstimate = b->bigSize;
 			b->bigUsed = sizeof(ArenaBlock);
 
 			if (FLOW_KNOBS && g_allocation_tracing_disabled == 0 &&
@@ -650,3 +671,59 @@ TEST_CASE("/flow/Arena/DefaultBoostHash") {
 
 	return Void();
 }
+
+TEST_CASE("/flow/Arena/Size") {
+	Arena a;
+
+	// Size estimates are accurate unless dependencies are added to an Arena via another Arena
+	// handle which points to a non-root node.
+	//
+	// Note that the ASSERT argument order matters, the estimate must be calculated first as
+	// the full accurate calculation will update the estimate
+	makeString(40, a);
+	ASSERT_EQ(a.getSize(true), a.getSize());
+
+	makeString(700, a);
+	ASSERT_EQ(a.getSize(true), a.getSize());
+
+	// Copy a at a point where it points to a large block with room for block references
+	Arena b = a;
+
+	// copy a at a point where there isn't room for more block references
+	makeString(1000, a);
+	Arena c = a;
+
+	makeString(1000, a);
+	makeString(1000, a);
+	ASSERT_EQ(a.getSize(true), a.getSize());
+
+	Standalone<StringRef> s = makeString(500);
+	a.dependsOn(s.arena());
+	ASSERT_EQ(a.getSize(true), a.getSize());
+
+	Standalone<StringRef> s2 = makeString(500);
+	a.dependsOn(s2.arena());
+	ASSERT_EQ(a.getSize(true), a.getSize());
+
+	// Add a dependency to b, which will fit in b's root and update b's size estimate
+	Standalone<StringRef> s3 = makeString(100);
+	b.dependsOn(s3.arena());
+	ASSERT_EQ(b.getSize(true), b.getSize());
+
+	// But now a's size estimate is out of date because the new reference in b's root is still
+	// in a's tree
+	ASSERT_LT(a.getSize(true), a.getSize());
+
+	// Now that a full size calc has been done on a, the estimate is up to date.
+	ASSERT_EQ(a.getSize(true), a.getSize());
+
+	// Add a dependency to c, which will NOT fit in c's root, so it will be added to a new
+	// root for c and that root will not be in a's tree so a's size and estimate remain
+	// unchanged and the same.  The size and estimate of c will also match.
+	Standalone<StringRef> s4 = makeString(100);
+	c.dependsOn(s4.arena());
+	ASSERT_EQ(c.getSize(true), c.getSize());
+	ASSERT_EQ(a.getSize(true), a.getSize());
+
+	return Void();
+}

flow/Arena.h

@@ -104,7 +104,13 @@ public:
 
 	void dependsOn(const Arena& p);
 	void* allocate4kAlignedBuffer(uint32_t size);
-	size_t getSize() const;
+
+	// If fastInaccurateEstimate is true this operation is O(1) but it is inaccurate in that it
+	// will omit memory added to this Arena's block tree using Arena handles which reference
+	// non-root nodes in this Arena's block tree.
+	// When fastInaccurateEstimate is false, all estimates in the block tree will be updated to
+	// be accurate.
+	size_t getSize(bool fastInaccurateEstimate = false) const;
 
 	bool hasFree(size_t size, const void* address);
 
@@ -156,6 +162,7 @@ struct ArenaBlock : NonCopyable, ThreadSafeReferenceCounted<ArenaBlock> {
 	// if tinySize != NOT_TINY, following variables aren't used
 	uint32_t bigSize, bigUsed; // include block header
 	uint32_t nextBlockOffset;
+	mutable size_t totalSizeEstimate; // Estimate of the minimum total size of arena blocks this one reaches
 
 	void addref();
 	void delref();
@@ -165,7 +172,8 @@ struct ArenaBlock : NonCopyable, ThreadSafeReferenceCounted<ArenaBlock> {
 	int unused() const;
 	const void* getData() const;
 	const void* getNextData() const;
-	size_t totalSize();
+	size_t totalSize() const;
+	size_t estimatedTotalSize() const;
 	// just for debugging:
 	void getUniqueBlocks(std::set<ArenaBlock*>& a);
 	int addUsed(int bytes);