-
Notifications
You must be signed in to change notification settings - Fork 70
/
unchoker.go
155 lines (137 loc) · 3.95 KB
/
unchoker.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
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
package unchoker
import (
"math/rand"
"sort"
)
// Unchoker implements an algorithm to select peers to unchoke based on their download speed.
type Unchoker struct {
numUnchoked int
numOptimisticUnchoked int
// Every 3rd round an optimistic unchoke logic is applied.
round uint8
peersUnchoked map[Peer]struct{}
peersUnchokedOptimistic map[Peer]struct{}
}
// Peer of a torrent.
type Peer interface {
// Sends messages and set choking status of local peeer
Choke()
Unchoke()
// Choking returns choke status of local peer
Choking() bool
// Interested returns interest status of remote peer
Interested() bool
// SetOptimistic sets the uptimistic unchoke status of peer
SetOptimistic(value bool)
// OptimisticUnchoked returns the value previously set by SetOptimistic
Optimistic() bool
DownloadSpeed() int
UploadSpeed() int
}
// New returns a new Unchoker.
func New(numUnchoked, numOptimisticUnchoked int) *Unchoker {
return &Unchoker{
numUnchoked: numUnchoked,
numOptimisticUnchoked: numOptimisticUnchoked,
peersUnchoked: make(map[Peer]struct{}, numUnchoked),
peersUnchokedOptimistic: make(map[Peer]struct{}, numUnchoked),
}
}
// HandleDisconnect must be called to remove the peer from internal indexes.
func (u *Unchoker) HandleDisconnect(pe Peer) {
delete(u.peersUnchoked, pe)
delete(u.peersUnchokedOptimistic, pe)
}
func (u *Unchoker) candidatesUnchoke(allPeers []Peer) []Peer {
peers := allPeers[:0]
for _, pe := range allPeers {
if pe.Interested() {
peers = append(peers, pe)
}
}
return peers
}
func (u *Unchoker) sortPeers(peers []Peer, completed bool) {
byUploadSpeed := func(i, j int) bool { return peers[i].UploadSpeed() > peers[j].UploadSpeed() }
byDownloadSpeed := func(i, j int) bool { return peers[i].DownloadSpeed() > peers[j].DownloadSpeed() }
if completed {
sort.Slice(peers, byUploadSpeed)
} else {
sort.Slice(peers, byDownloadSpeed)
}
}
// TickUnchoke must be called at every 10 seconds.
func (u *Unchoker) TickUnchoke(allPeers []Peer, torrentCompleted bool) {
optimistic := u.round == 0
peers := u.candidatesUnchoke(allPeers)
u.sortPeers(peers, torrentCompleted)
var i, unchoked int
for ; i < len(peers) && unchoked < u.numUnchoked; i++ {
if !optimistic && peers[i].Optimistic() {
continue
}
u.unchokePeer(peers[i])
unchoked++
}
peers = peers[i:]
if optimistic {
for i = 0; i < u.numOptimisticUnchoked && len(peers) > 0; i++ {
n := rand.Intn(len(peers))
pe := peers[n]
u.optimisticUnchokePeer(pe)
peers[n], peers = peers[len(peers)-1], peers[:len(peers)-1]
}
}
for _, pe := range peers {
u.chokePeer(pe)
}
u.round = (u.round + 1) % 3
}
func (u *Unchoker) chokePeer(pe Peer) {
if pe.Choking() {
return
}
pe.Choke()
pe.SetOptimistic(false)
delete(u.peersUnchoked, pe)
delete(u.peersUnchokedOptimistic, pe)
}
func (u *Unchoker) unchokePeer(pe Peer) {
if !pe.Choking() {
if pe.Optimistic() {
// Move into regular unchoked peers
pe.SetOptimistic(false)
delete(u.peersUnchokedOptimistic, pe)
u.peersUnchoked[pe] = struct{}{}
}
return
}
pe.Unchoke()
u.peersUnchoked[pe] = struct{}{}
pe.SetOptimistic(false)
}
func (u *Unchoker) optimisticUnchokePeer(pe Peer) {
if !pe.Choking() {
if !pe.Optimistic() {
// Move into optimistic unchoked peers
pe.SetOptimistic(true)
delete(u.peersUnchoked, pe)
u.peersUnchokedOptimistic[pe] = struct{}{}
}
return
}
pe.Unchoke()
u.peersUnchokedOptimistic[pe] = struct{}{}
pe.SetOptimistic(true)
}
// FastUnchoke must be called when remote peer is interested.
// Remote peer is unchoked immediately if there are not enough unchoked peers.
// Without this function, remote peer would have to wait for next unchoke period.
func (u *Unchoker) FastUnchoke(pe Peer) {
if pe.Choking() && pe.Interested() && len(u.peersUnchoked) < u.numUnchoked {
u.unchokePeer(pe)
}
if pe.Choking() && pe.Interested() && len(u.peersUnchokedOptimistic) < u.numOptimisticUnchoked {
u.optimisticUnchokePeer(pe)
}
}