← Back to branch summary

~cosmin.lupu/+junk/penguintv

« back to all changes in this revision

Viewing changes to penguintv/ptvbittorrent/Uploader.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/Uploader.py
 
1
# Written by Bram Cohen
 
2
# see LICENSE.txt for license information
 
3
 
 
4
from CurrentRateMeasure import Measure
 
5
 
 
6
class Upload:
 
7
    def __init__(self, connection, choker, storage, 
 
8
            max_slice_length, max_rate_period, fudge):
 
9
        self.connection = connection
 
10
        self.choker = choker
 
11
        self.storage = storage
 
12
        self.max_slice_length = max_slice_length
 
13
        self.max_rate_period = max_rate_period
 
14
        self.choked = True
 
15
        self.interested = False
 
16
        self.buffer = []
 
17
        self.measure = Measure(max_rate_period, fudge)
 
18
        if storage.do_I_have_anything():
 
19
            connection.send_bitfield(storage.get_have_list())
 
20
 
 
21
    def got_not_interested(self):
 
22
        if self.interested:
 
23
            self.interested = False
 
24
            del self.buffer[:]
 
25
            self.choker.not_interested(self.connection)
 
26
 
 
27
    def got_interested(self):
 
28
        if not self.interested:
 
29
            self.interested = True
 
30
            self.choker.interested(self.connection)
 
31
 
 
32
    def flushed(self):
 
33
        while len(self.buffer) > 0 and self.connection.is_flushed():
 
34
            index, begin, length = self.buffer[0]
 
35
            del self.buffer[0]
 
36
            piece = self.storage.get_piece(index, begin, length)
 
37
            if piece is None:
 
38
                self.connection.close()
 
39
                return
 
40
            self.measure.update_rate(len(piece))
 
41
            self.connection.send_piece(index, begin, piece)
 
42
 
 
43
    def got_request(self, index, begin, length):
 
44
        if not self.interested or length > self.max_slice_length:
 
45
            self.connection.close()
 
46
            return
 
47
        if not self.choked:
 
48
            self.buffer.append((index, begin, length))
 
49
            self.flushed()
 
50
 
 
51
    def got_cancel(self, index, begin, length):
 
52
        try:
 
53
            self.buffer.remove((index, begin, length))
 
54
        except ValueError:
 
55
            pass
 
56
 
 
57
    def choke(self):
 
58
        if not self.choked:
 
59
            self.choked = True
 
60
            del self.buffer[:]
 
61
            self.connection.send_choke()
 
62
 
 
63
    def unchoke(self):
 
64
        if self.choked:
 
65
            self.choked = False
 
66
            self.connection.send_unchoke()
 
67
        
 
68
    def is_choked(self):
 
69
        return self.choked
 
70
        
 
71
    def is_interested(self):
 
72
        return self.interested
 
73
 
 
74
    def has_queries(self):
 
75
        return len(self.buffer) > 0
 
76
 
 
77
    def get_rate(self):
 
78
        return self.measure.get_rate()
 
79
 
 
80
class DummyConnection:
 
81
    def __init__(self, events):
 
82
        self.events = events
 
83
        self.flushed = False
 
84
 
 
85
    def send_bitfield(self, bitfield):
 
86
        self.events.append(('bitfield', bitfield))
 
87
    
 
88
    def is_flushed(self):
 
89
        return self.flushed
 
90
 
 
91
    def close(self):
 
92
        self.events.append('closed')
 
93
 
 
94
    def send_piece(self, index, begin, piece):
 
95
        self.events.append(('piece', index, begin, piece))
 
96
 
 
97
    def send_choke(self):
 
98
        self.events.append('choke')
 
99
 
 
100
    def send_unchoke(self):
 
101
        self.events.append('unchoke')
 
102
 
 
103
class DummyChoker:
 
104
    def __init__(self, events):
 
105
        self.events = events
 
106
 
 
107
    def interested(self, connection):
 
108
        self.events.append('interested')
 
109
    
 
110
    def not_interested(self, connection):
 
111
        self.events.append('not interested')
 
112
 
 
113
class DummyStorage:
 
114
    def __init__(self, events):
 
115
        self.events = events
 
116
 
 
117
    def do_I_have_anything(self):
 
118
        self.events.append('do I have')
 
119
        return True
 
120
 
 
121
    def get_have_list(self):
 
122
        self.events.append('get have list')
 
123
        return [False, True]
 
124
 
 
125
    def get_piece(self, index, begin, length):
 
126
        self.events.append(('get piece', index, begin, length))
 
127
        if length == 4:
 
128
            return None
 
129
        return 'a' * length
 
130
 
 
131
def test_skip_over_choke():
 
132
    events = []
 
133
    dco = DummyConnection(events)
 
134
    dch = DummyChoker(events)
 
135
    ds = DummyStorage(events)
 
136
    u = Upload(dco, dch, ds, 100, 20, 5)
 
137
    assert u.is_choked()
 
138
    assert not u.is_interested()
 
139
    u.got_interested()
 
140
    assert u.is_interested()
 
141
    u.got_request(0, 0, 3)
 
142
    dco.flushed = True
 
143
    u.flushed()
 
144
    assert events == ['do I have', 'get have list', 
 
145
        ('bitfield', [False, True]), 'interested']
 
146
 
 
147
def test_bad_piece():
 
148
    events = []
 
149
    dco = DummyConnection(events)
 
150
    dch = DummyChoker(events)
 
151
    ds = DummyStorage(events)
 
152
    u = Upload(dco, dch, ds, 100, 20, 5)
 
