fix(state) : Update StorageState to load storage from database. (#1486)

dot/state/storage.go

@@ -204,19 +204,9 @@ func (s *StorageState) LoadFromDB(root common.Hash) (*trie.Trie, error) {
 
 // ExistsStorage check if the key exists in the storage trie with the given storage hash
 // If no hash is provided, the current chain head is used
-func (s *StorageState) ExistsStorage(hash *common.Hash, key []byte) (bool, error) {
-	if hash == nil {
-		sr, err := s.blockState.BestBlockStateRoot()
-		if err != nil {
-			return false, err
-		}
-		hash = &sr
-	}
-
-	s.lock.RLock()
-	defer s.lock.RUnlock()
-	val := s.tries[*hash].Get(key)
-	return val != nil, nil
+func (s *StorageState) ExistsStorage(root *common.Hash, key []byte) (bool, error) {
+	val, err := s.GetStorage(root, key)
+	return val != nil, err
 }
 
 // GetStorage gets the object from the trie using the given key and storage hash
@@ -233,8 +223,8 @@ func (s *StorageState) GetStorage(root *common.Hash, key []byte) ([]byte, error)
 	s.lock.RLock()
 	defer s.lock.RUnlock()
 
-	if s.tries[*root] != nil {
-		val := s.tries[*root].Get(key)
+	if trie, ok := s.tries[*root]; ok {
+		val := trie.Get(key)
 		return val, nil
 	}
 
@@ -262,19 +252,7 @@ func (s *StorageState) GetStateRootFromBlock(bhash *common.Hash) (*common.Hash,
 
 // StorageRoot returns the root hash of the current storage trie
 func (s *StorageState) StorageRoot() (common.Hash, error) {
-	sr, err := s.blockState.BestBlockStateRoot()
-	if err != nil {
-		return common.Hash{}, err
-	}
-
-	s.lock.RLock()
-	defer s.lock.RUnlock()
-
-	if s.tries[sr] == nil {
-		return common.Hash{}, errTrieDoesNotExist(sr)
-	}
-
-	return s.tries[sr].Hash()
+	return s.blockState.BestBlockStateRoot()
 }
 
 // EnumeratedTrieRoot not implemented
@@ -294,17 +272,19 @@ func (s *StorageState) Entries(root *common.Hash) (map[string][]byte, error) {
 	}
 
 	s.lock.RLock()
-	defer s.lock.RUnlock()
-
-	if s.tries[*root] != nil {
-		return s.tries[*root].Entries(), nil
-	}
+	tr, ok := s.tries[*root]
+	s.lock.RUnlock()
 
-	tr, err := s.LoadFromDB(*root)
-	if err != nil {
-		return nil, err
+	if !ok {
+		var err error
+		tr, err = s.LoadFromDB(*root)
+		if err != nil {
+			return nil, errTrieDoesNotExist(*root)
+		}
 	}
 
+	s.lock.RLock()
+	defer s.lock.RUnlock()
 	return tr.Entries(), nil
 }
 
@@ -317,11 +297,22 @@ func (s *StorageState) GetKeysWithPrefix(hash *common.Hash, prefix []byte) ([][]
 		}
 		hash = &sr
 	}
-	t := s.tries[*hash]
-	if t == nil {
-		return nil, fmt.Errorf("unable to retrieve trie with hash %x", *hash)
+
+	s.lock.RLock()
+	tr, ok := s.tries[*hash]
+	s.lock.RUnlock()
+
+	if !ok {
+		var err error
+		tr, err = s.LoadFromDB(*hash)
+		if err != nil {
+			return nil, errTrieDoesNotExist(*hash)
+		}
 	}
-	return t.GetKeysWithPrefix(prefix), nil
+
+	s.lock.RLock()
+	defer s.lock.RUnlock()
+	return tr.GetKeysWithPrefix(prefix), nil
 }
 
 // GetStorageChild return GetChild from the trie
@@ -335,13 +326,20 @@ func (s *StorageState) GetStorageChild(hash *common.Hash, keyToChild []byte) (*t
 	}
 
 	s.lock.RLock()
-	defer s.lock.RUnlock()
+	tr, ok := s.tries[*hash]
+	s.lock.RUnlock()
 
-	if s.tries[*hash] == nil {
-		return nil, errTrieDoesNotExist(*hash)
+	if !ok {
+		var err error
+		tr, err = s.LoadFromDB(*hash)
+		if err != nil {
+			return nil, errTrieDoesNotExist(*hash)
+		}
 	}
 
-	return s.tries[*hash].GetChild(keyToChild)
+	s.lock.RLock()
+	defer s.lock.RUnlock()
+	return tr.GetChild(keyToChild)
 }
 
 // GetStorageFromChild return GetFromChild from the trie
@@ -355,12 +353,20 @@ func (s *StorageState) GetStorageFromChild(hash *common.Hash, keyToChild, key []
 	}
 
 	s.lock.RLock()
-	defer s.lock.RUnlock()
+	tr, ok := s.tries[*hash]
+	s.lock.RUnlock()
 
-	if s.tries[*hash] == nil {
-		return nil, errTrieDoesNotExist(*hash)
+	if !ok {
+		var err error
+		tr, err = s.LoadFromDB(*hash)
+		if err != nil {
+			return nil, errTrieDoesNotExist(*hash)
+		}
 	}
-	return s.tries[*hash].GetFromChild(keyToChild, key)
+
+	s.lock.RLock()
+	defer s.lock.RUnlock()
+	return tr.GetFromChild(keyToChild, key)
 }
 
 // LoadCode returns the runtime code (located at :code)

dot/state/storage_test.go

@@ -69,3 +69,54 @@ func TestStorage_GetStorageByBlockHash(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, value, res)
 }
+
+func TestStorage_LoadFromDB(t *testing.T) {
+	storage := newTestStorageState(t)
+	ts, err := storage.TrieState(&trie.EmptyHash)
+	require.NoError(t, err)
+
+	trieKV := []struct {
+		key   []byte
+		value []byte
+	}{{},
+		{[]byte("key1"), []byte("value1")},
+		{[]byte("key2"), []byte("value2")},
+		{[]byte("xyzKey1"), []byte("xyzValue1")},
+	}
+
+	for _, kv := range trieKV {
+		ts.Set(kv.key, kv.value)
+	}
+
+	root, err := ts.Root()
+	require.NoError(t, err)
+
+	// Write trie to disk.
+	err = storage.StoreTrie(ts)
+	require.NoError(t, err)
+
+	// Clear trie from cache and fetch data from disk.
+	storage.lock.Lock()
+	delete(storage.tries, root)
+	storage.lock.Unlock()
+
+	data, err := storage.GetStorage(&root, trieKV[0].key)
+	require.NoError(t, err)
+	require.Equal(t, trieKV[0].value, data)
+
+	storage.lock.Lock()
+	delete(storage.tries, root)
+	storage.lock.Unlock()
+
+	prefixKeys, err := storage.GetKeysWithPrefix(&root, []byte("ke"))
+	require.NoError(t, err)
+	require.Equal(t, 2, len(prefixKeys))
+
+	storage.lock.Lock()
+	delete(storage.tries, root)
+	storage.lock.Unlock()
+
+	entries, err := storage.Entries(&root)
+	require.NoError(t, err)
+	require.Equal(t, 3, len(entries))
+}

lib/trie/child_storage.go

@@ -36,7 +36,7 @@ func (t *Trie) PutChild(keyToChild []byte, child *Trie) error {
 	value := [32]byte(childHash)
 
 	t.Put(key, value[:])
-	t.children[childHash] = child
+	t.childTries[childHash] = child
 	return nil
 }
 
@@ -50,7 +50,7 @@ func (t *Trie) GetChild(keyToChild []byte) (*Trie, error) {
 
 	hash := [32]byte{}
 	copy(hash[:], childHash)
-	return t.children[common.Hash(hash)], nil
+	return t.childTries[common.Hash(hash)], nil
 }
 
 // PutIntoChild puts a key-value pair into the child trie located in the main trie at key :child_storage:[keyToChild]
@@ -71,8 +71,8 @@ func (t *Trie) PutIntoChild(keyToChild, key, value []byte) error {
 		return err
 	}
 
-	t.children[origChildHash] = nil
-	t.children[childHash] = child
+	t.childTries[origChildHash] = nil
+	t.childTries[childHash] = child
 
 	return t.PutChild(keyToChild, child)
 }

lib/trie/trie.go

@@ -31,7 +31,7 @@ var EmptyHash, _ = NewEmptyTrie().Hash()
 type Trie struct {
 	generation uint64
 	root       node
-	children   map[common.Hash]*Trie // Used to store the child tries.
+	childTries map[common.Hash]*Trie // Used to store the child tries.
 }
 
 // NewEmptyTrie creates a trie with a nil root
@@ -43,7 +43,7 @@ func NewEmptyTrie() *Trie {
 func NewTrie(root node) *Trie {
 	return &Trie{
 		root:       root,
-		children:   make(map[common.Hash]*Trie),
+		childTries: make(map[common.Hash]*Trie),
 		generation: 0, // Initially zero but increases after every snapshot.
 	}
 }
@@ -53,7 +53,7 @@ func (t *Trie) Snapshot() *Trie {
 	oldTrie := &Trie{
 		generation: t.generation,
 		root:       t.root,
-		children:   t.children,
+		childTries: t.childTries,
 	}
 	t.generation++
 	return oldTrie