fix(dot/state): update *state.BlockState.AddBlockToBlockTree to sto…

…re block in `unfinalisedBlocksMap` (ChainSafe#2006)
timwu20 · Nov 9, 2021 · 55d997f · 55d997f
1 parent 6b153e9
commit 55d997f
Show file tree

Hide file tree

Showing 6 changed files with 17 additions and 17 deletions.
diff --git a/dot/state/block.go b/dot/state/block.go
@@ -221,10 +221,6 @@ func (bs *BlockState) getAndDeleteUnfinalisedBlock(hash common.Hash) (*types.Blo
 	return block.(*types.Block), true
 }
 
-func (bs *BlockState) deleteUnfinalisedBlock(hash common.Hash) {
-	bs.unfinalisedBlocks.Delete(hash)
-}
-
 // HasHeader returns if the db contains a header with the given hash
 func (bs *BlockState) HasHeader(hash common.Hash) (bool, error) {
 	if bs.hasUnfinalisedBlock(hash) {
@@ -434,16 +430,17 @@ func (bs *BlockState) AddBlockWithArrivalTime(block *types.Block, arrivalTime ti
 
 // AddBlockToBlockTree adds the given block to the blocktree. It does not write it to the database.
 // TODO: remove this func and usage from sync (after sync refactor?)
-func (bs *BlockState) AddBlockToBlockTree(header *types.Header) error {
+func (bs *BlockState) AddBlockToBlockTree(block *types.Block) error {
 	bs.Lock()
 	defer bs.Unlock()
 
-	arrivalTime, err := bs.GetArrivalTime(header.Hash())
+	arrivalTime, err := bs.GetArrivalTime(block.Header.Hash())
 	if err != nil {
 		arrivalTime = time.Now()
 	}
 
-	return bs.bt.AddBlock(header, arrivalTime)
+	bs.storeUnfinalisedBlock(block)
+	return bs.bt.AddBlock(&block.Header, arrivalTime)
 }
 
 // GetAllBlocksAtNumber returns all unfinalised blocks with the given number

diff --git a/dot/state/block_finalisation.go b/dot/state/block_finalisation.go
@@ -194,7 +194,6 @@ func (bs *BlockState) SetFinalisedHash(hash common.Hash, round, setID uint64) er
 		return fmt.Errorf("could not send 'notify.finalized' telemetry message, error: %s", err)
 	}
 
-	// return bs.setHighestRoundAndSetID(round, setID)
 	bs.lastFinalised = hash
 	return nil
 }
@@ -252,7 +251,8 @@ func (bs *BlockState) handleFinalisedBlock(curr common.Hash) error {
 			return err
 		}
 
-		bs.deleteUnfinalisedBlock(hash)
+		// the block will be deleted from the unfinalisedBlockMap in the pruning loop
+		// in `SetFinalisedHash()`, which calls this function
 	}
 
 	return batch.Flush()

diff --git a/dot/state/block_test.go b/dot/state/block_test.go
@@ -470,7 +470,10 @@ func TestAddBlockToBlockTree(t *testing.T) {
 	err := bs.setArrivalTime(header.Hash(), time.Now())
 	require.NoError(t, err)
 
-	err = bs.AddBlockToBlockTree(header)
+	err = bs.AddBlockToBlockTree(&types.Block{
+		Header: *header,
+		Body:   types.Body{},
+	})
 	require.NoError(t, err)
 	require.Equal(t, bs.BestBlockHash(), header.Hash())
 }

diff --git a/dot/sync/chain_processor.go b/dot/sync/chain_processor.go
@@ -130,7 +130,7 @@ func (s *chainProcessor) processBlockData(bd *types.BlockData) error {
 
 		logger.Debug("skipping block, already have", "hash", bd.Hash, "number", block.Header.Number)
 
-		err = s.blockState.AddBlockToBlockTree(&block.Header)
+		err = s.blockState.AddBlockToBlockTree(block)
 		if errors.Is(err, blocktree.ErrBlockExists) {
 			return nil
 		} else if err != nil {

diff --git a/dot/sync/interface.go b/dot/sync/interface.go
@@ -50,7 +50,7 @@ type BlockState interface {
 	GetJustification(common.Hash) ([]byte, error)
 	SetJustification(hash common.Hash, data []byte) error
 	SetFinalisedHash(hash common.Hash, round, setID uint64) error
-	AddBlockToBlockTree(header *types.Header) error
+	AddBlockToBlockTree(block *types.Block) error
 	GetHashByNumber(*big.Int) (common.Hash, error)
 	GetBlockByHash(common.Hash) (*types.Block, error)
 	GetRuntime(*common.Hash) (runtime.Instance, error)

diff --git a/dot/sync/mocks/block_state.go b/dot/sync/mocks/block_state.go