pkg/exchange: emit balance snapshot after authed

pkg/exchange/okex/exchange.go

@@ -375,7 +375,7 @@ func (e *Exchange) CancelOrders(ctx context.Context, orders ...types.Order) erro
 }
 
 func (e *Exchange) NewStream() types.Stream {
-	return NewStream(e.client)
+	return NewStream(e.client, e)
 }
 
 func (e *Exchange) QueryKLines(ctx context.Context, symbol string, interval types.Interval, options types.KLineQueryOptions) ([]types.KLine, error) {

pkg/exchange/okex/stream.go

@@ -8,6 +8,7 @@ import (
 	"time"
 
 	"github.com/c9s/bbgo/pkg/exchange/okex/okexapi"
+	"github.com/c9s/bbgo/pkg/exchange/retry"
 	"github.com/c9s/bbgo/pkg/types"
 )
 
@@ -32,6 +33,7 @@ type Stream struct {
 	types.StandardStream
 
 	client          *okexapi.RestClient
+	balanceProvider types.ExchangeAccountService
 
 	// public callbacks
 	kLineEventCallbacks        []func(candle KLineEvent)
@@ -41,9 +43,10 @@ type Stream struct {
 	marketTradeEventCallbacks  []func(tradeDetail []MarketTradeEvent)
 }
 
-func NewStream(client *okexapi.RestClient) *Stream {
+func NewStream(client *okexapi.RestClient, balanceProvider types.ExchangeAccountService) *Stream {
 	stream := &Stream{
 		client:          client,
+		balanceProvider: balanceProvider,
 		StandardStream:  types.NewStandardStream(),
 	}
 
@@ -57,7 +60,7 @@ func NewStream(client *okexapi.RestClient) *Stream {
 	stream.OnMarketTradeEvent(stream.handleMarketTradeEvent)
 	stream.OnOrderDetailsEvent(stream.handleOrderDetailsEvent)
 	stream.OnConnect(stream.handleConnect)
-	stream.OnAuth(stream.handleAuth)
+	stream.OnAuth(stream.subscribePrivateChannels(stream.emitBalanceSnapshot))
 	return stream
 }
 
@@ -151,7 +154,8 @@ func (s *Stream) handleConnect() {
 	}
 }
 
-func (s *Stream) handleAuth() {
+func (s *Stream) subscribePrivateChannels(next func()) func() {
+	return func() {
 		var subs = []WebsocketSubscription{
 			{Channel: ChannelAccount},
 			{Channel: "orders", InstrumentType: string(okexapi.InstrumentTypeSpot)},
@@ -164,7 +168,28 @@ func (s *Stream) handleAuth() {
 		})
 		if err != nil {
 			log.WithError(err).Error("private channel subscribe error")
+			return
 		}
+		next()
+	}
+}
+
+func (s *Stream) emitBalanceSnapshot() {
+	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
+	defer cancel()
+
+	var balancesMap types.BalanceMap
+	var err error
+	err = retry.GeneralBackoff(ctx, func() error {
+		balancesMap, err = s.balanceProvider.QueryAccountBalances(ctx)
+		return err
+	})
+	if err != nil {
+		log.WithError(err).Error("no more attempts to retrieve balances")
+		return
+	}
+
+	s.EmitBalanceSnapshot(balancesMap)
 }
 
 func (s *Stream) handleOrderDetailsEvent(orderDetails []okexapi.OrderDetails) {

pkg/exchange/okex/stream_test.go

@@ -25,7 +25,7 @@ func getTestClientOrSkip(t *testing.T) *Stream {
 	}
 
 	exchange := New(key, secret, passphrase)
-	return NewStream(exchange.client)
+	return NewStream(exchange.client, exchange)
 }
 
 func TestStream(t *testing.T) {
@@ -39,6 +39,9 @@ func TestStream(t *testing.T) {
 		s.OnBalanceUpdate(func(balances types.BalanceMap) {
 			t.Log("got snapshot", balances)
 		})
+		s.OnBalanceSnapshot(func(balances types.BalanceMap) {
+			t.Log("got snapshot", balances)
+		})
 		s.OnBookUpdate(func(book types.SliceOrderBook) {
 			t.Log("got update", book)
 		})