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

common/common.go

@@ -19,6 +19,7 @@ import (
 	"strconv"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"time"
 	"unicode"
 
@@ -672,3 +673,19 @@ func SortStrings[S ~[]E, E fmt.Stringer](x S) S {
 	})
 	return n
 }
+
+// Counter is a thread-safe counter.
+type Counter struct {
+	n int64 // privatised so you can't use counter as a value type
+}
+
+// IncrementAndGet returns the next count after incrementing.
+func (c *Counter) IncrementAndGet() int64 {
+	newID := atomic.AddInt64(&c.n, 1)
+	// Handle overflow by resetting the counter to 1 if it becomes negative
+	if newID < 0 {
+		atomic.StoreInt64(&c.n, 1)
+		return 1
+	}
+	return newID
+}

common/common_test.go

@@ -862,3 +862,18 @@ func (a A) String() string {
 func TestSortStrings(t *testing.T) {
 	assert.Equal(t, []A{1, 2, 5, 6}, SortStrings([]A{6, 2, 5, 1}))
 }
+
+func TestCounter(t *testing.T) {
+	t.Parallel()
+	c := Counter{n: -5}
+	require.Equal(t, int64(1), c.IncrementAndGet())
+	require.Equal(t, int64(2), c.IncrementAndGet())
+}
+
+// 683185328	         1.787 ns/op	       0 B/op	       0 allocs/op
+func BenchmarkCounter(b *testing.B) {
+	c := Counter{}
+	for i := 0; i < b.N; i++ {
+		c.IncrementAndGet()
+	}
+}

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
 }