From 49e8cb8297df6b3cb411ac1224c4c7588f85e264 Mon Sep 17 00:00:00 2001
From: Lazar <12626340+Lazar955@users.noreply.github.com>
Date: Mon, 16 Sep 2024 16:33:56 +0200
Subject: [PATCH] chore(btcclient): reduce rpc calls (#51)

Reduces the number of RPC calls in methods by combining them.

[References issue](https://github.com/babylonchain/vigilante/issues/218)
---
 btcclient/interface.go                 |  2 +-
 btcclient/query.go                     | 41 ++++++++++++++++++--------
 monitor/btcscanner/btc_scanner.go      |  2 +-
 monitor/btcscanner/btc_scanner_test.go |  2 +-
 reporter/bootstrapping.go              |  4 +--
 testutil/mocks/btcclient.go            |  9 +++---
 6 files changed, 38 insertions(+), 22 deletions(-)

diff --git a/btcclient/interface.go b/btcclient/interface.go
index f9d89aa..5324dea 100644
--- a/btcclient/interface.go
+++ b/btcclient/interface.go
@@ -14,7 +14,7 @@ import (
 type BTCClient interface {
 	Stop()
 	WaitForShutdown()
-	GetBestBlock() (*chainhash.Hash, uint64, error)
+	GetBestBlock() (uint64, error)
 	GetBlockByHash(blockHash *chainhash.Hash) (*types.IndexedBlock, *wire.MsgBlock, error)
 	FindTailBlocksByHeight(height uint64) ([]*types.IndexedBlock, error)
 	GetBlockByHeight(height uint64) (*types.IndexedBlock, *wire.MsgBlock, error)
diff --git a/btcclient/query.go b/btcclient/query.go
index 512cb57..ad5057f 100644
--- a/btcclient/query.go
+++ b/btcclient/query.go
@@ -12,20 +12,14 @@ import (
 	"github.com/babylonlabs-io/vigilante/types"
 )
 
-// GetBestBlock provides similar functionality with the btcd.rpcclient.GetBestBlock function
-// We implement this, because this function is only provided by btcd.
-// TODO: replace two rpc calls with only one c.GetBlockCount
-func (c *Client) GetBestBlock() (*chainhash.Hash, uint64, error) {
-	btcLatestBlockHash, err := c.getBestBlockHashWithRetry()
+// GetBestBlock returns the height of the best block
+func (c *Client) GetBestBlock() (uint64, error) {
+	height, err := c.getBlockCountWithRetry()
 	if err != nil {
-		return nil, 0, err
+		return 0, err
 	}
-	btcLatestBlock, err := c.getBlockVerboseWithRetry(btcLatestBlockHash)
-	if err != nil {
-		return nil, 0, err
-	}
-	btcLatestBlockHeight := uint64(btcLatestBlock.Height)
-	return btcLatestBlockHash, btcLatestBlockHeight, nil
+
+	return uint64(height), nil
 }
 
 func (c *Client) GetBlockByHash(blockHash *chainhash.Hash) (*types.IndexedBlock, *wire.MsgBlock, error) {
@@ -213,3 +207,26 @@ func (c *Client) FindTailBlocksByHeight(baseHeight uint64) ([]*types.IndexedBloc
 
 	return c.getChainBlocks(baseHeight, tipIb)
 }
+
+func (c *Client) getBlockCountWithRetry() (int64, error) {
+	var (
+		height int64
+		err    error
+	)
+
+	if err = retry.Do(func() error {
+		height, err = c.GetBlockCount()
+		if err != nil {
+			return err
+		}
+		return nil
+	},
+		retry.Delay(c.retrySleepTime),
+		retry.MaxDelay(c.maxRetrySleepTime),
+		retry.Attempts(c.maxRetryTimes),
+	); err != nil {
+		c.logger.Debug("failed to query get block count", zap.Error(err))
+	}
+
+	return height, nil
+}
diff --git a/monitor/btcscanner/btc_scanner.go b/monitor/btcscanner/btc_scanner.go
index fa50387..e18e239 100644
--- a/monitor/btcscanner/btc_scanner.go
+++ b/monitor/btcscanner/btc_scanner.go
@@ -152,7 +152,7 @@ func (bs *BtcScanner) Bootstrap() {
 	// clear all the blocks in the cache to avoid forks
 	bs.unconfirmedBlockCache.RemoveAll()
 
-	_, bestHeight, err := bs.btcClient.GetBestBlock()
+	bestHeight, err := bs.btcClient.GetBestBlock()
 	if err != nil {
 		panic(fmt.Errorf("cannot get the best BTC block %w", err))
 	}
diff --git a/monitor/btcscanner/btc_scanner_test.go b/monitor/btcscanner/btc_scanner_test.go
index 698dfd8..62b85f5 100644
--- a/monitor/btcscanner/btc_scanner_test.go
+++ b/monitor/btcscanner/btc_scanner_test.go
@@ -30,7 +30,7 @@ func FuzzBootStrap(f *testing.F) {
 		defer ctl.Finish()
 		mockBtcClient := mocks.NewMockBTCClient(ctl)
 		confirmedBlocks := chainIndexedBlocks[:numBlocks-k]
-		mockBtcClient.EXPECT().GetBestBlock().Return(nil, uint64(bestHeight), nil)
+		mockBtcClient.EXPECT().GetBestBlock().Return(uint64(bestHeight), nil)
 		for i := 0; i < int(numBlocks); i++ {
 			mockBtcClient.EXPECT().GetBlockByHeight(gomock.Eq(uint64(chainIndexedBlocks[i].Height))).
 				Return(chainIndexedBlocks[i], nil, nil).AnyTimes()
diff --git a/reporter/bootstrapping.go b/reporter/bootstrapping.go
index 729fb1e..24382fd 100644
--- a/reporter/bootstrapping.go
+++ b/reporter/bootstrapping.go
@@ -233,7 +233,7 @@ func (r *Reporter) waitUntilBTCSync() error {
 	)
 
 	// Retrieve hash/height of the latest block in BTC
-	btcLatestBlockHash, btcLatestBlockHeight, err = r.btcClient.GetBestBlock()
+	btcLatestBlockHeight, err = r.btcClient.GetBestBlock()
 	if err != nil {
 		return err
 	}
@@ -264,7 +264,7 @@ func (r *Reporter) waitUntilBTCSync() error {
 		// When BTC catches up, break and continue the bootstrapping process
 		ticker := time.NewTicker(5 * time.Second) // TODO: parameterise the polling interval
 		for range ticker.C {
-			_, btcLatestBlockHeight, err = r.btcClient.GetBestBlock()
+			btcLatestBlockHeight, err = r.btcClient.GetBestBlock()
 			if err != nil {
 				return err
 			}
diff --git a/testutil/mocks/btcclient.go b/testutil/mocks/btcclient.go
index 240994f..105f29f 100644
--- a/testutil/mocks/btcclient.go
+++ b/testutil/mocks/btcclient.go
@@ -56,13 +56,12 @@ func (mr *MockBTCClientMockRecorder) FindTailBlocksByHeight(height interface{})
 }
 
 // GetBestBlock mocks base method.
-func (m *MockBTCClient) GetBestBlock() (*chainhash.Hash, uint64, error) {
+func (m *MockBTCClient) GetBestBlock() (uint64, error) {
 	m.ctrl.T.Helper()
 	ret := m.ctrl.Call(m, "GetBestBlock")
-	ret0, _ := ret[0].(*chainhash.Hash)
-	ret1, _ := ret[1].(uint64)
-	ret2, _ := ret[2].(error)
-	return ret0, ret1, ret2
+	ret0, _ := ret[0].(uint64)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
 }
 
 // GetBestBlock indicates an expected call of GetBestBlock.