-
-
Notifications
You must be signed in to change notification settings - Fork 66
/
table.go
134 lines (120 loc) · 2.87 KB
/
table.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
package dht
import (
"errors"
"fmt"
"github.com/anacrolix/dht/v2/int160"
)
// Node table, with indexes on distance from root ID to bucket, and node addr.
type table struct {
rootID int160.T
k int
buckets [160]bucket
addrs map[string]map[int160.T]struct{}
}
func (tbl *table) K() int {
return tbl.k
}
func (tbl *table) randomIdForBucket(bucketIndex int) int160.T {
randomId := randomIdInBucket(tbl.rootID, bucketIndex)
if randomIdBucketIndex := tbl.bucketIndex(randomId); randomIdBucketIndex != bucketIndex {
panic(fmt.Sprintf("bucket index for random id %v == %v not %v", randomId, randomIdBucketIndex, bucketIndex))
}
return randomId
}
func (tbl *table) addrNodes(addr Addr) []*node {
a := tbl.addrs[addr.String()]
ret := make([]*node, 0, len(a))
for id := range a {
ret = append(ret, tbl.getNode(addr, id))
}
return ret
}
func (tbl *table) dropNode(n *node) {
as := n.Addr.String()
if _, ok := tbl.addrs[as][n.Id]; !ok {
panic("missing id for addr")
}
delete(tbl.addrs[as], n.Id)
if len(tbl.addrs[as]) == 0 {
delete(tbl.addrs, as)
}
b := tbl.bucketForID(n.Id)
if _, ok := b.nodes[n]; !ok {
panic("expected node in bucket")
}
delete(b.nodes, n)
}
func (tbl *table) bucketForID(id int160.T) *bucket {
return &tbl.buckets[tbl.bucketIndex(id)]
}
func (tbl *table) numNodes() (num int) {
for i := range tbl.buckets {
num += tbl.buckets[i].Len()
}
return
}
func (tbl *table) bucketIndex(id int160.T) int {
if id == tbl.rootID {
panic("nobody puts the root ID in a bucket")
}
var a int160.T
a.Xor(&tbl.rootID, &id)
index := 160 - a.BitLen()
return index
}
func (tbl *table) forNodes(f func(*node) bool) bool {
for i := range tbl.buckets {
if !tbl.buckets[i].EachNode(f) {
return false
}
}
return true
}
func (tbl *table) getNode(addr Addr, id int160.T) *node {
if id == tbl.rootID {
return nil
}
return tbl.buckets[tbl.bucketIndex(id)].GetNode(addr, id)
}
func (tbl *table) closestNodes(k int, target int160.T, filter func(*node) bool) (ret []*node) {
for bi := func() int {
if target == tbl.rootID {
return len(tbl.buckets) - 1
} else {
return tbl.bucketIndex(target)
}
}(); bi >= 0 && len(ret) < k; bi-- {
for n := range tbl.buckets[bi].nodes {
if filter(n) {
ret = append(ret, n)
}
}
}
// TODO: Keep only the closest.
if len(ret) > k {
ret = ret[:k]
}
return
}
func (tbl *table) addNode(n *node) error {
if n.Id == tbl.rootID {
return errors.New("is root id")
}
b := &tbl.buckets[tbl.bucketIndex(n.Id)]
if b.GetNode(n.Addr, n.Id) != nil {
return errors.New("already present")
}
if b.Len() >= tbl.k {
return errors.New("bucket is full")
}
b.AddNode(n, tbl.k)
if tbl.addrs == nil {
tbl.addrs = make(map[string]map[int160.T]struct{}, 160*tbl.k)
}
as := n.Addr.String()
if tbl.addrs[as] == nil {
tbl.addrs[as] = make(map[int160.T]struct{}, 1)
}
tbl.addrs[as][n.Id] = struct{}{}
return nil
}