aboutsummaryrefslogtreecommitdiffstats
path: root/p2p/discover
diff options
context:
space:
mode:
Diffstat (limited to 'p2p/discover')
-rw-r--r--p2p/discover/table.go49
-rw-r--r--p2p/discover/table_test.go35
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
}