open.source/magic-wormhole-transit-relay
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

use 'backends' for usage-recording

Browse Source
This commit is contained in:
meejah 2021-02-12 01:32:52 -07:00
parent 5ed572187b
commit 53864f57f0
2 changed files with 124 additions and 25 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

104
src/wormhole_transit_relay/server_state.py
View File

@ -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):

45
src/wormhole_transit_relay/test/test_transit_server.py
View File

@ -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")