[tor-commits] [nyx/master] Use cache for resolving addresses to relays

[atagar at torproject.org]

commit 451cda3f5456f3b3eaf84af1c86f6e618bcd0cae
Author: Damian Johnson <atagar at torproject.org>
Date:   Sat Sep 2 17:53:15 2017 -0700

    Use cache for resolving addresses to relays
---
 nyx/__init__.py | 23 ++++++++++++++++++++---
 nyx/tracker.py  | 10 +---------
 test/cache.py   | 18 ++++++++++++++++++
 3 files changed, 39 insertions(+), 12 deletions(-)

diff --git a/nyx/__init__.py b/nyx/__init__.py
index b3cbafb..2a5cf57 100644
--- a/nyx/__init__.py
+++ b/nyx/__init__.py
@@ -109,12 +109,13 @@ stem.control.CACHEABLE_GETINFO_PARAMS = list(stem.control.CACHEABLE_GETINFO_PARA
 
 stem.control.LOG_CACHE_FETCHES = False
 
-SCHEMA_VERSION = 1  # version of our scheme, bump this if you change the following
+SCHEMA_VERSION = 2  # version of our scheme, bump this if you change the following
 SCHEMA = (
-  'CREATE TABLE schema(version NUMBER)',
+  'CREATE TABLE schema(version INTEGER)',
   'INSERT INTO schema(version) VALUES (%i)' % SCHEMA_VERSION,
 
-  'CREATE TABLE relays(fingerprint TEXT PRIMARY KEY, address TEXT, or_port NUMBER, nickname TEXT)',
+  'CREATE TABLE relays(fingerprint TEXT PRIMARY KEY, address TEXT, or_port INTEGER, nickname TEXT)',
+  'CREATE INDEX addresses ON relays(address)',
 )
 
 
@@ -399,6 +400,22 @@ class Cache(object):
     with self._conn:
       yield CacheWriter(self)
 
+  def relays_for_address(self, address):
+    """
+    Provides the relays running at a given location.
+
+    :param str address: address to be checked
+
+    :returns: **dict** of ORPorts to their fingerprint
+    """
+
+    result = {}
+
+    for entry in self._query('SELECT or_port, fingerprint FROM relays WHERE address=?', address).fetchall():
+      result[entry[0]] = entry[1]
+
+    return result
+
   def relay_nickname(self, fingerprint, default = None):
     """
     Provides the nickname associated with the given relay.
diff --git a/nyx/tracker.py b/nyx/tracker.py
index cbfaedf..2530bf2 100644
--- a/nyx/tracker.py
+++ b/nyx/tracker.py
@@ -808,8 +808,6 @@ class ConsensusTracker(object):
   """
 
   def __init__(self):
-    self._fingerprint_cache = {}  # {address => [(port, fingerprint), ..]} for relays
-
     self._my_router_status_entry = None
     self._my_router_status_entry_time = 0
 
@@ -831,7 +829,6 @@ class ConsensusTracker(object):
             fingerprint = stem.descriptor.router_status_entry._base64_to_hex(r_comp[2])
             nickname = r_comp[1]
 
-            self._fingerprint_cache.setdefault(address, []).append((or_port, fingerprint))
             writer.record_relay(fingerprint, address, or_port, nickname)
 
       stem.util.log.info('Cached consensus data, took %0.2fs.' % (time.time() - start_time))
@@ -845,22 +842,17 @@ class ConsensusTracker(object):
     :param list router_status_entries: router status entries to populate our cache with
     """
 
-    new_fingerprint_cache = {}
-
     start_time = time.time()
     our_fingerprint = tor_controller().get_info('fingerprint', None)
 
     with nyx.cache().write() as writer:
       for desc in router_status_entries:
-        new_fingerprint_cache.setdefault(desc.address, []).append((desc.or_port, desc.fingerprint))
         writer.record_relay(desc.fingerprint, desc.address, desc.or_port, desc.nickname)
 
         if desc.fingerprint == our_fingerprint:
           self._my_router_status_entry = desc
           self._my_router_status_entry_time = time.time()
 
-    self._fingerprint_cache = new_fingerprint_cache
-
     stem.util.log.info('Updated consensus cache, took %0.2fs.' % (time.time() - start_time))
 
   def my_router_status_entry(self):
@@ -915,7 +907,7 @@ class ConsensusTracker(object):
       if fingerprint and ports:
         return dict([(port, fingerprint) for port in ports])
 
-    return dict([(port, fp) for (port, fp) in self._fingerprint_cache.get(address, [])])
+    return nyx.cache().relays_for_address(address)
 
   def get_relay_address(self, fingerprint, default):
     """
diff --git a/test/cache.py b/test/cache.py
index dd76c00..281ad86 100644
--- a/test/cache.py
+++ b/test/cache.py
@@ -47,6 +47,24 @@ class TestCache(unittest.TestCase):
         self.assertEqual('caersidi', cache.relay_nickname('3EA8E960F6B94CE30062AA8EF02894C00F8D1E66'))
 
   @patch('nyx.data_directory', Mock(return_value = None))
+  def test_relays_for_address(self):
+    """
+    Basic checks for registering and fetching nicknames.
+    """
+
+    cache = nyx.cache()
+
+    with cache.write() as writer:
+      writer.record_relay('3EA8E960F6B94CE30062AA8EF02894C00F8D1E66', '208.113.165.162', 1443, 'caersidi1')
+      writer.record_relay('9695DFC35FFEB861329B9F1AB04C46397020CE31', '128.31.0.34', 9101, 'moria1')
+      writer.record_relay('74A910646BCEEFBCD2E874FC1DC997430F968145', '208.113.165.162', 1543, 'caersidi2')
+
+    self.assertEqual({9101: '9695DFC35FFEB861329B9F1AB04C46397020CE31'}, cache.relays_for_address('128.31.0.34'))
+    self.assertEqual({1443: '3EA8E960F6B94CE30062AA8EF02894C00F8D1E66', 1543: '74A910646BCEEFBCD2E874FC1DC997430F968145'}, cache.relays_for_address('208.113.165.162'))
+
+    self.assertEqual({}, cache.relays_for_address('199.254.238.53'))
+
+  @patch('nyx.data_directory', Mock(return_value = None))
   def test_relay_nickname(self):
     """
     Basic checks for registering and fetching nicknames.