diff --git a/docs/server-statemachine.dot b/docs/server-statemachine.dot new file mode 100644 index 0000000..d3a2215 --- /dev/null +++ b/docs/server-statemachine.dot @@ -0,0 +1,22 @@ +/** +. thinking about state-machine from "hand-drawn" perspective +. will it look the same as an Automat one? +**/ + +digraph { + listening -> wait_relay [label="connection_made"] + + wait_relay -> wait_partner [label="please_relay\nFindPartner"] + wait_relay -> wait_partner [label="please_relay_for_side\nFindPartner"] + wait_relay -> done [label="invalid_token\nSend('bad handshake')\nDisconnect"] + wait_relay -> done [label="connection_lost"] + + wait_partner -> relaying [label="got_partner\nConnectPartner(partner)\nSend('ok')"] + wait_partner -> done [label="got_bytes\nDisconnect"] + wait_partner -> done [label="connection_lost"] + + relaying -> relaying [label="got_bytes\nSend(bytes)"] + relaying -> done [label="partner_connection_lost\nDisconnectMe"] + relaying -> done [label="connection_lost\nDisconnectPartner"] +} + diff --git a/src/wormhole_transit_relay/server_state.py b/src/wormhole_transit_relay/server_state.py new file mode 100644 index 0000000..a773f87 --- /dev/null +++ b/src/wormhole_transit_relay/server_state.py @@ -0,0 +1,688 @@ +import time +import json +from collections import defaultdict + +import automat +from zope.interface import ( + Interface, + Attribute, + implementer, +) +from twisted.python import log + + +class ITransitClient(Interface): + + started_time = Attribute("timestamp when the connection was established") + + def send(data): + """ + Send some byets to the client + """ + + def disconnect(): + """ + Disconnect the client transport + """ + + def connect_partner(other): + """ + Hook up to our partner. + :param ITransitClient other: our partner + """ + + def disconnect_partner(): + """ + Disconnect our partner's transport + """ + + +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) + "\n") + self._file.flush() + + +@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. + """ + self._db.execute( + "INSERT INTO `usage`" + " (`started`, `total_time`, `waiting_time`," + " `total_bytes`, `result`)" + " VALUES (?,?,?,?,?)", + (started, total_time, waiting_time, total_bytes, mood) + ) + # original code did "self._update_stats()" here, thus causing + # "global" stats update on every connection update .. should + # we repeat this behavior, or really only record every + # 60-seconds with the timer? + self._db.commit() + + +def round_to(size, coarseness): + return int(coarseness*(1+int((size-1)/coarseness))) + + +def blur_size(size): + if size == 0: + return 0 + if size < 1e6: + return round_to(size, 10e3) + if size < 1e9: + return round_to(size, 1e6) + return round_to(size, 100e6) + + +def create_usage_tracker(blur_usage, log_file, usage_db): + """ + :param int blur_usage: see UsageTracker + + :param log_file: None or a file-like object to write JSON-encoded + lines of usage information to. + + :param usage_db: None or an sqlite3 database connection + + :returns: a new UsageTracker instance configured with backends. + """ + tracker = UsageTracker(blur_usage) + if usage_db: + tracker.add_backend(DatabaseUsageRecorder(usage_db)) + if log_file: + tracker.add_backend(LogFileUsageRecorder(log_file)) + return tracker + + +class UsageTracker(object): + """ + Tracks usage statistics of connections + """ + + def __init__(self, blur_usage): + """ + :param int blur_usage: None or the number of seconds to use as a + window around which to blur time statistics (e.g. "60" means times + will be rounded to 1 minute intervals). When blur_usage is + non-zero, sizes will also be rounded into buckets of "one + megabyte", "one gigabyte" or "lots" + """ + self._backends = set() + self._blur_usage = blur_usage + if blur_usage: + log.msg("blurring access times to %d seconds" % self._blur_usage) +## XXX log.msg("not logging Transit connections to Twisted log") + else: + log.msg("not blurring access times") + + def add_backend(self, backend): + """ + Add a new backend. + + :param IUsageWriter backend: the backend to add + """ + self._backends.add(backend) + + def record(self, started, buddy_started, result, bytes_sent, buddy_bytes): + """ + :param int started: timestamp when our connection started + + :param int buddy_started: None, or the timestamp when our + partner's connection started (will be None if we don't yet + have a partner). + + :param str result: a label for the result of the connection + (one of the "moods"). + + :param int bytes_sent: number of bytes we sent + + :param int buddy_bytes: number of bytes our partner sent + """ + # ideally self._reactor.seconds() or similar, but .. + finished = time.time() + if buddy_started is not None: + starts = [started, buddy_started] + total_time = finished - min(starts) + waiting_time = max(starts) - min(starts) + total_bytes = bytes_sent + buddy_bytes + else: + total_time = finished - started + waiting_time = None + total_bytes = bytes_sent + + if self._blur_usage: + started = self._blur_usage * (started // self._blur_usage) + total_bytes = blur_size(total_bytes) + + # This is "a dict" instead of "kwargs" because we have to make + # it into a dict for the log use-case and in-memory/testing + # use-case anyway so this is less repeats of the names. + self._notify_backends({ + "started": started, + "total_time": total_time, + "waiting_time": waiting_time, + "total_bytes": total_bytes, + "mood": result, + }) + + def update_stats(self, rebooted, updated, connected, waiting, + incomplete_bytes): + """ + Update general statistics. + """ + # in original code, this is only recorded in the database + # .. perhaps a better way to do this, but .. + for backend in self._backends: + if isinstance(backend, DatabaseUsageRecorder): + backend._db.execute("DELETE FROM `current`") + backend._db.execute( + "INSERT INTO `current`" + " (`rebooted`, `updated`, `connected`, `waiting`," + " `incomplete_bytes`)" + " VALUES (?, ?, ?, ?, ?)", + (int(rebooted), int(updated), connected, waiting, + incomplete_bytes) + ) + + + def _notify_backends(self, data): + """ + Internal helper. Tell every backend we have about a new usage record. + """ + for backend in self._backends: + backend.record_usage(**data) + + +class ActiveConnections(object): + """ + Tracks active connections. + + A connection is 'active' when both sides have shown up and they + are glued together (and thus could be passing data back and forth + if any is flowing). + """ + def __init__(self): + self._connections = set() + + def register(self, side0, side1): + """ + A connection has become active so register both its sides + + :param TransitConnection side0: one side of the connection + :param TransitConnection side1: one side of the connection + """ + self._connections.add(side0) + self._connections.add(side1) + + def unregister(self, side): + """ + One side of a connection has become inactive. + + :param TransitConnection side: an inactive side of a connection + """ + self._connections.discard(side) + + +class PendingRequests(object): + """ + Tracks outstanding (non-"active") requests. + + We register client connections against the tokens we have + received. When the other side shows up we can thus match it to the + correct partner connection. At this point, the connection becomes + "active" is and is thus no longer "pending" and so will no longer + be in this collection. + """ + + def __init__(self, active_connections): + """ + :param active_connections: an instance of ActiveConnections where + connections are put when both sides arrive. + """ + self._requests = defaultdict(set) # token -> set((side, TransitConnection)) + self._active = active_connections + + def unregister(self, token, side, tc): + """ + We no longer care about a particular client (e.g. it has + disconnected). + """ + if token in self._requests: + self._requests[token].discard((side, tc)) + if not self._requests[token]: + # no more sides; token is dead + del self._requests[token] + self._active.unregister(tc) + + def register(self, token, new_side, new_tc): + """ + A client has connected and successfully offered a token (and + optional 'side' token). If this is the first one for this + token, we merely remember it. If it is the second side for + this token we connect them together. + + :param bytes token: the token for this connection. + + :param bytes new_side: None or the side token for this connection + + :param TransitServerState new_tc: the state-machine of the connection + + :returns bool: True if we are the first side to register this + token + """ + potentials = self._requests[token] + for old in potentials: + (old_side, old_tc) = old + if ((old_side is None) + or (new_side is None) + or (old_side != new_side)): + # we found a match + # FIXME: debug-log this + # print("transit relay 2: %s" % new_tc.get_token()) + + # drop and stop tracking the rest + potentials.remove(old) + for (_, leftover_tc) in potentials.copy(): + # Don't record this as errory. It's just a spare connection + # from the same side as a connection that got used. This + # can happen if the connection hint contains multiple + # addresses (we don't currently support those, but it'd + # probably be useful in the future). + ##leftover_tc.disconnect_redundant() + leftover_tc.partner_connection_lost() + self._requests.pop(token, None) + + # glue the two ends together + self._active.register(new_tc, old_tc) + new_tc.got_partner(old_tc) + old_tc.got_partner(new_tc) + return False + + # FIXME: debug-log this + # print("transit relay 1: %s" % new_tc.get_token()) + potentials.add((new_side, new_tc)) + return True + # TODO: timer + + +class TransitServerState(object): + """ + Encapsulates the state-machine of the server side of a transit + relay connection. + + Once the protocol has been told to relay (or to relay for a side) + it starts passing all received bytes to the other side until it + closes. + """ + + _machine = automat.MethodicalMachine() + _client = None + _buddy = None + _token = None + _side = None + _first = None + _mood = "empty" + _total_sent = 0 + + def __init__(self, pending_requests, usage_recorder): + self._pending_requests = pending_requests + self._usage = usage_recorder + + def get_token(self): + """ + :returns str: a string describing our token. This will be "-" if + we have no token yet, or "{16 chars}-" if we have + just a token or "{16 chars}-{16 chars}" if we have a token and + a side. + """ + d = "-" + if self._token is not None: + d = self._token[:16].decode("ascii") + if self._side is not None: + d += "-" + self._side.decode("ascii") + else: + d += "-" + return d + + @_machine.input() + def connection_made(self, client): + """ + A client has connected. May only be called once. + + :param ITransitClient client: our client. + """ + # NB: the "only called once" is enforced by the state-machine; + # this input is only valid for the "listening" state, to which + # we never return. + + @_machine.input() + def please_relay(self, token): + """ + A 'please relay X' message has been received (the original version + of the protocol). + """ + + @_machine.input() + def please_relay_for_side(self, token, side): + """ + A 'please relay X for side Y' message has been received (the + second version of the protocol). + """ + + @_machine.input() + def bad_token(self): + """ + A bad token / relay line was received (e.g. couldn't be parsed) + """ + + @_machine.input() + def got_partner(self, client): + """ + The partner for this relay session has been found + """ + + @_machine.input() + def connection_lost(self): + """ + Our transport has failed. + """ + + @_machine.input() + def partner_connection_lost(self): + """ + Our partner's transport has failed. + """ + + @_machine.input() + def got_bytes(self, data): + """ + Some bytes have arrived (that aren't part of the handshake) + """ + + @_machine.output() + def _remember_client(self, client): + self._client = client + + @_machine.output() + def _register_token(self, token): + return self._real_register_token_for_side(token, None) + + @_machine.output() + def _register_token_for_side(self, token, side): + return self._real_register_token_for_side(token, side) + + @_machine.output() + def _unregister(self): + """ + remove us from the thing that remembers tokens and sides + """ + return self._pending_requests.unregister(self._token, self._side, self) + + @_machine.output() + def _send_bad(self): + self._mood = "errory" + self._client.send(b"bad handshake\n") + + @_machine.output() + def _send_ok(self): + self._client.send(b"ok\n") + + @_machine.output() + def _send_impatient(self): + self._client.send(b"impatient\n") + + @_machine.output() + def _count_bytes(self, data): + self._total_sent += len(data) + + @_machine.output() + def _send(self, data): + self._client.send(data) + + @_machine.output() + def _send_to_partner(self, data): + self._buddy._client.send(data) + + @_machine.output() + def _connect_partner(self, client): + self._buddy = client + self._client.connect_partner(client) + + @_machine.output() + def _disconnect(self): + self._client.disconnect() + + @_machine.output() + def _disconnect_partner(self): + self._client.disconnect_partner() + + # some outputs to record "usage" information .. + @_machine.output() + def _record_usage(self): + if self._mood == "jilted": + if self._buddy: + if self._buddy._mood == "happy": + return + self._usage.record( + started=self._client.started_time, + buddy_started=self._buddy._client.started_time if self._buddy is not None else None, + result=self._mood, + bytes_sent=self._total_sent, + buddy_bytes=self._buddy._total_sent if self._buddy is not None else None + ) + + # some outputs to record the "mood" .. + @_machine.output() + def _mood_happy(self): + self._mood = "happy" + + @_machine.output() + def _mood_lonely(self): + self._mood = "lonely" + + @_machine.output() + def _mood_redundant(self): + self._mood = "redundant" + + @_machine.output() + def _mood_impatient(self): + self._mood = "impatient" + + @_machine.output() + def _mood_errory(self): + self._mood = "errory" + + @_machine.output() + def _mood_happy_if_first(self): + """ + We disconnected first so we're only happy if we also connected + first. + """ + if self._first: + self._mood = "happy" + else: + self._mood = "jilted" + + def _real_register_token_for_side(self, token, side): + """ + A client has connected and sent a valid version 1 or version 2 + handshake. If the former, `side` will be None. + + In either case, we remember the tokens and register + ourselves. This might result in 'got_partner' notifications to + two state-machines if this is the second side for a given token. + + :param bytes token: the token + :param bytes side: The side token (or None) + """ + self._token = token + self._side = side + self._first = self._pending_requests.register(token, side, self) + + @_machine.state(initial=True) + def listening(self): + """ + Initial state, awaiting connection. + """ + + @_machine.state() + def wait_relay(self): + """ + Waiting for a 'relay' message + """ + + @_machine.state() + def wait_partner(self): + """ + Waiting for our partner to connect + """ + + @_machine.state() + def relaying(self): + """ + Relaying bytes to our partner + """ + + @_machine.state() + def done(self): + """ + Terminal state + """ + + listening.upon( + connection_made, + enter=wait_relay, + outputs=[_remember_client], + ) + + wait_relay.upon( + please_relay, + enter=wait_partner, + outputs=[_mood_lonely, _register_token], + ) + wait_relay.upon( + please_relay_for_side, + enter=wait_partner, + outputs=[_mood_lonely, _register_token_for_side], + ) + wait_relay.upon( + bad_token, + enter=done, + outputs=[_mood_errory, _send_bad, _disconnect, _record_usage], + ) + wait_relay.upon( + got_bytes, + enter=done, + outputs=[_count_bytes, _mood_errory, _disconnect, _record_usage], + ) + wait_relay.upon( + connection_lost, + enter=done, + outputs=[_disconnect, _record_usage], + ) + + wait_partner.upon( + got_partner, + enter=relaying, + outputs=[_mood_happy, _send_ok, _connect_partner], + ) + wait_partner.upon( + connection_lost, + enter=done, + outputs=[_mood_lonely, _unregister, _record_usage], + ) + wait_partner.upon( + got_bytes, + enter=done, + outputs=[_mood_impatient, _send_impatient, _disconnect, _unregister, _record_usage], + ) + wait_partner.upon( + partner_connection_lost, + enter=done, + outputs=[_mood_redundant, _disconnect, _record_usage], + ) + + relaying.upon( + got_bytes, + enter=relaying, + outputs=[_count_bytes, _send_to_partner], + ) + relaying.upon( + connection_lost, + enter=done, + outputs=[_mood_happy_if_first, _disconnect_partner, _unregister, _record_usage], + ) + + done.upon( + connection_lost, + enter=done, + outputs=[], + ) + done.upon( + partner_connection_lost, + enter=done, + outputs=[], + ) diff --git a/src/wormhole_transit_relay/server_tap.py b/src/wormhole_transit_relay/server_tap.py index 8fbfde2..b4028d2 100644 --- a/src/wormhole_transit_relay/server_tap.py +++ b/src/wormhole_transit_relay/server_tap.py @@ -6,7 +6,9 @@ StreamServerEndpointService) from twisted.internet import endpoints from . import transit_server +from .server_state import create_usage_tracker from .increase_rlimits import increase_rlimits +from .database import get_db LONGDESC = """\ This plugin sets up a 'Transit Relay' server for magic-wormhole. This service @@ -32,13 +34,19 @@ def opt_blur_usage(self, arg): def makeService(config, reactor=reactor): increase_rlimits() ep = endpoints.serverFromString(reactor, config["port"]) # to listen - log_file = (os.fdopen(int(config["log-fd"]), "w") - if config["log-fd"] is not None - else None) - f = transit_server.Transit(blur_usage=config["blur-usage"], - log_file=log_file, - usage_db=config["usage-db"]) + log_file = ( + os.fdopen(int(config["log-fd"]), "w") + if config["log-fd"] is not None + else None + ) + db = None if config["usage-db"] is None else get_db(config["usage-db"]) + usage = create_usage_tracker( + blur_usage=config["blur-usage"], + log_file=log_file, + usage_db=db, + ) + factory = transit_server.Transit(usage, reactor.seconds) parent = MultiService() - StreamServerEndpointService(ep, f).setServiceParent(parent) - TimerService(5*60.0, f.timerUpdateStats).setServiceParent(parent) + StreamServerEndpointService(ep, factory).setServiceParent(parent) + TimerService(5*60.0, factory.update_stats).setServiceParent(parent) return parent diff --git a/src/wormhole_transit_relay/test/common.py b/src/wormhole_transit_relay/test/common.py index 53958fb..834caec 100644 --- a/src/wormhole_transit_relay/test/common.py +++ b/src/wormhole_transit_relay/test/common.py @@ -1,21 +1,22 @@ from twisted.test import proto_helpers from ..transit_server import Transit +from ..server_state import create_usage_tracker class ServerBase: log_requests = False def setUp(self): self._lp = None - if self.log_requests: - blur_usage = None - else: - blur_usage = 60.0 - self._setup_relay(blur_usage=blur_usage) + self._setup_relay(blur_usage=60.0 if self.log_requests else None) self._transit_server._debug_log = self.log_requests def _setup_relay(self, blur_usage=None, log_file=None, usage_db=None): - self._transit_server = Transit(blur_usage=blur_usage, - log_file=log_file, usage_db=usage_db) + usage = create_usage_tracker( + blur_usage=blur_usage, + log_file=log_file, + usage_db=usage_db, + ) + self._transit_server = Transit(usage, lambda: 123456789.0) def new_protocol(self): protocol = self._transit_server.buildProtocol(('127.0.0.1', 0)) diff --git a/src/wormhole_transit_relay/test/test_rlimits.py b/src/wormhole_transit_relay/test/test_rlimits.py index 10497e4..3ee23a9 100644 --- a/src/wormhole_transit_relay/test/test_rlimits.py +++ b/src/wormhole_transit_relay/test/test_rlimits.py @@ -1,5 +1,5 @@ from __future__ import print_function, unicode_literals -import mock +from unittest import mock from twisted.trial import unittest from ..increase_rlimits import increase_rlimits diff --git a/src/wormhole_transit_relay/test/test_service.py b/src/wormhole_transit_relay/test/test_service.py index dac642c..197d376 100644 --- a/src/wormhole_transit_relay/test/test_service.py +++ b/src/wormhole_transit_relay/test/test_service.py @@ -1,6 +1,6 @@ from __future__ import unicode_literals, print_function from twisted.trial import unittest -import mock +from unittest import mock from twisted.application.service import MultiService from .. import server_tap @@ -8,7 +8,7 @@ class Service(unittest.TestCase): def test_defaults(self): o = server_tap.Options() o.parseOptions([]) - with mock.patch("wormhole_transit_relay.server_tap.transit_server.Transit") as t: + with mock.patch("wormhole_transit_relay.server_tap.create_usage_tracker") as t: s = server_tap.makeService(o) self.assertEqual(t.mock_calls, [mock.call(blur_usage=None, @@ -18,7 +18,7 @@ def test_defaults(self): def test_blur(self): o = server_tap.Options() o.parseOptions(["--blur-usage=60"]) - with mock.patch("wormhole_transit_relay.server_tap.transit_server.Transit") as t: + with mock.patch("wormhole_transit_relay.server_tap.create_usage_tracker") as t: server_tap.makeService(o) self.assertEqual(t.mock_calls, [mock.call(blur_usage=60, @@ -28,7 +28,7 @@ def test_log_fd(self): o = server_tap.Options() o.parseOptions(["--log-fd=99"]) fd = object() - with mock.patch("wormhole_transit_relay.server_tap.transit_server.Transit") as t: + with mock.patch("wormhole_transit_relay.server_tap.create_usage_tracker") as t: with mock.patch("wormhole_transit_relay.server_tap.os.fdopen", return_value=fd) as f: server_tap.makeService(o) diff --git a/src/wormhole_transit_relay/test/test_stats.py b/src/wormhole_transit_relay/test/test_stats.py index f9433ef..5dca530 100644 --- a/src/wormhole_transit_relay/test/test_stats.py +++ b/src/wormhole_transit_relay/test/test_stats.py @@ -1,28 +1,39 @@ from __future__ import print_function, unicode_literals -import os, io, json, sqlite3 -import mock +import os, io, json +from unittest import mock from twisted.trial import unittest from ..transit_server import Transit +from ..server_state import create_usage_tracker from .. import database class DB(unittest.TestCase): - def open_db(self, dbfile): - db = sqlite3.connect(dbfile) - database._initialize_db_connection(db) - return db def test_db(self): + T = 1519075308.0 + + class Timer: + t = T + def __call__(self): + return self.t + get_time = Timer() + d = self.mktemp() os.mkdir(d) usage_db = os.path.join(d, "usage.sqlite") - with mock.patch("time.time", return_value=T+0): - t = Transit(blur_usage=None, log_file=None, usage_db=usage_db) - db = self.open_db(usage_db) + db = database.get_db(usage_db) + t = Transit( + create_usage_tracker(blur_usage=None, log_file=None, usage_db=db), + get_time, + ) + self.assertEqual(len(t.usage._backends), 1) + usage = list(t.usage._backends)[0] + + get_time.t = T + 1 + usage.record_usage(started=123, mood="happy", total_bytes=100, + total_time=10, waiting_time=2) + t.update_stats() - with mock.patch("time.time", return_value=T+1): - t.recordUsage(started=123, result="happy", total_bytes=100, - total_time=10, waiting_time=2) self.assertEqual(db.execute("SELECT * FROM `usage`").fetchall(), [dict(result="happy", started=123, total_bytes=100, total_time=10, waiting_time=2), @@ -32,9 +43,10 @@ def test_db(self): incomplete_bytes=0, waiting=0, connected=0)) - with mock.patch("time.time", return_value=T+2): - t.recordUsage(started=150, result="errory", total_bytes=200, - total_time=11, waiting_time=3) + get_time.t = T + 2 + usage.record_usage(started=150, mood="errory", total_bytes=200, + total_time=11, waiting_time=3) + t.update_stats() self.assertEqual(db.execute("SELECT * FROM `usage`").fetchall(), [dict(result="happy", started=123, total_bytes=100, total_time=10, waiting_time=2), @@ -46,27 +58,37 @@ def test_db(self): incomplete_bytes=0, waiting=0, connected=0)) - with mock.patch("time.time", return_value=T+3): - t.timerUpdateStats() + get_time.t = T + 3 + t.update_stats() self.assertEqual(db.execute("SELECT * FROM `current`").fetchone(), dict(rebooted=T+0, updated=T+3, incomplete_bytes=0, waiting=0, connected=0)) def test_no_db(self): - t = Transit(blur_usage=None, log_file=None, usage_db=None) + t = Transit( + create_usage_tracker(blur_usage=None, log_file=None, usage_db=None), + lambda: 0, + ) + self.assertEqual(0, len(t.usage._backends)) - t.recordUsage(started=123, result="happy", total_bytes=100, - total_time=10, waiting_time=2) - t.timerUpdateStats() class LogToStdout(unittest.TestCase): def test_log(self): # emit lines of JSON to log_file, if set log_file = io.StringIO() - t = Transit(blur_usage=None, log_file=log_file, usage_db=None) - t.recordUsage(started=123, result="happy", total_bytes=100, - total_time=10, waiting_time=2) + t = Transit( + create_usage_tracker(blur_usage=None, log_file=log_file, usage_db=None), + lambda: 0, + ) + with mock.patch("time.time", return_value=133): + t.usage.record( + started=123, + buddy_started=125, + result="happy", + bytes_sent=100, + buddy_bytes=0, + ) self.assertEqual(json.loads(log_file.getvalue()), {"started": 123, "total_time": 10, "waiting_time": 2, "total_bytes": 100, @@ -76,15 +98,34 @@ def test_log_blurred(self): # if blurring is enabled, timestamps should be rounded to the # requested amount, and sizes should be rounded up too log_file = io.StringIO() - t = Transit(blur_usage=60, log_file=log_file, usage_db=None) - t.recordUsage(started=123, result="happy", total_bytes=11999, - total_time=10, waiting_time=2) + t = Transit( + create_usage_tracker(blur_usage=60, log_file=log_file, usage_db=None), + lambda: 0, + ) + + with mock.patch("time.time", return_value=123 + 10): + t.usage.record( + started=123, + buddy_started=125, + result="happy", + bytes_sent=11999, + buddy_bytes=0, + ) + print(log_file.getvalue()) self.assertEqual(json.loads(log_file.getvalue()), {"started": 120, "total_time": 10, "waiting_time": 2, "total_bytes": 20000, "mood": "happy"}) def test_do_not_log(self): - t = Transit(blur_usage=60, log_file=None, usage_db=None) - t.recordUsage(started=123, result="happy", total_bytes=11999, - total_time=10, waiting_time=2) + t = Transit( + create_usage_tracker(blur_usage=60, log_file=None, usage_db=None), + lambda: 0, + ) + t.usage.record( + started=123, + buddy_started=124, + result="happy", + bytes_sent=11999, + buddy_bytes=12, + ) diff --git a/src/wormhole_transit_relay/test/test_transit_server.py b/src/wormhole_transit_relay/test/test_transit_server.py index a4763d9..70f49a8 100644 --- a/src/wormhole_transit_relay/test/test_transit_server.py +++ b/src/wormhole_transit_relay/test/test_transit_server.py @@ -2,7 +2,10 @@ from binascii import hexlify from twisted.trial import unittest from .common import ServerBase -from .. import transit_server +from ..server_state import ( + MemoryUsageRecorder, + blur_size, +) def handshake(token, side=None): hs = b"please relay " + hexlify(token) @@ -13,27 +16,28 @@ def handshake(token, side=None): class _Transit: def count(self): - return sum([len(potentials) - for potentials - in self._transit_server._pending_requests.values()]) + return sum([ + len(potentials) + for potentials + in self._transit_server.pending_requests._requests.values() + ]) def test_blur_size(self): - blur = transit_server.blur_size - self.failUnlessEqual(blur(0), 0) - self.failUnlessEqual(blur(1), 10e3) - self.failUnlessEqual(blur(10e3), 10e3) - self.failUnlessEqual(blur(10e3+1), 20e3) - self.failUnlessEqual(blur(15e3), 20e3) - self.failUnlessEqual(blur(20e3), 20e3) - self.failUnlessEqual(blur(1e6), 1e6) - self.failUnlessEqual(blur(1e6+1), 2e6) - self.failUnlessEqual(blur(1.5e6), 2e6) - self.failUnlessEqual(blur(2e6), 2e6) - self.failUnlessEqual(blur(900e6), 900e6) - self.failUnlessEqual(blur(1000e6), 1000e6) - self.failUnlessEqual(blur(1050e6), 1100e6) - self.failUnlessEqual(blur(1100e6), 1100e6) - self.failUnlessEqual(blur(1150e6), 1200e6) + self.failUnlessEqual(blur_size(0), 0) + self.failUnlessEqual(blur_size(1), 10e3) + self.failUnlessEqual(blur_size(10e3), 10e3) + self.failUnlessEqual(blur_size(10e3+1), 20e3) + self.failUnlessEqual(blur_size(15e3), 20e3) + self.failUnlessEqual(blur_size(20e3), 20e3) + self.failUnlessEqual(blur_size(1e6), 1e6) + self.failUnlessEqual(blur_size(1e6+1), 2e6) + self.failUnlessEqual(blur_size(1.5e6), 2e6) + self.failUnlessEqual(blur_size(2e6), 2e6) + self.failUnlessEqual(blur_size(900e6), 900e6) + self.failUnlessEqual(blur_size(1000e6), 1000e6) + self.failUnlessEqual(blur_size(1050e6), 1100e6) + self.failUnlessEqual(blur_size(1100e6), 1100e6) + self.failUnlessEqual(blur_size(1150e6), 1200e6) def test_register(self): p1 = self.new_protocol() @@ -48,7 +52,7 @@ def test_register(self): self.assertEqual(self.count(), 0) # the token should be removed too - self.assertEqual(len(self._transit_server._pending_requests), 0) + self.assertEqual(len(self._transit_server.pending_requests._requests), 0) def test_both_unsided(self): p1 = self.new_protocol() @@ -168,8 +172,8 @@ def test_ignore_same_side(self): side2 = b"\x02"*8 p3.dataReceived(handshake(token1, side=side2)) self.assertEqual(self.count(), 0) - self.assertEqual(len(self._transit_server._pending_requests), 0) - self.assertEqual(len(self._transit_server._active_connections), 2) + self.assertEqual(len(self._transit_server.pending_requests._requests), 0) + self.assertEqual(len(self._transit_server.active_connections._connections), 2) # That will trigger a disconnect on exactly one of (p1 or p2). # The other connection should still be connected self.assertEqual(sum([int(t.transport.connected) for t in [p1, p2]]), 1) @@ -301,20 +305,21 @@ def test_empty_handshake(self): # hang up before sending anything p1.transport.loseConnection() + class TransitWithLogs(_Transit, ServerBase, unittest.TestCase): log_requests = True + class TransitWithoutLogs(_Transit, ServerBase, unittest.TestCase): log_requests = False + class Usage(ServerBase, unittest.TestCase): + def setUp(self): super(Usage, self).setUp() - self._usage = [] - def record(started, result, total_bytes, total_time, waiting_time): - self._usage.append((started, result, total_bytes, - total_time, waiting_time)) - self._transit_server.recordUsage = record + self._usage = MemoryUsageRecorder() + self._transit_server.usage.add_backend(self._usage) def test_empty(self): p1 = self.new_protocol() @@ -322,9 +327,8 @@ def test_empty(self): p1.transport.loseConnection() # that will log the "empty" usage event - self.assertEqual(len(self._usage), 1, self._usage) - (started, result, total_bytes, total_time, waiting_time) = self._usage[0] - self.assertEqual(result, "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() @@ -333,9 +337,8 @@ def test_short(self): p1.transport.loseConnection() # that will log the "empty" usage event - self.assertEqual(len(self._usage), 1, self._usage) - (started, result, total_bytes, total_time, waiting_time) = self._usage[0] - self.assertEqual(result, "empty", self._usage) + 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() @@ -343,9 +346,8 @@ def test_errory(self): p1.dataReceived(b"this is a very bad handshake\n") # that will log the "errory" usage event, then drop the connection p1.transport.loseConnection() - self.assertEqual(len(self._usage), 1, self._usage) - (started, result, total_bytes, total_time, waiting_time) = self._usage[0] - self.assertEqual(result, "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() @@ -356,10 +358,9 @@ def test_lonely(self): # now we disconnect before the peer connects p1.transport.loseConnection() - self.assertEqual(len(self._usage), 1, self._usage) - (started, result, total_bytes, total_time, waiting_time) = self._usage[0] - self.assertEqual(result, "lonely", self._usage) - self.assertIdentical(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() @@ -371,18 +372,17 @@ def test_one_happy_one_jilted(self): p1.dataReceived(handshake(token1, side=side1)) p2.dataReceived(handshake(token1, side=side2)) - self.assertEqual(self._usage, []) # no events yet + self.assertEqual(self._usage.events, []) # no events yet p1.dataReceived(b"\x00" * 13) p2.dataReceived(b"\xff" * 7) p1.transport.loseConnection() - self.assertEqual(len(self._usage), 1, self._usage) - (started, result, total_bytes, total_time, waiting_time) = self._usage[0] - self.assertEqual(result, "happy", self._usage) - self.assertEqual(total_bytes, 20) - self.assertNotIdentical(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() @@ -402,19 +402,16 @@ def test_redundant(self): p1c.dataReceived(handshake(token1, side=side1)) p1c.transport.loseConnection() - self.assertEqual(len(self._usage), 1, self._usage) - (started, result, total_bytes, total_time, waiting_time) = self._usage[0] - self.assertEqual(result, "lonely", self._usage) + self.assertEqual(len(self._usage.events), 1, self._usage) + self.assertEqual(self._usage.events[0]["mood"], "lonely") p2.dataReceived(handshake(token1, side=side2)) - self.assertEqual(len(self._transit_server._pending_requests), 0) - self.assertEqual(len(self._usage), 2, self._usage) - (started, result, total_bytes, total_time, waiting_time) = self._usage[1] - self.assertEqual(result, "redundant", self._usage) + self.assertEqual(len(self._transit_server.pending_requests._requests), 0) + 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.transport.loseConnection() p1b.transport.loseConnection() - self.assertEqual(len(self._usage), 3, self._usage) - (started, result, total_bytes, total_time, waiting_time) = self._usage[2] - self.assertEqual(result, "happy", self._usage) + self.assertEqual(len(self._usage.events), 3, self._usage) + self.assertEqual(self._usage.events[2]["mood"], "happy") diff --git a/src/wormhole_transit_relay/transit_server.py b/src/wormhole_transit_relay/transit_server.py index 91d84e0..d9dabfd 100644 --- a/src/wormhole_transit_relay/transit_server.py +++ b/src/wormhole_transit_relay/transit_server.py @@ -1,10 +1,9 @@ from __future__ import print_function, unicode_literals -import re, time, json -from collections import defaultdict +import re +import time from twisted.python import log from twisted.internet import protocol from twisted.protocols.basic import LineReceiver -from .database import get_db SECONDS = 1.0 MINUTE = 60*SECONDS @@ -12,343 +11,185 @@ DAY = 24*HOUR MB = 1000*1000 -def round_to(size, coarseness): - return int(coarseness*(1+int((size-1)/coarseness))) -def blur_size(size): - if size == 0: - return 0 - if size < 1e6: - return round_to(size, 10e3) - if size < 1e9: - return round_to(size, 1e6) - return round_to(size, 100e6) +from wormhole_transit_relay.server_state import ( + TransitServerState, + PendingRequests, + ActiveConnections, + ITransitClient, +) +from zope.interface import implementer + +@implementer(ITransitClient) class TransitConnection(LineReceiver): delimiter = b'\n' # maximum length of a line we will accept before the handshake is complete. # This must be >= to the longest possible handshake message. MAX_LENGTH = 1024 + started_time = None + + def send(self, data): + """ + ITransitClient API + """ + self.transport.write(data) + + def disconnect(self): + """ + ITransitClient API + """ + self.transport.loseConnection() - def __init__(self): - self._got_token = False - self._got_side = False - self._sent_ok = False - self._mood = "empty" - self._buddy = None - self._total_sent = 0 - - def describeToken(self): - d = "-" - if self._got_token: - d = self._got_token[:16].decode("ascii") - if self._got_side: - d += "-" + self._got_side.decode("ascii") - else: - d += "-" - return d + def connect_partner(self, other): + """ + ITransitClient API + """ + self._buddy = other + + def disconnect_partner(self): + """ + ITransitClient API + """ + if self._buddy is not None: + log.msg("buddy_disconnected {}".format(self._buddy.get_token())) + self._buddy._client.transport.loseConnection() + self._buddy = None def connectionMade(self): - self._started = time.time() - self._log_requests = self.factory._log_requests + # ideally more like self._reactor.seconds() ... but Twisted + # doesn't have a good way to get the reactor for a protocol + # (besides "use the global one") + self.started_time = time.time() + self._state = TransitServerState( + self.factory.pending_requests, + self.factory.usage, + ) + self._state.connection_made(self) +## self._log_requests = self.factory._log_requests try: self.transport.setTcpKeepAlive(True) except AttributeError: pass def lineReceived(self, line): + """ + LineReceiver API + """ # old: "please relay {64}\n" + token = None old = re.search(br"^please relay (\w{64})$", line) if old: token = old.group(1) - return self._got_handshake(token, None) + self._state.please_relay(token) # new: "please relay {64} for side {16}\n" new = re.search(br"^please relay (\w{64}) for side (\w{16})$", line) if new: token = new.group(1) side = new.group(2) - return self._got_handshake(token, side) + self._state.please_relay_for_side(token, side) - self.sendLine(b"bad handshake") - if self._log_requests: - log.msg("transit handshake failure") - return self.disconnect_error() + if token is None: + self._state.bad_token() + else: + self.setRawMode() def rawDataReceived(self, data): + """ + LineReceiver API + """ # We are an IPushProducer to our buddy's IConsumer, so they'll # throttle us (by calling pauseProducing()) when their outbound # buffer is full (e.g. when their downstream pipe is full). In # practice, this buffers about 10MB per connection, after which # point the sender will only transmit data as fast as the # receiver can handle it. - if self._sent_ok: - if not self._buddy: - # Our buddy disconnected (we're "jilted"), so we hung up too, - # but our incoming data hasn't stopped yet (it will in a - # moment, after our disconnect makes a roundtrip through the - # kernel). This probably means the file receiver hung up, and - # this connection is the file sender. In may-2020 this - # happened 11 times in 40 days. - return - self._total_sent += len(data) - self._buddy.transport.write(data) - return - - # handshake is complete but not yet sent_ok - self.sendLine(b"impatient") - if self._log_requests: - log.msg("transit impatience failure") - return self.disconnect_error() # impatience yields failure - - def _got_handshake(self, token, side): - self._got_token = token - self._got_side = side - self._mood = "lonely" # until buddy connects - self.setRawMode() - self.factory.connection_got_token(token, side, self) - - def buddy_connected(self, them): - self._buddy = them - self._mood = "happy" - self.sendLine(b"ok") - self._sent_ok = True - # Connect the two as a producer/consumer pair. We use streaming=True, - # so this expects the IPushProducer interface, and uses - # pauseProducing() to throttle, and resumeProducing() to unthrottle. - self._buddy.transport.registerProducer(self.transport, True) - # The Transit object calls buddy_connected() on both protocols, so - # there will be two producer/consumer pairs. - - def buddy_disconnected(self): - if self._log_requests: - log.msg("buddy_disconnected %s" % self.describeToken()) - self._buddy = None - self._mood = "jilted" - self.transport.loseConnection() + self._state.got_bytes(data) def disconnect_error(self): # we haven't finished the handshake, so there are no tokens tracking # us - self._mood = "errory" self.transport.loseConnection() + # XXX probably should be logged by state? if self.factory._debug_log: log.msg("transitFailed %r" % self) def disconnect_redundant(self): # this is called if a buddy connected and we were found unnecessary. # Any token-tracking cleanup will have been done before we're called. - self._mood = "redundant" self.transport.loseConnection() def connectionLost(self, reason): - finished = time.time() - total_time = finished - self._started + self._state.connection_lost() +# XXX this probably resulted in a log message we've not refactored yet +# self.factory.transitFinished(self, self._got_token, self._got_side, +# self.describeToken()) +# XXX describeToken -> self._state.get_token() - # Record usage. There are eight cases: - # * n0: we haven't gotten a full handshake yet (empty) - # * n1: the handshake failed, not a real client (errory) - # * n2: real client disconnected before any buddy appeared (lonely) - # * n3: real client closed as redundant after buddy appears (redundant) - # * n4: real client connected first, buddy closes first (jilted) - # * n5: real client connected first, buddy close last (happy) - # * n6: real client connected last, buddy closes first (jilted) - # * n7: real client connected last, buddy closes last (happy) - - # * non-connected clients (0,1,2,3) always write a usage record - # * for connected clients, whoever disconnects first gets to write the - # usage record (5, 7). The last disconnect doesn't write a record. - - if self._mood == "empty": # 0 - assert not self._buddy - self.factory.recordUsage(self._started, "empty", 0, - total_time, None) - elif self._mood == "errory": # 1 - assert not self._buddy - self.factory.recordUsage(self._started, "errory", 0, - total_time, None) - elif self._mood == "redundant": # 3 - assert not self._buddy - self.factory.recordUsage(self._started, "redundant", 0, - total_time, None) - elif self._mood == "jilted": # 4 or 6 - # we were connected, but our buddy hung up on us. They record the - # usage event, we do not - pass - elif self._mood == "lonely": # 2 - assert not self._buddy - self.factory.recordUsage(self._started, "lonely", 0, - total_time, None) - else: # 5 or 7 - # we were connected, we hung up first. We record the event. - assert self._mood == "happy", self._mood - assert self._buddy - starts = [self._started, self._buddy._started] - total_time = finished - min(starts) - waiting_time = max(starts) - min(starts) - total_bytes = self._total_sent + self._buddy._total_sent - self.factory.recordUsage(self._started, "happy", total_bytes, - total_time, waiting_time) - if self._buddy: - self._buddy.buddy_disconnected() - self.factory.transitFinished(self, self._got_token, self._got_side, - self.describeToken()) class Transit(protocol.ServerFactory): - # I manage pairs of simultaneous connections to a secondary TCP port, - # both forwarded to the other. Clients must begin each connection with - # "please relay TOKEN for SIDE\n" (or a legacy form without the "for - # SIDE"). Two connections match if they use the same TOKEN and have - # different SIDEs (the redundant connections are dropped when a match is - # made). Legacy connections match any with the same TOKEN, ignoring SIDE - # (so two legacy connections will match each other). - - # I will send "ok\n" when the matching connection is established, or - # disconnect if no matching connection is made within MAX_WAIT_TIME - # seconds. I will disconnect if you send data before the "ok\n". All data - # you get after the "ok\n" will be from the other side. You will not - # receive "ok\n" until the other side has also connected and submitted a - # matching token (and differing SIDE). - - # In addition, the connections will be dropped after MAXLENGTH bytes have - # been sent by either side, or MAXTIME seconds have elapsed after the - # matching connections were established. A future API will reveal these - # limits to clients instead of causing mysterious spontaneous failures. - - # These relay connections are not half-closeable (unlike full TCP - # connections, applications will not receive any data after half-closing - # their outgoing side). Applications must negotiate shutdown with their - # peer and not close the connection until all data has finished - # transferring in both directions. Applications which only need to send - # data in one direction can use close() as usual. - + """ + I manage pairs of simultaneous connections to a secondary TCP port, + both forwarded to the other. Clients must begin each connection with + "please relay TOKEN for SIDE\n" (or a legacy form without the "for + SIDE"). Two connections match if they use the same TOKEN and have + different SIDEs (the redundant connections are dropped when a match is + made). Legacy connections match any with the same TOKEN, ignoring SIDE + (so two legacy connections will match each other). + + I will send "ok\n" when the matching connection is established, or + disconnect if no matching connection is made within MAX_WAIT_TIME + seconds. I will disconnect if you send data before the "ok\n". All data + you get after the "ok\n" will be from the other side. You will not + receive "ok\n" until the other side has also connected and submitted a + matching token (and differing SIDE). + + In addition, the connections will be dropped after MAXLENGTH bytes have + been sent by either side, or MAXTIME seconds have elapsed after the + matching connections were established. A future API will reveal these + limits to clients instead of causing mysterious spontaneous failures. + + These relay connections are not half-closeable (unlike full TCP + connections, applications will not receive any data after half-closing + their outgoing side). Applications must negotiate shutdown with their + peer and not close the connection until all data has finished + transferring in both directions. Applications which only need to send + data in one direction can use close() as usual. + """ + + # TODO: unused MAX_WAIT_TIME = 30*SECONDS + # TODO: unused MAXLENGTH = 10*MB + # TODO: unused MAXTIME = 60*SECONDS protocol = TransitConnection - def __init__(self, blur_usage, log_file, usage_db): - self._blur_usage = blur_usage - self._log_requests = blur_usage is None - if self._blur_usage: - log.msg("blurring access times to %d seconds" % self._blur_usage) - log.msg("not logging Transit connections to Twisted log") - else: - log.msg("not blurring access times") + def __init__(self, usage, get_timestamp): + self.active_connections = ActiveConnections() + self.pending_requests = PendingRequests(self.active_connections) + self.usage = usage self._debug_log = False - self._log_file = log_file - self._db = None - if usage_db: - self._db = get_db(usage_db) - self._rebooted = time.time() - # we don't track TransitConnections until they submit a token - self._pending_requests = defaultdict(set) # token -> set((side, TransitConnection)) - self._active_connections = set() # TransitConnection - - def connection_got_token(self, token, new_side, new_tc): - potentials = self._pending_requests[token] - for old in potentials: - (old_side, old_tc) = old - if ((old_side is None) - or (new_side is None) - or (old_side != new_side)): - # we found a match - if self._debug_log: - log.msg("transit relay 2: %s" % new_tc.describeToken()) + self._timestamp = get_timestamp + self._rebooted = self._timestamp() - # drop and stop tracking the rest - potentials.remove(old) - for (_, leftover_tc) in potentials.copy(): - # Don't record this as errory. It's just a spare connection - # from the same side as a connection that got used. This - # can happen if the connection hint contains multiple - # addresses (we don't currently support those, but it'd - # probably be useful in the future). - leftover_tc.disconnect_redundant() - self._pending_requests.pop(token, None) - - # glue the two ends together - self._active_connections.add(new_tc) - self._active_connections.add(old_tc) - new_tc.buddy_connected(old_tc) - old_tc.buddy_connected(new_tc) - return - if self._debug_log: - log.msg("transit relay 1: %s" % new_tc.describeToken()) - potentials.add((new_side, new_tc)) - # TODO: timer - - def transitFinished(self, tc, token, side, description): - if token in self._pending_requests: - side_tc = (side, tc) - self._pending_requests[token].discard(side_tc) - if not self._pending_requests[token]: # set is now empty - del self._pending_requests[token] - if self._debug_log: - log.msg("transitFinished %s" % (description,)) - self._active_connections.discard(tc) - # we could update the usage database "current" row immediately, or wait - # until the 5-minute timer updates it. If we update it now, just after - # losing a connection, we should probably also update it just after - # establishing one (at the end of connection_got_token). For now I'm - # going to omit these, but maybe someday we'll turn them both on. The - # consequence is that a manual execution of the munin scripts ("munin - # run wormhole_transit_active") will give the wrong value just after a - # connect/disconnect event. Actual munin graphs should accurately - # report connections that last longer than the 5-minute sampling - # window, which is what we actually care about. - #self.timerUpdateStats() - - def recordUsage(self, started, result, total_bytes, - total_time, waiting_time): - if self._debug_log: - log.msg(format="Transit.recordUsage {bytes}B", bytes=total_bytes) - if self._blur_usage: - started = self._blur_usage * (started // self._blur_usage) - total_bytes = blur_size(total_bytes) - if self._log_file is not None: - data = {"started": started, - "total_time": total_time, - "waiting_time": waiting_time, - "total_bytes": total_bytes, - "mood": result, - } - self._log_file.write(json.dumps(data)+"\n") - self._log_file.flush() - if self._db: - self._db.execute("INSERT INTO `usage`" - " (`started`, `total_time`, `waiting_time`," - " `total_bytes`, `result`)" - " VALUES (?,?,?, ?,?)", - (started, total_time, waiting_time, - total_bytes, result)) - self._update_stats() - self._db.commit() - - def timerUpdateStats(self): - if self._db: - self._update_stats() - self._db.commit() - - def _update_stats(self): - # current status: should be zero when idle - rebooted = self._rebooted - updated = time.time() - connected = len(self._active_connections) / 2 + def update_stats(self): # TODO: when a connection is half-closed, len(active) will be odd. a # moment later (hopefully) the other side will disconnect, but # _update_stats isn't updated until later. - waiting = len(self._pending_requests) + # "waiting" doesn't count multiple parallel connections from the same # side - incomplete_bytes = sum(tc._total_sent - for tc in self._active_connections) - self._db.execute("DELETE FROM `current`") - self._db.execute("INSERT INTO `current`" - " (`rebooted`, `updated`, `connected`, `waiting`," - " `incomplete_bytes`)" - " VALUES (?, ?, ?, ?, ?)", - (rebooted, updated, connected, waiting, - incomplete_bytes)) + self.usage.update_stats( + rebooted=self._rebooted, + updated=self._timestamp(), + connected=len(self.active_connections._connections), + waiting=len(self.pending_requests._requests), + incomplete_bytes=sum( + tc._total_sent + for tc in self.active_connections._connections + ), + )