From c3701b265e3568b5de8922bb7627d99185c42487 Mon Sep 17 00:00:00 2001 From: aaronbuchwald Date: Tue, 15 Mar 2022 04:28:26 -0400 Subject: [PATCH] core/rawdb: add specified key length iterator (#24535) --- core/rawdb/accessors_snapshot.go | 2 +- core/rawdb/key_length_iterator.go | 47 ++++++++++++++++++++ core/rawdb/key_length_iterator_test.go | 60 ++++++++++++++++++++++++++ core/state/snapshot/snapshot.go | 23 +++++----- 4 files changed, 119 insertions(+), 13 deletions(-) create mode 100644 core/rawdb/key_length_iterator.go create mode 100644 core/rawdb/key_length_iterator_test.go diff --git a/core/rawdb/accessors_snapshot.go b/core/rawdb/accessors_snapshot.go index 1c828662c..3c82b3f73 100644 --- a/core/rawdb/accessors_snapshot.go +++ b/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 diff --git a/core/rawdb/key_length_iterator.go b/core/rawdb/key_length_iterator.go new file mode 100644 index 000000000..d1c5af269 --- /dev/null +++ b/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 . + +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 +} diff --git a/core/rawdb/key_length_iterator_test.go b/core/rawdb/key_length_iterator_test.go new file mode 100644 index 000000000..654efc5b5 --- /dev/null +++ b/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 . + +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) + } +} diff --git a/core/state/snapshot/snapshot.go b/core/state/snapshot/snapshot.go index 6ee6b06bb..76200851e 100644 --- a/core/state/snapshot/snapshot.go +++ b/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()