websocket/gateio: Support multi connection management and integrate with GateIO (#1580)

* gateio: Add multi asset websocket support WIP.

* meow

* Add tests and shenanigans

* integrate flushing and for enabling/disabling pairs from rpc shenanigans

* some changes

* linter: fixes strikes again.

* Change name ConnectionAssociation -> ConnectionCandidate for better clarity on purpose. Change connections map to point to candidate to track subscriptions for future dynamic connections holder and drop struct ConnectionDetails.

* Add subscription tests (state functional)

* glorious:nits + proxy handling

* Spelling

* linter: fixerino

* instead of nil, dont do nil.

* clean up nils

* cya nils

* don't need to set URL or check if its running

* stop ping handler routine leak

* * Fix bug where reader routine on error that is not a disconnection error but websocket frame error or anything really makes the reader routine return and then connection never cycles and the buffer gets filled.
* Handle reconnection via an errors.Is check which is simpler and in that scope allow for quick disconnect reconnect without waiting for connection cycle.
* Dial now uses code from DialContext but just calls context.Background()
* Don't allow reader to return on parse binary response error. Just output error and return a non nil response

* Allow rollback on connect on any error across all connections

* fix shadow jutsu

* glorious/gk: nitters - adds in ws mock server

* linter: fix

* fix deadlock on connection as the previous channel had no reader and would hang connection reader for eternity.

* gk: nits

* Leak issue and edge case

* gk: nits

* gk: drain brain

* glorious: nits

* Update exchanges/stream/websocket.go

Co-authored-by: Scott <gloriousCode@users.noreply.github.com>

* glorious: nits

* add tests

* linter: fix

* After merge

* Add error connection info

* Fix edge case where it does not reconnect made by an already closed connection

* stream coverage

* glorious: nits

* glorious: nits removed asset error handling in stream package

* linter: fix

* rm block

* Add basic readme

* fix asset enabled flush cycle for multi connection

* spella: fix

* linter: fix

* Add glorious suggestions, fix some race thing

* reinstate name before any routine gets spawned

* stop on error in mock tests

* glorious: nits

* glorious: nits found in CI build

* Add test for drain, bumped wait times as there seems to be something happening on macos CI builds, used context.WithTimeout because its instant.

* mutex across shutdown and connect for protection

* lint: fix

* test time withoffset, reinstate stop

* fix whoops

* const trafficCheckInterval; rm testmain

* y

* fix lint

* bump time check window

* stream: fix intermittant test failures while testing routines and remove code that is not needed.

* spells

* cant do what I did

* protect race due to routine.

* update testURL

* use mock websocket connection instead of test URL's

* linter: fix

* remove url because its throwing errors on CI builds

* connections drop all the time, don't need to worry about not being able to echo back ws data as it can be easily reviewed _test file side.

* remove another superfluous url thats not really set up for this

* spawn overwatch routine when there is no errors, inline checker instead of waiting for a time period, add sleep inline with echo handler as this is really quick and wanted to ensure that latency is handing correctly

* linter: fixerino uperino

* glorious: panix

* linter: things

* whoops

* defer lock and use functions that don't require locking in SetProxyAddress

* lint: fix

* thrasher: nits

---------

Co-authored-by: shazbert <ryan.oharareid@thrasher.io>
Co-authored-by: Scott <gloriousCode@users.noreply.github.com>

exchanges/stream/README.md

@@ -0,0 +1,137 @@
+# GoCryptoTrader Exchange Stream Package
+
+This package is part of the GoCryptoTrader project and is responsible for handling exchange streaming data.
+
+## Overview
+
+The `stream` package uses Gorilla Websocket and provides functionalities to connect to various cryptocurrency exchanges and handle real-time data streams.
+
+## Features
+
+- Handle real-time market data streams
+- Unified interface for managing data streams
+- Multi-connection management - a system that can be used to manage multiple connections to the same exchange
+- Connection monitoring - a system that can be used to monitor the health of the websocket connections. This can be used to check if the connection is still alive and if it is not, it will attempt to reconnect
+- Traffic monitoring - will reconnect if no message is sent for a period of time defined in your config
+- Subscription management - a system that can be used to manage subscriptions to various data streams
+- Rate limiting - a system that can be used to rate limit the number of requests sent to the exchange
+- Message ID generation - a system that can be used to generate message IDs for websocket requests
+- Websocket message response matching - can be used to match websocket responses to the requests that were sent
+
+## Usage
+
+### Default single websocket connection
+Here is a basic example of how to setup the `stream` package for websocket:
+
+```go
+package main
+
+import (
+    "github.com/thrasher-corp/gocryptotrader/exchanges/stream"
+    exchange "github.com/thrasher-corp/gocryptotrader/exchanges"
+    "github.com/thrasher-corp/gocryptotrader/exchanges/request"
+)
+
+type Exchange struct {
+    exchange.Base
+}
+
+// In the exchange wrapper this will set up the initial pointer field provided by exchange.Base
+func (e *Exchange) SetDefault() {
+    e.Websocket = stream.NewWebsocket()
+	e.WebsocketResponseMaxLimit = exchange.DefaultWebsocketResponseMaxLimit
+	e.WebsocketResponseCheckTimeout = exchange.DefaultWebsocketResponseCheckTimeout
+	e.WebsocketOrderbookBufferLimit = exchange.DefaultWebsocketOrderbookBufferLimit
+}
+
+// In the exchange wrapper this is the original setup pattern for the websocket services 
+func (e *Exchange) Setup(exch *config.Exchange) error {
+    // This sets up global connection, sub, unsub and generate subscriptions for each connection defined below.
+    if err := e.Websocket.Setup(&stream.WebsocketSetup{
+		ExchangeConfig:                         exch,
+		DefaultURL:                             connectionURLString,
+		RunningURL:                             connectionURLString,
+		Connector:                              e.WsConnect,
+		Subscriber:                             e.Subscribe,
+		Unsubscriber:                           e.Unsubscribe,
+		GenerateSubscriptions:                  e.GenerateDefaultSubscriptions,
+		Features:                               &e.Features.Supports.WebsocketCapabilities,
+		MaxWebsocketSubscriptionsPerConnection: 240,
+		OrderbookBufferConfig: buffer.Config{ Checksum: e.CalculateUpdateOrderbookChecksum },
+	}); err != nil {
+		return err
+	}
+
+    // This is a public websocket connection
+	if err := ok.Websocket.SetupNewConnection(&stream.ConnectionSetup{
+		URL:                  connectionURLString,
+		ResponseCheckTimeout: exch.WebsocketResponseCheckTimeout,
+		ResponseMaxLimit:     exchangeWebsocketResponseMaxLimit,
+		RateLimit:            request.NewRateLimitWithWeight(time.Second, 2, 1),
+	}); err != nil {
+		return err
+	}
+
+    // This is a private websocket connection 
+	return ok.Websocket.SetupNewConnection(&stream.ConnectionSetup{
+		URL:                  privateConnectionURLString,
+		ResponseCheckTimeout: exch.WebsocketResponseCheckTimeout,
+		ResponseMaxLimit:     exchangeWebsocketResponseMaxLimit,
+		Authenticated:        true,
+		RateLimit:            request.NewRateLimitWithWeight(time.Second, 2, 1),
+	})
+}
+```
+
+### Multiple websocket connections
+ The example below provides the now optional multi connection management system which allows for more connections
+ to be maintained and established based off URL, connections types, asset types etc.
+```go
+func (e *Exchange) Setup(exch *config.Exchange) error {
+    // This sets up global connection, sub, unsub and generate subscriptions for each connection defined below.
+    if err := e.Websocket.Setup(&stream.WebsocketSetup{
+		ExchangeConfig:               exch,
+		Features:                     &e.Features.Supports.WebsocketCapabilities,
+		FillsFeed:                    e.Features.Enabled.FillsFeed,
+		TradeFeed:                    e.Features.Enabled.TradeFeed,
+		UseMultiConnectionManagement: true,
+	})
+	if err != nil {
+		return err
+	}
+	// Spot connection
+	err = g.Websocket.SetupNewConnection(&stream.ConnectionSetup{
+		URL:                      connectionURLStringForSpot,
+		RateLimit:                request.NewWeightedRateLimitByDuration(gateioWebsocketRateLimit),
+		ResponseCheckTimeout:     exch.WebsocketResponseCheckTimeout,
+		ResponseMaxLimit:         exch.WebsocketResponseMaxLimit,
+        // Custom handlers for the specific connection:
+		Handler:                  e.WsHandleSpotData,
+		Subscriber:               e.SpotSubscribe,
+		Unsubscriber:             e.SpotUnsubscribe,
+		GenerateSubscriptions:    e.GenerateDefaultSubscriptionsSpot,
+		Connector:                e.WsConnectSpot,
+		BespokeGenerateMessageID: e.GenerateWebsocketMessageID,
+	})
+	if err != nil {
+		return err
+	}
+	// Futures connection - USDT margined
+	err = g.Websocket.SetupNewConnection(&stream.ConnectionSetup{
+		URL:                  connectionURLStringForSpotForFutures,
+		RateLimit:            request.NewWeightedRateLimitByDuration(gateioWebsocketRateLimit),
+		ResponseCheckTimeout: exch.WebsocketResponseCheckTimeout,
+		ResponseMaxLimit:     exch.WebsocketResponseMaxLimit,
+        // Custom handlers for the specific connection:
+		Handler: func(ctx context.Context, incoming []byte) error {	return e.WsHandleFuturesData(ctx, incoming, asset.Futures)	},
+		Subscriber:               e.FuturesSubscribe,
+		Unsubscriber:             e.FuturesUnsubscribe,
+		GenerateSubscriptions:    func() (subscription.List, error) { return e.GenerateFuturesDefaultSubscriptions(currency.USDT) },
+		Connector:                e.WsFuturesConnect,
+		BespokeGenerateMessageID: e.GenerateWebsocketMessageID,
+	})
+	if err != nil {
+		return err
+	}
+}
+```

exchanges/stream/stream_types.go

@@ -10,11 +10,13 @@ import (
 	"github.com/thrasher-corp/gocryptotrader/exchanges/asset"
 	"github.com/thrasher-corp/gocryptotrader/exchanges/order"
 	"github.com/thrasher-corp/gocryptotrader/exchanges/request"
+	"github.com/thrasher-corp/gocryptotrader/exchanges/subscription"
 )
 
 // Connection defines a streaming services connection
 type Connection interface {
 	Dial(*websocket.Dialer, http.Header) error
+	DialContext(context.Context, *websocket.Dialer, http.Header) error
 	ReadMessage() Response
 	SetupPingHandler(request.EndpointLimit, PingHandler)
 	// GenerateMessageID generates a message ID for the individual connection. If a bespoke function is set
@@ -46,15 +48,50 @@ type ConnectionSetup struct {
 	ResponseCheckTimeout    time.Duration
 	ResponseMaxLimit        time.Duration
 	RateLimit               *request.RateLimiterWithWeight
-	URL                     string
 	Authenticated           bool
 	ConnectionLevelReporter Reporter
+
+	// URL defines the websocket server URL to connect to
+	URL string
+	// Connector is the function that will be called to connect to the
+	// exchange's websocket server. This will be called once when the stream
+	// service is started. Any bespoke connection logic should be handled here.
+	Connector func(ctx context.Context, conn Connection) error
+	// GenerateSubscriptions is a function that will be called to generate a
+	// list of subscriptions to be made to the exchange's websocket server.
+	GenerateSubscriptions func() (subscription.List, error)
+	// Subscriber is a function that will be called to send subscription
+	// messages based on the exchange's websocket server requirements to
+	// subscribe to specific channels.
+	Subscriber func(ctx context.Context, conn Connection, sub subscription.List) error
+	// Unsubscriber is a function that will be called to send unsubscription
+	// messages based on the exchange's websocket server requirements to
+	// unsubscribe from specific channels. NOTE: IF THE FEATURE IS ENABLED.
+	Unsubscriber func(ctx context.Context, conn Connection, unsub subscription.List) error
+	// Handler defines the function that will be called when a message is
+	// received from the exchange's websocket server. This function should
+	// handle the incoming message and pass it to the appropriate data handler.
+	Handler func(ctx context.Context, incoming []byte) error
 	// BespokeGenerateMessageID is a function that returns a unique message ID.
 	// This is useful for when an exchange connection requires a unique or
 	// structured message ID for each message sent.
 	BespokeGenerateMessageID func(highPrecision bool) int64
 }
 
+// ConnectionWrapper contains the connection setup details to be used when
+// attempting a new connection. It also contains the subscriptions that are
+// associated with the specific connection.
+type ConnectionWrapper struct {
+	// Setup contains the connection setup details
+	Setup *ConnectionSetup
+	// Subscriptions contains the subscriptions that are associated with the
+	// specific connection(s)
+	Subscriptions *subscription.Store
+	// Connection contains the active connection based off the connection
+	// details above.
+	Connection Connection // TODO: Upgrade to slice of connections.
+}
+
 // PingHandler container for ping handler settings
 type PingHandler struct {
 	Websocket         bool

exchanges/stream/websocket_connection.go

@@ -24,12 +24,19 @@ import (
 )
 
 var (
+	// errConnectionFault is a connection fault error which alerts the system that a connection cycle needs to take place.
+	errConnectionFault         = errors.New("connection fault")
 	errWebsocketIsDisconnected = errors.New("websocket connection is disconnected")
 	errRateLimitNotFound       = errors.New("rate limit definition not found")
 )
 
 // Dial sets proxy urls and then connects to the websocket
 func (w *WebsocketConnection) Dial(dialer *websocket.Dialer, headers http.Header) error {
+	return w.DialContext(context.Background(), dialer, headers)
+}
+
+// DialContext sets proxy urls and then connects to the websocket
+func (w *WebsocketConnection) DialContext(ctx context.Context, dialer *websocket.Dialer, headers http.Header) error {
 	if w.ProxyURL != "" {
 		proxy, err := url.Parse(w.ProxyURL)
 		if err != nil {
@@ -40,15 +47,15 @@ func (w *WebsocketConnection) Dial(dialer *websocket.Dialer, headers http.Header
 
 	var err error
 	var conStatus *http.Response
-
-	w.Connection, conStatus, err = dialer.Dial(w.URL, headers)
+	w.Connection, conStatus, err = dialer.DialContext(ctx, w.URL, headers)
 	if err != nil {
 		if conStatus != nil {
+			_ = conStatus.Body.Close()
 			return fmt.Errorf("%s websocket connection: %v %v %v Error: %w", w.ExchangeName, w.URL, conStatus, conStatus.StatusCode, err)
 		}
 		return fmt.Errorf("%s websocket connection: %v Error: %w", w.ExchangeName, w.URL, err)
 	}
-	defer conStatus.Body.Close()
+	_ = conStatus.Body.Close()
 
 	if w.Verbose {
 		log.Infof(log.WebsocketMgr, "%v Websocket connected to %s\n", w.ExchangeName, w.URL)
@@ -131,18 +138,18 @@ func (w *WebsocketConnection) SetupPingHandler(epl request.EndpointLimit, handle
 		return
 	}
 	w.Wg.Add(1)
-	defer w.Wg.Done()
 	go func() {
+		defer w.Wg.Done()
 		ticker := time.NewTicker(handler.Delay)
 		for {
 			select {
-			case <-w.ShutdownC:
+			case <-w.shutdown:
 				ticker.Stop()
 				return
 			case <-ticker.C:
 				err := w.SendRawMessage(context.TODO(), epl, handler.MessageType, handler.Message)
 				if err != nil {
-					log.Errorf(log.WebsocketMgr, "%v websocket connection: ping handler failed to send message [%s]", w.ExchangeName, handler.Message)
+					log.Errorf(log.WebsocketMgr, "%v websocket connection: ping handler failed to send message [%s]: %v", w.ExchangeName, handler.Message, err)
 					return
 				}
 			}
@@ -168,23 +175,24 @@ func (w *WebsocketConnection) IsConnected() bool {
 func (w *WebsocketConnection) ReadMessage() Response {
 	mType, resp, err := w.Connection.ReadMessage()
 	if err != nil {
-		if IsDisconnectionError(err) {
-			if w.setConnectedStatus(false) {
-				// NOTE: When w.setConnectedStatus() returns true the underlying
-				// state was changed and this infers that the connection was
-				// externally closed and an error is reported else Shutdown()
-				// method on WebsocketConnection type has been called and can
-				// be skipped.
-				select {
-				case w.readMessageErrors <- err:
-				default:
-					// bypass if there is no receiver, as this stops it returning
-					// when shutdown is called.
-					log.Warnf(log.WebsocketMgr,
-						"%s failed to relay error: %v",
-						w.ExchangeName,
-						err)
-				}
+		// If any error occurs, a Response{Raw: nil, Type: 0} is returned, causing the
+		// reader routine to exit. This leaves the connection without an active reader,
+		// leading to potential buffer issue from the ongoing websocket writes.
+		// Such errors are passed to `w.readMessageErrors` when the connection is active.
+		// The `connectionMonitor` handles these errors by flushing the buffer, reconnecting,
+		// and resubscribing to the websocket to restore the connection.
+		if w.setConnectedStatus(false) {
+			// NOTE: When w.setConnectedStatus() returns true the underlying
+			// state was changed and this infers that the connection was
+			// externally closed and an error is reported else Shutdown()
+			// method on WebsocketConnection type has been called and can
+			// be skipped.
+			select {
+			case w.readMessageErrors <- fmt.Errorf("%w: %w", err, errConnectionFault):
+			default:
+				// bypass if there is no receiver, as this stops it returning
+				// when shutdown is called.
+				log.Warnf(log.WebsocketMgr, "%s failed to relay error: %v", w.ExchangeName, err)
 			}
 		}
 		return Response{}
@@ -203,7 +211,7 @@ func (w *WebsocketConnection) ReadMessage() Response {
 		standardMessage, err = w.parseBinaryResponse(resp)
 		if err != nil {
 			log.Errorf(log.WebsocketMgr, "%v %v: Parse binary response error: %v", w.ExchangeName, removeURLQueryString(w.URL), err)
-			return Response{}
+			return Response{Raw: []byte(``)} // Non-nil response to avoid the reader returning on this case.
 		}
 	}
 	if w.Verbose {
@@ -264,7 +272,9 @@ func (w *WebsocketConnection) Shutdown() error {
 		return nil
 	}
 	w.setConnectedStatus(false)
-	return w.Connection.UnderlyingConn().Close()
+	w.writeControl.Lock()
+	defer w.writeControl.Unlock()
+	return w.Connection.NetConn().Close()
 }
 
 // SetURL sets connection URL

exchanges/stream/websocket_types.go

@@ -38,8 +38,6 @@ type Websocket struct {
 	state                        atomic.Uint32
 	verbose                      bool
 	connectionMonitorRunning     atomic.Bool
-	trafficMonitorRunning        atomic.Bool
-	dataMonitorRunning           atomic.Bool
 	trafficTimeout               time.Duration
 	connectionMonitorDelay       time.Duration
 	proxyAddr                    string
@@ -51,6 +49,15 @@ type Websocket struct {
 	m                            sync.Mutex
 	connector                    func() error
 
+	// connectionManager stores all *potential* connections for the exchange, organised within ConnectionWrapper structs.
+	// Each ConnectionWrapper one connection (will be expanded soon) tailored for specific exchange functionalities or asset types. // TODO: Expand this to support multiple connections per ConnectionWrapper
+	// For example, separate connections can be used for Spot, Margin, and Futures trading. This structure is especially useful
+	// for exchanges that differentiate between trading pairs by using different connection endpoints or protocols for various asset classes.
+	// If an exchange does not require such differentiation, all connections may be managed under a single ConnectionWrapper.
+	connectionManager []ConnectionWrapper
+	// connections holds a look up table for all connections to their corresponding ConnectionWrapper and subscription holder
+	connections map[Connection]*ConnectionWrapper
+
 	subscriptions *subscription.Store
 
 	// Subscriber function for exchange specific subscribe implementation
@@ -60,6 +67,8 @@ type Websocket struct {
 	// GenerateSubs function for exchange specific generating subscriptions from Features.Subscriptions, Pairs and Assets
 	GenerateSubs func() (subscription.List, error)
 
+	useMultiConnectionManagement bool
+
 	DataHandler chan interface{}
 	ToRoutine   chan interface{}
 
@@ -117,6 +126,11 @@ type WebsocketSetup struct {
 	// Local orderbook buffer config values
 	OrderbookBufferConfig buffer.Config
 
+	// UseMultiConnectionManagement allows the connections to be managed by the
+	// connection manager. If false, this will default to the global fields
+	// provided in this struct.
+	UseMultiConnectionManagement bool
+
 	TradeFeed bool
 
 	// Fill data config values
@@ -155,7 +169,9 @@ type WebsocketConnection struct {
 	ProxyURL     string
 	Wg           *sync.WaitGroup
 	Connection   *websocket.Conn
-	ShutdownC    chan struct{}
+
+	// shutdown synchronises shutdown event across routines associated with this connection only e.g. ping handler
+	shutdown chan struct{}
 
 	Match             *Match
 	ResponseMaxLimit  time.Duration