← Back to branch summary

~cosmin.lupu/+junk/penguintv

« back to all changes in this revision

Viewing changes to penguintv/ptvbittorrent/Choker.py

  • Committer: cosmin.lupu at gmail
  • Date: 2010-04-27 16:47:43 UTC
  • Revision ID: cosmin.lupu@gmail.com-20100427164743-ds8xrqonipp5ovdf
initial packaging

Show diffs side-by-side

added added

removed removed

Lines of Context:
penguintv/ptvbittorrent/Choker.py
 
1
# Written by Bram Cohen
 
2
# see LICENSE.txt for license information
 
3
 
 
4
from random import randrange
 
5
 
 
6
class Choker:
 
7
    def __init__(self, max_uploads, schedule, done = lambda: False, min_uploads = None):
 
8
        self.max_uploads = max_uploads
 
9
        if min_uploads is None:
 
10
            min_uploads = max_uploads
 
11
        self.min_uploads = min_uploads
 
12
        self.schedule = schedule
 
13
        self.connections = []
 
14
        self.count = 0
 
15
        self.done = done
 
16
        schedule(self._round_robin, 10)
 
17
    
 
18
    def _round_robin(self):
 
19
        self.schedule(self._round_robin, 10)
 
20
        self.count += 1
 
21
        if self.count % 3 == 0:
 
22
            for i in xrange(len(self.connections)):
 
23
                u = self.connections[i].get_upload()
 
24
                if u.is_choked() and u.is_interested():
 
25
                    self.connections = self.connections[i:] + self.connections[:i]
 
26
                    break
 
27
        self._rechoke()
 
28
 
 
29
    def _snubbed(self, c):
 
30
        if self.done():
 
31
            return False
 
32
        return c.get_download().is_snubbed()
 
33
 
 
34
    def _rate(self, c):
 
35
        if self.done():
 
36
            return c.get_upload().get_rate()
 
37
        else:
 
38
            return c.get_download().get_rate()
 
39
 
 
40
    def _rechoke(self):
 
41
        preferred = []
 
42
        for c in self.connections:
 
43
            if not self._snubbed(c) and c.get_upload().is_interested():
 
44
                preferred.append((-self._rate(c), c))
 
45
        preferred.sort()
 
46
        del preferred[self.max_uploads - 1:]
 
47
        preferred = [x[1] for x in preferred]
 
48
        count = len(preferred)
 
49
        hit = False
 
50
        for c in self.connections:
 
51
            u = c.get_upload()
 
52
            if c in preferred:
 
53
                u.unchoke()
 
54
            else:
 
55
                if count < self.min_uploads or not hit:
 
56
                    u.unchoke()
 
57
                    if u.is_interested():
 
58
                        count += 1
 
59
                        hit = True
 
60
                else:
 
61
                    u.choke()
 
62
 
 
63
    def connection_made(self, connection, p = None):
 
64
        if p is None:
 
65
            p = randrange(-2, len(self.connections) + 1)
 
66
        self.connections.insert(max(p, 0), connection)
 
67
        self._rechoke()
 
68
 
 
69
    def connection_lost(self, connection):
 
70
        self.connections.remove(connection)
 
71
        if connection.get_upload().is_interested() and not connection.get_upload().is_choked():
 
72
            self._rechoke()
 
73
 
 
74
    def interested(self, connection):
 
75
        if not connection.get_upload().is_choked():
 
76
            self._rechoke()
 
77
 
 
78
    def not_interested(self, connection):
 
79
        if not connection.get_upload().is_choked():
 
80
            self._rechoke()
 
81
 
 
82
    def change_max_uploads(self, newval):
 
83
        def foo(self=self, newval=newval):
 
84
            self._change_max_uploads(newval)
 
85
        self.schedule(foo, 0);
 
86
        
 
87
    def _change_max_uploads(self, newval):
 
88
        self.max_uploads = newval
 
89
        self._rechoke()
 
90
 
 
91
class DummyScheduler:
 
92
    def __init__(self):
 
