chainio: add method CurrentHeight

Add a new method `CurrentHeight` to query the current best height of the
dispatcher.

Author: yyforyongyu

chainio/dispatcher.go

@@ -49,14 +49,35 @@ type BlockbeatDispatcher struct {
 
 	// quit is used to signal the BlockbeatDispatcher to stop.
 	quit chan struct{}
+
+	// queryHeightChan is used to receive queries on the current height of
+	// the dispatcher.
+	queryHeightChan chan *query
+}
+
+// query is used to fetch the internal state of the dispatcher.
+type query struct {
+	// respChan is used to send back the current height back to the caller.
+	//
+	// NOTE: This channel must be buffered.
+	respChan chan int32
+}
+
+// newQuery creates a query to be used to fetch the internal state of the
+// dispatcher.
+func newQuery() *query {
+	return &query{
+		respChan: make(chan int32, 1),
+	}
 }
 
 // NewBlockbeatDispatcher returns a new blockbeat dispatcher instance.
 func NewBlockbeatDispatcher(n chainntnfs.ChainNotifier) *BlockbeatDispatcher {
 	return &BlockbeatDispatcher{
-		notifier:       n,
-		quit:           make(chan struct{}),
-		consumerQueues: make(map[uint32][]Consumer),
+		notifier:        n,
+		quit:            make(chan struct{}),
+		consumerQueues:  make(map[uint32][]Consumer),
+		queryHeightChan: make(chan *query, 1),
 	}
 }
 
@@ -161,6 +182,18 @@ func (b *BlockbeatDispatcher) dispatchBlocks(
 			b.log().Infof("Notified all consumers on new block "+
 				"in %v", time.Since(start))
 
+		// A query has been made to fetch the current height, we now
+		// send the height from its current beat.
+		case query := <-b.queryHeightChan:
+			// The beat may not be set yet, e.g., during the startup
+			// the query is made before the block epoch being sent.
+			height := int32(0)
+			if b.beat != nil {
+				height = b.beat.Height()
+			}
+
+			query.respChan <- height
+
 		case <-b.quit:
 			b.log().Debugf("BlockbeatDispatcher quit signal " +
 				"received")
@@ -170,6 +203,30 @@ func (b *BlockbeatDispatcher) dispatchBlocks(
 	}
 }
 
+// CurrentHeight returns the current best height known to the dispatcher. 0 is
+// returned if the dispatcher is shutting down.
+func (b *BlockbeatDispatcher) CurrentHeight() int32 {
+	query := newQuery()
+
+	select {
+	case b.queryHeightChan <- query:
+
+	case <-b.quit:
+		clog.Debugf("BlockbeatDispatcher quit before query")
+		return 0
+	}
+
+	select {
+	case height := <-query.respChan:
+		clog.Debugf("Responded current height: %v", height)
+		return height
+
+	case <-b.quit:
+		clog.Debugf("BlockbeatDispatcher quit before response")
+		return 0
+	}
+}
+
 // notifyQueues notifies each queue concurrently about the latest block epoch.
 func (b *BlockbeatDispatcher) notifyQueues() error {
 	// errChans is a map of channels that will be used to receive errors

chainio/dispatcher_test.go

@@ -381,3 +381,60 @@ func TestNotifyQueuesError(t *testing.T) {
 	err := b.notifyQueues()
 	require.ErrorIs(t, err, errDummy)
 }
+
+// TestCurrentHeight asserts `CurrentHeight` returns the expected block height.
+func TestCurrentHeight(t *testing.T) {
+	t.Parallel()
+
+	testHeight := int32(1000)
+
+	// Create a mock chain notifier.
+	mockNotifier := &chainntnfs.MockChainNotifier{}
+	defer mockNotifier.AssertExpectations(t)
+
+	// Create a mock beat.
+	mockBeat := &MockBlockbeat{}
+	defer mockBeat.AssertExpectations(t)
+	mockBeat.On("logger").Return(clog)
+	mockBeat.On("Height").Return(testHeight).Once()
+
+	// Create a mock consumer.
+	consumer := &MockConsumer{}
+	defer consumer.AssertExpectations(t)
+	consumer.On("Name").Return("mocker1")
+
+	// Create one queue.
+	queue := []Consumer{consumer}
+
+	// Create a new dispatcher.
+	b := NewBlockbeatDispatcher(mockNotifier)
+
+	// Register the queues.
+	b.RegisterQueue(queue)
+
+	// Attach the blockbeat.
+	b.beat = mockBeat
+
+	// Mock the chain notifier to return a valid notifier.
+	blockEpochs := &chainntnfs.BlockEpochEvent{
+		Cancel: func() {},
+	}
+	mockNotifier.On("RegisterBlockEpochNtfn",
+		mock.Anything).Return(blockEpochs, nil).Once()
+
+	// Start the dispatcher now should not return an error.
+	err := b.Start()
+	require.NoError(t, err)
+
+	// Make a query on the current height and assert it equals to
+	// testHeight.
+	height := b.CurrentHeight()
+	require.Equal(t, testHeight, height)
+
+	// Stop the dispatcher.
+	b.Stop()
+
+	// Make a query on the current height and assert it equals to 0.
+	height = b.CurrentHeight()
+	require.Zero(t, height)
+}