Add SizeEstimate to CDBBatch

sipa · random-zebra · commit 802f8c8748ba · 2020-08-18T15:17:51.000+02:00
This allows estimating the in-memory size of a LevelDB batch.
diff --git a/src/dbwrapper.h b/src/dbwrapper.h
@@ -48,7 +48,22 @@ class CDBBatch
 private:
     leveldb::WriteBatch batch;
 
+    CDataStream ssKey;
+    CDataStream ssValue;
+
+    size_t size_estimate;
+
 public:
+    /**
+     * @param[in] _parent   CDBWrapper that this batch is to be submitted to
+     */
+    CDBBatch() : ssKey(SER_DISK, CLIENT_VERSION), ssValue(SER_DISK, CLIENT_VERSION), size_estimate(0) { };
+
+    void Clear()
+    {
+        batch.Clear();
+        size_estimate = 0;
+    }
 
     template <typename K, typename V>
     void Write(const K& key, const V& value)
@@ -64,6 +79,17 @@ class CDBBatch
         leveldb::Slice slValue(&ssValue[0], ssValue.size());
 
         batch.Put(slKey, slValue);
+
+        // LevelDB serializes writes as:
+        // - byte: header
+        // - varint: key length (1 byte up to 127B, 2 bytes up to 16383B, ...)
+        // - byte[]: key
+        // - varint: value length
+        // - byte[]: value
+        // The formula below assumes the key and value are both less than 16k.
+        size_estimate += 3 + (slKey.size() > 127) + slKey.size() + (slValue.size() > 127) + slValue.size();
+        ssKey.clear();
+        ssValue.clear();
     }
 
     template <typename K>
@@ -75,7 +101,17 @@ class CDBBatch
         leveldb::Slice slKey(&ssKey[0], ssKey.size());
 
         batch.Delete(slKey);
+
+        // LevelDB serializes erases as:
+        // - byte: header
+        // - varint: key length
+        // - byte[]: key
+        // The formula below assumes the key is less than 16kB.
+        size_estimate += 2 + (slKey.size() > 127) + slKey.size();
+        ssKey.clear();
     }
+
+    size_t SizeEstimate() const { return size_estimate; }
 };
 
 class CDBIterator
