engine/websocket: subscribe to default channels only when actually needed (#610)

* if this is required by ws routines or sync manager
* restore previous subscriptions on reconnect

engine/engine.go

@@ -474,7 +474,7 @@ func (bot *Engine) Start() error {
 	}
 
 	if bot.Settings.EnableWebsocketRoutine {
-		go WebsocketRoutine()
+		go bot.WebsocketRoutine()
 	}
 
 	if bot.Settings.EnableGCTScriptManager {

engine/orders_test.go

@@ -11,18 +11,19 @@ import (
 
 var ordersSetupRan bool
 
-func OrdersSetup(t *testing.T) {
-	SetupTestHelpers(t)
+func OrdersSetup(t *testing.T) *Engine {
+	bot := SetupTestHelpers(t)
 	if !ordersSetupRan {
-		err := Bot.OrderManager.Start()
+		err := bot.OrderManager.Start()
 		if err != nil {
 			t.Fatal(err)
 		}
-		if !Bot.OrderManager.Started() {
+		if !bot.OrderManager.Started() {
 			t.Fatal("Order manager not started")
 		}
 		ordersSetupRan = true
 	}
+	return bot
 }
 
 func TestOrdersGet(t *testing.T) {

engine/routines.go

@@ -200,16 +200,16 @@ func relayWebsocketEvent(result interface{}, event, assetType, exchangeName stri
 }
 
 // WebsocketRoutine Initial routine management system for websocket
-func WebsocketRoutine() {
-	if Bot.Settings.Verbose {
+func (bot *Engine) WebsocketRoutine() {
+	if bot.Settings.Verbose {
 		log.Debugln(log.WebsocketMgr, "Connecting exchange websocket services...")
 	}
 
-	exchanges := Bot.GetExchanges()
+	exchanges := bot.GetExchanges()
 	for i := range exchanges {
 		go func(i int) {
 			if exchanges[i].SupportsWebsocket() {
-				if Bot.Settings.Verbose {
+				if bot.Settings.Verbose {
 					log.Debugf(log.WebsocketMgr,
 						"Exchange %s websocket support: Yes Enabled: %v\n",
 						exchanges[i].GetName(),
@@ -235,15 +235,19 @@ func WebsocketRoutine() {
 				}
 
 				// Data handler routine
-				go WebsocketDataReceiver(ws)
+				go bot.WebsocketDataReceiver(ws)
 
 				if ws.IsEnabled() {
 					err = ws.Connect()
 					if err != nil {
 						log.Errorf(log.WebsocketMgr, "%v\n", err)
 					}
+					err = ws.FlushChannels()
+					if err != nil {
+						log.Errorf(log.WebsocketMgr, "Failed to subscribe: %v\n", err)
+					}
 				}
-			} else if Bot.Settings.Verbose {
+			} else if bot.Settings.Verbose {
 				log.Debugf(log.WebsocketMgr,
 					"Exchange %s websocket support: No\n",
 					exchanges[i].GetName(),
@@ -258,7 +262,7 @@ var wg sync.WaitGroup
 
 // WebsocketDataReceiver handles websocket data coming from a websocket feed
 // associated with an exchange
-func WebsocketDataReceiver(ws *stream.Websocket) {
+func (bot *Engine) WebsocketDataReceiver(ws *stream.Websocket) {
 	wg.Add(1)
 	defer wg.Done()
 
@@ -267,7 +271,7 @@ func WebsocketDataReceiver(ws *stream.Websocket) {
 		case <-shutdowner:
 			return
 		case data := <-ws.ToRoutine:
-			err := WebsocketDataHandler(ws.GetName(), data)
+			err := bot.WebsocketDataHandler(ws.GetName(), data)
 			if err != nil {
 				log.Error(log.WebsocketMgr, err)
 			}
@@ -277,7 +281,7 @@ func WebsocketDataReceiver(ws *stream.Websocket) {
 
 // WebsocketDataHandler is a central point for exchange websocket implementations to send
 // processed data. WebsocketDataHandler will then pass that to an appropriate handler
-func WebsocketDataHandler(exchName string, data interface{}) error {
+func (bot *Engine) WebsocketDataHandler(exchName string, data interface{}) error {
 	if data == nil {
 		return fmt.Errorf("routines.go - exchange %s nil data sent to websocket",
 			exchName)
@@ -289,7 +293,7 @@ func WebsocketDataHandler(exchName string, data interface{}) error {
 	case error:
 		return fmt.Errorf("routines.go exchange %s websocket error - %s", exchName, data)
 	case stream.FundingData:
-		if Bot.Settings.Verbose {
+		if bot.Settings.Verbose {
 			log.Infof(log.WebsocketMgr, "%s websocket %s %s funding updated %+v",
 				exchName,
 				FormatCurrency(d.CurrencyPair),
@@ -297,8 +301,8 @@ func WebsocketDataHandler(exchName string, data interface{}) error {
 				d)
 		}
 	case *ticker.Price:
-		if Bot.Settings.EnableExchangeSyncManager && Bot.ExchangeCurrencyPairManager != nil {
-			Bot.ExchangeCurrencyPairManager.update(exchName,
+		if bot.Settings.EnableExchangeSyncManager && bot.ExchangeCurrencyPairManager != nil {
+			bot.ExchangeCurrencyPairManager.update(exchName,
 				d.Pair,
 				d.AssetType,
 				SyncItemTicker,
@@ -307,7 +311,7 @@ func WebsocketDataHandler(exchName string, data interface{}) error {
 		err := ticker.ProcessTicker(d)
 		printTickerSummary(d, "websocket", err)
 	case stream.KlineData:
-		if Bot.Settings.Verbose {
+		if bot.Settings.Verbose {
 			log.Infof(log.WebsocketMgr, "%s websocket %s %s kline updated %+v",
 				exchName,
 				FormatCurrency(d.Pair),
@@ -315,8 +319,8 @@ func WebsocketDataHandler(exchName string, data interface{}) error {
 				d)
 		}
 	case *orderbook.Base:
-		if Bot.Settings.EnableExchangeSyncManager && Bot.ExchangeCurrencyPairManager != nil {
-			Bot.ExchangeCurrencyPairManager.update(exchName,
+		if bot.Settings.EnableExchangeSyncManager && bot.ExchangeCurrencyPairManager != nil {
+			bot.ExchangeCurrencyPairManager.update(exchName,
 				d.Pair,
 				d.AssetType,
 				SyncItemOrderbook,
@@ -324,22 +328,22 @@ func WebsocketDataHandler(exchName string, data interface{}) error {
 		}
 		printOrderbookSummary(d, "websocket", nil)
 	case *order.Detail:
-		if !Bot.OrderManager.orderStore.exists(d) {
-			err := Bot.OrderManager.orderStore.Add(d)
+		if !bot.OrderManager.orderStore.exists(d) {
+			err := bot.OrderManager.orderStore.Add(d)
 			if err != nil {
 				return err
 			}
 		} else {
-			od, err := Bot.OrderManager.orderStore.GetByExchangeAndID(d.Exchange, d.ID)
+			od, err := bot.OrderManager.orderStore.GetByExchangeAndID(d.Exchange, d.ID)
 			if err != nil {
 				return err
 			}
 			od.UpdateOrderFromDetail(d)
 		}
 	case *order.Cancel:
-		return Bot.OrderManager.Cancel(d)
+		return bot.OrderManager.Cancel(d)
 	case *order.Modify:
-		od, err := Bot.OrderManager.orderStore.GetByExchangeAndID(d.Exchange, d.ID)
+		od, err := bot.OrderManager.orderStore.GetByExchangeAndID(d.Exchange, d.ID)
 		if err != nil {
 			return err
 		}
@@ -349,7 +353,7 @@ func WebsocketDataHandler(exchName string, data interface{}) error {
 	case stream.UnhandledMessageWarning:
 		log.Warn(log.WebsocketMgr, d.Message)
 	default:
-		if Bot.Settings.Verbose {
+		if bot.Settings.Verbose {
 			log.Warnf(log.WebsocketMgr,
 				"%s websocket Unknown type: %+v",
 				exchName,

engine/routines_test.go

@@ -15,33 +15,33 @@ import (
 
 func TestWebsocketDataHandlerProcess(t *testing.T) {
 	ws := sharedtestvalues.NewTestWebsocket()
-	go WebsocketDataReceiver(ws)
+	go Bot.WebsocketDataReceiver(ws)
 	ws.DataHandler <- "string"
 	time.Sleep(time.Second)
 	close(shutdowner)
 }
 
 func TestHandleData(t *testing.T) {
-	OrdersSetup(t)
+	b := OrdersSetup(t)
 	var exchName = "exch"
 	var orderID = "testOrder.Detail"
-	err := WebsocketDataHandler(exchName, errors.New("error"))
+	err := b.WebsocketDataHandler(exchName, errors.New("error"))
 	if err == nil {
 		t.Error("Error not handled correctly")
 	}
-	err = WebsocketDataHandler(exchName, nil)
+	err = b.WebsocketDataHandler(exchName, nil)
 	if err == nil {
 		t.Error("Expected nil data error")
 	}
-	err = WebsocketDataHandler(exchName, stream.FundingData{})
+	err = b.WebsocketDataHandler(exchName, stream.FundingData{})
 	if err != nil {
 		t.Error(err)
 	}
-	err = WebsocketDataHandler(exchName, &ticker.Price{})
+	err = b.WebsocketDataHandler(exchName, &ticker.Price{})
 	if err != nil {
 		t.Error(err)
 	}
-	err = WebsocketDataHandler(exchName, stream.KlineData{})
+	err = b.WebsocketDataHandler(exchName, stream.KlineData{})
 	if err != nil {
 		t.Error(err)
 	}
@@ -51,12 +51,12 @@ func TestHandleData(t *testing.T) {
 		Amount:   1337,
 		Price:    1337,
 	}
-	err = WebsocketDataHandler(exchName, origOrder)
+	err = b.WebsocketDataHandler(exchName, origOrder)
 	if err != nil {
 		t.Error(err)
 	}
 	// Send it again since it exists now
-	err = WebsocketDataHandler(exchName, &order.Detail{
+	err = b.WebsocketDataHandler(exchName, &order.Detail{
 		Exchange: fakePassExchange,
 		ID:       orderID,
 		Amount:   1338,
@@ -68,7 +68,7 @@ func TestHandleData(t *testing.T) {
 		t.Error("Bad pipeline")
 	}
 
-	err = WebsocketDataHandler(exchName, &order.Modify{
+	err = b.WebsocketDataHandler(exchName, &order.Modify{
 		Exchange: fakePassExchange,
 		ID:       orderID,
 		Status:   order.Active,
@@ -80,7 +80,7 @@ func TestHandleData(t *testing.T) {
 		t.Error("Expected order to be modified to Active")
 	}
 
-	err = WebsocketDataHandler(exchName, &order.Cancel{
+	err = b.WebsocketDataHandler(exchName, &order.Cancel{
 		Exchange: fakePassExchange,
 		ID:       orderID,
 	})
@@ -91,12 +91,12 @@ func TestHandleData(t *testing.T) {
 		t.Error("Expected order status to be cancelled")
 	}
 	// Send some gibberish
-	err = WebsocketDataHandler(exchName, order.Stop)
+	err = b.WebsocketDataHandler(exchName, order.Stop)
 	if err != nil {
 		t.Error(err)
 	}
 
-	err = WebsocketDataHandler(exchName, stream.UnhandledMessageWarning{
+	err = b.WebsocketDataHandler(exchName, stream.UnhandledMessageWarning{
 		Message: "there's an issue here's a tissue"},
 	)
 	if err != nil {
@@ -108,7 +108,7 @@ func TestHandleData(t *testing.T) {
 		OrderID:  "one",
 		Err:      errors.New("lol"),
 	}
-	err = WebsocketDataHandler(exchName, classificationError)
+	err = b.WebsocketDataHandler(exchName, classificationError)
 	if err == nil {
 		t.Error("Expected error")
 	}
@@ -116,14 +116,14 @@ func TestHandleData(t *testing.T) {
 		t.Errorf("Problem formatting error. Expected %v Received %v", classificationError.Error(), err.Error())
 	}
 
-	err = WebsocketDataHandler(exchName, &orderbook.Base{
+	err = b.WebsocketDataHandler(exchName, &orderbook.Base{
 		ExchangeName: fakePassExchange,
 		Pair:         currency.NewPair(currency.BTC, currency.USD),
 	})
 	if err != nil {
 		t.Error(err)
 	}
-	err = WebsocketDataHandler(exchName, "this is a test string")
+	err = b.WebsocketDataHandler(exchName, "this is a test string")
 	if err != nil {
 		t.Error(err)
 	}

engine/syncer.go

@@ -529,9 +529,12 @@ func (e *ExchangeCurrencyPairSyncer) Start() {
 			}
 
 			if !ws.IsConnected() && !ws.IsConnecting() {
-				go WebsocketDataReceiver(ws)
+				go Bot.WebsocketDataReceiver(ws)
 
 				err = ws.Connect()
+				if err == nil {
+					err = ws.FlushChannels()
+				}
 				if err != nil {
 					log.Errorf(log.SyncMgr,
 						"%s websocket failed to connect. Err: %s\n",

exchanges/binance/binance_websocket.go

@@ -86,11 +86,7 @@ func (b *Binance) WsConnect() error {
 
 	go b.wsReadData()
 
-	subs, err := b.GenerateSubscriptions()
-	if err != nil {
-		return err
-	}
-	return b.Websocket.SubscribeToChannels(subs)
+	return nil
 }
 
 // KeepAuthKeyAlive will continuously send messages to

exchanges/bitfinex/bitfinex_websocket.go

@@ -61,12 +61,8 @@ func (b *Bitfinex) WsConnect() error {
 		}
 	}
 
-	subs, err := b.GenerateDefaultSubscriptions()
-	if err != nil {
-		return err
-	}
 	go b.WsDataHandler()
-	return b.Websocket.SubscribeToChannels(subs)
+	return nil
 }
 
 // wsReadData receives and passes on websocket messages for processing

exchanges/bitmex/bitmex_websocket.go

@@ -95,15 +95,6 @@ func (b *Bitmex) WsConnect() error {
 	}
 
 	go b.wsReadData()
-	subs, err := b.GenerateDefaultSubscriptions()
-	if err != nil {
-		return err
-	}
-
-	err = b.Websocket.SubscribeToChannels(subs)
-	if err != nil {
-		return err
-	}
 
 	err = b.websocketSendAuth()
 	if err != nil {

exchanges/bitstamp/bitstamp_websocket.go

@@ -40,12 +40,8 @@ func (b *Bitstamp) WsConnect() error {
 	if err != nil {
 		b.Websocket.DataHandler <- err
 	}
-	subs, err := b.generateDefaultSubscriptions()
-	if err != nil {
-		return err
-	}
 	go b.wsReadData()
-	return b.Websocket.SubscribeToChannels(subs)
+	return nil
 }
 
 // wsReadData receives and passes on websocket messages for processing

exchanges/btcmarkets/btcmarkets_websocket.go

@@ -40,11 +40,7 @@ func (b *BTCMarkets) WsConnect() error {
 		log.Debugf(log.ExchangeSys, "%s Connected to Websocket.\n", b.Name)
 	}
 	go b.wsReadData()
-	subs, err := b.generateDefaultSubscriptions()
-	if err != nil {
-		return err
-	}
-	return b.Websocket.SubscribeToChannels(subs)
+	return nil
 }
 
 // wsReadData receives and passes on websocket messages for processing

exchanges/btse/btse_websocket.go

@@ -49,11 +49,7 @@ func (b *BTSE) WsConnect() error {
 		}
 	}
 
-	subs, err := b.GenerateDefaultSubscriptions()
-	if err != nil {
-		return err
-	}
-	return b.Websocket.SubscribeToChannels(subs)
+	return nil
 }
 
 // WsAuthenticate Send an authentication message to receive auth data

exchanges/coinbasepro/coinbasepro_websocket.go

@@ -38,12 +38,8 @@ func (c *CoinbasePro) WsConnect() error {
 		return err
 	}
 
-	subs, err := c.GenerateDefaultSubscriptions()
-	if err != nil {
-		return err
-	}
 	go c.wsReadData()
-	return c.Websocket.SubscribeToChannels(subs)
+	return nil
 }
 
 // wsReadData receives and passes on websocket messages for processing

exchanges/coinbene/coinbene_websocket.go

@@ -49,11 +49,7 @@ func (c *Coinbene) WsConnect() error {
 			c.Websocket.SetCanUseAuthenticatedEndpoints(false)
 		}
 	}
-	subs, err := c.GenerateDefaultSubscriptions()
-	if err != nil {
-		return err
-	}
-	return c.Websocket.SubscribeToChannels(subs)
+	return nil
 }
 
 // GenerateDefaultSubscriptions generates stuff

exchanges/coinut/coinut_websocket.go

@@ -61,15 +61,6 @@ func (c *COINUT) WsConnect() error {
 		c.Websocket.SetCanUseAuthenticatedEndpoints(false)
 		log.Error(log.WebsocketMgr, err)
 	}
-	subs, err := c.GenerateDefaultSubscriptions()
-	if err != nil {
-		return err
-	}
-
-	err = c.Websocket.SubscribeToChannels(subs)
-	if err != nil {
-		return err
-	}
 
 	// define bi-directional communication
 	channels = make(map[string]chan []byte)

exchanges/ftx/ftx_websocket.go

@@ -69,11 +69,7 @@ func (f *FTX) WsConnect() error {
 		}
 	}
 
-	subs, err := f.GenerateDefaultSubscriptions()
-	if err != nil {
-		return err
-	}
-	return f.Websocket.SubscribeToChannels(subs)
+	return nil
 }
 
 // WsAuth sends an authentication message to receive auth data

exchanges/gateio/gateio_websocket.go

@@ -62,11 +62,7 @@ func (g *Gateio) WsConnect() error {
 		}
 	}
 
-	subs, err := g.GenerateDefaultSubscriptions()
-	if err != nil {
-		return err
-	}
-	return g.Websocket.SubscribeToChannels(subs)
+	return nil
 }
 
 func (g *Gateio) wsServerSignIn() error {

exchanges/hitbtc/hitbtc_websocket.go

@@ -50,11 +50,7 @@ func (h *HitBTC) WsConnect() error {
 		log.Errorf(log.ExchangeSys, "%v - authentication failed: %v\n", h.Name, err)
 	}
 
-	subs, err := h.GenerateDefaultSubscriptions()
-	if err != nil {
-		return err
-	}
-	return h.Websocket.SubscribeToChannels(subs)
+	return nil
 }
 
 // wsReadData receives and passes on websocket messages for processing

exchanges/huobi/huobi_websocket.go

@@ -83,11 +83,7 @@ func (h *HUOBI) WsConnect() error {
 	}
 
 	go h.wsReadData()
-	subs, err := h.GenerateDefaultSubscriptions()
-	if err != nil {
-		return err
-	}
-	return h.Websocket.SubscribeToChannels(subs)
+	return nil
 }
 
 func (h *HUOBI) wsDial(dialer *websocket.Dialer) error {

exchanges/kraken/kraken_websocket.go

@@ -131,11 +131,7 @@ func (k *Kraken) WsConnect() error {
 			k.Name,
 			err)
 	}
-	gensubs, err := k.GenerateDefaultSubscriptions()
-	if err != nil {
-		return err
-	}
-	return k.Websocket.SubscribeToChannels(gensubs)
+	return nil
 }
 
 // wsFunnelConnectionData funnels both auth and public ws data into one manageable place

exchanges/lakebtc/lakebtc_websocket.go

@@ -50,11 +50,7 @@ func (l *LakeBTC) WsConnect() error {
 		return err
 	}
 	go l.wsHandleIncomingData()
-	subs, err := l.GenerateDefaultSubscriptions()
-	if err != nil {
-		return err
-	}
-	return l.Websocket.SubscribeToChannels(subs)
+	return nil
 }
 
 func (l *LakeBTC) listenToEndpoints() error {

exchanges/okgroup/okgroup_websocket.go

@@ -203,11 +203,7 @@ func (o *OKGroup) WsConnect() error {
 		}
 	}
 
-	subs, err := o.GenerateDefaultSubscriptions()
-	if err != nil {
-		return err
-	}
-	return o.Websocket.SubscribeToChannels(subs)
+	return nil
 }
 
 // WsLogin sends a login request to websocket to enable access to authenticated endpoints

exchanges/poloniex/poloniex_websocket.go

@@ -53,12 +53,8 @@ func (p *Poloniex) WsConnect() error {
 	}
 
 	go p.wsReadData()
-	subs, err := p.GenerateDefaultSubscriptions()
-	if err != nil {
-		return err
-	}
 
-	return p.Websocket.SubscribeToChannels(subs)
+	return nil
 }
 
 func (p *Poloniex) getCurrencyIDMap() error {

exchanges/stream/websocket.go

@@ -192,11 +192,6 @@ func (w *Websocket) Connect() error {
 	w.trafficMonitor()
 	w.setConnectingStatus(true)
 
-	// flush any subscriptions from last connection if needed
-	w.subscriptionMutex.Lock()
-	w.subscriptions = nil
-	w.subscriptionMutex.Unlock()
-
 	err := w.connector()
 	if err != nil {
 		w.setConnectingStatus(false)
@@ -211,6 +206,14 @@ func (w *Websocket) Connect() error {
 		w.connectionMonitor()
 	}
 
+	// Resubscribe after re-connection
+	if len(w.subscriptions) != 0 {
+		err = w.Subscriber(w.subscriptions)
+		if err != nil {
+			return fmt.Errorf("%v Error subscribing %s", w.exchangeName, err)
+		}
+	}
+
 	return nil
 }
 

exchanges/zb/zb_websocket.go

@@ -41,12 +41,8 @@ func (z *ZB) WsConnect() error {
 		return err
 	}
 
-	subs, err := z.GenerateDefaultSubscriptions()
-	if err != nil {
-		return err
-	}
 	go z.wsReadData()
-	return z.Websocket.SubscribeToChannels(subs)
+	return nil
 }
 
 // wsReadData handles all the websocket data coming from the websocket