From 2cea0894043d2b00f3d898a26ac0c0b8eb3f5aee Mon Sep 17 00:00:00 2001
From: Edwin <edwin@am.is>
Date: Wed, 1 Nov 2023 11:46:43 +0800
Subject: [PATCH] pkg/exchange: add rate limiter for query ticker, account

---
 pkg/exchange/bitget/convert_test.go | 32 +++++++++++++++++++++++-
 pkg/exchange/bitget/exchange.go     | 38 ++++++++++++++++++++---------
 2 files changed, 58 insertions(+), 12 deletions(-)

diff --git a/pkg/exchange/bitget/convert_test.go b/pkg/exchange/bitget/convert_test.go
index 6ab315eaa..5a80a045a 100644
--- a/pkg/exchange/bitget/convert_test.go
+++ b/pkg/exchange/bitget/convert_test.go
@@ -10,7 +10,37 @@ import (
 	"github.com/c9s/bbgo/pkg/types"
 )
 
-func TestToGlobalMarket(t *testing.T) {
+func Test_toGlobalBalance(t *testing.T) {
+	// sample:
+	// {
+	//        "coinId":"10012",
+	//        "coinName":"usdt",
+	//        "available":"0",
+	//        "frozen":"0",
+	//        "lock":"0",
+	//        "uTime":"1622697148"
+	//    }
+	asset := bitgetapi.AccountAsset{
+		CoinId:    2,
+		CoinName:  "USDT",
+		Available: fixedpoint.NewFromFloat(1.2),
+		Frozen:    fixedpoint.NewFromFloat(0.5),
+		Lock:      fixedpoint.NewFromFloat(0.5),
+		UTime:     types.NewMillisecondTimestampFromInt(1622697148),
+	}
+
+	assert.Equal(t, types.Balance{
+		Currency:          "USDT",
+		Available:         fixedpoint.NewFromFloat(1.2),
+		Locked:            fixedpoint.NewFromFloat(1), // frozen + lock
+		Borrowed:          fixedpoint.Zero,
+		Interest:          fixedpoint.Zero,
+		NetAsset:          fixedpoint.Zero,
+		MaxWithdrawAmount: fixedpoint.Zero,
+	}, toGlobalBalance(asset))
+}
+
+func Test_toGlobalMarket(t *testing.T) {
 	// sample:
 	//{
 	//            "symbol":"BTCUSDT_SPBL",
diff --git a/pkg/exchange/bitget/exchange.go b/pkg/exchange/bitget/exchange.go
index b7af3828b..a3a570016 100644
--- a/pkg/exchange/bitget/exchange.go
+++ b/pkg/exchange/bitget/exchange.go
@@ -23,6 +23,10 @@ var log = logrus.WithFields(logrus.Fields{
 var (
 	// queryMarketRateLimiter has its own rate limit. https://bitgetlimited.github.io/apidoc/en/spot/#get-symbols
 	queryMarketRateLimiter = rate.NewLimiter(rate.Every(time.Second/10), 5)
+	// queryAccountRateLimiter has its own rate limit. https://bitgetlimited.github.io/apidoc/en/spot/#get-account-assets
+	queryAccountRateLimiter = rate.NewLimiter(rate.Every(time.Second/5), 5)
+	// queryTickerRateLimiter has its own rate limit. https://bitgetlimited.github.io/apidoc/en/spot/#get-single-ticker
+	queryTickerRateLimiter = rate.NewLimiter(rate.Every(time.Second/10), 5)
 )
 
 type Exchange struct {
@@ -80,11 +84,15 @@ func (e *Exchange) QueryMarkets(ctx context.Context) (types.MarketMap, error) {
 }
 
 func (e *Exchange) QueryTicker(ctx context.Context, symbol string) (*types.Ticker, error) {
+	if err := queryTickerRateLimiter.Wait(ctx); err != nil {
+		return nil, fmt.Errorf("ticker rate limiter wait error: %w", err)
+	}
+
 	req := e.client.NewGetTickerRequest()
 	req.Symbol(symbol)
 	ticker, err := req.Do(ctx)
 	if err != nil {
-		return nil, err
+		return nil, fmt.Errorf("failed to query ticker: %w", err)
 	}
 
 	return &types.Ticker{
@@ -110,26 +118,34 @@ func (e *Exchange) QueryKLines(ctx context.Context, symbol string, interval type
 }
 
 func (e *Exchange) QueryAccount(ctx context.Context) (*types.Account, error) {
-	req := e.client.NewGetAccountAssetsRequest()
-	resp, err := req.Do(ctx)
+	bals, err := e.QueryAccountBalances(ctx)
 	if err != nil {
 		return nil, err
 	}
 
-	bals := types.BalanceMap{}
-	for _, asset := range resp {
-		b := toGlobalBalance(asset)
-		bals[asset.CoinName] = b
-	}
-
 	account := types.NewAccount()
 	account.UpdateBalances(bals)
 	return account, nil
 }
 
 func (e *Exchange) QueryAccountBalances(ctx context.Context) (types.BalanceMap, error) {
-	// TODO implement me
-	panic("implement me")
+	if err := queryAccountRateLimiter.Wait(ctx); err != nil {
+		return nil, fmt.Errorf("account rate limiter wait error: %w", err)
+	}
+
+	req := e.client.NewGetAccountAssetsRequest()
+	resp, err := req.Do(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("failed to query account assets: %w", err)
+	}
+
+	bals := types.BalanceMap{}
+	for _, asset := range resp {
+		b := toGlobalBalance(asset)
+		bals[asset.CoinName] = b
+	}
+
+	return bals, nil
 }
 
 func (e *Exchange) SubmitOrder(ctx context.Context, order types.SubmitOrder) (createdOrder *types.Order, err error) {