[tor-commits] [arm/master] Adding bandwidth_from_state() util

atagar at torproject.org atagar at torproject.org
Sun Sep 14 00:27:05 UTC 2014


commit b8e8e64a851baf78a493e12bf675266a2458414d
Author: Damian Johnson <atagar at torproject.org>
Date:   Sat Sep 13 17:24:03 2014 -0700

    Adding bandwidth_from_state() util
    
    Moving the bulk of our function for prepopulating bandwidth information from
    the state file to our util, and adding tests for it.
---
 arm/graphing/bandwidth_stats.py   |   72 +++---------------------
 arm/util/__init__.py              |   94 +++++++++++++++++++++++++++++++
 test/util/__init__.py             |    2 +-
 test/util/bandwidth_from_state.py |  111 +++++++++++++++++++++++++++++++++++++
 4 files changed, 214 insertions(+), 65 deletions(-)

diff --git a/arm/graphing/bandwidth_stats.py b/arm/graphing/bandwidth_stats.py
index 1067e90..8dcf203 100644
--- a/arm/graphing/bandwidth_stats.py
+++ b/arm/graphing/bandwidth_stats.py
@@ -4,17 +4,16 @@ stats if they're set.
 """
 
 import calendar
-import os
 import time
 import curses
 
 import arm.controller
 
 from arm.graphing import graph_panel
-from arm.util import tor_controller
+from arm.util import bandwidth_from_state, tor_controller
 
 from stem.control import State
-from stem.util import conf, str_tools, system
+from stem.util import conf, str_tools
 
 
 def conf_handler(key, value):
@@ -135,70 +134,15 @@ class BandwidthStats(graph_panel.GraphStats):
     returns True if successful and False otherwise.
     """
 
-    controller = tor_controller()
-
-    if not controller.is_localhost():
-      raise ValueError('we can only prepopulate bandwidth information for a local tor instance')
-
-    start_time = system.start_time(controller.get_pid(None))
-    uptime = time.time() - start_time if start_time else None
-
-    # Only attempt to prepopulate information if we've been running for a day.
-    # Reason is that the state file stores a day's worth of data, and we don't
-    # want to prepopulate with information from a prior tor instance.
-
-    if not uptime:
-      raise ValueError("unable to determine tor's uptime")
-    elif uptime < (24 * 60 * 60):
-      raise ValueError("insufficient uptime, tor must've been running for at least a day")
-
-    # read the user's state file in their data directory (usually '~/.tor')
-
-    data_dir = controller.get_conf('DataDirectory', None)
-
-    if not data_dir:
-      raise ValueError("unable to determine tor's data directory")
-
-    state_path = os.path.join(CONFIG['tor.chroot'] + data_dir, 'state')
-
-    try:
-      with open(state_path) as state_file:
-        state_content = state_file.readlines()
-    except IOError as exc:
-      raise ValueError('unable to read the state file at %s, %s' % (state_path, exc))
-
-    # We're interested in two types of entries from our state file...
-    #
-    # * BWHistory*Values - Comma separated list of bytes we read or wrote
-    #   during each fifteen minute period. The last value is an incremental
-    #   counter for our current period, so ignoring that.
-    #
-    # * BWHistory*Ends - When our last sampling was recorded, in UTC.
-
-    bw_read_entries, bw_write_entries = None, None
-    missing_read_entries, missing_write_entries = None, None
-
-    for line in state_content:
-      line = line.strip()
-
-      if line.startswith('BWHistoryReadValues '):
-        bw_read_entries = [int(entry) / 1024.0 / 900 for entry in line[20:].split(',')[:-1]]
-      elif line.startswith('BWHistoryWriteValues '):
-        bw_write_entries = [int(entry) / 1024.0 / 900 for entry in line[21:].split(',')[:-1]]
-      elif line.startswith('BWHistoryReadEnds '):
-        last_read_time = calendar.timegm(time.strptime(line[18:], '%Y-%m-%d %H:%M:%S')) - 900
-        missing_read_entries = int((time.time() - last_read_time) / 900)
-      elif line.startswith('BWHistoryWriteEnds '):
-        last_write_time = calendar.timegm(time.strptime(line[19:], '%Y-%m-%d %H:%M:%S')) - 900
-        missing_write_entries = int((time.time() - last_write_time) / 900)
+    stats = bandwidth_from_state()
 
-    if not bw_read_entries or not bw_write_entries or not last_read_time or not last_write_time:
-      raise ValueError('bandwidth stats missing from state file')
+    missing_read_entries = int((time.time() - stats.last_read_time) / 900)
+    missing_write_entries = int((time.time() - stats.last_write_time) / 900)
 
     # fills missing entries with the last value
 
-    bw_read_entries += [bw_read_entries[-1]] * missing_read_entries
-    bw_write_entries += [bw_write_entries[-1]] * missing_write_entries
+    bw_read_entries = stats.read_entries + [stats.read_entries[-1]] * missing_read_entries
+    bw_write_entries = stats.write_entries + [stats.write_entries[-1]] * missing_write_entries
 
     # crops starting entries so they're the same size
 
