From 855d21c37eb253a7266f9e3de13dcbeeed3d9b3b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Zolt=C3=A1n=20Papp?= Date: Mon, 2 Jun 2025 16:18:31 +0200 Subject: [PATCH] Refactor peer state change subscription mechanism Because the code generated new channel for every single event, was easy to miss notification. Use single channel. --- client/internal/peer/status.go | 78 ++++++++++++++++++++------ client/internal/peer/status_test.go | 6 +- client/internal/routemanager/client.go | 11 ++-- 3 files changed, 70 insertions(+), 25 deletions(-) diff --git a/client/internal/peer/status.go b/client/internal/peer/status.go index 40956e68e..52c7fb24b 100644 --- a/client/internal/peer/status.go +++ b/client/internal/peer/status.go @@ -146,11 +146,33 @@ type FullStatus struct { LazyConnectionEnabled bool } +type StatusChangeSubscription struct { + peerID string + id string + eventsChan chan struct{} +} + +func newStatusChangeSubscription(peerID string) *StatusChangeSubscription { + return &StatusChangeSubscription{ + peerID: peerID, + id: uuid.New().String(), + eventsChan: make(chan struct{}, 1), + } +} + +func (s *StatusChangeSubscription) Events() chan struct{} { + return s.eventsChan +} + +func (s *StatusChangeSubscription) close() { + close(s.eventsChan) +} + // Status holds a state of peers, signal, management connections and relays type Status struct { mux sync.Mutex peers map[string]State - changeNotify map[string]chan struct{} + changeNotify map[string]map[string]*StatusChangeSubscription // map[peerID]map[subscriptionID]*StatusChangeSubscription signalState bool signalError error managementState bool @@ -187,7 +209,7 @@ type Status struct { func NewRecorder(mgmAddress string) *Status { return &Status{ peers: make(map[string]State), - changeNotify: make(map[string]chan struct{}), + changeNotify: make(map[string]map[string]*StatusChangeSubscription), eventStreams: make(map[string]chan *proto.SystemEvent), eventQueue: NewEventQueue(eventQueueSize), offlinePeers: make([]State, 0), @@ -312,7 +334,6 @@ func (d *Status) UpdatePeerState(receivedState State) error { // when we close the connection we will not notify the router manager if receivedState.ConnStatus == StatusIdle { d.notifyPeerStateChangeListeners(receivedState.PubKey) - } return nil } @@ -552,19 +573,42 @@ func (d *Status) FinishPeerListModifications() { d.notifyPeerListChanged() } -// GetPeerStateChangeNotifier returns a change notifier channel for a peer -func (d *Status) GetPeerStateChangeNotifier(peer string) <-chan struct{} { +func (d *Status) SubscribeToPeerStateChanges(peerID string) *StatusChangeSubscription { d.mux.Lock() defer d.mux.Unlock() - ch, found := d.changeNotify[peer] - if found { - return ch + sub := newStatusChangeSubscription(peerID) + if _, ok := d.changeNotify[peerID]; !ok { + d.changeNotify[peerID] = make(map[string]*StatusChangeSubscription) + } + d.changeNotify[peerID][sub.id] = sub + + return sub +} + +func (d *Status) UnsubscribePeerStateChanges(subscription *StatusChangeSubscription) { + d.mux.Lock() + defer d.mux.Unlock() + + if subscription == nil { + return } - ch = make(chan struct{}) - d.changeNotify[peer] = ch - return ch + channels, ok := d.changeNotify[subscription.peerID] + if !ok { + return + } + + sub, exists := channels[subscription.id] + if !exists { + return + } + + sub.close() + delete(channels, subscription.id) + if len(channels) == 0 { + delete(d.changeNotify, sub.peerID) + } } // GetLocalPeerState returns the local peer state @@ -939,13 +983,15 @@ func (d *Status) onConnectionChanged() { // notifyPeerStateChangeListeners notifies route manager about the change in peer state func (d *Status) notifyPeerStateChangeListeners(peerID string) { - ch, found := d.changeNotify[peerID] - if !found { + subs, ok := d.changeNotify[peerID] + if !ok { return } - - close(ch) - delete(d.changeNotify, peerID) + for _, sub := range subs { + select { + case sub.eventsChan <- struct{}{}: + } + } } func (d *Status) notifyPeerListChanged() { diff --git a/client/internal/peer/status_test.go b/client/internal/peer/status_test.go index 8f28a9862..4d7a4ca78 100644 --- a/client/internal/peer/status_test.go +++ b/client/internal/peer/status_test.go @@ -86,8 +86,8 @@ func TestGetPeerStateChangeNotifierLogic(t *testing.T) { status := NewRecorder("https://mgm") _ = status.AddPeer(key, "abc.netbird", ip) - ch := status.GetPeerStateChangeNotifier(key) - assert.NotNil(t, ch, "channel shouldn't be nil") + sub := status.SubscribeToPeerStateChanges(key) + assert.NotNil(t, sub, "channel shouldn't be nil") peerState := State{ PubKey: key, @@ -100,7 +100,7 @@ func TestGetPeerStateChangeNotifierLogic(t *testing.T) { assert.NoError(t, err, "shouldn't return error") select { - case <-ch: + case <-sub.eventsChan: default: t.Errorf("channel wasn't closed after update") } diff --git a/client/internal/routemanager/client.go b/client/internal/routemanager/client.go index 137e00d31..797ef6c48 100644 --- a/client/internal/routemanager/client.go +++ b/client/internal/routemanager/client.go @@ -224,19 +224,18 @@ func (c *clientNetwork) getBestRouteFromStatuses(routePeerStatuses map[route.ID] } func (c *clientNetwork) watchPeerStatusChanges(ctx context.Context, peerKey string, peerStateUpdate chan struct{}, closer chan struct{}) { + subscription := c.statusRecorder.SubscribeToPeerStateChanges(peerKey) + defer c.statusRecorder.UnsubscribePeerStateChanges(subscription) + for { select { case <-ctx.Done(): return case <-closer: return - case <-c.statusRecorder.GetPeerStateChangeNotifier(peerKey): - state, err := c.statusRecorder.GetPeer(peerKey) - if err != nil { - continue - } + case <-subscription.Events(): peerStateUpdate <- struct{}{} - log.Debugf("triggered route state update for Peer %s, state: %s", peerKey, state.ConnStatus) + log.Debugf("triggered route state update for Peer: %s", peerKey) } } }