From 55d997fcbaefbca3c4f7337ff229d212e93eea92 Mon Sep 17 00:00:00 2001
From: noot <36753753+noot@users.noreply.github.com>
Date: Tue, 9 Nov 2021 12:36:54 -0500
Subject: [PATCH] fix(dot/state): update
 `*state.BlockState.AddBlockToBlockTree` to store block in
 `unfinalisedBlocksMap` (#2006)

---
 dot/state/block.go              | 11 ++++-------
 dot/state/block_finalisation.go |  4 ++--
 dot/state/block_test.go         |  5 ++++-
 dot/sync/chain_processor.go     |  2 +-
 dot/sync/interface.go           |  2 +-
 dot/sync/mocks/block_state.go   | 10 +++++-----
 6 files changed, 17 insertions(+), 17 deletions(-)

diff --git a/dot/state/block.go b/dot/state/block.go
index 09632c733a..d21108647e 100644
--- 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
index 2e85ac37d4..473a9674c1 100644
--- 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
index e6469531e3..bf27a2bf61 100644
--- 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
index bc7972ea45..7216e8ba4b 100644
--- 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
index 8d75922f47..5b2941d1b1 100644
--- 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
index 2a3dcd60bc..a8ef319dc5 100644
--- a/dot/sync/mocks/block_state.go
+++ b/dot/sync/mocks/block_state.go
@@ -32,13 +32,13 @@ func (_m *BlockState) AddBlock(_a0 *types.Block) error {
 	return r0
 }
 
-// AddBlockToBlockTree provides a mock function with given fields: header
-func (_m *BlockState) AddBlockToBlockTree(header *types.Header) error {
-	ret := _m.Called(header)
+// AddBlockToBlockTree provides a mock function with given fields: block
+func (_m *BlockState) AddBlockToBlockTree(block *types.Block) error {
+	ret := _m.Called(block)
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func(*types.Header) error); ok {
-		r0 = rf(header)
+	if rf, ok := ret.Get(0).(func(*types.Block) error); ok {
+		r0 = rf(block)
 	} else {
 		r0 = ret.Error(0)
 	}