93
        self.s = []
 
94
 
 
95
    def __call__(self, func, delay):
 
96
        self.s.append((func, delay))
 
97
 
 
98
class DummyConnection:
 
99
    def __init__(self, v = 0):
 
100
        self.u = DummyUploader()
 
101
        self.d = DummyDownloader(self)
 
102
        self.v = v
 
103
    
 
104
    def get_upload(self):
 
105
        return self.u
 
106
 
 
107
    def get_download(self):
 
108
        return self.d
 
109
 
 
110
class DummyDownloader:
 
111
    def __init__(self, c):
 
112
        self.s = False
 
113
        self.c = c
 
114
 
 
115
    def is_snubbed(self):
 
116
        return self.s
 
117
 
 
118
    def get_rate(self):
 
119
        return self.c.v
 
120
 
 
121
class DummyUploader:
 
122
    def __init__(self):
 
123
        self.i = False
 
124
        self.c = True
 
125
 
 
126
    def choke(self):
 
127
        if not self.c:
 
128
            self.c = True
 
129
 
 
130
    def unchoke(self):
 
131
        if self.c:
 
132
            self.c = False
 
133
 
 
134
    def is_choked(self):
 
135
        return self.c
 
136
 
 
137
    def is_interested(self):
 
138
        return self.i
 
139
 
 
140
def test_round_robin_with_no_downloads():
 
141
    s = DummyScheduler()
 
142
    Choker(2, s)
 
143
    assert len(s.s) == 1
 
144
    assert s.s[0][1] == 10
 
145
    s.s[0][0]()
 
146
    del s.s[0]
 
147
    assert len(s.s) == 1
 
148
    assert s.s[0][1] == 10
 
149
    s.s[0][0]()
 
150
    del s.s[0]
 
151
    s.s[0][0]()
 
152
    del s.s[0]
 
153
    s.s[0][0]()
 
154
    del s.s[0]
 
155
 
 
156
def test_resort():
 
157
    s = DummyScheduler()
 
158
    choker = Choker(1, s)
 
159
    c1 = DummyConnection()
 
160
    c2 = DummyConnection(1)
 
161
    c3 = DummyConnection(2)
 
162
    c4 = DummyConnection(3)
 
163
    c2.u.i = True
 
164
    c3.u.i = True
 
165
    choker.connection_made(c1)
 
166
    assert not c1.u.c
 
167
    choker.connection_made(c2, 1)
 
168
    assert not c1.u.c
 
169
    assert not c2.u.c
 
170
    choker.connection_made(c3, 1)
 
171
    assert not c1.u.c
 
172
    assert c2.u.c
 
173
    assert not c3.u.c
 
174
    c2.v = 2
 
175
    c3.v = 1
 
176
    choker.connection_made(c4, 1)
 
177
    assert not c1.u.c
 
178
    assert c2.u.c
 
179
    assert not c3.u.c
 
180
    assert not c4.u.c
 
181
    choker.connection_lost(c4)
 
182
    assert not c1.u.c
 
183
    assert c2.u.c
 
184
    assert not c3.u.c
 
185
    s.s[0][0]()
 
186
    assert not c1.u.c
 
187
    assert c2.u.c
 
188
    assert not c3.u.c
 
189
 
 
190
def test_interest():
 
191
    s = DummyScheduler()
 
192
    choker = Choker(1, s)
 
193
    c1 = DummyConnection()
 
194
    c2 = DummyConnection(1)
 
195
    c3 = DummyConnection(2)
 
196
    c2.u.i = True
 
197
    c3.u.i = True
 
198
    choker.connection_made(c1)
 
199
    assert not c1.u.c
 
200
    choker.connection_made(c2, 1)
 
201
    assert not c1.u.c
 
202
    assert not c2.u.c
 
203
    choker.connection_made(c3, 1)
 
204
    assert not c1.u.c
 
205
    assert c2.u.c
 
206
    assert not c3.u.c
 
207
    c3.u.i = False
 
208
    choker.not_interested(c3)
 
