| 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 PiecePicker:
|
|
|
13 |
def __init__(self, numpieces,
|
|
|
14 |
rarest_first_cutoff = 1, rarest_first_priority_cutoff = 3,
|
|
|
15 |
priority_step = 20):
|
|
|
16 |
self.rarest_first_cutoff = rarest_first_cutoff
|
|
|
17 |
self.rarest_first_priority_cutoff = rarest_first_priority_cutoff + priority_step
|
|
|
18 |
self.priority_step = priority_step
|
|
|
19 |
self.cutoff = rarest_first_priority_cutoff
|
|
|
20 |
self.numpieces = numpieces
|
|
|
21 |
self.started = []
|
|
|
22 |
self.totalcount = 0
|
|
|
23 |
self.numhaves = [0] * numpieces
|
|
|
24 |
self.priority = [1] * numpieces
|
|
|
25 |
self.removed_partials = {}
|
|
|
26 |
self.crosscount = [numpieces]
|
|
|
27 |
self.crosscount2 = [numpieces]
|
|
|
28 |
self.has = [0] * numpieces
|
|
|
29 |
self.numgot = 0
|
|
|
30 |
self.done = False
|
|
|
31 |
self.seed_connections = {}
|
|
|
32 |
self.past_ips = {}
|
|
|
33 |
self.seed_time = None
|
|
|
34 |
self.superseed = False
|
|
|
35 |
self.seeds_connected = 0
|
|
|
36 |
self._init_interests()
|
|
|
37 |
|
|
|
38 |
def _init_interests(self):
|
|
|
39 |
self.interests = [[] for x in xrange(self.priority_step)]
|
|
|
40 |
self.level_in_interests = [self.priority_step] * self.numpieces
|
|
|
41 |
interests = range(self.numpieces)
|
|
|
42 |
shuffle(interests)
|
|
|
43 |
self.pos_in_interests = [0] * self.numpieces
|
|
|
44 |
for i in xrange(self.numpieces):
|
|
|
45 |
self.pos_in_interests[interests[i]] = i
|
|
|
46 |
self.interests.append(interests)
|
|
|
47 |
|
|
|
48 |
|
|
|
49 |
def got_have(self, piece):
|
|
|
50 |
self.totalcount+=1
|
|
|
51 |
numint = self.numhaves[piece]
|
|
|
52 |
self.numhaves[piece] += 1
|
|
|
53 |
self.crosscount[numint] -= 1
|
|
|
54 |
if numint+1==len(self.crosscount):
|
|
|
55 |
self.crosscount.append(0)
|
|
|
56 |
self.crosscount[numint+1] += 1
|
|
|
57 |
if not self.done:
|
|
|
58 |
numintplus = numint+self.has[piece]
|
|
|
59 |
self.crosscount2[numintplus] -= 1
|
|
|
60 |
if numintplus+1 == len(self.crosscount2):
|
|
|
61 |
self.crosscount2.append(0)
|
|
|
62 |
self.crosscount2[numintplus+1] += 1
|
|
|
63 |
numint = self.level_in_interests[piece]
|
|
|
64 |
self.level_in_interests[piece] += 1
|
|
|
65 |
if self.superseed:
|
|
|
66 |
self.seed_got_haves[piece] += 1
|
|
|
67 |
numint = self.level_in_interests[piece]
|
|
|
68 |
self.level_in_interests[piece] += 1
|
|
|
69 |
elif self.has[piece] or self.priority[piece] == -1:
|
|
|
70 |
return
|
|
|
71 |
if numint == len(self.interests) - 1:
|
|
|
72 |
self.interests.append([])
|
|
|
73 |
self._shift_over(piece, self.interests[numint], self.interests[numint + 1])
|
|
|
74 |
|
|
|
75 |
def lost_have(self, piece):
|
|
|
76 |
self.totalcount-=1
|
|
|
77 |
numint = self.numhaves[piece]
|
|
|
78 |
self.numhaves[piece] -= 1
|
|
|
79 |
self.crosscount[numint] -= 1
|
|
|
80 |
self.crosscount[numint-1] += 1
|
|
|
81 |
if not self.done:
|
|
|
82 |
numintplus = numint+self.has[piece]
|
|
|
83 |
self.crosscount2[numintplus] -= 1
|
|
|
84 |
self.crosscount2[numintplus-1] += 1
|
|
|
85 |
numint = self.level_in_interests[piece]
|
|
|
86 |
self.level_in_interests[piece] -= 1
|
|
|
87 |
if self.superseed:
|
|
|
88 |
numint = self.level_in_interests[piece]
|
|
|
89 |
self.level_in_interests[piece] -= 1
|
|
|
90 |
elif self.has[piece] or self.priority[piece] == -1:
|
|
|
91 |
return
|
|
|
92 |
self._shift_over(piece, self.interests[numint], self.interests[numint - 1])
|
|
|
93 |
|
|
|
94 |
def _shift_over(self, piece, l1, l2):
|
|
|
95 |
assert self.superseed or (not self.has[piece] and self.priority[piece] >= 0)
|
|
|
96 |
parray = self.pos_in_interests
|
|
|
97 |
p = parray[piece]
|
|
|
98 |
assert l1[p] == piece
|
|
|
99 |
q = l1[-1]
|
|
|
100 |
l1[p] = q
|
|
|
101 |
parray[q] = p
|
|
|
102 |
del l1[-1]
|
|
|
103 |
newp = randrange(len(l2)+1)
|
|
|
104 |
if newp == len(l2):
|
|
|
105 |
parray[piece] = len(l2)
|
|
|
106 |
l2.append(piece)
|
|
|
107 |
else:
|
|
|
108 |
old = l2[newp]
|
|
|
109 |
parray[old] = len(l2)
|
|
|
110 |
l2.append(old)
|
|
|
111 |
l2[newp] = piece
|
|
|
112 |
parray[piece] = newp
|
|
|
113 |
|
|
|
114 |
|
|
|
115 |
def got_seed(self):
|
|
|
116 |
self.seeds_connected += 1
|
|
|
117 |
self.cutoff = max(self.rarest_first_priority_cutoff-self.seeds_connected,0)
|
|
|
118 |
|
|
|
119 |
def became_seed(self):
|
|
|
120 |
self.got_seed()
|
|
|
121 |
self.totalcount -= self.numpieces
|
|
|
122 |
self.numhaves = [i-1 for i in self.numhaves]
|
|
|
123 |
if self.superseed or not self.done:
|
|
|
124 |
self.level_in_interests = [i-1 for i in self.level_in_interests]
|
|
|
125 |
del self.interests[0]
|
|
|
126 |
del self.crosscount[0]
|
|
|
127 |
if not self.done:
|
|
|
128 |
del self.crosscount2[0]
|
|
|
129 |
|
|
|
130 |
def lost_seed(self):
|
|
|
131 |
self.seeds_connected -= 1
|
|
|
132 |
self.cutoff = max(self.rarest_first_priority_cutoff-self.seeds_connected,0)
|
|
|
133 |
|
|
|
134 |
|
|
|
135 |
def requested(self, piece):
|
|
|
136 |
if piece not in self.started:
|
|
|
137 |
self.started.append(piece)
|
|
|
138 |
|
|
|
139 |
def _remove_from_interests(self, piece, keep_partial = False):
|
|
|
140 |
l = self.interests[self.level_in_interests[piece]]
|
|
|
141 |
p = self.pos_in_interests[piece]
|
|
|
142 |
assert l[p] == piece
|
|
|
143 |
q = l[-1]
|
|
|
144 |
l[p] = q
|
|
|
145 |
self.pos_in_interests[q] = p
|
|
|
146 |
del l[-1]
|
|
|
147 |
try:
|
|
|
148 |
self.started.remove(piece)
|
|
|
149 |
if keep_partial:
|
|
|
150 |
self.removed_partials[piece] = 1
|
|
|
151 |
except ValueError:
|
|
|
152 |
pass
|
|
|
153 |
|
|
|
154 |
def complete(self, piece):
|
|
|
155 |
assert not self.has[piece]
|
|
|
156 |
self.has[piece] = 1
|
|
|
157 |
self.numgot += 1
|
|
|
158 |
if self.numgot == self.numpieces:
|
|
|
159 |
self.done = True
|
|
|
160 |
self.crosscount2 = self.crosscount
|
|
|
161 |
else:
|
|
|
162 |
numhaves = self.numhaves[piece]
|
|
|
163 |
self.crosscount2[numhaves] -= 1
|
|
|
164 |
if numhaves+1 == len(self.crosscount2):
|
|
|
165 |
self.crosscount2.append(0)
|
|
|
166 |
self.crosscount2[numhaves+1] += 1
|
|
|
167 |
self._remove_from_interests(piece)
|
|
|
168 |
|
|
|
169 |
|
|
|
170 |
def next(self, haves, wantfunc, complete_first = False):
|
|
|
171 |
cutoff = self.numgot < self.rarest_first_cutoff
|
|
|
172 |
complete_first = (complete_first or cutoff) and not haves.complete()
|
|
|
173 |
best = None
|
|
|
174 |
bestnum = 2 ** 30
|
|
|
175 |
for i in self.started:
|
|
|
176 |
if haves[i] and wantfunc(i):
|
|
|
177 |
if self.level_in_interests[i] < bestnum:
|
|
|
178 |
best = i
|
|
|
179 |
bestnum = self.level_in_interests[i]
|
|
|
180 |
if best is not None:
|
|
|
181 |
if complete_first or (cutoff and len(self.interests) > self.cutoff):
|
|
|
182 |
return best
|
|
|
183 |
if haves.complete():
|
|
|
184 |
r = [ (0, min(bestnum,len(self.interests))) ]
|
|
|
185 |
elif cutoff and len(self.interests) > self.cutoff:
|
|
|
186 |
r = [ (self.cutoff, min(bestnum,len(self.interests))),
|
|
|
187 |
(0, self.cutoff) ]
|
|
|
188 |
else:
|
|
|
189 |
r = [ (0, min(bestnum,len(self.interests))) ]
|
|
|
190 |
for lo,hi in r:
|
|
|
191 |
for i in xrange(lo,hi):
|
|
|
192 |
for j in self.interests[i]:
|
|
|
193 |
if haves[j] and wantfunc(j):
|
|
|
194 |
return j
|
|
|
195 |
if best is not None:
|
|
|
196 |
return best
|
|
|
197 |
return None
|
|
|
198 |
|
|
|
199 |
|
|
|
200 |
def am_I_complete(self):
|
|
|
201 |
return self.done
|
|
|
202 |
|
|
|
203 |
def bump(self, piece):
|
|
|
204 |
l = self.interests[self.level_in_interests[piece]]
|
|
|
205 |
pos = self.pos_in_interests[piece]
|
|
|
206 |
del l[pos]
|
|
|
207 |
l.append(piece)
|
|
|
208 |
for i in range(pos,len(l)):
|
|
|
209 |
self.pos_in_interests[l[i]] = i
|
|
|
210 |
try:
|
|
|
211 |
self.started.remove(piece)
|
|
|
212 |
except:
|
|
|
213 |
pass
|
|
|
214 |
|
|
|
215 |
def set_priority(self, piece, p):
|
|
|
216 |
if self.superseed:
|
|
|
217 |
return False # don't muck with this if you're a superseed
|
|
|
218 |
oldp = self.priority[piece]
|
|
|
219 |
if oldp == p:
|
|
|
220 |
return False
|
|
|
221 |
self.priority[piece] = p
|
|
|
222 |
if p == -1:
|
|
|
223 |
# when setting priority -1,
|
|
|
224 |
# make sure to cancel any downloads for this piece
|
|
|
225 |
if not self.has[piece]:
|
|
|
226 |
self._remove_from_interests(piece, True)
|
|
|
227 |
return True
|
|
|
228 |
if oldp == -1:
|
|
|
229 |
level = self.numhaves[piece] + (self.priority_step * p)
|
|
|
230 |
self.level_in_interests[piece] = level
|
|
|
231 |
if self.has[piece]:
|
|
|
232 |
return True
|
|
|
233 |
while len(self.interests) < level+1:
|
|
|
234 |
self.interests.append([])
|
|
|
235 |
l2 = self.interests[level]
|
|
|
236 |
parray = self.pos_in_interests
|
|
|
237 |
newp = randrange(len(l2)+1)
|
|
|
238 |
if newp == len(l2):
|
|
|
239 |
parray[piece] = len(l2)
|
|
|
240 |
l2.append(piece)
|
|
|
241 |
else:
|
|
|
242 |
old = l2[newp]
|
|
|
243 |
parray[old] = len(l2)
|
|
|
244 |
l2.append(old)
|
|
|
245 |
l2[newp] = piece
|
|
|
246 |
parray[piece] = newp
|
|
|
247 |
if self.removed_partials.has_key(piece):
|
|
|
248 |
del self.removed_partials[piece]
|
|
|
249 |
self.started.append(piece)
|
|
|
250 |
# now go to downloader and try requesting more
|
|
|
251 |
return True
|
|
|
252 |
numint = self.level_in_interests[piece]
|
|
|
253 |
newint = numint + ((p - oldp) * self.priority_step)
|
|
|
254 |
self.level_in_interests[piece] = newint
|
|
|
255 |
if self.has[piece]:
|
|
|
256 |
return False
|
|
|
257 |
while len(self.interests) < newint+1:
|
|
|
258 |
self.interests.append([])
|
|
|
259 |
self._shift_over(piece, self.interests[numint], self.interests[newint])
|
|
|
260 |
return False
|
|
|
261 |
|
|
|
262 |
def is_blocked(self, piece):
|
|
|
263 |
return self.priority[piece] < 0
|
|
|
264 |
|
|
|
265 |
|
|
|
266 |
def set_superseed(self):
|
|
|
267 |
assert self.done
|
|
|
268 |
self.superseed = True
|
|
|
269 |
self.seed_got_haves = [0] * self.numpieces
|
|
|
270 |
self._init_interests() # assume everyone is disconnected
|
|
|
271 |
|
|
|
272 |
def next_have(self, connection, looser_upload):
|
|
|
273 |
if self.seed_time is None:
|
|
|
274 |
self.seed_time = clock()
|
|
|
275 |
return None
|
|
|
276 |
if clock() < self.seed_time+10: # wait 10 seconds after seeing the first peers
|
|
|
277 |
return None # to give time to grab have lists
|
|
|
278 |
if not connection.upload.super_seeding:
|
|
|
279 |
return None
|
|
|
280 |
olddl = self.seed_connections.get(connection)
|
|
|
281 |
if olddl is not None:
|
|
|
282 |
ip = connection.get_ip()
|
|
|
283 |
olddl = self.past_ips.get(ip)
|
|
|
284 |
if olddl is not None: # peer reconnected
|
|
|
285 |
self.seed_connections[connection] = olddl
|
|
|
286 |
if olddl is not None:
|
|
|
287 |
if looser_upload:
|
|
|
288 |
num = 1 # send a new have even if it hasn't spread that piece elsewhere
|
|
|
289 |
else:
|
|
|
290 |
num = 2
|
|
|
291 |
if self.seed_got_haves[olddl] < num:
|
|
|
292 |
return None
|
|
|
293 |
if not connection.upload.was_ever_interested: # it never downloaded it?
|
|
|
294 |
connection.upload.skipped_count += 1
|
|
|
295 |
if connection.upload.skipped_count >= 3: # probably another stealthed seed
|
|
|
296 |
return -1 # signal to close it
|
|
|
297 |
for tier in self.interests:
|
|
|
298 |
for piece in tier:
|
|
|
299 |
if not connection.download.have[piece]:
|
|
|
300 |
seedint = self.level_in_interests[piece]
|
|
|
301 |
self.level_in_interests[piece] += 1 # tweak it up one, so you don't duplicate effort
|
|
|
302 |
if seedint == len(self.interests) - 1:
|
|
|
303 |
self.interests.append([])
|
|
|
304 |
self._shift_over(piece,
|
|
|
305 |
self.interests[seedint], self.interests[seedint + 1])
|
|
|
306 |
self.seed_got_haves[piece] = 0 # reset this
|
|
|
307 |
self.seed_connections[connection] = piece
|
|
|
308 |
connection.upload.seed_have_list.append(piece)
|
|
|
309 |
return piece
|
|
|
310 |
return -1 # something screwy; terminate connection
|
|
|
311 |
|
|
|
312 |
def lost_peer(self, connection):
|
|
|
313 |
olddl = self.seed_connections.get(connection)
|
|
|
314 |
if olddl is None:
|
|
|
315 |
return
|
|
|
316 |
del self.seed_connections[connection]
|
|
|
317 |
self.past_ips[connection.get_ip()] = olddl
|
|
|
318 |
if self.seed_got_haves[olddl] == 1:
|
|
|
319 |
self.seed_got_haves[olddl] = 0
|