Merge pull request #1424 from c9s/c9s/start-up-balance-query

IMPROVE: improve startup balance query process

pkg/bbgo/config.go

@@ -329,6 +329,11 @@ type ServiceConfig struct {
 type EnvironmentConfig struct {
 	DisableDefaultKLineSubscription bool `json:"disableDefaultKLineSubscription"`
 	DisableHistoryKLinePreload      bool `json:"disableHistoryKLinePreload"`
+
+	// DisableStartUpBalanceQuery disables the balance query in the startup process
+	// which initializes the session.Account with the QueryAccount method.
+	DisableStartupBalanceQuery bool `json:"disableStartupBalanceQuery"`
+
 	DisableSessionTradeBuffer bool `json:"disableSessionTradeBuffer"`
 	MaxSessionTradeBufferSize int  `json:"maxSessionTradeBufferSize"`
 }

pkg/bbgo/session.go

@@ -15,6 +15,7 @@ import (
 
 	"github.com/c9s/bbgo/pkg/cache"
 	"github.com/c9s/bbgo/pkg/core"
+	"github.com/c9s/bbgo/pkg/exchange/retry"
 	"github.com/c9s/bbgo/pkg/util/templateutil"
 
 	exchange2 "github.com/c9s/bbgo/pkg/exchange"
@@ -177,12 +178,16 @@ func (session *ExchangeSession) UpdateAccount(ctx context.Context) (*types.Accou
 		return nil, err
 	}
 
-	session.accountMutex.Lock()
-	session.Account = account
-	session.accountMutex.Unlock()
+	session.setAccount(account)
 	return account, nil
 }
 
+func (session *ExchangeSession) setAccount(a *types.Account) {
+	session.accountMutex.Lock()
+	session.Account = a
+	session.accountMutex.Unlock()
+}
+
 // Init initializes the basic data structure and market information by its exchange.
 // Note that the subscribed symbols are not loaded in this stage.
 func (session *ExchangeSession) Init(ctx context.Context, environ *Environment) error {
@@ -256,17 +261,22 @@ func (session *ExchangeSession) Init(ctx context.Context, environ *Environment)
 			}
 		}
 
+		disableStartupBalanceQuery := environ.environmentConfig != nil && environ.environmentConfig.DisableStartupBalanceQuery
+		if disableStartupBalanceQuery {
+			session.accountMutex.Lock()
+			session.Account = types.NewAccount()
+			session.accountMutex.Unlock()
+		} else {
 			logger.Infof("querying account balances...")
-		account, err := session.Exchange.QueryAccount(ctx)
+			account, err := retry.QueryAccountUntilSuccessful(ctx, session.Exchange)
 			if err != nil {
 				return err
 			}
 
-		session.accountMutex.Lock()
-		session.Account = account
-		session.accountMutex.Unlock()
-
+			session.setAccount(account)
+			session.metricsBalancesUpdater(account.Balances())
 			logger.Infof("account %s balances:\n%s", session.Name, account.Balances().String())
+		}
 
 		// forward trade updates and order updates to the order executor
 		session.UserDataStream.OnTradeUpdate(session.OrderExecutor.EmitTradeUpdate)
@@ -288,7 +298,6 @@ func (session *ExchangeSession) Init(ctx context.Context, environ *Environment)
 
 		// if metrics mode is enabled, we bind the callbacks to update metrics
 		if viper.GetBool("metrics") {
-			session.metricsBalancesUpdater(account.Balances())
 			session.bindUserDataStreamMetrics(session.UserDataStream)
 		}
 	}

pkg/exchange/retry/order.go

@@ -16,7 +16,9 @@ type advancedOrderCancelService interface {
 	CancelOrdersByGroupID(ctx context.Context, groupID uint32) ([]types.Order, error)
 }
 
-func QueryOrderUntilFilled(ctx context.Context, queryOrderService types.ExchangeOrderQueryService, symbol string, orderId uint64) (o *types.Order, err error) {
+func QueryOrderUntilFilled(
+	ctx context.Context, queryOrderService types.ExchangeOrderQueryService, symbol string, orderId uint64,
+) (o *types.Order, err error) {
 	var op = func() (err2 error) {
 		o, err2 = queryOrderService.QueryOrder(ctx, types.OrderQuery{
 			Symbol:  symbol,
@@ -56,7 +58,9 @@ func GeneralLiteBackoff(ctx context.Context, op backoff.Operation) (err error) {
 	return err
 }
 
-func QueryOpenOrdersUntilSuccessful(ctx context.Context, ex types.Exchange, symbol string) (openOrders []types.Order, err error) {
+func QueryOpenOrdersUntilSuccessful(
+	ctx context.Context, ex types.Exchange, symbol string,
+) (openOrders []types.Order, err error) {
 	var op = func() (err2 error) {
 		openOrders, err2 = ex.QueryOpenOrders(ctx, symbol)
 		return err2
@@ -66,7 +70,9 @@ func QueryOpenOrdersUntilSuccessful(ctx context.Context, ex types.Exchange, symb
 	return openOrders, err
 }
 
-func QueryOpenOrdersUntilSuccessfulLite(ctx context.Context, ex types.Exchange, symbol string) (openOrders []types.Order, err error) {
+func QueryOpenOrdersUntilSuccessfulLite(
+	ctx context.Context, ex types.Exchange, symbol string,
+) (openOrders []types.Order, err error) {
 	var op = func() (err2 error) {
 		openOrders, err2 = ex.QueryOpenOrders(ctx, symbol)
 		return err2
@@ -76,7 +82,21 @@ func QueryOpenOrdersUntilSuccessfulLite(ctx context.Context, ex types.Exchange,
 	return openOrders, err
 }
 
-func QueryOrderUntilSuccessful(ctx context.Context, query types.ExchangeOrderQueryService, opts types.OrderQuery) (order *types.Order, err error) {
+func QueryAccountUntilSuccessful(
+	ctx context.Context, ex types.ExchangeAccountService,
+) (account *types.Account, err error) {
+	var op = func() (err2 error) {
+		account, err2 = ex.QueryAccount(ctx)
+		return err2
+	}
+
+	err = GeneralBackoff(ctx, op)
+	return account, err
+}
+
+func QueryOrderUntilSuccessful(
+	ctx context.Context, query types.ExchangeOrderQueryService, opts types.OrderQuery,
+) (order *types.Order, err error) {
 	var op = func() (err2 error) {
 		order, err2 = query.QueryOrder(ctx, opts)
 		return err2

pkg/types/account.go

@@ -103,8 +103,26 @@ type IsolatedMarginAccountInfo struct {
 
 func NewAccount() *Account {
 	return &Account{
+		AccountType:        "spot",
+		FuturesInfo:        nil,
+		MarginInfo:         nil,
+		IsolatedMarginInfo: nil,
+		MarginLevel:        fixedpoint.Zero,
+		MarginTolerance:    fixedpoint.Zero,
+		BorrowEnabled:      false,
+		TransferEnabled:    false,
+		MarginRatio:        fixedpoint.Zero,
+		LiquidationPrice:   fixedpoint.Zero,
+		LiquidationRate:    fixedpoint.Zero,
+		MakerFeeRate:       fixedpoint.Zero,
+		TakerFeeRate:       fixedpoint.Zero,
+		TotalAccountValue:  fixedpoint.Zero,
+		CanDeposit:         false,
+		CanTrade:           false,
+		CanWithdraw:        false,
 		balances:           make(BalanceMap),
 	}
+
 }
 
 // Balances lock the balances and returned the copied balances