@@ -236,7 +180,7 @@ class BandwidthStats(graph_panel.GraphStats):
     del self.primary_counts[interval_index][self.max_column + 1:]
     del self.secondary_counts[interval_index][self.max_column + 1:]
 
-    return time.time() - min(last_read_time, last_write_time)
+    return time.time() - min(stats.last_read_time, stats.last_write_time)
 
   def bandwidth_event(self, event):
     if self.is_accounting and self.is_next_tick_redraw():
diff --git a/arm/util/__init__.py b/arm/util/__init__.py
index de40170..c410284 100644
--- a/arm/util/__init__.py
+++ b/arm/util/__init__.py
@@ -12,17 +12,28 @@ __all__ = [
   'ui_tools',
 ]
 
+import calendar
+import collections
 import os
 import sys
+import time
 
 import stem.connection
 import stem.util.conf
+import stem.util.system
 
 from arm.util import log
 
 TOR_CONTROLLER = None
 BASE_DIR = os.path.sep.join(__file__.split(os.path.sep)[:-2])
 
+StateBandwidth = collections.namedtuple('StateBandwidth', (
+  'read_entries',
+  'write_entries',
+  'last_read_time',
+  'last_write_time',
+))
+
 try:
   uses_settings = stem.util.conf.uses_settings('arm', os.path.join(BASE_DIR, 'config'), lazy_load = False)
 except IOError as exc:
@@ -69,3 +80,86 @@ def msg(message, config, **attr):
   except:
     log.notice('BUG: We attempted to use an undefined string resource (%s)' % message)
     return ''
