feat(dot/state): create Range to traverse the blocktree and the blocks in the disk

dot/state/block.go

@@ -38,6 +38,7 @@ var (
 	messageQueuePrefix  = []byte("mqp") // messageQueuePrefix + hash -> message queue
 	justificationPrefix = []byte("jcp") // justificationPrefix + hash -> justification
 
+	errNilBlockTree = errors.New("blocktree is nil")
 	errNilBlockBody = errors.New("block body is nil")
 
 	syncedBlocksGauge = promauto.NewGauge(prometheus.GaugeOpts{
@@ -542,10 +543,132 @@ func (bs *BlockState) GetSlotForBlock(hash common.Hash) (uint64, error) {
 	return types.GetSlotFromHeader(header)
 }
 
+func (bs *BlockState) loadHeaderFromDisk(hash common.Hash) (header *types.Header, err error) {
+	startHeaderData, err := bs.db.Get(headerKey(hash))
+	if err != nil {
+		return nil, fmt.Errorf("querying database: %w", err)
+	}
+
+	header = types.NewEmptyHeader()
+	err = scale.Unmarshal(startHeaderData, header)
+	if err != nil {
+		return nil, fmt.Errorf("unmarshaling start header: %w", err)
+	}
+
+	if header.Empty() {
+		return nil, fmt.Errorf("%w: %s", chaindb.ErrKeyNotFound, hash)
+	}
+
+	return header, nil
+}
+
+// Range returns the sub-blockchain between the starting hash and the
+// ending hash using both block tree and disk
+func (bs *BlockState) Range(startHash, endHash common.Hash) (hashes []common.Hash, err error) {
+	if bs.bt == nil {
+		return nil, errNilBlockTree
+	}
+
+	endHeader, err := bs.loadHeaderFromDisk(endHash)
+	if errors.Is(err, chaindb.ErrKeyNotFound) {
+		// end hash is not in the disk so we should lookup
+		// block that could be in memory and in the disk as well
+		return bs.retrieveRange(startHash, endHash)
+	} else if err != nil {
+		return nil, fmt.Errorf("retrieving end hash from disk: %w", err)
+	}
+
+	// end hash was found in the disk, that means all the blocks
+	// between start and end can be found in the disk
+	return bs.retrieveRangeFromDisk(startHash, endHeader)
+}
+
+func (bs *BlockState) retrieveRange(startHash, endHash common.Hash) (hashes []common.Hash, err error) {
+	inMemoryHashes, err := bs.bt.SubBlockchain(startHash, endHash)
+	if err != nil {
+		return nil, fmt.Errorf("retrieving range from in-memory blocktree: %w", err)
+	}
+
+	firstItem := inMemoryHashes[0]
+
+	// if the first item is equal to the startHash that means we got the range
+	// from the in-memory blocktree
+	if firstItem == startHash {
+		return inMemoryHashes, nil
+	}
+
+	// since we got as many blocks as we could from
+	// the block tree but still missing blocks to
+	// fulfil the range we should lookup in the
+	// disk for the remaining ones, the first item in the hashes array
+	// must be the block tree root that is also placed in the disk
+	//  so we will start from its parent since it is already in the array
+	blockTreeRootHeader, err := bs.loadHeaderFromDisk(firstItem)
+	if err != nil {
+		return nil, fmt.Errorf("loading block tree root from disk: %w", err)
+	}
+
+	startingAtParentHeader, err := bs.loadHeaderFromDisk(blockTreeRootHeader.ParentHash)
+	if err != nil {
+		return nil, fmt.Errorf("range end should be in database: %w", err)
+	}
+
+	inDiskHashes, err := bs.retrieveRangeFromDisk(startHash, startingAtParentHeader)
+	if err != nil {
+		return nil, fmt.Errorf("retrieving range from disk: %w", err)
+	}
+
+	hashes = append(inDiskHashes, inMemoryHashes...)
+	return hashes, nil
+}
+
+var ErrStartHashMismatch = errors.New("start hash mismatch")
+
+// retrieveRangeFromDisk takes the start and the end and will retrieve all block in between
+// where all blocks (start and end inclusive) are supposed to be placed at disk
+func (bs *BlockState) retrieveRangeFromDisk(startHash common.Hash,
+	endHeader *types.Header) (hashes []common.Hash, err error) {
+	startHeader, err := bs.loadHeaderFromDisk(startHash)
+	if err != nil {
+		return nil, fmt.Errorf("range start should be in database: %w", err)
+	}
+
+	// blocksInRange is the difference between the end number to start number
+	// but the difference don't includes the start item that is why we add 1
+	blocksInRange := (endHeader.Number - startHeader.Number) + 1
+
+	hashes = make([]common.Hash, blocksInRange)
+
+	lastPosition := blocksInRange - 1
+
+	hashes[0] = startHash
+	hashes[lastPosition] = endHeader.Hash()
+
+	inLoopHash := endHeader.ParentHash
+	for currentPosition := lastPosition - 1; currentPosition > 0; currentPosition-- {
+		hashes[currentPosition] = inLoopHash
+
+		inLoopHeader, err := bs.loadHeaderFromDisk(inLoopHash)
+		if err != nil {
+			return nil, fmt.Errorf("retrieving hash %s from disk: %w", inLoopHash.Short(), err)
+		}
+
+		inLoopHash = inLoopHeader.ParentHash
+	}
+
+	// here we ensure that we finished up the loop
+	// with the same hash as the startHash
+	if inLoopHash != startHash {
+		return nil, fmt.Errorf("%w: expecting %s, found: %s", ErrStartHashMismatch, startHash.Short(), inLoopHash.Short())
+	}
+
+	return hashes, err
+}
+
 // SubChain returns the sub-blockchain between the starting hash and the ending hash using the block tree
 func (bs *BlockState) SubChain(start, end common.Hash) ([]common.Hash, error) {
 	if bs.bt == nil {
-		return nil, fmt.Errorf("blocktree is nil")
+		return nil, errNilBlockTree
 	}
 
 	return bs.bt.SubBlockchain(start, end)
@@ -555,7 +678,7 @@ func (bs *BlockState) SubChain(start, end common.Hash) ([]common.Hash, error) {
 // it returns an error if parent or child are not in the blocktree.
 func (bs *BlockState) IsDescendantOf(parent, child common.Hash) (bool, error) {
 	if bs.bt == nil {
-		return false, fmt.Errorf("blocktree is nil")
+		return false, errNilBlockTree
 	}
 
 	return bs.bt.IsDescendantOf(parent, child)

dot/state/block_test.go

@@ -4,6 +4,7 @@
 package state
 
 import (
+	"errors"
 	"testing"
 	"time"
 
@@ -599,3 +600,162 @@ func TestNumberIsFinalised(t *testing.T) {
 	require.NoError(t, err)
 	require.False(t, fin)
 }
+
+func TestRange(t *testing.T) {
+	t.Parallel()
+
+	loadHeaderFromDiskErr := errors.New("[mocked] cannot read, database closed ex")
+	testcases := map[string]struct {
+		blocksToCreate        int
+		blocksToPersistAtDisk int
+
+		newBlockState func(t *testing.T, ctrl *gomock.Controller,
+			genesisHeader *types.Header) *BlockState
+		wantErr   error
+		stringErr string
+	}{
+		"all_blocks_stored_in_disk": {
+			blocksToCreate:        128,
+			blocksToPersistAtDisk: 128,
+			newBlockState: func(t *testing.T, ctrl *gomock.Controller,
+				genesisHeader *types.Header) *BlockState {
+				telemetryMock := NewMockClient(ctrl)
+				telemetryMock.EXPECT().SendMessage(gomock.Any()).AnyTimes()
+
+				db := NewInMemoryDB(t)
+
+				blockState, err := NewBlockStateFromGenesis(db, newTriesEmpty(), genesisHeader, telemetryMock)
+				require.NoError(t, err)
+
+				return blockState
+			},
+		},
+
+		"all_blocks_persisted_in_blocktree": {
+			blocksToCreate:        128,
+			blocksToPersistAtDisk: 0,
+			newBlockState: func(t *testing.T, ctrl *gomock.Controller,
+				genesisHeader *types.Header) *BlockState {
+				telemetryMock := NewMockClient(ctrl)
+				telemetryMock.EXPECT().SendMessage(gomock.Any()).AnyTimes()
+
+				db := NewInMemoryDB(t)
+
+				blockState, err := NewBlockStateFromGenesis(db, newTriesEmpty(), genesisHeader, telemetryMock)
+				require.NoError(t, err)
+
+				return blockState
+			},
+		},
+
+		"half_blocks_placed_in_blocktree_half_stored_in_disk": {
+			blocksToCreate:        128,
+			blocksToPersistAtDisk: 64,
+			newBlockState: func(t *testing.T, ctrl *gomock.Controller,
+				genesisHeader *types.Header) *BlockState {
+				telemetryMock := NewMockClient(ctrl)
+				telemetryMock.EXPECT().SendMessage(gomock.Any()).AnyTimes()
+
+				db := NewInMemoryDB(t)
+
+				blockState, err := NewBlockStateFromGenesis(db, newTriesEmpty(), genesisHeader, telemetryMock)
+				require.NoError(t, err)
+
+				return blockState
+			},
+		},
+
+		"error_while_loading_header_from_disk": {
+			blocksToCreate:        1,
+			blocksToPersistAtDisk: 0,
+			wantErr:               loadHeaderFromDiskErr,
+			stringErr: "retrieving end hash from disk: " +
+				"querying database: [mocked] cannot read, database closed ex",
+			newBlockState: func(t *testing.T, ctrl *gomock.Controller,
+				genesisHeader *types.Header) *BlockState {
+				telemetryMock := NewMockClient(ctrl)
+				telemetryMock.EXPECT().SendMessage(gomock.Any()).AnyTimes()
+
+				db := NewInMemoryDB(t)
+				blockState, err := NewBlockStateFromGenesis(db, newTriesEmpty(), genesisHeader, telemetryMock)
+
+				mockedDb := NewMockDatabase(ctrl)
+				// cannot assert the exact hash type since the block header
+				// hash is generate by the running test case
+				mockedDb.EXPECT().Get(gomock.AssignableToTypeOf([]byte{})).
+					Return(nil, loadHeaderFromDiskErr)
+
+				blockState.db = mockedDb
+
+				require.NoError(t, err)
+				return blockState
+			},
+		},
+	}
+
+	for tname, tt := range testcases {
+		tt := tt
+
+		t.Run(tname, func(t *testing.T) {
+			t.Parallel()
+
+			if tt.blocksToCreate < tt.blocksToPersistAtDisk {
+				require.Fail(t, "blocksToPersistAtDisk should be lower or equal blocksToCreate")
+			}
+
+			ctrl := gomock.NewController(t)
+			genesisHeader := &types.Header{
+				Number:    0,
+				StateRoot: trie.EmptyHash,
+				Digest:    types.NewDigest(),
+			}
+
+			blockState := tt.newBlockState(t, ctrl, genesisHeader)
+
+			testBlockBody := *types.NewBody([]types.Extrinsic{[]byte{0, 1, 2, 3, 4, 5, 6, 7, 8, 9}})
+			hashes := make([]common.Hash, 0, tt.blocksToCreate)
+			previousHeaderHash := genesisHeader.Hash()
+			for blockNumber := 1; blockNumber <= tt.blocksToCreate; blockNumber++ {
+				currentHeader := &types.Header{
+					Number:     uint(blockNumber),
+					Digest:     createPrimaryBABEDigest(t),
+					ParentHash: previousHeaderHash,
+				}
+
+				block := &types.Block{
+					Header: *currentHeader,
+					Body:   testBlockBody,
+				}
+
+				err := blockState.AddBlock(block)
+				require.NoError(t, err)
+
+				hashes = append(hashes, currentHeader.Hash())
+				previousHeaderHash = currentHeader.Hash()
+			}
+
+			if tt.blocksToPersistAtDisk > 0 {
+				hashIndexToSetAsFinalized := tt.blocksToPersistAtDisk - 1
+				selectedHash := hashes[hashIndexToSetAsFinalized]
+
+				err := blockState.SetFinalisedHash(selectedHash, 0, 0)
+				require.NoError(t, err)
+			}
+
+			// execute the Range call. All the values returned must
+			// match the hashes we previsouly created
+			startHash := hashes[0]
+			endHash := hashes[len(hashes)-1]
+
+			hashesInRange, err := blockState.Range(startHash, endHash)
+			require.ErrorIs(t, err, tt.wantErr)
+			if tt.stringErr != "" {
+				require.EqualError(t, err, tt.stringErr)
+				require.Empty(t, hashesInRange)
+				return
+			}
+
+			require.Equal(t, hashes, hashesInRange)
+		})
+	}
+}