netbird/signal/peer/peer_test.go
Misha Bragin 762a26dcea
Fix Register/Deregister race on Signal (#431)
This PR fixes a race condition that happens
when agents connect to a Signal stream, multiple
times within a short amount of time. Common on
slow and unstable internet connections.
Every time an agent establishes a new connection
to Signal, Signal creates a Stream and writes an entry
to the registry of connected peers storing the stream.
Every time an agent disconnects, Signal removes the
stream from the registry.
Due to unstable connections, the agent could detect
a broken connection, and attempt to reconnect to Signal.
Signal will override the stream, but it might detect
the old broken connection later, causing peer deregistration.
It will deregister the peer leaving the client thinking
it is still connected, rejecting any messages.
2022-08-22 12:21:19 +02:00

80 lines
1.7 KiB
Go

package peer
import (
"github.com/stretchr/testify/assert"
"testing"
"time"
)
func TestRegistry_ShouldNotDeregisterWhenHasNewerStreamRegistered(t *testing.T) {
r := NewRegistry()
peerID := "peer"
olderPeer := NewPeer(peerID, nil)
r.Register(olderPeer)
time.Sleep(time.Nanosecond)
newerPeer := NewPeer(peerID, nil)
r.Register(newerPeer)
registered, _ := r.Get(olderPeer.Id)
assert.NotNil(t, registered, "peer can't be nil")
assert.Equal(t, newerPeer, registered)
r.Deregister(olderPeer)
registered, _ = r.Get(olderPeer.Id)
assert.NotNil(t, registered, "peer can't be nil")
assert.Equal(t, newerPeer, registered)
}
func TestRegistry_GetNonExistentPeer(t *testing.T) {
r := NewRegistry()
peer, ok := r.Get("non_existent_peer")
if peer != nil {
t.Errorf("expected non_existent_peer not found in the registry")
}
if ok {
t.Errorf("expected non_existent_peer not found in the registry")
}
}
func TestRegistry_Register(t *testing.T) {
r := NewRegistry()
peer1 := NewPeer("test_peer_1", nil)
peer2 := NewPeer("test_peer_2", nil)
r.Register(peer1)
r.Register(peer2)
if _, ok := r.Get("test_peer_1"); !ok {
t.Errorf("expected test_peer_1 not found in the registry")
}
if _, ok := r.Get("test_peer_2"); !ok {
t.Errorf("expected test_peer_2 not found in the registry")
}
}
func TestRegistry_Deregister(t *testing.T) {
r := NewRegistry()
peer1 := NewPeer("test_peer_1", nil)
peer2 := NewPeer("test_peer_2", nil)
r.Register(peer1)
r.Register(peer2)
r.Deregister(peer1)
if _, ok := r.Get("test_peer_1"); ok {
t.Errorf("expected test_peer_1 to absent in the registry after deregistering")
}
if _, ok := r.Get("test_peer_2"); !ok {
t.Errorf("expected test_peer_2 not found in the registry")
}
}