cmd/devp2p: add support for -limit option in nodeset filter command

cmd/devp2p/README.md

@@ -30,6 +30,29 @@ Run `devp2p dns to-route53 <directory>` to publish a tree to Amazon Route53.
 
 You can find more information about these commands in the [DNS Discovery Setup Guide][dns-tutorial].
 
+### Node Set Utilities
+
+There are several commands for working with JSON node set files. These files are generated
+by the discovery crawlers and DNS client commands. Node sets also used as the input of the
+DNS deployer commands.
+
+Run `devp2p nodeset info <nodes.json>` to display statistics of a node set.
+
+Run `devp2p nodeset filter <nodes.json> <filter flags...>` to write a new, filtered node
+set to standard output. The following filters are supported:
+
+- `-limit <N>` limits the output set to N entries, taking the top N nodes by score
+- `-ip <CIDR>` filters nodes by IP subnet
+- `-min-age <duration>` filters nodes by 'first seen' time
+- `-eth-network <mainnet/rinkeby/goerli/ropsten>` filters nodes by "eth" ENR entry
+- `-les-server` filters nodes by LES server support
+- `-snap` filters nodes by snap protocol support
+
+For example, given a node set in `nodes.json`, you could create a filtered set containing
+up to 20 eth mainnet nodes which also support snap sync using this command:
+
+    devp2p nodeset filter nodes.json -eth-network mainnet -snap -limit 20
+
 ### Discovery v4 Utilities
 
 The `devp2p discv4 ...` command family deals with the [Node Discovery v4][discv4]
