[tor-commits] [stem/master] Implementing a use_directory_mirrors() method

[atagar at torproject.org]

commit 4005a88becb68858ec4742455ea4d80e1ff9f827
Author: Damian Johnson <atagar at torproject.org>
Date:   Fri Jul 19 10:12:39 2013 -0700

    Implementing a use_directory_mirrors() method
    
    Method so we can balance load against directory mirrors rather than hammering
    the authorities. We can either request this during construction (which fails
    silently) or call this method explicitly (to get the exception).
---
 stem/descriptor/remote.py       |   37 ++++++++++++++++++++++++++++++++++++-
 test/integ/descriptor/remote.py |   14 ++++++++++++++
 2 files changed, 50 insertions(+), 1 deletion(-)

diff --git a/stem/descriptor/remote.py b/stem/descriptor/remote.py
index 66e6d85..9df09f4 100644
--- a/stem/descriptor/remote.py
+++ b/stem/descriptor/remote.py
@@ -48,6 +48,7 @@ import urllib2
 
 import stem.descriptor
 
+from stem import Flag
 from stem.util import log
 
 # Tor has a limit on the number of descriptors we can fetch explicitly by their
@@ -250,6 +251,9 @@ class DescriptorDownloader(object):
   :var int retries: number of times to attempt the request if it fails
   :var float timeout: duration before we'll time out our request, no timeout is
     applied if **None**
+  :var bool use_mirrors: downloads the present consensus and uses the directory
+    mirrors to fetch future requests, this fails silently if the consensus
+    cannot be downloaded
   :var bool start_when_requested: issues requests when our methods are called
     if **True**, otherwise provides non-running
     :class:`~stem.descriptor.remote.Query` instances
@@ -257,13 +261,42 @@ class DescriptorDownloader(object):
     request to a directory authority if **True**
   """
 
-  def __init__(self, retries = 2, fall_back_to_authority = True, timeout = None, start_when_requested = True):
+  def __init__(self, retries = 2, use_mirrors = False, fall_back_to_authority = True, timeout = None, start_when_requested = True):
     self.retries = retries
     self.timeout = timeout
     self.start_when_requested = start_when_requested
     self.fall_back_to_authority = fall_back_to_authority
     self._endpoints = DIRECTORY_AUTHORITIES.values()
 
+    if use_mirrors:
+      try:
+        start_time = time.time()
+        self.use_directory_mirrors()
+        log.debug("Retrieve directory mirrors (took %0.2fs)" % (time.time() - start_time))
+      except Exception, exc:
+        log.debug("Unable to retrieve directory mirrors: %s" % exc)
+
+  def use_directory_mirrors(self):
+    """
+    Downloads the present consensus and configures ourselves to use directory
+    mirrors, in addition to authorities.
+
+    :raises: **Exception** if unable to determine the directory mirrors
+    """
+
+    new_endpoints = set(DIRECTORY_AUTHORITIES.values())
+
+    query = self.get_consensus()
+    query.run()  # running explicitly so we'll raise errors
+
+    for desc in query:
+      if Flag.V2DIR in desc.flags:
+        new_endpoints.add((desc.address, desc.dir_port))
+
+    # we need our endpoints to be a list rather than set for random.choice()
+
+    self._endpoints = list(new_endpoints)
+
   def get_server_descriptors(self, fingerprints = None):
     """
     Provides the server descriptors with the given fingerprints. If no
@@ -351,6 +384,8 @@ class DescriptorDownloader(object):
     Issues a request for the given resource.
     """
 
+    log.trace("Retrieving descriptors (resource: %s, type: %s)" % (resource, descriptor_type))
+
     return Query(
       resource,
       descriptor_type,
diff --git a/test/integ/descriptor/remote.py b/test/integ/descriptor/remote.py
index 9273b61..a705486 100644
--- a/test/integ/descriptor/remote.py
+++ b/test/integ/descriptor/remote.py
@@ -49,6 +49,20 @@ class TestDescriptorReader(unittest.TestCase):
       self.assertEqual(1, len(descriptors))
       self.assertEqual('moria1', descriptors[0].nickname)
 
+  def test_use_directory_mirrors(self):
+    """
+    Checks that we can fetch and use a list of directory mirrors.
+    """
+
+    if test.runner.require_online(self):
+      return
+    elif test.runner.only_run_once(self, "test_use_directory_mirrors"):
+      return
+
+    downloader = stem.descriptor.remote.DescriptorDownloader()
+    downloader.use_directory_mirrors()
+    self.assertTrue(len(downloader._endpoints) > 50)
+
   def test_get_server_descriptors(self):
     """
     Exercises the downloader's get_server_descriptors() method.