Websocket connection handling and subscription management (#297)

* Step one: Sets up  connection handler for websockets to always be connected until a shutdown event is received.
Sets up a vague subscription handler to ensure subscriptions are subscribed

* Adds support for resubscriptions for bitfinex, bitstamp, bitmex and btcc. Adds subscription params for special websocket subscription requirements. Removes subscription monitor from wait group so that it can exist despite a shutdown and continuously check

* Adds channel subscription support to bitmex, btse, coibasepro, coinut, gateio, gemini, hitbtc, huobi, hadax, kraken, okgroup, poloniex and zb

* Implements unsubscribe for bitfinex, btcc, btse, coinbasepro, gateio, gitbtc, huobi, hadax

* ManageSubscriptions now called from WSConnect and made private instead of inside individual exchanges. ManageSubscriptions can now unsubscribe. exchange_websocket_types.go now contains all exchange_websocket.go types to avoid clutter

* Adds it to websocket functionality so managesubscriptions will close when not supported

* Separates functions into testable functions to ensure logic works. Adds tests. Updates websocket setup to include verbosity (inherited from exchange). Adds no connection tolerance to fatal on failed reconnects

* More exchange_websocket tests. Updating to use pointers. Creation of equals func to make comparison easier

* Fixes okex, okcoin tests. Fixes race conditions. Removes pointer usage again.

* Adds subscribe and unsubscribe to wrappers

* Fixes deadlock. Fixes ws verbosity.

* Updates all exchanges to properly support subscription/connection feature. Also reintroduces race conditions....

* Moves connection varialbes to struct from package to allow each websocket to have their own reconnection checks. Neatens up logs

* Fixes lint/critic issues. Fixes tests. Removes unused function.

* Moves websocket ratelimiter to their own const variables. Fixes more race conditions with connecting variable

* Removes redundant subscribe functions. Ensuring only the exchange_websocket.go can manage subscriptions. Fixes debug logs to be verbose wrapped

* Fixes issue with slice copying. Re-adds okgroup default channels

* Adds nolint to append

* Adds comments and adds support for gateio auth request subscriptions

* Adds new test to ensure slices dont point to the same vars

* removes fatals. gofmt goimports

* more gofmts

* Addresses PR comments, removing empty and redundant lines

* Addresses PR comments. Ensures that writing to the websocket is single-threaded by adding a mutex to exchanges. Minimises wrapper code and moves subscription loops to exchange_websocket. Privatises ChannelsToSubscribe, Connecting properties and removeChannelToSubscribe func to prevent unnecessary tampering.

* Removes unused mutex. FMTS and IMPORTS

* Fixes request lock time change

* More specific logs

* Renames ws mutex. Fixes bitmex subscriptions. Increased gateio ratelimiter to 120ms. Removes ratelimiter from bitfinex, bitmex, bitstamp, btcc, btse, coibasepro, hitbtc, huobi, hadax, poloniex and zb

* changes recieved typo due to not being well received

* Fixes parsing issue with Huobi and hadax

* Fixes data race with more locks

* removes defer locks. fixes huobi/hadax verbose output

* Fixes double JSONEncode for coinut. Fixes verbose output for coinut

* gofmt,goimport for coinut

* Fixes issue where multiple connection monitors can spawn

* Removes defer exchange.WebsocketConn.Close() in defer handledata exit as connectionmonitor handles connections instead

* gofmt and go import

* More fmts

exchanges/alphapoint/alphapoint_wrapper.go

@@ -316,3 +316,15 @@ func (a *Alphapoint) GetOrderHistory(getOrdersRequest *exchange.GetOrdersRequest
 
 	return orders, nil
 }
+
+// SubscribeToWebsocketChannels appends to ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle subscribing
+func (a *Alphapoint) SubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	return common.ErrFunctionNotSupported
+}
+
+// UnsubscribeToWebsocketChannels removes from ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle unsubscribing
+func (a *Alphapoint) UnsubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	return common.ErrFunctionNotSupported
+}

exchanges/anx/anx_wrapper.go

@@ -445,3 +445,15 @@ func (a *ANX) GetOrderHistory(getOrdersRequest *exchange.GetOrdersRequest) ([]ex
 
 	return orders, nil
 }
+
+// SubscribeToWebsocketChannels appends to ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle subscribing
+func (a *ANX) SubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	return common.ErrFunctionNotSupported
+}
+
+// UnsubscribeToWebsocketChannels removes from ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle unsubscribing
+func (a *ANX) UnsubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	return common.ErrFunctionNotSupported
+}

exchanges/binance/binance.go

@@ -139,8 +139,11 @@ func (b *Binance) Setup(exch *config.ExchangeConfig) {
 			log.Fatal(err)
 		}
 		err = b.WebsocketSetup(b.WSConnect,
+			nil,
+			nil,
 			exch.Name,
 			exch.Websocket,
+			exch.Verbose,
 			binanceDefaultWebsocketURL,
 			exch.WebsocketURL)
 		if err != nil {

exchanges/binance/binance_websocket.go

@@ -199,11 +199,6 @@ func (b *Binance) WsHandleData() {
 	b.Websocket.Wg.Add(1)
 
 	defer func() {
-		err := b.WebsocketConn.Close()
-		if err != nil {
-			b.Websocket.DataHandler <- fmt.Errorf("binance_websocket.go - Unable to to close Websocket connection. Error: %s",
-				err)
-		}
 		b.Websocket.Wg.Done()
 	}()
 

exchanges/binance/binance_wrapper.go

@@ -413,3 +413,15 @@ func (b *Binance) GetOrderHistory(getOrdersRequest *exchange.GetOrdersRequest) (
 
 	return orders, nil
 }
+
+// SubscribeToWebsocketChannels appends to ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle subscribing
+func (b *Binance) SubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	return common.ErrFunctionNotSupported
+}
+
+// UnsubscribeToWebsocketChannels removes from ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle unsubscribing
+func (b *Binance) UnsubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	return common.ErrFunctionNotSupported
+}

exchanges/bitfinex/bitfinex.go

