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