From 638adc71ec46c97b86732dc72b2bc68d156553dd Mon Sep 17 00:00:00 2001
From: Brian Warner <warner@lothar.com>
Date: Sun, 26 Jun 2016 14:58:31 -0700
Subject: [PATCH] server: add transit to stats.json

---
 src/wormhole/server/server.py         |  2 ++
 src/wormhole/server/transit_server.py | 34 ++++++++++++++++++++++++++-
 2 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/src/wormhole/server/server.py b/src/wormhole/server/server.py
index 8903aea..c2baf20 100644
--- a/src/wormhole/server/server.py
+++ b/src/wormhole/server/server.py
@@ -95,6 +95,7 @@ class RelayServer(service.MultiService):
         self._root = root
         self._rendezvous_web_service = rendezvous_web_service
         self._rendezvous_websocket = wsrf
+        self._transit = None
         if transit_port:
             self._transit = transit
             self._transit_service = transit_service
@@ -125,6 +126,7 @@ class RelayServer(service.MultiService):
         data["valid_until"] = now + validity
 
         data["rendezvous"] = self._rendezvous.get_stats()
+        data["transit"] = self._transit.get_stats()
 
         with open(tmpfn, "wb") as f:
             json.dump(data, f, indent=1)
diff --git a/src/wormhole/server/transit_server.py b/src/wormhole/server/transit_server.py
index 17d42a0..44478be 100644
--- a/src/wormhole/server/transit_server.py
+++ b/src/wormhole/server/transit_server.py
@@ -1,5 +1,5 @@
 from __future__ import print_function, unicode_literals
-import re, time
+import re, time, collections
 from twisted.python import log
 from twisted.internet import protocol
 from twisted.application import service
@@ -166,6 +166,7 @@ class Transit(protocol.ServerFactory, service.MultiService):
         self._blur_usage = blur_usage
         self._pending_requests = {} # token -> TransitConnection
         self._active_connections = set() # TransitConnection
+        self._counts = collections.defaultdict(int)
 
     def connection_got_token(self, token, p):
         if token in self._pending_requests:
@@ -193,6 +194,8 @@ class Transit(protocol.ServerFactory, service.MultiService):
                          (started, total_time, waiting_time,
                           total_bytes, result))
         self._db.commit()
+        self._counts[result] += 1
+        self._counts["bytes"] += total_bytes
 
     def transitFinished(self, p, token, description):
         for token,tc in self._pending_requests.items():
@@ -205,3 +208,32 @@ class Transit(protocol.ServerFactory, service.MultiService):
     def transitFailed(self, p):
         log.msg("transitFailed %r" % p)
         pass
+
+    def get_stats(self):
+        stats = {}
+        def q(query, values=()):
+            row = self._db.execute(query, values).fetchone()
+            return list(row.values())[0]
+
+        # current status: expected to be zero most of the time
+        c = stats["active"] = {}
+        c["connected"] = len(self._active_connections) / 2
+        c["waiting"] = len(self._pending_requests)
+
+        # usage since last reboot
+        rb = stats["since_reboot"] = {}
+        for result, count in self._counts.items():
+            rb[result] = count
+
+        # historical usage (all-time)
+        u = stats["all_time"] = {}
+        u["total"] = q("SELECT COUNT() FROM `transit_usage`")
+        u["happy"] = q("SELECT COUNT() FROM `transit_usage`"
+                       " WHERE `result`='happy'")
+        u["lonely"] = q("SELECT COUNT() FROM `transit_usage`"
+                        " WHERE `result`='lonely'")
+        u["errory"] = q("SELECT COUNT() FROM `transit_usage`"
+                        " WHERE `result`='errory'")
+        u["bytes"] = q("SELECT SUM(`total_bytes`) FROM `transit_usage`") or 0
+
+        return stats