Choker.py 7.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207
  1. # The contents of this file are subject to the BitTorrent Open Source License
  2. # Version 1.1 (the License). You may not copy or use this file, in either
  3. # source code or executable form, except in compliance with the License. You
  4. # may obtain a copy of the License at http://www.bittorrent.com/license/.
  5. #
  6. # Software distributed under the License is distributed on an AS IS basis,
  7. # WITHOUT WARRANTY OF ANY KIND, either express or implied. See the License
  8. # for the specific language governing rights and limitations under the
  9. # License.
  10. # Written by Greg Hazel
  11. # based on code by Bram Cohen
  12. from __future__ import division
  13. import math
  14. import random
  15. from BTL.obsoletepythonsupport import set
  16. class Choker(object):
  17. def __init__(self, config, schedule):
  18. self.config = config
  19. self.schedule = schedule
  20. self.connections = []
  21. self.count = 0
  22. self.unchokes_since_last = 0
  23. self.interval = 5
  24. self.shutting_down = False
  25. #self.magic_number = 6 # magic 6 : (30 / self.interval)
  26. self.magic_number = (30 / self.interval)
  27. schedule(self.interval, self._round_robin)
  28. def _round_robin(self):
  29. self.schedule(self.interval, self._round_robin)
  30. self.count += 1
  31. # don't do more work than you have to
  32. if not self.connections:
  33. return
  34. # rotation for round-robin
  35. if self.count % self.magic_number == 0:
  36. for i, c in enumerate(self.connections):
  37. u = c.upload
  38. if u.choked and u.interested:
  39. self.connections = self.connections[i:] + self.connections[:i]
  40. break
  41. self._rechoke()
  42. ## new
  43. ############################################################################
  44. def _rechoke(self):
  45. # step 1:
  46. # get sorted in order of preference lists of peers
  47. # one for downloading torrents, and one for seeding torrents
  48. down_pref = []
  49. seed_pref = []
  50. for i, c in enumerate(self.connections):
  51. u = c.upload
  52. if c.download.have.numfalse == 0 or not u.interested:
  53. continue
  54. # I cry.
  55. if c.download.multidownload.storage.have.numfalse != 0:
  56. ## heuristic for downloading torrents
  57. if not c.download.is_snubbed():
  58. ## simple download rate based
  59. down_pref.append((-c.download.get_rate(), i))
  60. ## ratio based
  61. #dr = c.download.get_rate()
  62. #ur = max(1, u.get_rate())
  63. #ratio = dr / ur
  64. #down_pref.append((-ratio, i))
  65. else:
  66. ## heuristic for seeding torrents
  67. ## Uoti special
  68. ## if c._decrypt is not None:
  69. ## seed_pref.append((self.count, u.get_rate(), i))
  70. ## elif (u.unchoke_time > self.count - self.magic_number or
  71. ## u.buffer and c.connection.is_flushed()):
  72. ## seed_pref.append((u.unchoke_time, u.get_rate(), i))
  73. ## else:
  74. ## seed_pref.append((1, u.get_rate(), i))
  75. ## sliding, first pass (see below)
  76. r = u.get_rate()
  77. if c._decrypt is not None:
  78. seed_pref.append((2, r, i))
  79. else:
  80. seed_pref.append((1, r, i))
  81. down_pref.sort()
  82. seed_pref.sort()
  83. #pprint(down_pref)
  84. #pprint(seed_pref)
  85. down_pref = [ self.connections[i] for junk, i in down_pref ]
  86. seed_pref = [ self.connections[i] for junk, junk, i in seed_pref ]
  87. max_uploads = self._max_uploads()
  88. ## sliding, second pass
  89. ## # up-side-down sum for an idea of capacity
  90. ## uprate_sum = sum(rates[-max_uploads:])
  91. ## if max_uploads == 0:
  92. ## avg_uprate = 0
  93. ## else:
  94. ## avg_uprate = uprate_sum / max_uploads
  95. ## #print 'avg_uprate', avg_uprate, 'of', max_uploads
  96. ## self.extra_slots = max(self.extra_slots - 1, 0)
  97. ## if avg_uprate > self.arbitrary_min:
  98. ## for r in rates:
  99. ## if r < (avg_uprate * 0.80): # magic 80%
  100. ## self.extra_slots += 2
  101. ## break
  102. ## self.extra_slots = min(len(seed_pref), self.extra_slots)
  103. ## max_uploads += self.extra_slots
  104. ## #print 'plus', self.extra_slots
  105. # step 2:
  106. # split the peer lists by a ratio to fill the available upload slots
  107. d_uploads = max(1, int(round(max_uploads * 0.70)))
  108. s_uploads = max(1, int(round(max_uploads * 0.30)))
  109. #print 'original', 'ds', d_uploads, 'us', s_uploads
  110. extra = max(0, d_uploads - len(down_pref))
  111. if extra > 0:
  112. s_uploads += extra
  113. d_uploads -= extra
  114. extra = max(0, s_uploads - len(seed_pref))
  115. if extra > 0:
  116. s_uploads -= extra
  117. d_uploads = min(d_uploads + extra, len(down_pref))
  118. #print 'ds', d_uploads, 'us', s_uploads
  119. down_pref = down_pref[:d_uploads]
  120. seed_pref = seed_pref[:s_uploads]
  121. preferred = set(down_pref)
  122. preferred.update(seed_pref)
  123. # step 3:
  124. # enforce unchoke states
  125. count = 0
  126. to_choke = []
  127. for i, c in enumerate(self.connections):
  128. u = c.upload
  129. if c in preferred:
  130. u.unchoke(self.count)
  131. count += 1
  132. else:
  133. to_choke.append(c)
  134. # step 4:
  135. # enforce choke states and handle optimistics
  136. optimistics = max(self.config['min_uploads'],
  137. max_uploads - len(preferred))
  138. #print 'optimistics', optimistics
  139. for c in to_choke:
  140. u = c.upload
  141. if c.download.have.numfalse == 0:
  142. u.choke()
  143. elif count >= optimistics:
  144. u.choke()
  145. else:
  146. # this one's optimistic
  147. u.unchoke(self.count)
  148. if u.interested:
  149. count += 1
  150. ############################################################################
  151. def shutdown(self):
  152. self.shutting_down = True
  153. def connection_made(self, connection):
  154. p = random.randrange(len(self.connections) + 1)
  155. self.connections.insert(p, connection)
  156. def connection_lost(self, connection):
  157. self.connections.remove(connection)
  158. if (not self.shutting_down and
  159. connection.upload.interested and not connection.upload.choked):
  160. self._rechoke()
  161. def interested(self, connection):
  162. if not connection.upload.choked:
  163. self._rechoke()
  164. def not_interested(self, connection):
  165. if not connection.upload.choked:
  166. self._rechoke()
  167. def _max_uploads(self):
  168. uploads = self.config['max_uploads']
  169. rate = self.config['max_upload_rate'] / 1024
  170. if uploads > 0:
  171. pass
  172. elif rate <= 0:
  173. uploads = 7 # unlimited, just guess something here...
  174. elif rate < 9:
  175. uploads = 2
  176. elif rate < 15:
  177. uploads = 3
  178. elif rate < 42:
  179. uploads = 4
  180. else:
  181. uploads = int(math.sqrt(rate * .6))
  182. return uploads