fix(dot/state): add StorageState Lock/Unlock API for usage by babe and sync

dot/core/service.go

@@ -307,16 +307,17 @@ func (s *Service) handleCurrentSlot(header *types.Header) error {
 // does not need to be completed before the next block can be imported.
 func (s *Service) handleBlocksAsync() {
 	for {
+		prev := s.blockState.BestBlockHash()
+
 		select {
 		case block := <-s.blockAddCh:
 			if block == nil {
 				continue
 			}
 
-			// TODO: add inherent check
-			// if err := s.handleChainReorg(prev, block.Header.Hash()); err != nil {
-			// 	logger.Warn("failed to re-add transactions to chain upon re-org", "error", err)
-			// }
+			if err := s.handleChainReorg(prev, block.Header.Hash()); err != nil {
+				logger.Warn("failed to re-add transactions to chain upon re-org", "error", err)
+			}
 
 			if err := s.maintainTransactionPool(block); err != nil {
 				logger.Warn("failed to maintain transaction pool", "error", err)
@@ -422,12 +423,11 @@ func (s *Service) maintainTransactionPool(block *types.Block) error {
 	// re-validate transactions in the pool and move them to the queue
 	txs := s.transactionState.PendingInPool()
 	for _, tx := range txs {
-		// TODO: re-add this on update to v0.8
-
+		// TODO: re-add this
 		// val, err := s.rt.ValidateTransaction(tx.Extrinsic)
 		// if err != nil {
 		// 	// failed to validate tx, remove it from the pool or queue
-		// 	s.transactionState.RemoveExtrinsic(ext)
+		// 	s.transactionState.RemoveExtrinsic(tx.Extrinsic)
 		// 	continue
 		// }
 

dot/network/sync.go

@@ -344,7 +344,13 @@ func (q *syncQueue) benchmark() {
 				continue
 			}
 
-			logger.Info("💤 node waiting", "peer count", len(q.s.host.peers()), "head", before.Number, "finalised", finalised.Number)
+			logger.Info("💤 node waiting",
+				"peer count", len(q.s.host.peers()),
+				"head", before.Number,
+				"hash", before.Hash(),
+				"finalised", finalised.Number,
+				"hash", finalised.Hash(),
+			)
 			time.Sleep(time.Second * 5)
 			continue
 		}

dot/state/service.go

@@ -271,14 +271,13 @@ func (s *Service) Stop() error {
 		return err
 	}
 
-	s.Storage.lock.RLock()
-	t := s.Storage.tries[head]
-	s.Storage.lock.RUnlock()
-
-	if t == nil {
+	st, has := s.Storage.tries.Load(head)
+	if !has {
 		return errTrieDoesNotExist(head)
 	}
 
+	t := st.(*trie.Trie)
+
 	if err = s.Base.StoreLatestStorageHash(head); err != nil {
 		return err
 	}

dot/state/service_test.go

@@ -272,10 +272,8 @@ func TestService_PruneStorage(t *testing.T) {
 	time.Sleep(1 * time.Second)
 
 	for _, v := range prunedArr {
-		serv.Storage.lock.Lock()
-		_, ok := serv.Storage.tries[v.hash]
-		serv.Storage.lock.Unlock()
-		require.Equal(t, false, ok)
+		_, has := serv.Storage.tries.Load(v.hash)
+		require.Equal(t, false, has)
 	}
 }
 

dot/state/storage.go

@@ -44,10 +44,10 @@ 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 // map of root -> trie
+	tries      *sync.Map // map[common.Hash]*trie.Trie // map of root -> trie
 
-	db   chaindb.Database
-	lock sync.RWMutex
+	db chaindb.Database
+	sync.RWMutex
 
 	// change notifiers
 	changedLock  sync.RWMutex
@@ -66,8 +66,8 @@ func NewStorageState(db chaindb.Database, blockState *BlockState, t *trie.Trie,
 		return nil, fmt.Errorf("cannot have nil trie")
 	}
 
-	tries := make(map[common.Hash]*trie.Trie)
-	tries[t.MustHash()] = t
+	tries := new(sync.Map)
+	tries.Store(t.MustHash(), t)
 
 	storageTable := chaindb.NewTable(db, storagePrefix)
 
@@ -97,25 +97,23 @@ func (s *StorageState) SetSyncing(syncing bool) {
 }
 
 func (s *StorageState) pruneKey(keyHeader *types.Header) {
-	s.lock.Lock()
-	defer s.lock.Unlock()
-
-	delete(s.tries, keyHeader.StateRoot)
+	s.tries.Delete(keyHeader.StateRoot)
 }
 
 // StoreTrie stores the given trie in the StorageState and writes it to the database
 func (s *StorageState) StoreTrie(ts *rtstorage.TrieState, header *types.Header) error {
-	s.lock.Lock()
-	defer s.lock.Unlock()
-
 	root := ts.MustRoot()
+
 	if s.syncing {
 		// keep only the trie at the head of the chain when syncing
-		for key := range s.tries {
-			delete(s.tries, key)
-		}
+		// TODO: probably remove this when memory usage improves
+		s.tries.Range(func(k, _ interface{}) bool {
+			s.tries.Delete(k)
+			return true
+		})
 	}
-	s.tries[root] = ts.Trie()
+
+	_, _ = s.tries.LoadOrStore(root, ts.Trie())
 
 	if _, ok := s.pruner.(*pruner.FullNode); header == nil && ok {
 		return fmt.Errorf("block cannot be empty for Full node pruner")
@@ -136,7 +134,7 @@ func (s *StorageState) StoreTrie(ts *rtstorage.TrieState, header *types.Header)
 
 	logger.Trace("cached trie in storage state", "root", root)
 
-	if err := s.tries[root].WriteDirty(s.db); err != nil {
+	if err := ts.Trie().WriteDirty(s.db); err != nil {
 		logger.Warn("failed to write trie to database", "root", root, "error", err)
 		return err
 	}
@@ -156,20 +154,21 @@ func (s *StorageState) TrieState(root *common.Hash) (*rtstorage.TrieState, error
 		root = &sr
 	}
 
-	s.lock.RLock()
-	t := s.tries[*root]
-	s.lock.RUnlock()
-
-	if t != nil && t.MustHash() != *root {
-		panic("trie does not have expected root")
-	}
-
-	if t == nil {
+	st, has := s.tries.Load(*root)
+	if !has {
 		var err error
-		t, err = s.LoadFromDB(*root)
+		st, err = s.LoadFromDB(*root)
 		if err != nil {
 			return nil, err
 		}
+
+		_, _ = s.tries.LoadOrStore(*root, st)
+	}
+
+	t := st.(*trie.Trie)
+
+	if has && t.MustHash() != *root {
+		panic("trie does not have expected root")
 	}
 
 	nextTrie := t.Snapshot()
@@ -178,10 +177,20 @@ func (s *StorageState) TrieState(root *common.Hash) (*rtstorage.TrieState, error
 		return nil, err
 	}
 
-	logger.Trace("returning trie to be modified", "root", root, "next", next.MustRoot())
+	logger.Trace("returning trie to be modified", "root", root)
 	return next, nil
 }
 
+// BeginModifyTrie should be called before the trie is modified
+func (s *StorageState) BeginModifyTrie() {
+	s.Lock()
+}
+
+// FinishModifyTrie should be called after the trie is done being modified
+func (s *StorageState) FinishModifyTrie() {
+	s.Unlock()
+}
+
 // LoadFromDB loads an encoded trie from the DB where the key is `root`
 func (s *StorageState) LoadFromDB(root common.Hash) (*trie.Trie, error) {
 	t := trie.NewEmptyTrie()
@@ -190,13 +199,31 @@ func (s *StorageState) LoadFromDB(root common.Hash) (*trie.Trie, error) {
 		return nil, err
 	}
 
-	s.lock.Lock()
-	defer s.lock.Unlock()
-
-	s.tries[t.MustHash()] = t
+	_, _ = s.tries.LoadOrStore(t.MustHash(), t)
 	return t, nil
 }
 
+func (s *StorageState) loadTrie(root *common.Hash) (*trie.Trie, error) {
+	if root == nil {
+		sr, err := s.blockState.BestBlockStateRoot()
+		if err != nil {
+			return nil, err
+		}
+		root = &sr
+	}
+
+	if t, has := s.tries.Load(*root); !has && t != nil {
+		return t.(*trie.Trie), nil
+	}
+
+	tr, err := s.LoadFromDB(*root)
+	if err != nil {
+		return nil, errTrieDoesNotExist(*root)
+	}
+
+	return tr, nil
+}
+
 // 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(root *common.Hash, key []byte) (bool, error) {
@@ -215,11 +242,8 @@ func (s *StorageState) GetStorage(root *common.Hash, key []byte) ([]byte, error)
 		root = &sr
 	}
 
-	s.lock.RLock()
-	defer s.lock.RUnlock()
-
-	if trie, ok := s.tries[*root]; ok {
-		val := trie.Get(key)
+	if t, has := s.tries.Load(*root); has {
+		val := t.(*trie.Trie).Get(key)
 		return val, nil
 	}
 
@@ -258,109 +282,41 @@ func (s *StorageState) EnumeratedTrieRoot(values [][]byte) {
 
 // Entries returns Entries from the trie with the given state root
 func (s *StorageState) Entries(root *common.Hash) (map[string][]byte, error) {
-	if root == nil {
-		head, err := s.blockState.BestBlockStateRoot()
-		if err != nil {
-			return nil, err
-		}
-		root = &head
-	}
-
-	s.lock.RLock()
-	tr, ok := s.tries[*root]
-	s.lock.RUnlock()
-
-	if !ok {
-		var err error
-		tr, err = s.LoadFromDB(*root)
-		if err != nil {
-			return nil, errTrieDoesNotExist(*root)
-		}
+	tr, err := s.loadTrie(root)
+	if err != nil {
+		return nil, err
 	}
 
-	s.lock.RLock()
-	defer s.lock.RUnlock()
 	return tr.Entries(), nil
 }
 
 // GetKeysWithPrefix returns all that match the given prefix for the given hash (or best block state root if hash is nil) in lexicographic order
-func (s *StorageState) GetKeysWithPrefix(hash *common.Hash, prefix []byte) ([][]byte, error) {
-	if hash == nil {
-		sr, err := s.blockState.BestBlockStateRoot()
-		if err != nil {
-			return nil, err
-		}
-		hash = &sr
-	}
-
-	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)
-		}
+func (s *StorageState) GetKeysWithPrefix(root *common.Hash, prefix []byte) ([][]byte, error) {
+	tr, err := s.loadTrie(root)
+	if err != nil {
+		return nil, err
 	}
 
-	s.lock.RLock()
-	defer s.lock.RUnlock()
 	return tr.GetKeysWithPrefix(prefix), nil
 }
 
-// GetStorageChild return GetChild from the trie
-func (s *StorageState) GetStorageChild(hash *common.Hash, keyToChild []byte) (*trie.Trie, error) {
-	if hash == nil {
-		sr, err := s.blockState.BestBlockStateRoot()
-		if err != nil {
-			return nil, err
-		}
-		hash = &sr
-	}
-
-	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)
-		}
+// GetStorageChild returns a child trie, if it exists
+func (s *StorageState) GetStorageChild(root *common.Hash, keyToChild []byte) (*trie.Trie, error) {
+	tr, err := s.loadTrie(root)
+	if err != nil {
+		return nil, err
 	}
 
-	s.lock.RLock()
-	defer s.lock.RUnlock()
 	return tr.GetChild(keyToChild)
 }
 
-// GetStorageFromChild return GetFromChild from the trie
-func (s *StorageState) GetStorageFromChild(hash *common.Hash, keyToChild, key []byte) ([]byte, error) {
-	if hash == nil {
-		sr, err := s.blockState.BestBlockStateRoot()
-		if err != nil {
-			return nil, err
-		}
-		hash = &sr
-	}
-
-	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)
-		}
+// GetStorageFromChild get a value from a child trie
+func (s *StorageState) GetStorageFromChild(root *common.Hash, keyToChild, key []byte) ([]byte, error) {
+	tr, err := s.loadTrie(root)
+	if err != nil {
+		return nil, err
 	}
 
-	s.lock.RLock()
-	defer s.lock.RUnlock()
 	return tr.GetFromChild(keyToChild, key)
 }