Add standard attributes to kaya-otel websocket spans

packages/kaya-otel/README.md

@@ -69,9 +69,13 @@ async def home(ctx: HttpContext):
   `http.response.status_code` when the handler sends the response; 5xx
   statuses mark the span as failed. Exceptions escaping the handler are
   recorded on the span and also mark it as failed.
-- Every WebSocket connection gets one span for its whole lifetime. The close
+- Every WebSocket connection gets one span for its whole lifetime, carrying
-  code is recorded as `kaya.websocket.close_code`; exceptions and configured
+  the applicable standard attributes: `url.scheme`, `url.path`, `url.query`,
-  error close codes mark the span as failed.
+  `network.protocol.name` (`websocket`), `network.transport`,
+  `client.address`/`client.port` and `server.address`/`server.port`. The
+  close code is recorded as `kaya.websocket.close_code`; exceptions and
+  configured error close codes mark the span as failed and set `error.type`
+  (the exception class name, or the close code).
 - W3C `traceparent`/`tracestate` headers on incoming requests are honored,
   so traces propagate from upstream services.
 - Metrics:

packages/kaya-otel/src/kaya/otel/_mixin.py

@@ -84,6 +84,30 @@ def _server_attributes(ctx: HttpContext) -> Dict[str, Any]:
     return attributes
 
 
+def _websocket_attributes(ws: WebSocket) -> Dict[str, Any]:
+    """Standard attributes applicable to a websocket connection span.
+
+    OpenTelemetry defines no websocket-specific conventions; these are the
+    general URL, network and peer attributes that apply to the connection.
+    """
+    attributes: Dict[str, Any] = {
+        'url.scheme': ws.scheme,
+        'url.path': ws.path,
+        'network.protocol.name': 'websocket',
+        'network.transport': 'tcp',
+    }
+    if ws.query_string:
+        attributes['url.query'] = ws.query_string
+    if ws.client is not None:
+        attributes['client.address'] = ws.client[0]
+        attributes['client.port'] = ws.client[1]
+    if ws.server is not None:
+        attributes['server.address'] = ws.server[0]
+        if ws.server[1] is not None:
+            attributes['server.port'] = ws.server[1]
+    return attributes
+
+
 def _normalize_header_name(name: str) -> str:
     return name.lower().replace('-', '_')
 
@@ -201,7 +225,11 @@ class OTelMixin(KayaMixin):
     usual HTTP semantic attributes; the response status code is recorded when
     the handler sends the response and 5xx statuses mark the span as failed.
     Exceptions escaping the handler are recorded and mark the span as failed.
-    WebSocket connections get one span for the whole connection lifetime.
+    WebSocket connections get one span for the whole connection lifetime,
+    carrying the applicable standard websocket URL, network and peer
+    attributes (``url.scheme``, ``url.path``, ``network.protocol.name``,
+    ``client.address``/``server.address``, ...); exceptions and configured
+    error close codes set ``error.type`` and mark the span as failed.
     W3C ``traceparent``/``tracestate`` headers on incoming requests are
     honored, so traces propagate from upstream services.
 
@@ -455,14 +483,7 @@ class OTelMixin(KayaMixin):
     async def _before_websocket(self, ws: WebSocket) -> Optional[WebSocket]:
         if self._is_excluded(ws.path):
             return None
-        holder = self._start_span(
+        holder = self._start_span(f'WS {ws.path}', _websocket_attributes(ws), ws.headers)
-            f'WS {ws.path}',
-            {
-                'url.scheme': ws.scheme,
-                'url.path': ws.path,
-            },
-            ws.headers,
-        )
         self._spans[id(ws)] = holder
         self._call_hook(self._websocket_connect_hook, holder.span, ws)
         return _TracedWebSocket(ws, holder)
@@ -475,10 +496,16 @@ class OTelMixin(KayaMixin):
         error = isinstance(exception, BaseException)
         if isinstance(exception, BaseException):
             holder.span.record_exception(exception)
+            holder.span.set_attribute('error.type', type(exception).__qualname__)
         close_code = holder.close_code
         if close_code is not None:
             holder.span.set_attribute('kaya.websocket.close_code', close_code)
