Merge pull request #180 from canopy-network/issue-#164

Issue-#164:  Allow multiple subscribers per `ChainID` to ws updates

Author: rem1niscence

cmd/rpc/sock.go

@@ -2,15 +2,17 @@ package rpc
 
 import (
 	"fmt"
-	"github.com/canopy-network/canopy/controller"
-	"github.com/canopy-network/canopy/lib"
-	"github.com/gorilla/websocket"
-	"github.com/julienschmidt/httprouter"
 	"net/http"
 	"net/url"
+	"slices"
 	"strconv"
 	"sync"
 	"time"
+
+	"github.com/canopy-network/canopy/controller"
+	"github.com/canopy-network/canopy/lib"
+	"github.com/gorilla/websocket"
+	"github.com/julienschmidt/httprouter"
 )
 
 /* This file implements the client & server logic for the 'root-chain info' and corresponding 'on-demand' calls to the rpc */
@@ -23,7 +25,7 @@ const chainIdParamName = "chainId"
 type RCManager struct {
 	c             lib.Config                    // the global node config
 	subscriptions map[uint64]*RCSubscription    // chainId -> subscription
-	subscribers   map[uint64]*RCSubscriber      // chainId -> subscriber
+	subscribers   map[uint64][]*RCSubscriber    // chainId -> subscribers
 	l             *sync.Mutex                   // thread safety
 	afterRCUpdate func(info *lib.RootChainInfo) // callback after the root chain info update
 	upgrader      websocket.Upgrader            // upgrade http connection to ws
@@ -36,7 +38,7 @@ func NewRCManager(controller *controller.Controller, config lib.Config, logger l
 	manager = &RCManager{
 		c:             config,
 		subscriptions: make(map[uint64]*RCSubscription),
-		subscribers:   make(map[uint64]*RCSubscriber),
+		subscribers:   make(map[uint64][]*RCSubscriber),
 		l:             controller.Mutex,
 		afterRCUpdate: controller.UpdateRootChainInfo,
 		upgrader:      websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }},
@@ -65,14 +67,13 @@ func (r *RCManager) Publish(chainId uint64, info *lib.RootChainInfo) {
 		return
 	}
 	// for each ws client
-	for _, subscriber := range r.subscribers {
-		// skip if not this chain id
-		if subscriber.chainId != chainId {
-			continue
-		}
+	for _, subscriber := range r.subscribers[chainId] {
 		// publish to each client
 		if e := subscriber.conn.WriteMessage(websocket.BinaryMessage, protoBytes); e != nil {
-			subscriber.Stop(e)
+			// defer the Stop() call to prevent the slice modification during iteration.
+			// since Stop() removes the subscriber from r.subscribers, immediate execution
+			// would affect the slice that is currently being iterated.
+			defer subscriber.Stop(e)
 		}
 	}
 }
@@ -82,10 +83,17 @@ func (r *RCManager) ChainIds() (list []uint64) {
 	// de-duplicate the results
 	deDupe := lib.NewDeDuplicator[uint64]()
 	// for each client
-	for _, client := range r.subscribers {
+	for chainId, chainSubscribers := range r.subscribers {
 		// if the client chain id isn't empty and not duplicate
-		if client.chainId != 0 && !deDupe.Found(client.chainId) {
-			list = append(list, client.chainId)
+		for _, subscriber := range chainSubscribers {
+			if subscriber.chainId != chainId {
+				// remove subscriber with incorrect chain id
+				subscriber.Stop(lib.ErrWrongChainId())
+				continue
+			}
+			if subscriber.chainId != 0 && !deDupe.Found(subscriber.chainId) {
+				list = append(list, subscriber.chainId)
+			}
 		}
 	}
 	return
@@ -424,16 +432,18 @@ func (r *RCManager) AddSubscriber(subscriber *RCSubscriber) {
 	r.l.Lock()
 	defer r.l.Unlock()
 	// add to the map
-	r.subscribers[subscriber.chainId] = subscriber
+	r.subscribers[subscriber.chainId] = append(r.subscribers[subscriber.chainId], subscriber)
 }
 
 // RemoveSubscriber() gracefully deletes a RC subscriber
-func (r *RCManager) RemoveSubscriber(chainId uint64) {
+func (r *RCManager) RemoveSubscriber(chainId uint64, subscriber *RCSubscriber) {
 	// lock for thread safety
 	r.l.Lock()
 	defer r.l.Unlock()
-	// remove from the map
-	delete(r.subscribers, chainId)
+	// remove from the slice
+	r.subscribers[chainId] = slices.DeleteFunc(r.subscribers[chainId], func(sub *RCSubscriber) bool {
+		return sub == subscriber
+	})
 }
 
 // Stop() stops the client
@@ -445,5 +455,5 @@ func (r *RCSubscriber) Stop(err error) {
 		r.log.Error(err.Error())
 	}
 	// remove from the manager
-	r.manager.RemoveSubscriber(r.chainId)
+	r.manager.RemoveSubscriber(r.chainId, r)
 }