add Reset() for spm and ewl

state/interface.go

@@ -178,6 +178,7 @@ type DBRemoveCacher interface {
 	Put([]byte, common.ModifiedHashes) error
 	Evict([]byte) (common.ModifiedHashes, error)
 	ShouldKeepHash(hash string, identifier TriePruningIdentifier) (bool, error)
+	Reset()
 	IsInterfaceNil() bool
 	Close() error
 }
@@ -194,6 +195,7 @@ type StoragePruningManager interface {
 	MarkForEviction([]byte, []byte, common.ModifiedHashes, common.ModifiedHashes) error
 	PruneTrie(rootHash []byte, identifier TriePruningIdentifier, tsm common.StorageManager, handler PruningHandler)
 	CancelPrune(rootHash []byte, identifier TriePruningIdentifier, tsm common.StorageManager)
+	Reset()
 	Close() error
 	IsInterfaceNil() bool
 }

state/storagePruningManager/disabled/disabledStoragePruningManager.go

@@ -26,6 +26,9 @@ func (i *disabledStoragePruningManager) PruneTrie(_ []byte, _ state.TriePruningI
 func (i *disabledStoragePruningManager) CancelPrune(_ []byte, _ state.TriePruningIdentifier, _ common.StorageManager) {
 }
 
+// Reset does nothing for this implementation
+func (i *disabledStoragePruningManager) Reset() {}
+
 // Close does nothing for this implementation
 func (i *disabledStoragePruningManager) Close() error {
 	return nil

state/storagePruningManager/evictionWaitingList/memoryEvictionWaitingList.go

@@ -205,6 +205,19 @@ func (mewl *memoryEvictionWaitingList) ShouldKeepHash(hash string, identifier st
 	return false, nil
 }
 
+// Reset will reinitialize the eviction waiting list, by emptying the cache and reversed cache. It will not change the sizes of the caches.
+func (mewl *memoryEvictionWaitingList) Reset() {
+	mewl.opMutex.Lock()
+
+	for key := range mewl.cache {
+		log.Debug("trie nodes eviction waiting list reset", "rootHash", []byte(key))
+	}
+
+	mewl.cache = make(map[string]*rootHashData)
+	mewl.reversedCache = make(map[string]*hashInfo)
+	mewl.opMutex.Unlock()
+}
+
 // Close returns nil
 func (mewl *memoryEvictionWaitingList) Close() error {
 	return nil

state/storagePruningManager/storagePruningManager.go

@@ -234,6 +234,12 @@ func (spm *storagePruningManager) removeFromDb(
 	return nil
 }
 
+func (spm *storagePruningManager) Reset() {
+	log.Debug("storage pruning manager reset")
+	_ = spm.pruningBuffer.RemoveAll()
+	spm.dbEvictionWaitingList.Reset()
+}
+
 // Close will handle the closing of the underlying components
 func (spm *storagePruningManager) Close() error {
 	return spm.dbEvictionWaitingList.Close()

state/storagePruningManager/storagePruningManager_test.go

@@ -260,3 +260,37 @@ func TestStoragePruningManager_MarkForEviction_removeDuplicatedKeys(t *testing.T
 	_, ok = map2["hash4"]
 	assert.False(t, ok)
 }
+
+func TestStoreagePruningManager_Reset(t *testing.T) {
+	t.Parallel()
+
+	args := storage.GetStorageManagerArgs()
+	trieStorage, _ := trie.NewTrieStorageManager(args)
+	ewlArgs := evictionWaitingList.MemoryEvictionWaitingListArgs{
+		RootHashesSize: 100,
+		HashesSize:     10000,
+	}
+	ewl, _ := evictionWaitingList.NewMemoryEvictionWaitingList(ewlArgs)
+	spm, _ := NewStoragePruningManager(ewl, 1000)
+
+	err := spm.MarkForEviction([]byte("rootHash"), []byte("newRootHash"), map[string]struct{}{"hash1": {}, "hash2": {}}, map[string]struct{}{"hash3": {}, "hash4": {}})
+	assert.Nil(t, err)
+	err = spm.markForEviction([]byte("rootHash2"), map[string]struct{}{"hash5": {}, "hash6": {}}, state.NewRoot)
+	assert.Nil(t, err)
+
+	trieStorage.EnterPruningBufferingMode()
+	spm.PruneTrie([]byte("rootHash"), state.OldRoot, trieStorage, state.NewPruningHandler(state.EnableDataRemoval))
+	spm.CancelPrune([]byte("newRootHash"), state.NewRoot, trieStorage)
+	trieStorage.ExitPruningBufferingMode()
+
+	assert.Equal(t, 2, spm.pruningBuffer.Len())
+
+	spm.Reset()
+	assert.Equal(t, 0, spm.pruningBuffer.Len())
+
+	// rootHash2 will should not be added to the pruning buffer because ewl was also reset when spm.Reset() was called
+	trieStorage.EnterPruningBufferingMode()
+	spm.PruneTrie([]byte("rootHash2"), state.NewRoot, trieStorage, state.NewPruningHandler(state.EnableDataRemoval))
+	trieStorage.ExitPruningBufferingMode()
+	assert.Equal(t, 0, spm.pruningBuffer.Len())
+}

testscommon/state/evictionWaitingListMock.go

@@ -81,6 +81,13 @@ func (ewl *EvictionWaitingListMock) ShouldKeepHash(hash string, identifier state
 	return false, nil
 }
 
+// Reset will reinitialize the cache
+func (ewl *EvictionWaitingListMock) Reset() {
+	ewl.OpMutex.Lock()
+	ewl.Cache = make(map[string]common.ModifiedHashes)
+	ewl.OpMutex.Unlock()
+}
+
 // Close -
 func (ewl *EvictionWaitingListMock) Close() error {
 	return nil

testscommon/state/storagePruningManagerStub.go

@@ -10,6 +10,7 @@ type StoragePruningManagerStub struct {
 	MarkForEvictionCalled func(bytes []byte, bytes2 []byte, hashes common.ModifiedHashes, hashes2 common.ModifiedHashes) error
 	PruneTrieCalled       func(rootHash []byte, identifier state.TriePruningIdentifier, tsm common.StorageManager, handler state.PruningHandler)
 	CancelPruneCalled     func(rootHash []byte, identifier state.TriePruningIdentifier, tsm common.StorageManager)
+	ResetCalled           func()
 	CloseCalled           func() error
 }
 
@@ -36,6 +37,13 @@ func (stub *StoragePruningManagerStub) CancelPrune(rootHash []byte, identifier s
 	}
 }
 
+// Reset -
+func (stub *StoragePruningManagerStub) Reset() {
+	if stub.ResetCalled != nil {
+		stub.ResetCalled()
+	}
+}
+
 // Close -
 func (stub *StoragePruningManagerStub) Close() error {
 	if stub.CloseCalled != nil {