common/gateio/stream: add thread-safe counter and overide default GenerateMessageID with connection specific implementation (#1615)

* add counter and update gateio

* Update exchanges/gateio/gateio.go

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

* thrasher: nits

* add test case

* linter: fix

* revert change

* thrasher nits

---------

Co-authored-by: Ryan O'Hara-Reid <ryan.oharareid@thrasher.io>
Co-authored-by: Scott <gloriousCode@users.noreply.github.com>

exchanges/gateio/gateio.go

@@ -174,6 +174,7 @@ var (
 // Gateio is the overarching type across this package
 type Gateio struct {
 	exchange.Base
+	Counter common.Counter
 }
 
 // ***************************************** SubAccounts ********************************

exchanges/gateio/gateio_test.go

@@ -3606,3 +3606,8 @@ func TestGetUnifiedAccount(t *testing.T) {
 	require.NoError(t, err)
 	require.NotEmpty(t, payload)
 }
+
+func TestGenerateWebsocketMessageID(t *testing.T) {
+	t.Parallel()
+	require.NotEmpty(t, g.GenerateWebsocketMessageID(false))
+}

exchanges/gateio/gateio_websocket.go

@@ -868,3 +868,9 @@ func (g *Gateio) listOfAssetsCurrencyPairEnabledFor(cp currency.Pair) map[asset.
 	}
 	return assetPairEnabled
 }
+
+// GenerateWebsocketMessageID generates a message ID for the individual
+// connection.
+func (g *Gateio) GenerateWebsocketMessageID(bool) int64 {
+	return g.Counter.IncrementAndGet()
+}

exchanges/gateio/gateio_wrapper.go

@@ -226,10 +226,11 @@ func (g *Gateio) Setup(exch *config.Exchange) error {
 		return err
 	}
 	return g.Websocket.SetupNewConnection(stream.ConnectionSetup{
-		URL:                  gateioWebsocketEndpoint,
-		RateLimit:            gateioWebsocketRateLimit,
-		ResponseCheckTimeout: exch.WebsocketResponseCheckTimeout,
-		ResponseMaxLimit:     exch.WebsocketResponseMaxLimit,
+		URL:                      gateioWebsocketEndpoint,
+		RateLimit:                gateioWebsocketRateLimit,
+		ResponseCheckTimeout:     exch.WebsocketResponseCheckTimeout,
+		ResponseMaxLimit:         exch.WebsocketResponseMaxLimit,
+		BespokeGenerateMessageID: g.GenerateWebsocketMessageID,
 	})
 }
 

exchanges/stream/stream_types.go

