p2p/discover: add ReadRandomNodes

2 files changed, 83 insertions, 1 deletions

diff --git a/p2p/discover/table.go b/p2p/discover/table.go
index 91d617f05..b523a0684 100644
--- a/p2p/discover/table.go
+++ b/p2p/discover/table.go
@@ -8,6 +8,7 @@ package discover
 
 import (
 	"crypto/rand"
+	"encoding/binary"
 	"net"
 	"sort"
 	"sync"
@@ -90,10 +91,58 @@ func newTable(t transport, ourID NodeID, ourAddr *net.UDPAddr, nodeDBPath string
 }
 
 // Self returns the local node.
+// The returned node should not be modified by the caller.
 func (tab *Table) Self() *Node {
 	return tab.self
 }
 
+// ReadRandomNodes fills the given slice with random nodes from the
+// table. It will not write the same node more than once. The nodes in
+// the slice are copies and can be modified by the caller.
+func (tab *Table) ReadRandomNodes(buf []*Node) (n int) {
+	tab.mutex.Lock()
+	defer tab.mutex.Unlock()
+	// TODO: tree-based buckets would help here
+	// Find all non-empty buckets and get a fresh slice of their entries.
+	var buckets [][]*Node
+	for _, b := range tab.buckets {
+		if len(b.entries) > 0 {
+			buckets = append(buckets, b.entries[:])
+		}
+	}
+	if len(buckets) == 0 {
+		return 0
+	}
+	// Shuffle the buckets.
+	for i := uint32(len(buckets)) - 1; i > 0; i-- {
+		j := randUint(i)
+		buckets[i], buckets[j] = buckets[j], buckets[i]
+	}
+	// Move head of each bucket into buf, removing buckets that become empty.
+	var i, j int
+	for ; i < len(buf); i, j = i+1, (j+1)%len(buckets) {
+		b := buckets[j]
+		buf[i] = &(*b[0])
+		buckets[j] = b[1:]
+		if len(b) == 1 {
+			buckets = append(buckets[:j], buckets[j+1:]...)
+		}
+		if len(buckets) == 0 {
+			break
+		}
+	}
+	return i + 1
+}
+
+func randUint(max uint32) uint32 {
+	if max == 0 {
+		return 0
+	}
+	var b [4]byte
+	rand.Read(b[:])
+	return binary.BigEndian.Uint32(b[:]) % max
+}
+
 // Close terminates the network listener and flushes the node database.
 func (tab *Table) Close() {
 	tab.net.close()
diff --git a/p2p/discover/table_test.go b/p2p/discover/table_test.go
index aa5267928..da398d137 100644
--- a/p2p/discover/table_test.go
+++ b/p2p/discover/table_test.go
@@ -210,6 +210,36 @@ func TestTable_closest(t *testing.T) {
 	}
 }
 
+func TestTable_ReadRandomNodesGetAll(t *testing.T) {
+	cfg := &quick.Config{
+		MaxCount: 200,
+		Rand:     quickrand,
+		Values: func(args []reflect.Value, rand *rand.Rand) {
+			args[0] = reflect.ValueOf(make([]*Node, rand.Intn(1000)))
+		},
+	}
+	test := func(buf []*Node) bool {
+		tab := newTable(nil, NodeID{}, &net.UDPAddr{}, "")
+		for i := 0; i < len(buf); i++ {
+			ld := quickrand.Intn(len(tab.buckets))
+			tab.add([]*Node{nodeAtDistance(tab.self.sha, ld)})
+		}
+		gotN := tab.ReadRandomNodes(buf)
+		if gotN != tab.len() {
+			t.Errorf("wrong number of nodes, got %d, want %d", gotN, tab.len())
+			return false
+		}
+		if hasDuplicates(buf[:gotN]) {
+			t.Errorf("result contains duplicates")
+			return false
+		}
+		return true
+	}
+	if err := quick.Check(test, cfg); err != nil {
+		t.Error(err)
+	}
+}
+
 type closeTest struct {
 	Self   NodeID
 	Target common.Hash
@@ -517,7 +547,10 @@ func (n *preminedTestnet) mine(target NodeID) {
 
 func hasDuplicates(slice []*Node) bool {
 	seen := make(map[NodeID]bool)
-	for _, e := range slice {
+	for i, e := range slice {
+		if e == nil {
+			panic(fmt.Sprintf("nil *Node at %d", i))
+		}
 		if seen[e.ID] {
 			return true
 		}