apt_p2p/PeerManager.py

   1
   2 """Manage a set of peers and the requests to them."""
   3
   4 from random import choice
   5 from urlparse import urlparse, urlunparse
   6 from urllib import quote_plus
   7 from binascii import b2a_hex, a2b_hex
   8 import sha
   9
  10 from twisted.internet import reactor, defer
  11 from twisted.python import log, filepath
  12 from twisted.trial import unittest
  13 from twisted.web2 import stream
  14 from twisted.web2.http import Response, splitHostPort
  15
  16 from HTTPDownloader import Peer
  17 from util import uncompact
  18 from Hash import PIECE_SIZE
  19 from apt_p2p_Khashmir.bencode import bdecode
  20 from apt_p2p_conf import config
  21
  22
  23 class PeerError(Exception):
  24     """An error occurred downloading from peers."""
  25
  26 class GrowingFileStream(stream.FileStream):
  27     """Modified to stream data from a file as it becomes available.
  28
  29     @ivar CHUNK_SIZE: the maximum size of chunks of data to send at a time
  30     @ivar deferred: waiting for the result of the last read attempt
  31     @ivar available: the number of bytes that are currently available to read
  32     @ivar position: the current position in the file where the next read will begin
  33     @ivar finished: True when no more data will be coming available
  34     """
  35
  36     CHUNK_SIZE = 4*1024
  37
  38     def __init__(self, f, length = None):
  39         stream.FileStream.__init__(self, f)
  40         self.length = length
  41         self.deferred = None
  42         self.available = 0L
  43         self.position = 0L
  44         self.finished = False
  45
  46     def updateAvailable(self, newlyAvailable):
  47         """Update the number of bytes that are available.
  48
  49         Call it with 0 to trigger reading of a fully read file.
  50
  51         @param newlyAvailable: the number of bytes that just became available
  52         """
  53         assert not self.finished
  54         self.available += newlyAvailable
  55
  56         # If a read is pending, let it go
  57         if self.deferred and self.position < self.available:
  58             # Try to read some data from the file
  59             length = self.available - self.position
  60             readSize = min(length, self.CHUNK_SIZE)
  61             self.f.seek(self.position)
  62             b = self.f.read(readSize)
  63             bytesRead = len(b)
  64
  65             # Check if end of file was reached
  66             if bytesRead:
  67                 self.position += bytesRead
  68                 deferred = self.deferred
  69                 self.deferred = None
  70                 deferred.callback(b)
  71
  72     def allAvailable(self):
  73         """Indicate that no more data will be coming available."""
  74         self.finished = True
  75
  76         # If a read is pending, let it go
  77         if self.deferred:
  78             if self.position < self.available:
  79                 # Try to read some data from the file
  80                 length = self.available - self.position
  81                 readSize = min(length, self.CHUNK_SIZE)
  82                 self.f.seek(self.position)
  83                 b = self.f.read(readSize)
  84                 bytesRead = len(b)
  85
  86                 # Check if end of file was reached
  87                 if bytesRead:
  88                     self.position += bytesRead
  89                     deferred = self.deferred
  90                     self.deferred = None
  91                     deferred.callback(b)
  92                 else:
  93                     # We're done
  94                     self._close()
  95                     deferred = self.deferred
  96                     self.deferred = None
  97                     deferred.callback(None)
  98             else:
  99                 # We're done
 100                 self._close()
 101                 deferred = self.deferred
 102                 self.deferred = None
 103                 deferred.callback(None)
 104
 105     def read(self, sendfile=False):
 106         assert not self.deferred, "A previous read is still deferred."
 107
 108         if self.f is None:
 109             return None
 110
 111         length = self.available - self.position
 112         readSize = min(length, self.CHUNK_SIZE)
 113
 114         # If we don't have any available, we're done or deferred
 115         if readSize <= 0:
 116             if self.finished:
 117                 self._close()
 118                 return None
 119             else:
 120                 self.deferred = defer.Deferred()
 121                 return self.deferred
 122
 123         # Try to read some data from the file
 124         self.f.seek(self.position)
 125         b = self.f.read(readSize)
 126         bytesRead = len(b)
 127         if not bytesRead:
 128             # End of file was reached, we're done or deferred
 129             if self.finished:
 130                 self._close()
 131                 return None
 132             else:
 133                 self.deferred = defer.Deferred()
 134                 return self.deferred
 135         else:
 136             self.position += bytesRead
 137             return b
 138
 139     def _close(self):
 140         """Close the temporary file and remove it."""
 141         self.f.close()
 142         filepath.FilePath(self.f.name).remove()
 143         self.f = None
 144
 145 class StreamToFile:
 146     """Save a stream to a partial file and hash it.
 147
 148     @type stream: L{twisted.web2.stream.IByteStream}
 149     @ivar stream: the input stream being read
 150     @type outFile: L{twisted.python.filepath.FilePath}
 151     @ivar outFile: the file being written
 152     @type hash: C{sha1}
 153     @ivar hash: the hash object for the data
 154     @type position: C{int}
 155     @ivar position: the current file position to write the next data to
 156     @type length: C{int}
 157     @ivar length: the position in the file to not write beyond
 158     @type doneDefer: L{twisted.internet.defer.Deferred}
 159     @ivar doneDefer: the deferred that will fire when done writing
 160     """
 161
 162     def __init__(self, inputStream, outFile, start = 0, length = None):
 163         """Initializes the file.
 164
 165         @type inputStream: L{twisted.web2.stream.IByteStream}
 166         @param inputStream: the input stream to read from
 167         @type outFile: L{twisted.python.filepath.FilePath}
 168         @param outFile: the file to write to
 169         @type start: C{int}
 170         @param start: the file position to start writing at
 171             (optional, defaults to the start of the file)
 172         @type length: C{int}
 173         @param length: the maximum amount of data to write to the file
 174             (optional, defaults to not limiting the writing to the file
 175         """
 176         self.stream = inputStream
 177         self.outFile = outFile
 178         self.hash = sha.new()
 179         self.position = start
 180         self.length = None
 181         if length is not None:
 182             self.length = start + length
 183         self.doneDefer = None
 184
 185     def run(self):
 186         """Start the streaming.
 187
 188         @rtype: L{twisted.internet.defer.Deferred}
 189         """
 190         log.msg('Started streaming %r bytes to file at position %d' % (self.length, self.position))
 191         self.doneDefer = stream.readStream(self.stream, self._gotData)
 192         self.doneDefer.addCallbacks(self._done, self._error)
 193         return self.doneDefer
 194
 195     def _gotData(self, data):
 196         """Process the received data."""
 197         if self.outFile.closed:
 198             raise PeerError, "outFile was unexpectedly closed"
 199
 200         if data is None:
 201             raise PeerError, "Data is None?"
 202
 203         # Make sure we don't go too far
 204         if self.length is not None and self.position + len(data) > self.length:
 205             data = data[:(self.length - self.position)]
 206
 207         # Write and hash the streamed data
 208         self.outFile.seek(self.position)
 209         self.outFile.write(data)
 210         self.hash.update(data)
 211         self.position += len(data)
 212
 213     def _done(self, result):
 214         """Return the result."""
 215         log.msg('Streaming is complete')
 216         return self.hash.digest()
 217
 218     def _error(self, err):
 219         """Log the error."""
 220         log.msg('Streaming error')
 221         log.err(err)
 222         return err
 223
 224 class FileDownload:
 225     """Manage a download from a list of peers or a mirror.
 226
 227     @type manager: L{PeerManager}
 228     @ivar manager: the manager to send requests for peers to
 229     @type hash: L{Hash.HashObject}
 230     @ivar hash: the hash object containing the expected hash for the file
 231     @ivar mirror: the URI of the file on the mirror
 232     @type compact_peers: C{list} of C{dictionary}
 233     @ivar compact_peers: a list of the peer info where the file can be found
 234     @type file: C{file}
 235     @ivar file: the open file to right the download to
 236     @type path: C{string}
 237     @ivar path: the path to request from peers to access the file
 238     @type pieces: C{list} of C{string}
 239     @ivar pieces: the hashes of the pieces in the file
 240     @type started: C{boolean}
 241     @ivar started: whether the download has begun yet
 242     @type defer: L{twisted.internet.defer.Deferred}
 243     @ivar defer: the deferred that will callback with the result of the download
 244     @type peers: C{dictionary}
 245     @ivar peers: information about each of the peers available to download from
 246     @type outstanding: C{int}
 247     @ivar outstanding: the number of requests to peers currently outstanding
 248     @type peerlist: C{list} of L{HTTPDownloader.Peer}
 249     @ivar peerlist: the sorted list of peers for this download
 250     @type stream: L{GrowingFileStream}
 251     @ivar stream: the stream of resulting data from the download
 252     @type nextFinish: C{int}
 253     @ivar nextFinish: the next piece that is needed to finish for the stream
 254     @type completePieces: C{list} of C{boolean} or L{HTTPDownloader.Peer}
 255     @ivar completePieces: one per piece, will be False if no requests are
 256         outstanding for the piece, True if the piece has been successfully
 257         downloaded, or the Peer that a request for this piece has been sent
 258     """
 259
 260     def __init__(self, manager, hash, mirror, compact_peers, file):
 261         """Initialize the instance and check for piece hashes.
 262
 263         @type manager: L{PeerManager}
 264         @param manager: the manager to send requests for peers to
 265         @type hash: L{Hash.HashObject}
 266         @param hash: the hash object containing the expected hash for the file
 267         @param mirror: the URI of the file on the mirror
 268         @type compact_peers: C{list} of C{dictionary}
 269         @param compact_peers: a list of the peer info where the file can be found
 270         @type file: L{twisted.python.filepath.FilePath}
 271         @param file: the temporary file to use to store the downloaded file
 272         """
 273         self.manager = manager
 274         self.hash = hash
 275         self.mirror = mirror
 276         self.compact_peers = compact_peers
 277
 278         self.path = '/~/' + quote_plus(hash.expected())
 279         self.defer = None
 280         self.mirror_path = None
 281         self.pieces = None
 282         self.started = False
 283
 284         file.restat(False)
 285         if file.exists():
 286             file.remove()
 287         self.file = file.open('w+')
 288
 289     def run(self):
 290         """Start the downloading process."""
 291         log.msg('Checking for pieces for %s' % self.path)
 292         self.defer = defer.Deferred()
 293         self.peers = {}
 294         no_pieces = 0
 295         pieces_string = {0: 0}
 296         pieces_hash = {0: 0}
 297         pieces_dl_hash = {0: 0}
 298
 299         for compact_peer in self.compact_peers:
 300             # Build a list of all the peers for this download
 301             site = uncompact(compact_peer['c'])
 302             peer = self.manager.getPeer(site)
 303             self.peers.setdefault(site, {})['peer'] = peer
 304
 305             # Extract any piece information from the peers list
 306             if 't' in compact_peer:
 307                 self.peers[site]['t'] = compact_peer['t']['t']
 308                 pieces_string.setdefault(compact_peer['t']['t'], 0)
 309                 pieces_string[compact_peer['t']['t']] += 1
 310             elif 'h' in compact_peer:
 311                 self.peers[site]['h'] = compact_peer['h']
 312                 pieces_hash.setdefault(compact_peer['h'], 0)
 313                 pieces_hash[compact_peer['h']] += 1
 314             elif 'l' in compact_peer:
 315                 self.peers[site]['l'] = compact_peer['l']
 316                 pieces_dl_hash.setdefault(compact_peer['l'], 0)
 317                 pieces_dl_hash[compact_peer['l']] += 1
 318             else:
 319                 no_pieces += 1
 320
 321         # Select the most popular piece info
 322         max_found = max(no_pieces, max(pieces_string.values()),
 323                         max(pieces_hash.values()), max(pieces_dl_hash.values()))
 324
 325         if max_found < len(self.peers):
 326             log.msg('Misleading piece information found, using most popular %d of %d peers' %
 327                     (max_found, len(self.peers)))
 328
 329         if max_found == no_pieces:
 330             # The file is not split into pieces
 331             log.msg('No pieces were found for the file')
 332             self.pieces = [self.hash.expected()]
 333             self.startDownload()
 334         elif max_found == max(pieces_string.values()):
 335             # Small number of pieces in a string
 336             for pieces, num in pieces_string.items():
 337                 # Find the most popular piece string
 338                 if num == max_found:
 339                     self.pieces = [pieces[x:x+20] for x in xrange(0, len(pieces), 20)]
 340                     log.msg('Peer info contained %d piece hashes' % len(self.pieces))
 341                     self.startDownload()
 342                     break
 343         elif max_found == max(pieces_hash.values()):
 344             # Medium number of pieces stored in the DHT
 345             for pieces, num in pieces_hash.items():
 346                 # Find the most popular piece hash to lookup
 347                 if num == max_found:
 348                     log.msg('Found a hash for pieces to lookup in the DHT: %r' % pieces)
 349                     self.getDHTPieces(pieces)
 350                     break
 351         elif max_found == max(pieces_dl_hash.values()):
 352             # Large number of pieces stored in peers
 353             for pieces, num in pieces_dl_hash.items():
 354                 # Find the most popular piece hash to download
 355                 if num == max_found:
 356                     log.msg('Found a hash for pieces to lookup in peers: %r' % pieces)
 357                     self.getPeerPieces(pieces)
 358                     break
 359         return self.defer
 360
 361     #{ Downloading the piece hashes
 362     def getDHTPieces(self, key):
 363         """Retrieve the piece information from the DHT.
 364
 365         @param key: the key to lookup in the DHT
 366         """
 367         # Remove any peers with the wrong piece hash
 368         #for site in self.peers.keys():
 369         #    if self.peers[site].get('h', '') != key:
 370         #        del self.peers[site]
 371
 372         # Start the DHT lookup
 373         lookupDefer = self.manager.dht.getValue(key)
 374         lookupDefer.addBoth(self._getDHTPieces, key)
 375
 376     def _getDHTPieces(self, results, key):
 377         """Check the retrieved values."""
 378         if isinstance(results, list):
 379             for result in results:
 380                 # Make sure the hash matches the key
 381                 result_hash = sha.new(result.get('t', '')).digest()
 382                 if result_hash == key:
 383                     pieces = result['t']
 384                     self.pieces = [pieces[x:x+20] for x in xrange(0, len(pieces), 20)]
 385                     log.msg('Retrieved %d piece hashes from the DHT' % len(self.pieces))
 386                     self.startDownload()
 387                     return
 388
 389             log.msg('Could not retrieve the piece hashes from the DHT')
 390         else:
 391             log.msg('Looking up piece hashes in the DHT resulted in an error: %r' % (result, ))
 392
 393         # Continue without the piece hashes
 394         self.pieces = [None for x in xrange(0, self.hash.expSize, PIECE_SIZE)]
 395         self.startDownload()
 396
 397     def getPeerPieces(self, key, failedSite = None):
 398         """Retrieve the piece information from the peers.
 399
 400         @param key: the key to request from the peers
 401         """
 402         if failedSite is None:
 403             log.msg('Starting the lookup of piece hashes in peers')
 404             self.outstanding = 0
 405             # Remove any peers with the wrong piece hash
 406             #for site in self.peers.keys():
 407             #    if self.peers[site].get('l', '') != key:
 408             #        del self.peers[site]
 409         else:
 410             log.msg('Piece hash lookup failed for peer %r' % (failedSite, ))
 411             self.peers[failedSite]['failed'] = True
 412             self.outstanding -= 1
 413
 414         if self.pieces is None:
 415             # Send a request to one or more peers
 416             for site in self.peers:
 417                 if self.peers[site].get('failed', False) != True:
 418                     log.msg('Sending a piece hash request to %r' % (site, ))
 419                     path = '/~/' + quote_plus(key)
 420                     lookupDefer = self.peers[site]['peer'].get(path)
 421                     reactor.callLater(0, lookupDefer.addCallbacks,
 422                                       *(self._getPeerPieces, self._gotPeerError),
 423                                       **{'callbackArgs': (key, site),
 424                                          'errbackArgs': (key, site)})
 425                     self.outstanding += 1
 426                     if self.outstanding >= 4:
 427                         break
 428
 429         if self.pieces is None and self.outstanding <= 0:
 430             # Continue without the piece hashes
 431             log.msg('Could not retrieve the piece hashes from the peers')
 432             self.pieces = [None for x in xrange(0, self.hash.expSize, PIECE_SIZE)]
 433             self.startDownload()
 434
 435     def _getPeerPieces(self, response, key, site):
 436         """Process the retrieved response from the peer."""
 437         log.msg('Got a piece hash response %d from %r' % (response.code, site))
 438         if response.code != 200:
 439             # Request failed, try a different peer
 440             self.getPeerPieces(key, site)
 441         else:
 442             # Read the response stream to a string
 443             self.peers[site]['pieces'] = ''
 444             def _gotPeerPiece(data, self = self, site = site):
 445                 log.msg('Peer %r got %d bytes of piece hashes' % (site, len(data)))
 446                 self.peers[site]['pieces'] += data
 447             log.msg('Streaming piece hashes from peer')
 448             df = stream.readStream(response.stream, _gotPeerPiece)
 449             df.addCallbacks(self._gotPeerPieces, self._gotPeerError,
 450                             callbackArgs=(key, site), errbackArgs=(key, site))
 451
 452     def _gotPeerError(self, err, key, site):
 453         """Peer failed, try again."""
 454         log.msg('Peer piece hash request failed for %r' % (site, ))
 455         log.err(err)
 456         self.getPeerPieces(key, site)
 457
 458     def _gotPeerPieces(self, result, key, site):
 459         """Check the retrieved pieces from the peer."""
 460         log.msg('Finished streaming piece hashes from peer %r' % (site, ))
 461         if self.pieces is not None:
 462             # Already done
 463             log.msg('Already done')
 464             return
 465
 466         try:
 467             result = bdecode(self.peers[site]['pieces'])
 468         except:
 469             log.msg('Error bdecoding piece hashes')
 470             log.err()
 471             self.getPeerPieces(key, site)
 472             return
 473
 474         result_hash = sha.new(result.get('t', '')).digest()
 475         if result_hash == key:
 476             pieces = result['t']
 477             self.pieces = [pieces[x:x+20] for x in xrange(0, len(pieces), 20)]
 478             log.msg('Retrieved %d piece hashes from the peer %r' % (len(self.pieces), site))
 479             self.startDownload()
 480         else:
 481             log.msg('Peer returned a piece string that did not match')
 482             self.getPeerPieces(key, site)
 483
 484     #{ Downloading the file
 485     def sort(self):
 486         """Sort the peers by their rank (highest ranked at the end)."""
 487         def sort(a, b):
 488             """Sort peers by their rank."""
 489             if a.rank > b.rank:
 490                 return 1
 491             elif a.rank < b.rank:
 492                 return -1
 493             return 0
 494         self.peerlist.sort(sort)
 495
 496     def startDownload(self):
 497         """Start the download from the peers."""
 498         # Don't start twice
 499         if self.started:
 500             return
 501
 502         log.msg('Starting to download %s' % self.path)
 503         self.started = True
 504         assert self.pieces, "You must initialize the piece hashes first"
 505         self.peerlist = [self.peers[site]['peer'] for site in self.peers]
 506
 507         # Use the mirror if there are few peers
 508         if len(self.peerlist) < config.getint('DEFAULT', 'MIN_DOWNLOAD_PEERS'):
 509             parsed = urlparse(self.mirror)
 510             if parsed[0] == "http":
 511                 site = splitHostPort(parsed[0], parsed[1])
 512                 self.mirror_path = urlunparse(('', '') + parsed[2:])
 513                 peer = self.manager.getPeer(site, mirror = True)
 514                 self.peerlist.append(peer)
 515
 516         # Special case if there's only one good peer left
 517 #        if len(self.peerlist) == 1:
 518 #            log.msg('Downloading from peer %r' % (self.peerlist[0], ))
 519 #            self.defer.callback(self.peerlist[0].get(self.path))
 520 #            return
 521
 522         # Begin to download the pieces
 523         self.outstanding = 0
 524         self.nextFinish = 0
 525         self.completePieces = [False for piece in self.pieces]
 526         self.getPieces()
 527
 528     #{ Downloading the pieces
 529     def getPieces(self):
 530         """Download the next pieces from the peers."""
 531         self.sort()
 532         piece = self.nextFinish
 533         while self.outstanding < 4 and self.peerlist and piece < len(self.completePieces):
 534             if self.completePieces[piece] == False:
 535                 # Send a request to the highest ranked peer
 536                 peer = self.peerlist.pop()
 537                 self.completePieces[piece] = peer
 538                 log.msg('Sending a request for piece %d to peer %r' % (piece, peer))
 539
 540                 self.outstanding += 1
 541                 if peer.mirror:
 542                     df = peer.getRange(self.mirror_path, piece*PIECE_SIZE, (piece+1)*PIECE_SIZE - 1)
 543                 else:
 544                     df = peer.getRange(self.path, piece*PIECE_SIZE, (piece+1)*PIECE_SIZE - 1)
 545                 reactor.callLater(0, df.addCallbacks,
 546                                   *(self._getPiece, self._getError),
 547                                   **{'callbackArgs': (piece, peer),
 548                                      'errbackArgs': (piece, peer)})
 549             piece += 1
 550
 551         # Check if we're done
 552         if self.outstanding <= 0 and self.nextFinish >= len(self.completePieces):
 553             log.msg('We seem to be done with all pieces')
 554             self.stream.allAvailable()
 555
 556     def _getPiece(self, response, piece, peer):
 557         """Process the retrieved headers from the peer."""
 558         log.msg('Got response for piece %d from peer %r' % (piece, peer))
 559         if ((len(self.completePieces) > 1 and response.code != 206) or
 560             (response.code not in (200, 206))):
 561             # Request failed, try a different peer
 562             log.msg('Wrong response type %d for piece %d from peer %r' % (response.code, piece, peer))
 563             peer.hashError('Peer responded with the wrong type of download: %r' % response.code)
 564             self.completePieces[piece] = False
 565             if response.stream and response.stream.length:
 566                 stream.readAndDiscard(response.stream)
 567         else:
 568             if self.defer:
 569                 # Start sending the return file
 570                 df = self.defer
 571                 self.defer = None
 572                 self.stream = GrowingFileStream(self.file, self.hash.expSize)
 573
 574                 # Get the headers from the peer's response
 575                 headers = {}
 576                 if response.headers.hasHeader('last-modified'):
 577                     headers['last-modified'] = response.headers.getHeader('last-modified')
 578                 resp = Response(200, {}, self.stream)
 579                 df.callback(resp)
 580
 581             # Read the response stream to the file
 582             log.msg('Streaming piece %d from peer %r' % (piece, peer))
 583             if response.code == 206:
 584                 df = StreamToFile(response.stream, self.file, piece*PIECE_SIZE,
 585                                   PIECE_SIZE).run()
 586             else:
 587                 df = StreamToFile(response.stream, self.file).run()
 588             reactor.callLater(0, df.addCallbacks,
 589                               *(self._gotPiece, self._gotError),
 590                               **{'callbackArgs': (piece, peer),
 591                                  'errbackArgs': (piece, peer)})
 592
 593         self.outstanding -= 1
 594         self.peerlist.append(peer)
 595         self.getPieces()
 596
 597     def _getError(self, err, piece, peer):
 598         """Peer failed, try again."""
 599         log.msg('Got error for piece %d from peer %r' % (piece, peer))
 600         self.outstanding -= 1
 601         self.peerlist.append(peer)
 602         self.completePieces[piece] = False
 603         self.getPieces()
 604         log.err(err)
 605
 606     def _gotPiece(self, response, piece, peer):
 607         """Process the retrieved piece from the peer."""
 608         log.msg('Finished streaming piece %d from peer %r: %r' % (piece, peer, response))
 609         if self.pieces[piece] and response != self.pieces[piece]:
 610             # Hash doesn't match
 611             log.msg('Hash error for piece %d from peer %r' % (piece, peer))
 612             peer.hashError('Piece received from peer does not match expected')
 613             self.completePieces[piece] = False
 614         else:
 615             # Successfully completed one of several pieces
 616             log.msg('Finished with piece %d from peer %r' % (piece, peer))
 617             self.completePieces[piece] = True
 618             while (self.nextFinish < len(self.completePieces) and
 619                    self.completePieces[self.nextFinish] == True):
 620                 self.nextFinish += 1
 621                 self.stream.updateAvailable(PIECE_SIZE)
 622
 623         self.getPieces()
 624
 625     def _gotError(self, err, piece, peer):
 626         """Piece download failed, try again."""
 627         log.msg('Error streaming piece %d from peer %r: %r' % (piece, peer, response))
 628         log.err(err)
 629         self.completePieces[piece] = False
 630         self.getPieces()
 631
 632 class PeerManager:
 633     """Manage a set of peers and the requests to them.
 634
 635     @type cache_dir: L{twisted.python.filepath.FilePath}
 636     @ivar cache_dir: the directory to use for storing all files
 637     @type dht: L{interfaces.IDHT}
 638     @ivar dht: the DHT instance
 639     @type stats: L{stats.StatsLogger}
 640     @ivar stats: the statistics logger to record sent data to
 641     @type clients: C{dictionary}
 642     @ivar clients: the available peers that have been previously contacted
 643     """
 644
 645     def __init__(self, cache_dir, dht, stats):
 646         """Initialize the instance."""
 647         self.cache_dir = cache_dir
 648         self.cache_dir.restat(False)
 649         if not self.cache_dir.exists():
 650             self.cache_dir.makedirs()
 651         self.dht = dht
 652         self.stats = stats
 653         self.clients = {}
 654
 655     def get(self, hash, mirror, peers = [], method="GET", modtime=None):
 656         """Download from a list of peers or fallback to a mirror.
 657
 658         @type hash: L{Hash.HashObject}
 659         @param hash: the hash object containing the expected hash for the file
 660         @param mirror: the URI of the file on the mirror
 661         @type peers: C{list} of C{string}
 662         @param peers: a list of the peer info where the file can be found
 663             (optional, defaults to downloading from the mirror)
 664         @type method: C{string}
 665         @param method: the HTTP method to use, 'GET' or 'HEAD'
 666             (optional, defaults to 'GET')
 667         @type modtime: C{int}
 668         @param modtime: the modification time to use for an 'If-Modified-Since'
 669             header, as seconds since the epoch
 670             (optional, defaults to not sending that header)
 671         """
 672         if not peers or method != "GET" or modtime is not None:
 673             log.msg('Downloading (%s) from mirror %s' % (method, mirror))
 674             parsed = urlparse(mirror)
 675             assert parsed[0] == "http", "Only HTTP is supported, not '%s'" % parsed[0]
 676             site = splitHostPort(parsed[0], parsed[1])
 677             path = urlunparse(('', '') + parsed[2:])
 678             peer = self.getPeer(site, mirror = True)
 679             return peer.get(path, method, modtime)
 680 #        elif len(peers) == 1:
 681 #            site = uncompact(peers[0]['c'])
 682 #            log.msg('Downloading from peer %r' % (site, ))
 683 #            path = '/~/' + quote_plus(hash.expected())
 684 #            peer = self.getPeer(site)
 685 #            return peer.get(path)
 686         else:
 687             tmpfile = self.cache_dir.child(hash.hexexpected())
 688             return FileDownload(self, hash, mirror, peers, tmpfile).run()
 689
 690     def getPeer(self, site, mirror = False):
 691         """Create a new peer if necessary and return it.
 692
 693         @type site: (C{string}, C{int})
 694         @param site: the IP address and port of the peer
 695         @param mirror: whether the peer is actually a mirror
 696             (optional, defaults to False)
 697         """
 698         if site not in self.clients:
 699             self.clients[site] = Peer(site[0], site[1], self.stats)
 700             if mirror:
 701                 self.clients[site].mirror = True
 702         return self.clients[site]
 703
 704     def close(self):
 705         """Close all the connections to peers."""
 706         for site in self.clients:
 707             self.clients[site].close()
 708         self.clients = {}
 709
 710 class TestPeerManager(unittest.TestCase):
 711     """Unit tests for the PeerManager."""
 712
 713     manager = None
 714     pending_calls = []
 715
 716     def tearDown(self):
 717         for p in self.pending_calls:
 718             if p.active():
 719                 p.cancel()
 720         self.pending_calls = []
 721         if self.manager:
 722             self.manager.close()
 723             self.manager = None