@@ -17,6 +17,10 @@ type Connection interface {
 	ReadMessage() Response
 	SendJSONMessage(any) error
 	SetupPingHandler(PingHandler)
+	// GenerateMessageID generates a message ID for the individual connection.
+	// If a bespoke function is set (by using SetupNewConnection) it will use
+	// that, otherwise it will use the defaultGenerateMessageID function defined
+	// in websocket_connection.go.
 	GenerateMessageID(highPrecision bool) int64
 	SendMessageReturnResponse(ctx context.Context, signature any, request any) ([]byte, error)
 	SendMessageReturnResponses(ctx context.Context, signature any, request any, expected int) ([][]byte, error)
@@ -41,6 +45,10 @@ type ConnectionSetup struct {
 	URL                     string
 	Authenticated           bool
 	ConnectionLevelReporter Reporter
+	// 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
 }
 
 // PingHandler container for ping handler settings

exchanges/stream/websocket.go

@@ -204,7 +204,13 @@ func (w *Websocket) SetupNewConnection(c ConnectionSetup) error {
 	if w == nil {
 		return fmt.Errorf("%w: %w", errConnSetup, errWebsocketIsNil)
 	}
-	if c == (ConnectionSetup{}) {
+
+	if c.ResponseCheckTimeout == 0 &&
+		c.ResponseMaxLimit == 0 &&
+		c.RateLimit == 0 &&
+		c.URL == "" &&
+		c.ConnectionLevelReporter == nil &&
+		c.BespokeGenerateMessageID == nil {
 		return fmt.Errorf("%w: %w", errConnSetup, errExchangeConfigEmpty)
 	}
 
@@ -234,18 +240,19 @@ func (w *Websocket) SetupNewConnection(c ConnectionSetup) error {
 	}
 
 	newConn := &WebsocketConnection{
-		ExchangeName:      w.exchangeName,
-		URL:               connectionURL,
-		ProxyURL:          w.GetProxyAddress(),
-		Verbose:           w.verbose,
-		ResponseMaxLimit:  c.ResponseMaxLimit,
-		Traffic:           w.TrafficAlert,
-		readMessageErrors: w.ReadMessageErrors,
-		ShutdownC:         w.ShutdownC,
-		Wg:                &w.Wg,
-		Match:             w.Match,
-		RateLimit:         c.RateLimit,
-		Reporter:          c.ConnectionLevelReporter,
+		ExchangeName:             w.exchangeName,
+		URL:                      connectionURL,
+		ProxyURL:                 w.GetProxyAddress(),
+		Verbose:                  w.verbose,
+		ResponseMaxLimit:         c.ResponseMaxLimit,
+		Traffic:                  w.TrafficAlert,
+		readMessageErrors:        w.ReadMessageErrors,
+		ShutdownC:                w.ShutdownC,
+		Wg:                       &w.Wg,
+		Match:                    w.Match,
+		RateLimit:                c.RateLimit,
+		Reporter:                 c.ConnectionLevelReporter,
+		bespokeGenerateMessageID: c.BespokeGenerateMessageID,
 	}
 
 	if c.Authenticated {

exchanges/stream/websocket_connection.go

@@ -229,8 +229,18 @@ func (w *WebsocketConnection) parseBinaryResponse(resp []byte) ([]byte, error) {
 	return standardMessage, reader.Close()
 }
 
-// GenerateMessageID Creates a random message ID
+// GenerateMessageID generates a message ID for the individual connection.
+// If a bespoke function is set (by using SetupNewConnection) it will use that,
+// otherwise it will use the defaultGenerateMessageID function.
 func (w *WebsocketConnection) GenerateMessageID(highPrec bool) int64 {
+	if w.bespokeGenerateMessageID != nil {
+		return w.bespokeGenerateMessageID(highPrec)
+	}
+	return w.defaultGenerateMessageID(highPrec)
+}
+
+// defaultGenerateMessageID generates the default message ID
+func (w *WebsocketConnection) defaultGenerateMessageID(highPrec bool) int64 {
 	var minValue int64 = 1e8
 	var maxValue int64 = 2e8
 	if highPrec {

exchanges/stream/websocket_test.go

@@ -911,6 +911,9 @@ func TestGenerateMessageID(t *testing.T) {
 		assert.NotContains(t, ids, id, "GenerateMessageID must not generate the same ID twice")
 		ids[i] = id
 	}
+
+	wc.bespokeGenerateMessageID = func(bool) int64 { return 42 }
+	assert.EqualValues(t, 42, wc.GenerateMessageID(true), "GenerateMessageID must use bespokeGenerateMessageID")
 }
 
 // BenchmarkGenerateMessageID-8   	 2850018	       408 ns/op	      56 B/op	       4 allocs/op

exchanges/stream/websocket_types.go

@@ -145,5 +145,10 @@ type WebsocketConnection struct {
 	Traffic           chan struct{}
 	readMessageErrors chan error
 
+	// bespokeGenerateMessageID is a function that returns a unique message ID
+	// defined externally. This is used for exchanges that require a unique
+	// message ID for each message sent.
+	bespokeGenerateMessageID func(highPrecision bool) int64
+
 	Reporter Reporter
 }