| 36 |
kaklik |
1 |
# Written by Bram Cohen
|
|
|
2 |
# see LICENSE.txt for license information
|
|
|
3 |
|
|
|
4 |
from random import randrange, shuffle
|
|
|
5 |
from BitTornado.clock import clock
|
|
|
6 |
try:
|
|
|
7 |
True
|
|
|
8 |
except:
|
|
|
9 |
True = 1
|
|
|
10 |
False = 0
|
|
|
11 |
|
|
|
12 |
class Choker:
|
|
|
13 |
def __init__(self, config, schedule, picker, done = lambda: False):
|
|
|
14 |
self.config = config
|
|
|
15 |
self.round_robin_period = config['round_robin_period']
|
|
|
16 |
self.schedule = schedule
|
|
|
17 |
self.picker = picker
|
|
|
18 |
self.connections = []
|
|
|
19 |
self.last_preferred = 0
|
|
|
20 |
self.last_round_robin = clock()
|
|
|
21 |
self.done = done
|
|
|
22 |
self.super_seed = False
|
|
|
23 |
self.paused = False
|
|
|
24 |
schedule(self._round_robin, 5)
|
|
|
25 |
|
|
|
26 |
def set_round_robin_period(self, x):
|
|
|
27 |
self.round_robin_period = x
|
|
|
28 |
|
|
|
29 |
def _round_robin(self):
|
|
|
30 |
self.schedule(self._round_robin, 5)
|
|
|
31 |
if self.super_seed:
|
|
|
32 |
cons = range(len(self.connections))
|
|
|
33 |
to_close = []
|
|
|
34 |
count = self.config['min_uploads']-self.last_preferred
|
|
|
35 |
if count > 0: # optimization
|
|
|
36 |
shuffle(cons)
|
|
|
37 |
for c in cons:
|
|
|
38 |
i = self.picker.next_have(self.connections[c], count > 0)
|
|
|
39 |
if i is None:
|
|
|
40 |
continue
|
|
|
41 |
if i < 0:
|
|
|
42 |
to_close.append(self.connections[c])
|
|
|
43 |
continue
|
|
|
44 |
self.connections[c].send_have(i)
|
|
|
45 |
count -= 1
|
|
|
46 |
for c in to_close:
|
|
|
47 |
c.close()
|
|
|
48 |
if self.last_round_robin + self.round_robin_period < clock():
|
|
|
49 |
self.last_round_robin = clock()
|
|
|
50 |
for i in xrange(1, len(self.connections)):
|
|
|
51 |
c = self.connections[i]
|
|
|
52 |
u = c.get_upload()
|
|
|
53 |
if u.is_choked() and u.is_interested():
|
|
|
54 |
self.connections = self.connections[i:] + self.connections[:i]
|
|
|
55 |
break
|
|
|
56 |
self._rechoke()
|
|
|
57 |
|
|
|
58 |
def _rechoke(self):
|
|
|
59 |
preferred = []
|
|
|
60 |
maxuploads = self.config['max_uploads']
|
|
|
61 |
if self.paused:
|
|
|
62 |
for c in self.connections:
|
|
|
63 |
c.get_upload().choke()
|
|
|
64 |
return
|
|
|
65 |
if maxuploads > 1:
|
|
|
66 |
for c in self.connections:
|
|
|
67 |
u = c.get_upload()
|
|
|
68 |
if not u.is_interested():
|
|
|
69 |
continue
|
|
|
70 |
if self.done():
|
|
|
71 |
r = u.get_rate()
|
|
|
72 |
else:
|
|
|
73 |
d = c.get_download()
|
|
|
74 |
r = d.get_rate()
|
|
|
75 |
if r < 1000 or d.is_snubbed():
|
|
|
76 |
continue
|
|
|
77 |
preferred.append((-r, c))
|
|
|
78 |
self.last_preferred = len(preferred)
|
|
|
79 |
preferred.sort()
|
|
|
80 |
del preferred[maxuploads-1:]
|
|
|
81 |
preferred = [x[1] for x in preferred]
|
|
|
82 |
count = len(preferred)
|
|
|
83 |
hit = False
|
|
|
84 |
to_unchoke = []
|
|
|
85 |
for c in self.connections:
|
|
|
86 |
u = c.get_upload()
|
|
|
87 |
if c in preferred:
|
|
|
88 |
to_unchoke.append(u)
|
|
|
89 |
else:
|
|
|
90 |
if count < maxuploads or not hit:
|
|
|
91 |
to_unchoke.append(u)
|
|
|
92 |
if u.is_interested():
|
|
|
93 |
count += 1
|
|
|
94 |
hit = True
|
|
|
95 |
else:
|
|
|
96 |
u.choke()
|
|
|
97 |
for u in to_unchoke:
|
|
|
98 |
u.unchoke()
|
|
|
99 |
|
|
|
100 |
def connection_made(self, connection, p = None):
|
|
|
101 |
if p is None:
|
|
|
102 |
p = randrange(-2, len(self.connections) + 1)
|
|
|
103 |
self.connections.insert(max(p, 0), connection)
|
|
|
104 |
self._rechoke()
|
|
|
105 |
|
|
|
106 |
def connection_lost(self, connection):
|
|
|
107 |
self.connections.remove(connection)
|
|
|
108 |
self.picker.lost_peer(connection)
|
|
|
109 |
if connection.get_upload().is_interested() and not connection.get_upload().is_choked():
|
|
|
110 |
self._rechoke()
|
|
|
111 |
|
|
|
112 |
def interested(self, connection):
|
|
|
113 |
if not connection.get_upload().is_choked():
|
|
|
114 |
self._rechoke()
|
|
|
115 |
|
|
|
116 |
def not_interested(self, connection):
|
|
|
117 |
if not connection.get_upload().is_choked():
|
|
|
118 |
self._rechoke()
|
|
|
119 |
|
|
|
120 |
def set_super_seed(self):
|
|
|
121 |
while self.connections: # close all connections
|
|
|
122 |
self.connections[0].close()
|
|
|
123 |
self.picker.set_superseed()
|
|
|
124 |
self.super_seed = True
|
|
|
125 |
|
|
|
126 |
def pause(self, flag):
|
|
|
127 |
self.paused = flag
|
|
|
128 |
self._rechoke()
|