Source Code for Module Anomos.Choker

  1  # This program is free software: you can redistribute it and/or modify 
  2  # it under the terms of the GNU General Public License as published by 
  3  # the Free Software Foundation, either version 3 of the License, or 
  4  # (at your option) any later version. 
  5  # 
  6  # This program is distributed in the hope that it will be useful, 
  7  # but WITHOUT ANY WARRANTY; without even the implied warranty of 
  8  # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the 
  9  # GNU General Public License for more details. 
 10  # 
 11  # You should have received a copy of the GNU General Public License 
 12  # along with this program.  If not, see <http://www.gnu.org/licenses/>. 
 13   
 14  # Written by Bram Cohen 
 15   
 16  from random import randrange 
 17   
 18 -class Choker(object): 
 19   
 20 -    def __init__(self, config, schedule, done = lambda: False): 
 21          self.config = config 
 22          self.schedule = schedule 
 23          self.connections = [] 
 24          self.count = 0 
 25          self.done = done 
 26          self.unchokes_since_last = 0 
 27          schedule(10, self._round_robin) 
 28   
 29 -    def _round_robin(self): 
 30          self.schedule(10, self._round_robin) 
 31          self.count += 1 
 32          if self.done(): 
 33              self._rechoke_seed(True) 
 34              return 
 35          if self.count % 3 == 0: 
 36              for i in xrange(len(self.connections)): 
 37                  u = self.connections[i].upload 
 38                  if u.choked and u.interested: 
 39                      self.connections = self.connections[i:] + self.connections[:i] 
 40                      break 
 41          self._rechoke() 
 42   
 43 -    def _rechoke(self): 
 44          if self.done(): 
 45              self._rechoke_seed() 
 46              return 
 47          preferred = [] 
 48          for i in xrange(len(self.connections)): 
 49              c = self.connections[i] 
 50              if c.upload.interested and not c.download.is_snubbed(): 
 51                  preferred.append((-c.download.get_rate(), i)) 
 52          preferred.sort() 
 53          prefcount = min(len(preferred), self.config['max_uploads_internal'] -1) 
 54          mask = [0] * len(self.connections) 
 55          for _, i in preferred[:prefcount]: 
 56              mask[i] = 1 
 57          count = max(1, self.config['min_uploads'] - prefcount) 
 58          for i in xrange(len(self.connections)): 
 59              c = self.connections[i] 
 60              u = c.upload 
 61              if mask[i]: 
 62                  u.unchoke(self.count) 
 63              elif count > 0: 
 64                  u.unchoke(self.count) 
 65                  if u.interested: 
 66                      count -= 1 
 67              else: 
 68                  u.choke() 
 69   
 70 -    def _rechoke_seed(self, force_new_unchokes = False): 
 71          if force_new_unchokes: 
 72              # number of unchokes per 30 second period 
 73              i = (self.config['max_uploads_internal'] + 2) // 3 
 74              # this is called 3 times in 30 seconds, if i==4 then unchoke 1+1+2 
 75              # and so on; substract unchokes recently triggered by disconnects 
 76              num_force_unchokes = max(0, (i + self.count % 3) // 3 - \ 
 77                                   self.unchokes_since_last) 
 78          else: 
 79              num_force_unchokes = 0 
 80          preferred = [] 
 81          new_limit = self.count - 3 
 82          for i in xrange(len(self.connections)): 
 83              c = self.connections[i] 
 84              u = c.upload 
 85              if not u.choked and u.interested: 
 86                  if u.unchoke_time > new_limit or ( 
 87                          u.buffer and c.is_flushed()): 
 88                      preferred.append((-u.unchoke_time, -u.get_rate(), i)) 
 89                  else: 
 90                      preferred.append((1, -u.get_rate(), i)) 
 91          num_kept = self.config['max_uploads_internal'] - num_force_unchokes 
 92          assert num_kept >= 0 
 93          preferred.sort() 
 94          preferred = preferred[:num_kept] 
 95          mask = [0] * len(self.connections) 
 96          for _, _, i in preferred: 
 97              mask[i] = 1 
 98          num_nonpref = self.config['max_uploads_internal'] - len(preferred) 
 99          if force_new_unchokes: 
100              self.unchokes_since_last = 0 
101          else: 
102              self.unchokes_since_last += num_nonpref 
103          last_unchoked = None 
104          for i in xrange(len(self.connections)): 
105              c = self.connections[i] 
106              u = c.upload 
107              if not mask[i]: 
108                  if not u.interested: 
109                      u.choke() 
110                  elif u.choked: 
111                      if num_nonpref > 0 and c.is_flushed(): 
112                          u.unchoke(self.count) 
113                          num_nonpref -= 1 
114                          if num_nonpref == 0: 
115                              last_unchoked = i 
116                  else: 
117                      if num_nonpref == 0: 
118                          u.choke() 
119                      else: 
120                          num_nonpref -= 1 
121                          if num_nonpref == 0: 
122                              last_unchoked = i 
123          if last_unchoked is not None: 
124              self.connections = self.connections[last_unchoked + 1:] + \ 
125                                 self.connections[:last_unchoked + 1] 
126   
127 -    def connection_made(self, connection): 
128          if connection in self.connections: 
129              return 
130          p = randrange(len(self.connections) + 1) 
131          self.connections.insert(p, connection) 
132   
133 -    def connection_lost(self, connection): 
134          self.connections.remove(connection) 
135          if connection.upload.interested and not connection.upload.choked: 
136              self._rechoke() 
137   
138 -    def interested(self, connection): 
139          if not connection.upload.choked: 
140              self._rechoke() 
141   
142 -    def not_interested(self, connection): 
143          if not connection.upload.choked: 
144              self._rechoke() 
145