153
    assert u.is_choked()
 
154
    assert not u.is_interested()
 
155
    u.got_interested()
 
156
    assert u.is_interested()
 
157
    u.unchoke()
 
158
    assert not u.is_choked()
 
159
    u.got_request(0, 0, 4)
 
160
    dco.flushed = True
 
161
    u.flushed()
 
162
    assert events == ['do I have', 'get have list', 
 
163
        ('bitfield', [False, True]), 'interested', 'unchoke', 
 
164
        ('get piece', 0, 0, 4), 'closed']
 
165
 
 
166
def test_still_rejected_after_unchoke():
 
167
    events = []
 
168
    dco = DummyConnection(events)
 
169
    dch = DummyChoker(events)
 
170
    ds = DummyStorage(events)
 
171
    u = Upload(dco, dch, ds, 100, 20, 5)
 
172
    assert u.is_choked()
 
173
    assert not u.is_interested()
 
174
    u.got_interested()
 
175
    assert u.is_interested()
 
176
    u.unchoke()
 
177
    assert not u.is_choked()
 
178
    u.got_request(0, 0, 3)
 
179
    u.choke()
 
180
    u.unchoke()
 
181
    dco.flushed = True
 
182
    u.flushed()
 
183
    assert events == ['do I have', 'get have list', 
 
184
        ('bitfield', [False, True]), 'interested', 'unchoke', 
 
185
        'choke', 'unchoke']
 
186
 
 
187
def test_sends_when_flushed():
 
188
    events = []
 
189
    dco = DummyConnection(events)
 
190
    dch = DummyChoker(events)
 
191
    ds = DummyStorage(events)
 
192
    u = Upload(dco, dch, ds, 100, 20, 5)
 
193
    u.unchoke()
 
194
    u.got_interested()
 
195
    u.got_request(0, 1, 3)
 
196
    dco.flushed = True
 
197
    u.flushed()
 
198
    u.flushed()
 
199
    assert events == ['do I have', 'get have list', 
 
200
        ('bitfield', [False, True]), 'unchoke', 'interested', 
 
201
        ('get piece', 0, 1, 3), ('piece', 0, 1, 'aaa')]
 
202
 
 
203
def test_sends_immediately():
 
204
    events = []
 
205
    dco = DummyConnection(events)
 
206
    dch = DummyChoker(events)
 
207
    ds = DummyStorage(events)
 
208
    u = Upload(dco, dch, ds, 100, 20, 5)
 
209
    u.unchoke()
 
210
    u.got_interested()
 
211
    dco.flushed = True
 
212
    u.got_request(0, 1, 3)
 
213
    assert events == ['do I have', 'get have list', 
 
214
        ('bitfield', [False, True]), 'unchoke', 'interested', 
 
215
        ('get piece', 0, 1, 3), ('piece', 0, 1, 'aaa')]
 
216
 
 
217
def test_cancel():
 
218
    events = []
 
219
    dco = DummyConnection(events)
 
220
    dch = DummyChoker(events)
 
221
    ds = DummyStorage(events)
 
222
    u = Upload(dco, dch, ds, 100, 20, 5)
 
223
    u.unchoke()
 
224
    u.got_interested()
 
225
    u.got_request(0, 1, 3)
 
226
    u.got_cancel(0, 1, 3)
 
227
    u.got_cancel(0, 1, 2)
 
228
    u.flushed()
 
229
    dco.flushed = True
 
230
    assert events == ['do I have', 'get have list', 
 
231
        ('bitfield', [False, True]), 'unchoke', 'interested']
 
232
 
 
233
def test_clears_on_not_interested():
 
234
    events = []
 
235
    dco = DummyConnection(events)
 
236
    dch = DummyChoker(events)
 
237
    ds = DummyStorage(events)
 
238
    u = Upload(dco, dch, ds, 100, 20, 5)
 
239
    u.unchoke()
 
240
    u.got_interested()
 
241
    u.got_request(0, 1, 3)
 
242
    u.got_not_interested()
 
243
    dco.flushed = True
 
244
    u.flushed()
 
245
    assert events == ['do I have', 'get have list', 
 
246
        ('bitfield', [False, True]), 'unchoke', 'interested', 
 
247
        'not interested']
 
248
 
 
249
def test_close_when_sends_on_not_interested():
 
250
    events = []
 
251
    dco = DummyConnection(events)
 
252
    dch = DummyChoker(events)
 
253
    ds = DummyStorage(events)
 
254
    u = Upload(dco, dch, ds, 100, 20, 5)
 
255
    u.got_request(0, 1, 3)
 
256
    assert events == ['do I have', 'get have list', 
 
257
        ('bitfield', [False, True]), 'closed']
 
258
 
 
259
def test_close_over_max_length():
 
260
    events = []
 
261
    dco = DummyConnection(events)
 
262
    dch = DummyChoker(events)
 
263
    ds = DummyStorage(events)
 
264
    u = Upload(dco, dch, ds, 100, 20, 5)
 
265
    u.got_interested()
 
266
    u.got_request(0, 1, 101)
 
267
    assert events == ['do I have', 'get have list', 
 
268
        ('bitfield', [False, True]), 'interested', 'closed']
 
269
 
 
270
def test_no_bitfield_on_start_empty():
 
271
    events = []
 
272
    dco = DummyConnection(events)
 
273
    dch = DummyChoker(events)
 
274
    ds = DummyStorage(events)
 
275
    ds.do_I_have_anything = lambda: False
 
276
    u = Upload(dco, dch, ds, 100, 20, 5)
 
277
    assert events == []