Merge pull request #193 from project-receptor/diagnostics

Adding a diagnostics subsystem

Reviewed-by: https://github.com/apps/ansible-zuul

Author: ansible-zuul[bot]

receptor/__main__.py

@@ -1,10 +1,10 @@
 import asyncio
 import logging
 import logging.config
-import signal
 import sys
 
 from .config import ReceptorConfig
+from .diagnostics import log_buffer
 from .logstash_formatter.logstash import LogstashFormatter
 
 logger = logging.getLogger(__name__)
@@ -48,17 +48,13 @@ def main(args=None):
 
     def _f(record):
         record.node_id = config.default_node_id
+        if record.levelno == logging.ERROR:
+            log_buffer.appendleft(record)
         return True
 
     for h in logging.getLogger("receptor").handlers:
         h.addFilter(_f)
 
-    def dump_stacks(signum, frame):
-        for t in asyncio.Task.all_tasks():
-            t.print_stack(file=sys.stderr)
-
-    signal.signal(signal.SIGHUP, dump_stacks)
-
     try:
         config.go()
     except asyncio.CancelledError:

receptor/config.py

@@ -246,7 +246,7 @@ def __init__(self, args=None):
             long_option="--ws_heartbeat",
             default_value=None,
             value_type="int",
-            hint="Set heartbeat interval for websocket connections."
+            hint="Set heartbeat interval for websocket connections.",
         )
         # ping options
         self.add_config_option(
@@ -295,7 +295,7 @@ def __init__(self, args=None):
             long_option="--ws_heartbeat",
             default_value=None,
             value_type="int",
-            hint="Set heartbeat interval for websocket connections."
+            hint="Set heartbeat interval for websocket connections.",
         )
         # send options
         self.add_config_option(
@@ -344,7 +344,7 @@ def __init__(self, args=None):
             long_option="--ws_heartbeat",
             default_value=None,
             value_type="int",
-            hint="Set heartbeat interval for websocket connections."
+            hint="Set heartbeat interval for websocket connections.",
         )
         # status options
         self.add_config_option(
@@ -378,7 +378,7 @@ def __init__(self, args=None):
             long_option="--ws_heartbeat",
             default_value=None,
             value_type="int",
-            hint="Set heartbeat interval for websocket connections."
+            hint="Set heartbeat interval for websocket connections.",
         )
         self.parse_options(args)
 

receptor/connection/base.py

@@ -7,6 +7,7 @@
 from .. import fileio
 from ..bridgequeue import BridgeQueue
 from ..messages.framed import FramedBuffer
+from ..stats import bytes_recv
 
 logger = logging.getLogger(__name__)
 
@@ -59,6 +60,7 @@ async def receive(self):
             async for msg in self.conn:
                 if self.conn.closed:
                     break
+                bytes_recv.inc(len(msg))
                 await self.buf.put(msg)
         except ConnectionResetError:
             logger.debug("receive: other side closed the connection")

receptor/connection/sock.py

@@ -1,5 +1,6 @@
 import asyncio
 import logging
+
 from .base import Transport, log_ssl_detail
 
 logger = logging.getLogger(__name__)
@@ -31,6 +32,20 @@ async def send(self, q):
             self.writer.write(chunk)
         await self.writer.drain()
 
+    def _diagnostics(self):
+        t = self.writer._transport.get_extra_info
+        addr, port = t("peername", (None, None))
+        return {
+            "address": addr,
+            "port": port,
+            "compression": t("compression"),
+            "cipher": t("cipher"),
+            "peercert": t("peercert"),
+            "sslcontext": t("sslcontext"),
+            "closed": self.closed,
+            "chunk_size": self.chunk_size,
+        }
+
 
 async def connect(host, port, factory, loop=None, ssl=None, reconnect=True):
     if not loop:

receptor/connection/ws.py