209
    assert not c1.u.c
 
210
    assert not c2.u.c
 
211
    assert not c3.u.c
 
212
    c3.u.i = True
 
213
    choker.interested(c3)
 
214
    assert not c1.u.c
 
215
    assert c2.u.c
 
216
    assert not c3.u.c
 
217
    choker.connection_lost(c3)
 
218
    assert not c1.u.c
 
219
    assert not c2.u.c
 
220
 
 
221
def test_robin_interest():
 
222
    s = DummyScheduler()
 
223
    choker = Choker(1, s)
 
224
    c1 = DummyConnection(0)
 
225
    c2 = DummyConnection(1)
 
226
    c1.u.i = True
 
227
    choker.connection_made(c2)
 
228
    assert not c2.u.c
 
229
    choker.connection_made(c1, 0)
 
230
    assert not c1.u.c
 
231
    assert c2.u.c
 
232
    c1.u.i = False
 
233
    choker.not_interested(c1)
 
234
    assert not c1.u.c
 
235
    assert not c2.u.c
 
236
    c1.u.i = True
 
237
    choker.interested(c1)
 
238
    assert not c1.u.c
 
239
    assert c2.u.c
 
240
    choker.connection_lost(c1)
 
241
    assert not c2.u.c
 
242
 
 
243
def test_skip_not_interested():
 
244
    s = DummyScheduler()
 
245
    choker = Choker(1, s)
 
246
    c1 = DummyConnection(0)
 
247
    c2 = DummyConnection(1)
 
248
    c3 = DummyConnection(2)
 
249
    c1.u.i = True
 
250
    c3.u.i = True
 
251
    choker.connection_made(c2)
 
252
    assert not c2.u.c
 
253
    choker.connection_made(c1, 0)
 
254
    assert not c1.u.c
 
255
    assert c2.u.c
 
256
    choker.connection_made(c3, 2)
 
257
    assert not c1.u.c
 
258
    assert c2.u.c
 
259
    assert c3.u.c
 
260
    f = s.s[0][0]
 
261
    f()
 
262
    assert not c1.u.c
 
263
    assert c2.u.c
 
264
    assert c3.u.c
 
265
    f()
 
266
    assert not c1.u.c
 
267
    assert c2.u.c
 
268
    assert c3.u.c
 
269
    f()
 
270
    assert c1.u.c
 
271
    assert c2.u.c
 
272
    assert not c3.u.c
 
273
 
 
274
def test_connection_lost_no_interrupt():
 
275
    s = DummyScheduler()
 
276
    choker = Choker(1, s)
 
277
    c1 = DummyConnection(0)
 
278
    c2 = DummyConnection(1)
 
279
    c3 = DummyConnection(2)
 
280
    c1.u.i = True
 
281
    c2.u.i = True
 
282
    c3.u.i = True
 
283
    choker.connection_made(c1)
 
284
    choker.connection_made(c2, 1)
 
285
    choker.connection_made(c3, 2)
 
286
    f = s.s[0][0]
 
287
    f()
 
288
    assert not c1.u.c
 
289
    assert c2.u.c
 
290
    assert c3.u.c
 
291
    f()
 
292
    assert not c1.u.c
 
293
    assert c2.u.c
 
294
    assert c3.u.c
 
295
    f()
 
296
    assert c1.u.c
 
297
    assert not c2.u.c
 
298
    assert c3.u.c
 
299
    f()
 
300
    assert c1.u.c
 
301
    assert not c2.u.c
 
302
    assert c3.u.c
 
303
    f()
 
304
    assert c1.u.c
 
305
    assert not c2.u.c
 
306
    assert c3.u.c
 
307
    choker.connection_lost(c3)
 
308
    assert c1.u.c
 
309
    assert not c2.u.c
 
310
    f()
 
311
    assert not c1.u.c
 
312
    assert c2.u.c
 
313
    choker.connection_lost(c2)
 
314
    assert not c1.u.c
 
315
 
 
316
def test_connection_made_no_interrupt():
 
