Move Subnet Functionality to its Own File

beacon-chain/p2p/BUILD.bazel

@@ -20,6 +20,7 @@ go_library(
         "rpc_topic_mappings.go",
         "sender.go",
         "service.go",
+        "subnets.go",
         "utils.go",
         "watch_peers.go",
     ],
@@ -89,6 +90,7 @@ go_test(
         "parameter_test.go",
         "sender_test.go",
         "service_test.go",
+        "subnets_test.go",
     ],
     embed = [":go_default_library"],
     flaky = True,

beacon-chain/p2p/discovery.go

@@ -14,12 +14,8 @@ import (
 	"github.com/libp2p/go-libp2p-core/peer"
 	ma "github.com/multiformats/go-multiaddr"
 	"github.com/pkg/errors"
-	"github.com/prysmaticlabs/go-bitfield"
 )
 
-const attestationSubnetCount = 64
-const attSubnetEnrKey = "attnets"
-
 // Listener defines the discovery V5 network interface that is used
 // to communicate with other peers.
 type Listener interface {
@@ -113,29 +109,6 @@ func startDHTDiscovery(host core.Host, bootstrapAddr string) error {
 	return err
 }
 
-func intializeAttSubnets(node *enode.LocalNode) *enode.LocalNode {
-	bitV := bitfield.NewBitvector64()
-	entry := enr.WithEntry(attSubnetEnrKey, bitV.Bytes())
-	node.Set(entry)
-	return node
-}
-
-func retrieveAttSubnets(record *enr.Record) ([]uint64, error) {
-	bitV := bitfield.NewBitvector64()
-	entry := enr.WithEntry(attSubnetEnrKey, &bitV)
-	err := record.Load(entry)
-	if err != nil {
-		return nil, err
-	}
-	committeeIdxs := []uint64{}
-	for i := uint64(0); i < 64; i++ {
-		if bitV.BitAt(i) {
-			committeeIdxs = append(committeeIdxs, i)
-		}
-	}
-	return committeeIdxs, nil
-}
-
 func parseBootStrapAddrs(addrs []string) (discv5Nodes []string, kadDHTNodes []string) {
 	discv5Nodes, kadDHTNodes = parseGenericAddrs(addrs)
 	if len(discv5Nodes) == 0 && len(kadDHTNodes) == 0 {

beacon-chain/p2p/discovery_test.go

@@ -13,13 +13,11 @@ import (
 
 	"github.com/ethereum/go-ethereum/p2p/discover"
 	"github.com/ethereum/go-ethereum/p2p/enode"
-	"github.com/ethereum/go-ethereum/p2p/enr"
 	"github.com/libp2p/go-libp2p-core/host"
-	"github.com/prysmaticlabs/go-bitfield"
-	"github.com/prysmaticlabs/prysm/beacon-chain/cache"
+	logTest "github.com/sirupsen/logrus/hooks/test"
+
 	"github.com/prysmaticlabs/prysm/shared/iputils"
 	"github.com/prysmaticlabs/prysm/shared/testutil"
-	logTest "github.com/sirupsen/logrus/hooks/test"
 )
 
 var discoveryWaitTime = 1 * time.Second
@@ -110,90 +108,6 @@ func TestStartDiscV5_DiscoverAllPeers(t *testing.T) {
 	}
 }
 
-func TestStartDiscV5_DiscoverPeersWithSubnets(t *testing.T) {
-	port := 2000
-	ipAddr, pkey := createAddrAndPrivKey(t)
-	bootListener := createListener(ipAddr, pkey, &Config{UDPPort: uint(port)})
-	defer bootListener.Close()
-
-	bootNode := bootListener.Self()
-	cfg := &Config{
-		BootstrapNodeAddr:   []string{bootNode.String()},
-		Discv5BootStrapAddr: []string{bootNode.String()},
-		Encoding:            "ssz",
-		MaxPeers:            30,
-	}
-	// Use shorter period for testing.
-	currentPeriod := pollingPeriod
-	pollingPeriod = 1 * time.Second
-	defer func() {
-		pollingPeriod = currentPeriod
-	}()
-
-	var listeners []*discover.UDPv5
-	for i := 1; i <= 3; i++ {
-		port = 3000 + i
-		cfg.UDPPort = uint(port)
-		ipAddr, pkey := createAddrAndPrivKey(t)
-		listener, err := startDiscoveryV5(ipAddr, pkey, cfg)
-		if err != nil {
-			t.Errorf("Could not start discovery for node: %v", err)
-		}
-		bitV := bitfield.NewBitvector64()
-		bitV.SetBitAt(uint64(i), true)
-
-		entry := enr.WithEntry(attSubnetEnrKey, &bitV)
-		listener.LocalNode().Set(entry)
-		listeners = append(listeners, listener)
-	}
-
-	// Make one service on port 3001.
-	port = 4000
-	cfg.UDPPort = uint(port)
-	s, err := NewService(cfg)
-	if err != nil {
-		t.Fatal(err)
-	}
-	s.Start()
-	defer s.Stop()
-
-	// Wait for the nodes to have their local routing tables to be populated with the other nodes
-	time.Sleep(discoveryWaitTime)
-
-	// look up 3 different subnets
-	exists, err := s.FindPeersWithSubnet(1)
-	if err != nil {
-		t.Fatal(err)
-	}
-	exists2, err := s.FindPeersWithSubnet(2)
-	if err != nil {
-		t.Fatal(err)
-	}
-	exists3, err := s.FindPeersWithSubnet(3)
-	if err != nil {
-		t.Fatal(err)
-	}
-	if !exists || !exists2 || !exists3 {
-		t.Fatal("Peer with subnet doesn't exist")
-	}
-
-	// update ENR of a peer
-	testService := &Service{dv5Listener: listeners[0]}
-	cache.CommitteeIDs.AddIDs([]uint64{10}, 0)
-	testService.RefreshENR(0)
-	time.Sleep(2 * time.Second)
-
-	exists, err = s.FindPeersWithSubnet(2)
-	if err != nil {
-		t.Fatal(err)
-	}
-
-	if !exists {
-		t.Fatal("Peer with subnet doesn't exist")
-	}
-
-}
-
 func TestMultiAddrsConversion_InvalidIPAddr(t *testing.T) {
 	addr := net.ParseIP("invalidIP")
 	_, pkey := createAddrAndPrivKey(t)

beacon-chain/p2p/subnets.go

@@ -0,0 +1,33 @@
+package p2p
+
+import (
+	"github.com/ethereum/go-ethereum/p2p/enode"
+	"github.com/ethereum/go-ethereum/p2p/enr"
+	"github.com/prysmaticlabs/go-bitfield"
+)
+
+const attestationSubnetCount = 64
+const attSubnetEnrKey = "attnets"
+
+func intializeAttSubnets(node *enode.LocalNode) *enode.LocalNode {
+	bitV := bitfield.NewBitvector64()
+	entry := enr.WithEntry(attSubnetEnrKey, bitV.Bytes())
+	node.Set(entry)
+	return node
+}
+
+func retrieveAttSubnets(record *enr.Record) ([]uint64, error) {
+	bitV := bitfield.NewBitvector64()
+	entry := enr.WithEntry(attSubnetEnrKey, &bitV)
+	err := record.Load(entry)
+	if err != nil {
+		return nil, err
+	}
+	committeeIdxs := []uint64{}
+	for i := uint64(0); i < 64; i++ {
+		if bitV.BitAt(i) {
+			committeeIdxs = append(committeeIdxs, i)
+		}
+	}
+	return committeeIdxs, nil
+}

beacon-chain/p2p/subnets_test.go

@@ -0,0 +1,95 @@
+package p2p
+
+import (
+	"testing"
+	"time"
+
+	"github.com/ethereum/go-ethereum/p2p/discover"
+	"github.com/ethereum/go-ethereum/p2p/enr"
+	"github.com/prysmaticlabs/go-bitfield"
+	"github.com/prysmaticlabs/prysm/beacon-chain/cache"
+)
+
+func TestStartDiscV5_DiscoverPeersWithSubnets(t *testing.T) {
+	port := 2000
+	ipAddr, pkey := createAddrAndPrivKey(t)
+	bootListener := createListener(ipAddr, pkey, &Config{UDPPort: uint(port)})
+	defer bootListener.Close()
+
+	bootNode := bootListener.Self()
+	cfg := &Config{
+		BootstrapNodeAddr:   []string{bootNode.String()},
+		Discv5BootStrapAddr: []string{bootNode.String()},
+		Encoding:            "ssz",
+		MaxPeers:            30,
+	}
+	// Use shorter period for testing.
+	currentPeriod := pollingPeriod
+	pollingPeriod = 1 * time.Second
+	defer func() {
+		pollingPeriod = currentPeriod
+	}()
+
+	var listeners []*discover.UDPv5
+	for i := 1; i <= 3; i++ {
+		port = 3000 + i
+		cfg.UDPPort = uint(port)
+		ipAddr, pkey := createAddrAndPrivKey(t)
+		listener, err := startDiscoveryV5(ipAddr, pkey, cfg)
+		if err != nil {
+			t.Errorf("Could not start discovery for node: %v", err)
+		}
+		bitV := bitfield.NewBitvector64()
+		bitV.SetBitAt(uint64(i), true)
+
+		entry := enr.WithEntry(attSubnetEnrKey, &bitV)
+		listener.LocalNode().Set(entry)
+		listeners = append(listeners, listener)
+	}
+
+	// Make one service on port 3001.
+	port = 4000
+	cfg.UDPPort = uint(port)
+	s, err := NewService(cfg)
+	if err != nil {
+		t.Fatal(err)
+	}
+	s.Start()
+	defer s.Stop()
+
+	// Wait for the nodes to have their local routing tables to be populated with the other nodes
+	time.Sleep(discoveryWaitTime)
+
+	// look up 3 different subnets
+	exists, err := s.FindPeersWithSubnet(1)
+	if err != nil {
+		t.Fatal(err)
+	}
+	exists2, err := s.FindPeersWithSubnet(2)
+	if err != nil {
+		t.Fatal(err)
+	}
+	exists3, err := s.FindPeersWithSubnet(3)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !exists || !exists2 || !exists3 {
+		t.Fatal("Peer with subnet doesn't exist")
+	}
+
+	// update ENR of a peer.
+	testService := &Service{dv5Listener: listeners[0]}
+	cache.CommitteeIDs.AddIDs([]uint64{10}, 0)
+	testService.RefreshENR(0)
+	time.Sleep(2 * time.Second)
+
+	exists, err = s.FindPeersWithSubnet(2)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if !exists {
+		t.Fatal("Peer with subnet doesn't exist")
+	}
+
+}