stream/match: allow a single connection to maintain its own match lookup for multi-connection

exchanges/gateio/gateio_websocket.go

@@ -107,15 +107,15 @@ func (g *Gateio) generateWsSignature(secret, event, channel string, t int64) (st
 }
 
 // WsHandleSpotData handles spot data
-func (g *Gateio) WsHandleSpotData(_ context.Context, respRaw []byte) error {
+func (g *Gateio) WsHandleSpotData(_ context.Context, conn stream.Connection, respRaw []byte) error {
 	var push WsResponse
 	err := json.Unmarshal(respRaw, &push)
 	if err != nil {
 		return err
 	}
 
 	if push.Event == subscribeEvent || push.Event == unsubscribeEvent {
-		if !g.Websocket.Match.IncomingWithData(push.ID, respRaw) {
+		if !conn.RouteIncomingWebsocketData(push.ID, respRaw) {
 			return fmt.Errorf("couldn't match subscription message with ID: %d", push.ID)
 		}
 		return nil

exchanges/gateio/gateio_wrapper.go

@@ -229,8 +229,8 @@ func (g *Gateio) Setup(exch *config.Exchange) error {
 		RateLimit:            request.NewWeightedRateLimitByDuration(gateioWebsocketRateLimit),
 		ResponseCheckTimeout: exch.WebsocketResponseCheckTimeout,
 		ResponseMaxLimit:     exch.WebsocketResponseMaxLimit,
-		Handler: func(ctx context.Context, incoming []byte) error {
-			return g.WsHandleFuturesData(ctx, incoming, asset.Futures)
+		Handler: func(ctx context.Context, conn stream.Connection, incoming []byte) error {
+			return g.WsHandleFuturesData(ctx, conn, incoming, asset.Futures)
 		},
 		Subscriber:               g.FuturesSubscribe,
 		Unsubscriber:             g.FuturesUnsubscribe,
@@ -248,8 +248,8 @@ func (g *Gateio) Setup(exch *config.Exchange) error {
 		RateLimit:            request.NewWeightedRateLimitByDuration(gateioWebsocketRateLimit),
 		ResponseCheckTimeout: exch.WebsocketResponseCheckTimeout,
 		ResponseMaxLimit:     exch.WebsocketResponseMaxLimit,
-		Handler: func(ctx context.Context, incoming []byte) error {
-			return g.WsHandleFuturesData(ctx, incoming, asset.Futures)
+		Handler: func(ctx context.Context, conn stream.Connection, incoming []byte) error {
+			return g.WsHandleFuturesData(ctx, conn, incoming, asset.Futures)
 		},
 		Subscriber:               g.FuturesSubscribe,
 		Unsubscriber:             g.FuturesUnsubscribe,
@@ -268,8 +268,8 @@ func (g *Gateio) Setup(exch *config.Exchange) error {
 		RateLimit:            request.NewWeightedRateLimitByDuration(gateioWebsocketRateLimit),
 		ResponseCheckTimeout: exch.WebsocketResponseCheckTimeout,
 		ResponseMaxLimit:     exch.WebsocketResponseMaxLimit,
-		Handler: func(ctx context.Context, incoming []byte) error {
-			return g.WsHandleFuturesData(ctx, incoming, asset.DeliveryFutures)
+		Handler: func(ctx context.Context, conn stream.Connection, incoming []byte) error {
+			return g.WsHandleFuturesData(ctx, conn, incoming, asset.DeliveryFutures)
 		},
 		Subscriber:               g.DeliveryFuturesSubscribe,
 		Unsubscriber:             g.DeliveryFuturesUnsubscribe,

exchanges/gateio/gateio_ws_futures.go

@@ -162,15 +162,15 @@ func (g *Gateio) FuturesUnsubscribe(ctx context.Context, conn stream.Connection,
 }
 
 // WsHandleFuturesData handles futures websocket data
-func (g *Gateio) WsHandleFuturesData(_ context.Context, respRaw []byte, a asset.Item) error {
+func (g *Gateio) WsHandleFuturesData(_ context.Context, conn stream.Connection, respRaw []byte, a asset.Item) error {
 	var push WsResponse
 	err := json.Unmarshal(respRaw, &push)
 	if err != nil {
 		return err
 	}
 
 	if push.Event == subscribeEvent || push.Event == unsubscribeEvent {
-		if !g.Websocket.Match.IncomingWithData(push.ID, respRaw) {
+		if !conn.RouteIncomingWebsocketData(push.ID, respRaw) {
 			return fmt.Errorf("couldn't match subscription message with ID: %d", push.ID)
 		}
 		return nil

exchanges/gateio/gateio_ws_option.go

@@ -293,15 +293,15 @@ func (g *Gateio) OptionsUnsubscribe(ctx context.Context, conn stream.Connection,
 }
 
 // WsHandleOptionsData handles options websocket data
-func (g *Gateio) WsHandleOptionsData(_ context.Context, respRaw []byte) error {
+func (g *Gateio) WsHandleOptionsData(_ context.Context, conn stream.Connection, respRaw []byte) error {
 	var push WsResponse
 	err := json.Unmarshal(respRaw, &push)
 	if err != nil {
 		return err
 	}
 
 	if push.Event == subscribeEvent || push.Event == unsubscribeEvent {
-		if !g.Websocket.Match.IncomingWithData(push.ID, respRaw) {
+		if !conn.RouteIncomingWebsocketData(push.ID, respRaw) {
 			return fmt.Errorf("couldn't match subscription message with ID: %d", push.ID)
 		}
 		return nil

exchanges/stream/stream_types.go

@@ -35,6 +35,10 @@ type Connection interface {
 	SetProxy(string)
 	GetURL() string
 	Shutdown() error
+
+	// RouteIncomingWebsocketData routes incoming websocket data to the correct handler.
+	// Returns true if a handler was found and data was passed to it.
+	RouteIncomingWebsocketData(signature any, incoming []byte) (matched bool)
 }
 
 // Response defines generalised data from the stream connection
@@ -71,7 +75,7 @@ type ConnectionSetup struct {
 	// 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
+	Handler func(ctx context.Context, conn Connection, 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.

exchanges/stream/websocket.go

@@ -288,6 +288,13 @@ func (w *Websocket) getConnectionFromSetup(c *ConnectionSetup) *WebsocketConnect
 	if c.URL != "" {
 		connectionURL = c.URL
 	}
+	match := w.Match
+	if w.useMultiConnectionManagement {
+		// If we are using multi connection management, we can decouple
+		// the match from the global match and have a match per connection.
+		match = NewMatch()
+	}
+
 	return &WebsocketConnection{
 		ExchangeName:             w.exchangeName,
 		URL:                      connectionURL,
@@ -298,7 +305,7 @@ func (w *Websocket) getConnectionFromSetup(c *ConnectionSetup) *WebsocketConnect
 		readMessageErrors:        w.ReadMessageErrors,
 		shutdown:                 w.ShutdownC,
 		Wg:                       &w.Wg,
-		Match:                    w.Match,
+		Match:                    match,
 		RateLimit:                c.RateLimit,
 		Reporter:                 c.ConnectionLevelReporter,
 		bespokeGenerateMessageID: c.BespokeGenerateMessageID,
@@ -1064,7 +1071,7 @@ func (w *Websocket) checkSubscriptions(conn Connection, subs subscription.List)
 		if s.State() == subscription.ResubscribingState {
 			continue
 		}
-		if found := w.subscriptions.Get(s); found != nil {
+		if found := subscriptionStore.Get(s); found != nil {
 			return fmt.Errorf("%w: %s", subscription.ErrDuplicate, s)
 		}
 	}
@@ -1073,14 +1080,14 @@ func (w *Websocket) checkSubscriptions(conn Connection, subs subscription.List)
 }
 
 // Reader reads and handles data from a specific connection
-func (w *Websocket) Reader(ctx context.Context, conn Connection, handler func(ctx context.Context, message []byte) error) {
+func (w *Websocket) Reader(ctx context.Context, conn Connection, handler func(ctx context.Context, conn Connection, message []byte) error) {
 	defer w.Wg.Done()
 	for {
 		resp := conn.ReadMessage()
 		if resp.Raw == nil {
 			return // Connection has been closed
 		}
-		if err := handler(ctx, resp.Raw); err != nil {
+		if err := handler(ctx, conn, resp.Raw); err != nil {
 			w.DataHandler <- fmt.Errorf("connection URL:[%v] error: %w", conn.GetURL(), err)
 		}
 	}

exchanges/stream/websocket_connection.go

@@ -358,3 +358,9 @@ func removeURLQueryString(url string) string {
 	}
 	return url
 }
+
+// RouteIncomingWebsocketData routes incoming websocket data to the correct handler.
+// Returns true if a handler was found and data was passed to it.
+func (w *WebsocketConnection) RouteIncomingWebsocketData(signature any, incoming []byte) (matched bool) {
+	return w.Match.IncomingWithData(signature, incoming)
+}

exchanges/stream/websocket_test.go

@@ -248,7 +248,7 @@ func TestConnectionMessageErrors(t *testing.T) {
 	err = ws.Connect()
 	require.ErrorIs(t, err, errWebsocketDataHandlerUnset)
 
-	ws.connectionManager[0].Setup.Handler = func(context.Context, []byte) error {
+	ws.connectionManager[0].Setup.Handler = func(context.Context, Connection, []byte) error {
 		return errDastardlyReason
 	}
 	err = ws.Connect()
@@ -266,7 +266,7 @@ func TestConnectionMessageErrors(t *testing.T) {
 	err = ws.Connect()
 	require.ErrorIs(t, err, errDastardlyReason)
 
-	ws.connectionManager[0].Setup.Handler = func(context.Context, []byte) error {
+	ws.connectionManager[0].Setup.Handler = func(context.Context, Connection, []byte) error {
 		return errDastardlyReason
 	}
 	err = ws.Connect()
@@ -458,7 +458,7 @@ func TestSubscribeUnsubscribe(t *testing.T) {
 		Unsubscriber: func(ctx context.Context, c Connection, s subscription.List) error {
 			return currySimpleUnsubConn(multi)(ctx, c, s)
 		},
-		Handler: func(context.Context, []byte) error { return nil },
+		Handler: func(context.Context, Connection, []byte) error { return nil },
 	}
 	require.NoError(t, multi.SetupNewConnection(amazingCandidate))
 
@@ -1130,7 +1130,7 @@ func TestFlushChannels(t *testing.T) {
 		Unsubscriber: func(ctx context.Context, c Connection, s subscription.List) error {
 			return currySimpleUnsubConn(w)(ctx, c, s)
 		},
-		Handler: func(context.Context, []byte) error { return nil },
+		Handler: func(context.Context, Connection, []byte) error { return nil },
 	}
 	require.NoError(t, w.SetupNewConnection(amazingCandidate))
 	require.NoError(t, w.FlushChannels(), "FlushChannels must not error")
@@ -1228,7 +1228,7 @@ func TestSetupNewConnection(t *testing.T) {
 	err = multi.SetupNewConnection(connSetup)
 	require.ErrorIs(t, err, errWebsocketDataHandlerUnset)
 
-	connSetup.Handler = func(context.Context, []byte) error { return nil }
+	connSetup.Handler = func(context.Context, Connection, []byte) error { return nil }
 	err = multi.SetupNewConnection(connSetup)
 	require.NoError(t, err)