feat: limit number of light nodes (#1898)

pkg/p2p/libp2p/connections_test.go

@@ -18,6 +18,8 @@ import (
 	"github.com/ethersphere/bee/pkg/p2p/libp2p/internal/handshake"
 	"github.com/ethersphere/bee/pkg/statestore/mock"
 	"github.com/ethersphere/bee/pkg/swarm"
+	"github.com/ethersphere/bee/pkg/swarm/test"
+	"github.com/ethersphere/bee/pkg/topology/lightnode"
 	"github.com/libp2p/go-libp2p-core/mux"
 	libp2ppeer "github.com/libp2p/go-libp2p-core/peer"
 	ma "github.com/multiformats/go-multiaddr"
@@ -82,6 +84,46 @@ func TestConnectToLightPeer(t *testing.T) {
 	expectPeersEventually(t, s1)
 }
 
+func TestLightPeerLimit(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	var (
+		limit     = 3
+		container = lightnode.NewContainer(test.RandomAddress())
+		sf, _     = newService(t, 1, libp2pServiceOpts{lightNodes: container,
+			libp2pOpts: libp2p.Options{
+				LightNodeLimit: limit,
+				FullNode:       true,
+			}})
+
+		notifier = mockNotifier(noopCf, noopDf, true)
+	)
+	sf.SetPickyNotifier(notifier)
+
+	addr := serviceUnderlayAddress(t, sf)
+
+	for i := 0; i < 5; i++ {
+		sl, _ := newService(t, 1, libp2pServiceOpts{
+			libp2pOpts: libp2p.Options{
+				FullNode: false,
+			}})
+		_, err := sl.Connect(ctx, addr)
+		if err != nil {
+			t.Fatal(err)
+		}
+	}
+
+	for i := 0; i < 20; i++ {
+		if cnt := container.Count(); cnt == limit {
+			return
+		}
+		time.Sleep(50 * time.Millisecond)
+	}
+
+	t.Fatal("timed out waiting for correct number of lightnodes")
+}
+
 func TestDoubleConnect(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
@@ -754,3 +796,9 @@ func mockNotifier(c cFunc, d dFunc, pick bool) p2p.PickyNotifier {
 
 type cFunc func(context.Context, p2p.Peer) error
 type dFunc func(p2p.Peer)
+
+var noopCf = func(_ context.Context, _ p2p.Peer) error {
+	return nil
+}
+
+var noopDf = func(p p2p.Peer) {}

pkg/p2p/libp2p/libp2p_test.go

@@ -31,6 +31,7 @@ type libp2pServiceOpts struct {
 	PrivateKey  *ecdsa.PrivateKey
 	MockPeerKey *ecdsa.PrivateKey
 	libp2pOpts  libp2p.Options
+	lightNodes  *lightnode.Container
 }
 
 // newService constructs a new libp2p service.
@@ -69,14 +70,15 @@ func newService(t *testing.T, networkID uint64, o libp2pServiceOpts) (s *libp2p.
 
 	ctx, cancel := context.WithCancel(context.Background())
 
-	lightnodes := lightnode.NewContainer(overlay)
-
+	if o.lightNodes == nil {
+		o.lightNodes = lightnode.NewContainer(overlay)
+	}
 	opts := o.libp2pOpts
 	opts.Transaction = []byte(hexutil.EncodeUint64(o.PrivateKey.Y.Uint64()))
 
 	senderMatcher := &MockSenderMatcher{}
 
-	s, err = libp2p.New(ctx, crypto.NewDefaultSigner(swarmKey), networkID, overlay, addr, o.Addressbook, statestore, lightnodes, senderMatcher, o.Logger, nil, opts)
+	s, err = libp2p.New(ctx, crypto.NewDefaultSigner(swarmKey), networkID, overlay, addr, o.Addressbook, statestore, o.lightNodes, senderMatcher, o.Logger, nil, opts)
 	if err != nil {
 		t.Fatal(err)
 	}

pkg/topology/lightnode/container.go

@@ -6,6 +6,8 @@ package lightnode
 
 import (
 	"context"
+	"crypto/rand"
+	"math/big"
 	"sync"
 
 	"github.com/ethersphere/bee/pkg/p2p"
@@ -49,6 +51,43 @@ func (c *Container) Disconnected(peer p2p.Peer) {
 	}
 }
 
+func (c *Container) Count() int {
+	return c.connectedPeers.Length()
+}
+
+func (c *Container) RandomPeer(not swarm.Address) (swarm.Address, error) {
+	c.peerMu.Lock()
+	defer c.peerMu.Unlock()
+	var (
+		cnt   = big.NewInt(int64(c.Count()))
+		addr  = swarm.ZeroAddress
+		count = int64(0)
+	)
+
+PICKPEER:
+	i, e := rand.Int(rand.Reader, cnt)
+	if e != nil {
+		return swarm.ZeroAddress, e
+	}
+	i64 := i.Int64()
+
+	count = 0
+	_ = c.connectedPeers.EachBinRev(func(peer swarm.Address, _ uint8) (bool, bool, error) {
+		if count == i64 {
+			addr = peer
+			return true, false, nil
+		}
+		count++
+		return false, false, nil
+	})
+
+	if addr.Equal(not) {
+		goto PICKPEER
+	}
+
+	return addr, nil
+}
+
 func (c *Container) PeerInfo() topology.BinInfo {
 	return topology.BinInfo{
 		BinPopulation:     uint(c.connectedPeers.Length()),

pkg/topology/lightnode/container_test.go

@@ -33,14 +33,35 @@ func TestContainer(t *testing.T) {
 	t.Run("can add peers to container", func(t *testing.T) {
 		c := lightnode.NewContainer(base)
 
-		c.Connected(context.Background(), p2p.Peer{Address: swarm.NewAddress([]byte("123"))})
-		c.Connected(context.Background(), p2p.Peer{Address: swarm.NewAddress([]byte("456"))})
+		p1 := swarm.NewAddress([]byte("123"))
+		p2 := swarm.NewAddress([]byte("456"))
+		c.Connected(context.Background(), p2p.Peer{Address: p1})
+		c.Connected(context.Background(), p2p.Peer{Address: p2})
 
 		peerCount := len(c.PeerInfo().ConnectedPeers)
 
 		if peerCount != 2 {
 			t.Errorf("expected %d connected peer, got %d", 2, peerCount)
 		}
+
+		if cc := c.Count(); cc != 2 {
+			t.Errorf("expected count 2 got %d", cc)
+		}
+		p, err := c.RandomPeer(p1)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if !p.Equal(p2) {
+			t.Fatalf("expected p1 but got %s", p.String())
+		}
+
+		p, err = c.RandomPeer(p2)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if !p.Equal(p1) {
+			t.Fatalf("expected p2 but got %s", p.String())
+		}
 	})
 	t.Run("empty container after peer disconnect", func(t *testing.T) {
 		c := lightnode.NewContainer(base)
@@ -59,5 +80,8 @@ func TestContainer(t *testing.T) {
 		if connPeerCount != 0 {
 			t.Errorf("expected %d connected peer, got %d", 0, connPeerCount)
 		}
+		if cc := c.Count(); cc != 0 {
+			t.Errorf("expected count 0 got %d", cc)
+		}
 	})
 }