-            error = error or close_code in self._websocket_error_close_codes
+            if close_code in self._websocket_error_close_codes:
+                error = True
+                if not isinstance(exception, BaseException):
+                    # Mirror the HTTP conventions, where server-side error
+                    # statuses are recorded as ``error.type`` values.
+                    holder.span.set_attribute('error.type', str(close_code))
         self._call_hook(self._websocket_close_hook, holder.span, ws)
         self._finish_span(holder, error)
 

packages/kaya-otel/tests/test_otel.py

@@ -279,7 +279,7 @@ class MetricsTest(unittest.TestCase):
 class WebSocketTracingTest(unittest.TestCase):
 
     @staticmethod
-    def _make_ws(headers=()):
+    def _make_ws(headers=(), query_string: bytes = b''):
         async def send(message):
             pass
 
@@ -289,7 +289,7 @@ class WebSocketTracingTest(unittest.TestCase):
         scope = {
             'type': 'websocket',
             'path': '/ws/games/abc',
-            'query_string': b'',
+            'query_string': query_string,
             'scheme': 'ws',
             'client': ('127.0.0.1', 12345),
             'server': ('127.0.0.1', 80),
@@ -308,7 +308,29 @@ class WebSocketTracingTest(unittest.TestCase):
         self.assertEqual(0, len(mixin._spans))
         spans = span_exporter.get_finished_spans()
         self.assertEqual(1, len(spans))
-        self.assertEqual('WS /ws/games/abc', spans[0].name)
+        span = spans[0]
+        self.assertEqual('WS /ws/games/abc', span.name)
+        assert span.attributes is not None
+        self.assertEqual('ws', span.attributes['url.scheme'])
+        self.assertEqual('/ws/games/abc', span.attributes['url.path'])
+        self.assertNotIn('url.query', span.attributes)
+        self.assertEqual('websocket', span.attributes['network.protocol.name'])
+        self.assertEqual('tcp', span.attributes['network.transport'])
+        self.assertEqual('127.0.0.1', span.attributes['client.address'])
+        self.assertEqual(12345, span.attributes['client.port'])
+        self.assertEqual('127.0.0.1', span.attributes['server.address'])
+        self.assertEqual(80, span.attributes['server.port'])
+
+    @async_test
+    async def test_websocket_query_string_recorded(self):
+        mixin, span_exporter, _ = make_mixin()
+        ws = self._make_ws(query_string=b'room=1&token=abc')
+        wrapped = await mixin._before_websocket(ws)
+        await mixin._after_websocket(wrapped)  # type: ignore[arg-type]
+        spans = span_exporter.get_finished_spans()
+        self.assertEqual(1, len(spans))
+        assert spans[0].attributes is not None
+        self.assertEqual('room=1&token=abc', spans[0].attributes['url.query'])
 
     @async_test
     async def test_websocket_traceparent_propagates(self):
@@ -344,6 +366,7 @@ class WebSocketTracingTest(unittest.TestCase):
         self.assertEqual(StatusCode.ERROR, spans[0].status.status_code)
         assert spans[0].attributes is not None
         self.assertEqual(1011, spans[0].attributes['kaya.websocket.close_code'])
+        self.assertEqual('1011', spans[0].attributes['error.type'])
 
     @async_test
     async def test_websocket_exception_marks_span_error(self):
@@ -358,6 +381,8 @@ class WebSocketTracingTest(unittest.TestCase):
         self.assertEqual(1, len(spans))
         self.assertEqual(StatusCode.ERROR, spans[0].status.status_code)
         self.assertTrue(any(event.name == 'exception' for event in spans[0].events))
+        assert spans[0].attributes is not None
+        self.assertEqual('RuntimeError', spans[0].attributes['error.type'])
 
     @async_test
     async def test_websocket_hooks_are_called(self):
@@ -386,6 +411,7 @@ class WebSocketTracingTest(unittest.TestCase):
         self.assertEqual(1, len(spans))
         self.assertEqual(StatusCode.UNSET, spans[0].status.status_code)
         assert spans[0].attributes is not None
+        self.assertNotIn('error.type', spans[0].attributes)
         self.assertTrue(spans[0].attributes['test.websocket_connect_hook'])
         self.assertTrue(spans[0].attributes['test.websocket_close_hook'])