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/gateio/gateio_websocket.go

@@ -1,11 +1,9 @@
 package gateio
 
 import (
-	"encoding/json"
 	"errors"
 	"fmt"
 	"net/http"
-	"net/url"
 	"strconv"
 	"strings"
 	"time"
@@ -15,40 +13,28 @@ import (
 	"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"
 )
 
 const (
 	gateioWebsocketEndpoint  = "wss://ws.gate.io/v3/"
 	gatioWsMethodPing        = "ping"
-	gateioWebsocketRateLimit = 120 * time.Millisecond
+	gateioWebsocketRateLimit = 120
 )
 
 // WsConnect initiates a websocket connection
 func (g *Gateio) WsConnect() error {
 	if !g.Websocket.IsEnabled() || !g.IsEnabled() {
-		return errors.New(exchange.WebsocketNotEnabled)
+		return errors.New(wshandler.WebsocketNotEnabled)
 	}
-
 	var dialer websocket.Dialer
-	if g.Websocket.GetProxyAddress() != "" {
-		proxy, err := url.Parse(g.Websocket.GetProxyAddress())
-		if err != nil {
-			return err
-		}
-
-		dialer.Proxy = http.ProxyURL(proxy)
-	}
-
-	var err error
-	g.WebsocketConn, _, err = dialer.Dial(g.Websocket.GetWebsocketURL(),
-		http.Header{})
+	err := g.WebsocketConn.Dial(&dialer, http.Header{})
 	if err != nil {
 		return err
 	}
 	go g.WsHandleData()
-
-	err = g.wsServerSignIn()
+	_, err = g.wsServerSignIn()
 	if err != nil {
 		log.Errorf("%v - authentication failed: %v", g.Name, err)
 	}
@@ -57,37 +43,33 @@ func (g *Gateio) WsConnect() error {
 	return nil
 }
 
-func (g *Gateio) wsServerSignIn() error {
+func (g *Gateio) wsServerSignIn() (*WebsocketAuthenticationResponse, error) {
 	if !g.GetAuthenticatedAPISupport(exchange.WebsocketAuthentication) {
-		return fmt.Errorf("%v AuthenticatedWebsocketAPISupport not enabled", g.Name)
+		return nil, fmt.Errorf("%v AuthenticatedWebsocketAPISupport not enabled", g.Name)
 	}
 	nonce := int(time.Now().Unix() * 1000)
 	sigTemp := g.GenerateSignature(strconv.Itoa(nonce))
 	signature := common.Base64Encode(sigTemp)
 	signinWsRequest := WebsocketRequest{
-		ID:     IDSignIn,
+		ID:     g.WebsocketConn.GenerateMessageID(true),
 		Method: "server.sign",
 		Params: []interface{}{g.APIKey, signature, nonce},
 	}
-	err := g.wsSend(signinWsRequest)
+	resp, err := g.WebsocketConn.SendMessageReturnResponse(signinWsRequest.ID, signinWsRequest)
 	if err != nil {
 		g.Websocket.SetCanUseAuthenticatedEndpoints(false)
-		return err
+		return nil, err
 	}
-	time.Sleep(time.Second * 2) // sleep to allow server to complete sign-on if further authenticated requests are sent prior to this they will fail
-	return nil
-}
-
-// WsReadData reads from the websocket connection and returns the websocket
-// response
-func (g *Gateio) WsReadData() (exchange.WebsocketResponse, error) {
-	_, resp, err := g.WebsocketConn.ReadMessage()
+	var response WebsocketAuthenticationResponse
+	err = common.JSONDecode(resp, &response)
 	if err != nil {
-		return exchange.WebsocketResponse{}, err
+		g.Websocket.SetCanUseAuthenticatedEndpoints(false)
+		return nil, err
 	}
-
-	g.Websocket.TrafficAlert <- struct{}{}
-	return exchange.WebsocketResponse{Raw: resp}, nil
+	if response.Result.Status == "success" {
+		g.Websocket.SetCanUseAuthenticatedEndpoints(true)
+	}
+	return &response, nil
 }
 
 // WsHandleData handles all the websocket data coming from the websocket
@@ -105,25 +87,28 @@ func (g *Gateio) WsHandleData() {
 			return
 
 		default:
-			resp, err := g.WsReadData()
+			resp, err := g.WebsocketConn.ReadMessage()
 			if err != nil {
 				g.Websocket.DataHandler <- err
-				// Read data error messages can overwhelm and panic the application
-				time.Sleep(time.Second)
-				continue
+				return
 			}
-
+			g.Websocket.TrafficAlert <- struct{}{}
 			var result WebsocketResponse
 			err = common.JSONDecode(resp.Raw, &result)
 			if err != nil {
 				g.Websocket.DataHandler <- err
 				continue
 			}
+
+			if result.ID > 0 {
+				g.WebsocketConn.AddResponseWithID(result.ID, resp.Raw)
+				continue
+			}
+
 			if result.Error.Code != 0 {
 				if common.StringContains(result.Error.Message, "authentication") {
 					g.Websocket.DataHandler <- fmt.Errorf("%v - authentication failed: %v", g.Name, err)
 					g.Websocket.SetCanUseAuthenticatedEndpoints(false)
-
 					continue
 				}
 				g.Websocket.DataHandler <- fmt.Errorf("%v error %s",
@@ -131,49 +116,6 @@ func (g *Gateio) WsHandleData() {
 				continue
 			}
 
-			switch result.ID {
-			case IDSignIn:
-				g.Websocket.SetCanUseAuthenticatedEndpoints(true)
-				g.Websocket.DataHandler <- string(result.Result)
-			case IDBalance:
-				var balance WebsocketBalance
-				var balanceInterface interface{}
-				err = json.Unmarshal(result.Result, &balanceInterface)
-				if err != nil {
-					g.Websocket.DataHandler <- err
-				}
-				var p WebsocketBalanceCurrency
-				switch x := balanceInterface.(type) {
-				case map[string]interface{}:
-					for xx := range x {
-						switch kk := x[xx].(type) {
-						case map[string]interface{}:
-							p = WebsocketBalanceCurrency{
-								Currency:  xx,
-								Available: kk["available"].(string),
-								Locked:    kk["freeze"].(string),
-							}
-							balance.Currency = append(balance.Currency, p)
-						default:
-							break
-						}
-					}
-				default:
-					break
-				}
-				g.Websocket.DataHandler <- balance
-			case IDOrderQuery:
-				var orderQuery WebSocketOrderQueryResult
-				err = common.JSONDecode(result.Result, &orderQuery)
-				if err != nil {
-					g.Websocket.DataHandler <- err
-					continue
-				}
-				g.Websocket.DataHandler <- orderQuery
-			default:
-				break
-			}
-
 			switch {
 			case common.StringContains(result.Method, "ticker"):
 				var ticker WebsocketTicker
@@ -190,7 +132,7 @@ func (g *Gateio) WsHandleData() {
 					continue
 				}
 
-				g.Websocket.DataHandler <- exchange.TickerData{
+				g.Websocket.DataHandler <- wshandler.TickerData{
 					Timestamp:  time.Now(),
 					Pair:       currency.NewPairFromString(c),
 					AssetType:  "SPOT",
@@ -218,7 +160,7 @@ func (g *Gateio) WsHandleData() {
 				}
 
 				for _, trade := range trades {
-					g.Websocket.DataHandler <- exchange.TradeData{
+					g.Websocket.DataHandler <- wshandler.TradeData{
 						Timestamp:    time.Now(),
 						CurrencyPair: currency.NewPairFromString(c),
 						AssetType:    "SPOT",
@@ -310,7 +252,7 @@ func (g *Gateio) WsHandleData() {
 					}
 				}
 
-				g.Websocket.DataHandler <- exchange.WebsocketOrderbookUpdate{
+				g.Websocket.DataHandler <- wshandler.WebsocketOrderbookUpdate{
 					Pair:     currency.NewPairFromString(c),
 					Asset:    "SPOT",
 					Exchange: g.GetName(),
@@ -330,7 +272,7 @@ func (g *Gateio) WsHandleData() {
 				low, _ := strconv.ParseFloat(data[4].(string), 64)
 				volume, _ := strconv.ParseFloat(data[5].(string), 64)
 
-				g.Websocket.DataHandler <- exchange.KlineData{
+				g.Websocket.DataHandler <- wshandler.KlineData{
 					Timestamp:  time.Now(),
 					Pair:       currency.NewPairFromString(data[7].(string)),
 					AssetType:  "SPOT",
@@ -352,11 +294,11 @@ func (g *Gateio) GenerateAuthenticatedSubscriptions() {
 		return
 	}
 	var channels = []string{"balance.subscribe", "order.subscribe"}
-	var subscriptions []exchange.WebsocketChannelSubscription
+	var subscriptions []wshandler.WebsocketChannelSubscription
 	enabledCurrencies := g.GetEnabledCurrencies()
 	for i := range channels {
 		for j := range enabledCurrencies {
-			subscriptions = append(subscriptions, exchange.WebsocketChannelSubscription{
+			subscriptions = append(subscriptions, wshandler.WebsocketChannelSubscription{
 				Channel:  channels[i],
 				Currency: enabledCurrencies[j],
 			})
@@ -368,7 +310,7 @@ func (g *Gateio) GenerateAuthenticatedSubscriptions() {
 // 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"}
-	var subscriptions []exchange.WebsocketChannelSubscription
+	var subscriptions []wshandler.WebsocketChannelSubscription
 	enabledCurrencies := g.GetEnabledCurrencies()
 	for i := range channels {
 		for j := range enabledCurrencies {
@@ -379,7 +321,7 @@ func (g *Gateio) GenerateDefaultSubscriptions() {
 			} else if strings.EqualFold(channels[i], "kline.subscribe") {
 				params["interval"] = 1800
 			}
-			subscriptions = append(subscriptions, exchange.WebsocketChannelSubscription{
+			subscriptions = append(subscriptions, wshandler.WebsocketChannelSubscription{
 				Channel:  channels[i],
 				Currency: enabledCurrencies[j],
 				Params:   params,
@@ -390,54 +332,84 @@ func (g *Gateio) GenerateDefaultSubscriptions() {
 }
 
 // Subscribe sends a websocket message to receive data from the channel
-func (g *Gateio) Subscribe(channelToSubscribe exchange.WebsocketChannelSubscription) error {
+func (g *Gateio) Subscribe(channelToSubscribe wshandler.WebsocketChannelSubscription) error {
 	params := []interface{}{channelToSubscribe.Currency.String()}
 	for _, paramValue := range channelToSubscribe.Params {
 		params = append(params, paramValue)
 	}
 
 	subscribe := WebsocketRequest{
-		ID:     IDGeneric,
+		ID:     g.WebsocketConn.GenerateMessageID(true),
 		Method: channelToSubscribe.Channel,
 		Params: params,
 	}
 
-	if strings.EqualFold(channelToSubscribe.Channel, "balance.subscribe") {
-		subscribe.ID = IDBalance
+	resp, err := g.WebsocketConn.SendMessageReturnResponse(subscribe.ID, subscribe)
+	if err != nil {
+		return err
 	}
-
-	return g.wsSend(subscribe)
+	var response WebsocketAuthenticationResponse
+	err = common.JSONDecode(resp, &response)
+	if err != nil {
+		return err
+	}
+	if response.Result.Status != "success" {
+		return fmt.Errorf("%v could not subscribe to %v", g.Name, channelToSubscribe.Channel)
+	}
+	return nil
 }
 
 // Unsubscribe sends a websocket message to stop receiving data from the channel
-func (g *Gateio) Unsubscribe(channelToSubscribe exchange.WebsocketChannelSubscription) error {
+func (g *Gateio) Unsubscribe(channelToSubscribe wshandler.WebsocketChannelSubscription) error {
 	unsbuscribeText := strings.Replace(channelToSubscribe.Channel, "subscribe", "unsubscribe", 1)
 	subscribe := WebsocketRequest{
-		ID:     IDGeneric,
+		ID:     g.WebsocketConn.GenerateMessageID(true),
 		Method: unsbuscribeText,
 		Params: []interface{}{channelToSubscribe.Currency.String(), 1800},
 	}
-	return g.wsSend(subscribe)
+	resp, err := g.WebsocketConn.SendMessageReturnResponse(subscribe.ID, subscribe)
+	if err != nil {
+		return err
+	}
+	var response WebsocketAuthenticationResponse
+	err = common.JSONDecode(resp, &response)
+	if err != nil {
+		return err
+	}
+	if response.Result.Status != "success" {
+		return fmt.Errorf("%v could not subscribe to %v", g.Name, channelToSubscribe.Channel)
+	}
+	return nil
 }
 
-func (g *Gateio) wsGetBalance() error {
+func (g *Gateio) wsGetBalance(currencies []string) (*WsGetBalanceResponse, error) {
 	if !g.Websocket.CanUseAuthenticatedEndpoints() {
-		return fmt.Errorf("%v not authorised to get balance", g.Name)
+		return nil, fmt.Errorf("%v not authorised to get balance", g.Name)
 	}
-	balanceWsRequest := WebsocketRequest{
-		ID:     IDBalance,
+	balanceWsRequest := wsGetBalanceRequest{
+		ID:     g.WebsocketConn.GenerateMessageID(false),
 		Method: "balance.query",
-		Params: []interface{}{},
+		Params: currencies,
 	}
-	return g.wsSend(balanceWsRequest)
+	resp, err := g.WebsocketConn.SendMessageReturnResponse(balanceWsRequest.ID, balanceWsRequest)
+	if err != nil {
+		return nil, err
+	}
+	var balance WsGetBalanceResponse
+	err = common.JSONDecode(resp, &balance)
+	if err != nil {
+		return &balance, err
+	}
+
+	return &balance, nil
 }
 
-func (g *Gateio) wsGetOrderInfo(market string, offset, limit int) error {
+func (g *Gateio) wsGetOrderInfo(market string, offset, limit int) (*WebSocketOrderQueryResult, error) {
 	if !g.Websocket.CanUseAuthenticatedEndpoints() {
-		return fmt.Errorf("%v not authorised to get order info", g.Name)
+		return nil, fmt.Errorf("%v not authorised to get order info", g.Name)
 	}
 	order := WebsocketRequest{
-		ID:     IDOrderQuery,
+		ID:     g.WebsocketConn.GenerateMessageID(true),
 		Method: "order.query",
 		Params: []interface{}{
 			market,
@@ -445,17 +417,14 @@ func (g *Gateio) wsGetOrderInfo(market string, offset, limit int) error {
 			limit,
 		},
 	}
-	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)
+	resp, err := g.WebsocketConn.SendMessageReturnResponse(order.ID, order)
+	if err != nil {
+		return nil, err
 	}
-	// Basic rate limiter
-	time.Sleep(gateioWebsocketRateLimit)
-	return g.WebsocketConn.WriteJSON(data)
+	var orderQuery WebSocketOrderQueryResult
+	err = common.JSONDecode(resp, &orderQuery)
+	if err != nil {
+		return &orderQuery, err
+	}
+	return &orderQuery, nil
 }