317
    s = DummyScheduler()
 
318
    choker = Choker(1, s)
 
319
    c1 = DummyConnection(0)
 
320
    c2 = DummyConnection(1)
 
321
    c3 = DummyConnection(2)
 
322
    c1.u.i = True
 
323
    c2.u.i = True
 
324
    c3.u.i = True
 
325
    choker.connection_made(c1)
 
326
    choker.connection_made(c2, 1)
 
327
    f = s.s[0][0]
 
328
    assert not c1.u.c
 
329
    assert c2.u.c
 
330
    f()
 
331
    assert not c1.u.c
 
332
    assert c2.u.c
 
333
    f()
 
334
    assert not c1.u.c
 
335
    assert c2.u.c
 
336
    choker.connection_made(c3, 1)
 
337
    assert not c1.u.c
 
338
    assert c2.u.c
 
339
    assert c3.u.c
 
340
    f()
 
341
    assert c1.u.c
 
342
    assert c2.u.c
 
343
    assert not c3.u.c
 
344
 
 
345
def test_round_robin():
 
346
    s = DummyScheduler()
 
347
    choker = Choker(1, s)
 
348
    c1 = DummyConnection(0)
 
349
    c2 = DummyConnection(1)
 
350
    c1.u.i = True
 
351
    c2.u.i = True
 
352
    choker.connection_made(c1)
 
353
    choker.connection_made(c2, 1)
 
354
    f = s.s[0][0]
 
355
    assert not c1.u.c
 
356
    assert c2.u.c
 
357
    f()
 
358
    assert not c1.u.c
 
359
    assert c2.u.c
 
360
    f()
 
361
    assert not c1.u.c
 
362
    assert c2.u.c
 
363
    f()
 
364
    assert c1.u.c
 
365
    assert not c2.u.c
 
366
    f()
 
367
    assert c1.u.c
 
368
    assert not c2.u.c
 
369
    f()
 
370
    assert c1.u.c
 
371
    assert not c2.u.c
 
372
    f()
 
373
    assert not c1.u.c
 
374
    assert c2.u.c
 
375
    
 
376
def test_multi():
 
377
    s = DummyScheduler()
 
378
    choker = Choker(4, s)
 
379
    c1 = DummyConnection(0)
 
380
    c2 = DummyConnection(0)
 
381
    c3 = DummyConnection(0)
 
382
    c4 = DummyConnection(8)
 
383
    c5 = DummyConnection(0)
 
384
    c6 = DummyConnection(0)
 
385
    c7 = DummyConnection(6)
 
386
    c8 = DummyConnection(0)
 
387
    c9 = DummyConnection(9)
 
388
    c10 = DummyConnection(7)
 
389
    c11 = DummyConnection(10)
 
390
    choker.connection_made(c1, 0)
 
391
    choker.connection_made(c2, 1)
 
392
    choker.connection_made(c3, 2)
 
393
    choker.connection_made(c4, 3)
 
394
    choker.connection_made(c5, 4)
 
395
    choker.connection_made(c6, 5)
 
396
    choker.connection_made(c7, 6)
 
397
    choker.connection_made(c8, 7)
 
398
    choker.connection_made(c9, 8)
 
399
    choker.connection_made(c10, 9)
 
400
    choker.connection_made(c11, 10)
 
401
    c2.u.i = True
 
402
    c4.u.i = True
 
403
    c6.u.i = True
 
404
    c8.u.i = True
 
405
    c10.u.i = True
 
406
    c2.d.s = True
 
407
    c6.d.s = True
 
408
    c8.d.s = True
 
409
    s.s[0][0]()
 
410
    assert not c1.u.c
 
411
    assert not c2.u.c
 
412
    assert not c3.u.c
 
413
    assert not c4.u.c
 
414
    assert not c5.u.c
 
415
    assert not c6.u.c
 
416
    assert c7.u.c
 
417
    assert c8.u.c
 
418
    assert c9.u.c
 
419
    assert not c10.u.c
 
420
    assert c11.u.c
 
421
 
 
422