diff options
Diffstat (limited to 'BitTorrent/Choker.py')
| -rw-r--r-- | BitTorrent/Choker.py | 157 |
1 files changed, 157 insertions, 0 deletions
diff --git a/BitTorrent/Choker.py b/BitTorrent/Choker.py new file mode 100644 index 0000000..ed3d500 --- /dev/null +++ b/BitTorrent/Choker.py @@ -0,0 +1,157 @@ +# The contents of this file are subject to the BitTorrent Open Source License +# Version 1.1 (the License). You may not copy or use this file, in either +# source code or executable form, except in compliance with the License. You +# may obtain a copy of the License at http://www.bittorrent.com/license/. +# +# Software distributed under the License is distributed on an AS IS basis, +# WITHOUT WARRANTY OF ANY KIND, either express or implied. See the License +# for the specific language governing rights and limitations under the +# License. + +# Written by Bram Cohen + +from random import randrange +from math import sqrt + +class Choker(object): + + def __init__(self, config, schedule, done = lambda: False): + self.config = config + self.schedule = schedule + self.connections = [] + self.count = 0 + self.done = done + self.unchokes_since_last = 0 + schedule(self._round_robin, 10) + + def _round_robin(self): + self.schedule(self._round_robin, 10) + self.count += 1 + if self.done(): + self._rechoke_seed(True) + return + if self.count % 3 == 0: + for i in xrange(len(self.connections)): + u = self.connections[i].upload + if u.choked and u.interested: + self.connections = self.connections[i:] + self.connections[:i] + break + self._rechoke() + + def _rechoke(self): + if self.done(): + self._rechoke_seed() + return + preferred = [] + for i in xrange(len(self.connections)): + c = self.connections[i] + if c.upload.interested and not c.download.is_snubbed() and c.download.have.numfalse: + preferred.append((-c.download.get_rate(), i)) + preferred.sort() + prefcount = min(len(preferred), self._max_uploads() -1) + mask = [0] * len(self.connections) + for _, i in preferred[:prefcount]: + mask[i] = 1 + count = max(1, self.config['min_uploads'] - prefcount) + for i in xrange(len(self.connections)): + c = self.connections[i] + u = c.upload + if mask[i]: + u.unchoke(self.count) + elif count > 0 and c.download.have.numfalse: + u.unchoke(self.count) + if u.interested: + count -= 1 + else: + u.choke() + + def _rechoke_seed(self, force_new_unchokes = False): + if force_new_unchokes: + # number of unchokes per 30 second period + i = (self._max_uploads() + 2) // 3 + # this is called 3 times in 30 seconds, if i==4 then unchoke 1+1+2 + # and so on; substract unchokes recently triggered by disconnects + num_force_unchokes = max(0, (i + self.count % 3) // 3 - \ + self.unchokes_since_last) + else: + num_force_unchokes = 0 + preferred = [] + new_limit = self.count - 3 + for i in xrange(len(self.connections)): + c = self.connections[i] + u = c.upload + if not u.choked and u.interested and c.download.have.numfalse: + if u.unchoke_time > new_limit or ( + u.buffer and c.connection.is_flushed()): + preferred.append((-u.unchoke_time, -u.get_rate(), i)) + else: + preferred.append((1, -u.get_rate(), i)) + num_kept = self._max_uploads() - num_force_unchokes + assert num_kept >= 0 + preferred.sort() + preferred = preferred[:num_kept] + mask = [0] * len(self.connections) + for _, _, i in preferred: + mask[i] = 1 + num_nonpref = self._max_uploads() - len(preferred) + if force_new_unchokes: + self.unchokes_since_last = 0 + else: + self.unchokes_since_last += num_nonpref + last_unchoked = None + for i in xrange(len(self.connections)): + c = self.connections[i] + u = c.upload + if not mask[i]: + if not u.interested: + u.choke() + elif u.choked: + if num_nonpref > 0 and c.connection.is_flushed() and c.download.have.numfalse: + u.unchoke(self.count) + num_nonpref -= 1 + if num_nonpref == 0: + last_unchoked = i + else: + if num_nonpref == 0 or not c.download.have.numfalse: + u.choke() + else: + num_nonpref -= 1 + if num_nonpref == 0: + last_unchoked = i + if last_unchoked is not None: + self.connections = self.connections[last_unchoked + 1:] + \ + self.connections[:last_unchoked + 1] + + def connection_made(self, connection): + p = randrange(len(self.connections) + 1) + self.connections.insert(p, connection) + + def connection_lost(self, connection): + self.connections.remove(connection) + if connection.upload.interested and not connection.upload.choked: + self._rechoke() + + def interested(self, connection): + if not connection.upload.choked: + self._rechoke() + + def not_interested(self, connection): + if not connection.upload.choked: + self._rechoke() + + def _max_uploads(self): + uploads = self.config['max_uploads'] + rate = self.config['max_upload_rate'] + if uploads > 0: + pass + elif rate <= 0: + uploads = 7 # unlimited, just guess something here... + elif rate < 9: + uploads = 2 + elif rate < 15: + uploads = 3 + elif rate < 42: + uploads = 4 + else: + uploads = int(sqrt(rate * .6)) + return uploads |