@@ -6,6 +6,7 @@ import (
 	"net/http"
 	"net/url"
 	"strconv"
+	"sync"
 	"time"
 
 	"github.com/gorilla/websocket"
@@ -87,6 +88,7 @@ type Bitfinex struct {
 	exchange.Base
 	WebsocketConn         *websocket.Conn
 	WebsocketSubdChannels map[int]WebsocketChanInfo
+	wsRequestMtx          sync.Mutex
 }
 
 // SetDefaults sets the basic defaults for bitfinex
@@ -114,7 +116,9 @@ func (b *Bitfinex) SetDefaults() {
 	b.WebsocketInit()
 	b.Websocket.Functionality = exchange.WebsocketTickerSupported |
 		exchange.WebsocketTradeDataSupported |
-		exchange.WebsocketOrderbookSupported
+		exchange.WebsocketOrderbookSupported |
+		exchange.WebsocketSubscribeSupported |
+		exchange.WebsocketUnsubscribeSupported
 }
 
 // Setup takes in the supplied exchange configuration details and sets params
@@ -155,8 +159,11 @@ func (b *Bitfinex) Setup(exch *config.ExchangeConfig) {
 			log.Fatal(err)
 		}
 		err = b.WebsocketSetup(b.WsConnect,
+			b.Subscribe,
+			b.Unsubscribe,
 			exch.Name,
 			exch.Websocket,
+			exch.Verbose,
 			bitfinexWebsocket,
 			exch.WebsocketURL)
 		if err != nil {
@@ -617,7 +624,7 @@ func (b *Bitfinex) WithdrawCryptocurrency(withdrawType, wallet, address, payment
 			&response)
 }
 
-// WithdrawFIAT requests a withdrawal from a designated fiat wallet
+// WithdrawFIAT Sends an authenticated request to withdraw FIAT currency
 func (b *Bitfinex) WithdrawFIAT(withdrawalType, walletType string, withdrawRequest *exchange.WithdrawRequest) ([]Withdrawal, error) {
 	var response []Withdrawal
 	req := make(map[string]interface{})

exchanges/bitfinex/bitfinex_websocket.go

@@ -57,30 +57,21 @@ func (b *Bitfinex) WsPingHandler() error {
 	req := make(map[string]string)
 	req["event"] = "ping"
 
-	return b.WsSend(req)
+	return b.wsSend(req)
 }
 
 // WsSend sends data to the websocket server
-func (b *Bitfinex) WsSend(data interface{}) error {
+func (b *Bitfinex) wsSend(data interface{}) error {
+	b.wsRequestMtx.Lock()
+	defer b.wsRequestMtx.Unlock()
 	json, err := common.JSONEncode(data)
 	if err != nil {
 		return err
 	}
-	return b.WebsocketConn.WriteMessage(websocket.TextMessage, json)
-}
-
-// WsSubscribe subscribes to the websocket channel
-func (b *Bitfinex) WsSubscribe(channel string, params map[string]string) error {
-	req := make(map[string]string)
-	req["event"] = "subscribe"
-	req["channel"] = channel
-
-	if len(params) > 0 {
-		for k, v := range params {
-			req[k] = v
-		}
+	if b.Verbose {
+		log.Debugf("%v sending message to websocket %v", b.Name, data)
 	}
-	return b.WsSend(req)
+	return b.WebsocketConn.WriteMessage(websocket.TextMessage, json)
 }
 
 // WsSendAuth sends a autheticated event payload
@@ -98,7 +89,7 @@ func (b *Bitfinex) WsSendAuth() error {
 
 	req["authPayload"] = payload
 
-	return b.WsSend(req)
+	return b.wsSend(req)
 }
 
 // WsSendUnauth sends an unauthenticated payload
@@ -106,7 +97,7 @@ func (b *Bitfinex) WsSendUnauth() error {
 	req := make(map[string]string)
 	req["event"] = "unauth"
 
-	return b.WsSend(req)
+	return b.wsSend(req)
 }
 
 // WsAddSubscriptionChannel adds a new subscription channel to the
@@ -130,7 +121,6 @@ func (b *Bitfinex) WsConnect() error {
 		return errors.New(exchange.WebsocketNotEnabled)
 	}
 
-	var channels = []string{"book", "trades", "ticker"}
 	var Dialer websocket.Dialer
 	var err error
 
@@ -145,12 +135,12 @@ func (b *Bitfinex) WsConnect() error {
 
 	b.WebsocketConn, _, err = Dialer.Dial(b.Websocket.GetWebsocketURL(), http.Header{})
 	if err != nil {
-		return fmt.Errorf("unable to connect to Websocket. Error: %s", err)
+		return fmt.Errorf("%v unable to connect to Websocket. Error: %s", b.Name, err)
 	}
 
 	_, resp, err := b.WebsocketConn.ReadMessage()
 	if err != nil {
-		return fmt.Errorf("unable to read from Websocket. Error: %s", err)
+		return fmt.Errorf("%v unable to read from Websocket. Error: %s", b.Name, err)
 	}
 
 	var hs WebsocketHandshake
@@ -159,26 +149,13 @@ func (b *Bitfinex) WsConnect() error {
 		return err
 	}
 
+	b.GenerateDefaultSubscriptions()
 	if hs.Event == "info" {
 		if b.Verbose {
 			log.Debugf("%s Connected to Websocket.\n", b.GetName())
 		}
 	}
 
-	for _, x := range channels {
-		for _, y := range b.EnabledPairs {
-			params := make(map[string]string)
-			if x == "book" {
-				params["prec"] = "P0"
-			}
-			params["pair"] = y.String()
-			err = b.WsSubscribe(x, params)
-			if err != nil {
-				return err
-			}
-		}
-	}
-
 	if b.AuthenticatedAPISupport {
 		err = b.WsSendAuth()
 		if err != nil {
@@ -214,11 +191,6 @@ func (b *Bitfinex) WsDataHandler() {
 	b.Websocket.Wg.Add(1)
 
 	defer func() {
-		err := b.WebsocketConn.Close()
-		if err != nil {
-			b.Websocket.DataHandler <- fmt.Errorf("bitfinex_websocket.go - closing websocket connection error %s",
-				err)
-		}
 		b.Websocket.Wg.Done()
 	}()
 
@@ -237,12 +209,13 @@ func (b *Bitfinex) WsDataHandler() {
 			if stream.Type == websocket.TextMessage {
 				var result interface{}
 				common.JSONDecode(stream.Raw, &result)
-
 				switch reflect.TypeOf(result).String() {
 				case "map[string]interface {}":
 					eventData := result.(map[string]interface{})
 					event := eventData["event"]
-
+					if b.Verbose {
+						log.Debugf("%v Received message. Type '%v' Message: %v", b.Name, event, eventData)
+					}
 					switch event {
 					case "subscribed":
 						b.WsAddSubscriptionChannel(int(eventData["chanId"].(float64)),
@@ -624,3 +597,51 @@ func (b *Bitfinex) WsUpdateOrderbook(p currency.Pair, assetType string, book Web
 
 	return nil
 }
+
+// GenerateDefaultSubscriptions Adds default subscriptions to websocket to be handled by ManageSubscriptions()
+func (b *Bitfinex) GenerateDefaultSubscriptions() {
+	var channels = []string{"book", "trades", "ticker"}
+	subscriptions := []exchange.WebsocketChannelSubscription{}
+	for i := range channels {
+		for j := range b.EnabledPairs {
+			params := make(map[string]interface{})
+			if channels[i] == "book" {
+				params["prec"] = "P0"
+			}
+			subscriptions = append(subscriptions, exchange.WebsocketChannelSubscription{
+				Channel:  channels[i],
+				Currency: b.EnabledPairs[j],
+				Params:   params,
+			})
+		}
+	}
+	b.Websocket.SubscribeToChannels(subscriptions)
+}
+
+// Subscribe sends a websocket message to receive data from the channel
+func (b *Bitfinex) Subscribe(channelToSubscribe exchange.WebsocketChannelSubscription) error {
+	req := make(map[string]interface{})
+	req["event"] = "subscribe"
+	req["channel"] = channelToSubscribe.Channel
+	req["pair"] = channelToSubscribe.Currency.String()
+	if len(channelToSubscribe.Params) > 0 {
+		for k, v := range channelToSubscribe.Params {
+			req[k] = v
+		}
+	}
+	return b.wsSend(req)
+}
+
+// Unsubscribe sends a websocket message to stop receiving data from the channel
+func (b *Bitfinex) Unsubscribe(channelToSubscribe exchange.WebsocketChannelSubscription) error {
+	req := make(map[string]interface{})
+	req["event"] = "unsubscribe"
+	req["channel"] = channelToSubscribe.Channel
+
+	if len(channelToSubscribe.Params) > 0 {
+		for k, v := range channelToSubscribe.Params {
+			req[k] = v
+		}
+	}
+	return b.wsSend(req)
+}

exchanges/bitfinex/bitfinex_wrapper.go

@@ -454,3 +454,17 @@ func (b *Bitfinex) GetOrderHistory(getOrdersRequest *exchange.GetOrdersRequest)
 
 	return orders, nil
 }
+
+// SubscribeToWebsocketChannels appends to ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle subscribing
+func (b *Bitfinex) SubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	b.Websocket.SubscribeToChannels(channels)
+	return nil
+}
+
+// UnsubscribeToWebsocketChannels removes from ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle unsubscribing
+func (b *Bitfinex) UnsubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	b.Websocket.UnsubscribeToChannels(channels)
+	return nil
+}

exchanges/bitflyer/bitflyer_wrapper.go

@@ -242,3 +242,15 @@ func (b *Bitflyer) GetFeeByType(feeBuilder *exchange.FeeBuilder) (float64, error
 	}
 	return b.GetFee(feeBuilder)
 }
+
+// SubscribeToWebsocketChannels appends to ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle subscribing
+func (b *Bitflyer) SubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	return common.ErrFunctionNotSupported
+}
+
+// UnsubscribeToWebsocketChannels removes from ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle unsubscribing
+func (b *Bitflyer) UnsubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	return common.ErrFunctionNotSupported
+}

exchanges/bithumb/bithumb_wrapper.go

@@ -420,3 +420,15 @@ func (b *Bithumb) GetOrderHistory(getOrdersRequest *exchange.GetOrdersRequest) (
 
 	return orders, nil
 }
+
+// SubscribeToWebsocketChannels appends to ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle subscribing
+func (b *Bithumb) SubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	return common.ErrFunctionNotSupported
+}
+
+// UnsubscribeToWebsocketChannels removes from ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle unsubscribing
+func (b *Bithumb) UnsubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	return common.ErrFunctionNotSupported
+}

exchanges/bitmex/bitmex.go

@@ -7,6 +7,7 @@ import (
 	"net/http"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/gorilla/websocket"
@@ -23,6 +24,7 @@ import (
 type Bitmex struct {
 	exchange.Base
 	WebsocketConn *websocket.Conn
+	wsRequestMtx  sync.Mutex
 }
 
 const (
@@ -135,7 +137,9 @@ func (b *Bitmex) SetDefaults() {
 	b.SupportsAutoPairUpdating = true
 	b.WebsocketInit()
 	b.Websocket.Functionality = exchange.WebsocketTradeDataSupported |
-		exchange.WebsocketOrderbookSupported
+		exchange.WebsocketOrderbookSupported |
+		exchange.WebsocketSubscribeSupported |
+		exchange.WebsocketUnsubscribeSupported
 }
 
 // Setup takes in the supplied exchange configuration details and sets params
@@ -174,8 +178,11 @@ func (b *Bitmex) Setup(exch *config.ExchangeConfig) {
 			log.Fatal(err)
 		}
 		err = b.WebsocketSetup(b.WsConnector,
+			b.Subscribe,
+			b.Unsubscribe,
 			exch.Name,
 			exch.Websocket,
+			exch.Verbose,
 			bitmexWSURL,
 			exch.WebsocketURL)
 		if err != nil {

exchanges/bitmex/bitmex_websocket.go

@@ -105,16 +105,7 @@ func (b *Bitmex) WsConnector() error {
 	}
 
 	go b.wsHandleIncomingData()
-
-	err = b.websocketSubscribe()
-	if err != nil {
-		closeError := b.WebsocketConn.Close()
-		if closeError != nil {
-			return fmt.Errorf("bitmex_websocket.go error - Websocket connection could not close %s",
-				closeError)
-		}
-		return err
-	}
+	b.GenerateDefaultSubscriptions()
 
 	if b.AuthenticatedAPISupport {
 		err := b.websocketSendAuth()
@@ -143,11 +134,6 @@ func (b *Bitmex) wsHandleIncomingData() {
 	b.Websocket.Wg.Add(1)
 
 	defer func() {
-		err := b.WebsocketConn.Close()
-		if err != nil {
-			b.Websocket.DataHandler <- fmt.Errorf("bitmex_websocket.go - Unable to close Websocket connection. Error: %s",
-				err)
-		}
 		b.Websocket.Wg.Done()
 	}()
 
@@ -170,7 +156,7 @@ func (b *Bitmex) wsHandleIncomingData() {
 			}
 
 			if common.StringContains(message, "ping") {
-				err = b.WebsocketConn.WriteJSON("pong")
+				err = b.wsSend("pong")
 				if err != nil {
 					b.Websocket.DataHandler <- err
 					continue
@@ -398,26 +384,42 @@ func (b *Bitmex) processOrderbook(data []OrderBookL2, action string, currencyPai
 	return nil
 }
 
-// WebsocketSubscribe subscribes to a websocket channel
-func (b *Bitmex) websocketSubscribe() error {
+// GenerateDefaultSubscriptions Adds default subscriptions to websocket to be handled by ManageSubscriptions()
+func (b *Bitmex) GenerateDefaultSubscriptions() {
 	contracts := b.GetEnabledCurrencies()
+	channels := []string{bitmexWSOrderbookL2, bitmexWSTrade}
+	subscriptions := []exchange.WebsocketChannelSubscription{
+		{
+			Channel: bitmexWSAnnouncement,
+		},
+	}
+	for i := range channels {
+		for j := range contracts {
+			subscriptions = append(subscriptions, exchange.WebsocketChannelSubscription{
+				Channel:  fmt.Sprintf("%v:%v", channels[i], contracts[j].String()),
+				Currency: contracts[j],
+			})
+		}
+	}
+	b.Websocket.SubscribeToChannels(subscriptions)
+}
 
-	// Subscriber
+// Subscribe subscribes to a websocket channel
+func (b *Bitmex) Subscribe(channelToSubscribe exchange.WebsocketChannelSubscription) error {
 	var subscriber WebsocketRequest
 	subscriber.Command = "subscribe"
+	subscriber.Arguments = append(subscriber.Arguments, channelToSubscribe.Channel)
+	return b.wsSend(subscriber)
+}
 
-	// Announcement subscribe
-	subscriber.Arguments = append(subscriber.Arguments, bitmexWSAnnouncement)
-
-	for _, contract := range contracts {
-		// Orderbook and Trade subscribe
-		// NOTE more added here in future
-		subscriber.Arguments = append(subscriber.Arguments,
-			bitmexWSOrderbookL2+":"+contract.String(),
-			bitmexWSTrade+":"+contract.String())
-	}
-
-	return b.WebsocketConn.WriteJSON(subscriber)
+// Unsubscribe sends a websocket message to stop receiving data from the channel
+func (b *Bitmex) Unsubscribe(channelToSubscribe exchange.WebsocketChannelSubscription) error {
+	var subscriber WebsocketRequest
+	subscriber.Command = "unsubscribe"
+	subscriber.Arguments = append(subscriber.Arguments,
+		channelToSubscribe.Params["args"],
+		channelToSubscribe.Channel+":"+channelToSubscribe.Currency.String())
+	return b.wsSend(subscriber)
 }
 
 // WebsocketSendAuth sends an authenticated subscription
@@ -433,5 +435,15 @@ func (b *Bitmex) websocketSendAuth() error {
 	sendAuth.Command = "authKeyExpires"
 	sendAuth.Arguments = append(sendAuth.Arguments, b.APIKey, timestamp,
 		signature)
-	return b.WebsocketConn.WriteJSON(sendAuth)
+	return b.wsSend(sendAuth)
+}
+
+// WsSend sends data to the websocket server
+func (b *Bitmex) wsSend(data interface{}) error {
+	b.wsRequestMtx.Lock()
+	defer b.wsRequestMtx.Unlock()
+	if b.Verbose {
+		log.Debugf("%v sending message to websocket %v", b.Name, data)
+	}
+	return b.WebsocketConn.WriteJSON(data)
 }

exchanges/bitmex/bitmex_wrapper.go

@@ -396,3 +396,17 @@ func (b *Bitmex) GetOrderHistory(getOrdersRequest *exchange.GetOrdersRequest) ([
 
 	return orders, nil
 }
+
+// SubscribeToWebsocketChannels appends to ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle subscribing
+func (b *Bitmex) SubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	b.Websocket.SubscribeToChannels(channels)
+	return nil
+}
+
+// UnsubscribeToWebsocketChannels removes from ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle unsubscribing
+func (b *Bitmex) UnsubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	b.Websocket.UnsubscribeToChannels(channels)
+	return nil
+}

exchanges/bitstamp/bitstamp.go

@@ -9,6 +9,7 @@ import (
 	"reflect"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/thrasher-/gocryptotrader/common"
@@ -63,6 +64,7 @@ type Bitstamp struct {
 	exchange.Base
 	Balance       Balances
 	WebsocketConn WebsocketConn
+	wsRequestMtx  sync.Mutex
 }
 
 // SetDefaults sets default for Bitstamp
@@ -88,7 +90,9 @@ func (b *Bitstamp) SetDefaults() {
 	b.APIUrl = b.APIUrlDefault
 	b.WebsocketInit()
 	b.Websocket.Functionality = exchange.WebsocketOrderbookSupported |
-		exchange.WebsocketTradeDataSupported
+		exchange.WebsocketTradeDataSupported |
+		exchange.WebsocketSubscribeSupported |
+		exchange.WebsocketUnsubscribeSupported
 }
 
 // Setup sets configuration values to bitstamp
@@ -133,8 +137,11 @@ func (b *Bitstamp) Setup(exch *config.ExchangeConfig) {
 			log.Fatal(err)
 		}
 		err = b.WebsocketSetup(b.WsConnect,
+			b.Subscribe,
+			b.Unsubscribe,
 			exch.Name,
 			exch.Websocket,
+			exch.Verbose,
 			BitstampPusherKey,
 			exch.WebsocketURL)
 		if err != nil {

exchanges/bitstamp/bitstamp_types.go

@@ -1,5 +1,7 @@
 package bitstamp
 
+import pusher "github.com/toorop/go-pusher"
+
 // Ticker holds ticker information
 type Ticker struct {
 	Last      float64 `json:"last,string"`
@@ -157,3 +159,35 @@ const (
 	internationalWithdrawal string = "international"
 	errStr                  string = "error"
 )
+
+// WebsocketConn defines a pusher websocket connection
+type WebsocketConn struct {
+	Client *pusher.Client
+	Data   chan *pusher.Event
+	Trade  chan *pusher.Event
+}
+
+// PusherOrderbook holds order book information to be pushed
+type PusherOrderbook struct {
+	Asks      [][]string `json:"asks"`
+	Bids      [][]string `json:"bids"`
+	Timestamp int64      `json:"timestamp,string"`
+}
+
+// PusherTrade holds trade information to be pushed
+type PusherTrade struct {
+	Price       float64 `json:"price"`
+	Amount      float64 `json:"amount"`
+	ID          int64   `json:"id"`
+	Type        int64   `json:"type"`
+	Timestamp   int64   `json:"timestamp,string"`
+	BuyOrderID  int64   `json:"buy_order_id"`
+	SellOrderID int64   `json:"sell_order_id"`
+}
+
+// PusherOrders defines order information
+type PusherOrders struct {
+	ID     int64   `json:"id"`
+	Amount float64 `json:"amount"`
+	Price  float64 `json:""`
+}

exchanges/bitstamp/bitstamp_websocket.go

@@ -15,38 +15,6 @@ import (
 	pusher "github.com/toorop/go-pusher"
 )
 
-// WebsocketConn defins a pusher websocket connection
-type WebsocketConn struct {
-	Client *pusher.Client
-	Data   chan *pusher.Event
-	Trade  chan *pusher.Event
-}
-
-// PusherOrderbook holds order book information to be pushed
-type PusherOrderbook struct {
-	Asks      [][]string `json:"asks"`
-	Bids      [][]string `json:"bids"`
-	Timestamp int64      `json:"timestamp,string"`
-}
-
-// PusherTrade holds trade information to be pushed
-type PusherTrade struct {
-	Price       float64 `json:"price"`
-	Amount      float64 `json:"amount"`
-	ID          int64   `json:"id"`
-	Type        int64   `json:"type"`
-	Timestamp   int64   `json:"timestamp,string"`
-	BuyOrderID  int64   `json:"buy_order_id"`
-	SellOrderID int64   `json:"sell_order_id"`
-}
-
-// PusherOrders defines order information
-type PusherOrders struct {
-	ID     int64   `json:"id"`
-	Amount float64 `json:"amount"`
-	Price  float64 `json:""`
-}
-
 const (
 	// BitstampPusherKey holds the current pusher key
 	BitstampPusherKey = "de504dc5763aeef9ff52"
@@ -98,7 +66,7 @@ func (b *Bitstamp) WsConnect() error {
 	if err != nil {
 		return fmt.Errorf("%s Websocket Bind error: %s", b.GetName(), err)
 	}
-
+	b.GenerateDefaultSubscriptions()
 	go b.WsReadData()
 
 	for _, p := range b.GetEnabledCurrencies() {
@@ -141,32 +109,49 @@ func (b *Bitstamp) WsConnect() error {
 			Exchange: b.GetName(),
 		}
 
-		err = b.WebsocketConn.Client.Subscribe(fmt.Sprintf("live_trades_%s",
-			p.Lower().String()))
-
-		if err != nil {
-			return fmt.Errorf("%s Websocket Trade subscription error: %s",
-				b.GetName(),
-				err)
-		}
-
-		err = b.WebsocketConn.Client.Subscribe(fmt.Sprintf("diff_order_book_%s",
-			p.Lower().String()))
-
-		if err != nil {
-			return fmt.Errorf("%s Websocket Trade subscription error: %s",
-				b.GetName(),
-				err)
-		}
-
 	}
 	return nil
 }
 
+// GenerateDefaultSubscriptions Adds default subscriptions to websocket to be handled by ManageSubscriptions()
+func (b *Bitstamp) GenerateDefaultSubscriptions() {
+	var channels = []string{"live_trades_", "diff_order_book_"}
+	enabledCurrencies := b.GetEnabledCurrencies()
+	subscriptions := []exchange.WebsocketChannelSubscription{}
+	for i := range channels {
+		for j := range enabledCurrencies {
+			subscriptions = append(subscriptions, exchange.WebsocketChannelSubscription{
+				Channel:  fmt.Sprintf("%v%v", channels[i], enabledCurrencies[j].Lower().String()),
+				Currency: enabledCurrencies[j],
+			})
+		}
+	}
+	b.Websocket.SubscribeToChannels(subscriptions)
+}
+
+// Subscribe sends a websocket message to receive data from the channel
+func (b *Bitstamp) Subscribe(channelToSubscribe exchange.WebsocketChannelSubscription) error {
+	b.wsRequestMtx.Lock()
+	defer b.wsRequestMtx.Unlock()
+	if b.Verbose {
+		log.Debugf("%v sending message to websocket %v", b.Name, channelToSubscribe)
+	}
+	return b.WebsocketConn.Client.Subscribe(channelToSubscribe.Channel)
+}
+
+// Unsubscribe sends a websocket message to stop receiving data from the channel
+func (b *Bitstamp) Unsubscribe(channelToSubscribe exchange.WebsocketChannelSubscription) error {
+	b.wsRequestMtx.Lock()
+	defer b.wsRequestMtx.Unlock()
+	if b.Verbose {
+		log.Debugf("%v sending message to websocket %v", b.Name, channelToSubscribe)
+	}
+	return b.WebsocketConn.Client.Unsubscribe(channelToSubscribe.Channel)
+}
+
 // WsReadData reads data coming from bitstamp websocket connection
 func (b *Bitstamp) WsReadData() {
 	b.Websocket.Wg.Add(1)
-
 	defer func() {
 		err := b.WebsocketConn.Client.Close()
 		if err != nil {

exchanges/bitstamp/bitstamp_wrapper.go

@@ -399,3 +399,17 @@ func (b *Bitstamp) GetOrderHistory(getOrdersRequest *exchange.GetOrdersRequest)
 
 	return orders, nil
 }
+
+// SubscribeToWebsocketChannels appends to ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle subscribing
+func (b *Bitstamp) SubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	b.Websocket.SubscribeToChannels(channels)
+	return nil
+}
+
+// UnsubscribeToWebsocketChannels removes from ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle unsubscribing
+func (b *Bitstamp) UnsubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	b.Websocket.UnsubscribeToChannels(channels)
+	return nil
+}

exchanges/bittrex/bittrex_wrapper.go

@@ -397,3 +397,15 @@ func (b *Bittrex) GetOrderHistory(getOrdersRequest *exchange.GetOrdersRequest) (
 
 	return orders, nil
 }
+
+// SubscribeToWebsocketChannels appends to ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle subscribing
+func (b *Bittrex) SubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	return common.ErrFunctionNotSupported
+}
+
+// UnsubscribeToWebsocketChannels removes from ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle unsubscribing
+func (b *Bittrex) UnsubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	return common.ErrFunctionNotSupported
+}

exchanges/btcc/btcc.go

@@ -1,6 +1,7 @@
 package btcc
 
 import (
+	"sync"
 	"time"
 
 	"github.com/gorilla/websocket"
@@ -23,7 +24,8 @@ const (
 // been dropped
 type BTCC struct {
 	exchange.Base
-	Conn *websocket.Conn
+	Conn         *websocket.Conn
+	wsRequestMtx sync.Mutex
 }
 
 // SetDefaults sets default values for the exchange
@@ -46,6 +48,9 @@ func (b *BTCC) SetDefaults() {
 		request.NewRateLimit(time.Second, btccUnauthRate),
 		common.NewHTTPClientWithTimeout(exchange.DefaultHTTPTimeout))
 	b.WebsocketInit()
+	b.Websocket.Functionality =
+		exchange.WebsocketSubscribeSupported |
+			exchange.WebsocketUnsubscribeSupported
 }
 
 // Setup is run on startup to setup exchange with config values
@@ -86,8 +91,11 @@ func (b *BTCC) Setup(exch *config.ExchangeConfig) {
 			log.Fatal(err)
 		}
 		err = b.WebsocketSetup(b.WsConnect,
+			b.Subscribe,
+			b.Unsubscribe,
 			exch.Name,
 			exch.Websocket,
+			exch.Verbose,
 			btccSocketioAddress,
 			exch.WebsocketURL)
 		if err != nil {

exchanges/btcc/btcc_websocket.go

@@ -70,18 +70,9 @@ func (b *BTCC) WsConnect() error {
 	}
 
 	go b.WsHandleData()
+	b.GenerateDefaultSubscriptions()
 
-	err = b.WsSubscribeToOrderbook()
-	if err != nil {
-		return err
-	}
-
-	err = b.WsSubcribeToTicker()
-	if err != nil {
-		return err
-	}
-
-	return b.WsSubcribeToTrades()
+	return nil
 }
 
 // WsReadData reads data from the websocket connection
@@ -251,33 +242,8 @@ func (b *BTCC) WsHandleData() {
 	}
 }
 
-// WsSubscribeAllTickers subscribes to a ticker channel
-func (b *BTCC) WsSubscribeAllTickers() error {
-	mtx.Lock()
-	defer mtx.Unlock()
-
-	return b.Conn.WriteJSON(WsOutgoing{
-		Action: "SubscribeAllTickers",
-	})
-}
-
-// WsUnSubscribeAllTickers unsubscribes from a ticker channel
-func (b *BTCC) WsUnSubscribeAllTickers() error {
-	mtx.Lock()
-	defer mtx.Unlock()
-
-	return b.Conn.WriteJSON(WsOutgoing{
-		Action: "UnSubscribeAllTickers",
-	})
-}
-
 // WsUpdateCurrencyPairs updates currency pairs from the websocket connection
 func (b *BTCC) WsUpdateCurrencyPairs() error {
-	err := b.WsSubscribeAllTickers()
-	if err != nil {
-		return err
-	}
-
 	var currencyResponse WsResponseMain
 	for {
 		_, resp, err := b.Conn.ReadMessage()
@@ -313,8 +279,6 @@ func (b *BTCC) WsUpdateCurrencyPairs() error {
 					err)
 			}
 
-			return b.WsUnSubscribeAllTickers()
-
 		case "Heartbeat":
 
 		default:
@@ -324,61 +288,6 @@ func (b *BTCC) WsUpdateCurrencyPairs() error {
 	}
 }
 
-// WsSubscribeToOrderbook subscribes to an orderbook channel
-func (b *BTCC) WsSubscribeToOrderbook() error {
-	mtx.Lock()
-	defer mtx.Unlock()
-
-	for _, pair := range b.GetEnabledCurrencies() {
-		formattedPair := exchange.FormatExchangeCurrency(b.GetName(), pair)
-		err := b.Conn.WriteJSON(WsOutgoing{
-			Action: "SubOrderBook",
-			Symbol: formattedPair.String(),
-			Len:    100})
-		if err != nil {
-			return err
-		}
-	}
-	return nil
-}
-
-// WsSubcribeToTicker subscribes to a ticker channel
-func (b *BTCC) WsSubcribeToTicker() error {
-	mtx.Lock()
-	defer mtx.Unlock()
-
-	for _, pair := range b.GetEnabledCurrencies() {
-		formattedPair := exchange.FormatExchangeCurrency(b.GetName(), pair)
-		err := b.Conn.WriteJSON(WsOutgoing{
-			Action: "Subscribe",
-			Symbol: formattedPair.String(),
-		})
-		if err != nil {
-			return err
-		}
-	}
-	return nil
-}
-
-// WsSubcribeToTrades subscribes to a trade channel
-func (b *BTCC) WsSubcribeToTrades() error {
-	mtx.Lock()
-	defer mtx.Unlock()
-
-	for _, pair := range b.GetEnabledCurrencies() {
-		formattedPair := exchange.FormatExchangeCurrency(b.GetName(), pair)
-		err := b.Conn.WriteJSON(WsOutgoing{
-			Action: "GetTrades",
-			Symbol: formattedPair.String(),
-			Count:  100,
-		})
-		if err != nil {
-			return err
-		}
-	}
-	return nil
-}
-
 // WsProcessOrderbookSnapshot processes a new orderbook snapshot
 func (b *BTCC) WsProcessOrderbookSnapshot(ob *WsOrderbookSnapshot) error {
 	var asks, bids []orderbook.Item
@@ -559,3 +468,72 @@ func (b *BTCC) WsProcessOldOrderbookSnapshot(ob WsOrderbookSnapshotOld, symbol s
 
 	return nil
 }
+
+// GenerateDefaultSubscriptions Adds default subscriptions to websocket to be handled by ManageSubscriptions()
+func (b *BTCC) GenerateDefaultSubscriptions() {
+	subscriptions := []exchange.WebsocketChannelSubscription{}
+	subscriptions = append(subscriptions, exchange.WebsocketChannelSubscription{
+		Channel: "SubscribeAllTickers",
+	})
+
+	var channels = []string{"SubOrderBook", "GetTrades", "Subscribe"}
+	enabledCurrencies := b.GetEnabledCurrencies()
+	for i := range channels {
+		for j := range enabledCurrencies {
+			params := make(map[string]interface{})
+			if channels[i] == "SubOrderBook" {
+				params["len"] = "100"
+			} else if channels[i] == "GetTrades" {
+				params["count"] = "100"
+			}
+			subscriptions = append(subscriptions, exchange.WebsocketChannelSubscription{
+				Channel:  channels[i],
+				Currency: enabledCurrencies[j],
+				Params:   params,
+			})
+		}
+	}
+	b.Websocket.SubscribeToChannels(subscriptions)
+}
+
+// Subscribe sends a websocket message to receive data from the channel
+func (b *BTCC) Subscribe(channelToSubscribe exchange.WebsocketChannelSubscription) error {
+	subscription := WsOutgoing{
+		Action: channelToSubscribe.Channel,
+		Symbol: channelToSubscribe.Currency.String(),
+	}
+	if subscription.Action == "SubOrderBook" {
+		subscription.Len = 100
+	} else if subscription.Action == "GetTrades" {
+		subscription.Count = 100
+	}
+
+	return b.wsSend(subscription)
+}
+
+// Unsubscribe sends a websocket message to stop receiving data from the channel
+func (b *BTCC) Unsubscribe(channelToSubscribe exchange.WebsocketChannelSubscription) error {
+	subscription := WsOutgoing{}
+	switch channelToSubscribe.Channel {
+	case "SubOrderBook":
+		subscription.Action = "UnSubOrderBook"
+		subscription.Symbol = channelToSubscribe.Currency.String()
+	case "Subscribe":
+		subscription.Action = "UnSubscribe"
+		subscription.Symbol = channelToSubscribe.Currency.String()
+	case "SubscribeAllTickers":
+		subscription.Action = "UnSubscribeAllTickers"
+	}
+
+	return b.wsSend(subscription)
+}
+
+// WsSend sends data to the websocket server
+func (b *BTCC) wsSend(data interface{}) error {
+	b.wsRequestMtx.Lock()
+	defer b.wsRequestMtx.Unlock()
+	if b.Verbose {
+		log.Debugf("%v sending message to websocket %v", b.Name, data)
+	}
+	return b.Conn.WriteJSON(data)
+}

exchanges/btcc/btcc_wrapper.go

@@ -175,3 +175,17 @@ func (b *BTCC) GetActiveOrders(getOrdersRequest *exchange.GetOrdersRequest) ([]e
 func (b *BTCC) GetOrderHistory(getOrdersRequest *exchange.GetOrdersRequest) ([]exchange.OrderDetail, error) {
 	return nil, common.ErrNotYetImplemented
 }
+
+// SubscribeToWebsocketChannels appends to ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle subscribing
+func (b *BTCC) SubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	b.Websocket.SubscribeToChannels(channels)
+	return nil
+}
+
+// UnsubscribeToWebsocketChannels removes from ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle unsubscribing
+func (b *BTCC) UnsubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	b.Websocket.UnsubscribeToChannels(channels)
+	return nil
+}

exchanges/btcmarkets/btcmarkets_wrapper.go

@@ -464,3 +464,15 @@ func (b *BTCMarkets) GetOrderHistory(getOrdersRequest *exchange.GetOrdersRequest
 
 	return orders, nil
 }
+
+// SubscribeToWebsocketChannels appends to ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle subscribing
+func (b *BTCMarkets) SubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	return common.ErrFunctionNotSupported
+}
+
+// UnsubscribeToWebsocketChannels removes from ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle unsubscribing
+func (b *BTCMarkets) UnsubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	return common.ErrFunctionNotSupported
+}

exchanges/btse/btse.go

@@ -6,6 +6,7 @@ import (
 	"net/http"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/gorilla/websocket"
@@ -22,6 +23,7 @@ import (
 type BTSE struct {
 	exchange.Base
 	WebsocketConn *websocket.Conn
+	wsRequestMtx  sync.Mutex
 }
 
 const (
@@ -66,7 +68,9 @@ func (b *BTSE) SetDefaults() {
 	b.SupportsRESTTickerBatching = false
 	b.WebsocketInit()
 	b.Websocket.Functionality = exchange.WebsocketOrderbookSupported |
-		exchange.WebsocketTickerSupported
+		exchange.WebsocketTickerSupported |
+		exchange.WebsocketSubscribeSupported |
+		exchange.WebsocketUnsubscribeSupported
 }
 
 // Setup takes in the supplied exchange configuration details and sets params
@@ -106,8 +110,11 @@ func (b *BTSE) Setup(exch *config.ExchangeConfig) {
 			log.Fatal(err)
 		}
 		err = b.WebsocketSetup(b.WsConnect,
+			b.Subscribe,
+			b.Unsubscribe,
 			exch.Name,
 			exch.Websocket,
+			exch.Verbose,
 			btseWebsocket,
 			exch.WebsocketURL)
 		if err != nil {

exchanges/btse/btse_websocket.go

@@ -21,31 +21,6 @@ const (
 	btseWebsocket = "wss://ws.btse.com/api/ws-feed"
 )
 
-// WebsocketSubscriber subscribes to websocket channels with respect to enabled
-// currencies
-func (b *BTSE) WebsocketSubscriber() error {
-	subscribe := websocketSubscribe{
-		Type: "subscribe",
-		Channels: []websocketChannel{
-			{
-				Name:       "snapshot",
-				ProductIDs: b.EnabledPairs.Strings(),
-			},
-			{
-				Name:       "ticker",
-				ProductIDs: b.EnabledPairs.Strings(),
-			},
-		},
-	}
-
-	data, err := common.JSONEncode(subscribe)
-	if err != nil {
-		return err
-	}
-
-	return b.WebsocketConn.WriteMessage(websocket.TextMessage, data)
-}
-
 // WsConnect connects the websocket client
 func (b *BTSE) WsConnect() error {
 	if !b.Websocket.IsEnabled() || !b.IsEnabled() {
@@ -72,12 +47,12 @@ func (b *BTSE) WsConnect() error {
 			b.Name, err)
 	}
 
-	err = b.WebsocketSubscriber()
 	if err != nil {
 		return err
 	}
 
 	go b.WsHandleData()
+	b.GenerateDefaultSubscriptions()
 
 	return nil
 }
@@ -98,11 +73,6 @@ func (b *BTSE) WsHandleData() {
 	b.Websocket.Wg.Add(1)
 
 	defer func() {
-		err := b.WebsocketConn.Close()
-		if err != nil {
-			b.Websocket.DataHandler <- fmt.Errorf("%s - Unable to to close Websocket connection. Error: %s",
-				b.Name, err)
-		}
 		b.Websocket.Wg.Done()
 	}()
 
@@ -137,7 +107,6 @@ func (b *BTSE) WsHandleData() {
 				b.Websocket.DataHandler <- err
 				continue
 			}
-
 			switch msgType.Type {
 			case "ticker":
 				var t wsTicker
@@ -234,3 +203,61 @@ func (b *BTSE) wsProcessSnapshot(snapshot *websocketOrderbookSnapshot) error {
 
 	return nil
 }
+
+// GenerateDefaultSubscriptions Adds default subscriptions to websocket to be handled by ManageSubscriptions()
+func (b *BTSE) GenerateDefaultSubscriptions() {
+	var channels = []string{"snapshot", "ticker"}
+	enabledCurrencies := b.GetEnabledCurrencies()
+	subscriptions := []exchange.WebsocketChannelSubscription{}
+	for i := range channels {
+		for j := range enabledCurrencies {
+			subscriptions = append(subscriptions, exchange.WebsocketChannelSubscription{
+				Channel:  channels[i],
+				Currency: enabledCurrencies[j],
+			})
+		}
+	}
+	b.Websocket.SubscribeToChannels(subscriptions)
+}
+
+// Subscribe sends a websocket message to receive data from the channel
+func (b *BTSE) Subscribe(channelToSubscribe exchange.WebsocketChannelSubscription) error {
+	subscribe := websocketSubscribe{
+		Type: "subscribe",
+		Channels: []websocketChannel{
+			{
+				Name:       channelToSubscribe.Channel,
+				ProductIDs: []string{channelToSubscribe.Currency.String()},
+			},
+		},
+	}
+	return b.wsSend(subscribe)
+}
+
+// Unsubscribe sends a websocket message to stop receiving data from the channel
+func (b *BTSE) Unsubscribe(channelToSubscribe exchange.WebsocketChannelSubscription) error {
+	subscribe := websocketSubscribe{
+		Type: "unsubscribe",
+		Channels: []websocketChannel{
+			{
+				Name:       channelToSubscribe.Channel,
+				ProductIDs: []string{channelToSubscribe.Currency.String()},
+			},
+		},
+	}
+	return b.wsSend(subscribe)
+}
+
+// WsSend sends data to the websocket server
+func (b *BTSE) wsSend(data interface{}) error {
+	b.wsRequestMtx.Lock()
+	defer b.wsRequestMtx.Unlock()
+	if b.Verbose {
+		log.Debugf("%v sending message to websocket %v", b.Name, data)
+	}
+	json, err := common.JSONEncode(data)
+	if err != nil {
+		return err
+	}
+	return b.WebsocketConn.WriteMessage(websocket.TextMessage, json)
+}

exchanges/btse/btse_wrapper.go

@@ -358,3 +358,17 @@ func (b *BTSE) GetFeeByType(feeBuilder *exchange.FeeBuilder) (float64, error) {
 	}
 	return b.GetFee(feeBuilder)
 }
+
+// SubscribeToWebsocketChannels appends to ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle subscribing
+func (b *BTSE) SubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	b.Websocket.SubscribeToChannels(channels)
+	return nil
+}
+
+// UnsubscribeToWebsocketChannels removes from ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle unsubscribing
+func (b *BTSE) UnsubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	b.Websocket.UnsubscribeToChannels(channels)
+	return nil
+}

exchanges/coinbasepro/coinbasepro.go

@@ -8,6 +8,7 @@ import (
 	"net/url"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/gorilla/websocket"
@@ -61,6 +62,7 @@ const (
 type CoinbasePro struct {
 	exchange.Base
 	WebsocketConn *websocket.Conn
+	wsRequestMtx  sync.Mutex
 }
 
 // SetDefaults sets default values for the exchange
@@ -88,7 +90,9 @@ func (c *CoinbasePro) SetDefaults() {
 	c.APIUrl = c.APIUrlDefault
 	c.WebsocketInit()
 	c.Websocket.Functionality = exchange.WebsocketTickerSupported |
-		exchange.WebsocketOrderbookSupported
+		exchange.WebsocketOrderbookSupported |
+		exchange.WebsocketSubscribeSupported |
+		exchange.WebsocketUnsubscribeSupported
 }
 
 // Setup initialises the exchange parameters with the current configuration
@@ -132,8 +136,11 @@ func (c *CoinbasePro) Setup(exch *config.ExchangeConfig) {
 			log.Fatal(err)
 		}
 		err = c.WebsocketSetup(c.WsConnect,
+			c.Subscribe,
+			c.Unsubscribe,
 			exch.Name,
 			exch.Websocket,
+			exch.Verbose,
 			coinbaseproWebsocketURL,
 			exch.WebsocketURL)
 		if err != nil {

exchanges/coinbasepro/coinbasepro_websocket.go

@@ -13,46 +13,13 @@ import (
 	"github.com/thrasher-/gocryptotrader/currency"
 	exchange "github.com/thrasher-/gocryptotrader/exchanges"
 	"github.com/thrasher-/gocryptotrader/exchanges/orderbook"
+	log "github.com/thrasher-/gocryptotrader/logger"
 )
 
 const (
 	coinbaseproWebsocketURL = "wss://ws-feed.pro.coinbase.com"
 )
 
-// WebsocketSubscriber subscribes to websocket channels with respect to enabled
-// currencies
-func (c *CoinbasePro) WebsocketSubscriber() error {
-	var currencies []string
-	for _, x := range c.EnabledPairs.Strings() {
-		currency := x[0:3] + "-" + x[3:]
-		currencies = append(currencies, currency)
-	}
-
-	var channels = []WsChannels{
-		{
-			Name:       "heartbeat",
-			ProductIDs: currencies,
-		},
-		{
-			Name:       "ticker",
-			ProductIDs: currencies,
-		},
-		{
-			Name:       "level2",
-			ProductIDs: currencies,
-		},
-	}
-
-	subscribe := WebsocketSubscribe{Type: "subscribe", Channels: channels}
-
-	data, err := common.JSONEncode(subscribe)
-	if err != nil {
-		return err
-	}
-
-	return c.WebsocketConn.WriteMessage(websocket.TextMessage, data)
-}
-
 // WsConnect initiates a websocket connection
 func (c *CoinbasePro) WsConnect() error {
 	if !c.Websocket.IsEnabled() || !c.IsEnabled() {
@@ -79,11 +46,7 @@ func (c *CoinbasePro) WsConnect() error {
 			err)
 	}
 
-	err = c.WebsocketSubscriber()
-	if err != nil {
-		return err
-	}
-
+	c.GenerateDefaultSubscriptions()
 	go c.WsHandleData()
 
 	return nil
@@ -95,7 +58,6 @@ func (c *CoinbasePro) WsReadData() (exchange.WebsocketResponse, error) {
 	if err != nil {
 		return exchange.WebsocketResponse{}, err
 	}
-
 	c.Websocket.TrafficAlert <- struct{}{}
 	return exchange.WebsocketResponse{Raw: resp}, nil
 }
@@ -105,11 +67,6 @@ func (c *CoinbasePro) WsHandleData() {
 	c.Websocket.Wg.Add(1)
 
 	defer func() {
-		err := c.WebsocketConn.Close()
-		if err != nil {
-			c.Websocket.DataHandler <- fmt.Errorf("coinbasepro_websocket.go - Unable to to close Websocket connection. Error: %s",
-				err)
-		}
 		c.Websocket.Wg.Done()
 	}()
 
@@ -117,14 +74,12 @@ func (c *CoinbasePro) WsHandleData() {
 		select {
 		case <-c.Websocket.ShutdownC:
 			return
-
 		default:
 			resp, err := c.WsReadData()
 			if err != nil {
 				c.Websocket.DataHandler <- err
 				return
 			}
-
 			type MsgType struct {
 				Type      string `json:"type"`
 				Sequence  int64  `json:"sequence"`
@@ -283,3 +238,66 @@ func (c *CoinbasePro) ProcessUpdate(update WebsocketL2Update) error {
 
 	return nil
 }
+
+// GenerateDefaultSubscriptions Adds default subscriptions to websocket to be handled by ManageSubscriptions()
+func (c *CoinbasePro) GenerateDefaultSubscriptions() {
+	var channels = []string{"heartbeat", "level2", "ticker"}
+	enabledCurrencies := c.GetEnabledCurrencies()
+	subscriptions := []exchange.WebsocketChannelSubscription{}
+	for i := range channels {
+		for j := range enabledCurrencies {
+			enabledCurrencies[j].Delimiter = "-"
+			subscriptions = append(subscriptions, exchange.WebsocketChannelSubscription{
+				Channel:  channels[i],
+				Currency: enabledCurrencies[j],
+			})
+		}
+	}
+	c.Websocket.SubscribeToChannels(subscriptions)
+}
+
+// Subscribe sends a websocket message to receive data from the channel
+func (c *CoinbasePro) Subscribe(channelToSubscribe exchange.WebsocketChannelSubscription) error {
+	subscribe := WebsocketSubscribe{
+		Type: "subscribe",
+		Channels: []WsChannels{
+			{
+				Name: channelToSubscribe.Channel,
+				ProductIDs: []string{
+					channelToSubscribe.Currency.String(),
+				},
+			},
+		},
+	}
+	return c.wsSend(subscribe)
+}
+
+// Unsubscribe sends a websocket message to stop receiving data from the channel
+func (c *CoinbasePro) Unsubscribe(channelToSubscribe exchange.WebsocketChannelSubscription) error {
+	subscribe := WebsocketSubscribe{
+		Type: "unsubscribe",
+		Channels: []WsChannels{
+			{
+				Name: channelToSubscribe.Channel,
+				ProductIDs: []string{
+					channelToSubscribe.Currency.String(),
+				},
+			},
+		},
+	}
+	return c.wsSend(subscribe)
+}
+
+// WsSend sends data to the websocket server
+func (c *CoinbasePro) wsSend(data interface{}) error {
+	c.wsRequestMtx.Lock()
+	defer c.wsRequestMtx.Unlock()
+	if c.Verbose {
+		log.Debugf("%v sending message to websocket %v", c.Name, data)
+	}
+	json, err := common.JSONEncode(data)
+	if err != nil {
+		return err
+	}
+	return c.WebsocketConn.WriteMessage(websocket.TextMessage, json)
+}

exchanges/coinbasepro/coinbasepro_wrapper.go

@@ -382,3 +382,17 @@ func (c *CoinbasePro) GetOrderHistory(getOrdersRequest *exchange.GetOrdersReques
 
 	return orders, nil
 }
+
+// SubscribeToWebsocketChannels appends to ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle subscribing
+func (c *CoinbasePro) SubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	c.Websocket.SubscribeToChannels(channels)
+	return nil
+}
+
+// UnsubscribeToWebsocketChannels removes from ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle unsubscribing
+func (c *CoinbasePro) UnsubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	c.Websocket.UnsubscribeToChannels(channels)
+	return nil
+}

exchanges/coinut/coinut.go

@@ -6,6 +6,7 @@ import (
 	"errors"
 	"fmt"
 	"net/http"
+	"sync"
 	"time"
 
 	"github.com/gorilla/websocket"
@@ -48,6 +49,7 @@ type COINUT struct {
 	exchange.Base
 	WebsocketConn *websocket.Conn
 	InstrumentMap map[string]int
+	wsRequestMtx  sync.Mutex
 }
 
 // SetDefaults sets current default values
@@ -77,7 +79,9 @@ func (c *COINUT) SetDefaults() {
 	c.WebsocketInit()
 	c.Websocket.Functionality = exchange.WebsocketTickerSupported |
 		exchange.WebsocketOrderbookSupported |
-		exchange.WebsocketTradeDataSupported
+		exchange.WebsocketTradeDataSupported |
+		exchange.WebsocketSubscribeSupported |
+		exchange.WebsocketUnsubscribeSupported
 }
 
 // Setup sets the current exchange configuration
@@ -118,8 +122,11 @@ func (c *COINUT) Setup(exch *config.ExchangeConfig) {
 			log.Fatal(err)
 		}
 		err = c.WebsocketSetup(c.WsConnect,
+			c.Subscribe,
+			c.Unsubscribe,
 			exch.Name,
 			exch.Websocket,
+			exch.Verbose,
 			coinutWebsocketURL,
 			exch.WebsocketURL)
 		if err != nil {

exchanges/coinut/coinut_websocket.go

@@ -2,7 +2,6 @@ package coinut
 
 import (
 	"errors"
-	"fmt"
 	"net/http"
 	"net/url"
 	"time"
@@ -12,9 +11,11 @@ import (
 	"github.com/thrasher-/gocryptotrader/currency"
 	exchange "github.com/thrasher-/gocryptotrader/exchanges"
 	"github.com/thrasher-/gocryptotrader/exchanges/orderbook"
+	log "github.com/thrasher-/gocryptotrader/logger"
 )
 
 const coinutWebsocketURL = "wss://wsapi.coinut.com"
+const coinutWebsocketRateLimit = 30 * time.Millisecond
 
 var nNonce map[int64]string
 var channels map[string]chan []byte
@@ -43,11 +44,6 @@ func (c *COINUT) WsHandleData() {
 	c.Websocket.Wg.Add(1)
 
 	defer func() {
-		err := c.WebsocketConn.Close()
-		if err != nil {
-			c.Websocket.DataHandler <- fmt.Errorf("coinut_websocket.go - Unable to to close Websocket connection. Error: %s",
-				err)
-		}
 		c.Websocket.Wg.Done()
 	}()
 
@@ -69,7 +65,6 @@ func (c *COINUT) WsHandleData() {
 				c.Websocket.DataHandler <- err
 				continue
 			}
-
 			switch incoming.Reply {
 			case "hb":
 				channels["hb"] <- resp.Raw
@@ -203,10 +198,7 @@ func (c *COINUT) WsConnect() error {
 		populatedList = true
 	}
 
-	err = c.WsSubscribe()
-	if err != nil {
-		return err
-	}
+	c.GenerateDefaultSubscriptions()
 
 	// define bi-directional communication
 	channels = make(map[string]chan []byte)
@@ -230,17 +222,11 @@ func (c *COINUT) GetNonce() int64 {
 
 // WsSetInstrumentList fetches instrument list and propagates a local cache
 func (c *COINUT) WsSetInstrumentList() error {
-	req, err := common.JSONEncode(wsRequest{
+	err := c.wsSend(wsRequest{
 		Request: "inst_list",
 		SecType: "SPOT",
 		Nonce:   c.GetNonce(),
 	})
-
-	if err != nil {
-		return err
-	}
-
-	err = c.WebsocketConn.WriteMessage(websocket.TextMessage, req)
 	if err != nil {
 		return err
 	}
@@ -270,48 +256,6 @@ func (c *COINUT) WsSetInstrumentList() error {
 	return nil
 }
 
-// WsSubscribe subscribes to websocket streams
-func (c *COINUT) WsSubscribe() error {
-	pairs := c.GetEnabledCurrencies()
-
-	for _, p := range pairs {
-		ticker := wsRequest{
-			Request:   "inst_tick",
-			InstID:    instrumentListByString[p.String()],
-			Subscribe: true,
-			Nonce:     c.GetNonce(),
-		}
-
-		tickjson, err := common.JSONEncode(ticker)
-		if err != nil {
-			return err
-		}
-
-		err = c.WebsocketConn.WriteMessage(websocket.TextMessage, tickjson)
-		if err != nil {
-			return err
-		}
-
-		ob := wsRequest{
-			Request:   "inst_order_book",
-			InstID:    instrumentListByString[p.String()],
-			Subscribe: true,
-			Nonce:     c.GetNonce(),
-		}
-
-		objson, err := common.JSONEncode(ob)
-		if err != nil {
-			return err
-		}
-
-		err = c.WebsocketConn.WriteMessage(websocket.TextMessage, objson)
-		if err != nil {
-			return err
-		}
-	}
-	return nil
-}
-
 // WsProcessOrderbookSnapshot processes the orderbook snapshot
 func (c *COINUT) WsProcessOrderbookSnapshot(ob *WsOrderbookSnapshot) error {
 	var bids []orderbook.Item
@@ -361,3 +305,58 @@ func (c *COINUT) WsProcessOrderbookUpdate(ob *WsOrderbookUpdate) error {
 		c.GetName(),
 		"SPOT")
 }
+
+// GenerateDefaultSubscriptions Adds default subscriptions to websocket to be handled by ManageSubscriptions()
+func (c *COINUT) GenerateDefaultSubscriptions() {
+	var channels = []string{"inst_tick", "inst_order_book"}
+	subscriptions := []exchange.WebsocketChannelSubscription{}
+	enabledCurrencies := c.GetEnabledCurrencies()
+	for i := range channels {
+		for j := range enabledCurrencies {
+			subscriptions = append(subscriptions, exchange.WebsocketChannelSubscription{
+				Channel:  channels[i],
+				Currency: enabledCurrencies[j],
+			})
+		}
+	}
+	c.Websocket.SubscribeToChannels(subscriptions)
+}
+
+// Subscribe sends a websocket message to receive data from the channel
+func (c *COINUT) Subscribe(channelToSubscribe exchange.WebsocketChannelSubscription) error {
+	subscribe := wsRequest{
+		Request:   channelToSubscribe.Channel,
+		InstID:    instrumentListByString[channelToSubscribe.Currency.String()],
+		Subscribe: true,
+		Nonce:     c.GetNonce(),
+	}
+	return c.wsSend(subscribe)
+}
+
+// Unsubscribe sends a websocket message to stop receiving data from the channel
+func (c *COINUT) Unsubscribe(channelToSubscribe exchange.WebsocketChannelSubscription) error {
+	subscribe := wsRequest{
+		Request:   channelToSubscribe.Channel,
+		InstID:    instrumentListByString[channelToSubscribe.Currency.String()],
+		Subscribe: false,
+		Nonce:     c.GetNonce(),
+	}
+	return c.wsSend(subscribe)
+}
+
+// WsSend sends data to the websocket server
+func (c *COINUT) wsSend(data interface{}) error {
+	c.wsRequestMtx.Lock()
+	defer c.wsRequestMtx.Unlock()
+
+	json, err := common.JSONEncode(data)
+	if err != nil {
+		return err
+	}
+	if c.Verbose {
+		log.Debugf("%v sending message to websocket %v", c.Name, string(json))
+	}
+	// Basic rate limiter
+	time.Sleep(coinutWebsocketRateLimit)
+	return c.WebsocketConn.WriteMessage(websocket.TextMessage, json)
+}

exchanges/coinut/coinut_wrapper.go

@@ -503,3 +503,17 @@ func (c *COINUT) GetOrderHistory(getOrdersRequest *exchange.GetOrdersRequest) ([
 
 	return orders, nil
 }
+
+// SubscribeToWebsocketChannels appends to ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle subscribing
+func (c *COINUT) SubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	c.Websocket.SubscribeToChannels(channels)
+	return nil
+}
+
+// UnsubscribeToWebsocketChannels removes from ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle unsubscribing
+func (c *COINUT) UnsubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	c.Websocket.UnsubscribeToChannels(channels)
+	return nil
+}

exchanges/exchange.go

@@ -323,6 +323,8 @@ type IBotExchange interface {
 	WithdrawFiatFunds(withdrawRequest *WithdrawRequest) (string, error)
 	WithdrawFiatFundsToInternationalBank(withdrawRequest *WithdrawRequest) (string, error)
 	GetWebsocket() (*Websocket, error)
+	SubscribeToWebsocketChannels(channels []WebsocketChannelSubscription) error
+	UnsubscribeToWebsocketChannels(channels []WebsocketChannelSubscription) error
 }
 
 // SupportsRESTTickerBatchUpdates returns whether or not the

exchanges/exchange_websocket.go

@@ -10,37 +10,7 @@ import (
 	"github.com/thrasher-/gocryptotrader/config"
 	"github.com/thrasher-/gocryptotrader/currency"
 	"github.com/thrasher-/gocryptotrader/exchanges/orderbook"
-)
-
-// Websocket functionality list and state consts
-const (
-	NoWebsocketSupport       uint32 = 0
-	WebsocketTickerSupported uint32 = 1 << (iota - 1)
-	WebsocketOrderbookSupported
-	WebsocketKlineSupported
-	WebsocketTradeDataSupported
-	WebsocketAccountSupported
-	WebsocketAllowsRequests
-
-	WebsocketTickerSupportedText    = "TICKER STREAMING SUPPORTED"
-	WebsocketOrderbookSupportedText = "ORDERBOOK STREAMING SUPPORTED"
-	WebsocketKlineSupportedText     = "KLINE STREAMING SUPPORTED"
-	WebsocketTradeDataSupportedText = "TRADE STREAMING SUPPORTED"
-	WebsocketAccountSupportedText   = "ACCOUNT STREAMING SUPPORTED"
-	WebsocketAllowsRequestsText     = "WEBSOCKET REQUESTS SUPPORTED"
-	NoWebsocketSupportText          = "WEBSOCKET NOT SUPPORTED"
-	UnknownWebsocketFunctionality   = "UNKNOWN FUNCTIONALITY BITMASK"
-
-	// WebsocketNotEnabled alerts of a disabled websocket
-	WebsocketNotEnabled = "exchange_websocket_not_enabled"
-	// WebsocketTrafficLimitTime defines a standard time for no traffic from the
-	// websocket connection
-	WebsocketTrafficLimitTime = 5 * time.Second
-	// WebsocketStateTimeout defines a const for when a websocket connection
-	// times out, will be handled by the routine management system
-	WebsocketStateTimeout = "TIMEOUT"
-
-	websocketRestablishConnection = 1 * time.Second
+	log "github.com/thrasher-/gocryptotrader/logger"
 )
 
 // WebsocketInit initialises the websocket struct
@@ -56,8 +26,11 @@ func (e *Base) WebsocketInit() {
 
 // WebsocketSetup sets main variables for websocket connection
 func (e *Base) WebsocketSetup(connector func() error,
+	subscriber func(channelToSubscribe WebsocketChannelSubscription) error,
+	unsubscriber func(channelToUnsubscribe WebsocketChannelSubscription) error,
 	exchangeName string,
-	wsEnabled bool,
+	wsEnabled,
+	verbose bool,
 	defaultURL,
 	runningURL string) error {
 
@@ -65,118 +38,26 @@ func (e *Base) WebsocketSetup(connector func() error,
 	e.Websocket.Connected = make(chan struct{}, 1)
 	e.Websocket.Disconnected = make(chan struct{}, 1)
 	e.Websocket.TrafficAlert = make(chan struct{}, 1)
+	e.Websocket.verbose = verbose
 
+	e.Websocket.SetChannelSubscriber(subscriber)
+	e.Websocket.SetChannelUnsubscriber(unsubscriber)
 	err := e.Websocket.SetWsStatusAndConnection(wsEnabled)
 	if err != nil {
 		return err
 	}
-
 	e.Websocket.SetDefaultURL(defaultURL)
 	e.Websocket.SetConnector(connector)
 	e.Websocket.SetWebsocketURL(runningURL)
 	e.Websocket.SetExchangeName(exchangeName)
 
 	e.Websocket.init = false
+	e.Websocket.noConnectionCheckLimit = 5
+	e.Websocket.reconnectionLimit = 10
 
 	return nil
 }
 
-// Websocket defines a return type for websocket connections via the interface
-// wrapper for routine processing in routines.go
-type Websocket struct {
-	proxyAddr    string
-	defaultURL   string
-	runningURL   string
-	exchangeName string
-	enabled      bool
-	init         bool
-	connected    bool
-	connector    func() error
-	m            sync.Mutex
-
-	// Connected denotes a channel switch for diversion of request flow
-	Connected chan struct{}
-
-	// Disconnected denotes a channel switch for diversion of request flow
-	Disconnected chan struct{}
-
-	// DataHandler pipes websocket data to an exchange websocket data handler
-	DataHandler chan interface{}
-
-	// ShutdownC is the main shutdown channel used within an exchange package
-	// called by its own defined Shutdown function
-	ShutdownC chan struct{}
-
-	// Orderbook is a local cache of orderbooks
-	Orderbook WebsocketOrderbookLocal
-
-	// Wg defines a wait group for websocket routines for cleanly shutting down
-	// routines
-	Wg sync.WaitGroup
-
-	// TrafficAlert monitors if there is a halt in traffic throughput
-	TrafficAlert chan struct{}
-
-	// Functionality defines websocket stream capabilities
-	Functionality uint32
-}
-
-// trafficMonitor monitors traffic and switches connection modes for websocket
-func (w *Websocket) trafficMonitor(wg *sync.WaitGroup) {
-	w.Wg.Add(1)
-	wg.Done() // Makes sure we are unlocking after we add to waitgroup
-
-	defer func() {
-		if w.connected {
-			w.Disconnected <- struct{}{}
-		}
-		w.Wg.Done()
-	}()
-
-	// Define an initial traffic timer which will be a delay then fall over to
-	// WebsocketTrafficLimitTime after first response
-	trafficTimer := time.NewTimer(5 * time.Second)
-
-	for {
-		select {
-		case <-w.ShutdownC: // Returns on shutdown channel close
-			return
-		case <-w.TrafficAlert: // Resets timer on traffic
-			if !w.connected {
-				w.Connected <- struct{}{}
-				w.connected = true
-			}
-
-			trafficTimer.Reset(WebsocketTrafficLimitTime)
-
-		case <-trafficTimer.C: // Falls through when timer runs out
-			newtimer := time.NewTimer(10 * time.Second) // New secondary timer set
-			if w.connected {
-				// If connected divert traffic to rest
-				w.Disconnected <- struct{}{}
-				w.connected = false
-			}
-
-			select {
-			case <-w.ShutdownC: // Returns on shutdown channel close
-				return
-
-			case <-newtimer.C: // If secondary timer runs state timeout is sent to the data handler
-				w.DataHandler <- WebsocketStateTimeout
-				return
-
-			case <-w.TrafficAlert: // If in this time response traffic comes through
-				trafficTimer.Reset(WebsocketTrafficLimitTime)
-				if !w.connected {
-					// If not connected divert traffic from REST to websocket
-					w.Connected <- struct{}{}
-					w.connected = true
-				}
-			}
-		}
-	}
-}
-
 // Connect intiates a websocket connection by using a package defined connection
 // function
 func (w *Websocket) Connect() error {
@@ -188,13 +69,15 @@ func (w *Websocket) Connect() error {
 	}
 
 	if w.connected {
+		w.connecting = false
 		return errors.New("exchange_websocket.go error - already connected, cannot connect again")
 	}
 
+	w.connecting = true
 	w.ShutdownC = make(chan struct{}, 1)
-
 	err := w.connector()
 	if err != nil {
+		w.connecting = false
 		return fmt.Errorf("exchange_websocket.go connection error %s",
 			err)
 	}
@@ -202,41 +85,130 @@ func (w *Websocket) Connect() error {
 	if !w.connected {
 		w.Connected <- struct{}{}
 		w.connected = true
+		w.connecting = false
 	}
 
 	var anotherWG sync.WaitGroup
 	anotherWG.Add(1)
 	go w.trafficMonitor(&anotherWG)
 	anotherWG.Wait()
+	if !w.connectionMonitorRunning {
+		go w.wsConnectionMonitor()
+	}
+	go w.manageSubscriptions()
 
 	return nil
 }
 
+// WsConnectionMonitor ensures that the WS keeps connecting
+func (w *Websocket) wsConnectionMonitor() {
+	w.m.Lock()
+	w.connectionMonitorRunning = true
+	w.m.Unlock()
+	defer func() {
+		w.connectionMonitorRunning = false
+	}()
+
+	for {
+		time.Sleep(connectionMonitorDelay)
+		w.m.Lock()
+		if !w.enabled {
+			w.m.Unlock()
+			w.DataHandler <- fmt.Errorf("%v WsConnectionMonitor: websocket disabled, shutting down", w.exchangeName)
+			err := w.Shutdown()
+			if err != nil {
+				log.Error(err)
+			}
+			if w.verbose {
+				log.Debugf("%v WsConnectionMonitor exiting", w.exchangeName)
+			}
+			return
+		}
+		w.m.Unlock()
+		err := w.checkConnection()
+		if err != nil {
+			log.Error(err)
+		}
+	}
+}
+
+// checkConnection ensures the connection is maintained
+// Will reconnect on disconnect
+func (w *Websocket) checkConnection() error {
+	if w.verbose {
+		log.Debugf("%v checking connection", w.exchangeName)
+	}
+	switch {
+	case !w.IsConnected() && !w.IsConnecting():
+		w.m.Lock()
+		defer w.m.Unlock()
+		if w.verbose {
+			log.Debugf("%v no connection. Attempt %v/%v", w.exchangeName, w.noConnectionChecks, w.noConnectionCheckLimit)
+		}
+		if w.noConnectionChecks >= w.noConnectionCheckLimit {
+			if w.verbose {
+				log.Debugf("%v resetting connection", w.exchangeName)
+			}
+			w.connecting = true
+			go w.WebsocketReset()
+			w.noConnectionChecks = 0
+		}
+		w.noConnectionChecks++
+	case w.IsConnecting():
+		if w.reconnectionChecks >= w.reconnectionLimit {
+			return fmt.Errorf("%v websocket failed to reconnect after %v seconds",
+				w.exchangeName,
+				w.reconnectionLimit*int(connectionMonitorDelay.Seconds()))
+		}
+		if w.verbose {
+			log.Debugf("%v Busy reconnecting", w.exchangeName)
+		}
+		w.reconnectionChecks++
+	default:
+		w.noConnectionChecks = 0
+		w.reconnectionChecks = 0
+	}
+	return nil
+}
+
 // IsConnected exposes websocket connection status
 func (w *Websocket) IsConnected() bool {
+	w.m.Lock()
+	defer w.m.Unlock()
 	return w.connected
 }
 
+// IsConnecting checks whether websocket is busy connecting
+func (w *Websocket) IsConnecting() bool {
+	w.m.Lock()
+	defer w.m.Unlock()
+	return w.connecting
+}
+
 // Shutdown attempts to shut down a websocket connection and associated routines
 // by using a package defined shutdown function
 func (w *Websocket) Shutdown() error {
 	w.m.Lock()
-
 	defer func() {
 		w.Orderbook.FlushCache()
 		w.m.Unlock()
 	}()
-
-	if !w.connected {
-		return errors.New("exchange_websocket.go error - System not connected to shut down")
+	if !w.connected && w.ShutdownC == nil {
+		return fmt.Errorf("%v cannot shutdown a disconnected websocket", w.exchangeName)
 	}
 
-	timer := time.NewTimer(5 * time.Second)
+	if w.verbose {
+		log.Debugf("%v shutting down websocket channels", w.exchangeName)
+	}
+	timer := time.NewTimer(15 * time.Second)
 	c := make(chan struct{}, 1)
 
 	go func(c chan struct{}) {
 		close(w.ShutdownC)
 		w.Wg.Wait()
+		if w.verbose {
+			log.Debugf("%v completed websocket channel shutdown", w.exchangeName)
+		}
 		c <- struct{}{}
 	}(c)
 
@@ -245,11 +217,105 @@ func (w *Websocket) Shutdown() error {
 		w.connected = false
 		return nil
 	case <-timer.C:
-		return fmt.Errorf("%s - Websocket routines failed to shutdown",
+		return fmt.Errorf("%s websocket routines failed to shutdown after 15 seconds",
 			w.GetName())
 	}
 }
 
+// WebsocketReset sends the shutdown command, waits for channel/func closure and then reconnects
+func (w *Websocket) WebsocketReset() error {
+	err := w.Shutdown()
+	if err != nil {
+		// does not return here to allow connection to be made if already shut down
+		log.Errorf("%v shutdown error: %v", w.exchangeName, err)
+	}
+	log.Infof("%v reconnecting to websocket", w.exchangeName)
+	w.m.Lock()
+	w.init = true
+	w.m.Unlock()
+	err = w.Connect()
+	if err != nil {
+		log.Errorf("%v connection error: %v", w.exchangeName, err)
+	}
+	return err
+}
+
+// trafficMonitor monitors traffic and switches connection modes for websocket
+func (w *Websocket) trafficMonitor(wg *sync.WaitGroup) {
+	w.Wg.Add(1)
+	wg.Done() // Makes sure we are unlocking after we add to waitgroup
+	defer func() {
+		if w.connected {
+			w.Disconnected <- struct{}{}
+		}
+		w.Wg.Done()
+	}()
+
+	// Define an initial traffic timer which will be a delay then fall over to
+	// WebsocketTrafficLimitTime after first response
+	trafficTimer := time.NewTimer(5 * time.Second)
+	for {
+		select {
+		case <-w.ShutdownC: // Returns on shutdown channel close
+			if w.verbose {
+				log.Debugf("%v trafficMonitor shutdown message received", w.exchangeName)
+			}
+			return
+		case <-w.TrafficAlert: // Resets timer on traffic
+			w.m.Lock()
+			if !w.connected {
+				w.Connected <- struct{}{}
+				w.connected = true
+			}
+			w.m.Unlock()
+			if w.verbose {
+				log.Debugf("%v received a traffic alert", w.exchangeName)
+			}
+			trafficTimer.Reset(WebsocketTrafficLimitTime)
+		case <-trafficTimer.C: // Falls through when timer runs out
+			newtimer := time.NewTimer(10 * time.Second) // New secondary timer set
+			if w.verbose {
+				log.Debugf("%v has not received a traffic alert in 5 seconds.", w.exchangeName)
+			}
+			w.m.Lock()
+			if w.connected {
+				// If connected divert traffic to rest
+				w.Disconnected <- struct{}{}
+				w.connected = false
+			}
+			w.m.Unlock()
+
+			select {
+			case <-w.ShutdownC: // Returns on shutdown channel close
+				w.m.Lock()
+				w.connected = false
+				w.m.Unlock()
+				return
+
+			case <-newtimer.C: // If secondary timer runs state timeout is sent to the data handler
+				if w.verbose {
+					log.Debugf("%v has not received a traffic alert in 15 seconds, exiting", w.exchangeName)
+				}
+				w.DataHandler <- fmt.Errorf("trafficMonitor %v", WebsocketStateTimeout)
+				return
+
+			case <-w.TrafficAlert: // If in this time response traffic comes through
+				trafficTimer.Reset(WebsocketTrafficLimitTime)
+				w.m.Lock()
+				if !w.connected {
+					// If not connected dive rt traffic from REST to websocket
+					w.Connected <- struct{}{}
+					if w.verbose {
+						log.Debugf("%v has received a traffic alert. Setting status to connected", w.exchangeName)
+					}
+					w.connected = true
+				}
+				w.m.Unlock()
+			}
+		}
+	}
+}
+
 // SetWebsocketURL sets websocket URL
 func (w *Websocket) SetWebsocketURL(websocketURL string) {
 	if websocketURL == "" || websocketURL == config.WebsocketURLNonDefaultMessage {
@@ -267,29 +333,36 @@ func (w *Websocket) GetWebsocketURL() string {
 // SetWsStatusAndConnection sets if websocket is enabled
 // it will also connect/disconnect the websocket connection
 func (w *Websocket) SetWsStatusAndConnection(enabled bool) error {
+	w.m.Lock()
 	if w.enabled == enabled {
 		if w.init {
+			w.m.Unlock()
 			return nil
 		}
+		w.m.Unlock()
 		return fmt.Errorf("exchange_websocket.go error - already set as %t",
 			enabled)
 	}
-
 	w.enabled = enabled
 
 	if !w.init {
 		if enabled {
 			if w.connected {
+				w.m.Unlock()
 				return nil
 			}
+			w.m.Unlock()
 			return w.Connect()
 		}
 
 		if !w.connected {
+			w.m.Unlock()
 			return nil
 		}
+		w.m.Unlock()
 		return w.Shutdown()
 	}
+	w.m.Unlock()
 	return nil
 }
 
@@ -305,14 +378,12 @@ func (w *Websocket) SetProxyAddress(proxyAddr string) error {
 	}
 
 	w.proxyAddr = proxyAddr
-
 	if !w.init && w.enabled {
 		if w.connected {
 			err := w.Shutdown()
 			if err != nil {
 				return err
 			}
-			return w.Connect()
 		}
 		return w.Connect()
 	}
@@ -349,14 +420,6 @@ func (w *Websocket) GetName() string {
 	return w.exchangeName
 }
 
-// WebsocketOrderbookLocal defines a local cache of orderbooks for amending,
-// appending and deleting changes and updates the main store in orderbook.go
-type WebsocketOrderbookLocal struct {
-	ob          []*orderbook.Base
-	lastUpdated time.Time
-	m           sync.Mutex
-}
-
 // Update updates a local cache using bid targets and ask targets then updates
 // main cache in orderbook.go
 // Volume == 0; deletion at price target
@@ -567,70 +630,6 @@ func (w *WebsocketOrderbookLocal) FlushCache() {
 	w.m.Unlock()
 }
 
-// WebsocketResponse defines generalised data from the websocket connection
-type WebsocketResponse struct {
-	Type int
-	Raw  []byte
-}
-
-// WebsocketOrderbookUpdate defines a websocket event in which the orderbook
-// has been updated in the orderbook package
-type WebsocketOrderbookUpdate struct {
-	Pair     currency.Pair
-	Asset    string
-	Exchange string
-}
-
-// TradeData defines trade data
-type TradeData struct {
-	Timestamp    time.Time
-	CurrencyPair currency.Pair
-	AssetType    string
-	Exchange     string
-	EventType    string
-	EventTime    int64
-	Price        float64
-	Amount       float64
-	Side         string
-}
-
-// TickerData defines ticker feed
-type TickerData struct {
-	Timestamp  time.Time
-	Pair       currency.Pair
-	AssetType  string
-	Exchange   string
-	ClosePrice float64
-	Quantity   float64
-	OpenPrice  float64
-	HighPrice  float64
-	LowPrice   float64
-}
-
-// KlineData defines kline feed
-type KlineData struct {
-	Timestamp  time.Time
-	Pair       currency.Pair
-	AssetType  string
-	Exchange   string
-	StartTime  time.Time
-	CloseTime  time.Time
-	Interval   string
-	OpenPrice  float64
-	ClosePrice float64
-	HighPrice  float64
-	LowPrice   float64
-	Volume     float64
-}
-
-// WebsocketPositionUpdated reflects a change in orders/contracts on an exchange
-type WebsocketPositionUpdated struct {
-	Timestamp time.Time
-	Pair      currency.Pair
-	AssetType string
-	Exchange  string
-}
-
 // GetFunctionality returns a functionality bitmask for the websocket
 // connection
 func (w *Websocket) GetFunctionality() uint32 {
@@ -668,6 +667,12 @@ func (w *Websocket) FormatFunctionality() string {
 			case WebsocketAllowsRequests:
 				functionality = append(functionality, WebsocketAllowsRequestsText)
 
+			case WebsocketSubscribeSupported:
+				functionality = append(functionality, WebsocketSubscribeSupportedText)
+
+			case WebsocketUnsubscribeSupported:
+				functionality = append(functionality, WebsocketUnsubscribeSupportedText)
+
 			default:
 				functionality = append(functionality,
 					fmt.Sprintf("%s[1<<%v]", UnknownWebsocketFunctionality, i))
@@ -681,3 +686,172 @@ func (w *Websocket) FormatFunctionality() string {
 
 	return NoWebsocketSupportText
 }
+
+// SetChannelSubscriber sets the function to use the base subscribe func
+func (w *Websocket) SetChannelSubscriber(subscriber func(channelToSubscribe WebsocketChannelSubscription) error) {
+	w.channelSubscriber = subscriber
+}
+
+// SetChannelUnsubscriber sets the function to use the base unsubscribe func
+func (w *Websocket) SetChannelUnsubscriber(unsubscriber func(channelToUnsubscribe WebsocketChannelSubscription) error) {
+	w.channelUnsubscriber = unsubscriber
+}
+
+// ManageSubscriptions ensures the subscriptions specified continue to be subscribed to
+func (w *Websocket) manageSubscriptions() error {
+	if !w.SupportsFunctionality(WebsocketSubscribeSupported) && !w.SupportsFunctionality(WebsocketUnsubscribeSupported) {
+		return fmt.Errorf("%v does not support channel subscriptions, exiting ManageSubscriptions()", w.exchangeName)
+	}
+	w.Wg.Add(1)
+	defer func() {
+		if w.verbose {
+			log.Debugf("%v ManageSubscriptions exiting", w.exchangeName)
+		}
+		w.Wg.Done()
+	}()
+	for {
+		select {
+		case <-w.ShutdownC:
+			w.subscribedChannels = []WebsocketChannelSubscription{}
+			if w.verbose {
+				log.Debugf("%v shutdown manageSubscriptions", w.exchangeName)
+			}
+			return nil
+		default:
+			time.Sleep(manageSubscriptionsDelay)
+			if w.verbose {
+				log.Debugf("%v checking subscriptions", w.exchangeName)
+			}
+			// Subscribe to channels Pending a subscription
+			if w.SupportsFunctionality(WebsocketSubscribeSupported) {
+				err := w.subscribeToChannels()
+				if err != nil {
+					w.DataHandler <- err
+				}
+			}
+			if w.SupportsFunctionality(WebsocketUnsubscribeSupported) {
+				err := w.unsubscribeToChannels()
+				if err != nil {
+					w.DataHandler <- err
+				}
+			}
+		}
+	}
+}
+
+// subscribeToChannels compares channelsToSubscribe to subscribedChannels
+// and subscribes to any channels not present in subscribedChannels
+func (w *Websocket) subscribeToChannels() error {
+	w.subscriptionLock.Lock()
+	defer w.subscriptionLock.Unlock()
+	for i := 0; i < len(w.channelsToSubscribe); i++ {
+		channelIsSubscribed := false
+		for j := 0; j < len(w.subscribedChannels); j++ {
+			if w.subscribedChannels[j].Equal(&w.channelsToSubscribe[i]) {
+				channelIsSubscribed = true
+				break
+			}
+		}
+		if !channelIsSubscribed {
+			if w.verbose {
+				log.Debugf("%v Subscribing to %v %v", w.exchangeName, w.channelsToSubscribe[i].Channel, w.channelsToSubscribe[i].Currency.String())
+			}
+			err := w.channelSubscriber(w.channelsToSubscribe[i])
+			if err != nil {
+				return err
+			}
+			w.subscribedChannels = append(w.subscribedChannels, w.channelsToSubscribe[i])
+		}
+	}
+	return nil
+}
+
+// unsubscribeToChannels compares subscribedChannels to channelsToSubscribe
+// and unsubscribes to any channels not present in  channelsToSubscribe
+func (w *Websocket) unsubscribeToChannels() error {
+	w.subscriptionLock.Lock()
+	defer w.subscriptionLock.Unlock()
+	for i := 0; i < len(w.subscribedChannels); i++ {
+		subscriptionFound := false
+		for j := 0; j < len(w.channelsToSubscribe); j++ {
+			if w.channelsToSubscribe[j].Equal(&w.subscribedChannels[i]) {
+				subscriptionFound = true
+				break
+			}
+		}
+		if !subscriptionFound {
+			err := w.channelUnsubscriber(w.subscribedChannels[i])
+			if err != nil {
+				return err
+			}
+		}
+	}
+	// Now that the slices should match, assign rather than looping and appending the differences
+	w.subscribedChannels = append(w.channelsToSubscribe[:0:0], w.channelsToSubscribe...) //nolint:gocritic
+
+	return nil
+}
+
+// removeChannelToSubscribe removes an entry from w.channelsToSubscribe
+// so an unsubscribe event can be triggered
+func (w *Websocket) removeChannelToSubscribe(subscribedChannel WebsocketChannelSubscription) {
+	w.subscriptionLock.Lock()
+	defer w.subscriptionLock.Unlock()
+	channelLength := len(w.channelsToSubscribe)
+	i := 0
+	for j := 0; j < len(w.channelsToSubscribe); j++ {
+		if !w.channelsToSubscribe[j].Equal(&subscribedChannel) {
+			w.channelsToSubscribe[i] = w.channelsToSubscribe[j]
+			i++
+		}
+	}
+	w.channelsToSubscribe = w.channelsToSubscribe[:i]
+	if channelLength == len(w.channelsToSubscribe) {
+		w.DataHandler <- fmt.Errorf("%v removeChannelToSubscribe() Channel %v Currency %v could not be removed because it was not found",
+			w.exchangeName,
+			subscribedChannel.Channel,
+			subscribedChannel.Currency)
+	}
+}
+
+// ResubscribeToChannel calls unsubscribe func and
+// removes it from subscribedChannels to trigger a subscribe event
+func (w *Websocket) ResubscribeToChannel(subscribedChannel WebsocketChannelSubscription) {
+	w.subscriptionLock.Lock()
+	defer w.subscriptionLock.Unlock()
+	err := w.channelUnsubscriber(subscribedChannel)
+	if err != nil {
+		w.DataHandler <- err
+	}
+	// Remove the channel from the list of subscribed channels
+	// ManageSubscriptions will automatically resubscribe
+	i := 0
+	for j := 0; j < len(w.subscribedChannels); j++ {
+		if !w.subscribedChannels[j].Equal(&subscribedChannel) {
+			w.subscribedChannels[i] = w.subscribedChannels[j]
+			i++
+		}
+	}
+	w.subscribedChannels = w.subscribedChannels[:i]
+}
+
+// SubscribeToChannels appends supplied channels to channelsToSubscribe
+func (w *Websocket) SubscribeToChannels(channels []WebsocketChannelSubscription) {
+	for i := range channels {
+		w.channelsToSubscribe = append(w.channelsToSubscribe, channels[i])
+	}
+	w.noConnectionChecks = 0
+}
+
+// UnsubscribeToChannels removes supplied channels from channelsToSubscribe
+func (w *Websocket) UnsubscribeToChannels(channels []WebsocketChannelSubscription) {
+	for i := range channels {
+		w.removeChannelToSubscribe(channels[i])
+	}
+}
+
+// Equal two WebsocketChannelSubscription to determine equality
+func (w *WebsocketChannelSubscription) Equal(subscribedChannel *WebsocketChannelSubscription) bool {
+	return strings.EqualFold(w.Channel, subscribedChannel.Channel) &&
+		strings.EqualFold(w.Currency.String(), subscribedChannel.Currency.String())
+}

exchanges/exchange_websocket_test.go

@@ -1,6 +1,8 @@
 package exchange
 
 import (
+	"fmt"
+	"strings"
 	"testing"
 	"time"
 
@@ -28,8 +30,11 @@ func TestWebsocket(t *testing.T) {
 	}
 
 	wsTest.WebsocketSetup(func() error { return nil },
+		func(test WebsocketChannelSubscription) error { return nil },
+		func(test WebsocketChannelSubscription) error { return nil },
 		"testName",
 		true,
+		false,
 		"testDefaultURL",
 		"testRunningURL")
 
@@ -71,13 +76,12 @@ func TestWebsocket(t *testing.T) {
 			}
 		}
 	}()
-
 	// -- Not connected shutdown
 	err := wsTest.Websocket.Shutdown()
 	if err == nil {
 		t.Fatal("test failed - should not be connected to able to shut down")
 	}
-
+	wsTest.Websocket.Wg.Wait()
 	// -- Normal connect
 	err = wsTest.Websocket.Connect()
 	if err != nil {
@@ -255,7 +259,6 @@ func TestUpdate(t *testing.T) {
 		{Price: 1337, Amount: 100}, // Append
 		{Price: 1336, Amount: 0},   // Ghost delete
 	}
-
 	err := wsTest.Websocket.Orderbook.Update(bidTargets,
 		askTargets,
 		LTCUSDPAIR,
@@ -328,3 +331,238 @@ func TestFunctionality(t *testing.T) {
 		t.Fatal("Test Failed - SupportsFunctionality error should be true")
 	}
 }
+
+// placeholderSubscriber basic function to test subscriptions
+func placeholderSubscriber(channelToSubscribe WebsocketChannelSubscription) error {
+	return nil
+}
+
+// TestSubscribe logic test
+func TestSubscribe(t *testing.T) {
+	w := Websocket{
+		channelsToSubscribe: []WebsocketChannelSubscription{
+			{
+				Channel: "hello",
+			},
+		},
+		subscribedChannels: []WebsocketChannelSubscription{},
+	}
+	w.SetChannelSubscriber(placeholderSubscriber)
+	w.subscribeToChannels()
+	if len(w.subscribedChannels) != 1 {
+		t.Errorf("Subscription did not occur")
+	}
+}
+
+// TestUnsubscribe logic test
+func TestUnsubscribe(t *testing.T) {
+	w := Websocket{
+		channelsToSubscribe: []WebsocketChannelSubscription{},
+		subscribedChannels: []WebsocketChannelSubscription{
+			{
+				Channel: "hello",
+			},
+		},
+	}
+	w.SetChannelUnsubscriber(placeholderSubscriber)
+	w.unsubscribeToChannels()
+	if len(w.subscribedChannels) != 0 {
+		t.Errorf("Unsubscription did not occur")
+	}
+}
+
+// TestSubscriptionWithExistingEntry logic test
+func TestSubscriptionWithExistingEntry(t *testing.T) {
+	w := Websocket{
+		channelsToSubscribe: []WebsocketChannelSubscription{
+			{
+				Channel: "hello",
+			},
+		},
+		subscribedChannels: []WebsocketChannelSubscription{
+			{
+				Channel: "hello",
+			},
+		},
+	}
+	w.SetChannelSubscriber(placeholderSubscriber)
+	w.subscribeToChannels()
+	if len(w.subscribedChannels) != 1 {
+		t.Errorf("Subscription should not have occured")
+	}
+}
+
+// TestUnsubscriptionWithExistingEntry logic test
+func TestUnsubscriptionWithExistingEntry(t *testing.T) {
+	w := Websocket{
+		channelsToSubscribe: []WebsocketChannelSubscription{
+			{
+				Channel: "hello",
+			},
+		},
+		subscribedChannels: []WebsocketChannelSubscription{
+			{
+				Channel: "hello",
+			},
+		},
+	}
+	w.SetChannelUnsubscriber(placeholderSubscriber)
+	w.unsubscribeToChannels()
+	if len(w.subscribedChannels) != 1 {
+		t.Errorf("Unsubscription should not have occured")
+	}
+}
+
+// TestManageSubscriptionsWithoutFunctionality logic test
+func TestManageSubscriptionsWithoutFunctionality(t *testing.T) {
+	w := Websocket{
+		ShutdownC: make(chan struct{}, 1),
+	}
+	err := w.manageSubscriptions()
+	if err == nil {
+		t.Error("Requires functionality to work")
+	}
+}
+
+// TestManageSubscriptionsStartStop logic test
+func TestManageSubscriptionsStartStop(t *testing.T) {
+	w := Websocket{
+		ShutdownC:     make(chan struct{}, 1),
+		Functionality: WebsocketSubscribeSupported | WebsocketUnsubscribeSupported,
+	}
+	go w.manageSubscriptions()
+	time.Sleep(time.Second)
+	close(w.ShutdownC)
+}
+
+// TestWsConnectionMonitorNoConnection logic test
+func TestWsConnectionMonitorNoConnection(t *testing.T) {
+	w := Websocket{}
+	w.DataHandler = make(chan interface{}, 1)
+	w.ShutdownC = make(chan struct{}, 1)
+	w.exchangeName = "hello"
+	go w.wsConnectionMonitor()
+	err := <-w.DataHandler
+	if !strings.EqualFold(err.(error).Error(),
+		fmt.Sprintf("%v WsConnectionMonitor: websocket disabled, shutting down", w.exchangeName)) {
+		t.Errorf("expecting error 'WsConnectionMonitor: websocket disabled, shutting down', received '%v'", err)
+	}
+}
+
+// TestWsNoConnectionTolerance logic test
+func TestWsNoConnectionTolerance(t *testing.T) {
+	w := Websocket{}
+	w.DataHandler = make(chan interface{}, 1)
+	w.ShutdownC = make(chan struct{}, 1)
+	w.enabled = true
+	w.noConnectionCheckLimit = 500
+	w.checkConnection()
+	if w.noConnectionChecks == 0 {
+		t.Errorf("Expected noConnectionTolerance to increment, received '%v'", w.noConnectionChecks)
+	}
+}
+
+// TestConnecting logic test
+func TestConnecting(t *testing.T) {
+	w := Websocket{}
+	w.DataHandler = make(chan interface{}, 1)
+	w.ShutdownC = make(chan struct{}, 1)
+	w.enabled = true
+	w.connecting = true
+	w.reconnectionLimit = 500
+	w.checkConnection()
+	if w.reconnectionChecks != 1 {
+		t.Errorf("Expected reconnectionLimit to increment, received '%v'", w.reconnectionChecks)
+	}
+}
+
+// TestReconnectionLimit logic test
+func TestReconnectionLimit(t *testing.T) {
+	w := Websocket{}
+	w.DataHandler = make(chan interface{}, 1)
+	w.ShutdownC = make(chan struct{}, 1)
+	w.enabled = true
+	w.connecting = true
+	w.reconnectionChecks = 99
+	w.reconnectionLimit = 1
+	err := w.checkConnection()
+	if err == nil {
+		t.Error("Expected error")
+	}
+}
+
+// TestRemoveChannelToSubscribe logic test
+func TestRemoveChannelToSubscribe(t *testing.T) {
+	subscription := WebsocketChannelSubscription{
+		Channel: "hello",
+	}
+	w := Websocket{
+		channelsToSubscribe: []WebsocketChannelSubscription{
+			subscription,
+		},
+	}
+	w.SetChannelUnsubscriber(placeholderSubscriber)
+	w.removeChannelToSubscribe(subscription)
+	if len(w.subscribedChannels) != 0 {
+		t.Errorf("Unsubscription did not occur")
+	}
+}
+
+// TestRemoveChannelToSubscribeWithNoSubscription logic test
+func TestRemoveChannelToSubscribeWithNoSubscription(t *testing.T) {
+	subscription := WebsocketChannelSubscription{
+		Channel: "hello",
+	}
+	w := Websocket{
+		channelsToSubscribe: []WebsocketChannelSubscription{},
+	}
+	w.DataHandler = make(chan interface{}, 1)
+	w.SetChannelUnsubscriber(placeholderSubscriber)
+	go w.removeChannelToSubscribe(subscription)
+	err := <-w.DataHandler
+	if !strings.Contains(err.(error).Error(), "could not be removed because it was not found") {
+		t.Error("Expected not found error")
+	}
+}
+
+// TestResubscribeToChannel logic test
+func TestResubscribeToChannel(t *testing.T) {
+	subscription := WebsocketChannelSubscription{
+		Channel: "hello",
+	}
+	w := Websocket{
+		channelsToSubscribe: []WebsocketChannelSubscription{},
+	}
+	w.DataHandler = make(chan interface{}, 1)
+	w.SetChannelUnsubscriber(placeholderSubscriber)
+	w.SetChannelSubscriber(placeholderSubscriber)
+	w.ResubscribeToChannel(subscription)
+}
+
+// TestSliceCopyDoesntImpactBoth logic test
+func TestSliceCopyDoesntImpactBoth(t *testing.T) {
+	w := Websocket{
+		channelsToSubscribe: []WebsocketChannelSubscription{
+			{
+				Channel: "hello1",
+			},
+			{
+				Channel: "hello2",
+			},
+		},
+		subscribedChannels: []WebsocketChannelSubscription{
+			{
+				Channel: "hello3",
+			},
+		},
+	}
+	w.SetChannelUnsubscriber(placeholderSubscriber)
+	w.unsubscribeToChannels()
+	if len(w.subscribedChannels) != 2 {
+		t.Errorf("Unsubscription did not occur")
+	}
+	w.subscribedChannels[0].Channel = "test"
+	if strings.EqualFold(w.subscribedChannels[0].Channel, w.channelsToSubscribe[0].Channel) {
+		t.Errorf("Slice has not been copies appropriately")
+	}
+}

exchanges/exchange_websocket_types.go

@@ -0,0 +1,172 @@
+package exchange
+
+import (
+	"sync"
+	"time"
+
+	"github.com/thrasher-/gocryptotrader/currency"
+	"github.com/thrasher-/gocryptotrader/exchanges/orderbook"
+)
+
+// Websocket functionality list and state consts
+const (
+	NoWebsocketSupport       uint32 = 0
+	WebsocketTickerSupported uint32 = 1 << (iota - 1)
+	WebsocketOrderbookSupported
+	WebsocketKlineSupported
+	WebsocketTradeDataSupported
+	WebsocketAccountSupported
+	WebsocketAllowsRequests
+	WebsocketSubscribeSupported
+	WebsocketUnsubscribeSupported
+
+	WebsocketTickerSupportedText      = "TICKER STREAMING SUPPORTED"
+	WebsocketOrderbookSupportedText   = "ORDERBOOK STREAMING SUPPORTED"
+	WebsocketKlineSupportedText       = "KLINE STREAMING SUPPORTED"
+	WebsocketTradeDataSupportedText   = "TRADE STREAMING SUPPORTED"
+	WebsocketAccountSupportedText     = "ACCOUNT STREAMING SUPPORTED"
+	WebsocketAllowsRequestsText       = "WEBSOCKET REQUESTS SUPPORTED"
+	NoWebsocketSupportText            = "WEBSOCKET NOT SUPPORTED"
+	UnknownWebsocketFunctionality     = "UNKNOWN FUNCTIONALITY BITMASK"
+	WebsocketSubscribeSupportedText   = "WEBSOCKET SUBSCRIBE SUPPORTED"
+	WebsocketUnsubscribeSupportedText = "WEBSOCKET UNSUBSCRIBE SUPPORTED"
+
+	// WebsocketNotEnabled alerts of a disabled websocket
+	WebsocketNotEnabled = "exchange_websocket_not_enabled"
+	// WebsocketTrafficLimitTime defines a standard time for no traffic from the
+	// websocket connection
+	WebsocketTrafficLimitTime     = 5 * time.Second
+	websocketRestablishConnection = time.Second
+	manageSubscriptionsDelay      = 5 * time.Second
+	// connection monitor time delays and limits
+	connectionMonitorDelay = 2 * time.Second
+	// WebsocketStateTimeout defines a const for when a websocket connection
+	// times out, will be handled by the routine management system
+	WebsocketStateTimeout = "TIMEOUT"
+)
+
+// Websocket defines a return type for websocket connections via the interface
+// wrapper for routine processing in routines.go
+type Websocket struct {
+	proxyAddr                string
+	defaultURL               string
+	runningURL               string
+	exchangeName             string
+	enabled                  bool
+	init                     bool
+	connected                bool
+	connecting               bool
+	verbose                  bool
+	connector                func() error
+	m                        sync.Mutex
+	subscriptionLock         sync.Mutex
+	connectionMonitorRunning bool
+	reconnectionLimit        int
+	noConnectionChecks       int
+	reconnectionChecks       int
+	noConnectionCheckLimit   int
+	// Subscriptions stuff
+	subscribedChannels  []WebsocketChannelSubscription
+	channelsToSubscribe []WebsocketChannelSubscription
+	channelSubscriber   func(channelToSubscribe WebsocketChannelSubscription) error
+	channelUnsubscriber func(channelToUnsubscribe WebsocketChannelSubscription) error
+	// Connected denotes a channel switch for diversion of request flow
+	Connected chan struct{}
+	// Disconnected denotes a channel switch for diversion of request flow
+	Disconnected chan struct{}
+	// DataHandler pipes websocket data to an exchange websocket data handler
+	DataHandler chan interface{}
+	// ShutdownC is the main shutdown channel which controls all websocket go funcs
+	ShutdownC                 chan struct{}
+	ShutdownConnectionMonitor chan struct{}
+	// Orderbook is a local cache of orderbooks
+	Orderbook WebsocketOrderbookLocal
+
+	// Wg defines a wait group for websocket routines for cleanly shutting down
+	// routines
+	Wg sync.WaitGroup
+	// TrafficAlert monitors if there is a halt in traffic throughput
+	TrafficAlert chan struct{}
+	// Functionality defines websocket stream capabilities
+	Functionality uint32
+}
+
+// WebsocketChannelSubscription container for websocket subscriptions
+// Currently only a one at a time thing to avoid complexity
+type WebsocketChannelSubscription struct {
+	Channel  string
+	Currency currency.Pair
+	Params   map[string]interface{}
+}
+
+// WebsocketOrderbookLocal defines a local cache of orderbooks for amending,
+// appending and deleting changes and updates the main store in orderbook.go
+type WebsocketOrderbookLocal struct {
+	ob          []*orderbook.Base
+	lastUpdated time.Time
+	m           sync.Mutex
+}
+
+// WebsocketResponse defines generalised data from the websocket connection
+type WebsocketResponse struct {
+	Type int
+	Raw  []byte
+}
+
+// WebsocketOrderbookUpdate defines a websocket event in which the orderbook
+// has been updated in the orderbook package
+type WebsocketOrderbookUpdate struct {
+	Pair     currency.Pair
+	Asset    string
+	Exchange string
+}
+
+// TradeData defines trade data
+type TradeData struct {
+	Timestamp    time.Time
+	CurrencyPair currency.Pair
+	AssetType    string
+	Exchange     string
+	EventType    string
+	EventTime    int64
+	Price        float64
+	Amount       float64
+	Side         string
+}
+
+// TickerData defines ticker feed
+type TickerData struct {
+	Timestamp  time.Time
+	Pair       currency.Pair
+	AssetType  string
+	Exchange   string
+	ClosePrice float64
+	Quantity   float64
+	OpenPrice  float64
+	HighPrice  float64
+	LowPrice   float64
+}
+
+// KlineData defines kline feed
+type KlineData struct {
+	Timestamp  time.Time
+	Pair       currency.Pair
+	AssetType  string
+	Exchange   string
+	StartTime  time.Time
+	CloseTime  time.Time
+	Interval   string
+	OpenPrice  float64
+	ClosePrice float64
+	HighPrice  float64
+	LowPrice   float64
+	Volume     float64
+}
+
+// WebsocketPositionUpdated reflects a change in orders/contracts on an exchange
+type WebsocketPositionUpdated struct {
+	Timestamp time.Time
+	Pair      currency.Pair
+	AssetType string
+	Exchange  string
+}

exchanges/exmo/exmo_wrapper.go

@@ -392,3 +392,15 @@ func (e *EXMO) GetOrderHistory(getOrdersRequest *exchange.GetOrdersRequest) ([]e
 
 	return orders, nil
 }
+
+// SubscribeToWebsocketChannels appends to ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle subscribing
+func (e *EXMO) SubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	return common.ErrFunctionNotSupported
+}
+
+// UnsubscribeToWebsocketChannels removes from ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle unsubscribing
+func (e *EXMO) UnsubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	return common.ErrFunctionNotSupported
+}

exchanges/gateio/gateio.go

@@ -7,6 +7,7 @@ import (
 	"net/http"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/gorilla/websocket"
@@ -49,6 +50,7 @@ const (
 type Gateio struct {
 	WebsocketConn *websocket.Conn
 	exchange.Base
+	wsRequestMtx sync.Mutex
 }
 
 // SetDefaults sets default values for the exchange
@@ -78,7 +80,9 @@ func (g *Gateio) SetDefaults() {
 	g.Websocket.Functionality = exchange.WebsocketTickerSupported |
 		exchange.WebsocketTradeDataSupported |
 		exchange.WebsocketOrderbookSupported |
-		exchange.WebsocketKlineSupported
+		exchange.WebsocketKlineSupported |
+		exchange.WebsocketSubscribeSupported |
+		exchange.WebsocketUnsubscribeSupported
 }
 
 // Setup sets user configuration
@@ -120,8 +124,11 @@ func (g *Gateio) Setup(exch *config.ExchangeConfig) {
 			log.Fatal(err)
 		}
 		err = g.WebsocketSetup(g.WsConnect,
+			g.Subscribe,
+			g.Unsubscribe,
 			exch.Name,
 			exch.Websocket,
+			exch.Verbose,
 			gateioWebsocketEndpoint,
 			exch.WebsocketURL)
 		if err != nil {
@@ -461,6 +468,7 @@ func (g *Gateio) GetTradeHistory(symbol string) (TradHistoryResponse, error) {
 	return result, nil
 }
 
+// GenerateSignature returns hash for authenticated requests
 func (g *Gateio) GenerateSignature(message string) []byte {
 	return common.GetHMAC(common.HashSHA512, []byte(message), []byte(g.APISecret))
 }

exchanges/gateio/gateio_types.go

@@ -35,6 +35,7 @@ var (
 	TimeIntervalDay            = TimeInterval(60 * 60 * 24)
 )
 
+// IDs for requests
 const (
 	IDGeneric    = 0000
 	IDSignIn     = 1010

exchanges/gateio/gateio_websocket.go

@@ -7,6 +7,7 @@ import (
 	"net/http"
 	"net/url"
 	"strconv"
+	"strings"
 	"time"
 
 	"github.com/gorilla/websocket"
@@ -18,8 +19,9 @@ import (
 )
 
 const (
-	gateioWebsocketEndpoint = "wss://ws.gate.io/v3/"
-	gatioWsMethodPing       = "ping"
+	gateioWebsocketEndpoint  = "wss://ws.gate.io/v3/"
+	gatioWsMethodPing        = "ping"
+	gateioWebsocketRateLimit = 120 * time.Millisecond
 )
 
 // WsConnect initiates a websocket connection
@@ -54,8 +56,9 @@ func (g *Gateio) WsConnect() error {
 	}
 
 	go g.WsHandleData()
+	g.GenerateDefaultSubscriptions()
 
-	return g.WsSubscribe()
+	return nil
 }
 
 func (g *Gateio) wsServerSignIn() error {
@@ -67,84 +70,7 @@ func (g *Gateio) wsServerSignIn() error {
 		Method: "server.sign",
 		Params: []interface{}{g.APIKey, signature, nonce},
 	}
-	return g.WebsocketConn.WriteJSON(signinWsRequest)
-}
-
-// WsSubscribe subscribes to the full websocket suite on ZB exchange
-func (g *Gateio) WsSubscribe() error {
-	enabled := g.GetEnabledCurrencies()
-
-	for _, c := range enabled {
-		ticker := WebsocketRequest{
-			ID:     1337,
-			Method: "ticker.subscribe",
-			Params: []interface{}{c.String()},
-		}
-
-		err := g.WebsocketConn.WriteJSON(ticker)
-		if err != nil {
-			return err
-		}
-
-		trade := WebsocketRequest{
-			ID:     1337,
-			Method: "trades.subscribe",
-			Params: []interface{}{c.String()},
-		}
-
-		err = g.WebsocketConn.WriteJSON(trade)
-		if err != nil {
-			return err
-		}
-
-		depth := WebsocketRequest{
-			ID:     1337,
-			Method: "depth.subscribe",
-			Params: []interface{}{c.String(), 30, "0.1"},
-		}
-
-		err = g.WebsocketConn.WriteJSON(depth)
-		if err != nil {
-			return err
-		}
-
-		kline := WebsocketRequest{
-			ID:     1337,
-			Method: "kline.subscribe",
-			Params: []interface{}{c.String(), 1800},
-		}
-
-		err = g.WebsocketConn.WriteJSON(kline)
-		if err != nil {
-			return err
-		}
-	}
-
-	if g.AuthenticatedAPISupport {
-		balance := WebsocketRequest{
-			ID:     IDBalance,
-			Method: "balance.subscribe",
-			Params: []interface{}{},
-		}
-
-		err := g.WebsocketConn.WriteJSON(balance)
-		if err != nil {
-			return err
-		}
-
-		for _, c := range enabled {
-			orderNotification := WebsocketRequest{
-				ID:     IDGeneric,
-				Method: "order.subscribe",
-				Params: []interface{}{c.String()},
-			}
-			err := g.WebsocketConn.WriteJSON(orderNotification)
-			if err != nil {
-				return err
-			}
-		}
-	}
-	return nil
+	return g.wsSend(signinWsRequest)
 }
 
 // WsReadData reads from the websocket connection and returns the websocket
@@ -165,11 +91,6 @@ func (g *Gateio) WsHandleData() {
 	g.Websocket.Wg.Add(1)
 
 	defer func() {
-		err := g.WebsocketConn.Close()
-		if err != nil {
-			g.Websocket.DataHandler <- fmt.Errorf("gateio_websocket.go - Unable to to close Websocket connection. Error: %s",
-				err)
-		}
 		g.Websocket.Wg.Done()
 	}()
 
@@ -182,6 +103,8 @@ func (g *Gateio) WsHandleData() {
 			resp, err := g.WsReadData()
 			if err != nil {
 				g.Websocket.DataHandler <- err
+				// Read data error messages can overwhelm and panic the application
+				time.Sleep(time.Second)
 				continue
 			}
 
@@ -416,13 +339,72 @@ func (g *Gateio) WsHandleData() {
 	}
 }
 
+// GenerateDefaultSubscriptions Adds default subscriptions to websocket to be handled by ManageSubscriptions()
+func (g *Gateio) GenerateDefaultSubscriptions() {
+	var channels = []string{"ticker.subscribe", "trades.subscribe", "depth.subscribe", "kline.subscribe"}
+	if g.AuthenticatedAPISupport {
+		channels = append(channels, "balance.subscribe", "order.subscribe")
+	}
+
+	subscriptions := []exchange.WebsocketChannelSubscription{}
+	enabledCurrencies := g.GetEnabledCurrencies()
+	for i := range channels {
+		for j := range enabledCurrencies {
+			params := make(map[string]interface{})
+			if strings.EqualFold(channels[i], "depth.subscribe") {
+				params["limit"] = 30
+				params["interval"] = "0.1"
+			} else if strings.EqualFold(channels[i], "kline.subscribe") {
+				params["interval"] = 1800
+			}
+			subscriptions = append(subscriptions, exchange.WebsocketChannelSubscription{
+				Channel:  channels[i],
+				Currency: enabledCurrencies[j],
+				Params:   params,
+			})
+		}
+	}
+	g.Websocket.SubscribeToChannels(subscriptions)
+}
+
+// Subscribe sends a websocket message to receive data from the channel
+func (g *Gateio) Subscribe(channelToSubscribe exchange.WebsocketChannelSubscription) error {
+	params := []interface{}{channelToSubscribe.Currency.String()}
+	for _, paramValue := range channelToSubscribe.Params {
+		params = append(params, paramValue)
+	}
+
+	subscribe := WebsocketRequest{
+		ID:     IDGeneric,
+		Method: channelToSubscribe.Channel,
+		Params: params,
+	}
+
+	if strings.EqualFold(channelToSubscribe.Channel, "balance.subscribe") {
+		subscribe.ID = IDBalance
+	}
+
+	return g.wsSend(subscribe)
+}
+
+// Unsubscribe sends a websocket message to stop receiving data from the channel
+func (g *Gateio) Unsubscribe(channelToSubscribe exchange.WebsocketChannelSubscription) error {
+	unsbuscribeText := strings.Replace(channelToSubscribe.Channel, "subscribe", "unsubscribe", 1)
+	subscribe := WebsocketRequest{
+		ID:     IDGeneric,
+		Method: unsbuscribeText,
+		Params: []interface{}{channelToSubscribe.Currency.String(), 1800},
+	}
+	return g.wsSend(subscribe)
+}
+
 func (g *Gateio) wsGetBalance() error {
 	balanceWsRequest := WebsocketRequest{
 		ID:     IDBalance,
 		Method: "balance.query",
 		Params: []interface{}{},
 	}
-	return g.WebsocketConn.WriteJSON(balanceWsRequest)
+	return g.wsSend(balanceWsRequest)
 }
 
 func (g *Gateio) wsGetOrderInfo(market string, offset, limit int) error {
@@ -435,5 +417,17 @@ func (g *Gateio) wsGetOrderInfo(market string, offset, limit int) error {
 			limit,
 		},
 	}
-	return g.WebsocketConn.WriteJSON(order)
+	return g.wsSend(order)
+}
+
+// WsSend sends data to the websocket server
+func (g *Gateio) wsSend(data interface{}) error {
+	g.wsRequestMtx.Lock()
+	defer g.wsRequestMtx.Unlock()
+	if g.Verbose {
+		log.Debugf("%v sending message to websocket %v", g.Name, data)
+	}
+	// Basic rate limiter
+	time.Sleep(gateioWebsocketRateLimit)
+	return g.WebsocketConn.WriteJSON(data)
 }

exchanges/gateio/gateio_wrapper.go

@@ -445,3 +445,17 @@ func (g *Gateio) GetOrderHistory(getOrdersRequest *exchange.GetOrdersRequest) ([
 
 	return orders, nil
 }
+
+// SubscribeToWebsocketChannels appends to ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle subscribing
+func (g *Gateio) SubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	g.Websocket.SubscribeToChannels(channels)
+	return nil
+}
+
+// UnsubscribeToWebsocketChannels removes from ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle unsubscribing
+func (g *Gateio) UnsubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	g.Websocket.UnsubscribeToChannels(channels)
+	return nil
+}

exchanges/gemini/gemini.go

@@ -167,8 +167,11 @@ func (g *Gemini) Setup(exch *config.ExchangeConfig) {
 			log.Fatal(err)
 		}
 		err = g.WebsocketSetup(g.WsConnect,
+			nil,
+			nil,
 			exch.Name,
 			exch.Websocket,
+			exch.Verbose,
 			geminiWebsocketEndpoint,
 			exch.WebsocketURL)
 		if err != nil {

exchanges/gemini/gemini_types.go

@@ -214,7 +214,7 @@ type Event struct {
 	Remaining float64 `json:"remaining,string"`
 	Side      string  `json:"side"`
 	MakerSide string  `json:"makerSide"`
-	Amount    float64 `json:"amount"`
+	Amount    float64 `json:"amount,string"`
 }
 
 // ReadData defines read data from the websocket connection

exchanges/gemini/gemini_wrapper.go

@@ -357,3 +357,15 @@ func (g *Gemini) GetOrderHistory(getOrdersRequest *exchange.GetOrdersRequest) ([
 
 	return orders, nil
 }
+
+// SubscribeToWebsocketChannels appends to ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle subscribing
+func (g *Gemini) SubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	return common.ErrFunctionNotSupported
+}
+
+// UnsubscribeToWebsocketChannels removes from ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle unsubscribing
+func (g *Gemini) UnsubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	return common.ErrFunctionNotSupported
+}

exchanges/hitbtc/hitbtc.go

@@ -7,6 +7,7 @@ import (
 	"net/http"
 	"net/url"
 	"strconv"
+	"sync"
 	"time"
 
 	"github.com/gorilla/websocket"
@@ -54,6 +55,7 @@ const (
 type HitBTC struct {
 	exchange.Base
 	WebsocketConn *websocket.Conn
+	wsRequestMtx  sync.Mutex
 }
 
 // SetDefaults sets default settings for hitbtc
@@ -80,7 +82,9 @@ func (h *HitBTC) SetDefaults() {
 	h.APIUrl = h.APIUrlDefault
 	h.WebsocketInit()
 	h.Websocket.Functionality = exchange.WebsocketTickerSupported |
-		exchange.WebsocketOrderbookSupported
+		exchange.WebsocketOrderbookSupported |
+		exchange.WebsocketSubscribeSupported |
+		exchange.WebsocketUnsubscribeSupported
 }
 
 // Setup sets user exchange configuration settings
@@ -121,8 +125,11 @@ func (h *HitBTC) Setup(exch *config.ExchangeConfig) {
 			log.Fatal(err)
 		}
 		err = h.WebsocketSetup(h.WsConnect,
+			h.Subscribe,
+			h.Unsubscribe,
 			exch.Name,
 			exch.Websocket,
+			exch.Verbose,
 			hitbtcWebsocketAddress,
 			exch.WebsocketURL)
 		if err != nil {

exchanges/hitbtc/hitbtc_types.go

@@ -293,3 +293,80 @@ type LendingHistory struct {
 	Open     string  `json:"open"`
 	Close    string  `json:"close"`
 }
+
+type capture struct {
+	Method string `json:"method"`
+	Result bool   `json:"result"`
+	Error  struct {
+		Code    int    `json:"code"`
+		Message string `json:"message"`
+	} `json:"error"`
+}
+
+// WsRequest defines a request obj for the JSON-RPC and gets a websocket
+// response
+type WsRequest struct {
+	Method string      `json:"method"`
+	Params interface{} `json:"params,omitempty"`
+	ID     interface{} `json:"id"`
+}
+
+// WsNotification defines a notification obj for the JSON-RPC this does not get
+// a websocket response
+type WsNotification struct {
+	JSONRPCVersion string      `json:"jsonrpc"`
+	Method         string      `json:"method"`
+	Params         interface{} `json:"params"`
+}
+
+type params struct {
+	Symbol string `json:"symbol"`
+	Period string `json:"period,omitempty"`
+	Limit  int64  `json:"limit,omitempty"`
+}
+
+// WsTicker defines websocket ticker feed return params
+type WsTicker struct {
+	Params struct {
+		Ask         float64 `json:"ask,string"`
+		Bid         float64 `json:"bid,string"`
+		Last        float64 `json:"last,string"`
+		Open        float64 `json:"open,string"`
+		Low         float64 `json:"low,string"`
+		High        float64 `json:"high,string"`
+		Volume      float64 `json:"volume,string"`
+		VolumeQuote float64 `json:"volumeQuote,string"`
+		Timestamp   string  `json:"timestamp"`
+		Symbol      string  `json:"symbol"`
+	} `json:"params"`
+}
+
+// WsOrderbook defines websocket orderbook feed return params
+type WsOrderbook struct {
+	Params struct {
+		Ask []struct {
+			Price float64 `json:"price,string"`
+			Size  float64 `json:"size,string"`
+		} `json:"ask"`
+		Bid []struct {
+			Price float64 `json:"price,string"`
+			Size  float64 `json:"size,string"`
+		} `json:"bid"`
+		Symbol   string `json:"symbol"`
+		Sequence int64  `json:"sequence"`
+	} `json:"params"`
+}
+
+// WsTrade defines websocket trade feed return params
+type WsTrade struct {
+	Params struct {
+		Data []struct {
+			ID        int64   `json:"id"`
+			Price     float64 `json:"price,string"`
+			Quantity  float64 `json:"quantity,string"`
+			Side      string  `json:"side"`
+			Timestamp string  `json:"timestamp"`
+		} `json:"data"`
+		Symbol string `json:"symbol"`
+	} `json:"params"`
+}

exchanges/hitbtc/hitbtc_websocket.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"net/http"
 	"net/url"
+	"strings"
 	"time"
 
 	"github.com/gorilla/websocket"
@@ -12,6 +13,7 @@ import (
 	"github.com/thrasher-/gocryptotrader/currency"
 	exchange "github.com/thrasher-/gocryptotrader/exchanges"
 	"github.com/thrasher-/gocryptotrader/exchanges/orderbook"
+	log "github.com/thrasher-/gocryptotrader/logger"
 )
 
 const (
@@ -43,58 +45,8 @@ func (h *HitBTC) WsConnect() error {
 	}
 
 	go h.WsHandleData()
+	h.GenerateDefaultSubscriptions()
 
-	return h.WsSubscribe()
-}
-
-// WsSubscribe subscribes to the relevant channels
-func (h *HitBTC) WsSubscribe() error {
-	enabledPairs := h.GetEnabledCurrencies()
-	for _, p := range enabledPairs {
-		pF := exchange.FormatExchangeCurrency(h.GetName(), p)
-
-		tickerSubReq, err := common.JSONEncode(WsNotification{
-			JSONRPCVersion: rpcVersion,
-			Method:         "subscribeTicker",
-			Params:         params{Symbol: pF.String()},
-		})
-		if err != nil {
-			return err
-		}
-
-		err = h.WebsocketConn.WriteMessage(websocket.TextMessage, tickerSubReq)
-		if err != nil {
-			return nil
-		}
-
-		orderbookSubReq, err := common.JSONEncode(WsNotification{
-			JSONRPCVersion: rpcVersion,
-			Method:         "subscribeOrderbook",
-			Params:         params{Symbol: pF.String()},
-		})
-		if err != nil {
-			return err
-		}
-
-		err = h.WebsocketConn.WriteMessage(websocket.TextMessage, orderbookSubReq)
-		if err != nil {
-			return nil
-		}
-
-		tradeSubReq, err := common.JSONEncode(WsNotification{
-			JSONRPCVersion: rpcVersion,
-			Method:         "subscribeTrades",
-			Params:         params{Symbol: pF.String()},
-		})
-		if err != nil {
-			return err
-		}
-
-		err = h.WebsocketConn.WriteMessage(websocket.TextMessage, tradeSubReq)
-		if err != nil {
-			return nil
-		}
-	}
 	return nil
 }
 
@@ -114,11 +66,6 @@ func (h *HitBTC) WsHandleData() {
 	h.Websocket.Wg.Add(1)
 
 	defer func() {
-		err := h.WebsocketConn.Close()
-		if err != nil {
-			h.Websocket.DataHandler <- fmt.Errorf("hitbtc_websocket.go - Unable to to close Websocket connection. Error: %s",
-				err)
-		}
 		h.Websocket.Wg.Done()
 	}()
 
@@ -290,77 +237,84 @@ func (h *HitBTC) WsProcessOrderbookUpdate(ob WsOrderbook) error {
 	return nil
 }
 
-type capture struct {
-	Method string `json:"method"`
-	Result bool   `json:"result"`
-	Error  struct {
-		Code    int    `json:"code"`
-		Message string `json:"message"`
-	} `json:"error"`
+// GenerateDefaultSubscriptions Adds default subscriptions to websocket to be handled by ManageSubscriptions()
+func (h *HitBTC) GenerateDefaultSubscriptions() {
+	var channels = []string{"subscribeTicker", "subscribeOrderbook", "subscribeTrades", "subscribeCandles"}
+	subscriptions := []exchange.WebsocketChannelSubscription{}
+	enabledCurrencies := h.GetEnabledCurrencies()
+	for i := range channels {
+		for j := range enabledCurrencies {
+			enabledCurrencies[j].Delimiter = ""
+			subscriptions = append(subscriptions, exchange.WebsocketChannelSubscription{
+				Channel:  channels[i],
+				Currency: enabledCurrencies[j],
+			})
+		}
+	}
+	h.Websocket.SubscribeToChannels(subscriptions)
 }
 
-// WsRequest defines a request obj for the JSON-RPC and gets a websocket
-// response
-type WsRequest struct {
-	Method string      `json:"method"`
-	Params interface{} `json:"params,omitempty"`
-	ID     interface{} `json:"id"`
+// Subscribe sends a websocket message to receive data from the channel
+func (h *HitBTC) Subscribe(channelToSubscribe exchange.WebsocketChannelSubscription) error {
+	subscribe := WsNotification{
+		JSONRPCVersion: rpcVersion,
+		Method:         channelToSubscribe.Channel,
+		Params: params{
+			Symbol: channelToSubscribe.Currency.String(),
+		},
+	}
+	if strings.EqualFold(channelToSubscribe.Channel, "subscribeTrades") {
+		subscribe.Params = params{
+			Symbol: channelToSubscribe.Currency.String(),
+			Limit:  100,
+		}
+	} else if strings.EqualFold(channelToSubscribe.Channel, "subscribeCandles") {
+		subscribe.Params = params{
+			Symbol: channelToSubscribe.Currency.String(),
+			Period: "M30",
+			Limit:  100,
+		}
+	}
+
+	return h.wsSend(subscribe)
 }
 
-// WsNotification defines a notification obj for the JSON-RPC this does not get
-// a websocket response
-type WsNotification struct {
-	JSONRPCVersion string      `json:"jsonrpc"`
-	Method         string      `json:"method"`
-	Params         interface{} `json:"params"`
+// Unsubscribe sends a websocket message to stop receiving data from the channel
+func (h *HitBTC) Unsubscribe(channelToSubscribe exchange.WebsocketChannelSubscription) error {
+	unsubscribeChannel := strings.Replace(channelToSubscribe.Channel, "subscribe", "unsubscribe", 1)
+	subscribe := WsNotification{
+		JSONRPCVersion: rpcVersion,
+		Method:         unsubscribeChannel,
+		Params: params{
+			Symbol: channelToSubscribe.Currency.String(),
+		},
+	}
+	if strings.EqualFold(unsubscribeChannel, "unsubscribeTrades") {
+		subscribe.Params = params{
+			Symbol: channelToSubscribe.Currency.String(),
+			Limit:  100,
+		}
+	} else if strings.EqualFold(unsubscribeChannel, "unsubscribeCandles") {
+		subscribe.Params = params{
+			Symbol: channelToSubscribe.Currency.String(),
+			Period: "M30",
+			Limit:  100,
+		}
+	}
+
+	return h.wsSend(subscribe)
 }
 
-type params struct {
-	Symbol string `json:"symbol"`
-}
-
-// WsTicker defines websocket ticker feed return params
-type WsTicker struct {
-	Params struct {
-		Ask         float64 `json:"ask,string"`
-		Bid         float64 `json:"bid,string"`
-		Last        float64 `json:"last,string"`
-		Open        float64 `json:"open,string"`
-		Low         float64 `json:"low,string"`
-		High        float64 `json:"high,string"`
-		Volume      float64 `json:"volume,string"`
-		VolumeQuote float64 `json:"volumeQuote,string"`
-		Timestamp   string  `json:"timestamp"`
-		Symbol      string  `json:"symbol"`
-	} `json:"params"`
-}
-
-// WsOrderbook defines websocket orderbook feed return params
-type WsOrderbook struct {
-	Params struct {
-		Ask []struct {
-			Price float64 `json:"price,string"`
-			Size  float64 `json:"size,string"`
-		} `json:"ask"`
-		Bid []struct {
-			Price float64 `json:"price,string"`
-			Size  float64 `json:"size,string"`
-		} `json:"bid"`
-		Symbol   string `json:"symbol"`
-		Sequence int64  `json:"sequence"`
-	} `json:"params"`
-}
-
-// WsTrade defines websocket trade feed return params
-type WsTrade struct {
-	Params struct {
-		Data []struct {
-			ID        int64   `json:"id"`
-			Price     float64 `json:"price,string"`
-			Quantity  float64 `json:"quantity,string"`
-			Side      string  `json:"side"`
-			Timestamp string  `json:"timestamp"`
-		} `json:"data"`
-		Symbol string `json:"symbol"`
-	} `json:"params"`
+// WsSend sends data to the websocket server
+func (h *HitBTC) wsSend(data interface{}) error {
+	h.wsRequestMtx.Lock()
+	defer h.wsRequestMtx.Unlock()
+	json, err := common.JSONEncode(data)
+	if err != nil {
+		return err
+	}
+	if h.Verbose {
+		log.Debugf("%v sending message to websocket %v", h.Name, data)
+	}
+	return h.WebsocketConn.WriteMessage(websocket.TextMessage, json)
 }

exchanges/hitbtc/hitbtc_wrapper.go

@@ -380,3 +380,17 @@ func (h *HitBTC) GetOrderHistory(getOrdersRequest *exchange.GetOrdersRequest) ([
 
 	return orders, nil
 }
+
+// SubscribeToWebsocketChannels appends to ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle subscribing
+func (h *HitBTC) SubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	h.Websocket.SubscribeToChannels(channels)
+	return nil
+}
+
+// UnsubscribeToWebsocketChannels removes from ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle unsubscribing
+func (h *HitBTC) UnsubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	h.Websocket.UnsubscribeToChannels(channels)
+	return nil
+}

exchanges/huobi/huobi.go

@@ -14,6 +14,7 @@ import (
 	"net/url"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/gorilla/websocket"
@@ -68,6 +69,7 @@ type HUOBI struct {
 	exchange.Base
 	AccountID     string
 	WebsocketConn *websocket.Conn
+	wsRequestMtx  sync.Mutex
 }
 
 // SetDefaults sets default values for the exchange
@@ -95,7 +97,9 @@ func (h *HUOBI) SetDefaults() {
 	h.WebsocketInit()
 	h.Websocket.Functionality = exchange.WebsocketKlineSupported |
 		exchange.WebsocketOrderbookSupported |
-		exchange.WebsocketTradeDataSupported
+		exchange.WebsocketTradeDataSupported |
+		exchange.WebsocketSubscribeSupported |
+		exchange.WebsocketUnsubscribeSupported
 }
 
 // Setup sets user configuration
@@ -138,8 +142,11 @@ func (h *HUOBI) Setup(exch *config.ExchangeConfig) {
 			log.Fatal(err)
 		}
 		err = h.WebsocketSetup(h.WsConnect,
+			h.Subscribe,
+			h.Unsubscribe,
 			exch.Name,
 			exch.Websocket,
+			exch.Verbose,
 			huobiSocketIOAddress,
 			exch.WebsocketURL)
 		if err != nil {

exchanges/huobi/huobi_types.go

@@ -259,3 +259,73 @@ var (
 	TimeIntervalMohth          = TimeInterval("1mon")
 	TimeIntervalYear           = TimeInterval("1year")
 )
+
+// WsRequest defines a request data structure
+type WsRequest struct {
+	Topic             string `json:"req,omitempty"`
+	Subscribe         string `json:"sub,omitempty"`
+	Unsubscribe       string `json:"unsub,omitempty"`
+	ClientGeneratedID string `json:"id,omitempty"`
+}
+
+// WsResponse defines a response from the websocket connection when there
+// is an error
+type WsResponse struct {
+	TS           int64  `json:"ts"`
+	Status       string `json:"status"`
+	ErrorCode    string `json:"err-code"`
+	ErrorMessage string `json:"err-msg"`
+	Ping         int64  `json:"ping"`
+	Channel      string `json:"ch"`
+	Subscribed   string `json:"subbed"`
+}
+
+// WsHeartBeat defines a heartbeat request
+type WsHeartBeat struct {
+	ClientNonce int64 `json:"ping"`
+}
+
+// WsDepth defines market depth websocket response
+type WsDepth struct {
+	Channel   string `json:"ch"`
+	Timestamp int64  `json:"ts"`
+	Tick      struct {
+		Bids      []interface{} `json:"bids"`
+		Asks      []interface{} `json:"asks"`
+		Timestamp int64         `json:"ts"`
+		Version   int64         `json:"version"`
+	} `json:"tick"`
+}
+
+// WsKline defines market kline websocket response
+type WsKline struct {
+	Channel   string `json:"ch"`
+	Timestamp int64  `json:"ts"`
+	Tick      struct {
+		ID     int64   `json:"id"`
+		Open   float64 `json:"open"`
+		Close  float64 `json:"close"`
+		Low    float64 `json:"low"`
+		High   float64 `json:"high"`
+		Amount float64 `json:"amount"`
+		Volume float64 `json:"vol"`
+		Count  int64   `json:"count"`
+	}
+}
+
+// WsTrade defines market trade websocket response
+type WsTrade struct {
+	Channel   string `json:"ch"`
+	Timestamp int64  `json:"ts"`
+	Tick      struct {
+		ID        int64 `json:"id"`
+		Timestamp int64 `json:"ts"`
+		Data      []struct {
+			Amount    float64 `json:"amount"`
+			Timestamp int64   `json:"ts"`
+			ID        float64 `json:"id,string"`
+			Price     float64 `json:"price"`
+			Direction string  `json:"direction"`
+		} `json:"data"`
+	}
+}

exchanges/huobi/huobi_websocket.go

@@ -6,7 +6,6 @@ import (
 	"errors"
 	"fmt"
 	"io/ioutil"
-	"math/big"
 	"net/http"
 	"net/url"
 	"time"
@@ -51,7 +50,8 @@ func (h *HUOBI) WsConnect() error {
 
 	go h.WsHandleData()
 
-	return h.WsSubscribe()
+	h.GenerateDefaultSubscriptions()
+	return nil
 }
 
 // WsReadData reads data from the websocket connection
@@ -83,11 +83,6 @@ func (h *HUOBI) WsHandleData() {
 	h.Websocket.Wg.Add(1)
 
 	defer func() {
-		err := h.WebsocketConn.Close()
-		if err != nil {
-			h.Websocket.DataHandler <- fmt.Errorf("huobi_websocket.go - Unable to to close Websocket connection. Error: %s",
-				err)
-		}
 		h.Websocket.Wg.Done()
 	}()
 
@@ -222,115 +217,52 @@ func (h *HUOBI) WsProcessOrderbook(ob *WsDepth, symbol string) error {
 	return nil
 }
 
-// WsSubscribe susbcribes to the current websocket streams based on the enabled
-// pair
-func (h *HUOBI) WsSubscribe() error {
-	pairs := h.GetEnabledCurrencies()
-
-	for _, p := range pairs {
-		fPair := exchange.FormatExchangeCurrency(h.GetName(), p)
-
-		depthTopic := fmt.Sprintf(wsMarketDepth, fPair.String())
-		depthJSON, err := common.JSONEncode(WsRequest{Subscribe: depthTopic})
-		if err != nil {
-			return err
-		}
-
-		err = h.WebsocketConn.WriteMessage(websocket.TextMessage, depthJSON)
-		if err != nil {
-			return err
-		}
-
-		klineTopic := fmt.Sprintf(wsMarketKline, fPair.String())
-		KlineJSON, err := common.JSONEncode(WsRequest{Subscribe: klineTopic})
-		if err != nil {
-			return err
-		}
-
-		err = h.WebsocketConn.WriteMessage(websocket.TextMessage, KlineJSON)
-		if err != nil {
-			return err
-		}
-
-		tradeTopic := fmt.Sprintf(wsMarketTrade, fPair.String())
-		tradeJSON, err := common.JSONEncode(WsRequest{Subscribe: tradeTopic})
-		if err != nil {
-			return err
-		}
-
-		err = h.WebsocketConn.WriteMessage(websocket.TextMessage, tradeJSON)
-		if err != nil {
-			return err
+// GenerateDefaultSubscriptions Adds default subscriptions to websocket to be handled by ManageSubscriptions()
+func (h *HUOBI) GenerateDefaultSubscriptions() {
+	var channels = []string{wsMarketKline, wsMarketDepth, wsMarketTrade}
+	enabledCurrencies := h.GetEnabledCurrencies()
+	subscriptions := []exchange.WebsocketChannelSubscription{}
+	for i := range channels {
+		for j := range enabledCurrencies {
+			enabledCurrencies[j].Delimiter = ""
+			channel := fmt.Sprintf(channels[i], enabledCurrencies[j].Lower().String())
+			subscriptions = append(subscriptions, exchange.WebsocketChannelSubscription{
+				Channel:  channel,
+				Currency: enabledCurrencies[j],
+			})
 		}
 	}
-	return nil
+	h.Websocket.SubscribeToChannels(subscriptions)
 }
 
-// WsRequest defines a request data structure
-type WsRequest struct {
-	Topic             string `json:"req,omitempty"`
-	Subscribe         string `json:"sub,omitempty"`
-	ClientGeneratedID string `json:"id,omitempty"`
-}
-
-// WsResponse defines a response from the websocket connection when there
-// is an error
-type WsResponse struct {
-	TS           int64  `json:"ts"`
-	Status       string `json:"status"`
-	ErrorCode    string `json:"err-code"`
-	ErrorMessage string `json:"err-msg"`
-	Ping         int64  `json:"ping"`
-	Channel      string `json:"ch"`
-	Subscribed   string `json:"subbed"`
-}
-
-// WsHeartBeat defines a heartbeat request
-type WsHeartBeat struct {
-	ClientNonce int64 `json:"ping"`
-}
-
-// WsDepth defines market depth websocket response
-type WsDepth struct {
-	Channel   string `json:"ch"`
-	Timestamp int64  `json:"ts"`
-	Tick      struct {
-		Bids      []interface{} `json:"bids"`
-		Asks      []interface{} `json:"asks"`
-		Timestamp int64         `json:"ts"`
-		Version   int64         `json:"version"`
-	} `json:"tick"`
-}
-
-// WsKline defines market kline websocket response
-type WsKline struct {
-	Channel   string `json:"ch"`
-	Timestamp int64  `json:"ts"`
-	Tick      struct {
-		ID     int64   `json:"id"`
-		Open   float64 `json:"open"`
-		Close  float64 `json:"close"`
-		Low    float64 `json:"low"`
-		High   float64 `json:"high"`
-		Amount float64 `json:"amount"`
-		Volume float64 `json:"vol"`
-		Count  int64   `json:"count"`
+// Subscribe sends a websocket message to receive data from the channel
+func (h *HUOBI) Subscribe(channelToSubscribe exchange.WebsocketChannelSubscription) error {
+	subscriptionRequest := WsRequest{Subscribe: channelToSubscribe.Channel}
+	if h.Verbose {
+		log.Debugf("Subscription: %v", subscriptionRequest)
 	}
+	subscription, err := common.JSONEncode(subscriptionRequest)
+	if err != nil {
+		return err
+	}
+	return h.wsSend(subscription)
 }
 
-// WsTrade defines market trade websocket response
-type WsTrade struct {
-	Channel   string `json:"ch"`
-	Timestamp int64  `json:"ts"`
-	Tick      struct {
-		ID        int64 `json:"id"`
-		Timestamp int64 `json:"ts"`
-		Data      []struct {
-			Amount    float64 `json:"amount"`
-			Timestamp int64   `json:"ts"`
-			ID        big.Int `json:"id,number"`
-			Price     float64 `json:"price"`
-			Direction string  `json:"direction"`
-		} `json:"data"`
+// Unsubscribe sends a websocket message to stop receiving data from the channel
+func (h *HUOBI) Unsubscribe(channelToSubscribe exchange.WebsocketChannelSubscription) error {
+	subscription, err := common.JSONEncode(WsRequest{Unsubscribe: channelToSubscribe.Channel})
+	if err != nil {
+		return err
 	}
+	return h.wsSend(subscription)
+}
+
+// WsSend sends data to the websocket server
+func (h *HUOBI) wsSend(data []byte) error {
+	h.wsRequestMtx.Lock()
+	defer h.wsRequestMtx.Unlock()
+	if h.Verbose {
+		log.Debugf("%v sending message to websocket %s", h.Name, string(data))
+	}
+	return h.WebsocketConn.WriteMessage(websocket.TextMessage, data)
 }

exchanges/huobi/huobi_wrapper.go

@@ -511,3 +511,17 @@ func setOrderSideAndType(requestType string, orderDetail *exchange.OrderDetail)
 		orderDetail.OrderType = exchange.LimitOrderType
 	}
 }
+
+// SubscribeToWebsocketChannels appends to ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle subscribing
+func (h *HUOBI) SubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	h.Websocket.SubscribeToChannels(channels)
+	return nil
+}
+
+// UnsubscribeToWebsocketChannels removes from ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle unsubscribing
+func (h *HUOBI) UnsubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	h.Websocket.UnsubscribeToChannels(channels)
+	return nil
+}

exchanges/huobihadax/huobihadax.go

@@ -8,6 +8,7 @@ import (
 	"net/http"
 	"net/url"
 	"strconv"
+	"sync"
 	"time"
 
 	"github.com/gorilla/websocket"
@@ -63,6 +64,7 @@ const (
 type HUOBIHADAX struct {
 	WebsocketConn *websocket.Conn
 	exchange.Base
+	wsRequestMtx sync.Mutex
 }
 
 // SetDefaults sets default values for the exchange
@@ -90,7 +92,9 @@ func (h *HUOBIHADAX) SetDefaults() {
 	h.WebsocketInit()
 	h.Websocket.Functionality = exchange.WebsocketKlineSupported |
 		exchange.WebsocketTradeDataSupported |
-		exchange.WebsocketOrderbookSupported
+		exchange.WebsocketOrderbookSupported |
+		exchange.WebsocketSubscribeSupported |
+		exchange.WebsocketUnsubscribeSupported
 }
 
 // Setup sets user configuration
@@ -132,8 +136,11 @@ func (h *HUOBIHADAX) Setup(exch *config.ExchangeConfig) {
 			log.Fatal(err)
 		}
 		err = h.WebsocketSetup(h.WsConnect,
+			h.Subscribe,
+			h.Unsubscribe,
 			exch.Name,
 			exch.Websocket,
+			exch.Verbose,
 			huobiGlobalWebsocketEndpoint,
 			exch.WebsocketURL)
 		if err != nil {

exchanges/huobihadax/huobihadax_types.go

@@ -1,7 +1,5 @@
 package huobihadax
 
-import "math/big"
-
 // Response stores the Huobi response information
 type Response struct {
 	Status       string `json:"status"`
@@ -258,6 +256,7 @@ type History struct {
 type WsRequest struct {
 	Topic             string `json:"req,omitempty"`
 	Subscribe         string `json:"sub,omitempty"`
+	Unsubscribe       string `json:"unsub,omitempty"`
 	ClientGeneratedID string `json:"id,omitempty"`
 }
 
@@ -316,7 +315,7 @@ type WsTrade struct {
 		Data      []struct {
 			Amount    float64 `json:"amount"`
 			Timestamp int64   `json:"ts"`
-			ID        big.Int `json:"id,number"`
+			ID        float64 `json:"id,string"`
 			Price     float64 `json:"price"`
 			Direction string  `json:"direction"`
 		} `json:"data"`

exchanges/huobihadax/huobihadax_websocket.go

@@ -15,6 +15,7 @@ import (
 	"github.com/thrasher-/gocryptotrader/currency"
 	exchange "github.com/thrasher-/gocryptotrader/exchanges"
 	"github.com/thrasher-/gocryptotrader/exchanges/orderbook"
+	log "github.com/thrasher-/gocryptotrader/logger"
 )
 
 const (
@@ -51,7 +52,8 @@ func (h *HUOBIHADAX) WsConnect() error {
 
 	go h.WsHandleData()
 
-	return h.WsSubscribe()
+	h.GenerateDefaultSubscriptions()
+	return nil
 }
 
 // WsReadData reads data from the websocket connection
@@ -83,11 +85,6 @@ func (h *HUOBIHADAX) WsHandleData() {
 	h.Websocket.Wg.Add(1)
 
 	defer func() {
-		err := h.WebsocketConn.Close()
-		if err != nil {
-			h.Websocket.DataHandler <- fmt.Errorf("huobi_websocket.go - Unable to to close Websocket connection. Error: %s",
-				err)
-		}
 		h.Websocket.Wg.Done()
 	}()
 
@@ -223,46 +220,48 @@ func (h *HUOBIHADAX) WsProcessOrderbook(ob *WsDepth, symbol string) error {
 	return nil
 }
 
-// WsSubscribe susbcribes to the current websocket streams based on the enabled
-// pair
-func (h *HUOBIHADAX) WsSubscribe() error {
-	pairs := h.GetEnabledCurrencies()
-
-	for _, p := range pairs {
-		fPair := exchange.FormatExchangeCurrency(h.GetName(), p)
-
-		depthTopic := fmt.Sprintf(wsMarketDepth, fPair.String())
-		depthJSON, err := common.JSONEncode(WsRequest{Subscribe: depthTopic})
-		if err != nil {
-			return err
-		}
-
-		err = h.WebsocketConn.WriteMessage(websocket.TextMessage, depthJSON)
-		if err != nil {
-			return err
-		}
-
-		klineTopic := fmt.Sprintf(wsMarketKline, fPair.String())
-		KlineJSON, err := common.JSONEncode(WsRequest{Subscribe: klineTopic})
-		if err != nil {
-			return err
-		}
-
-		err = h.WebsocketConn.WriteMessage(websocket.TextMessage, KlineJSON)
-		if err != nil {
-			return err
-		}
-
-		tradeTopic := fmt.Sprintf(wsMarketTrade, fPair.String())
-		tradeJSON, err := common.JSONEncode(WsRequest{Subscribe: tradeTopic})
-		if err != nil {
-			return err
-		}
-
-		err = h.WebsocketConn.WriteMessage(websocket.TextMessage, tradeJSON)
-		if err != nil {
-			return err
+// GenerateDefaultSubscriptions Adds default subscriptions to websocket to be handled by ManageSubscriptions()
+func (h *HUOBIHADAX) GenerateDefaultSubscriptions() {
+	var channels = []string{wsMarketKline, wsMarketDepth, wsMarketTrade}
+	enabledCurrencies := h.GetEnabledCurrencies()
+	subscriptions := []exchange.WebsocketChannelSubscription{}
+	for i := range channels {
+		for j := range enabledCurrencies {
+			enabledCurrencies[j].Delimiter = ""
+			channel := fmt.Sprintf(channels[i], enabledCurrencies[j].Lower().String())
+			subscriptions = append(subscriptions, exchange.WebsocketChannelSubscription{
+				Channel:  channel,
+				Currency: enabledCurrencies[j],
+			})
 		}
 	}
-	return nil
+	h.Websocket.SubscribeToChannels(subscriptions)
+}
+
+// Subscribe sends a websocket message to receive data from the channel
+func (h *HUOBIHADAX) Subscribe(channelToSubscribe exchange.WebsocketChannelSubscription) error {
+	subscription, err := common.JSONEncode(WsRequest{Subscribe: channelToSubscribe.Channel})
+	if err != nil {
+		return err
+	}
+	return h.wsSend(subscription)
+}
+
+// Unsubscribe sends a websocket message to stop receiving data from the channel
+func (h *HUOBIHADAX) Unsubscribe(channelToSubscribe exchange.WebsocketChannelSubscription) error {
+	subscription, err := common.JSONEncode(WsRequest{Unsubscribe: channelToSubscribe.Channel})
+	if err != nil {
+		return err
+	}
+	return h.wsSend(subscription)
+}
+
+// WsSend sends data to the websocket server
+func (h *HUOBIHADAX) wsSend(data []byte) error {
+	h.wsRequestMtx.Lock()
+	defer h.wsRequestMtx.Unlock()
+	if h.Verbose {
+		log.Debugf("%v sending message to websocket %s", h.Name, string(data))
+	}
+	return h.WebsocketConn.WriteMessage(websocket.TextMessage, data)
 }

exchanges/huobihadax/huobihadax_wrapper.go

@@ -450,3 +450,17 @@ func (h *HUOBIHADAX) GetOrderHistory(getOrdersRequest *exchange.GetOrdersRequest
 
 	return orders, nil
 }
+
+// SubscribeToWebsocketChannels appends to ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle subscribing
+func (h *HUOBIHADAX) SubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	h.Websocket.SubscribeToChannels(channels)
+	return nil
+}
+
+// UnsubscribeToWebsocketChannels removes from ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle unsubscribing
+func (h *HUOBIHADAX) UnsubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	h.Websocket.UnsubscribeToChannels(channels)
+	return nil
+}

exchanges/itbit/itbit_wrapper.go

@@ -406,3 +406,15 @@ func (i *ItBit) GetOrderHistory(getOrdersRequest *exchange.GetOrdersRequest) ([]
 
 	return orders, nil
 }
+
+// SubscribeToWebsocketChannels appends to ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle subscribing
+func (i *ItBit) SubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	return common.ErrFunctionNotSupported
+}
+
+// UnsubscribeToWebsocketChannels removes from ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle unsubscribing
+func (i *ItBit) UnsubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	return common.ErrFunctionNotSupported
+}

exchanges/kraken/kraken.go

@@ -60,7 +60,7 @@ type Kraken struct {
 	exchange.Base
 	WebsocketConn      *websocket.Conn
 	CryptoFee, FiatFee float64
-	mu                 sync.Mutex
+	wsRequestMtx       sync.Mutex
 }
 
 // SetDefaults sets current default settings
@@ -94,7 +94,9 @@ func (k *Kraken) SetDefaults() {
 	k.Websocket.Functionality = exchange.WebsocketTickerSupported |
 		exchange.WebsocketTradeDataSupported |
 		exchange.WebsocketKlineSupported |
-		exchange.WebsocketOrderbookSupported
+		exchange.WebsocketOrderbookSupported |
+		exchange.WebsocketSubscribeSupported |
+		exchange.WebsocketUnsubscribeSupported
 
 }
 
@@ -136,8 +138,11 @@ func (k *Kraken) Setup(exch *config.ExchangeConfig) {
 			log.Fatal(err)
 		}
 		err = k.WebsocketSetup(k.WsConnect,
+			k.Subscribe,
+			k.Unsubscribe,
 			exch.Name,
 			exch.Websocket,
+			exch.Verbose,
 			krakenWSURL,
 			exch.WebsocketURL)
 		if err != nil {

exchanges/kraken/kraken_test.go

@@ -748,167 +748,3 @@ func TestOrderBookOutOfOrder(t *testing.T) {
 		t.Error("Expected out of order orderbook error")
 	}
 }
-
-// TestSubscribeToChannel websocket test
-func TestSubscribeToChannel(t *testing.T) {
-	if k.Name == "" {
-		k.SetDefaults()
-		TestSetup(t)
-	}
-	if !k.Websocket.IsEnabled() {
-		t.Skip("Websocket not enabled, skipping")
-	}
-	if k.WebsocketConn == nil {
-		k.Websocket.Connect()
-	}
-
-	err := k.WsSubscribeToChannel("ticker", []string{"XTZ/USD"}, 1)
-	if err != nil {
-		t.Error(err)
-	}
-}
-
-// TestSubscribeToNonExistentChannel websocket test
-func TestSubscribeToNonExistentChannel(t *testing.T) {
-	if k.Name == "" {
-		k.SetDefaults()
-		TestSetup(t)
-	}
-	if !k.Websocket.IsEnabled() {
-		t.Skip("Websocket not enabled, skipping")
-	}
-	if k.WebsocketConn == nil {
-		k.Websocket.Connect()
-	}
-	err := k.WsSubscribeToChannel("ticker", []string{"pewdiepie"}, 1)
-	if err != nil {
-		t.Error(err)
-	}
-	subscriptionError := false
-	for i := 0; i < 7; i++ {
-		response := <-k.Websocket.DataHandler
-		if err, ok := response.(error); ok && err != nil {
-			subscriptionError = true
-			break
-		}
-	}
-	if !subscriptionError {
-		t.Error("Expected error")
-	}
-}
-
-// TestSubscribeUnsubscribeToChannel websocket test
-func TestSubscribeUnsubscribeToChannel(t *testing.T) {
-	if k.Name == "" {
-		k.SetDefaults()
-		TestSetup(t)
-	}
-	if !k.Websocket.IsEnabled() {
-		t.Skip("Websocket not enabled, skipping")
-	}
-	if k.WebsocketConn == nil {
-		k.Websocket.Connect()
-	}
-	err := k.WsSubscribeToChannel("ticker", []string{"XRP/JPY"}, 1)
-	if err != nil {
-		t.Error(err)
-	}
-	err = k.WsUnsubscribeToChannel("ticker", []string{"XRP/JPY"}, 2)
-	if err != nil {
-		t.Error(err)
-	}
-}
-
-// TestUnsubscribeWithoutSubscription websocket test
-func TestUnsubscribeWithoutSubscription(t *testing.T) {
-	if k.Name == "" {
-		k.SetDefaults()
-		TestSetup(t)
-	}
-	if !k.Websocket.IsEnabled() {
-		t.Skip("Websocket not enabled, skipping")
-	}
-	if k.WebsocketConn == nil {
-		k.Websocket.Connect()
-	}
-	err := k.WsUnsubscribeToChannel("ticker", []string{"QTUM/EUR"}, 3)
-	if err != nil {
-		t.Error(err)
-	}
-	unsubscriptionError := false
-	for i := 0; i < 5; i++ {
-		response := <-k.Websocket.DataHandler
-		t.Log(response)
-		if err, ok := response.(error); ok && err != nil {
-			if err.Error() == "requestID: '3'. Error: Subscription Not Found" {
-				unsubscriptionError = true
-				break
-			}
-		}
-	}
-	if !unsubscriptionError {
-		t.Error("Expected error")
-	}
-}
-
-// TestUnsubscribeWithChannelID websocket test
-func TestUnsubscribeWithChannelID(t *testing.T) {
-	if k.Name == "" {
-		k.SetDefaults()
-		TestSetup(t)
-	}
-	if !k.Websocket.IsEnabled() {
-		t.Skip("Websocket not enabled, skipping")
-	}
-	if k.WebsocketConn == nil {
-		k.Websocket.Connect()
-	}
-	err := k.WsUnsubscribeToChannelByChannelID(100)
-	if err != nil {
-		t.Error(err)
-	}
-	unsubscriptionError := false
-	for i := 0; i < 5; i++ {
-		response := <-k.Websocket.DataHandler
-		if err, ok := response.(error); ok && err != nil {
-			if err.Error() == "Not subscribed to the requested channelID" {
-				unsubscriptionError = true
-				break
-			}
-		}
-	}
-	if !unsubscriptionError {
-		t.Error("Expected error")
-	}
-}
-
-// TestUnsubscribeFromNonExistentChannel websocket test
-func TestUnsubscribeFromNonExistentChannel(t *testing.T) {
-	if k.Name == "" {
-		k.SetDefaults()
-		TestSetup(t)
-	}
-	if !k.Websocket.IsEnabled() {
-		t.Skip("Websocket not enabled, skipping")
-	}
-	if k.WebsocketConn == nil {
-		k.Websocket.Connect()
-	}
-	err := k.WsUnsubscribeToChannel("ticker", []string{"tseries"}, 0)
-	if err != nil {
-		t.Error(err)
-	}
-	unsubscriptionError := false
-	for i := 0; i < 5; i++ {
-		response := <-k.Websocket.DataHandler
-		if err, ok := response.(error); ok && err != nil {
-			if err.Error() == "Currency pair not in ISO 4217-A3 format tseries" {
-				unsubscriptionError = true
-				break
-			}
-		}
-	}
-	if !unsubscriptionError {
-		t.Error("Expected error")
-	}
-}

exchanges/kraken/kraken_websocket.go

@@ -11,7 +11,6 @@ import (
 	"net/url"
 	"sort"
 	"strconv"
-	"strings"
 	"sync"
 	"time"
 
@@ -47,6 +46,7 @@ const (
 	// Only supported asset type
 	krakenWsAssetType    = "SPOT"
 	orderbookBufferLimit = 3
+	krakenWsRateLimit    = 50 * time.Millisecond
 )
 
 // orderbookMutex Ensures if two entries arrive at once, only one can be processed at a time
@@ -62,16 +62,20 @@ var krakenOrderBooks map[int64]orderbook.Base
 var orderbookBuffer map[int64][]orderbook.Base
 var subscribeToDefaultChannels = true
 
+// Channels require a topic and a currency
+// Format [[ticker,but-t4u],[orderbook,nce-btt]]
+var defaultSubscribedChannels = []string{krakenWsTicker, krakenWsTrade, krakenWsOrderbook, krakenWsOHLC, krakenWsSpread}
+
 // writeToWebsocket sends a message to the websocket endpoint
 func (k *Kraken) writeToWebsocket(message []byte) error {
-	k.mu.Lock()
-	defer k.mu.Unlock()
+	k.wsRequestMtx.Lock()
+	defer k.wsRequestMtx.Unlock()
 	if k.Verbose {
 		log.Debugf("Sending message to WS: %v",
 			string(message))
 	}
 	// Really basic WS rate limit
-	time.Sleep(30 * time.Millisecond)
+	time.Sleep(krakenWsRateLimit)
 	return k.WebsocketConn.WriteMessage(websocket.TextMessage, message)
 }
 
@@ -110,24 +114,10 @@ func (k *Kraken) WsConnect() error {
 	go k.WsHandleData()
 	go k.wsPingHandler()
 	if subscribeToDefaultChannels {
-		k.WsSubscribeToDefaults()
+		k.GenerateDefaultSubscriptions()
 	}
-	return nil
-}
 
-// WsSubscribeToDefaults subscribes to the websocket channels
-func (k *Kraken) WsSubscribeToDefaults() {
-	channelsToSubscribe := []string{krakenWsTicker, krakenWsTrade, krakenWsOrderbook, krakenWsOHLC, krakenWsSpread}
-	for _, pair := range k.EnabledPairs {
-		// Kraken WS formats pairs with / but config and REST use -
-		formattedPair := strings.ToUpper(strings.Replace(pair.String(), "-", "/", 1))
-		for _, channel := range channelsToSubscribe {
-			err := k.WsSubscribeToChannel(channel, []string{formattedPair}, 0)
-			if err != nil {
-				k.Websocket.DataHandler <- err
-			}
-		}
-	}
+	return nil
 }
 
 // WsReadData reads data from the websocket connection
@@ -164,6 +154,7 @@ func (k *Kraken) wsPingHandler() {
 	k.Websocket.Wg.Add(1)
 	defer k.Websocket.Wg.Done()
 	ticker := time.NewTicker(time.Second * 27)
+
 	for {
 		select {
 		case <-k.Websocket.ShutdownC:
@@ -187,11 +178,6 @@ func (k *Kraken) wsPingHandler() {
 func (k *Kraken) WsHandleData() {
 	k.Websocket.Wg.Add(1)
 	defer func() {
-		err := k.WebsocketConn.Close()
-		if err != nil {
-			k.Websocket.DataHandler <- fmt.Errorf("%v unable to to close Websocket connection. Error: %s",
-				k.GetName(), err)
-		}
 		k.Websocket.Wg.Done()
 	}()
 
@@ -202,8 +188,10 @@ func (k *Kraken) WsHandleData() {
 		default:
 			resp, err := k.WsReadData()
 			if err != nil {
-				k.Websocket.DataHandler <- err
-				return
+				k.Websocket.DataHandler <- fmt.Errorf("%v WsHandleData: %v",
+					k.Name,
+					err)
+				time.Sleep(time.Second)
 			}
 			// event response handling
 			var eventResponse WebsocketEventResponse
@@ -219,8 +207,6 @@ func (k *Kraken) WsHandleData() {
 				k.WsHandleDataResponse(dataResponse)
 				continue
 			}
-			// Unknown data handling
-			k.Websocket.DataHandler <- fmt.Errorf("unrecognised response: %v", string(resp.Raw))
 			continue
 		}
 	}
@@ -312,64 +298,23 @@ func (k *Kraken) WsHandleEventResponse(response *WebsocketEventResponse) {
 	}
 }
 
-// WsSubscribeToChannel sends a request to WS to subscribe to supplied channel name and pairs
-func (k *Kraken) WsSubscribeToChannel(topic string, currencies []string, requestID int64) error {
-	resp := WebsocketSubscriptionEventRequest{
-		Event: krakenWsSubscribe,
-		Pairs: currencies,
-		Subscription: WebsocketSubscriptionData{
-			Name: topic,
-		},
-	}
-	if requestID > 0 {
-		resp.RequestID = requestID
-	}
-	json, err := common.JSONEncode(resp)
-	if err != nil {
-		return err
-	}
-	return k.writeToWebsocket(json)
-}
-
-// WsUnsubscribeToChannel sends a request to WS to unsubscribe to supplied channel name and pairs
-func (k *Kraken) WsUnsubscribeToChannel(topic string, currencies []string, requestID int64) error {
-	resp := WebsocketSubscriptionEventRequest{
-		Event: krakenWsUnsubscribe,
-		Pairs: currencies,
-		Subscription: WebsocketSubscriptionData{
-			Name: topic,
-		},
-	}
-	if requestID > 0 {
-		resp.RequestID = requestID
-	}
-	json, err := common.JSONEncode(resp)
-	if err != nil {
-		return err
-	}
-	return k.writeToWebsocket(json)
-}
-
-// WsUnsubscribeToChannelByChannelID sends a request to WS to unsubscribe to supplied channel ID
-func (k *Kraken) WsUnsubscribeToChannelByChannelID(channelID int64) error {
-	resp := WebsocketUnsubscribeByChannelIDEventRequest{
-		Event:     krakenWsUnsubscribe,
-		ChannelID: channelID,
-	}
-	json, err := common.JSONEncode(resp)
-	if err != nil {
-		return err
-	}
-	return k.writeToWebsocket(json)
-}
-
 // addNewSubscriptionChannelData stores channel ids, pairs and subscription types to an array
 // allowing correlation between subscriptions and returned data
 func addNewSubscriptionChannelData(response *WebsocketEventResponse) {
 	for i := range subscriptionChannelPair {
-		if response.ChannelID == subscriptionChannelPair[i].ChannelID {
-			return
+		if response.ChannelID != subscriptionChannelPair[i].ChannelID {
+			continue
 		}
+		// kill the stale orderbooks due to resubscribing
+		if orderbookBuffer == nil {
+			orderbookBuffer = make(map[int64][]orderbook.Base)
+		}
+		orderbookBuffer[response.ChannelID] = []orderbook.Base{}
+		if krakenOrderBooks == nil {
+			krakenOrderBooks = make(map[int64]orderbook.Base)
+		}
+		krakenOrderBooks[response.ChannelID] = orderbook.Base{}
+		return
 	}
 
 	// We change the / to - to maintain compatibility with REST/config
@@ -391,47 +336,6 @@ func getSubscriptionChannelData(id int64) WebsocketChannelData {
 	return WebsocketChannelData{}
 }
 
-// ResubscribeToChannel will attempt to unsubscribe and resubscribe to a channel
-func (k *Kraken) ResubscribeToChannel(channel string, pair currency.Pair) {
-	// Kraken WS formats pairs with / but config and REST use -
-	formattedPair := strings.ToUpper(strings.Replace(pair.String(), "-", "/", 1))
-	if krakenWsResubscribeFailureLimit > 0 {
-		var successfulUnsubscribe bool
-		for i := 0; i < krakenWsResubscribeFailureLimit; i++ {
-			err := k.WsUnsubscribeToChannel(channel, []string{formattedPair}, 0)
-			if err != nil {
-				log.Error(err)
-				time.Sleep(krakenWsResubscribeDelayInSeconds * time.Second)
-				continue
-			}
-			successfulUnsubscribe = true
-			break
-		}
-		if !successfulUnsubscribe {
-			log.Fatalf("%v websocket channel %v failed to unsubscribe after %v attempts",
-				k.GetName(), channel, krakenWsResubscribeFailureLimit)
-		}
-		successfulSubscribe := true
-		for i := 0; i < krakenWsResubscribeFailureLimit; i++ {
-			err := k.WsSubscribeToChannel(channel, []string{formattedPair}, 0)
-			if err != nil {
-				log.Error(err)
-				time.Sleep(krakenWsResubscribeDelayInSeconds * time.Second)
-				continue
-			}
-			successfulSubscribe = true
-			break
-		}
-		if !successfulSubscribe {
-			log.Fatalf("%v websocket channel %v failed to resubscribe after %v attempts",
-				k.GetName(), channel, krakenWsResubscribeFailureLimit)
-		}
-	} else {
-		log.Fatalf("%v websocket channel %v cannot resubscribe. Limit: %v",
-			k.GetName(), channel, krakenWsResubscribeFailureLimit)
-	}
-}
-
 // wsProcessTickers converts ticker data and sends it to the datahandler
 func (k *Kraken) wsProcessTickers(channelData *WebsocketChannelData, data interface{}) {
 	tickerData := data.(map[string]interface{})
@@ -509,13 +413,17 @@ func (k *Kraken) wsProcessOrderBook(channelData *WebsocketChannelData, data inte
 		_, asksExist := obData["a"]
 		_, bidsExist := obData["b"]
 		if asksExist || bidsExist {
-			k.mu.Lock()
-			defer k.mu.Unlock()
+			k.wsRequestMtx.Lock()
+			defer k.wsRequestMtx.Unlock()
 			k.wsProcessOrderBookBuffer(channelData, obData)
 			if len(orderbookBuffer[channelData.ChannelID]) >= orderbookBufferLimit {
 				err := k.wsProcessOrderBookUpdate(channelData)
 				if err != nil {
-					k.ResubscribeToChannel(channelData.Subscription, channelData.Pair)
+					subscriptionToRemove := exchange.WebsocketChannelSubscription{
+						Channel:  krakenWsOrderbook,
+						Currency: channelData.Pair,
+					}
+					k.Websocket.ResubscribeToChannel(subscriptionToRemove)
 				}
 			}
 		}
@@ -838,3 +746,57 @@ func (k *Kraken) wsProcessCandles(channelData *WebsocketChannelData, data interf
 		Volume:     volume,
 	}
 }
+
+// GenerateDefaultSubscriptions Adds default subscriptions to websocket to be handled by ManageSubscriptions()
+func (k *Kraken) GenerateDefaultSubscriptions() {
+	enabledCurrencies := k.GetEnabledCurrencies()
+	subscriptions := []exchange.WebsocketChannelSubscription{}
+	for i := range defaultSubscribedChannels {
+		for j := range enabledCurrencies {
+			enabledCurrencies[j].Delimiter = "/"
+			subscriptions = append(subscriptions, exchange.WebsocketChannelSubscription{
+				Channel:  defaultSubscribedChannels[i],
+				Currency: enabledCurrencies[j],
+			})
+		}
+	}
+	k.Websocket.SubscribeToChannels(subscriptions)
+}
+
+// Subscribe sends a websocket message to receive data from the channel
+func (k *Kraken) Subscribe(channelToSubscribe exchange.WebsocketChannelSubscription) error {
+	resp := WebsocketSubscriptionEventRequest{
+		Event: krakenWsSubscribe,
+		Pairs: []string{channelToSubscribe.Currency.String()},
+		Subscription: WebsocketSubscriptionData{
+			Name: channelToSubscribe.Channel,
+		},
+	}
+	json, err := common.JSONEncode(resp)
+	if err != nil {
+		if k.Verbose {
+			log.Debugf("%v subscribe error: %v", k.Name, err)
+		}
+		return err
+	}
+	return k.writeToWebsocket(json)
+}
+
+// Unsubscribe sends a websocket message to stop receiving data from the channel
+func (k *Kraken) Unsubscribe(channelToSubscribe exchange.WebsocketChannelSubscription) error {
+	resp := WebsocketSubscriptionEventRequest{
+		Event: krakenWsUnsubscribe,
+		Pairs: []string{channelToSubscribe.Currency.String()},
+		Subscription: WebsocketSubscriptionData{
+			Name: channelToSubscribe.Channel,
+		},
+	}
+	json, err := common.JSONEncode(resp)
+	if err != nil {
+		if k.Verbose {
+			log.Debugf("%v unsubscribe error: %v", k.Name, err)
+		}
+		return err
+	}
+	return k.writeToWebsocket(json)
+}

exchanges/kraken/kraken_wrapper.go

@@ -396,3 +396,17 @@ func (k *Kraken) GetOrderHistory(getOrdersRequest *exchange.GetOrdersRequest) ([
 
 	return orders, nil
 }
+
+// SubscribeToWebsocketChannels appends to ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle subscribing
+func (k *Kraken) SubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	k.Websocket.SubscribeToChannels(channels)
+	return nil
+}
+
+// UnsubscribeToWebsocketChannels removes from ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle unsubscribing
+func (k *Kraken) UnsubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	k.Websocket.UnsubscribeToChannels(channels)
+	return nil
+}

exchanges/lakebtc/lakebtc_wrapper.go

@@ -351,3 +351,15 @@ func (l *LakeBTC) GetOrderHistory(getOrdersRequest *exchange.GetOrdersRequest) (
 
 	return orders, nil
 }
+
+// SubscribeToWebsocketChannels appends to ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle subscribing
+func (l *LakeBTC) SubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	return common.ErrFunctionNotSupported
+}
+
+// UnsubscribeToWebsocketChannels removes from ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle unsubscribing
+func (l *LakeBTC) UnsubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	return common.ErrFunctionNotSupported
+}

exchanges/localbitcoins/localbitcoins_wrapper.go

@@ -428,3 +428,15 @@ func (l *LocalBitcoins) GetOrderHistory(getOrdersRequest *exchange.GetOrdersRequ
 
 	return orders, nil
 }
+
+// SubscribeToWebsocketChannels appends to ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle subscribing
+func (l *LocalBitcoins) SubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	return common.ErrFunctionNotSupported
+}
+
+// UnsubscribeToWebsocketChannels removes from ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle unsubscribing
+func (l *LocalBitcoins) UnsubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	return common.ErrFunctionNotSupported
+}

exchanges/okcoin/okcoin.go

@@ -54,5 +54,7 @@ func (o *OKCoin) SetDefaults() {
 	o.Websocket.Functionality = exchange.WebsocketTickerSupported |
 		exchange.WebsocketTradeDataSupported |
 		exchange.WebsocketKlineSupported |
-		exchange.WebsocketOrderbookSupported
+		exchange.WebsocketOrderbookSupported |
+		exchange.WebsocketSubscribeSupported |
+		exchange.WebsocketUnsubscribeSupported
 }

exchanges/okex/okex.go

@@ -79,7 +79,9 @@ func (o *OKEX) SetDefaults() {
 	o.Websocket.Functionality = exchange.WebsocketTickerSupported |
 		exchange.WebsocketTradeDataSupported |
 		exchange.WebsocketKlineSupported |
-		exchange.WebsocketOrderbookSupported
+		exchange.WebsocketOrderbookSupported |
+		exchange.WebsocketSubscribeSupported |
+		exchange.WebsocketUnsubscribeSupported
 }
 
 // GetFuturesPostions Get the information of all holding positions in futures trading.
@@ -252,7 +254,7 @@ func (o *OKEX) GetFuturesTokenInfoForCurrency(instrumentID string) (resp okgroup
 }
 
 // GetFuturesFilledOrder Get the recent 300 transactions of all contracts. Pagination is not supported here.
-// The whole book will be returned for one request. WebSocket is recommended here.
+// The whole book will be returned for one request. Websocket is recommended here.
 func (o *OKEX) GetFuturesFilledOrder(request okgroup.GetFuturesFilledOrderRequest) (resp []okgroup.GetFuturesFilledOrdersResponse, _ error) {
 	requestURL := fmt.Sprintf("%v/%v/%v%v", okgroup.OKGroupInstruments, request.InstrumentID, okgroup.OKGroupTrades, okgroup.FormatParameters(request))
 	return resp, o.SendHTTPRequest(http.MethodGet, okGroupFuturesSubsection, requestURL, nil, &resp, true)

exchanges/okgroup/okgroup.go

@@ -88,7 +88,7 @@ type OKGroup struct {
 	exchange.Base
 	ExchangeName  string
 	WebsocketConn *websocket.Conn
-	mu            sync.Mutex
+	wsRequestMtx  sync.Mutex
 	// Spot and contract market error codes as per https://www.okex.com/rest_request.html
 	ErrorCodes map[string]error
 	// Stores for corresponding variable checks
@@ -141,8 +141,11 @@ func (o *OKGroup) Setup(exch *config.ExchangeConfig) {
 			log.Fatal(err)
 		}
 		err = o.WebsocketSetup(o.WsConnect,
+			o.Subscribe,
+			o.Unsubscribe,
 			exch.Name,
 			exch.Websocket,
+			exch.Verbose,
 			o.WebsocketURL,
 			exch.WebsocketURL)
 		if err != nil {
@@ -359,7 +362,7 @@ func (o *OKGroup) GetSpotTokenPairDetails() (resp []GetSpotTokenPairDetailsRespo
 }
 
 // GetSpotOrderBook Getting the order book of a trading pair. Pagination is not supported here.
-// The whole book will be returned for one request. WebSocket is recommended here.
+// The whole book will be returned for one request. Websocket is recommended here.
 func (o *OKGroup) GetSpotOrderBook(request GetSpotOrderBookRequest) (resp GetSpotOrderBookResponse, _ error) {
 	requestURL := fmt.Sprintf("%v/%v/%v%v", OKGroupInstruments, request.InstrumentID, OKGroupGetSpotOrderBook, FormatParameters(request))
 	return resp, o.SendHTTPRequest(http.MethodGet, okGroupTokenSubsection, requestURL, nil, &resp, false)

exchanges/okgroup/okgroup_websocket.go

@@ -143,18 +143,23 @@ const (
 	okGroupWsFuturesAccount        = okGroupWsFuturesSubsection + okGroupWsAccount
 	okGroupWsFuturesPosition       = okGroupWsFuturesSubsection + okGroupWsPosition
 	okGroupWsFuturesOrder          = okGroupWsFuturesSubsection + okGroupWsOrder
+
+	okGroupWsRateLimit = 30 * time.Millisecond
 )
 
 // orderbookMutex Ensures if two entries arrive at once, only one can be processed at a time
 var orderbookMutex sync.Mutex
+var defaultSubscribedChannels = []string{okGroupWsSpotDepth, okGroupWsSpotCandle300s, okGroupWsSpotTicker, okGroupWsSpotTrade}
 
 // writeToWebsocket sends a message to the websocket endpoint
 func (o *OKGroup) writeToWebsocket(message string) error {
-	o.mu.Lock()
-	defer o.mu.Unlock()
+	o.wsRequestMtx.Lock()
+	defer o.wsRequestMtx.Unlock()
 	if o.Verbose {
-		log.Debugf("Sending message to WS: %v", message)
+		log.Debugf("%v sending message to WS: %v", o.Name, message)
 	}
+	// Really basic WS rate limit
+	time.Sleep(okGroupWsRateLimit)
 	return o.WebsocketConn.WriteMessage(websocket.TextMessage, []byte(message))
 }
 
@@ -165,13 +170,11 @@ func (o *OKGroup) WsConnect() error {
 	}
 
 	var dialer websocket.Dialer
-
 	if o.Websocket.GetProxyAddress() != "" {
 		proxy, err := url.Parse(o.Websocket.GetProxyAddress())
 		if err != nil {
 			return err
 		}
-
 		dialer.Proxy = http.ProxyURL(proxy)
 	}
 
@@ -187,41 +190,20 @@ func (o *OKGroup) WsConnect() error {
 			err)
 	}
 	if o.Verbose {
-		log.Debugf("Successful connection to %v", o.Websocket.GetWebsocketURL())
+		log.Debugf("Successful connection to %v",
+			o.Websocket.GetWebsocketURL())
 	}
-
-	var wg sync.WaitGroup
+	wg := sync.WaitGroup{}
 	wg.Add(2)
 	go o.WsHandleData(&wg)
 	go o.wsPingHandler(&wg)
-
-	err = o.WsSubscribeToDefaults()
-	if err != nil {
-		return fmt.Errorf("error: Could not subscribe to the OKEX websocket %s",
-			err)
-	}
+	o.GenerateDefaultSubscriptions()
 
 	// Ensures that we start the routines and we dont race when shutdown occurs
 	wg.Wait()
 	return nil
 }
 
-// WsSubscribeToDefaults subscribes to the websocket channels
-func (o *OKGroup) WsSubscribeToDefaults() (err error) {
-	channelsToSubscribe := []string{okGroupWsSpotDepth, okGroupWsSpotCandle300s, okGroupWsSpotTicker, okGroupWsSpotTrade}
-	for _, pair := range o.EnabledPairs {
-		formattedPair := strings.ToUpper(strings.Replace(pair.String(), "_", "-", 1))
-		for _, channel := range channelsToSubscribe {
-			err = o.WsSubscribeToChannel(fmt.Sprintf("%v:%s", channel, formattedPair))
-			if err != nil {
-				return
-			}
-		}
-	}
-
-	return nil
-}
-
 // WsReadData reads data from the websocket connection
 func (o *OKGroup) WsReadData() (exchange.WebsocketResponse, error) {
 	mType, resp, err := o.WebsocketConn.ReadMessage()
@@ -255,7 +237,7 @@ func (o *OKGroup) wsPingHandler(wg *sync.WaitGroup) {
 	o.Websocket.Wg.Add(1)
 	defer o.Websocket.Wg.Done()
 
-	ticker := time.NewTicker(time.Second * 27)
+	ticker := time.NewTicker(time.Second * 10)
 
 	wg.Done()
 
@@ -271,7 +253,6 @@ func (o *OKGroup) wsPingHandler(wg *sync.WaitGroup) {
 			}
 			if err != nil {
 				o.Websocket.DataHandler <- err
-				return
 			}
 		}
 	}
@@ -281,11 +262,6 @@ func (o *OKGroup) wsPingHandler(wg *sync.WaitGroup) {
 func (o *OKGroup) WsHandleData(wg *sync.WaitGroup) {
 	o.Websocket.Wg.Add(1)
 	defer func() {
-		err := o.WebsocketConn.Close()
-		if err != nil {
-			o.Websocket.DataHandler <- fmt.Errorf("okex_websocket.go - Unable to to close Websocket connection. Error: %s",
-				err)
-		}
 		o.Websocket.Wg.Done()
 	}()
 
@@ -295,11 +271,12 @@ func (o *OKGroup) WsHandleData(wg *sync.WaitGroup) {
 		select {
 		case <-o.Websocket.ShutdownC:
 			return
+
 		default:
 			resp, err := o.WsReadData()
 			if err != nil {
+				time.Sleep(time.Second)
 				o.Websocket.DataHandler <- err
-				return
 			}
 			var dataResponse WebsocketDataResponse
 			err = common.JSONDecode(resp.Raw, &dataResponse)
@@ -326,46 +303,10 @@ func (o *OKGroup) WsHandleData(wg *sync.WaitGroup) {
 				}
 				continue
 			}
-			o.Websocket.DataHandler <- fmt.Errorf("unrecognised response: %v", resp.Raw)
-			continue
 		}
 	}
 }
 
-// WsSubscribeToChannel sends a request to WS to subscribe to supplied channel
-func (o *OKGroup) WsSubscribeToChannel(topic string) error {
-	resp := WebsocketEventRequest{
-		Operation: "subscribe",
-		Arguments: []string{topic},
-	}
-	json, err := common.JSONEncode(resp)
-	if err != nil {
-		return err
-	}
-	err = o.writeToWebsocket(string(json))
-	if err != nil {
-		return err
-	}
-	return nil
-}
-
-// WsUnsubscribeToChannel sends a request to WS to unsubscribe to supplied channel
-func (o *OKGroup) WsUnsubscribeToChannel(topic string) error {
-	resp := WebsocketEventRequest{
-		Operation: "unsubscribe",
-		Arguments: []string{topic},
-	}
-	json, err := common.JSONEncode(resp)
-	if err != nil {
-		return err
-	}
-	err = o.writeToWebsocket(string(json))
-	if err != nil {
-		return err
-	}
-	return nil
-}
-
 // WsLogin sends a login request to websocket to enable access to authenticated endpoints
 func (o *OKGroup) WsLogin() error {
 	utcTime := time.Now().UTC()
@@ -440,9 +381,12 @@ func (o *OKGroup) WsHandleDataResponse(response *WebsocketDataResponse) {
 		orderbookMutex.Lock()
 		err := o.WsProcessOrderBook(response)
 		if err != nil {
-			log.Error(err)
-			subscriptionChannel := fmt.Sprintf("%v:%v", response.Table, response.Data[0].InstrumentID)
-			o.ResubscribeToChannel(subscriptionChannel)
+			pair := currency.NewPairDelimiter(response.Data[0].InstrumentID, "-")
+			channelToResubscribe := exchange.WebsocketChannelSubscription{
+				Channel:  response.Table,
+				Currency: pair,
+			}
+			o.Websocket.ResubscribeToChannel(channelToResubscribe)
 		}
 		orderbookMutex.Unlock()
 	case okGroupWsTicker:
@@ -460,42 +404,6 @@ func (o *OKGroup) WsHandleDataResponse(response *WebsocketDataResponse) {
 	}
 }
 
-// ResubscribeToChannel will attempt to unsubscribe and resubscribe to a channel
-func (o *OKGroup) ResubscribeToChannel(channel string) {
-	if okGroupWsResubscribeFailureLimit > 0 {
-		var successfulUnsubscribe bool
-		for i := 0; i < okGroupWsResubscribeFailureLimit; i++ {
-			err := o.WsUnsubscribeToChannel(channel)
-			if err != nil {
-				log.Error(err)
-				time.Sleep(okGroupWsResubscribeDelayInSeconds * time.Second)
-				continue
-			}
-			successfulUnsubscribe = true
-			break
-		}
-		if !successfulUnsubscribe {
-			log.Fatalf("%v websocket channel %v failed to unsubscribe after %v attempts", o.GetName(), channel, okGroupWsResubscribeFailureLimit)
-		}
-		successfulSubscribe := true
-		for i := 0; i < okGroupWsResubscribeFailureLimit; i++ {
-			err := o.WsSubscribeToChannel(channel)
-			if err != nil {
-				log.Error(err)
-				time.Sleep(okGroupWsResubscribeDelayInSeconds * time.Second)
-				continue
-			}
-			successfulSubscribe = true
-			break
-		}
-		if !successfulSubscribe {
-			log.Fatalf("%v websocket channel %v failed to resubscribe after %v attempts", o.GetName(), channel, okGroupWsResubscribeFailureLimit)
-		}
-	} else {
-		log.Fatalf("%v websocket channel %v cannot resubscribe. Limit: %v", o.GetName(), channel, okGroupWsResubscribeFailureLimit)
-	}
-}
-
 // logDataResponse will log the details of any websocket data event
 // where there is no websocket datahandler for it
 func logDataResponse(response *WebsocketDataResponse) {
@@ -767,3 +675,51 @@ func (o *OKGroup) CalculateUpdateOrderbookChecksum(orderbookData *orderbook.Base
 	checksum = strings.TrimSuffix(checksum, ":")
 	return int32(crc32.ChecksumIEEE([]byte(checksum)))
 }
+
+// GenerateDefaultSubscriptions Adds default subscriptions to websocket to be handled by ManageSubscriptions()
+func (o *OKGroup) GenerateDefaultSubscriptions() {
+	enabledCurrencies := o.GetEnabledCurrencies()
+	subscriptions := []exchange.WebsocketChannelSubscription{}
+	for i := range defaultSubscribedChannels {
+		for j := range enabledCurrencies {
+			enabledCurrencies[j].Delimiter = "-"
+			subscriptions = append(subscriptions, exchange.WebsocketChannelSubscription{
+				Channel:  defaultSubscribedChannels[i],
+				Currency: enabledCurrencies[j],
+			})
+		}
+	}
+	o.Websocket.SubscribeToChannels(subscriptions)
+}
+
+// Subscribe sends a websocket message to receive data from the channel
+func (o *OKGroup) Subscribe(channelToSubscribe exchange.WebsocketChannelSubscription) error {
+	resp := WebsocketEventRequest{
+		Operation: "subscribe",
+		Arguments: []string{fmt.Sprintf("%v:%v", channelToSubscribe.Channel, channelToSubscribe.Currency.String())},
+	}
+	json, err := common.JSONEncode(resp)
+	if err != nil {
+		if o.Verbose {
+			log.Debugf("%v subscribe error: %v", o.Name, err)
+		}
+		return err
+	}
+	return o.writeToWebsocket(string(json))
+}
+
+// Unsubscribe sends a websocket message to stop receiving data from the channel
+func (o *OKGroup) Unsubscribe(channelToSubscribe exchange.WebsocketChannelSubscription) error {
+	resp := WebsocketEventRequest{
+		Operation: "unsubscribe",
+		Arguments: []string{fmt.Sprintf("%v:%v", channelToSubscribe.Channel, channelToSubscribe.Currency.String())},
+	}
+	json, err := common.JSONEncode(resp)
+	if err != nil {
+		if o.Verbose {
+			log.Debugf("%v unsubscribe error: %v", o.Name, err)
+		}
+		return err
+	}
+	return o.writeToWebsocket(string(json))
+}

exchanges/okgroup/okgroup_wrapper.go

@@ -429,3 +429,17 @@ func (o *OKGroup) GetFeeByType(feeBuilder *exchange.FeeBuilder) (float64, error)
 func (o *OKGroup) GetWithdrawCapabilities() uint32 {
 	return o.GetWithdrawPermissions()
 }
+
+// SubscribeToWebsocketChannels appends to ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle subscribing
+func (o *OKGroup) SubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	o.Websocket.SubscribeToChannels(channels)
+	return nil
+}
+
+// UnsubscribeToWebsocketChannels removes from ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle unsubscribing
+func (o *OKGroup) UnsubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	o.Websocket.UnsubscribeToChannels(channels)
+	return nil
+}

exchanges/poloniex/poloniex.go

@@ -8,6 +8,7 @@ import (
 	"net/http"
 	"net/url"
 	"strconv"
+	"sync"
 	"time"
 
 	"github.com/gorilla/websocket"
@@ -62,6 +63,7 @@ const (
 type Poloniex struct {
 	exchange.Base
 	WebsocketConn *websocket.Conn
+	wsRequestMtx  sync.Mutex
 }
 
 // SetDefaults sets default settings for poloniex
@@ -89,7 +91,9 @@ func (p *Poloniex) SetDefaults() {
 	p.WebsocketInit()
 	p.Websocket.Functionality = exchange.WebsocketTradeDataSupported |
 		exchange.WebsocketOrderbookSupported |
-		exchange.WebsocketTickerSupported
+		exchange.WebsocketTickerSupported |
+		exchange.WebsocketSubscribeSupported |
+		exchange.WebsocketUnsubscribeSupported
 }
 
 // Setup sets user exchange configuration settings
@@ -130,8 +134,11 @@ func (p *Poloniex) Setup(exch *config.ExchangeConfig) {
 			log.Fatal(err)
 		}
 		err = p.WebsocketSetup(p.WsConnect,
+			p.Subscribe,
+			p.Unsubscribe,
 			exch.Name,
 			exch.Websocket,
+			exch.Verbose,
 			poloniexWebsocketAddress,
 			exch.WebsocketURL)
 		if err != nil {

exchanges/poloniex/poloniex_websocket.go

@@ -6,6 +6,7 @@ import (
 	"net/http"
 	"net/url"
 	"strconv"
+	"strings"
 	"time"
 
 	"github.com/gorilla/websocket"
@@ -65,41 +66,8 @@ func (p *Poloniex) WsConnect() error {
 	}
 
 	go p.WsHandleData()
+	p.GenerateDefaultSubscriptions()
 
-	return p.WsSubscribe()
-}
-
-// WsSubscribe subscribes to the websocket feeds
-func (p *Poloniex) WsSubscribe() error {
-	tickerJSON, err := common.JSONEncode(WsCommand{
-		Command: "subscribe",
-		Channel: wsTickerDataID})
-	if err != nil {
-		return err
-	}
-
-	err = p.WebsocketConn.WriteMessage(websocket.TextMessage, tickerJSON)
-	if err != nil {
-		return err
-	}
-
-	pairs := p.GetEnabledCurrencies()
-	for _, nextPair := range pairs {
-		fPair := exchange.FormatExchangeCurrency(p.GetName(), nextPair)
-
-		orderbookJSON, err := common.JSONEncode(WsCommand{
-			Command: "subscribe",
-			Channel: fPair.String(),
-		})
-		if err != nil {
-			return err
-		}
-
-		err = p.WebsocketConn.WriteMessage(websocket.TextMessage, orderbookJSON)
-		if err != nil {
-			return err
-		}
-	}
 	return nil
 }
 
@@ -129,11 +97,6 @@ func (p *Poloniex) WsHandleData() {
 	p.Websocket.Wg.Add(1)
 
 	defer func() {
-		err := p.WebsocketConn.Close()
-		if err != nil {
-			p.Websocket.DataHandler <- fmt.Errorf("poloniex_websocket.go - Unable to to close Websocket connection. Error: %s",
-				err)
-		}
 		p.Websocket.Wg.Done()
 	}()
 
@@ -155,119 +118,122 @@ func (p *Poloniex) WsHandleData() {
 				p.Websocket.DataHandler <- err
 				continue
 			}
-
-			data := result.([]interface{})
-			chanID := int(data[0].(float64))
-
-			if len(data) == 2 && chanID != wsHeartbeat {
-				if checkSubscriptionSuccess(data) {
-					if p.Verbose {
-						log.Debugf("poloniex websocket subscribed to channel successfully. %d", chanID)
+			switch data := result.(type) {
+			case map[string]interface{}:
+				// subscription error
+				p.Websocket.DataHandler <- errors.New(data["error"].(string))
+			case []interface{}:
+				chanID := int(data[0].(float64))
+				if len(data) == 2 && chanID != wsHeartbeat {
+					if checkSubscriptionSuccess(data) {
+						if p.Verbose {
+							log.Debugf("poloniex websocket subscribed to channel successfully. %d", chanID)
+						}
+					} else {
+						if p.Verbose {
+							log.Debugf("poloniex websocket subscription to channel failed. %d", chanID)
+						}
 					}
-				} else {
-					if p.Verbose {
-						log.Debugf("poloniex websocket subscription to channel failed. %d", chanID)
+					continue
+				}
+
+				switch chanID {
+				case wsAccountNotificationID:
+				case wsTickerDataID:
+					tickerData := data[2].([]interface{})
+					var t WsTicker
+					t.LastPrice, _ = strconv.ParseFloat(tickerData[1].(string), 64)
+					t.LowestAsk, _ = strconv.ParseFloat(tickerData[2].(string), 64)
+					t.HighestBid, _ = strconv.ParseFloat(tickerData[3].(string), 64)
+					t.PercentageChange, _ = strconv.ParseFloat(tickerData[4].(string), 64)
+					t.BaseCurrencyVolume24H, _ = strconv.ParseFloat(tickerData[5].(string), 64)
+					t.QuoteCurrencyVolume24H, _ = strconv.ParseFloat(tickerData[6].(string), 64)
+					isFrozen := false
+					if tickerData[7].(float64) == 1 {
+						isFrozen = true
 					}
-				}
-				continue
-			}
+					t.IsFrozen = isFrozen
+					t.HighestTradeIn24H, _ = strconv.ParseFloat(tickerData[8].(string), 64)
+					t.LowestTradePrice24H, _ = strconv.ParseFloat(tickerData[9].(string), 64)
 
-			switch chanID {
-			case wsAccountNotificationID:
-			case wsTickerDataID:
-				tickerData := data[2].([]interface{})
-				var t WsTicker
-				t.LastPrice, _ = strconv.ParseFloat(tickerData[1].(string), 64)
-				t.LowestAsk, _ = strconv.ParseFloat(tickerData[2].(string), 64)
-				t.HighestBid, _ = strconv.ParseFloat(tickerData[3].(string), 64)
-				t.PercentageChange, _ = strconv.ParseFloat(tickerData[4].(string), 64)
-				t.BaseCurrencyVolume24H, _ = strconv.ParseFloat(tickerData[5].(string), 64)
-				t.QuoteCurrencyVolume24H, _ = strconv.ParseFloat(tickerData[6].(string), 64)
-				isFrozen := false
-				if tickerData[7].(float64) == 1 {
-					isFrozen = true
-				}
-				t.IsFrozen = isFrozen
-				t.HighestTradeIn24H, _ = strconv.ParseFloat(tickerData[8].(string), 64)
-				t.LowestTradePrice24H, _ = strconv.ParseFloat(tickerData[9].(string), 64)
+					p.Websocket.DataHandler <- exchange.TickerData{
+						Timestamp: time.Now(),
+						Exchange:  p.GetName(),
+						AssetType: "SPOT",
+						LowPrice:  t.LowestAsk,
+						HighPrice: t.HighestBid,
+					}
+				case ws24HourExchangeVolumeID:
+				case wsHeartbeat:
+				default:
+					if len(data) > 2 {
+						subData := data[2].([]interface{})
 
-				p.Websocket.DataHandler <- exchange.TickerData{
-					Timestamp: time.Now(),
-					Exchange:  p.GetName(),
-					AssetType: "SPOT",
-					LowPrice:  t.LowestAsk,
-					HighPrice: t.HighestBid,
-				}
-			case ws24HourExchangeVolumeID:
-			case wsHeartbeat:
-			default:
-				if len(data) > 2 {
-					subData := data[2].([]interface{})
+						for x := range subData {
+							dataL2 := subData[x]
+							dataL3 := dataL2.([]interface{})
 
-					for x := range subData {
-						dataL2 := subData[x]
-						dataL3 := dataL2.([]interface{})
+							switch getWSDataType(dataL2) {
+							case "i":
+								dataL3map := dataL3[1].(map[string]interface{})
+								currencyPair, ok := dataL3map["currencyPair"].(string)
+								if !ok {
+									p.Websocket.DataHandler <- errors.New("poloniex.go error - could not find currency pair in map")
+									continue
+								}
 
-						switch getWSDataType(dataL2) {
-						case "i":
-							dataL3map := dataL3[1].(map[string]interface{})
-							currencyPair, ok := dataL3map["currencyPair"].(string)
-							if !ok {
-								p.Websocket.DataHandler <- errors.New("poloniex.go error - could not find currency pair in map")
-								continue
-							}
+								orderbookData, ok := dataL3map["orderBook"].([]interface{})
+								if !ok {
+									p.Websocket.DataHandler <- errors.New("poloniex.go error - could not find orderbook data in map")
+									continue
+								}
 
-							orderbookData, ok := dataL3map["orderBook"].([]interface{})
-							if !ok {
-								p.Websocket.DataHandler <- errors.New("poloniex.go error - could not find orderbook data in map")
-								continue
-							}
+								err := p.WsProcessOrderbookSnapshot(orderbookData, currencyPair)
+								if err != nil {
+									p.Websocket.DataHandler <- err
+									continue
+								}
 
-							err := p.WsProcessOrderbookSnapshot(orderbookData, currencyPair)
-							if err != nil {
-								p.Websocket.DataHandler <- err
-								continue
-							}
+								p.Websocket.DataHandler <- exchange.WebsocketOrderbookUpdate{
+									Exchange: p.GetName(),
+									Asset:    "SPOT",
+									Pair:     currency.NewPairFromString(currencyPair),
+								}
+							case "o":
+								currencyPair := CurrencyPairID[chanID]
+								err := p.WsProcessOrderbookUpdate(dataL3, currencyPair)
+								if err != nil {
+									p.Websocket.DataHandler <- err
+									continue
+								}
 
-							p.Websocket.DataHandler <- exchange.WebsocketOrderbookUpdate{
-								Exchange: p.GetName(),
-								Asset:    "SPOT",
-								Pair:     currency.NewPairFromString(currencyPair),
-							}
-						case "o":
-							currencyPair := CurrencyPairID[chanID]
-							err := p.WsProcessOrderbookUpdate(dataL3, currencyPair)
-							if err != nil {
-								p.Websocket.DataHandler <- err
-								continue
-							}
+								p.Websocket.DataHandler <- exchange.WebsocketOrderbookUpdate{
+									Exchange: p.GetName(),
+									Asset:    "SPOT",
+									Pair:     currency.NewPairFromString(currencyPair),
+								}
+							case "t":
+								currencyPair := CurrencyPairID[chanID]
+								var trade WsTrade
+								trade.Symbol = CurrencyPairID[chanID]
+								trade.TradeID, _ = strconv.ParseInt(dataL3[1].(string), 10, 64)
+								// 1 for buy 0 for sell
+								side := "buy"
+								if dataL3[2].(float64) != 1 {
+									side = "sell"
+								}
+								trade.Side = side
+								trade.Volume, _ = strconv.ParseFloat(dataL3[3].(string), 64)
+								trade.Price, _ = strconv.ParseFloat(dataL3[4].(string), 64)
+								trade.Timestamp = int64(dataL3[5].(float64))
 
-							p.Websocket.DataHandler <- exchange.WebsocketOrderbookUpdate{
-								Exchange: p.GetName(),
-								Asset:    "SPOT",
-								Pair:     currency.NewPairFromString(currencyPair),
-							}
-						case "t":
-							currencyPair := CurrencyPairID[chanID]
-							var trade WsTrade
-							trade.Symbol = CurrencyPairID[chanID]
-							trade.TradeID, _ = strconv.ParseInt(dataL3[1].(string), 10, 64)
-							// 1 for buy 0 for sell
-							side := "buy"
-							if dataL3[2].(float64) != 1 {
-								side = "sell"
-							}
-							trade.Side = side
-							trade.Volume, _ = strconv.ParseFloat(dataL3[3].(string), 64)
-							trade.Price, _ = strconv.ParseFloat(dataL3[4].(string), 64)
-							trade.Timestamp = int64(dataL3[5].(float64))
-
-							p.Websocket.DataHandler <- exchange.TradeData{
-								Timestamp:    time.Unix(trade.Timestamp, 0),
-								CurrencyPair: currency.NewPairFromString(currencyPair),
-								Side:         trade.Side,
-								Amount:       trade.Volume,
-								Price:        trade.Price,
+								p.Websocket.DataHandler <- exchange.TradeData{
+									Timestamp:    time.Unix(trade.Timestamp, 0),
+									CurrencyPair: currency.NewPairFromString(currencyPair),
+									Side:         trade.Side,
+									Amount:       trade.Volume,
+									Price:        trade.Price,
+								}
 							}
 						}
 					}
@@ -468,3 +434,62 @@ var CurrencyPairID = map[int]string{
 	226: "USDC_USDT",
 	225: "USDC_ETH",
 }
+
+// GenerateDefaultSubscriptions Adds default subscriptions to websocket to be handled by ManageSubscriptions()
+func (p *Poloniex) GenerateDefaultSubscriptions() {
+	subscriptions := []exchange.WebsocketChannelSubscription{}
+	// Tickerdata is its own channel
+	subscriptions = append(subscriptions, exchange.WebsocketChannelSubscription{
+		Channel: fmt.Sprintf("%v", wsTickerDataID),
+	})
+
+	enabledCurrencies := p.GetEnabledCurrencies()
+	for j := range enabledCurrencies {
+		enabledCurrencies[j].Delimiter = "_"
+		subscriptions = append(subscriptions, exchange.WebsocketChannelSubscription{
+			Channel:  "orderbook",
+			Currency: enabledCurrencies[j],
+		})
+	}
+	p.Websocket.SubscribeToChannels(subscriptions)
+}
+
+// Subscribe sends a websocket message to receive data from the channel
+func (p *Poloniex) Subscribe(channelToSubscribe exchange.WebsocketChannelSubscription) error {
+	subscriptionRequest := WsCommand{
+		Command: "subscribe",
+	}
+	if strings.EqualFold(fmt.Sprintf("%v", wsTickerDataID), channelToSubscribe.Channel) {
+		subscriptionRequest.Channel = wsTickerDataID
+	} else {
+		subscriptionRequest.Channel = channelToSubscribe.Currency.String()
+	}
+	return p.wsSend(subscriptionRequest)
+}
+
+// Unsubscribe sends a websocket message to stop receiving data from the channel
+func (p *Poloniex) Unsubscribe(channelToSubscribe exchange.WebsocketChannelSubscription) error {
+	unsubscriptionRequest := WsCommand{
+		Command: "unsubscribe",
+	}
+	if strings.EqualFold(fmt.Sprintf("%v", wsTickerDataID), channelToSubscribe.Channel) {
+		unsubscriptionRequest.Channel = wsTickerDataID
+	} else {
+		unsubscriptionRequest.Channel = channelToSubscribe.Currency.String()
+	}
+	return p.wsSend(unsubscriptionRequest)
+}
+
+// WsSend sends data to the websocket server
+func (p *Poloniex) wsSend(data interface{}) error {
+	p.wsRequestMtx.Lock()
+	defer p.wsRequestMtx.Unlock()
+	json, err := common.JSONEncode(data)
+	if err != nil {
+		return err
+	}
+	if p.Verbose {
+		log.Debugf("%v sending message to websocket %v", p.Name, data)
+	}
+	return p.WebsocketConn.WriteMessage(websocket.TextMessage, json)
+}

exchanges/poloniex/poloniex_wrapper.go

@@ -396,3 +396,17 @@ func (p *Poloniex) GetOrderHistory(getOrdersRequest *exchange.GetOrdersRequest)
 
 	return orders, nil
 }
+
+// SubscribeToWebsocketChannels appends to ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle subscribing
+func (p *Poloniex) SubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	p.Websocket.SubscribeToChannels(channels)
+	return nil
+}
+
+// UnsubscribeToWebsocketChannels removes from ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle unsubscribing
+func (p *Poloniex) UnsubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	p.Websocket.UnsubscribeToChannels(channels)
+	return nil
+}

exchanges/request/request.go

@@ -330,7 +330,6 @@ func (r *Requester) DoRequest(req *http.Request, path string, body io.Reader, re
 
 		if resp.StatusCode != 200 && resp.StatusCode != 201 && resp.StatusCode != 202 {
 			err = fmt.Errorf("unsuccessful HTTP status code: %d", resp.StatusCode)
-
 			if verbose {
 				err = fmt.Errorf("%s\n%s", err.Error(),
 					fmt.Sprintf("%s exchange raw response: %s", r.Name, string(contents)))

exchanges/yobit/yobit_wrapper.go

@@ -363,3 +363,15 @@ func (y *Yobit) GetOrderHistory(getOrdersRequest *exchange.GetOrdersRequest) ([]
 
 	return orders, nil
 }
+
+// SubscribeToWebsocketChannels appends to ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle subscribing
+func (y *Yobit) SubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	return common.ErrFunctionNotSupported
+}
+
+// UnsubscribeToWebsocketChannels removes from ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle unsubscribing
+func (y *Yobit) UnsubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	return common.ErrFunctionNotSupported
+}

exchanges/zb/zb.go

@@ -8,6 +8,7 @@ import (
 	"net/url"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/thrasher-/gocryptotrader/currency"
@@ -49,6 +50,7 @@ const (
 type ZB struct {
 	WebsocketConn *websocket.Conn
 	exchange.Base
+	wsRequestMtx sync.Mutex
 }
 
 // SetDefaults sets default values for the exchange
@@ -78,7 +80,8 @@ func (z *ZB) SetDefaults() {
 	z.WebsocketInit()
 	z.Websocket.Functionality = exchange.WebsocketTickerSupported |
 		exchange.WebsocketOrderbookSupported |
-		exchange.WebsocketTradeDataSupported
+		exchange.WebsocketTradeDataSupported |
+		exchange.WebsocketSubscribeSupported
 }
 
 // Setup sets user configuration
@@ -120,8 +123,11 @@ func (z *ZB) Setup(exch *config.ExchangeConfig) {
 			log.Fatal(err)
 		}
 		err = z.WebsocketSetup(z.WsConnect,
+			z.Subscribe,
+			nil,
 			exch.Name,
 			exch.Websocket,
+			exch.Verbose,
 			zbWebsocketAPI,
 			exch.WebsocketURL)
 		if err != nil {

exchanges/zb/zb_websocket.go

@@ -12,6 +12,7 @@ import (
 	"github.com/thrasher-/gocryptotrader/currency"
 	exchange "github.com/thrasher-/gocryptotrader/exchanges"
 	"github.com/thrasher-/gocryptotrader/exchanges/orderbook"
+	log "github.com/thrasher-/gocryptotrader/logger"
 )
 
 const (
@@ -42,75 +43,7 @@ func (z *ZB) WsConnect() error {
 	}
 
 	go z.WsHandleData()
-
-	return z.WsSubscribe()
-}
-
-// WsSubscribe subscribes to the full websocket suite on ZB exchange
-func (z *ZB) WsSubscribe() error {
-	markets := Subscription{
-		Event:   "addChannel",
-		Channel: "markets",
-	}
-
-	reqMarkets, err := common.JSONEncode(markets)
-	if err != nil {
-		return err
-	}
-
-	err = z.WebsocketConn.WriteMessage(websocket.TextMessage, reqMarkets)
-	if err != nil {
-		return err
-	}
-
-	for _, c := range z.GetEnabledCurrencies() {
-		cPair := c.Base.Lower().String() + c.Quote.Lower().String()
-
-		ticker := Subscription{
-			Event:   "addChannel",
-			Channel: fmt.Sprintf("%s_ticker", cPair),
-		}
-
-		reqTicker, err := common.JSONEncode(ticker)
-		if err != nil {
-			return err
-		}
-
-		err = z.WebsocketConn.WriteMessage(websocket.TextMessage, reqTicker)
-		if err != nil {
-			return err
-		}
-
-		depth := Subscription{
-			Event:   "addChannel",
-			Channel: fmt.Sprintf("%s_depth", cPair),
-		}
-
-		reqDepth, err := common.JSONEncode(depth)
-		if err != nil {
-			return err
-		}
-
-		err = z.WebsocketConn.WriteMessage(websocket.TextMessage, reqDepth)
-		if err != nil {
-			return err
-		}
-
-		trades := Subscription{
-			Event:   "addChannel",
-			Channel: fmt.Sprintf("%s_trades", cPair),
-		}
-
-		reqTrades, err := common.JSONEncode(trades)
-		if err != nil {
-			return err
-		}
-
-		err = z.WebsocketConn.WriteMessage(websocket.TextMessage, reqTrades)
-		if err != nil {
-			return err
-		}
-	}
+	z.GenerateDefaultSubscriptions()
 
 	return nil
 }
@@ -133,22 +66,18 @@ func (z *ZB) WsHandleData() {
 	z.Websocket.Wg.Add(1)
 
 	defer func() {
-		err := z.WebsocketConn.Close()
-		if err != nil {
-			z.Websocket.DataHandler <- fmt.Errorf("zb_websocket.go - Unable to to close Websocket connection. Error: %s",
-				err)
-		}
 		z.Websocket.Wg.Done()
 	}()
 
 	for {
 		select {
 		case <-z.Websocket.ShutdownC:
-
+			return
 		default:
 			resp, err := z.WsReadData()
 			if err != nil {
 				z.Websocket.DataHandler <- err
+				time.Sleep(time.Second)
 				continue
 			}
 
@@ -158,7 +87,6 @@ func (z *ZB) WsHandleData() {
 				z.Websocket.DataHandler <- err
 				continue
 			}
-
 			switch {
 			case common.StringContains(result.Channel, "markets"):
 				if !result.Success {
@@ -309,3 +237,47 @@ var wsErrCodes = map[int64]string{
 	4001: "API interface is locked",
 	4002: "Request too frequently",
 }
+
+// GenerateDefaultSubscriptions Adds default subscriptions to websocket to be handled by ManageSubscriptions()
+func (z *ZB) GenerateDefaultSubscriptions() {
+	subscriptions := []exchange.WebsocketChannelSubscription{}
+	// Tickerdata is its own channel
+	subscriptions = append(subscriptions, exchange.WebsocketChannelSubscription{
+		Channel: "markets",
+	})
+	channels := []string{"%s_ticker", "%s_depth", "%s_trades"}
+	enabledCurrencies := z.GetEnabledCurrencies()
+	for i := range channels {
+		for j := range enabledCurrencies {
+			enabledCurrencies[j].Delimiter = ""
+			subscriptions = append(subscriptions, exchange.WebsocketChannelSubscription{
+				Channel:  fmt.Sprintf(channels[i], enabledCurrencies[j].Lower().String()),
+				Currency: enabledCurrencies[j].Lower(),
+			})
+		}
+	}
+	z.Websocket.SubscribeToChannels(subscriptions)
+}
+
+// Subscribe sends a websocket message to receive data from the channel
+func (z *ZB) Subscribe(channelToSubscribe exchange.WebsocketChannelSubscription) error {
+	subscriptionRequest := Subscription{
+		Event:   "addChannel",
+		Channel: channelToSubscribe.Channel,
+	}
+	return z.wsSend(subscriptionRequest)
+}
+
+// WsSend sends data to the websocket server
+func (z *ZB) wsSend(data interface{}) error {
+	z.wsRequestMtx.Lock()
+	defer z.wsRequestMtx.Unlock()
+	json, err := common.JSONEncode(data)
+	if err != nil {
+		return err
+	}
+	if z.Verbose {
+		log.Debugf("%v sending message to websocket %v", z.Name, data)
+	}
+	return z.WebsocketConn.WriteMessage(websocket.TextMessage, json)
+}

exchanges/zb/zb_websocket_types.go

@@ -14,7 +14,7 @@ type Generic struct {
 	Success bool            `json:"success"`
 	Channel string          `json:"channel"`
 	Message string          `json:"message"`
-	No      int64           `json:"no"`
+	No      string          `json:"no"`
 	Data    json.RawMessage `json:"data"`
 }
 

exchanges/zb/zb_wrapper.go

@@ -413,3 +413,16 @@ func (z *ZB) GetOrderHistory(getOrdersRequest *exchange.GetOrdersRequest) ([]exc
 
 	return orders, nil
 }
+
+// SubscribeToWebsocketChannels appends to ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle subscribing
+func (z *ZB) SubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	z.Websocket.SubscribeToChannels(channels)
+	return nil
+}
+
+// UnsubscribeToWebsocketChannels removes from ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle unsubscribing
+func (z *ZB) UnsubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	return common.ErrFunctionNotSupported
+}

routines.go

@@ -178,8 +178,8 @@ func relayWebsocketEvent(result interface{}, event, assetType, exchangeName stri
 	}
 	err := BroadcastWebsocketMessage(evt)
 	if err != nil {
-		log.Errorf("Failed to broadcast websocket event. Error: %s",
-			err)
+		log.Errorf("Failed to broadcast websocket event %v. Error: %s",
+			event, err)
 	}
 }
 
@@ -405,7 +405,7 @@ func WebsocketDataHandler(ws *exchange.Websocket, verbose bool) {
 			case error:
 				switch {
 				case common.StringContains(d.Error(), "close 1006"):
-					go WebsocketReconnect(ws, verbose)
+					go ws.WebsocketReset()
 					continue
 				default:
 					log.Errorf("routines.go exchange %s websocket error - %s", ws.GetName(), data)
@@ -440,33 +440,3 @@ func WebsocketDataHandler(ws *exchange.Websocket, verbose bool) {
 		}
 	}
 }
-
-// WebsocketReconnect tries to reconnect to a websocket stream
-func WebsocketReconnect(ws *exchange.Websocket, verbose bool) {
-	if verbose {
-		log.Debugf("Websocket reconnection requested for %s", ws.GetName())
-	}
-
-	err := ws.Shutdown()
-	if err != nil {
-		log.Error(err)
-		return
-	}
-
-	wg.Add(1)
-	defer wg.Done()
-
-	tick := time.NewTicker(3 * time.Second)
-	for {
-		select {
-		case <-shutdowner:
-			return
-
-		case <-tick.C:
-			err = ws.Connect()
-			if err == nil {
-				return
-			}
-		}
-	}
-}

testdata/configtest.json

@@ -1085,7 +1085,7 @@
    "name": "OKCOIN International",
    "enabled": true,
    "verbose": false,
-   "websocket": false,
+   "websocket": true,
    "useSandbox": false,
    "restPollingDelay": 10,
    "httpTimeout": 15000000000,

tools/exchange_template/wrapper_file.tmpl

@@ -185,4 +185,18 @@ func ({{.Variable}} *{{.CapitalName}}) GetFeeByType(feeBuilder *exchange.FeeBuil
 	return 0, common.ErrNotYetImplemented
 }
 
+// SubscribeToWebsocketChannels appends to ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle subscribing
+func ({{.Variable}} *{{.CapitalName}}) SubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	{{.Variable}}.Websocket.SubscribeToChannels(channels)
+	return nil
+}
+
+// UnsubscribeToWebsocketChannels removes from ChannelsToSubscribe
+// which lets websocket.manageSubscriptions handle unsubscribing
+func ({{.Variable}} *{{.CapitalName}}) UnsubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+	{{.Variable}}.Websocket.UnubscribeToChannels(channels)
+	return nil
+}
+
 {{end}}