diff --git a/src/wormhole_transit_relay/server_state.py b/src/wormhole_transit_relay/server_state.py index 0e7c446..b98615e 100644 --- a/src/wormhole_transit_relay/server_state.py +++ b/src/wormhole_transit_relay/server_state.py @@ -67,11 +67,105 @@ class TestClient(object): print("disconnect_partner: {}".format(id(self._partner))) +class IUsageWriter(Interface): + """ + Records actual usage statistics in some way + """ + + def record_usage(started=None, total_time=None, waiting_time=None, total_bytes=None, mood=None): + """ + :param int started: timestemp when this connection began + + :param float total_time: total seconds this connection lasted + + :param float waiting_time: None or the total seconds one side + waited for the other + + :param int total_bytes: the total bytes sent. In case the + connection was concluded successfully, only one side will + record the total bytes (but count both). + + :param str mood: the 'mood' of the connection + """ + + +@implementer(IUsageWriter) +class MemoryUsageRecorder: + + def __init__(self): + self.events = [] + + def record_usage(self, started=None, total_time=None, waiting_time=None, total_bytes=None, mood=None): + """ + IUsageWriter. + """ + data = { + "started": started, + "total_time": total_time, + "waiting_time": waiting_time, + "total_bytes": total_bytes, + "mood": mood, + } + self.events.append(data) + + +@implementer(IUsageWriter) +class LogFileUsageRecorder: + + def __init__(self, writable_file): + self._file = writable_file + + def record_usage(self, started=None, total_time=None, waiting_time=None, total_bytes=None, mood=None): + """ + IUsageWriter. + """ + data = { + "started": started, + "total_time": total_time, + "waiting_time": waiting_time, + "total_bytes": total_bytes, + "mood": mood, + } + self._file.write(json.dumps(data)) + + + +@implementer(IUsageWriter) +class DatabaseUsageRecorder: + + def __init__(self, _db): + self._db = db + + def record_usage(self, started=None, total_time=None, waiting_time=None, total_bytes=None, mood=None): + """ + IUsageWriter. + """ + + class UsageRecorder(object): """ Tracks usage statistics of connections """ + def __init__(self): + self._backends = set() + + def add_backend(self, backend): + """ + Add a new backend. + + :param IUsageWriter backend: the backend to add + """ + self._backends.add(backend) + + def remove_backend(self, backend): + """ + Remove an existing backend + + :param IUsageWriter backend: the backend to remove + """ + self._backends.remove(backend) + def record(self, started, buddy_started, result, bytes_sent, buddy_bytes): """ :param int started: timestamp when our connection started @@ -102,7 +196,7 @@ class UsageRecorder(object): # probably want like "backends" here or something? original # code logs some JSON (maybe) and writes to a database (maybe) # and tests record in memory. - self.json_record({ + self._notify_backends({ "started": started, "total_time": total_time, "waiting_time": waiting_time, @@ -110,8 +204,12 @@ class UsageRecorder(object): "mood": result, }) - def json_record(self, data): - pass + def _notify_backends(self, data): + """ + Internal helper. Tell every backend we have about a new usage. + """ + for backend in self._backends: + backend.record_usage(**data) class ActiveConnections(object): diff --git a/src/wormhole_transit_relay/test/test_transit_server.py b/src/wormhole_transit_relay/test/test_transit_server.py index 09615c7..22c8ee4 100644 --- a/src/wormhole_transit_relay/test/test_transit_server.py +++ b/src/wormhole_transit_relay/test/test_transit_server.py @@ -3,6 +3,7 @@ from binascii import hexlify from twisted.trial import unittest from .common import ServerBase from .. import transit_server +from ..server_state import MemoryUsageRecorder def handshake(token, side=None): hs = b"please relay " + hexlify(token) @@ -341,8 +342,8 @@ class Usage(ServerBase, unittest.TestCase): def setUp(self): super(Usage, self).setUp() - self._usage = [] - self._transit_server.usage.json_record = self._usage.append + self._usage = MemoryUsageRecorder() + self._transit_server.usage.add_backend(self._usage) def test_empty(self): p1 = self.new_protocol() @@ -351,8 +352,8 @@ class Usage(ServerBase, unittest.TestCase): self.flush() # that will log the "empty" usage event - self.assertEqual(len(self._usage), 1, self._usage) - self.assertEqual(self._usage[0]["mood"], "empty", self._usage) + self.assertEqual(len(self._usage.events), 1, self._usage) + self.assertEqual(self._usage.events[0]["mood"], "empty", self._usage) def test_short(self): p1 = self.new_protocol() @@ -362,8 +363,8 @@ class Usage(ServerBase, unittest.TestCase): self.flush() # that will log the "empty" usage event - self.assertEqual(len(self._usage), 1, self._usage) - self.assertEqual("empty", self._usage[0]["mood"]) + self.assertEqual(len(self._usage.events), 1, self._usage) + self.assertEqual("empty", self._usage.events[0]["mood"]) def test_errory(self): p1 = self.new_protocol() @@ -372,8 +373,8 @@ class Usage(ServerBase, unittest.TestCase): self.flush() # that will log the "errory" usage event, then drop the connection p1.disconnect() - self.assertEqual(len(self._usage), 1, self._usage) - self.assertEqual(self._usage[0]["mood"], "errory", self._usage) + self.assertEqual(len(self._usage.events), 1, self._usage) + self.assertEqual(self._usage.events[0]["mood"], "errory", self._usage) def test_lonely(self): p1 = self.new_protocol() @@ -386,9 +387,9 @@ class Usage(ServerBase, unittest.TestCase): p1.disconnect() self.flush() - self.assertEqual(len(self._usage), 1, self._usage) - self.assertEqual(self._usage[0]["mood"], "lonely", self._usage) - self.assertIdentical(self._usage[0]["waiting_time"], None) + self.assertEqual(len(self._usage.events), 1, self._usage) + self.assertEqual(self._usage.events[0]["mood"], "lonely", self._usage) + self.assertIdentical(self._usage.events[0]["waiting_time"], None) def test_one_happy_one_jilted(self): p1 = self.new_protocol() @@ -402,7 +403,7 @@ class Usage(ServerBase, unittest.TestCase): p2.send(handshake(token1, side=side2)) self.flush() - self.assertEqual(self._usage, []) # no events yet + self.assertEqual(self._usage.events, []) # no events yet p1.send(b"\x00" * 13) self.flush() @@ -412,10 +413,10 @@ class Usage(ServerBase, unittest.TestCase): p1.disconnect() self.flush() - self.assertEqual(len(self._usage), 1, self._usage) - self.assertEqual(self._usage[0]["mood"], "happy", self._usage) - self.assertEqual(self._usage[0]["total_bytes"], 20) - self.assertNotIdentical(self._usage[0]["waiting_time"], None) + self.assertEqual(len(self._usage.events), 1, self._usage) + self.assertEqual(self._usage.events[0]["mood"], "happy", self._usage) + self.assertEqual(self._usage.events[0]["total_bytes"], 20) + self.assertNotIdentical(self._usage.events[0]["waiting_time"], None) def test_redundant(self): p1a = self.new_protocol() @@ -438,18 +439,18 @@ class Usage(ServerBase, unittest.TestCase): p1c.disconnect() self.flush() - self.assertEqual(len(self._usage), 1, self._usage) - self.assertEqual(self._usage[0]["mood"], "lonely") + self.assertEqual(len(self._usage.events), 1, self._usage) + self.assertEqual(self._usage.events[0]["mood"], "lonely") p2.send(handshake(token1, side=side2)) self.flush() self.assertEqual(len(self._transit_server.pending_requests._requests), 0) - self.assertEqual(len(self._usage), 2, self._usage) - self.assertEqual(self._usage[1]["mood"], "redundant") + self.assertEqual(len(self._usage.events), 2, self._usage) + self.assertEqual(self._usage.events[1]["mood"], "redundant") # one of the these is unecessary, but probably harmless p1a.disconnect() p1b.disconnect() self.flush() - self.assertEqual(len(self._usage), 3, self._usage) - self.assertEqual(self._usage[2]["mood"], "happy") + self.assertEqual(len(self._usage.events), 3, self._usage) + self.assertEqual(self._usage.events[2]["mood"], "happy")