Websocket request-response correlation (#328)

* Establishes new websocket functionality. Begins the creation of the websocket request

* Adding a wrapper over gorilla websocket connect,send,receive to handle ID messages. Doesn't work

* Successfully moved exchange_websocket into its own wshandler namespace. oof

* Sets up ZB to use a round trip WS request system

* Adds Kraken ID support to subscriptions. Renames duplicate func name UnsubscribeToChannels to RemoveSubscribedChannels. Adds some helper methods in the WebsocketConn to reduce duplicate code. Cleans up ZB implementation

* Fixes double locking which caused no websocket data to be read. Fixes requestid for kraken subscriptions

* Completes Huobi and Hadax implementation. Extends ZB error handling. Adds GZip support for reading messages

* Adds HitBTC support. Adds GetCurrencies, GetSymbols, GetTrades WS funcs. Adds super fun new parameter to GenerateMessageID for Unix and UnixNano

* Adds GateIO id support

* Adds Coinut support. Prevents nil reference error in constatus when there isnt one

* Standardises all Exchange websockets to use the wshandler websocket. Removes the wsRequestMtx as wshandler handles that now. Makes the Dialer a dialer, its not externally referenced that I can see.

* Fixes issue with coinut implementation. Updates bitmex currencies. Removes redundant log messages which are used to log messages

* Starts testing. Renames files

* Adds tests for websocket connection

* Reverts request.go change

* Linting everything

* Fixes rebase issue

* Final changes. Fixes variable names, removes log.Debug, removes lines, rearranges test types, removes order correlation websocket type

* Final final commit, fixing ZB issues.

* Adds traffic alerts where missed. Changes empty struct pointer addresses to nil instead. Removes empty lines

* Fixed string conversion

* Fixes issue with ZB not sending success codes

* Fixes issue with coinut processing due to nonce handling with subscriptions

* Fixes issue where ZB test failure was not caught. Removes unnecessary error handling from other ZB tests

* Removes unused interface

* Renames wshandler.Init() to wshandler.Run()

* Updates template file

* Capitalises cryptocurrencies in struct. Moves websocketResponseCheckTimeout and websocketResponseMaxLimit into config options. Moves connection configuration to main exchange Setup (where appropriate). Reverts currencylastupdated ticks. Improves reader close error checking

* Fixes two inconsistent websocket delay times

* Creates a default variable for websocket ResponseMaxLimit and ResponseCheckTimeout, then applies it to setdefaults and all tests

* Updates exchange template to set and use default websocket response limits

exchanges/okgroup/okgroup.go

@@ -10,14 +10,13 @@ import (
 	"reflect"
 	"strconv"
 	"strings"
-	"sync"
 	"time"
 
 	"github.com/google/go-querystring/query"
-	"github.com/gorilla/websocket"
 	"github.com/thrasher-/gocryptotrader/common"
 	"github.com/thrasher-/gocryptotrader/config"
 	exchange "github.com/thrasher-/gocryptotrader/exchanges"
+	"github.com/thrasher-/gocryptotrader/exchanges/wshandler"
 	log "github.com/thrasher-/gocryptotrader/logger"
 )
 
@@ -87,8 +86,7 @@ var errMissValue = errors.New("warning - resp value is missing from exchange")
 type OKGroup struct {
 	exchange.Base
 	ExchangeName  string
-	WebsocketConn *websocket.Conn
-	wsRequestMtx  sync.Mutex
+	WebsocketConn *wshandler.WebsocketConnection
 	// 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,17 +139,27 @@ func (o *OKGroup) Setup(exch *config.ExchangeConfig) {
 		if err != nil {
 			log.Fatal(err)
 		}
-		err = o.WebsocketSetup(o.WsConnect,
+		err = o.Websocket.Setup(o.WsConnect,
 			o.Subscribe,
 			o.Unsubscribe,
 			exch.Name,
 			exch.Websocket,
 			exch.Verbose,
 			o.WebsocketURL,
-			exch.WebsocketURL)
+			exch.WebsocketURL,
+			exch.AuthenticatedWebsocketAPISupport)
 		if err != nil {
 			log.Fatal(err)
 		}
+		o.WebsocketConn = &wshandler.WebsocketConnection{
+			ExchangeName:         o.Name,
+			URL:                  o.Websocket.GetWebsocketURL(),
+			ProxyURL:             o.Websocket.GetProxyAddress(),
+			Verbose:              o.Verbose,
+			RateLimit:            okGroupWsRateLimit,
+			ResponseCheckTimeout: exch.WebsocketResponseCheckTimeout,
+			ResponseMaxLimit:     exch.WebsocketResponseMaxLimit,
+		}
 	}
 }
 

