From 61ff3e86b2d33f32566cae4f9bd1be973d84d757 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?P=C3=A9ter=20Szil=C3=A1gyi?= <peterke@gmail.com>
Date: Tue, 30 Mar 2021 19:04:22 +0300
Subject: [PATCH] core/state/snapshot, ethdb: track deletions more accurately
 (#22582)

* core/state/snapshot, ethdb: track deletions more accurately

* core/state/snapshot: don't reset the iterator, leveldb's screwy

* ethdb: don't mess with the insert batches for now
---
 core/state/snapshot/snapshot.go | 21 ++++++++++++++++++++-
 ethdb/leveldb/leveldb.go        |  2 +-
 ethdb/memorydb/memorydb.go      |  2 +-
 3 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/core/state/snapshot/snapshot.go b/core/state/snapshot/snapshot.go
index 810f1354e92c..eccf377264a5 100644
--- a/core/state/snapshot/snapshot.go
+++ b/core/state/snapshot/snapshot.go
@@ -484,8 +484,17 @@ func diffToDisk(bottom *diffLayer) *diskLayer {
 			if key := it.Key(); len(key) == 65 { // TODO(karalabe): Yuck, we should move this into the iterator
 				batch.Delete(key)
 				base.cache.Del(key[1:])
-
 				snapshotFlushStorageItemMeter.Mark(1)
+
+				// Ensure we don't delete too much data blindly (contract can be
+				// huge). It's ok to flush, the root will go missing in case of a
+				// crash and we'll detect and regenerate the snapshot.
+				if batch.ValueSize() > ethdb.IdealBatchSize {
+					if err := batch.Write(); err != nil {
+						log.Crit("Failed to write storage deletions", "err", err)
+					}
+					batch.Reset()
+				}
 			}
 		}
 		it.Release()
@@ -503,6 +512,16 @@ func diffToDisk(bottom *diffLayer) *diskLayer {
 
 		snapshotFlushAccountItemMeter.Mark(1)
 		snapshotFlushAccountSizeMeter.Mark(int64(len(data)))
+
+		// Ensure we don't write too much data blindly. It's ok to flush, the
+		// root will go missing in case of a crash and we'll detect and regen
+		// the snapshot.
+		if batch.ValueSize() > ethdb.IdealBatchSize {
+			if err := batch.Write(); err != nil {
+				log.Crit("Failed to write storage deletions", "err", err)
+			}
+			batch.Reset()
+		}
 	}
 	// Push all the storage slots into the database
 	for accountHash, storage := range bottom.storageData {
diff --git a/ethdb/leveldb/leveldb.go b/ethdb/leveldb/leveldb.go
index d3011212aa32..5d19cc3577de 100644
--- a/ethdb/leveldb/leveldb.go
+++ b/ethdb/leveldb/leveldb.go
@@ -461,7 +461,7 @@ func (b *batch) Put(key, value []byte) error {
 // Delete inserts the a key removal into the batch for later committing.
 func (b *batch) Delete(key []byte) error {
 	b.b.Delete(key)
-	b.size++
+	b.size += len(key)
 	return nil
 }
 
diff --git a/ethdb/memorydb/memorydb.go b/ethdb/memorydb/memorydb.go
index 4c5e1a84de8f..fedc9e326cf8 100644
--- a/ethdb/memorydb/memorydb.go
+++ b/ethdb/memorydb/memorydb.go
@@ -211,7 +211,7 @@ func (b *batch) Put(key, value []byte) error {
 // Delete inserts the a key removal into the batch for later committing.
 func (b *batch) Delete(key []byte) error {
 	b.writes = append(b.writes, keyvalue{common.CopyBytes(key), nil, true})
-	b.size += 1
+	b.size += len(key)
 	return nil
 }