Bitfinex: Websocket subscription improvements (#1353)

* Websockets: Add keys to websocket subscriptions * This switches all RO uses of the mutex to use a RLock method. * The mutex used for discrete field access has had scope drift from name 'connectionMutex' so rename to more appropriate fieldsMutex * The mutex used for Set/CanUseAuthEndpoints moves from the subscriptions endpoint to the fieldsMutex * Add GetSubscription by key * Expose stream.Matcher type * Bitfinex: Subscribe and Unsubscribe atomicly * Fix Auth failures ignored * This change makes it so that Subscribe and Unsubscribe wait for success ** Tells the DataHandler about errors ** Errors are returned to consumers * Subscribes concurrently to the channels * It also simplifies the chanId to stream mapping * Removes unable to locate chanID: %d errors which are just noise * Paves the way for unified channelSubscription id handling * Adds support for subId for Book subscriptions, which is more robust * Vastly simplifies what we need to test TestWsSubscribedResponse This test was working to ensure that the various fancy key parsing mechanisms all worked. Now that we use subId, we just need a thorough test of that * Expose Match.Set in order to capture websocket incoming data Can't see another way of doing this. Doesn't seem too bad * Allow tests to run with auth or WS These flags made it difficult to run the tests whilst working on websockets * Enable API auth and WS in testconfig This change minimises the changes requires for a full test run against live endpoints, so that new contributors have a clearer testing path. I cannot see any reason to turn WS off and Auth endpoints off when we're not going to run API tests without Creds being set, and we're not going to do live fire tests without canManipulateRealOrders * TestWsSubscribe and various fixes ** Enables the websocket for live non-authed integration tests by default ** Adds an integration test for subscriptions ** Changes the Ws tests to respect canManipulateRealOrders ** Uses WsConnect instead of setupWs; fixes seqNo config not sent for WS tests ** Allows api creds to live in config/testdata.json which might be less likely to accidentally commit, and less obtrusive * Bitfinex: Support period and timeframe for Candles * Fixes manual Subscribe() symbol or key formatting * Unifies handling of params for DefaultSubscriptions and manual subsrciptions * Bitfinex: Handle conf and info WS channel events * Bitfinex: Better tests for subscriptions * fixup! Websockets: Add keys to websocket subscriptions * fixup! Bitfinex: Subscribe and Unsubscribe atomicly * fixup! Websockets: Add keys to websocket subscriptions * Websockets: Add Pending subscription status Add a status tracker so that Sub/Unsub can prevent duplicates, and also fixes when first message comes before we have added the sub to the tracker * Websockets: Add State instead of pending This change allows more clarity about the current state and checks for specifically already Unsubing * Bitfinex: Fix first sub message maybe lost The only link we have between a sub req and the sub resp is the subID. And the only link we have between a sub message and the sub is the chanID. We can't derive a link using Pair or anything else. This meant that by sending the resp and its chanID down the IncomingData channel, we allowed the channel reader to maybe process the next message, the first message on the channel, before the runtime executed the switch back to subscribeToChan waiting on the chan. To fix this, we key initially on subId.(string), and then replace it with chanId.(int64) when we have it *inside* the wsHandleData so we know we've procedurally handled it before the next message. subscribeToChan is then free to remove the subId keyed Sub regardless of error or not If there's an error, we don't need to inline handling because there won't be any second update. Expands test coverage to make sure those subId keyed subscriptions are removed. * Websocket: Validate state in SetChanState * fixup! Bitfinex: Fix first sub message maybe lost * Websockets: Rename RemoveUnsuccessfulSubs Implementation doesn't imply Unsuccessful or need to. This change supports the registering of Pending subs * Bitfinex: Fix race in Tests
2026-06-02 07:26:53 +00:00 · 2023-11-02 02:10:43 +01:00
parent dbe41a7412
commit f9437dbd08
39 changed files with 1004 additions and 668 deletions
--- a/exchanges/stream/websocket.go
+++ b/exchanges/stream/websocket.go
@@ -5,7 +5,6 @@ import (
 	"fmt"
 	"net"
 	"net/url"
-	"strings"
 	"sync"
 	"time"

@@ -25,8 +24,16 @@ const (
 )

 var (
+	// ErrSubscriptionNotFound defines an error when a subscription is not found
+	ErrSubscriptionNotFound = errors.New("subscription not found")
+	// ErrSubscribedAlready defines an error when a channel is already subscribed
+	ErrSubscribedAlready = errors.New("duplicate subscription")
 	// ErrSubscriptionFailure defines an error when a subscription fails
 	ErrSubscriptionFailure = errors.New("subscription failure")
+	// ErrUnsubscribeFailure defines an error when a unsubscribe fails
+	ErrUnsubscribeFailure = errors.New("unsubscribe failure")
+	// ErrChannelInStateAlready defines an error when a subscription channel is already in a new state
+	ErrChannelInStateAlready = errors.New("channel already in state")
 	// ErrAlreadyDisabled is returned when you double-disable the websocket
 	ErrAlreadyDisabled = errors.New("websocket already disabled")
 	// ErrNotConnected defines an error when websocket is not connected
@@ -52,7 +59,8 @@ var (
 	errSubscriptionsExceedsLimit            = errors.New("subscriptions exceeds limit")
 	errInvalidMaxSubscriptions              = errors.New("max subscriptions cannot be less than 0")
 	errNoSubscriptionsSupplied              = errors.New("no subscriptions supplied")
-	errChannelSubscriptionAlreadySubscribed = errors.New("channel subscription already subscribed")
+	errChannelAlreadySubscribed             = errors.New("channel already subscribed")
+	errInvalidChannelState                  = errors.New("invalid Channel state")
 )

 var globalReporter Reporter
@@ -373,9 +381,9 @@ func (w *Websocket) connectionMonitor() error {
 	if w.checkAndSetMonitorRunning() {
 		return errAlreadyRunning
 	}
-	w.connectionMutex.RLock()
+	w.fieldMutex.RLock()
 	delay := w.connectionMonitorDelay
-	w.connectionMutex.RUnlock()
+	w.fieldMutex.RUnlock()

 	go func() {
 		timer := time.NewTimer(delay)
@@ -477,7 +485,7 @@ func (w *Websocket) Shutdown() error {

 	// flush any subscriptions from last connection if needed
 	w.subscriptionMutex.Lock()
-	w.subscriptions = nil
+	w.subscriptions = subscriptionMap{}
 	w.subscriptionMutex.Unlock()

 	close(w.ShutdownC)
@@ -537,7 +545,7 @@ func (w *Websocket) FlushChannels() error {
 		if len(newsubs) != 0 {
 			// Purge subscription list as there will be conflicts
 			w.subscriptionMutex.Lock()
-			w.subscriptions = nil
+			w.subscriptions = subscriptionMap{}
 			w.subscriptionMutex.Unlock()
 			return w.SubscribeToChannels(newsubs)
 		}
@@ -629,73 +637,73 @@ func (w *Websocket) trafficMonitor() {
 }

 func (w *Websocket) setConnectedStatus(b bool) {
-	w.connectionMutex.Lock()
+	w.fieldMutex.Lock()
 	w.connected = b
-	w.connectionMutex.Unlock()
+	w.fieldMutex.Unlock()
 }

 // IsConnected returns status of connection
 func (w *Websocket) IsConnected() bool {
-	w.connectionMutex.RLock()
-	defer w.connectionMutex.RUnlock()
+	w.fieldMutex.RLock()
+	defer w.fieldMutex.RUnlock()
 	return w.connected
 }

 func (w *Websocket) setConnectingStatus(b bool) {
-	w.connectionMutex.Lock()
+	w.fieldMutex.Lock()
 	w.connecting = b
-	w.connectionMutex.Unlock()
+	w.fieldMutex.Unlock()
 }

 // IsConnecting returns status of connecting
 func (w *Websocket) IsConnecting() bool {
-	w.connectionMutex.RLock()
-	defer w.connectionMutex.RUnlock()
+	w.fieldMutex.RLock()
+	defer w.fieldMutex.RUnlock()
 	return w.connecting
 }

 func (w *Websocket) setEnabled(b bool) {
-	w.connectionMutex.Lock()
+	w.fieldMutex.Lock()
 	w.enabled = b
-	w.connectionMutex.Unlock()
+	w.fieldMutex.Unlock()
 }

 // IsEnabled returns status of enabled
 func (w *Websocket) IsEnabled() bool {
-	w.connectionMutex.RLock()
-	defer w.connectionMutex.RUnlock()
+	w.fieldMutex.RLock()
+	defer w.fieldMutex.RUnlock()
 	return w.enabled
 }

 func (w *Websocket) setInit(b bool) {
-	w.connectionMutex.Lock()
+	w.fieldMutex.Lock()
 	w.Init = b
-	w.connectionMutex.Unlock()
+	w.fieldMutex.Unlock()
 }

 // IsInit returns status of init
 func (w *Websocket) IsInit() bool {
-	w.connectionMutex.RLock()
-	defer w.connectionMutex.RUnlock()
+	w.fieldMutex.RLock()
+	defer w.fieldMutex.RUnlock()
 	return w.Init
 }

 func (w *Websocket) setTrafficMonitorRunning(b bool) {
-	w.connectionMutex.Lock()
+	w.fieldMutex.Lock()
 	w.trafficMonitorRunning = b
-	w.connectionMutex.Unlock()
+	w.fieldMutex.Unlock()
 }

 // IsTrafficMonitorRunning returns status of the traffic monitor
 func (w *Websocket) IsTrafficMonitorRunning() bool {
-	w.connectionMutex.RLock()
-	defer w.connectionMutex.RUnlock()
+	w.fieldMutex.RLock()
+	defer w.fieldMutex.RUnlock()
 	return w.trafficMonitorRunning
 }

 func (w *Websocket) checkAndSetMonitorRunning() (alreadyRunning bool) {
-	w.connectionMutex.Lock()
-	defer w.connectionMutex.Unlock()
+	w.fieldMutex.Lock()
+	defer w.fieldMutex.Unlock()
 	if w.connectionMonitorRunning {
 		return true
 	}
@@ -704,28 +712,28 @@ func (w *Websocket) checkAndSetMonitorRunning() (alreadyRunning bool) {
 }

 func (w *Websocket) setConnectionMonitorRunning(b bool) {
-	w.connectionMutex.Lock()
+	w.fieldMutex.Lock()
 	w.connectionMonitorRunning = b
-	w.connectionMutex.Unlock()
+	w.fieldMutex.Unlock()
 }

 // IsConnectionMonitorRunning returns status of connection monitor
 func (w *Websocket) IsConnectionMonitorRunning() bool {
-	w.connectionMutex.RLock()
-	defer w.connectionMutex.RUnlock()
+	w.fieldMutex.RLock()
+	defer w.fieldMutex.RUnlock()
 	return w.connectionMonitorRunning
 }

 func (w *Websocket) setDataMonitorRunning(b bool) {
-	w.connectionMutex.Lock()
+	w.fieldMutex.Lock()
 	w.dataMonitorRunning = b
-	w.connectionMutex.Unlock()
+	w.fieldMutex.Unlock()
 }

 // IsDataMonitorRunning returns status of data monitor
 func (w *Websocket) IsDataMonitorRunning() bool {
-	w.connectionMutex.RLock()
-	defer w.connectionMutex.RUnlock()
+	w.fieldMutex.RLock()
+	defer w.fieldMutex.RUnlock()
 	return w.dataMonitorRunning
 }

@@ -862,52 +870,44 @@ func (w *Websocket) GetName() string {
 // GetChannelDifference finds the difference between the subscribed channels
 // and the new subscription list when pairs are disabled or enabled.
 func (w *Websocket) GetChannelDifference(genSubs []ChannelSubscription) (sub, unsub []ChannelSubscription) {
-	w.subscriptionMutex.Lock()
-	defer w.subscriptionMutex.Unlock()
+	w.subscriptionMutex.RLock()
+	unsubMap := make(map[any]ChannelSubscription, len(w.subscriptions))
+	for k, c := range w.subscriptions {
+		unsubMap[k] = *c
+	}
+	w.subscriptionMutex.RUnlock()

-oldsubs:
-	for x := range w.subscriptions {
-		for y := range genSubs {
-			if w.subscriptions[x].Equal(&genSubs[y]) {
-				continue oldsubs
-			}
+	for i := range genSubs {
+		key := genSubs[i].EnsureKeyed()
+		if _, ok := unsubMap[key]; ok {
+			delete(unsubMap, key) // If it's in both then we remove it from the unsubscribe list
+		} else {
+			sub = append(sub, genSubs[i]) // If it's in genSubs but not existing subs we want to subscribe
 		}
-		unsub = append(unsub, w.subscriptions[x])
 	}

-newsubs:
-	for x := range genSubs {
-		for y := range w.subscriptions {
-			if genSubs[x].Equal(&w.subscriptions[y]) {
-				continue newsubs
-			}
-		}
-		sub = append(sub, genSubs[x])
+	for _, c := range unsubMap {
+		unsub = append(unsub, c)
 	}
+
 	return
 }

 // UnsubscribeChannels unsubscribes from a websocket channel
 func (w *Websocket) UnsubscribeChannels(channels []ChannelSubscription) error {
 	if len(channels) == 0 {
-		return fmt.Errorf("%s websocket: channels not populated cannot remove",
-			w.exchangeName)
+		return fmt.Errorf("%s websocket: %w", w.exchangeName, errNoSubscriptionsSupplied)
 	}
-	w.subscriptionMutex.Lock()
+	w.subscriptionMutex.RLock()

-channels:
-	for x := range channels {
-		for y := range w.subscriptions {
-			if channels[x].Equal(&w.subscriptions[y]) {
-				continue channels
-			}
+	for i := range channels {
+		key := channels[i].EnsureKeyed()
+		if _, ok := w.subscriptions[key]; !ok {
+			w.subscriptionMutex.RUnlock()
+			return fmt.Errorf("%s websocket: %w: %+v", w.exchangeName, ErrSubscriptionNotFound, channels[i])
 		}
-		w.subscriptionMutex.Unlock()
-		return fmt.Errorf("%s websocket: subscription not found in list: %+v",
-			w.exchangeName,
-			channels[x])
 	}
-	w.subscriptionMutex.Unlock()
+	w.subscriptionMutex.RUnlock()
 	return w.Unsubscriber(channels)
 }

@@ -922,69 +922,138 @@ func (w *Websocket) ResubscribeToChannel(subscribedChannel *ChannelSubscription)

 // SubscribeToChannels appends supplied channels to channelsToSubscribe
 func (w *Websocket) SubscribeToChannels(channels []ChannelSubscription) error {
-	err := w.checkSubscriptions(channels)
-	if err != nil {
+	if err := w.checkSubscriptions(channels); err != nil {
 		return fmt.Errorf("%s websocket: %w", w.exchangeName, common.AppendError(ErrSubscriptionFailure, err))
 	}
-	err = w.Subscriber(channels)
-	if err != nil {
+	if err := w.Subscriber(channels); err != nil {
 		return fmt.Errorf("%s websocket: %w", w.exchangeName, common.AppendError(ErrSubscriptionFailure, err))
 	}
 	return nil
 }

+// AddSubscription adds a subscription to the subscription lists
+// Unlike AddSubscriptions this method will error if the subscription already exists
+func (w *Websocket) AddSubscription(c *ChannelSubscription) error {
+	w.subscriptionMutex.Lock()
+	defer w.subscriptionMutex.Unlock()
+	if w.subscriptions == nil {
+		w.subscriptions = subscriptionMap{}
+	}
+	key := c.EnsureKeyed()
+	if _, ok := w.subscriptions[key]; ok {
+		return ErrSubscribedAlready
+	}
+
+	n := *c // Fresh copy; we don't want to use the pointer we were given and allow encapsulation/locks to be bypassed
+	w.subscriptions[key] = &n
+
+	return nil
+}
+
+// SetSubscriptionState sets an existing subscription state
+// returns an error if the subscription is not found, or the new state is already set
+func (w *Websocket) SetSubscriptionState(c *ChannelSubscription, state ChannelState) error {
+	w.subscriptionMutex.Lock()
+	defer w.subscriptionMutex.Unlock()
+	if w.subscriptions == nil {
+		w.subscriptions = subscriptionMap{}
+	}
+	key := c.EnsureKeyed()
+	p, ok := w.subscriptions[key]
+	if !ok {
+		return ErrSubscriptionNotFound
+	}
+	if state == p.State {
+		return ErrChannelInStateAlready
+	}
+	if state > ChannelUnsubscribing {
+		return errInvalidChannelState
+	}
+	p.State = state
+	return nil
+}
+
 // AddSuccessfulSubscriptions adds subscriptions to the subscription lists that
 // has been successfully subscribed
 func (w *Websocket) AddSuccessfulSubscriptions(channels ...ChannelSubscription) {
-	w.subscriptionMutex.Lock()
-	w.subscriptions = append(w.subscriptions, channels...)
-	w.subscriptionMutex.Unlock()
-}
-
-// RemoveSuccessfulUnsubscriptions removes subscriptions from the subscription
-// list that has been successfulling unsubscribed
-func (w *Websocket) RemoveSuccessfulUnsubscriptions(channels ...ChannelSubscription) {
 	w.subscriptionMutex.Lock()
 	defer w.subscriptionMutex.Unlock()
-	for x := range channels {
-		for y := range w.subscriptions {
-			if channels[x].Equal(&w.subscriptions[y]) {
-				w.subscriptions[y] = w.subscriptions[len(w.subscriptions)-1]
-				w.subscriptions[len(w.subscriptions)-1] = ChannelSubscription{}
-				w.subscriptions = w.subscriptions[:len(w.subscriptions)-1]
-				break
-			}
-		}
+	if w.subscriptions == nil {
+		w.subscriptions = subscriptionMap{}
+	}
+	for _, cN := range channels {
+		c := cN // cN is an iteration var; Not safe to make a pointer to
+		key := c.EnsureKeyed()
+		c.State = ChannelSubscribed
+		w.subscriptions[key] = &c
 	}
 }

-// Equal two WebsocketChannelSubscription to determine equality
-func (w *ChannelSubscription) Equal(s *ChannelSubscription) bool {
-	return strings.EqualFold(w.Channel, s.Channel) &&
-		w.Currency.Equal(s.Currency)
-}
-
-// GetSubscriptions returns a copied list of subscriptions
-// and is a private member that cannot be manipulated
-func (w *Websocket) GetSubscriptions() []ChannelSubscription {
+// RemoveSubscriptions removes subscriptions from the subscription list
+func (w *Websocket) RemoveSubscriptions(channels ...ChannelSubscription) {
 	w.subscriptionMutex.Lock()
 	defer w.subscriptionMutex.Unlock()
-	return append(w.subscriptions[:0:0], w.subscriptions...)
+	if w.subscriptions == nil {
+		w.subscriptions = subscriptionMap{}
+	}
+	for i := range channels {
+		key := channels[i].EnsureKeyed()
+		delete(w.subscriptions, key)
+	}
+}
+
+// EnsureKeyed sets the default key on a channel if it doesn't have one
+// Returns key for convenience
+func (c *ChannelSubscription) EnsureKeyed() any {
+	if c.Key == nil {
+		c.Key = DefaultChannelKey{
+			Channel:  c.Channel,
+			Asset:    c.Asset,
+			Currency: c.Currency,
+		}
+	}
+	return c.Key
+}
+
+// GetSubscription returns a pointer to a copy of the subscription at the key provided
+// returns nil if no subscription is at that key or the key is nil
+func (w *Websocket) GetSubscription(key any) *ChannelSubscription {
+	if key == nil || w == nil || w.subscriptions == nil {
+		return nil
+	}
+	w.subscriptionMutex.RLock()
+	defer w.subscriptionMutex.RUnlock()
+	if s, ok := w.subscriptions[key]; ok {
+		c := *s
+		return &c
+	}
+	return nil
+}
+
+// GetSubscriptions returns a new slice of the subscriptions
+func (w *Websocket) GetSubscriptions() []ChannelSubscription {
+	w.subscriptionMutex.RLock()
+	defer w.subscriptionMutex.RUnlock()
+	subs := make([]ChannelSubscription, 0, len(w.subscriptions))
+	for _, c := range w.subscriptions {
+		subs = append(subs, *c)
+	}
+	return subs
 }

 // SetCanUseAuthenticatedEndpoints sets canUseAuthenticatedEndpoints val in
 // a thread safe manner
 func (w *Websocket) SetCanUseAuthenticatedEndpoints(val bool) {
-	w.subscriptionMutex.Lock()
-	defer w.subscriptionMutex.Unlock()
+	w.fieldMutex.Lock()
+	defer w.fieldMutex.Unlock()
 	w.canUseAuthenticatedEndpoints = val
 }

 // CanUseAuthenticatedEndpoints gets canUseAuthenticatedEndpoints val in
 // a thread safe manner
 func (w *Websocket) CanUseAuthenticatedEndpoints() bool {
-	w.subscriptionMutex.Lock()
-	defer w.subscriptionMutex.Unlock()
+	w.fieldMutex.RLock()
+	defer w.fieldMutex.RUnlock()
 	return w.canUseAuthenticatedEndpoints
 }

@@ -1018,8 +1087,8 @@ func (w *Websocket) checkSubscriptions(subs []ChannelSubscription) error {
 		return errNoSubscriptionsSupplied
 	}

-	w.subscriptionMutex.Lock()
-	defer w.subscriptionMutex.Unlock()
+	w.subscriptionMutex.RLock()
+	defer w.subscriptionMutex.RUnlock()

 	if w.MaxSubscriptionsPerConnection > 0 && len(w.subscriptions)+len(subs) > w.MaxSubscriptionsPerConnection {
 		return fmt.Errorf("%w: current subscriptions: %v, incoming subscriptions: %v, max subscriptions per connection: %v - please reduce enabled pairs",
@@ -1029,12 +1098,12 @@ func (w *Websocket) checkSubscriptions(subs []ChannelSubscription) error {
 			w.MaxSubscriptionsPerConnection)
 	}

-	for x := range subs {
-		for y := range w.subscriptions {
-			if subs[x].Equal(&w.subscriptions[y]) {
-				return fmt.Errorf("%w for %+v", errChannelSubscriptionAlreadySubscribed, subs[x])
-			}
+	for i := range subs {
+		key := subs[i].EnsureKeyed()
+		if _, ok := w.subscriptions[key]; ok {
+			return fmt.Errorf("%w for %+v", errChannelAlreadySubscribed, subs[i])
 		}
 	}
+
 	return nil
 }