core/rawdb: add specified key length iterator (#24535)

core/rawdb/accessors_snapshot.go

@@ -115,7 +115,7 @@ func DeleteStorageSnapshot(db ethdb.KeyValueWriter, accountHash, storageHash com
 // IterateStorageSnapshots returns an iterator for walking the entire storage
 // space of a specific account.
 func IterateStorageSnapshots(db ethdb.Iteratee, accountHash common.Hash) ethdb.Iterator {
-	return db.NewIterator(storageSnapshotsKey(accountHash), nil)
+	return NewKeyLengthIterator(db.NewIterator(storageSnapshotsKey(accountHash), nil), len(SnapshotStoragePrefix)+2*common.HashLength)
 }
 
 // ReadSnapshotJournal retrieves the serialized in-memory diff layers saved at

core/rawdb/key_length_iterator.go

@@ -0,0 +1,47 @@
+// Copyright 2022 The go-ethereum Authors
+// This file is part of the go-ethereum library.
+//
+// The go-ethereum library is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Lesser General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// The go-ethereum library is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Lesser General Public License for more details.
+//
+// You should have received a copy of the GNU Lesser General Public License
+// along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.
+
+package rawdb
+
+import "github.com/ethereum/go-ethereum/ethdb"
+
+// KeyLengthIterator is a wrapper for a database iterator that ensures only key-value pairs
+// with a specific key length will be returned.
+type KeyLengthIterator struct {
+	requiredKeyLength int
+	ethdb.Iterator
+}
+
+// NewKeyLengthIterator returns a wrapped version of the iterator that will only return key-value
+// pairs where keys with a specific key length will be returned.
+func NewKeyLengthIterator(it ethdb.Iterator, keyLen int) ethdb.Iterator {
+	return &KeyLengthIterator{
+		Iterator:          it,
+		requiredKeyLength: keyLen,
+	}
+}
+
+func (it *KeyLengthIterator) Next() bool {
+	// Return true as soon as a key with the required key length is discovered
+	for it.Iterator.Next() {
+		if len(it.Iterator.Key()) == it.requiredKeyLength {
+			return true
+		}
+	}
+
+	// Return false when we exhaust the keys in the underlying iterator.
+	return false
+}

core/rawdb/key_length_iterator_test.go

@@ -0,0 +1,60 @@
+// Copyright 2022 The go-ethereum Authors
+// This file is part of the go-ethereum library.
+//
+// The go-ethereum library is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Lesser General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// The go-ethereum library is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Lesser General Public License for more details.
+//
+// You should have received a copy of the GNU Lesser General Public License
+// along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.
+
+package rawdb
+
+import (
+	"encoding/binary"
+	"testing"
+)
+
+func TestKeyLengthIterator(t *testing.T) {
+	db := NewMemoryDatabase()
+
+	keyLen := 8
+	expectedKeys := make(map[string]struct{})
+	for i := 0; i < 100; i++ {
+		key := make([]byte, keyLen)
+		binary.BigEndian.PutUint64(key, uint64(i))
+		if err := db.Put(key, []byte{0x1}); err != nil {
+			t.Fatal(err)
+		}
+		expectedKeys[string(key)] = struct{}{}
+
+		longerKey := make([]byte, keyLen*2)
+		binary.BigEndian.PutUint64(longerKey, uint64(i))
+		if err := db.Put(longerKey, []byte{0x1}); err != nil {
+			t.Fatal(err)
+		}
+	}
+
+	it := NewKeyLengthIterator(db.NewIterator(nil, nil), keyLen)
+	for it.Next() {
+		key := it.Key()
+		_, exists := expectedKeys[string(key)]
+		if !exists {
+			t.Fatalf("Found unexpected key %d", binary.BigEndian.Uint64(key))
+		}
+		delete(expectedKeys, string(key))
+		if len(key) != keyLen {
+			t.Fatalf("Found unexpected key in key length iterator with length %d", len(key))
+		}
+	}
+
+	if len(expectedKeys) != 0 {
+		t.Fatalf("Expected all keys of length %d to be removed from expected keys during iteration", keyLen)
+	}
+}

core/state/snapshot/snapshot.go

@@ -546,20 +546,19 @@ func diffToDisk(bottom *diffLayer) *diskLayer {
 
 		it := rawdb.IterateStorageSnapshots(base.diskdb, hash)
 		for it.Next() {
-			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)
+			key := it.Key()
+			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()
+			// 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()