[server, relay] Fix/relay race disconnection (#4174)

Avoid invalid disconnection notifications in case the closed race dials.
In this PR resolve multiple race condition questions. Easier to understand the fix based on commit by commit.

- Remove store dependency from notifier
- Enforce the notification orders
- Fix invalid disconnection notification
- Ensure the order of the events on the consumer side

relay/server/listener/quic/listener.go

@@ -18,12 +18,9 @@ type Listener struct {
 	TLSConfig *tls.Config
 
 	listener *quic.Listener
-	acceptFn func(conn net.Conn)
 }
 
 func (l *Listener) Listen(acceptFn func(conn net.Conn)) error {
-	l.acceptFn = acceptFn
-
 	quicCfg := &quic.Config{
 		EnableDatagrams:   true,
 		InitialPacketSize: 1452,
@@ -49,7 +46,7 @@ func (l *Listener) Listen(acceptFn func(conn net.Conn)) error {
 
 		log.Infof("QUIC client connected from: %s", session.RemoteAddr())
 		conn := NewConn(session)
-		l.acceptFn(conn)
+		acceptFn(conn)
 	}
 }
 

relay/server/peer.go

@@ -32,6 +32,9 @@ type Peer struct {
 	notifier *store.PeerNotifier
 
 	peersListener *store.Listener
+
+	// between the online peer collection step and the notification sending should not be sent offline notifications from another thread
+	notificationMutex sync.Mutex
 }
 
 // NewPeer creates a new Peer instance and prepare custom logging
@@ -241,10 +244,16 @@ func (p *Peer) handleSubscribePeerState(msg []byte) {
 	}
 
 	p.log.Debugf("received subscription message for %d peers", len(peerIDs))
-	onlinePeers := p.peersListener.AddInterestedPeers(peerIDs)
+
+	// collect online peers to response back to the caller
+	p.notificationMutex.Lock()
+	defer p.notificationMutex.Unlock()
+
+	onlinePeers := p.store.GetOnlinePeersAndRegisterInterest(peerIDs, p.peersListener)
 	if len(onlinePeers) == 0 {
 		return
 	}
+
 	p.log.Debugf("response with %d online peers", len(onlinePeers))
 	p.sendPeersOnline(onlinePeers)
 }
@@ -274,6 +283,9 @@ func (p *Peer) sendPeersOnline(peers []messages.PeerID) {
 }
 
 func (p *Peer) sendPeersWentOffline(peers []messages.PeerID) {
+	p.notificationMutex.Lock()
+	defer p.notificationMutex.Unlock()
+
 	msgs, err := messages.MarshalPeersWentOffline(peers)
 	if err != nil {
 		p.log.Errorf("failed to marshal peer location message: %s", err)

relay/server/relay.go

@@ -86,14 +86,13 @@ func NewRelay(config Config) (*Relay, error) {
 		return nil, fmt.Errorf("creating app metrics: %v", err)
 	}
 
-	peerStore := store.NewStore()
 	r := &Relay{
 		metrics:       m,
 		metricsCancel: metricsCancel,
 		validator:     config.AuthValidator,
 		instanceURL:   config.instanceURL,
-		store:         peerStore,
-		notifier:      store.NewPeerNotifier(peerStore),
+		store:         store.NewStore(),
+		notifier:      store.NewPeerNotifier(),
 	}
 
 	r.preparedMsg, err = newPreparedMsg(r.instanceURL)
@@ -131,15 +130,18 @@ func (r *Relay) Accept(conn net.Conn) {
 	peer := NewPeer(r.metrics, *peerID, conn, r.store, r.notifier)
 	peer.log.Infof("peer connected from: %s", conn.RemoteAddr())
 	storeTime := time.Now()
-	r.store.AddPeer(peer)
+	if isReconnection := r.store.AddPeer(peer); isReconnection {
+		r.metrics.RecordPeerReconnection()
+	}
 	r.notifier.PeerCameOnline(peer.ID())
 
 	r.metrics.RecordPeerStoreTime(time.Since(storeTime))
 	r.metrics.PeerConnected(peer.String())
 	go func() {
 		peer.Work()
-		r.notifier.PeerWentOffline(peer.ID())
-		r.store.DeletePeer(peer)
+		if deleted := r.store.DeletePeer(peer); deleted {
+			r.notifier.PeerWentOffline(peer.ID())
+		}
 		peer.log.Debugf("relay connection closed")
 		r.metrics.PeerDisconnected(peer.String())
 	}()

relay/server/store/listener.go

@@ -7,24 +7,27 @@ import (
 	"github.com/netbirdio/netbird/relay/messages"
 )
 
-type Listener struct {
-	ctx   context.Context
-	store *Store
+type event struct {
+	peerID messages.PeerID
+	online bool
+}
 
-	onlineChan                chan messages.PeerID
-	offlineChan               chan messages.PeerID
+type Listener struct {
+	ctx context.Context
+
+	eventChan                 chan *event
 	interestedPeersForOffline map[messages.PeerID]struct{}
 	interestedPeersForOnline  map[messages.PeerID]struct{}
 	mu                        sync.RWMutex
 }
 
-func newListener(ctx context.Context, store *Store) *Listener {
+func newListener(ctx context.Context) *Listener {
 	l := &Listener{
-		ctx:   ctx,
-		store: store,
+		ctx: ctx,
 
-		onlineChan:                make(chan messages.PeerID, 244), //244 is the message size limit in the relay protocol
-		offlineChan:               make(chan messages.PeerID, 244), //244 is the message size limit in the relay protocol
+		// important to use a single channel for offline and online events because with it we can ensure all events
+		// will be processed in the order they were sent
+		eventChan:                 make(chan *event, 244), //244 is the message size limit in the relay protocol
 		interestedPeersForOffline: make(map[messages.PeerID]struct{}),
 		interestedPeersForOnline:  make(map[messages.PeerID]struct{}),
 	}
@@ -32,8 +35,7 @@ func newListener(ctx context.Context, store *Store) *Listener {
 	return l
 }
 
-func (l *Listener) AddInterestedPeers(peerIDs []messages.PeerID) []messages.PeerID {
-	availablePeers := make([]messages.PeerID, 0)
+func (l *Listener) AddInterestedPeers(peerIDs []messages.PeerID) {
 	l.mu.Lock()
 	defer l.mu.Unlock()
 
@@ -41,17 +43,6 @@ func (l *Listener) AddInterestedPeers(peerIDs []messages.PeerID) []messages.Peer
 		l.interestedPeersForOnline[id] = struct{}{}
 		l.interestedPeersForOffline[id] = struct{}{}
 	}
-
-	// collect online peers to response back to the caller
-	for _, id := range peerIDs {
-		_, ok := l.store.Peer(id)
-		if !ok {
-			continue
-		}
-
-		availablePeers = append(availablePeers, id)
-	}
-	return availablePeers
 }
 
 func (l *Listener) RemoveInterestedPeer(peerIDs []messages.PeerID) {
@@ -61,7 +52,6 @@ func (l *Listener) RemoveInterestedPeer(peerIDs []messages.PeerID) {
 	for _, id := range peerIDs {
 		delete(l.interestedPeersForOffline, id)
 		delete(l.interestedPeersForOnline, id)
-
 	}
 }
 
@@ -70,26 +60,31 @@ func (l *Listener) listenForEvents(onPeersComeOnline, onPeersWentOffline func([]
 		select {
 		case <-l.ctx.Done():
 			return
-		case pID := <-l.onlineChan:
-			peers := make([]messages.PeerID, 0)
-			peers = append(peers, pID)
-
-			for len(l.onlineChan) > 0 {
-				pID = <-l.onlineChan
-				peers = append(peers, pID)
+		case e := <-l.eventChan:
+			peersOffline := make([]messages.PeerID, 0)
+			peersOnline := make([]messages.PeerID, 0)
+			if e.online {
+				peersOnline = append(peersOnline, e.peerID)
+			} else {
+				peersOffline = append(peersOffline, e.peerID)
 			}
 
-			onPeersComeOnline(peers)
-		case pID := <-l.offlineChan:
-			peers := make([]messages.PeerID, 0)
-			peers = append(peers, pID)
-
-			for len(l.offlineChan) > 0 {
-				pID = <-l.offlineChan
-				peers = append(peers, pID)
+			// Drain the channel to collect all events
+			for len(l.eventChan) > 0 {
+				e = <-l.eventChan
+				if e.online {
+					peersOnline = append(peersOnline, e.peerID)
+				} else {
+					peersOffline = append(peersOffline, e.peerID)
+				}
 			}
 
-			onPeersWentOffline(peers)
+			if len(peersOnline) > 0 {
+				onPeersComeOnline(peersOnline)
+			}
+			if len(peersOffline) > 0 {
+				onPeersWentOffline(peersOffline)
+			}
 		}
 	}
 }
@@ -100,7 +95,10 @@ func (l *Listener) peerWentOffline(peerID messages.PeerID) {
 
 	if _, ok := l.interestedPeersForOffline[peerID]; ok {
 		select {
-		case l.offlineChan <- peerID:
+		case l.eventChan <- &event{
+			peerID: peerID,
+			online: false,
+		}:
 		case <-l.ctx.Done():
 		}
 	}
@@ -112,9 +110,13 @@ func (l *Listener) peerComeOnline(peerID messages.PeerID) {
 
 	if _, ok := l.interestedPeersForOnline[peerID]; ok {
 		select {
-		case l.onlineChan <- peerID:
+		case l.eventChan <- &event{
+			peerID: peerID,
+			online: true,
+		}:
 		case <-l.ctx.Done():
 		}
+
 		delete(l.interestedPeersForOnline, peerID)
 	}
 }

relay/server/store/notifier.go

@@ -8,15 +8,12 @@ import (
 )
 
 type PeerNotifier struct {
-	store *Store
-
 	listeners      map[*Listener]context.CancelFunc
 	listenersMutex sync.RWMutex
 }
 
-func NewPeerNotifier(store *Store) *PeerNotifier {
+func NewPeerNotifier() *PeerNotifier {
 	pn := &PeerNotifier{
-		store:     store,
 		listeners: make(map[*Listener]context.CancelFunc),
 	}
 	return pn
@@ -24,7 +21,7 @@ func NewPeerNotifier(store *Store) *PeerNotifier {
 
 func (pn *PeerNotifier) NewListener(onPeersComeOnline, onPeersWentOffline func([]messages.PeerID)) *Listener {
 	ctx, cancel := context.WithCancel(context.Background())
-	listener := newListener(ctx, pn.store)
+	listener := newListener(ctx)
 	go listener.listenForEvents(onPeersComeOnline, onPeersWentOffline)
 
 	pn.listenersMutex.Lock()

relay/server/store/store.go

@@ -26,7 +26,9 @@ func NewStore() *Store {
 }
 
 // AddPeer adds a peer to the store
-func (s *Store) AddPeer(peer IPeer) {
+// If the peer already exists, it will be replaced and the old peer will be closed
+// Returns true if the peer was replaced, false if it was added for the first time.
+func (s *Store) AddPeer(peer IPeer) bool {
 	s.peersLock.Lock()
 	defer s.peersLock.Unlock()
 	odlPeer, ok := s.peers[peer.ID()]
@@ -35,22 +37,24 @@ func (s *Store) AddPeer(peer IPeer) {
 	}
 
 	s.peers[peer.ID()] = peer
+	return ok
 }
 
 // DeletePeer deletes a peer from the store
-func (s *Store) DeletePeer(peer IPeer) {
+func (s *Store) DeletePeer(peer IPeer) bool {
 	s.peersLock.Lock()
 	defer s.peersLock.Unlock()
 
 	dp, ok := s.peers[peer.ID()]
 	if !ok {
-		return
+		return false
 	}
 	if dp != peer {
-		return
+		return false
 	}
 
 	delete(s.peers, peer.ID())
+	return true
 }
 
 // Peer returns a peer by its ID
@@ -73,3 +77,21 @@ func (s *Store) Peers() []IPeer {
 	}
 	return peers
 }
+
+func (s *Store) GetOnlinePeersAndRegisterInterest(peerIDs []messages.PeerID, listener *Listener) []messages.PeerID {
+	s.peersLock.RLock()
+	defer s.peersLock.RUnlock()
+
+	onlinePeers := make([]messages.PeerID, 0, len(peerIDs))
+
+	listener.AddInterestedPeers(peerIDs)
+
+	// Check for currently online peers
+	for _, id := range peerIDs {
+		if _, ok := s.peers[id]; ok {
+			onlinePeers = append(onlinePeers, id)
+		}
+	}
+
+	return onlinePeers
+}