[relay] Unit test ForeignRelaysStore and FallbackOpener

# Conflicts:
#	shared/relay/client/manager.go

client/internal/peer/worker_relay.go

@@ -54,19 +54,15 @@ func (w *WorkerRelay) OnNewOffer(remoteOfferAnswer *OfferAnswer) {
 	w.relaySupportedOnRemotePeer.Store(true)
 
 	// the relayManager will return with error in case if the connection has lost with relay server
-	currentRelayAddress, _, err := w.relayManager.RelayInstanceAddress()
+	_, _, err := w.relayManager.RelayInstanceAddress()
 	if err != nil {
 		w.log.Errorf("failed to handle new offer: %s", err)
 		return
 	}
 
-	srv := w.preferredRelayServer(currentRelayAddress, remoteOfferAnswer.RelaySrvAddress)
-	var serverIP netip.Addr
-	if srv == remoteOfferAnswer.RelaySrvAddress {
-		serverIP = remoteOfferAnswer.RelaySrvIP
-	}
-
-	relayedConn, err := w.relayManager.OpenConn(w.peerCtx, srv, w.config.Key, serverIP)
+	preferForeign := !w.isController
+	remoteRelayServer := relayClient.RelayServer{Addr: remoteOfferAnswer.RelaySrvAddress, IP: remoteOfferAnswer.RelaySrvIP}
+	relayedConn, err := w.relayManager.OpenConn(w.peerCtx, remoteRelayServer, w.config.Key, preferForeign)
 	if err != nil {
 		if errors.Is(err, relayClient.ErrConnAlreadyExists) {
 			w.log.Debugf("handled offer by reusing existing relay connection")
@@ -80,14 +76,13 @@ func (w *WorkerRelay) OnNewOffer(remoteOfferAnswer *OfferAnswer) {
 	w.relayedConn = relayedConn
 	w.relayLock.Unlock()
 
-	err = w.relayManager.AddCloseListener(srv, w.onRelayClientDisconnected)
-	if err != nil {
-		log.Errorf("failed to add close listener: %s", err)
+	if err := w.relayManager.AddCloseListener(relayedConn.RemoteAddr().String(), w.onRelayClientDisconnected); err != nil {
+		w.log.Errorf("failed to add close listener: %s", err)
 		_ = relayedConn.Close()
 		return
 	}
 
-	w.log.Debugf("peer conn opened via Relay: %s", srv)
+	w.log.Debugf("peer conn opened via Relay: %s", relayedConn.RemoteAddr())
 	go w.conn.onRelayConnectionIsReady(RelayConnInfo{
 		relayedConn:     relayedConn,
 		rosenpassPubKey: remoteOfferAnswer.RosenpassPubKey,
@@ -126,13 +121,6 @@ func (w *WorkerRelay) isRelaySupported(answer *OfferAnswer) bool {
 	return answer.RelaySrvAddress != ""
 }
 
-func (w *WorkerRelay) preferredRelayServer(myRelayAddress, remoteRelayAddress string) string {
-	if w.isController {
-		return myRelayAddress
-	}
-	return remoteRelayAddress
-}
-
 func (w *WorkerRelay) onRelayClientDisconnected() {
 	go w.conn.onRelayDisconnected()
 }

shared/relay/client/fallback_opener.go

@@ -0,0 +1,165 @@
+package client
+
+import (
+	"context"
+	"errors"
+	"net"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+)
+
+const (
+	raceTotalTimeout  = 40 * time.Second
+	raceFallbackDelay = 10 * time.Second
+)
+
+type raceAttempt struct {
+	conn net.Conn
+	err  error
+}
+
+type raceOutcome struct {
+	conn net.Conn
+	err  error
+	done bool
+}
+
+type connRace struct {
+	opener            *FallbackOpener
+	peerKey           string
+	remoteRelayServer RelayServer
+	preferForeign     bool
+
+	raceCtx         context.Context
+	otherCtx        context.Context
+	cancelPreferred context.CancelFunc
+	cancelOther     context.CancelFunc
+	results         chan raceAttempt
+	fallbackTimer   *time.Timer
+
+	otherStarted bool
+	settled      int
+	lastErr      error
+}
+
+type FallbackOpener struct {
+	home         *Client
+	foreignStore *ForeignRelaysStore
+}
+
+func NewFallbackOpener(home *Client, foreignStore *ForeignRelaysStore) *FallbackOpener {
+	return &FallbackOpener{
+		home:         home,
+		foreignStore: foreignStore,
+	}
+}
+
+func (r *FallbackOpener) Run(ctx context.Context, peerKey string, remoteRelayServer RelayServer, preferForeign bool) (net.Conn, error) {
+	raceCtx, cancel := context.WithTimeout(ctx, raceTotalTimeout)
+	defer cancel()
+
+	preferredCtx, cancelPreferred := context.WithCancel(raceCtx)
+	otherCtx, cancelOther := context.WithCancel(raceCtx)
+
+	race := &connRace{
+		opener:            r,
+		peerKey:           peerKey,
+		remoteRelayServer: remoteRelayServer,
+		preferForeign:     preferForeign,
+		raceCtx:           raceCtx,
+		otherCtx:          otherCtx,
+		cancelPreferred:   cancelPreferred,
+		cancelOther:       cancelOther,
+		results:           make(chan raceAttempt, 2),
+		fallbackTimer:     time.NewTimer(raceFallbackDelay),
+	}
+	defer race.fallbackTimer.Stop()
+
+	go func() {
+		race.results <- r.open(preferredCtx, peerKey, remoteRelayServer, preferForeign)
+	}()
+
+	for {
+		select {
+		case <-race.fallbackTimer.C:
+			race.startOther()
+		case res := <-race.results:
+			if o := race.handleResult(res); o.done {
+				return o.conn, o.err
+			}
+		case <-raceCtx.Done():
+			return race.onTimeout()
+		}
+	}
+}
+
+func (c *connRace) startOther() {
+	if c.otherStarted {
+		return
+	}
+	c.otherStarted = true
+	c.fallbackTimer.Stop()
+	go func() {
+		c.results <- c.opener.open(c.otherCtx, c.peerKey, c.remoteRelayServer, !c.preferForeign)
+	}()
+}
+
+func (c *connRace) handleResult(res raceAttempt) raceOutcome {
+	if (res.err == nil && res.conn != nil) || errors.Is(res.err, ErrConnAlreadyExists) {
+		c.settled++
+		c.stop()
+		return raceOutcome{conn: res.conn, err: res.err, done: true}
+	}
+
+	c.lastErr = res.err
+	c.settled++
+	if !c.otherStarted {
+		c.startOther()
+		return raceOutcome{}
+	}
+	if c.settled == 2 {
+		c.cancelPreferred()
+		c.cancelOther()
+		return raceOutcome{err: c.lastErr, done: true}
+	}
+	return raceOutcome{}
+}
+
+func (c *connRace) onTimeout() (net.Conn, error) {
+	c.stop()
+	if c.lastErr != nil {
+		return nil, c.lastErr
+	}
+	return nil, c.raceCtx.Err()
+}
+
+func (c *connRace) stop() {
+	c.cancelPreferred()
+	c.cancelOther()
+	go c.opener.drainLoser(c.results, c.settled, c.otherStarted)
+}
+
+func (r *FallbackOpener) open(ctx context.Context, peerKey string, remoteRelayServer RelayServer, foreign bool) raceAttempt {
+	if foreign {
+		conn, err := r.foreignStore.OpenConn(ctx, peerKey, remoteRelayServer)
+		return raceAttempt{conn: conn, err: err}
+	}
+	conn, err := r.home.OpenConn(ctx, peerKey)
+	return raceAttempt{conn: conn, err: err}
+}
+
+func (r *FallbackOpener) drainLoser(results chan raceAttempt, settled int, otherStarted bool) {
+	started := 1
+	if otherStarted {
+		started = 2
+	}
+	for i := settled; i < started; i++ {
+		res := <-results
+		if res.conn != nil {
+			if err := res.conn.Close(); err != nil {
+				log.Debugf("failed to close losing relay connection: %v", err)
+			}
+		}
+	}
+}

shared/relay/client/fallback_opener_test.go

@@ -0,0 +1,175 @@
+package client
+
+import (
+	"context"
+	"errors"
+	"net"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/netbirdio/netbird/relay/server"
+)
+
+type fakeConn struct {
+	net.Conn
+	closed chan struct{}
+}
+
+func newFakeConn() *fakeConn {
+	return &fakeConn{closed: make(chan struct{})}
+}
+
+func (c *fakeConn) Close() error {
+	close(c.closed)
+	return nil
+}
+
+func newTestConnRace(t *testing.T) *connRace {
+	t.Helper()
+	raceCtx, cancel := context.WithCancel(context.Background())
+	t.Cleanup(cancel)
+	_, cancelPreferred := context.WithCancel(raceCtx)
+	otherCtx, cancelOther := context.WithCancel(raceCtx)
+	timer := time.NewTimer(time.Hour)
+	timer.Stop()
+	return &connRace{
+		opener:          &FallbackOpener{},
+		peerKey:         "peerKey",
+		raceCtx:         raceCtx,
+		otherCtx:        otherCtx,
+		cancelPreferred: cancelPreferred,
+		cancelOther:     cancelOther,
+		results:         make(chan raceAttempt, 2),
+		fallbackTimer:   timer,
+	}
+}
+
+func TestHandleResult_PreferredSucceeds(t *testing.T) {
+	c := newTestConnRace(t)
+
+	conn := newFakeConn()
+	o := c.handleResult(raceAttempt{conn: conn})
+
+	require.True(t, o.done)
+	require.NoError(t, o.err)
+	require.Same(t, net.Conn(conn), o.conn)
+	require.False(t, c.otherStarted, "fallback must not start once the preferred attempt wins")
+}
+
+func TestHandleResult_ConnAlreadyExistsIsSuccess(t *testing.T) {
+	c := newTestConnRace(t)
+
+	o := c.handleResult(raceAttempt{err: ErrConnAlreadyExists})
+
+	require.True(t, o.done)
+	require.ErrorIs(t, o.err, ErrConnAlreadyExists)
+	require.False(t, c.otherStarted)
+}
+
+func TestHandleResult_PreferredFailsStartsOther(t *testing.T) {
+	c := newTestConnRace(t)
+	// The fallback attempt opens against a stalling listener so startOther's
+	// goroutine blocks on Connect until raceCtx is cancelled by t.Cleanup.
+	serverAddr, _ := stallingRelayListener(t)
+	c.opener.foreignStore = NewForeignRelaysStore(c.raceCtx, hmacTokenStore, "alice", 1280, newTransportFallback(), func(string) {}, keepUnusedServerTime)
+	c.remoteRelayServer = RelayServer{Addr: serverAddr}
+	c.preferForeign = false
+
+	o := c.handleResult(raceAttempt{err: errors.New("boom")})
+
+	require.False(t, o.done, "a single failure must not settle the race")
+	require.True(t, c.otherStarted, "the fallback attempt must start after the preferred one fails")
+	require.EqualError(t, c.lastErr, "boom")
+}
+
+func TestHandleResult_BothFailReturnsLastErr(t *testing.T) {
+	c := newTestConnRace(t)
+	c.otherStarted = true
+	c.settled = 1
+	c.lastErr = errors.New("first")
+
+	o := c.handleResult(raceAttempt{err: errors.New("second")})
+
+	require.True(t, o.done)
+	require.EqualError(t, o.err, "second")
+}
+
+func TestOnTimeout_PrefersLastErr(t *testing.T) {
+	c := newTestConnRace(t)
+	c.lastErr = errors.New("dial failed")
+
+	_, err := c.onTimeout()
+	require.EqualError(t, err, "dial failed")
+}
+
+func TestOnTimeout_FallsBackToCtxErr(t *testing.T) {
+	c := newTestConnRace(t)
+	raceCtx, cancel := context.WithCancel(context.Background())
+	cancel()
+	c.raceCtx = raceCtx
+
+	_, err := c.onTimeout()
+	require.ErrorIs(t, err, context.Canceled)
+}
+
+func TestDrainLoser_ClosesLateWinner(t *testing.T) {
+	r := &FallbackOpener{}
+	results := make(chan raceAttempt, 2)
+
+	loser := newFakeConn()
+	results <- raceAttempt{conn: loser}
+
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		r.drainLoser(results, 1, true)
+	}()
+
+	select {
+	case <-loser.closed:
+	case <-time.After(2 * time.Second):
+		t.Fatal("losing connection was not closed")
+	}
+	select {
+	case <-done:
+	case <-time.After(2 * time.Second):
+		t.Fatal("drainLoser did not return")
+	}
+}
+
+func TestDrainLoser_NoOtherAttempt(t *testing.T) {
+	r := &FallbackOpener{}
+	results := make(chan raceAttempt)
+
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		r.drainLoser(results, 1, false)
+	}()
+
+	select {
+	case <-done:
+	case <-time.After(2 * time.Second):
+		t.Fatal("drainLoser blocked with no outstanding attempt")
+	}
+}
+
+func startTestRelayServer(t *testing.T, addr string) string {
+	t.Helper()
+
+	srv, err := server.NewServer(newManagerTestServerConfig(addr))
+	require.NoError(t, err)
+
+	errChan := make(chan error, 1)
+	go func() {
+		if err := srv.Listen(server.ListenerConfig{Address: addr}); err != nil {
+			errChan <- err
+		}
+	}()
+	t.Cleanup(func() { _ = srv.Shutdown(context.Background()) })
+
+	require.NoError(t, waitForServerToStart(errChan))
+	return addr
+}

shared/relay/client/foreign_relays_store.go

@@ -0,0 +1,155 @@
+package client
+
+import (
+	"context"
+	"net"
+	"sync"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+	"golang.org/x/sync/singleflight"
+
+	relayAuth "github.com/netbirdio/netbird/shared/relay/auth/hmac"
+)
+
+type foreignRelay struct {
+	client  *Client
+	created time.Time
+	inUse   int
+}
+
+type ForeignRelaysStore struct {
+	mu      sync.RWMutex
+	clients map[string]*foreignRelay
+
+	group singleflight.Group
+
+	ctx                  context.Context
+	tokenStore           *relayAuth.TokenStore
+	peerID               string
+	mtu                  uint16
+	transportFallback    *transportFallback
+	onDisconnect         func(string)
+	keepUnusedServerTime time.Duration
+}
+
+func NewForeignRelaysStore(ctx context.Context, tokenStore *relayAuth.TokenStore, peerID string, mtu uint16, transportFallback *transportFallback, onDisconnect func(string), keepUnusedServerTime time.Duration) *ForeignRelaysStore {
+	return &ForeignRelaysStore{
+		clients:              make(map[string]*foreignRelay),
+		ctx:                  ctx,
+		tokenStore:           tokenStore,
+		peerID:               peerID,
+		mtu:                  mtu,
+		transportFallback:    transportFallback,
+		onDisconnect:         onDisconnect,
+		keepUnusedServerTime: keepUnusedServerTime,
+	}
+}
+
+func (f *ForeignRelaysStore) OpenConn(ctx context.Context, peerKey string, remoteRelayServer RelayServer) (net.Conn, error) {
+	fr, err := f.acquire(remoteRelayServer)
+	if err != nil {
+		return nil, err
+	}
+	defer f.release(fr)
+
+	return fr.client.OpenConn(ctx, peerKey)
+}
+
+func (f *ForeignRelaysStore) acquire(remoteRelayServer RelayServer) (*foreignRelay, error) {
+	f.mu.Lock()
+	if fr, ok := f.clients[remoteRelayServer.Addr]; ok {
+		fr.inUse++
+		f.mu.Unlock()
+		return fr, nil
+	}
+	f.mu.Unlock()
+
+	v, err, _ := f.group.Do(remoteRelayServer.Addr, func() (any, error) {
+		f.mu.RLock()
+		fr, ok := f.clients[remoteRelayServer.Addr]
+		f.mu.RUnlock()
+		if ok {
+			return fr, nil
+		}
+
+		relayClient := NewClientWithServerIP(remoteRelayServer.Addr, remoteRelayServer.IP, f.tokenStore, f.peerID, f.mtu)
+		relayClient.SetTransportFallback(f.transportFallback)
+		if err := relayClient.Connect(f.ctx); err != nil {
+			return nil, err
+		}
+		relayClient.SetOnDisconnectListener(f.onDisconnect)
+
+		f.mu.Lock()
+		fr = &foreignRelay{client: relayClient, created: time.Now()}
+		f.clients[remoteRelayServer.Addr] = fr
+		f.mu.Unlock()
+		return fr, nil
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	fr := v.(*foreignRelay)
+	f.mu.Lock()
+	if cur, ok := f.clients[remoteRelayServer.Addr]; !ok || cur != fr {
+		f.mu.Unlock()
+		return f.acquire(remoteRelayServer)
+	}
+	fr.inUse++
+	f.mu.Unlock()
+	return fr, nil
+}
+
+func (f *ForeignRelaysStore) release(fr *foreignRelay) {
+	f.mu.Lock()
+	fr.inUse--
+	f.mu.Unlock()
+}
+
+func (f *ForeignRelaysStore) evict(serverAddress string) {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	if _, ok := f.clients[serverAddress]; ok {
+		delete(f.clients, serverAddress)
+		log.Debugf("evicted disconnected foreign relay client: %s", serverAddress)
+	}
+}
+
+func (f *ForeignRelaysStore) cleanupUnused() {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+
+	for addr, fr := range f.clients {
+		if time.Since(fr.created) <= f.keepUnusedServerTime {
+			continue
+		}
+		if fr.inUse > 0 {
+			continue
+		}
+		if fr.client.HasConns() {
+			continue
+		}
+		fr.client.SetOnDisconnectListener(nil)
+		go func() {
+			_ = fr.client.Close()
+		}()
+		log.Debugf("clean up unused relay server connection: %s", addr)
+		delete(f.clients, addr)
+	}
+}
+
+func (f *ForeignRelaysStore) states() []RelayConnState {
+	f.mu.RLock()
+	clients := make([]*Client, 0, len(f.clients))
+	for _, fr := range f.clients {
+		clients = append(clients, fr.client)
+	}
+	f.mu.RUnlock()
+
+	states := make([]RelayConnState, 0, len(clients))
+	for _, c := range clients {
+		states = append(states, relayConnState(c))
+	}
+	return states
+}

shared/relay/client/foreign_relays_store_test.go

@@ -0,0 +1,184 @@
+package client
+
+import (
+	"context"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+)
+
+func newTestForeignStore(t *testing.T, ctx context.Context) *ForeignRelaysStore {
+	t.Helper()
+	return NewForeignRelaysStore(ctx, hmacTokenStore, "alice", 1280, newTransportFallback(), func(string) {}, keepUnusedServerTime)
+}
+
+func TestForeignStore_AcquireDedupsConcurrentOpens(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	t.Cleanup(cancel)
+
+	addr := startTestRelayServer(t, "127.0.0.1:52601")
+	server := RelayServer{Addr: "rel://" + addr}
+
+	store := newTestForeignStore(t, ctx)
+
+	const n = 8
+	var wg sync.WaitGroup
+	results := make([]*foreignRelay, n)
+	for i := 0; i < n; i++ {
+		wg.Add(1)
+		go func(i int) {
+			defer wg.Done()
+			fr, err := store.acquire(server)
+			require.NoError(t, err)
+			results[i] = fr
+		}(i)
+	}
+	wg.Wait()
+
+	first := results[0]
+	require.NotNil(t, first)
+	for _, fr := range results {
+		require.Same(t, first, fr, "all acquires must share the same foreign relay")
+	}
+
+	store.mu.RLock()
+	require.Len(t, store.clients, 1, "only one client entry must be stored")
+	require.Equal(t, n, first.inUse, "every acquire must be counted")
+	store.mu.RUnlock()
+}
+
+func TestForeignStore_AcquireReleaseRefcount(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	t.Cleanup(cancel)
+
+	addr := startTestRelayServer(t, "127.0.0.1:52602")
+	server := RelayServer{Addr: "rel://" + addr}
+
+	store := newTestForeignStore(t, ctx)
+
+	fr, err := store.acquire(server)
+	require.NoError(t, err)
+	_, err = store.acquire(server)
+	require.NoError(t, err)
+
+	store.mu.RLock()
+	require.Equal(t, 2, fr.inUse)
+	store.mu.RUnlock()
+
+	store.release(fr)
+	store.mu.RLock()
+	require.Equal(t, 1, fr.inUse)
+	require.Len(t, store.clients, 1, "release must not evict the client")
+	store.mu.RUnlock()
+}
+
+func TestForeignStore_AcquireConnectFailure(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	t.Cleanup(cancel)
+
+	store := newTestForeignStore(t, ctx)
+
+	// Nothing is listening on this port, so Connect fails.
+	_, err := store.acquire(RelayServer{Addr: "rel://127.0.0.1:1"})
+	require.Error(t, err)
+
+	store.mu.RLock()
+	require.Empty(t, store.clients, "a failed connect must not leave a client behind")
+	store.mu.RUnlock()
+}
+
+func TestForeignStore_Evict(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	t.Cleanup(cancel)
+
+	store := newTestForeignStore(t, ctx)
+	store.clients["rel://a"] = &foreignRelay{}
+	store.clients["rel://b"] = &foreignRelay{}
+
+	store.evict("rel://a")
+	store.evict("rel://missing")
+
+	require.NotContains(t, store.clients, "rel://a")
+	require.Contains(t, store.clients, "rel://b")
+}
+
+func TestForeignStore_CleanupUnused_KeepsRecent(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	t.Cleanup(cancel)
+
+	addr := startTestRelayServer(t, "127.0.0.1:52603")
+	store := newTestForeignStore(t, ctx)
+
+	fr, err := store.acquire(RelayServer{Addr: "rel://" + addr})
+	require.NoError(t, err)
+	store.release(fr)
+
+	store.cleanupUnused()
+
+	store.mu.RLock()
+	require.Len(t, store.clients, 1, "a freshly created client must be kept")
+	store.mu.RUnlock()
+}
+
+func TestForeignStore_CleanupUnused_KeepsInUse(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	t.Cleanup(cancel)
+
+	addr := startTestRelayServer(t, "127.0.0.1:52604")
+	store := newTestForeignStore(t, ctx)
+
+	fr, err := store.acquire(RelayServer{Addr: "rel://" + addr})
+	require.NoError(t, err)
+
+	store.mu.Lock()
+	fr.created = time.Now().Add(-2 * keepUnusedServerTime)
+	store.mu.Unlock()
+
+	store.cleanupUnused()
+
+	store.mu.RLock()
+	require.Len(t, store.clients, 1, "an in-use client must be kept even when aged")
+	store.mu.RUnlock()
+}
+
+func TestForeignStore_CleanupUnused_EvictsAgedIdle(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	t.Cleanup(cancel)
+
+	addr := startTestRelayServer(t, "127.0.0.1:52605")
+	store := newTestForeignStore(t, ctx)
+
+	fr, err := store.acquire(RelayServer{Addr: "rel://" + addr})
+	require.NoError(t, err)
+	store.release(fr)
+
+	store.mu.Lock()
+	fr.created = time.Now().Add(-2 * keepUnusedServerTime)
+	store.mu.Unlock()
+
+	require.False(t, fr.client.HasConns(), "no peer connections were opened")
+
+	store.cleanupUnused()
+
+	store.mu.RLock()
+	require.Empty(t, store.clients, "an aged idle client must be evicted")
+	store.mu.RUnlock()
+}
+
+func TestForeignStore_States(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	t.Cleanup(cancel)
+
+	addr := startTestRelayServer(t, "127.0.0.1:52606")
+	store := newTestForeignStore(t, ctx)
+
+	fr, err := store.acquire(RelayServer{Addr: "rel://" + addr})
+	require.NoError(t, err)
+	store.release(fr)
+
+	states := store.states()
+	require.Len(t, states, 1)
+	require.NotEmpty(t, states[0].URL)
+}

shared/relay/client/manager.go

@@ -22,27 +22,6 @@ var (
 	ErrRelayClientNotConnected = fmt.Errorf("relay client not connected")
 )
 
-// RelayTrack hold the relay clients for the foreign relay servers.
-// With the mutex can ensure we can open new connection in case the relay connection has been established with
-// the relay server.
-type RelayTrack struct {
-	sync.RWMutex
-	relayClient *Client
-	err         error
-	created     time.Time
-	// ready is closed once the dial started by openConnVia finishes (relayClient
-	// or err is set). Callers reusing a track wait on this instead of the track
-	// lock, so the dial never runs under rt.Lock.
-	ready chan struct{}
-}
-
-func NewRelayTrack() *RelayTrack {
-	return &RelayTrack{
-		created: time.Now(),
-		ready:   make(chan struct{}),
-	}
-}
-
 type OnServerCloseListener func()
 
 // ManagerOption configures a Manager at construction time.
@@ -59,6 +38,11 @@ type RelayConnState struct {
 	Err error
 }
 
+type RelayServer struct {
+	Addr string
+	IP   netip.Addr
+}
+
 // WithMaxBackoffInterval caps the exponential backoff between reconnect
 // attempts to the home relay. A non-positive value keeps the default.
 func WithMaxBackoffInterval(d time.Duration) ManagerOption {
@@ -83,8 +67,7 @@ type Manager struct {
 	relayClientMu  sync.RWMutex
 	reconnectGuard *Guard
 
-	relayClients      map[string]*RelayTrack
-	relayClientsMutex sync.RWMutex
+	foreign *ForeignRelaysStore
 
 	onDisconnectedListeners map[string]*list.List
 	onReconnectedListenerFn func()
@@ -120,7 +103,6 @@ func NewManager(ctx context.Context, serverURLs []string, peerID string, mtu uin
 			ConnectionTimeout: defaultConnectionTimeout,
 			TransportFallback: tf,
 		},
-		relayClients:            make(map[string]*RelayTrack),
 		onDisconnectedListeners: make(map[string]*list.List),
 		cleanupInterval:         relayCleanupInterval,
 		keepUnusedServerTime:    keepUnusedServerTime,
@@ -128,6 +110,7 @@ func NewManager(ctx context.Context, serverURLs []string, peerID string, mtu uin
 	for _, opt := range opts {
 		opt(m)
 	}
+	m.foreign = NewForeignRelaysStore(ctx, tokenStore, peerID, mtu, tf, m.onServerDisconnected, m.keepUnusedServerTime)
 	m.serverPicker.ServerURLs.Store(serverURLs)
 	m.reconnectGuard = NewGuard(m.serverPicker, m.maxBackoffInterval)
 	return m
@@ -159,40 +142,26 @@ func (m *Manager) Serve() error {
 	return err
 }
 
-// OpenConn opens a connection to the given peer key. If the peer is on the same relay server, the connection will be
-// established via the relay server. If the peer is on a different relay server, the manager will establish a new
-// connection to the relay server. It returns back with a net.Conn what represent the remote peer connection.
-//
-// serverIP, when valid and serverAddress is foreign, is used as a dial target if the FQDN-based dial fails.
-// Ignored for the local home-server path. TLS verification still uses the FQDN via SNI.
-func (m *Manager) OpenConn(ctx context.Context, serverAddress, peerKey string, serverIP netip.Addr) (net.Conn, error) {
+func (m *Manager) OpenConn(ctx context.Context, remoteRelayServer RelayServer, peerKey string, preferForeign bool) (net.Conn, error) {
 	m.relayClientMu.RLock()
-	defer m.relayClientMu.RUnlock()
+	relayClient := m.relayClient
+	m.relayClientMu.RUnlock()
 
-	if m.relayClient == nil {
+	if relayClient == nil {
 		return nil, ErrRelayClientNotConnected
 	}
 
-	foreign, err := m.isForeignServer(serverAddress)
+	foreign, err := m.isForeignServer(relayClient, remoteRelayServer.Addr)
 	if err != nil {
 		return nil, err
 	}
 
-	var (
-		netConn net.Conn
-	)
 	if !foreign {
-		log.Debugf("open peer connection via permanent server: %s", peerKey)
-		netConn, err = m.relayClient.OpenConn(ctx, peerKey)
-	} else {
-		log.Debugf("open peer connection via foreign server: %s", serverAddress)
-		netConn, err = m.openConnVia(ctx, serverAddress, peerKey, serverIP)
-	}
-	if err != nil {
-		return nil, err
+		return relayClient.OpenConn(ctx, peerKey)
 	}
 
-	return netConn, err
+	opener := NewFallbackOpener(relayClient, m.foreign)
+	return opener.Run(ctx, peerKey, remoteRelayServer, preferForeign)
 }
 
 // Ready returns true if the home Relay client is connected to the relay server.
@@ -223,7 +192,7 @@ func (m *Manager) AddCloseListener(serverAddress string, onClosedListener OnServ
 		return ErrRelayClientNotConnected
 	}
 
-	foreign, err := m.isForeignServer(serverAddress)
+	foreign, err := m.isForeignServer(m.relayClient, serverAddress)
 	if err != nil {
 		return err
 	}
@@ -287,26 +256,7 @@ func (m *Manager) RelayStates() []RelayConnState {
 		states = append(states, st)
 	}
 
-	// Snapshot the tracks, then query each outside the map lock: a track can be
-	// held by an in-progress Connect, and blocking on it must not stall other
-	// relay operations.
-	m.relayClientsMutex.RLock()
-	tracks := make([]*RelayTrack, 0, len(m.relayClients))
-	for _, rt := range m.relayClients {
-		tracks = append(tracks, rt)
-	}
-	m.relayClientsMutex.RUnlock()
-
-	// Only connected foreign relays carry state; a failed connect is evicted
-	// immediately (openConnVia), so there is no error state to surface.
-	for _, rt := range tracks {
-		rt.RLock()
-		rc := rt.relayClient
-		rt.RUnlock()
-		if rc != nil {
-			states = append(states, relayConnState(rc))
-		}
-	}
+	states = append(states, m.foreign.states()...)
 
 	return states
 }
@@ -327,76 +277,6 @@ func (m *Manager) UpdateToken(token *relayAuth.Token) error {
 	return m.tokenStore.UpdateToken(token)
 }
 
-func (m *Manager) openConnVia(ctx context.Context, serverAddress, peerKey string, serverIP netip.Addr) (net.Conn, error) {
-	// check if already has a connection to the desired relay server
-	m.relayClientsMutex.RLock()
-	rt, ok := m.relayClients[serverAddress]
-	m.relayClientsMutex.RUnlock()
-	if ok {
-		return m.openConnOnTrack(ctx, rt, peerKey)
-	}
-
-	// if not, establish a new connection but check it again (because changed the lock type) before starting the
-	// connection
-	m.relayClientsMutex.Lock()
-	rt, ok = m.relayClients[serverAddress]
-	if ok {
-		m.relayClientsMutex.Unlock()
-		return m.openConnOnTrack(ctx, rt, peerKey)
-	}
-
-	// Publish the track and release the map lock BEFORE dialing, so the dial does
-	// not run under rt.Lock (which would block RelayStates and the cleanup loop
-	// for the full dial). Concurrent callers find this track and wait on rt.ready.
-	rt = NewRelayTrack()
-	m.relayClients[serverAddress] = rt
-	m.relayClientsMutex.Unlock()
-
-	relayClient := NewClientWithServerIP(serverAddress, serverIP, m.tokenStore, m.peerID, m.mtu)
-	relayClient.SetTransportFallback(m.transportFallback)
-	err := relayClient.Connect(m.ctx)
-	if err != nil {
-		rt.Lock()
-		rt.err = err
-		rt.Unlock()
-		close(rt.ready)
-		m.relayClientsMutex.Lock()
-		delete(m.relayClients, serverAddress)
-		m.relayClientsMutex.Unlock()
-		return nil, err
-	}
-	// if connection closed then delete the relay client from the list
-	relayClient.SetOnDisconnectListener(m.onServerDisconnected)
-	rt.Lock()
-	rt.relayClient = relayClient
-	rt.Unlock()
-	close(rt.ready)
-
-	return relayClient.OpenConn(ctx, peerKey)
-}
-
-// openConnOnTrack opens a peer connection through an existing relay track,
-// waiting for the dial started by another openConnVia call to finish. It waits
-// on rt.ready rather than the track lock, so it neither holds nor contends the
-// track lock across the dial.
-func (m *Manager) openConnOnTrack(ctx context.Context, rt *RelayTrack, peerKey string) (net.Conn, error) {
-	select {
-	case <-rt.ready:
-	case <-ctx.Done():
-		return nil, ctx.Err()
-	}
-
-	rt.RLock()
-	defer rt.RUnlock()
-	if rt.err != nil {
-		return nil, rt.err
-	}
-	if rt.relayClient == nil {
-		return nil, ErrRelayClientNotConnected
-	}
-	return rt.relayClient.OpenConn(ctx, peerKey)
-}
-
 func (m *Manager) onServerConnected() {
 	m.listenerLock.Lock()
 	defer m.listenerLock.Unlock()
@@ -422,21 +302,12 @@ func (m *Manager) onServerDisconnected(serverAddress string) {
 	m.relayClientMu.Unlock()
 
 	if !isHome {
-		m.evictForeignRelay(serverAddress)
+		m.foreign.evict(serverAddress)
 	}
 
 	m.notifyOnDisconnectListeners(serverAddress)
 }
 
-func (m *Manager) evictForeignRelay(serverAddress string) {
-	m.relayClientsMutex.Lock()
-	defer m.relayClientsMutex.Unlock()
-	if _, ok := m.relayClients[serverAddress]; ok {
-		delete(m.relayClients, serverAddress)
-		log.Debugf("evicted disconnected foreign relay client: %s", serverAddress)
-	}
-}
-
 func (m *Manager) listenGuardEvent(ctx context.Context) {
 	for {
 		select {
@@ -459,8 +330,8 @@ func (m *Manager) storeClient(client *Client) {
 	m.relayClient.SetOnDisconnectListener(m.onServerDisconnected)
 }
 
-func (m *Manager) isForeignServer(address string) (bool, error) {
-	rAddr, err := m.relayClient.ServerInstanceURL()
+func (m *Manager) isForeignServer(relayClient *Client, address string) (bool, error) {
+	rAddr, err := relayClient.ServerInstanceURL()
 	if err != nil {
 		return false, fmt.Errorf("relay client not connected")
 	}
@@ -475,50 +346,11 @@ func (m *Manager) startCleanupLoop() {
 		case <-m.ctx.Done():
 			return
 		case <-ticker.C:
-			m.cleanUpUnusedRelays()
+			m.foreign.cleanupUnused()
 		}
 	}
 }
 
-func (m *Manager) cleanUpUnusedRelays() {
-	m.relayClientsMutex.Lock()
-	defer m.relayClientsMutex.Unlock()
-
-	for addr, rt := range m.relayClients {
-		rt.Lock()
-		// if the connection failed to the server the relay client will be nil
-		// but the instance will be kept in the relayClients until the next locking
-		if rt.err != nil {
-			rt.Unlock()
-			continue
-		}
-
-		// dial still in progress (openConnVia publishes the track before Connect
-		// completes and no longer holds rt.Lock during it), nothing to clean up.
-		if rt.relayClient == nil {
-			rt.Unlock()
-			continue
-		}
-
-		if time.Since(rt.created) <= m.keepUnusedServerTime {
-			rt.Unlock()
-			continue
-		}
-
-		if rt.relayClient.HasConns() {
-			rt.Unlock()
-			continue
-		}
-		rt.relayClient.SetOnDisconnectListener(nil)
-		go func() {
-			_ = rt.relayClient.Close()
-		}()
-		log.Debugf("clean up unused relay server connection: %s", addr)
-		delete(m.relayClients, addr)
-		rt.Unlock()
-	}
-}
-
 func (m *Manager) addListener(serverAddress string, onClosedListener OnServerCloseListener) {
 	m.listenerLock.Lock()
 	defer m.listenerLock.Unlock()

shared/relay/client/manager_cleanup_test.go

@@ -2,17 +2,14 @@ package client
 
 import (
 	"context"
-	"net/netip"
 	"testing"
 	"time"
-
-	"github.com/stretchr/testify/require"
 )
 
 // TestCleanUpUnusedRelays_DoesNotBlockOnRealHangingDial drives a real, hanging foreign
-// relay dial and asserts cleanUpUnusedRelays does not stall behind it.
+// relay dial and asserts the foreign store cleanup does not stall behind it.
 func TestCleanUpUnusedRelays_DoesNotBlockOnRealHangingDial(t *testing.T) {
-	serverAddr := stallingRelayListener(t)
+	serverAddr, accepted := stallingRelayListener(t)
 
 	mCtx, mCancel := context.WithCancel(context.Background())
 	t.Cleanup(mCancel)
@@ -22,39 +19,32 @@ func TestCleanUpUnusedRelays_DoesNotBlockOnRealHangingDial(t *testing.T) {
 	dialDone := make(chan struct{})
 	go func() {
 		defer close(dialDone)
-		_, _ = m.openConnVia(mCtx, serverAddr, "peerKey", netip.Addr{})
+		_, _ = m.foreign.OpenConn(mCtx, "peerKey", RelayServer{Addr: serverAddr})
 	}()
 
-	// The track appears in the map once the dial is in flight.
-	require.Eventually(t, func() bool {
-		m.relayClientsMutex.RLock()
-		defer m.relayClientsMutex.RUnlock()
-		_, ok := m.relayClients[serverAddr]
-		return ok
-	}, 5*time.Second, 5*time.Millisecond, "relay dial did not start")
+	select {
+	case <-accepted:
+	case <-time.After(5 * time.Second):
+		t.Fatal("relay dial did not reach the listener")
+	}
 
 	cleanupDone := make(chan struct{})
 	go func() {
 		defer close(cleanupDone)
-		m.cleanUpUnusedRelays()
+		m.foreign.cleanupUnused()
 	}()
 
 	select {
 	case <-cleanupDone:
 	case <-time.After(2 * time.Second):
-		t.Fatal("cleanUpUnusedRelays blocked on an in-progress relay dial while holding the relay map lock")
+		t.Fatal("cleanupUnused blocked on an in-progress relay dial")
 	}
 
-	m.relayClientsMutex.RLock()
-	_, stillTracked := m.relayClients[serverAddr]
-	m.relayClientsMutex.RUnlock()
-	require.True(t, stillTracked, "an in-progress relay dial must not be evicted by cleanup")
-
 	// Release the hanging dial so the goroutine can exit cleanly.
 	mCancel()
 	select {
 	case <-dialDone:
 	case <-time.After(5 * time.Second):
-		t.Fatal("openConnVia did not return after context cancellation")
+		t.Fatal("foreign OpenConn did not return after context cancellation")
 	}
 }

shared/relay/client/manager_relaystates_test.go

@@ -3,7 +3,6 @@ package client
 import (
 	"context"
 	"net"
-	"net/netip"
 	"sync"
 	"testing"
 	"time"
@@ -13,13 +12,16 @@ import (
 
 // stallingRelayListener accepts TCP connections and holds them open without ever
 // responding, so a relay handshake dialed against it blocks until its context is
-// cancelled. It returns the "rel://host:port" URL to dial.
-func stallingRelayListener(t *testing.T) string {
+// cancelled. accepted is signalled once per incoming connection so a caller can
+// wait until a dial has actually reached the listener. It returns the
+// "rel://host:port" URL to dial.
+func stallingRelayListener(t *testing.T) (string, <-chan struct{}) {
 	t.Helper()
 
 	ln, err := net.Listen("tcp", "127.0.0.1:0")
 	require.NoError(t, err)
 
+	accepted := make(chan struct{}, 1)
 	var mu sync.Mutex
 	var conns []net.Conn
 	go func() {
@@ -31,6 +33,10 @@ func stallingRelayListener(t *testing.T) string {
 			mu.Lock()
 			conns = append(conns, c)
 			mu.Unlock()
+			select {
+			case accepted <- struct{}{}:
+			default:
+			}
 		}
 	}()
 	t.Cleanup(func() {
@@ -42,14 +48,14 @@ func stallingRelayListener(t *testing.T) string {
 		mu.Unlock()
 	})
 
-	return "rel://" + ln.Addr().String()
+	return "rel://" + ln.Addr().String(), accepted
 }
 
 // TestRelayStates_DoesNotBlockOnRealHangingDial is a regression test for
 // RelayStates() called by a "status -d command" hanging behind an in-progress
-// relay dial.
+// foreign relay dial.
 func TestRelayStates_DoesNotBlockOnRealHangingDial(t *testing.T) {
-	serverAddr := stallingRelayListener(t)
+	serverAddr, accepted := stallingRelayListener(t)
 
 	mCtx, mCancel := context.WithCancel(context.Background())
 	t.Cleanup(mCancel)
@@ -59,15 +65,14 @@ func TestRelayStates_DoesNotBlockOnRealHangingDial(t *testing.T) {
 	dialDone := make(chan struct{})
 	go func() {
 		defer close(dialDone)
-		_, _ = m.openConnVia(mCtx, serverAddr, "peerKey", netip.Addr{})
+		_, _ = m.foreign.OpenConn(mCtx, "peerKey", RelayServer{Addr: serverAddr})
 	}()
 
-	require.Eventually(t, func() bool {
-		m.relayClientsMutex.RLock()
-		defer m.relayClientsMutex.RUnlock()
-		_, ok := m.relayClients[serverAddr]
-		return ok
-	}, 5*time.Second, 5*time.Millisecond, "relay dial did not start")
+	select {
+	case <-accepted:
+	case <-time.After(5 * time.Second):
+		t.Fatal("relay dial did not reach the listener")
+	}
 
 	done := make(chan []RelayConnState, 1)
 	go func() {
@@ -86,6 +91,6 @@ func TestRelayStates_DoesNotBlockOnRealHangingDial(t *testing.T) {
 	select {
 	case <-dialDone:
 	case <-time.After(5 * time.Second):
-		t.Fatal("openConnVia did not return after context cancellation")
+		t.Fatal("foreign OpenConn did not return after context cancellation")
 	}
 }

shared/relay/client/manager_serverip_test.go

@@ -3,7 +3,6 @@ package client
 import (
 	"context"
 	"io"
-	"net/netip"
 	"testing"
 	"time"
 
@@ -85,7 +84,7 @@ func TestManager_ForeignRelayServerIP(t *testing.T) {
 	t.Run("no server IP, dial fails", func(t *testing.T) {
 		dialCtx, dialCancel := context.WithTimeout(ctx, 5*time.Second)
 		defer dialCancel()
-		_, err := mgrAlice.OpenConn(dialCtx, brokenFQDN, "bob", netip.Addr{})
+		_, err := mgrAlice.OpenConn(dialCtx, RelayServer{Addr: brokenFQDN}, "bob", true)
 		if err == nil {
 			t.Fatalf("expected OpenConn to fail without server IP, got success")
 		}
@@ -95,7 +94,7 @@ func TestManager_ForeignRelayServerIP(t *testing.T) {
 		// Bob waits for Alice's incoming peer connection on his side.
 		bobSideCh := make(chan error, 1)
 		go func() {
-			conn, err := mgrBob.OpenConn(ctx, bobRealAddr, "alice", netip.Addr{})
+			conn, err := mgrBob.OpenConn(ctx, RelayServer{Addr: bobRealAddr}, "alice", false)
 			if err != nil {
 				bobSideCh <- err
 				return
@@ -113,7 +112,7 @@ func TestManager_ForeignRelayServerIP(t *testing.T) {
 			bobSideCh <- nil
 		}()
 
-		aliceConn, err := mgrAlice.OpenConn(ctx, brokenFQDN, "bob", bobAdvertisedIP)
+		aliceConn, err := mgrAlice.OpenConn(ctx, RelayServer{Addr: brokenFQDN, IP: bobAdvertisedIP}, "bob", true)
 		if err != nil {
 			t.Fatalf("alice OpenConn with server IP: %s", err)
 		}

shared/relay/client/manager_test.go

@@ -3,7 +3,6 @@ package client
 import (
 	"context"
 	"fmt"
-	"net/netip"
 	"testing"
 	"time"
 
@@ -106,11 +105,11 @@ func TestForeignConn(t *testing.T) {
 	if err != nil {
 		t.Fatalf("failed to get relay address: %s", err)
 	}
-	connAliceToBob, err := clientAlice.OpenConn(ctx, bobsSrvAddr, "bob", netip.Addr{})
+	connAliceToBob, err := clientAlice.OpenConn(ctx, RelayServer{Addr: bobsSrvAddr}, "bob", true)
 	if err != nil {
 		t.Fatalf("failed to bind channel: %s", err)
 	}
-	connBobToAlice, err := clientBob.OpenConn(ctx, bobsSrvAddr, "alice", netip.Addr{})
+	connBobToAlice, err := clientBob.OpenConn(ctx, RelayServer{Addr: bobsSrvAddr}, "alice", false)
 	if err != nil {
 		t.Fatalf("failed to bind channel: %s", err)
 	}
@@ -210,7 +209,7 @@ func TestForeginConnClose(t *testing.T) {
 	if err != nil {
 		t.Fatalf("failed to serve manager: %s", err)
 	}
-	conn, err := mgr.OpenConn(ctx, toURL(srvCfg2)[0], "bob", netip.Addr{})
+	conn, err := mgr.OpenConn(ctx, RelayServer{Addr: toURL(srvCfg2)[0]}, "bob", true)
 	if err != nil {
 		t.Fatalf("failed to bind channel: %s", err)
 	}
@@ -302,7 +301,7 @@ func TestForeignAutoClose(t *testing.T) {
 	}
 
 	t.Log("open connection to another peer")
-	if _, err = mgr.OpenConn(ctx, foreignServerURL, "anotherpeer", netip.Addr{}); err == nil {
+	if _, err = mgr.OpenConn(ctx, RelayServer{Addr: foreignServerURL}, "anotherpeer", true); err == nil {
 		t.Fatalf("should have failed to open connection to another peer")
 	}
 
@@ -372,7 +371,7 @@ func TestAutoReconnect(t *testing.T) {
 	if err != nil {
 		t.Errorf("failed to get relay address: %s", err)
 	}
-	conn, err := clientAlice.OpenConn(ctx, ra, "bob", netip.Addr{})
+	conn, err := clientAlice.OpenConn(ctx, RelayServer{Addr: ra}, "bob", false)
 	if err != nil {
 		t.Errorf("failed to bind channel: %s", err)
 	}
@@ -392,7 +391,7 @@ func TestAutoReconnect(t *testing.T) {
 	}
 
 	log.Infof("reopent the connection")
-	_, err = clientAlice.OpenConn(ctx, ra, "bob", netip.Addr{})
+	_, err = clientAlice.OpenConn(ctx, RelayServer{Addr: ra}, "bob", false)
 	if err != nil {
 		t.Errorf("failed to open channel: %s", err)
 	}
@@ -454,7 +453,7 @@ func TestNotifierDoubleAdd(t *testing.T) {
 		t.Fatalf("failed to serve manager: %s", err)
 	}
 
-	conn1, err := clientAlice.OpenConn(ctx, clientAlice.ServerURLs()[0], "bob", netip.Addr{})
+	conn1, err := clientAlice.OpenConn(ctx, RelayServer{Addr: clientAlice.ServerURLs()[0]}, "bob", false)
 	if err != nil {
 		t.Fatalf("failed to bind channel: %s", err)
 	}