refactor(filter): unsubscribe waitgroup, execute async, and guard against calling funtions while the protocol is not started

waku/v2/protocol/filter/client.go

@@ -34,6 +34,9 @@ var (
 )
 
 type WakuFilterLightNode struct {
+	sync.RWMutex
+	started bool
+
 	cancel        context.CancelFunc
 	ctx           context.Context
 	h             host.Host
@@ -56,6 +59,9 @@ type WakuFilterPushResult struct {
 	PeerID peer.ID
 }
 
+var errNotStarted = errors.New("not started")
+var errAlreadyStarted = errors.New("already started")
+
 // NewWakuFilterLightnode returns a new instance of Waku Filter struct setup according to the chosen parameter and options
 // Takes an optional peermanager if WakuFilterLightnode is being created along with WakuNode.
 // If using libp2p host, then pass peermanager as nil
@@ -78,12 +84,20 @@ func (wf *WakuFilterLightNode) SetHost(h host.Host) {
 }
 
 func (wf *WakuFilterLightNode) Start(ctx context.Context) error {
+	wf.Lock()
+	defer wf.Unlock()
+
+	if wf.started {
+		return errAlreadyStarted
+	}
+
 	wf.wg.Wait() // Wait for any goroutines to stop
 
 	ctx, cancel := context.WithCancel(ctx)
 	wf.cancel = cancel
 	wf.ctx = ctx
 	wf.subscriptions = NewSubscriptionMap(wf.log)
+	wf.started = true
 
 	wf.h.SetStreamHandlerMatch(FilterPushID_v20beta1, protocol.PrefixTextMatch(string(FilterPushID_v20beta1)), wf.onRequest(ctx))
 
@@ -94,18 +108,34 @@ func (wf *WakuFilterLightNode) Start(ctx context.Context) error {
 
 // Stop unmounts the filter protocol
 func (wf *WakuFilterLightNode) Stop() {
-	if wf.cancel == nil {
+	wf.Lock()
+	defer wf.Unlock()
+
+	if !wf.started {
 		return
 	}
 
 	wf.cancel()
 
 	wf.h.RemoveStreamHandler(FilterPushID_v20beta1)
 
-	_, _ = wf.UnsubscribeAll(wf.ctx)
+	res, err := wf.unsubscribeAll(wf.ctx)
+	if err != nil {
+		wf.log.Warn("unsubscribing from full nodes", zap.Error(err))
+	}
+
+	for r := range res {
+		if r.Err != nil {
+			wf.log.Warn("unsubscribing from full nodes", zap.Error(r.Err), logging.HostID("peerID", r.PeerID))
+		}
+
+	}
 
 	wf.subscriptions.Clear()
 
+	wf.started = false
+	wf.cancel = nil
+
 	wf.wg.Wait()
 }
 
@@ -206,6 +236,13 @@ func (wf *WakuFilterLightNode) request(ctx context.Context, params *FilterSubscr
 
 // Subscribe setups a subscription to receive messages that match a specific content filter
 func (wf *WakuFilterLightNode) Subscribe(ctx context.Context, contentFilter ContentFilter, opts ...FilterSubscribeOption) (*SubscriptionDetails, error) {
+	wf.RLock()
+	defer wf.RUnlock()
+
+	if !wf.started {
+		return nil, errNotStarted
+	}
+
 	if contentFilter.Topic == "" {
 		return nil, errors.New("topic is required")
 	}
@@ -244,6 +281,13 @@ func (wf *WakuFilterLightNode) Subscribe(ctx context.Context, contentFilter Cont
 
 // FilterSubscription is used to obtain an object from which you could receive messages received via filter protocol
 func (wf *WakuFilterLightNode) FilterSubscription(peerID peer.ID, contentFilter ContentFilter) (*SubscriptionDetails, error) {
+	wf.RLock()
+	defer wf.RUnlock()
+
+	if !wf.started {
+		return nil, errNotStarted
+	}
+
 	if !wf.subscriptions.Has(peerID, contentFilter.Topic, contentFilter.ContentTopics...) {
 		return nil, errors.New("subscription does not exist")
 	}
@@ -263,6 +307,13 @@ func (wf *WakuFilterLightNode) getUnsubscribeParameters(opts ...FilterUnsubscrib
 }
 
 func (wf *WakuFilterLightNode) Ping(ctx context.Context, peerID peer.ID) error {
+	wf.RLock()
+	defer wf.RUnlock()
+
+	if !wf.started {
+		return errNotStarted
+	}
+
 	return wf.request(
 		ctx,
 		&FilterSubscribeParameters{selectedPeer: peerID},
@@ -271,10 +322,24 @@ func (wf *WakuFilterLightNode) Ping(ctx context.Context, peerID peer.ID) error {
 }
 
 func (wf *WakuFilterLightNode) IsSubscriptionAlive(ctx context.Context, subscription *SubscriptionDetails) error {
+	wf.RLock()
+	defer wf.RUnlock()
+
+	if !wf.started {
+		return errNotStarted
+	}
+
 	return wf.Ping(ctx, subscription.PeerID)
 }
 
 func (wf *WakuFilterLightNode) Subscriptions() []*SubscriptionDetails {
+	wf.RLock()
+	defer wf.RUnlock()
+
+	if !wf.started {
+		return nil
+	}
+
 	wf.subscriptions.RLock()
 	defer wf.subscriptions.RUnlock()
 
@@ -324,6 +389,13 @@ func (wf *WakuFilterLightNode) cleanupSubscriptions(peerID peer.ID, contentFilte
 
 // Unsubscribe is used to stop receiving messages from a peer that match a content filter
 func (wf *WakuFilterLightNode) Unsubscribe(ctx context.Context, contentFilter ContentFilter, opts ...FilterUnsubscribeOption) (<-chan WakuFilterPushResult, error) {
+	wf.RLock()
+	defer wf.RUnlock()
+
+	if !wf.started {
+		return nil, errNotStarted
+	}
+
 	if contentFilter.Topic == "" {
 		return nil, errors.New("topic is required")
 	}
@@ -341,17 +413,33 @@ func (wf *WakuFilterLightNode) Unsubscribe(ctx context.Context, contentFilter Co
 		return nil, err
 	}
 
-	localWg := sync.WaitGroup{}
 	resultChan := make(chan WakuFilterPushResult, len(wf.subscriptions.items))
-	var peersUnsubscribed []peer.ID
 	for peerID := range wf.subscriptions.items {
 		if params.selectedPeer != "" && peerID != params.selectedPeer {
 			continue
 		}
-		peersUnsubscribed = append(peersUnsubscribed, peerID)
-		localWg.Add(1)
+
+		subscriptions, ok := wf.subscriptions.items[peerID]
+		if !ok || subscriptions == nil {
+			continue
+		}
+
+		wf.cleanupSubscriptions(peerID, contentFilter)
+		if len(subscriptions.subscriptionsPerTopic) == 0 {
+			delete(wf.subscriptions.items, peerID)
+		}
+
+		if params.wg != nil {
+			params.wg.Add(1)
+		}
+
 		go func(peerID peer.ID) {
-			defer localWg.Done()
+			defer func() {
+				if params.wg != nil {
+					params.wg.Done()
+				}
+			}()
+
 			err := wf.request(
 				ctx,
 				&FilterSubscribeParameters{selectedPeer: peerID, requestID: params.requestID},
@@ -367,27 +455,33 @@ func (wf *WakuFilterLightNode) Unsubscribe(ctx context.Context, contentFilter Co
 				}
 			}
 
-			wf.cleanupSubscriptions(peerID, contentFilter)
-
-			resultChan <- WakuFilterPushResult{
-				Err:    err,
-				PeerID: peerID,
+			if params.wg != nil {
+				resultChan <- WakuFilterPushResult{
+					Err:    err,
+					PeerID: peerID,
+				}
 			}
 		}(peerID)
 	}
 
-	localWg.Wait()
-	close(resultChan)
-	for _, peerID := range peersUnsubscribed {
-		if wf.subscriptions != nil && wf.subscriptions.items != nil && wf.subscriptions.items[peerID] != nil && len(wf.subscriptions.items[peerID].subscriptionsPerTopic) == 0 {
-			delete(wf.subscriptions.items, peerID)
-		}
+	if params.wg != nil {
+		params.wg.Wait()
 	}
+
+	close(resultChan)
+
 	return resultChan, nil
 }
 
 // Unsubscribe is used to stop receiving messages from a peer that match a content filter
 func (wf *WakuFilterLightNode) UnsubscribeWithSubscription(ctx context.Context, sub *SubscriptionDetails, opts ...FilterUnsubscribeOption) (<-chan WakuFilterPushResult, error) {
+	wf.RLock()
+	defer wf.RUnlock()
+
+	if !wf.started {
+		return nil, errNotStarted
+	}
+
 	var contentTopics []string
 	for k := range sub.ContentTopics {
 		contentTopics = append(contentTopics, k)
@@ -398,8 +492,7 @@ func (wf *WakuFilterLightNode) UnsubscribeWithSubscription(ctx context.Context,
 	return wf.Unsubscribe(ctx, ContentFilter{Topic: sub.PubsubTopic, ContentTopics: contentTopics}, opts...)
 }
 
-// UnsubscribeAll is used to stop receiving messages from peer(s). It does not close subscriptions
-func (wf *WakuFilterLightNode) UnsubscribeAll(ctx context.Context, opts ...FilterUnsubscribeOption) (<-chan WakuFilterPushResult, error) {
+func (wf *WakuFilterLightNode) unsubscribeAll(ctx context.Context, opts ...FilterUnsubscribeOption) (<-chan WakuFilterPushResult, error) {
 	params, err := wf.getUnsubscribeParameters(opts...)
 	if err != nil {
 		return nil, err
@@ -408,19 +501,26 @@ func (wf *WakuFilterLightNode) UnsubscribeAll(ctx context.Context, opts ...Filte
 	wf.subscriptions.Lock()
 	defer wf.subscriptions.Unlock()
 
-	localWg := sync.WaitGroup{}
 	resultChan := make(chan WakuFilterPushResult, len(wf.subscriptions.items))
-	var peersUnsubscribed []peer.ID
 
 	for peerID := range wf.subscriptions.items {
 		if params.selectedPeer != "" && peerID != params.selectedPeer {
 			continue
 		}
-		peersUnsubscribed = append(peersUnsubscribed, peerID)
 
-		localWg.Add(1)
+		delete(wf.subscriptions.items, peerID)
+
+		if params.wg != nil {
+			params.wg.Add(1)
+		}
+
 		go func(peerID peer.ID) {
-			defer localWg.Done()
+			defer func() {
+				if params.wg != nil {
+					params.wg.Done()
+				}
+			}()
+
 			err := wf.request(
 				ctx,
 				&FilterSubscribeParameters{selectedPeer: peerID, requestID: params.requestID},
@@ -429,17 +529,32 @@ func (wf *WakuFilterLightNode) UnsubscribeAll(ctx context.Context, opts ...Filte
 			if err != nil {
 				wf.log.Error("could not unsubscribe from peer", logging.HostID("peerID", peerID), zap.Error(err))
 			}
-			resultChan <- WakuFilterPushResult{
-				Err:    err,
-				PeerID: peerID,
+			if params.wg != nil {
+				resultChan <- WakuFilterPushResult{
+					Err:    err,
+					PeerID: peerID,
+				}
 			}
 		}(peerID)
 	}
 
-	localWg.Wait()
-	close(resultChan)
-	for _, peerID := range peersUnsubscribed {
-		delete(wf.subscriptions.items, peerID)
+	if params.wg != nil {
+		params.wg.Wait()
 	}
+
+	close(resultChan)
+
 	return resultChan, nil
 }
+
+// UnsubscribeAll is used to stop receiving messages from peer(s). It does not close subscriptions
+func (wf *WakuFilterLightNode) UnsubscribeAll(ctx context.Context, opts ...FilterUnsubscribeOption) (<-chan WakuFilterPushResult, error) {
+	wf.RLock()
+	defer wf.RUnlock()
+
+	if !wf.started {
+		return nil, errNotStarted
+	}
+
+	return wf.unsubscribeAll(ctx, opts...)
+}