+
+
+ at uses_settings
+def bandwidth_from_state(config):
+  """
+  Read Tor's state file to determine its recent bandwidth usage. These
+  samplings are at fifteen minute granularity, and can only provide results if
+  we've been running for at least a day. This provides a named tuple with the
+  following...
+
+    * read_entries and write_entries
+
+      List of the average kilobytes read or written during each fifteen minute
+      period, oldest to newest.
+
+    * last_read_time and last_write_time
+
+      Unix timestamp for when the last entry was recorded.
+
+  :returns: **namedtuple** with the state file's bandwidth informaiton
+
+  :raises: **ValueError** if unable to get the bandwidth information from our
+    state file
+  """
+
+  controller = tor_controller()
+
+  if not controller.is_localhost():
+    raise ValueError('we can only prepopulate bandwidth information for a local tor instance')
+
+  start_time = stem.util.system.start_time(controller.get_pid(None))
+  uptime = time.time() - start_time if start_time else None
+
+  # Only attempt to prepopulate information if we've been running for a day.
+  # Reason is that the state file stores a day's worth of data, and we don't
+  # want to prepopulate with information from a prior tor instance.
+
+  if not uptime:
+    raise ValueError("unable to determine tor's uptime")
+  elif uptime < (24 * 60 * 60):
+    raise ValueError("insufficient uptime, tor must've been running for at least a day")
+
+  # read the user's state file in their data directory (usually '~/.tor')
+
+  data_dir = controller.get_conf('DataDirectory', None)
+
+  if not data_dir:
+    raise ValueError("unable to determine tor's data directory")
+
+  state_path = os.path.join(config.get('tor.chroot', '') + data_dir, 'state')
+
+  try:
+    with open(state_path) as state_file:
+      state_content = state_file.readlines()
+  except IOError as exc:
+    raise ValueError('unable to read the state file at %s, %s' % (state_path, exc))
+
+  # We're interested in two types of entries from our state file...
+  #
+  # * BWHistory*Values - Comma separated list of bytes we read or wrote
+  #   during each fifteen minute period. The last value is an incremental
+  #   counter for our current period, so ignoring that.
+  #
+  # * BWHistory*Ends - When our last sampling was recorded, in UTC.
+
+  attr = {}
+
+  for line in state_content:
+    line = line.strip()
+
+    if line.startswith('BWHistoryReadValues '):
+      attr['read_entries'] = [int(entry) / 1024.0 / 900 for entry in line[20:].split(',')[:-1]]
+    elif line.startswith('BWHistoryWriteValues '):
+      attr['write_entries'] = [int(entry) / 1024.0 / 900 for entry in line[21:].split(',')[:-1]]
+    elif line.startswith('BWHistoryReadEnds '):
+      attr['last_read_time'] = calendar.timegm(time.strptime(line[18:], '%Y-%m-%d %H:%M:%S')) - 900
+    elif line.startswith('BWHistoryWriteEnds '):
+      attr['last_write_time'] = calendar.timegm(time.strptime(line[19:], '%Y-%m-%d %H:%M:%S')) - 900
+
+  if len(attr) != 4:
+    raise ValueError('bandwidth stats missing from state file')
+
+  return StateBandwidth(**attr)
diff --git a/test/util/__init__.py b/test/util/__init__.py
index 3795845..f338157 100644
--- a/test/util/__init__.py
+++ b/test/util/__init__.py
@@ -2,4 +2,4 @@
 Unit tests for arm's utilities.
 """
 
-__all__ = []
+__all__ = ['bandwidth_from_state']
diff --git a/test/util/bandwidth_from_state.py b/test/util/bandwidth_from_state.py
new file mode 100644
index 0000000..b84c251
--- /dev/null
+++ b/test/util/bandwidth_from_state.py
@@ -0,0 +1,111 @@
+import datetime
+import io
+import time
+import unittest
+
+from mock import Mock, patch
+
+from arm.util import bandwidth_from_state
+
+STATE_FILE = """\
+# Tor state file last generated on 2014-07-20 13:05:10 local time
+# Other times below are in UTC
+# You *do not* need to edit this file.
+
+EntryGuard mullbinde7 2546FD2B50165C1567A297B02AD73F62DEA127A0 DirCache
+EntryGuardAddedBy 2546FD2B50165C1567A297B02AD73F62DEA127A0 0.2.4.10-alpha-dev 2014-07-11 01:18:47
+EntryGuardPathBias 9.000000 9.000000 9.000000 0.000000 0.000000 1.000000
+TorVersion Tor 0.2.4.10-alpha-dev (git-8be6058d8f31e578)
+LastWritten 2014-07-20 20:05:10
+TotalBuildTimes 68
+CircuitBuildTimeBin 525 1
+CircuitBuildTimeBin 575 1
+CircuitBuildTimeBin 675 1
+"""
+
+STATE_FILE_WITH_ENTRIES = STATE_FILE + """\
+BWHistoryReadValues 921600,1843200,2764800,3686400,4608000
+BWHistoryWriteValues 46080000,46080000,92160000,92160000,92160000
+BWHistoryReadEnds %s
+BWHistoryWriteEnds %s
+"""
+
+
+class TestBandwidthFromState(unittest.TestCase):
+  @patch('arm.util.tor_controller')
+  def test_when_not_localhost(self, tor_controller_mock):
+    tor_controller_mock().is_localhost.return_value = False
+
+    try:
+      bandwidth_from_state()
+      self.fail('expected a ValueError')
+    except ValueError as exc:
+      self.assertEqual('we can only prepopulate bandwidth information for a local tor instance', str(exc))
+
+  @patch('arm.util.tor_controller')
+  def test_unknown_pid(self, tor_controller_mock):
+    tor_controller_mock().is_localhost.return_value = True
+    tor_controller_mock().get_pid.return_value = None
+
+    try:
+      bandwidth_from_state()
+      self.fail('expected a ValueError')
+    except ValueError as exc:
+      self.assertEqual("unable to determine tor's uptime", str(exc))
+
+  @patch('arm.util.tor_controller')
+  @patch('stem.util.system.start_time')
+  def test_insufficient_uptime(self, start_time_mock, tor_controller_mock):
+    tor_controller_mock().is_localhost.return_value = True
+    start_time_mock.return_value = time.time() - 60  # one minute of uptime
+
+    try:
+      bandwidth_from_state()
+      self.fail('expected a ValueError')
+    except ValueError as exc:
+      self.assertEqual("insufficient uptime, tor must've been running for at least a day", str(exc))
+
+  @patch('arm.util.tor_controller')
+  @patch('stem.util.system.start_time', Mock(return_value = 50))
+  def test_no_data_dir(self, tor_controller_mock):
+    tor_controller_mock().is_localhost.return_value = True
+    tor_controller_mock().get_conf.return_value = None
+
+    try:
+      bandwidth_from_state()
+      self.fail('expected a ValueError')
+    except ValueError as exc:
+      self.assertEqual("unable to determine tor's data directory", str(exc))
+
+  @patch('arm.util.tor_controller')
+  @patch('arm.util.open', create = True)
+  @patch('stem.util.system.start_time', Mock(return_value = 50))
+  def test_no_bandwidth_entries(self, open_mock, tor_controller_mock):
+    tor_controller_mock().is_localhost.return_value = True
+    tor_controller_mock().get_conf.return_value = '/home/atagar/.tor'
+    open_mock.return_value = io.BytesIO(STATE_FILE)
+
+    try:
+      bandwidth_from_state()
+      self.fail('expected a ValueError')
+    except ValueError as exc:
+      self.assertEqual('bandwidth stats missing from state file', str(exc))
+
+    open_mock.assert_called_once_with('/home/atagar/.tor/state')
+
+  @patch('arm.util.tor_controller')
+  @patch('arm.util.open', create = True)
+  @patch('stem.util.system.start_time', Mock(return_value = 50))
+  def test_when_successful(self, open_mock, tor_controller_mock):
+    tor_controller_mock().is_localhost.return_value = True
+    tor_controller_mock().get_conf.return_value = '/home/atagar/.tor'
+
+    now = int(time.time())
+    timestamp = datetime.datetime.utcfromtimestamp(now + 900).strftime('%Y-%m-%d %H:%M:%S')
+    open_mock.return_value = io.BytesIO(STATE_FILE_WITH_ENTRIES % (timestamp, timestamp))
+
+    stats = bandwidth_from_state()
+    self.assertEqual([1, 2, 3, 4], stats.read_entries)
+    self.assertEqual([50, 50, 100, 100], stats.write_entries)
+    self.assertEqual(now, stats.last_read_time)
+    self.assertEqual(now, stats.last_write_time)



More information about the tor-commits mailing list