chore(dot/state): replace sync.Map with map+mutex

dot/state/block.go

@@ -59,7 +59,7 @@ type BlockState struct {
 	sync.RWMutex
 	genesisHash       common.Hash
 	lastFinalised     common.Hash
-	unfinalisedBlocks *sync.Map // map[common.Hash]*types.Block
+	unfinalisedBlocks *hashToBlock
 
 	// block notifiers
 	imported                       map[chan *types.Block]struct{}
@@ -80,7 +80,7 @@ func NewBlockState(db chaindb.Database, telemetry telemetry.Client) (*BlockState
 		dbPath:                     db.Path(),
 		baseState:                  NewBaseState(db),
 		db:                         chaindb.NewTable(db, blockPrefix),
-		unfinalisedBlocks:          new(sync.Map),
+		unfinalisedBlocks:          newHashToBlock(),
 		imported:                   make(map[chan *types.Block]struct{}),
 		finalised:                  make(map[chan *types.FinalisationInfo]struct{}),
 		pruneKeyCh:                 make(chan *types.Header, pruneKeyBufferSize),
@@ -113,7 +113,7 @@ func NewBlockStateFromGenesis(db chaindb.Database, header *types.Header,
 		bt:                         blocktree.NewBlockTreeFromRoot(header),
 		baseState:                  NewBaseState(db),
 		db:                         chaindb.NewTable(db, blockPrefix),
-		unfinalisedBlocks:          new(sync.Map),
+		unfinalisedBlocks:          newHashToBlock(),
 		imported:                   make(map[chan *types.Block]struct{}),
 		finalised:                  make(map[chan *types.FinalisationInfo]struct{}),
 		pruneKeyCh:                 make(chan *types.Header, pruneKeyBufferSize),
@@ -187,12 +187,12 @@ func (bs *BlockState) GenesisHash() common.Hash {
 }
 
 func (bs *BlockState) storeUnfinalisedBlock(block *types.Block) {
-	bs.unfinalisedBlocks.Store(block.Header.Hash(), block)
+	bs.unfinalisedBlocks.set(block.Header.Hash(), block)
 }
 
 func (bs *BlockState) hasUnfinalisedBlock(hash common.Hash) bool {
-	_, has := bs.unfinalisedBlocks.Load(hash)
-	return has
+	block := bs.unfinalisedBlocks.get(hash)
+	return block != nil
 }
 
 func (bs *BlockState) getUnfinalisedHeader(hash common.Hash) (*types.Header, bool) {
@@ -205,22 +205,13 @@ func (bs *BlockState) getUnfinalisedHeader(hash common.Hash) (*types.Header, boo
 }
 
 func (bs *BlockState) getUnfinalisedBlock(hash common.Hash) (*types.Block, bool) {
-	block, has := bs.unfinalisedBlocks.Load(hash)
-	if !has {
+	block := bs.unfinalisedBlocks.get(hash)
+	if block == nil {
 		return nil, false
 	}
 
 	// TODO: dot/core tx re-org test seems to abort here due to block body being invalid?
-	return block.(*types.Block), true
-}
-
-func (bs *BlockState) getAndDeleteUnfinalisedBlock(hash common.Hash) (*types.Block, bool) {
-	block, has := bs.unfinalisedBlocks.LoadAndDelete(hash)
-	if !has {
-		return nil, false
-	}
-
-	return block.(*types.Block), true
+	return block, true
 }
 
 // HasHeader returns if the db contains a header with the given hash

dot/state/block_finalisation.go

@@ -146,13 +146,13 @@ func (bs *BlockState) SetFinalisedHash(hash common.Hash, round, setID uint64) er
 
 	pruned := bs.bt.Prune(hash)
 	for _, hash := range pruned {
-		block, has := bs.getAndDeleteUnfinalisedBlock(hash)
-		if !has {
+		blockHeader := bs.unfinalisedBlocks.delete(hash)
+		if blockHeader == nil {
 			continue
 		}
 
-		logger.Tracef("pruned block number %s with hash %s", block.Header.Number, hash)
-		bs.pruneKeyCh <- &block.Header
+		logger.Tracef("pruned block number %s with hash %s", blockHeader.Number, hash)
+		bs.pruneKeyCh <- blockHeader
 	}
 
 	// if nothing was previously finalised, set the first slot of the network to the
@@ -233,13 +233,13 @@ func (bs *BlockState) handleFinalisedBlock(curr common.Hash) error {
 		}
 
 		// delete from the unfinalisedBlockMap and delete reference to in-memory trie
-		block, has = bs.getAndDeleteUnfinalisedBlock(hash)
-		if !has {
+		blockHeader := bs.unfinalisedBlocks.delete(hash)
+		if blockHeader == nil {
 			continue
 		}
 
-		logger.Tracef("cleaned out finalised block from memory; block number %s with hash %s", block.Header.Number, hash)
-		bs.pruneKeyCh <- &block.Header
+		logger.Tracef("cleaned out finalised block from memory; block number %s with hash %s", blockHeader.Number, hash)
+		bs.pruneKeyCh <- blockHeader
 	}
 
 	return batch.Flush()

dot/state/hashtoblock.go

@@ -0,0 +1,45 @@
+// Copyright 2022 ChainSafe Systems (ON)
+// SPDX-License-Identifier: LGPL-3.0-only
+
+package state
+
+import (
+	"sync"
+
+	"github.com/ChainSafe/gossamer/dot/types"
+	"github.com/ChainSafe/gossamer/lib/common"
+)
+
+type hashToBlock struct {
+	mutex   sync.RWMutex
+	mapping map[common.Hash]*types.Block
+}
+
+func newHashToBlock() *hashToBlock {
+	return &hashToBlock{
+		mapping: make(map[common.Hash]*types.Block),
+	}
+}
+
+func (h *hashToBlock) get(hash common.Hash) (block *types.Block) {
+	h.mutex.RLock()
+	defer h.mutex.RUnlock()
+	return h.mapping[hash]
+}
+
+func (h *hashToBlock) set(hash common.Hash, block *types.Block) {
+	h.mutex.Lock()
+	defer h.mutex.Unlock()
+	h.mapping[hash] = block
+}
+
+func (h *hashToBlock) delete(hash common.Hash) (deletedHeader *types.Header) {
+	h.mutex.Lock()
+	defer h.mutex.Unlock()
+	block := h.mapping[hash]
+	delete(h.mapping, hash)
+	if block == nil {
+		return nil
+	}
+	return &block.Header
+}