diff --git a/chainio/dispatcher.go b/chainio/dispatcher.go index 269bb1892..3f22948bd 100644 --- a/chainio/dispatcher.go +++ b/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 diff --git a/chainio/dispatcher_test.go b/chainio/dispatcher_test.go index 1cab5b998..eb1dedb6f 100644 --- a/chainio/dispatcher_test.go +++ b/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) +}