@@ -94,7 +117,7 @@ To run the eth protocol test suite against your implementation, the node needs t
 geth --datadir <datadir> --nodiscover --nat=none --networkid 19763 --verbosity 5
 ```
 
-Then, run the following command, replacing `<enode>` with the enode of the geth node: 
+Then, run the following command, replacing `<enode>` with the enode of the geth node:
  ```
  devp2p rlpx eth-test <enode> cmd/devp2p/internal/ethtest/testdata/chain.rlp cmd/devp2p/internal/ethtest/testdata/genesis.json
 ```
@@ -103,7 +126,7 @@ Repeat the above process (re-initialising the node) in order to run the Eth Prot
 
 #### Eth66 Test Suite
 
-The Eth66 test suite is also a conformance test suite for the eth 66 protocol version specifically. 
+The Eth66 test suite is also a conformance test suite for the eth 66 protocol version specifically.
 To run the eth66 protocol test suite, initialize a geth node as described above and run the following command,
 replacing `<enode>` with the enode of the geth node:
 

cmd/devp2p/nodeset.go

@@ -71,6 +71,7 @@ func writeNodesJSON(file string, nodes nodeSet) {
 	}
 }
 
+// nodes returns the node records contained in the set.
 func (ns nodeSet) nodes() []*enode.Node {
 	result := make([]*enode.Node, 0, len(ns))
 	for _, n := range ns {
@@ -83,12 +84,37 @@ func (ns nodeSet) nodes() []*enode.Node {
 	return result
 }
 
+// add ensures the given nodes are present in the set.
 func (ns nodeSet) add(nodes ...*enode.Node) {
 	for _, n := range nodes {
-		ns[n.ID()] = nodeJSON{Seq: n.Seq(), N: n}
+		v := ns[n.ID()]
+		v.N = n
+		v.Seq = n.Seq()
+		ns[n.ID()] = v
 	}
 }
 
+// topN returns the top n nodes by score as a new set.
+func (ns nodeSet) topN(n int) nodeSet {
+	if n >= len(ns) {
+		return ns
+	}
+
+	byscore := make([]nodeJSON, 0, len(ns))
+	for _, v := range ns {
+		byscore = append(byscore, v)
+	}
+	sort.Slice(byscore, func(i, j int) bool {
+		return byscore[i].Score >= byscore[j].Score
+	})
+	result := make(nodeSet, n)
+	for _, v := range byscore[:n] {
+		result[v.N.ID()] = v
+	}
+	return result
+}
+
+// verify performs integrity checks on the node set.
 func (ns nodeSet) verify() error {
 	for id, n := range ns {
 		if n.N.ID() != id {

cmd/devp2p/nodesetcmd.go

@@ -17,8 +17,12 @@
 package main
 
 import (
+	"errors"
 	"fmt"
 	"net"
+	"sort"
+	"strconv"
+	"strings"
 	"time"
 
 	"github.com/ethereum/go-ethereum/core/forkid"
@@ -60,25 +64,64 @@ func nodesetInfo(ctx *cli.Context) error {
 
 	ns := loadNodesJSON(ctx.Args().First())
 	fmt.Printf("Set contains %d nodes.\n", len(ns))
+	showAttributeCounts(ns)
 	return nil
 }
 
+// showAttributeCounts prints the distribution of ENR attributes in a node set.
+func showAttributeCounts(ns nodeSet) {
+	attrcount := make(map[string]int)
+	var attrlist []interface{}
+	for _, n := range ns {
+		r := n.N.Record()
+		attrlist = r.AppendElements(attrlist[:0])[1:]
+		for i := 0; i < len(attrlist); i += 2 {
+			key := attrlist[i].(string)
+			attrcount[key]++
+		}
+	}
+
+	var keys []string
+	var maxlength int
+	for key := range attrcount {
+		keys = append(keys, key)
+		if len(key) > maxlength {
+			maxlength = len(key)
+		}
+	}
+	sort.Strings(keys)
+	fmt.Println("ENR attribute counts:")
+	for _, key := range keys {
+		fmt.Printf("%s%s: %d\n", strings.Repeat(" ", maxlength-len(key)+1), key, attrcount[key])
+	}
+}
+
 func nodesetFilter(ctx *cli.Context) error {
 	if ctx.NArg() < 1 {
 		return fmt.Errorf("need nodes file as argument")
 	}
-	ns := loadNodesJSON(ctx.Args().First())
+	// Parse -limit.
+	limit, err := parseFilterLimit(ctx.Args().Tail())
+	if err != nil {
+		return err
+	}
+	// Parse the filters.
 	filter, err := andFilter(ctx.Args().Tail())
 	if err != nil {
 		return err
 	}
 
+	// Load nodes and apply filters.
+	ns := loadNodesJSON(ctx.Args().First())
 	result := make(nodeSet)
 	for id, n := range ns {
 		if filter(n) {
 			result[id] = n
 		}
 	}
+	if limit >= 0 {
+		result = result.topN(limit)
+	}
 	writeNodesJSON("-", result)
 	return nil
 }
@@ -91,13 +134,15 @@ type nodeFilterC struct {
 }
 
 var filterFlags = map[string]nodeFilterC{
+	"-limit":       {1, trueFilter}, // needed to skip over -limit
 	"-ip":          {1, ipFilter},
 	"-min-age":     {1, minAgeFilter},
 	"-eth-network": {1, ethFilter},
 	"-les-server":  {0, lesFilter},
 	"-snap":        {0, snapFilter},
 }
 
+// parseFilters parses nodeFilters from args.
 func parseFilters(args []string) ([]nodeFilter, error) {
 	var filters []nodeFilter
 	for len(args) > 0 {
@@ -118,6 +163,26 @@ func parseFilters(args []string) ([]nodeFilter, error) {
 	return filters, nil
 }
 
+// parseFilterLimit parses the -limit option in args. It returns -1 if there is no limit.
+func parseFilterLimit(args []string) (int, error) {
+	limit := -1
+	for i, arg := range args {
+		if arg == "-limit" {
+			if i == len(args)-1 {
+				return -1, errors.New("-limit requires an argument")
+			}
+			n, err := strconv.Atoi(args[i+1])
+			if err != nil {
+				return -1, fmt.Errorf("invalid -limit %q", args[i+1])
+			}
+			limit = n
+		}
+	}
+	return limit, nil
+}
+
+// andFilter parses node filters in args and and returns a single filter that requires all
+// of them to match.
 func andFilter(args []string) (nodeFilter, error) {
 	checks, err := parseFilters(args)
 	if err != nil {
@@ -134,6 +199,10 @@ func andFilter(args []string) (nodeFilter, error) {
 	return f, nil
 }
 
+func trueFilter(args []string) (nodeFilter, error) {
+	return func(n nodeJSON) bool { return true }, nil
+}
+
 func ipFilter(args []string) (nodeFilter, error) {
 	_, cidr, err := net.ParseCIDR(args[0])
 	if err != nil {