exchanges/okgroup/okgroup_websocket.go

@@ -1,34 +1,27 @@
 package okgroup
 
 import (
-	"bytes"
-	"compress/flate"
 	"errors"
 	"fmt"
 	"hash/crc32"
-	"io/ioutil"
 	"net/http"
-	"net/url"
 	"sort"
 	"strconv"
 	"strings"
 	"sync"
 	"time"
 
-	"github.com/thrasher-/gocryptotrader/currency"
-
 	"github.com/gorilla/websocket"
 	"github.com/thrasher-/gocryptotrader/common"
+	"github.com/thrasher-/gocryptotrader/currency"
 	exchange "github.com/thrasher-/gocryptotrader/exchanges"
 	"github.com/thrasher-/gocryptotrader/exchanges/orderbook"
+	"github.com/thrasher-/gocryptotrader/exchanges/wshandler"
 	log "github.com/thrasher-/gocryptotrader/logger"
 )
 
 // List of all websocket channels to subscribe to
 const (
-	// If a checksum fails, then resubscribing to the channel fails, fatal after these attempts
-	okGroupWsResubscribeFailureLimit   = 3
-	okGroupWsResubscribeDelayInSeconds = 3
 	// Orderbook events
 	okGroupWsOrderbookUpdate  = "update"
 	okGroupWsOrderbookPartial = "partial"
@@ -144,50 +137,22 @@ const (
 	okGroupWsFuturesPosition       = okGroupWsFuturesSubsection + okGroupWsPosition
 	okGroupWsFuturesOrder          = okGroupWsFuturesSubsection + okGroupWsOrder
 
-	okGroupWsRateLimit = 30 * time.Millisecond
+	okGroupWsRateLimit = 30
 )
 
 // 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.wsRequestMtx.Lock()
-	defer o.wsRequestMtx.Unlock()
-	if o.Verbose {
-		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))
-}
-
 // WsConnect initiates a websocket connection
 func (o *OKGroup) WsConnect() error {
 	if !o.Websocket.IsEnabled() || !o.IsEnabled() {
-		return errors.New(exchange.WebsocketNotEnabled)
+		return errors.New(wshandler.WebsocketNotEnabled)
 	}
-
 	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)
