aboutsummaryrefslogtreecommitdiff
path: root/dht/routing_table.go
blob: 3c1f2d2b5b73658b14d1214f31228af3ace7dd47 (plain)
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
package dht

import (
	"container/heap"
	"sync"
)

type rItem struct {
	value    *remoteNode
	distance int
	index    int // Index in heap
}

type priorityQueue []*rItem

type routingTable struct {
	id        Infohash
	max       int
	items     priorityQueue
	addresses map[string]*remoteNode
	sync.Mutex
}

func newRoutingTable(id Infohash, max int) (*routingTable, error) {
	k := &routingTable{
		id:  id,
		max: max,
	}
	k.flush()
	heap.Init(&k.items)
	return k, nil
}

// Len implements sort.Interface
func (pq priorityQueue) Len() int { return len(pq) }

// Less implements sort.Interface
func (pq priorityQueue) Less(i, j int) bool {
	return pq[i].distance > pq[j].distance
}

// Swap implements sort.Interface
func (pq priorityQueue) Swap(i, j int) {
	pq[i], pq[j] = pq[j], pq[i]
	pq[i].index = i
	pq[j].index = j
}

// Push implements heap.Interface
func (pq *priorityQueue) Push(x interface{}) {
	n := len(*pq)
	item := x.(*rItem)
	item.index = n
	*pq = append(*pq, item)
}

// Pop implements heap.Interface
func (pq *priorityQueue) Pop() interface{} {
	old := *pq
	n := len(old)
	item := old[n-1]
	item.index = -1 // for safety
	*pq = old[0 : n-1]
	return item
}

func (k *routingTable) add(rn *remoteNode) {
	// Check IP and ports are valid and not self
	if !rn.id.Valid() || rn.id.Equal(k.id) {
		return
	}

	k.Lock()
	defer k.Unlock()

	if _, ok := k.addresses[rn.address.String()]; ok {
		return
	}
	k.addresses[rn.address.String()] = rn

	item := &rItem{
		value:    rn,
		distance: k.id.Distance(rn.id),
	}

	heap.Push(&k.items, item)

	if len(k.items) > k.max {
		for i := k.max - 1; i < len(k.items); i++ {
			old := k.items[i]
			delete(k.addresses, old.value.address.String())
			heap.Remove(&k.items, i)
		}
	}
}

func (k *routingTable) get(n int) (out []*remoteNode) {
	if n == 0 {
		n = len(k.items)
	}
	for i := 0; i < n && i < len(k.items); i++ {
		out = append(out, k.items[i].value)
	}
	return out
}

func (k *routingTable) flush() {
	k.Lock()
	defer k.Unlock()

	k.items = make(priorityQueue, 0)
	k.addresses = make(map[string]*remoteNode, k.max)
}

func (k *routingTable) isEmpty() bool {
	k.Lock()
	defer k.Unlock()
	return len(k.items) == 0
}