apt_dht/PeerManager.py

   1
   2 from random import choice
   3 from urlparse import urlparse, urlunparse
   4 from urllib import quote_plus
   5
   6 from twisted.internet import reactor, defer
   7 from twisted.python import log
   8 from twisted.trial import unittest
   9 from twisted.web2 import stream as stream_mod
  10 from twisted.web2.http import splitHostPort
  11
  12 from HTTPDownloader import HTTPClientManager
  13 from util import uncompact
  14
  15 class PeerManager:
  16     def __init__(self):
  17         self.clients = {}
  18
  19     def get(self, hash, mirror, peers = [], method="GET", modtime=None):
  20         """Download from a list of peers or fallback to a mirror.
  21
  22         @type peers: C{list} of C{string}
  23         @param peers: a list of the peers where the file can be found
  24         """
  25         if peers:
  26             compact_peer = choice(peers)
  27             peer = uncompact(compact_peer['c'])
  28             log.msg('Downloading from peer %r' % (peer, ))
  29             host, port = peer
  30             path = '/~/' + quote_plus(hash.expected())
  31         else:
  32             log.msg('Downloading (%s) from mirror %s' % (method, mirror))
  33             parsed = urlparse(mirror)
  34             assert parsed[0] == "http", "Only HTTP is supported, not '%s'" % parsed[0]
  35             host, port = splitHostPort(parsed[0], parsed[1])
  36             path = urlunparse(('', '') + parsed[2:])
  37
  38         return self.getPeer(host, port, path, method, modtime)
  39
  40     def getPeer(self, host, port, path, method="GET", modtime=None):
  41         if not port:
  42             port = 80
  43         site = host + ":" + str(port)
  44         if site not in self.clients:
  45             self.clients[site] = HTTPClientManager(host, port)
  46         return self.clients[site].get(path, method, modtime)
  47
  48     def close(self):
  49         for site in self.clients:
  50             self.clients[site].close()
  51         self.clients = {}
  52
  53 class TestPeerManager(unittest.TestCase):
  54     """Unit tests for the PeerManager."""
  55
  56     manager = None
  57     pending_calls = []
  58
  59     def gotResp(self, resp, num, expect):
  60         self.failUnless(resp.code >= 200 and resp.code < 300, "Got a non-200 response: %r" % resp.code)
  61         if expect is not None:
  62             self.failUnless(resp.stream.length == expect, "Length was incorrect, got %r, expected %r" % (resp.stream.length, expect))
  63         def print_(n):
  64             pass
  65         def printdone(n):
  66             pass
  67         stream_mod.readStream(resp.stream, print_).addCallback(printdone)
  68
  69     def test_download(self):
  70         self.manager = PeerManager()
  71         self.timeout = 10
  72
  73         host = 'www.ietf.org'
  74         d = self.manager.get('', 'http://' + host + '/rfc/rfc0013.txt')
  75         d.addCallback(self.gotResp, 1, 1070)
  76         return d
  77
  78     def test_head(self):
  79         self.manager = PeerManager()
  80         self.timeout = 10
  81
  82         host = 'www.ietf.org'
  83         d = self.manager.get('', 'http://' + host + '/rfc/rfc0013.txt', method = "HEAD")
  84         d.addCallback(self.gotResp, 1, 0)
  85         return d
  86
  87     def test_multiple_downloads(self):
  88         self.manager = PeerManager()
  89         self.timeout = 120
  90         lastDefer = defer.Deferred()
  91
  92         def newRequest(host, path, num, expect, last=False):
  93             d = self.manager.get('', 'http://' + host + ':' + str(80) + path)
  94             d.addCallback(self.gotResp, num, expect)
  95             if last:
  96                 d.addBoth(lastDefer.callback)
  97
  98         newRequest('www.ietf.org', "/rfc/rfc0006.txt", 1, 1776)
  99         newRequest('www.ietf.org', "/rfc/rfc2362.txt", 2, 159833)
 100         newRequest('www.google.ca', "/", 3, None)
 101         self.pending_calls.append(reactor.callLater(1, newRequest, 'www.sfu.ca', '/', 4, None))
 102         self.pending_calls.append(reactor.callLater(10, newRequest, 'www.ietf.org', '/rfc/rfc0048.txt', 5, 41696))
 103         self.pending_calls.append(reactor.callLater(30, newRequest, 'www.ietf.org', '/rfc/rfc0022.txt', 6, 4606))
 104         self.pending_calls.append(reactor.callLater(31, newRequest, 'www.sfu.ca', '/studentcentral/index.html', 7, None))
 105         self.pending_calls.append(reactor.callLater(32, newRequest, 'www.ietf.org', '/rfc/rfc0014.txt', 8, 27))
 106         self.pending_calls.append(reactor.callLater(32, newRequest, 'www.ietf.org', '/rfc/rfc0001.txt', 9, 21088))
 107         self.pending_calls.append(reactor.callLater(62, newRequest, 'www.google.ca', '/intl/en/options/', 0, None, True))
 108         return lastDefer
 109
 110     def tearDown(self):
 111         for p in self.pending_calls:
 112             if p.active():
 113                 p.cancel()
 114         self.pending_calls = []
 115         if self.manager:
 116             self.manager.close()
 117             self.manager = None