-	}
-
-	var err error
-	if o.Verbose {
-		log.Debugf("Attempting to connect to %v", o.Websocket.GetWebsocketURL())
-	}
-	o.WebsocketConn, _, err = dialer.Dial(o.Websocket.GetWebsocketURL(),
-		http.Header{})
+	err := o.WebsocketConn.Dial(&dialer, http.Header{})
 	if err != nil {
-		return fmt.Errorf("%s Unable to connect to Websocket. Error: %s",
-			o.Name,
-			err)
+		return err
 	}
 	if o.Verbose {
 		log.Debugf("Successful connection to %v",
@@ -210,34 +175,6 @@ func (o *OKGroup) WsConnect() error {
 	return nil
 }
 
-// WsReadData reads data from the websocket connection
-func (o *OKGroup) WsReadData() (exchange.WebsocketResponse, error) {
-	mType, resp, err := o.WebsocketConn.ReadMessage()
-	if err != nil {
-		return exchange.WebsocketResponse{}, err
-	}
-
-	o.Websocket.TrafficAlert <- struct{}{}
-	var standardMessage []byte
-	switch mType {
-	case websocket.TextMessage:
-		standardMessage = resp
-
-	case websocket.BinaryMessage:
-		reader := flate.NewReader(bytes.NewReader(resp))
-		standardMessage, err = ioutil.ReadAll(reader)
-		reader.Close()
-		if err != nil {
-			return exchange.WebsocketResponse{}, err
-		}
-	}
-	if o.Verbose {
-		log.Debugf("%v Websocket message received: %v", o.Name, string(standardMessage))
-	}
-
-	return exchange.WebsocketResponse{Raw: standardMessage}, nil
-}
-
 // wsPingHandler sends a message "ping" every 27 to maintain the connection to the websocket
 func (o *OKGroup) wsPingHandler(wg *sync.WaitGroup) {
 	o.Websocket.Wg.Add(1)
@@ -254,7 +191,7 @@ func (o *OKGroup) wsPingHandler(wg *sync.WaitGroup) {
 			return
 
 		case <-ticker.C:
-			err := o.writeToWebsocket("ping")
+			err := o.WebsocketConn.SendMessage("ping")
 			if o.Verbose {
 				log.Debugf("%v sending ping", o.GetName())
 			}
@@ -280,11 +217,12 @@ func (o *OKGroup) WsHandleData(wg *sync.WaitGroup) {
 			return
 
 		default:
-			resp, err := o.WsReadData()
+			resp, err := o.WebsocketConn.ReadMessage()
 			if err != nil {
-				time.Sleep(time.Second)
 				o.Websocket.DataHandler <- err
+				return
 			}
+			o.Websocket.TrafficAlert <- struct{}{}
 			var dataResponse WebsocketDataResponse
 			err = common.JSONDecode(resp.Raw, &dataResponse)
 			if err == nil && dataResponse.Table != "" {
@@ -326,16 +264,11 @@ func (o *OKGroup) WsLogin() error {
 	signPath := "/users/self/verify"
 	hmac := common.GetHMAC(common.HashSHA256, []byte(fmt.Sprintf("%v", unixTime)+http.MethodGet+signPath), []byte(o.APISecret))
 	base64 := common.Base64Encode(hmac)
-	resp := WebsocketEventRequest{
+	request := WebsocketEventRequest{
 		Operation: "login",
 		Arguments: []string{o.APIKey, o.ClientID, fmt.Sprintf("%v", unixTime), base64},
 	}
-	json, err := common.JSONEncode(resp)
-	if err != nil {
-		o.Websocket.SetCanUseAuthenticatedEndpoints(false)
-		return err
-	}
-	err = o.writeToWebsocket(string(json))
+	err := o.WebsocketConn.SendMessage(request)
 	if err != nil {
 		o.Websocket.SetCanUseAuthenticatedEndpoints(false)
 		return err
@@ -384,20 +317,14 @@ func (o *OKGroup) WsHandleDataResponse(response *WebsocketDataResponse) {
 	case okGroupWsCandle60s, okGroupWsCandle180s, okGroupWsCandle300s, okGroupWsCandle900s,
 		okGroupWsCandle1800s, okGroupWsCandle3600s, okGroupWsCandle7200s, okGroupWsCandle14400s,
 		okGroupWsCandle21600s, okGroupWsCandle43200s, okGroupWsCandle86400s, okGroupWsCandle604900s:
-		if o.Verbose {
-			log.Debugf("%v Websocket candle data received", o.GetName())
-		}
 		o.wsProcessCandles(response)
 	case okGroupWsDepth, okGroupWsDepth5:
-		if o.Verbose {
-			log.Debugf("%v Websocket orderbook data received", o.GetName())
-		}
 		// Locking, orderbooks cannot be processed out of order
 		orderbookMutex.Lock()
 		err := o.WsProcessOrderBook(response)
 		if err != nil {
 			pair := currency.NewPairDelimiter(response.Data[0].InstrumentID, "-")
-			channelToResubscribe := exchange.WebsocketChannelSubscription{
+			channelToResubscribe := wshandler.WebsocketChannelSubscription{
 				Channel:  response.Table,
 				Currency: pair,
 			}
@@ -405,14 +332,8 @@ func (o *OKGroup) WsHandleDataResponse(response *WebsocketDataResponse) {
 		}
 		orderbookMutex.Unlock()
 	case okGroupWsTicker:
-		if o.Verbose {
-			log.Debugf("%v Websocket ticker data received", o.GetName())
-		}
 		o.wsProcessTickers(response)
 	case okGroupWsTrade:
-		if o.Verbose {
-			log.Debugf("%v Websocket trade data received", o.GetName())
-		}
 		o.wsProcessTrades(response)
 	default:
 		logDataResponse(response)
@@ -435,7 +356,7 @@ func logDataResponse(response *WebsocketDataResponse) {
 func (o *OKGroup) wsProcessTickers(response *WebsocketDataResponse) {
 	for i := range response.Data {
 		instrument := currency.NewPairDelimiter(response.Data[i].InstrumentID, "-")
-		o.Websocket.DataHandler <- exchange.TickerData{
+		o.Websocket.DataHandler <- wshandler.TickerData{
 			Timestamp:  response.Data[i].Timestamp,
 			Exchange:   o.GetName(),
 			AssetType:  o.GetAssetTypeFromTableName(response.Table),
@@ -451,7 +372,7 @@ func (o *OKGroup) wsProcessTickers(response *WebsocketDataResponse) {
 func (o *OKGroup) wsProcessTrades(response *WebsocketDataResponse) {
 	for i := range response.Data {
 		instrument := currency.NewPairDelimiter(response.Data[i].InstrumentID, "-")
-		o.Websocket.DataHandler <- exchange.TradeData{
+		o.Websocket.DataHandler <- wshandler.TradeData{
 			Amount:       response.Data[i].Size,
 			AssetType:    o.GetAssetTypeFromTableName(response.Table),
 			CurrencyPair: instrument,
@@ -480,7 +401,7 @@ func (o *OKGroup) wsProcessCandles(response *WebsocketDataResponse) {
 			candleInterval = response.Table[candleIndex+len(okGroupWsCandle) : secondIndex]
 		}
 
-		klineData := exchange.KlineData{
+		klineData := wshandler.KlineData{
 			AssetType: o.GetAssetTypeFromTableName(response.Table),
 			Pair:      instrument,
 			Exchange:  o.GetName(),
@@ -548,7 +469,7 @@ func (o *OKGroup) WsProcessPartialOrderBook(wsEventData *WebsocketDataWrapper, i
 	if err != nil {
 		return err
 	}
-	o.Websocket.DataHandler <- exchange.WebsocketOrderbookUpdate{
+	o.Websocket.DataHandler <- wshandler.WebsocketOrderbookUpdate{
 		Exchange: o.GetName(),
 		Asset:    o.GetAssetTypeFromTableName(tableName),
 		Pair:     instrument,
@@ -591,7 +512,7 @@ func (o *OKGroup) WsProcessUpdateOrderbook(wsEventData *WebsocketDataWrapper, in
 		if err != nil {
 			log.Error(err)
 		}
-		o.Websocket.DataHandler <- exchange.WebsocketOrderbookUpdate{
+		o.Websocket.DataHandler <- wshandler.WebsocketOrderbookUpdate{
 			Exchange: o.GetName(),
 			Asset:    o.GetAssetTypeFromTableName(tableName),
 			Pair:     instrument,
@@ -694,14 +615,14 @@ func (o *OKGroup) CalculateUpdateOrderbookChecksum(orderbookData *orderbook.Base
 // GenerateDefaultSubscriptions Adds default subscriptions to websocket to be handled by ManageSubscriptions()
 func (o *OKGroup) GenerateDefaultSubscriptions() {
 	enabledCurrencies := o.GetEnabledCurrencies()
-	var subscriptions []exchange.WebsocketChannelSubscription
+	var subscriptions []wshandler.WebsocketChannelSubscription
 	if o.GetAuthenticatedAPISupport(exchange.WebsocketAuthentication) {
 		defaultSubscribedChannels = append(defaultSubscribedChannels, okGroupWsSpotMarginAccount, okGroupWsSpotAccount, okGroupWsSpotOrder)
 	}
 	for i := range defaultSubscribedChannels {
 		for j := range enabledCurrencies {
 			enabledCurrencies[j].Delimiter = "-"
-			subscriptions = append(subscriptions, exchange.WebsocketChannelSubscription{
+			subscriptions = append(subscriptions, wshandler.WebsocketChannelSubscription{
 				Channel:  defaultSubscribedChannels[i],
 				Currency: enabledCurrencies[j],
 			})
@@ -711,37 +632,23 @@ func (o *OKGroup) GenerateDefaultSubscriptions() {
 }
 
 // Subscribe sends a websocket message to receive data from the channel
-func (o *OKGroup) Subscribe(channelToSubscribe exchange.WebsocketChannelSubscription) error {
-	resp := WebsocketEventRequest{
+func (o *OKGroup) Subscribe(channelToSubscribe wshandler.WebsocketChannelSubscription) error {
+	request := WebsocketEventRequest{
 		Operation: "subscribe",
 		Arguments: []string{fmt.Sprintf("%v:%v", channelToSubscribe.Channel, channelToSubscribe.Currency.String())},
 	}
 	if strings.EqualFold(channelToSubscribe.Channel, okGroupWsSpotAccount) {
-		resp.Arguments = []string{fmt.Sprintf("%v:%v", channelToSubscribe.Channel, channelToSubscribe.Currency.Base.String())}
+		request.Arguments = []string{fmt.Sprintf("%v:%v", channelToSubscribe.Channel, channelToSubscribe.Currency.Base.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))
+	return o.WebsocketConn.SendMessage(request)
 }
 
 // Unsubscribe sends a websocket message to stop receiving data from the channel
-func (o *OKGroup) Unsubscribe(channelToSubscribe exchange.WebsocketChannelSubscription) error {
-	resp := WebsocketEventRequest{
+func (o *OKGroup) Unsubscribe(channelToSubscribe wshandler.WebsocketChannelSubscription) error {
+	request := 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))
+	return o.WebsocketConn.SendMessage(request)
 }

exchanges/okgroup/okgroup_wrapper.go

@@ -11,6 +11,7 @@ import (
 	exchange "github.com/thrasher-/gocryptotrader/exchanges"
 	"github.com/thrasher-/gocryptotrader/exchanges/orderbook"
 	"github.com/thrasher-/gocryptotrader/exchanges/ticker"
+	"github.com/thrasher-/gocryptotrader/exchanges/wshandler"
 	log "github.com/thrasher-/gocryptotrader/logger"
 )
 
@@ -414,7 +415,7 @@ func (o *OKGroup) GetOrderHistory(getOrdersRequest *exchange.GetOrdersRequest) (
 }
 
 // GetWebsocket returns a pointer to the exchange websocket
-func (o *OKGroup) GetWebsocket() (*exchange.Websocket, error) {
+func (o *OKGroup) GetWebsocket() (*wshandler.Websocket, error) {
 	return o.Websocket, nil
 }
 
@@ -434,20 +435,20 @@ func (o *OKGroup) GetWithdrawCapabilities() uint32 {
 
 // SubscribeToWebsocketChannels appends to ChannelsToSubscribe
 // which lets websocket.manageSubscriptions handle subscribing
-func (o *OKGroup) SubscribeToWebsocketChannels(channels []exchange.WebsocketChannelSubscription) error {
+func (o *OKGroup) SubscribeToWebsocketChannels(channels []wshandler.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)
+func (o *OKGroup) UnsubscribeToWebsocketChannels(channels []wshandler.WebsocketChannelSubscription) error {
+	o.Websocket.RemoveSubscribedChannels(channels)
 	return nil
 }
 
 // GetSubscriptions returns a copied list of subscriptions
-func (o *OKGroup) GetSubscriptions() ([]exchange.WebsocketChannelSubscription, error) {
+func (o *OKGroup) GetSubscriptions() ([]wshandler.WebsocketChannelSubscription, error) {
 	return o.Websocket.GetSubscriptions(), nil
 }