fix(dot/state): fix usage of trie.Snapshot

dot/state/storage.go

@@ -43,7 +43,7 @@ func errTrieDoesNotExist(hash common.Hash) error {
 // StorageState is the struct that holds the trie, db and lock
 type StorageState struct {
 	blockState *BlockState
-	tries      map[common.Hash]*trie.Trie
+	tries      map[common.Hash]*trie.Trie // map of root -> trie
 
 	baseDB chaindb.Database
 	db     chaindb.Database
@@ -126,16 +126,27 @@ func (s *StorageState) TrieState(root *common.Hash) (*rtstorage.TrieState, error
 	t := s.tries[*root]
 	s.lock.RUnlock()
 
-	if t != nil && t.MustHash() == *root { // TODO: figure out why it seems like snapshotted tries are getting modified
-		return rtstorage.NewTrieState(t)
+	if t != nil && t.MustHash() != *root {
+		panic("trie does not have expected root")
 	}
 
-	tr, err := s.LoadFromDB(*root)
+	if t == nil {
+		var err error
+		t, err = s.LoadFromDB(*root)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	curr, err := rtstorage.NewTrieState(t)
 	if err != nil {
 		return nil, err
 	}
 
-	return rtstorage.NewTrieState(tr)
+	s.lock.Lock()
+	s.tries[*root] = curr.Snapshot()
+	s.lock.Unlock()
+	return curr, nil
 }
 
 // StoreInDB encodes the entire trie and writes it to the DB

dot/state/storage_test.go

@@ -37,6 +37,7 @@ func TestStorage_StoreAndLoadTrie(t *testing.T) {
 	require.NoError(t, err)
 	ts2, err := runtime.NewTrieState(trie)
 	require.NoError(t, err)
+	ts2.Snapshot()
 	require.Equal(t, ts.Trie(), ts2.Trie())
 }
 
@@ -70,6 +71,28 @@ func TestStorage_GetStorageByBlockHash(t *testing.T) {
 	require.Equal(t, value, res)
 }
 
+func TestStorage_TrieState(t *testing.T) {
+	storage := newTestStorageState(t)
+	ts, err := storage.TrieState(&trie.EmptyHash)
+	require.NoError(t, err)
+	ts.Set([]byte("noot"), []byte("washere"))
+
+	root, err := ts.Root()
+	require.NoError(t, err)
+	err = storage.StoreTrie(ts)
+	require.NoError(t, err)
+
+	time.Sleep(time.Millisecond * 100)
+
+	// get trie from db
+	storage.lock.Lock()
+	delete(storage.tries, root)
+	storage.lock.Unlock()
+	ts3, err := storage.TrieState(&root)
+	require.NoError(t, err)
+	require.Equal(t, ts.Trie().MustHash(), ts3.Trie().MustHash())
+}
+
 func TestStorage_LoadFromDB(t *testing.T) {
 	storage := newTestStorageState(t)
 	ts, err := storage.TrieState(&trie.EmptyHash)

dot/sync/syncer.go

@@ -303,7 +303,6 @@ func (s *Service) handleBlock(block *types.Block) error {
 		return err
 	}
 
-	ts.Snapshot()
 	root := ts.MustRoot()
 	if !bytes.Equal(parent.StateRoot[:], root[:]) {
 		panic("parent state root does not match snapshot state root")

lib/babe/babe.go

@@ -463,7 +463,6 @@ func (b *Service) handleSlot(slotNum uint64) error {
 		return err
 	}
 
-	ts.Snapshot()
 	b.rt.SetContextStorage(ts)
 
 	block, err := b.buildBlock(parent, currentSlot)

lib/runtime/storage/trie.go

@@ -52,8 +52,9 @@ func (s *TrieState) Trie() *trie.Trie {
 
 // Snapshot creates a new "version" of the trie. The trie before Snapshot is called
 // can no longer be modified, all further changes are on a new "version" of the trie.
-func (s *TrieState) Snapshot() {
-	_ = s.t.Snapshot()
+// It returns the previous version of the trie.
+func (s *TrieState) Snapshot() *trie.Trie {
+	return s.t.Snapshot()
 }
 
 // BeginStorageTransaction begins a new nested storage transaction which will either be committed or rolled back at a later time.
@@ -79,7 +80,7 @@ func (s *TrieState) RollbackStorageTransaction() {
 }
 
 // Set sets a key-value pair in the trie
-func (s *TrieState) Set(key []byte, value []byte) {
+func (s *TrieState) Set(key, value []byte) {
 	s.lock.Lock()
 	defer s.lock.Unlock()
 	s.t.Put(key, value)