@@ -33,15 +33,20 @@ async def send(self, q):
 
 
 async def connect(
-    uri, factory, loop=None, ssl_context=None, reconnect=True,
-    ws_extra_headers=None, ws_heartbeat=None
+    uri,
+    factory,
+    loop=None,
+    ssl_context=None,
+    reconnect=True,
+    ws_extra_headers=None,
+    ws_heartbeat=None,
 ):
     if not loop:
         loop = asyncio.get_event_loop()
 
     worker = factory()
     try:
-        proxy_scheme = {'ws': 'http', 'wss': 'https'}[urlparse(uri).scheme]
+        proxy_scheme = {"ws": "http", "wss": "https"}[urlparse(uri).scheme]
         proxies = proxies_from_env()
         if proxy_scheme in proxies:
             proxy = proxies[proxy_scheme].proxy
@@ -50,8 +55,12 @@ async def connect(
             proxy = None
             proxy_auth = None
         async with aiohttp.ClientSession().ws_connect(
-            uri, ssl=ssl_context, headers=ws_extra_headers,
-            proxy=proxy, proxy_auth=proxy_auth, heartbeat=ws_heartbeat
+            uri,
+            ssl=ssl_context,
+            headers=ws_extra_headers,
+            proxy=proxy,
+            proxy_auth=proxy_auth,
+            heartbeat=ws_heartbeat,
         ) as ws:
             log_ssl_detail(ws)
             t = WebSocket(ws)

receptor/controller.py

@@ -8,6 +8,7 @@
 
 from .connection.base import Worker
 from .connection.manager import Manager
+from .diagnostics import status
 from .messages.framed import FileBackedBuffer, FramedMessage
 from .receptor import Receptor
 
@@ -39,6 +40,7 @@ def __init__(self, config, loop=asyncio.get_event_loop(), queue=None):
         if self.queue is None:
             self.queue = asyncio.Queue(loop=loop)
         self.receptor.response_queue = self.queue
+        self.status_task = loop.create_task(status(self.receptor))
 
     async def shutdown_loop(self):
         tasks = [

receptor/diagnostics.py

@@ -0,0 +1,141 @@
+import asyncio
+import inspect
+import json
+import logging
+import os
+import signal
+import traceback as tb
+import types
+from collections import defaultdict, deque
+from datetime import datetime
+from functools import singledispatch
+
+from prometheus_client import generate_latest
+
+from . import fileio
+from .logstash_formatter.logstash import LogstashFormatter
+
+logger = logging.getLogger(__name__)
+
+log_buffer = deque(maxlen=10)
+fmt = LogstashFormatter()
+trigger = asyncio.Event()
+
+signal.signal(signal.SIGHUP, lambda n, h: trigger.set())
+
+
+@singledispatch
+def extract_module(o):
+    return o.__module__
+
+
+@extract_module.register(types.CoroutineType)
+def extract_coro(coro):
+    return inspect.getmodule(coro, coro.cr_code.co_filename).__name__
+
+
+@extract_module.register(types.GeneratorType)
+def extract_gen(gen):
+    return inspect.getmodule(gen, gen.gi_code.co_filename).__name__
+
+
+@singledispatch
+def encode(o):
+    return json.JSONEncoder().default(o)
+
+
+@encode.register(set)
+def encode_set(s):
+    return list(s)
+
+
+@encode.register(types.FunctionType)
+def encode_function_type(func):
+    return f"{func.__module__}.{func.__qualname__}"
+
+
+@encode.register(datetime)
+def encode_datetime(o):
+    return o.isoformat()
+
+
+def structure_task(task):
+    coro = task._coro
+    try:
+        mod = extract_module(coro)
+    except Exception:
+        mod = "<unknown>"
+
+    out = {"state": task._state, "name": f"{mod}.{coro.__qualname__}", "stack": []}
+
+    try:
+        stack = tb.extract_stack(task.get_stack()[0])
+        out["stack"] = [
+            {"filename": fs.filename, "line": fs.line, "lineno": fs.lineno} for fs in stack
+        ]
+    except IndexError:
+        pass
+
+    return out
+
+
+def tasks():
+    d = defaultdict(list)
+    for t in asyncio.Task.all_tasks():
+        st = structure_task(t)
+        state = st.pop("state")
+        d[state].append(st)
+    return [{"state": state, "items": tasks} for state, tasks in d.items()]
+
+
+def format_connection(node_id, connection, capabilities):
+    d = connection._diagnostics()
+    d["node_id"] = node_id
+    d["capabilities"] = capabilities
+    return d
+
+
+def format_router(router):
+    edges = [
+        {"left": edge[0], "right": edge[1], "cost": cost} for edge, cost in router._edges.items()
+    ]
+
+    neighbors = [
+        {"node_id": node_id, "items": values} for node_id, values in router._neighbors.items()
+    ]
+
+    table = [
+        {"destination_node_id": node_id, "next_hop": v[0], "cost": v[1]}
+        for node_id, v in router.routing_table.items()
+    ]
+
+    return {"nodes": router._nodes, "edges": edges, "neighbors": neighbors, "table": table}
+
+
+async def status(receptor_object):
+    path = os.path.join(receptor_object.base_path, "diagnostics.json")
+    doc = {}
+    doc["config"] = receptor_object.config._parsed_args.__dict__
+    doc["node_id"] = receptor_object.node_id
+    while True:
+        trigger.clear()
+        doc["datetime"] = datetime.utcnow()
+        doc["recent_errors"] = list(fmt._record_to_dict(r) for r in log_buffer)
+        doc["connections"] = [
+            format_connection(node_id, conn, receptor_object.node_capabilities[node_id])
+            for node_id, connections in receptor_object.connections.items()
+            for conn in connections
+        ]
+        doc["routes"] = format_router(receptor_object.router)
+        doc["tasks"] = tasks()
+        doc["metrics"] = generate_latest()
+        try:
+            await fileio.write(path, json.dumps(doc, default=encode), mode="w")
+        except Exception:
+            logger.exception("failed to dump diagnostic data")
+
+        # run every 30 seconds, or when triggered
+        try:
+            await asyncio.wait_for(trigger.wait(), 30)
+        except asyncio.TimeoutError:
+            pass

receptor/entrypoints.py

@@ -35,7 +35,7 @@ async def node_keepalive():
             controller.add_peer(
                 peer,
                 ws_extra_headers=config.node_ws_extra_headers,
-                ws_heartbeat=config.node_ws_heartbeat
+                ws_heartbeat=config.node_ws_heartbeat,
             )
         if config.node_keepalive_interval > 1:
             controller.loop.create_task(node_keepalive())
@@ -179,7 +179,7 @@ async def status_entrypoint():
             config.status_ws_extra_headers,
             config.status_ws_heartbeat,
             print_status,
-            noop
+            noop,
         )
 
     async def print_status():

receptor/logstash_formatter/logstash.py

@@ -58,13 +58,7 @@ def __init__(
             except Exception:
                 self.source_host = ""
 
-    def format(self, record):
-        """
-        Format a log record to JSON, if the message is a dict
-        assume an empty message and use the dict as additional
-        fields.
-        """
-
+    def _record_to_dict(self, record):
         fields = record.__dict__.copy()
 
         if isinstance(record.msg, dict):
@@ -107,7 +101,16 @@ def format(self, record):
                 "@fields": self._build_fields(logr, fields),
             }
         )
+        return logr
+
+    def format(self, record):
+        """
+        Format a log record to JSON, if the message is a dict
+        assume an empty message and use the dict as additional
+        fields.
+        """
 
+        logr = self._record_to_dict(record)
         return json.dumps(logr, default=self.json_default, cls=self.json_cls)
 
     def _build_fields(self, defaults, fields):