commit b8e8e64a851baf78a493e12bf675266a2458414d
Author: Damian Johnson <atagar(a)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 ''
+
+
+@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)
