[Merged by Bors] - fetch: smarter peer selection

fetch/fetch.go

@@ -12,6 +12,7 @@
 	"time"
 
 	"github.com/libp2p/go-libp2p/core/network"
+	"github.com/libp2p/go-libp2p/core/protocol"
 	"go.uber.org/zap"
 	"golang.org/x/sync/errgroup"
 	"golang.org/x/sync/semaphore"
@@ -297,7 +298,16 @@
 	// there is one test that covers this part.
 	if host != nil {
 		connectedf := func(peer p2p.Peer) {
-			if f.peers.Add(peer) {
+			protocols := func() []protocol.ID {
+				ps, err := host.Peerstore().GetProtocols(peer)
+				if err != nil {
+					f.logger.Debug("failed to get protocols for peer",
+						zap.Stringer("id", peer), zap.Error(err))
+					return nil
+				}
+				return ps
+			}
+			if f.peers.Add(peer, protocols) {
 				f.logger.Debug("adding peer", zap.Stringer("id", peer))
 			}
 		}
@@ -703,7 +713,9 @@
 	rng := rand.New(rand.NewChaCha8(seed))
 	peer2requests := make(map[p2p.Peer][]RequestMessage)
 
-	best := f.peers.SelectBest(RedundantPeers)
+	// When selecting peers, provide protocol IDs so that peers that aren't yet fully
+	// initialized are not picked for the request, avoiding unnecessary errors.
+	best := f.peers.SelectBestWithProtocols(RedundantPeers, []protocol.ID{hashProtocol, activeSetProtocol})
 	if len(best) == 0 {
 		f.logger.Warn("cannot send batch: no peers found")
 		f.mu.Lock()

fetch/fetch_test.go

@@ -7,6 +7,7 @@ import (
 	"time"
 
 	"github.com/libp2p/go-libp2p/core/peer"
+	"github.com/libp2p/go-libp2p/core/protocol"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"go.uber.org/mock/gomock"
@@ -187,7 +188,9 @@ func TestFetch_RequestHashBatchFromPeers(t *testing.T) {
 			f := createFetch(t)
 			f.cfg.MaxRetriesForRequest = 0
 			peer := p2p.Peer("buddy")
-			f.peers.Add(peer)
+			f.peers.Add(peer, func() []protocol.ID {
+				return []protocol.ID{hashProtocol, activeSetProtocol}
+			})
 
 			hsh0 := types.RandomHash()
 			res0 := ResponseMessage{
@@ -259,8 +262,9 @@ func TestFetch_Loop_BatchRequestMax(t *testing.T) {
 	f.cfg.BatchTimeout = 1
 	f.cfg.BatchSize = 2
 	peer := p2p.Peer("buddy")
-	f.peers.Add(peer)
-
+	f.peers.Add(peer, func() []protocol.ID {
+		return []protocol.ID{hashProtocol, activeSetProtocol}
+	})
 	h1 := types.RandomHash()
 	h2 := types.RandomHash()
 	h3 := types.RandomHash()

fetch/mesh_data_test.go

@@ -8,6 +8,7 @@ import (
 	"testing"
 
 	p2phost "github.com/libp2p/go-libp2p/core/host"
+	"github.com/libp2p/go-libp2p/core/protocol"
 	mocknet "github.com/libp2p/go-libp2p/p2p/net/mock"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -137,7 +138,9 @@ func TestFetch_getHashes(t *testing.T) {
 		f.Start()
 		tb.Cleanup(f.Stop)
 		for _, peer := range peers {
-			f.peers.Add(peer)
+			f.peers.Add(peer, func() []protocol.ID {
+				return []protocol.ID{hashProtocol, activeSetProtocol}
+			})
 		}
 		f.mh.EXPECT().ID().Return("self").AnyTimes()
 		f.RegisterPeerHashes(peers[0], hashes[:2])
@@ -249,7 +252,9 @@ func TestFetch_getHashesStreaming(t *testing.T) {
 		f.Start()
 		tb.Cleanup(f.Stop)
 		for _, peer := range peers {
-			f.peers.Add(peer)
+			f.peers.Add(peer, func() []protocol.ID {
+				return []protocol.ID{hashProtocol, activeSetProtocol}
+			})
 		}
 		f.mh.EXPECT().ID().Return("self").AnyTimes()
 		f.RegisterPeerHashes(peers[0], hashes[:2])

fetch/peers/peers.go

@@ -1,11 +1,13 @@
 package peers
 
 import (
+	"slices"
 	"strings"
 	"sync"
 	"time"
 
 	"github.com/libp2p/go-libp2p/core/peer"
+	"github.com/libp2p/go-libp2p/core/protocol"
 	"go.uber.org/zap/zapcore"
 
 	"github.com/spacemeshos/go-spacemesh/p2p"
@@ -16,6 +18,7 @@ type data struct {
 	success, failures int
 	failRate          float64
 	averageLatency    float64
+	protocols         func() []protocol.ID
 }
 
 func (d *data) latency(global float64) float64 {
@@ -61,14 +64,14 @@ func (p *Peers) Contains(id peer.ID) bool {
 	return exist
 }
 
-func (p *Peers) Add(id peer.ID) bool {
+func (p *Peers) Add(id peer.ID, protocols func() []protocol.ID) bool {
 	p.mu.Lock()
 	defer p.mu.Unlock()
 	_, exist := p.peers[id]
 	if exist {
 		return false
 	}
-	p.peers[id] = &data{id: id}
+	p.peers[id] = &data{id: id, protocols: protocols}
 	return true
 }
 
@@ -151,12 +154,39 @@ func (p *Peers) SelectBestFrom(peers []peer.ID) peer.ID {
 func (p *Peers) SelectBest(n int) []peer.ID {
 	p.mu.Lock()
 	defer p.mu.Unlock()
+	return p.selectBest(n, nil)
+}
+
+// SelectBestWithProtocols is similar to SelectBest but filters peers by supported protocols.
+// If protocols is empty, it returns the best peers regardless of the protocol.
+// If protocols is not empty, it returns the best peers that support at least one of the protocols.
+func (p *Peers) SelectBestWithProtocols(n int, protocols []protocol.ID) []peer.ID {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	return p.selectBest(n, protocols)
+}
+
+func (p *Peers) selectBest(n int, protocols []protocol.ID) []peer.ID {
+	slices.Sort(protocols)
+	protocols = slices.Compact(protocols)
 	lth := min(len(p.peers), n)
 	if lth == 0 {
 		return nil
 	}
 	best := make([]*data, 0, lth)
 	for _, peer := range p.peers {
+		if len(protocols) > 0 {
+			found := false
+			for _, proto := range peer.protocols() {
+				if slices.Contains(protocols, proto) {
+					found = true
+					break
+				}
+			}
+			if !found {
+				continue
+			}
+		}
 		for i := range best {
 			if peer.less(best[i], p.globalLatency) {
 				best[i], peer = peer, best[i]

fetch/peers/peers_test.go

@@ -8,6 +8,7 @@ import (
 	"time"
 
 	"github.com/libp2p/go-libp2p/core/peer"
+	"github.com/libp2p/go-libp2p/core/protocol"
 	"github.com/stretchr/testify/require"
 )
 
@@ -22,6 +23,7 @@ type event struct {
 	success     int
 	failure     int
 	latency     time.Duration
+	protocols   []protocol.ID
 }
 
 func withEvents(events []event) *Peers {
@@ -30,7 +32,7 @@ func withEvents(events []event) *Peers {
 		if ev.delete {
 			tracker.Delete(ev.id)
 		} else if ev.add {
-			tracker.Add(ev.id)
+			tracker.Add(ev.id, func() []protocol.ID { return ev.protocols })
 		}
 		for i := 0; i < ev.failure; i++ {
 			tracker.OnFailure(ev.id, 0, ev.latency)
@@ -210,6 +212,73 @@ func TestSelect(t *testing.T) {
 	}
 }
 
+func TestSelectBestWithProtocols(t *testing.T) {
+	for _, tc := range []struct {
+		desc   string
+		events []event
+
+		n         int
+		protocols []protocol.ID
+		expect    []peer.ID
+	}{
+		{
+			desc: "no protocols required and no peer protocols",
+			events: []event{
+				{id: "a", success: 1, latency: 8, add: true},
+				{id: "b", success: 1, latency: 9, add: true},
+				{id: "c", success: 3, latency: 14, add: true},
+			},
+			n:         2,
+			expect:    []peer.ID{"a", "b"},
+			protocols: nil,
+		},
+		{
+			desc: "no protocols required, peers have protocols",
+			events: []event{
+				{id: "a", success: 1, latency: 8, add: true, protocols: []protocol.ID{"a", "b"}},
+				{id: "b", success: 1, latency: 9, add: true, protocols: []protocol.ID{"b", "c"}},
+				{id: "c", success: 3, latency: 14, add: true, protocols: []protocol.ID{"c", "d"}},
+			},
+			n:         2,
+			expect:    []peer.ID{"a", "b"},
+			protocols: nil,
+		},
+		{
+			desc: "single protocol required, peers have protocols",
+			events: []event{
+				{id: "a", success: 1, latency: 8, add: true, protocols: []protocol.ID{"a", "b"}},
+				{id: "b", success: 1, latency: 9, add: true, protocols: []protocol.ID{"b", "c"}},
+				{id: "c", success: 3, latency: 14, add: true, protocols: []protocol.ID{"c", "d"}},
+			},
+			n:         2,
+			expect:    []peer.ID{"b", "c"},
+			protocols: []protocol.ID{"c"},
+		},
+		{
+			desc: "multiple protocols required, peers have protocols",
+			events: []event{
+				{id: "a", success: 1, latency: 8, add: true, protocols: []protocol.ID{"a", "b"}},
+				{id: "b", success: 1, latency: 9, add: true, protocols: []protocol.ID{"b", "c"}},
+				{id: "c", success: 3, latency: 14, add: true, protocols: []protocol.ID{"c", "d"}},
+				{id: "d", success: 3, latency: 12, add: true, protocols: []protocol.ID{"a", "e"}},
+			},
+			n:         3,
+			expect:    []peer.ID{"a", "b", "c"},
+			protocols: []protocol.ID{"b", "c"},
+		},
+	} {
+		t.Run(tc.desc, func(t *testing.T) {
+			require.Equal(
+				t,
+				tc.expect,
+				withEvents(tc.events).SelectBestWithProtocols(tc.n, tc.protocols),
+				"select best %d",
+				tc.n,
+			)
+		})
+	}
+}
+
 func TestTotal(t *testing.T) {
 	const total = 100
 	events := []event{}

sync2/multipeer/multipeer.go

@@ -7,6 +7,7 @@ import (
 	"time"
 
 	"github.com/jonboulle/clockwork"
+	"github.com/libp2p/go-libp2p/core/protocol"
 	"go.uber.org/zap"
 	"golang.org/x/sync/errgroup"
 
@@ -15,6 +16,10 @@ import (
 	"github.com/spacemeshos/go-spacemesh/sync2/rangesync"
 )
 
+const (
+	Protocol = "sync/2"
+)
+
 type syncability struct {
 	// peers that were probed successfully
 	syncable []p2p.Peer
@@ -283,7 +288,7 @@ func (mpr *MultiPeerReconciler) fullSync(ctx context.Context, syncPeers []p2p.Pe
 func (mpr *MultiPeerReconciler) syncOnce(ctx context.Context, lastWasSplit bool) (full bool, err error) {
 	var s syncability
 	for {
-		syncPeers := mpr.peers.SelectBest(mpr.cfg.SyncPeerCount)
+		syncPeers := mpr.peers.SelectBestWithProtocols(mpr.cfg.SyncPeerCount, []protocol.ID{Protocol})
 		mpr.logger.Debug("selected best peers for sync",
 			zap.Int("syncPeerCount", mpr.cfg.SyncPeerCount),
 			zap.Int("totalPeers", mpr.peers.Total()),

sync2/multipeer/multipeer_test.go

@@ -10,6 +10,7 @@ import (
 	"time"
 
 	"github.com/jonboulle/clockwork"
+	"github.com/libp2p/go-libp2p/core/protocol"
 	"github.com/stretchr/testify/require"
 	"go.uber.org/mock/gomock"
 	"go.uber.org/zap/zaptest"
@@ -101,7 +102,7 @@ func (mt *multiPeerSyncTester) addPeers(n int) []p2p.Peer {
 	r := make([]p2p.Peer, n)
 	for i := 0; i < n; i++ {
 		p := p2p.Peer(fmt.Sprintf("peer%d", i+1))
-		mt.peers.Add(p)
+		mt.peers.Add(p, func() []protocol.ID { return []protocol.ID{multipeer.Protocol} })
 		r[i] = p
 	}
 	return r

sync2/multipeer/split_sync_test.go

@@ -8,6 +8,7 @@ import (
 	"time"
 
 	"github.com/jonboulle/clockwork"
+	"github.com/libp2p/go-libp2p/core/protocol"
 	"github.com/stretchr/testify/require"
 	gomock "go.uber.org/mock/gomock"
 	"go.uber.org/zap/zaptest"
@@ -108,7 +109,7 @@ func newTestSplitSync(t testing.TB) *splitSyncTester {
 			AnyTimes()
 	}
 	for _, p := range tst.syncPeers {
-		tst.peers.Add(p)
+		tst.peers.Add(p, func() []protocol.ID { return []protocol.ID{multipeer.Protocol} })
 	}
 	tst.splitSync = multipeer.NewSplitSync(
 		zaptest.NewLogger(t),

sync2/p2p_test.go

@@ -7,6 +7,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/libp2p/go-libp2p/core/protocol"
 	mocknet "github.com/libp2p/go-libp2p/p2p/net/mock"
 	"github.com/stretchr/testify/require"
 	"go.uber.org/zap/zaptest"
@@ -16,6 +17,7 @@ import (
 	"github.com/spacemeshos/go-spacemesh/p2p"
 	"github.com/spacemeshos/go-spacemesh/p2p/server"
 	"github.com/spacemeshos/go-spacemesh/sync2"
+	"github.com/spacemeshos/go-spacemesh/sync2/multipeer"
 	"github.com/spacemeshos/go-spacemesh/sync2/rangesync"
 )
 
@@ -88,7 +90,9 @@ func TestP2P(t *testing.T) {
 		ps := peers.New()
 		for m := 0; m < numNodes; m++ {
 			if m != n {
-				ps.Add(mesh.Hosts()[m].ID())
+				ps.Add(mesh.Hosts()[m].ID(), func() []protocol.ID {
+					return []protocol.ID{multipeer.Protocol}
+				})
 			}
 		}
 		cfg := sync2.DefaultConfig()