Add tornado WebSocketHandler instrumentation support (open-telemetry#3498)

* Add tornado WebSocketHandler instrumentation support. (open-telemetry#2761)

* Linting

* Update CHANGELOG.md

* Apply refactor changes from open-telemetry#3582

---------

Co-authored-by: Riccardo Magliocchetti <[email protected]>

Author: devmonkey22

CHANGELOG.md

@@ -16,6 +16,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - `opentelemetry-instrumentation-pika` Added instrumentation for All `SelectConnection` adapters
   ([#3584](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/3584))
+- `opentelemetry-instrumentation-tornado` Add support for `WebSocketHandler` instrumentation
+  ([#3498](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/3498))
 
 ### Fixed
 

instrumentation/opentelemetry-instrumentation-tornado/src/opentelemetry/instrumentation/tornado/__init__.py

@@ -162,6 +162,7 @@ def client_response_hook(span, future):
 from typing import Collection, Dict
 
 import tornado.web
+import tornado.websocket
 import wrapt
 from wrapt import wrap_function_wrapper
 
@@ -363,12 +364,20 @@ def patch_handler_class(tracer, server_histograms, cls, request_hook=None):
         "prepare",
         partial(_prepare, tracer, server_histograms, request_hook),
     )
-    _wrap(cls, "on_finish", partial(_on_finish, tracer, server_histograms))
     _wrap(
         cls,
         "log_exception",
         partial(_log_exception, tracer, server_histograms),
     )
+
+    if issubclass(cls, tornado.websocket.WebSocketHandler):
+        _wrap(
+            cls,
+            "on_close",
+            partial(_websockethandler_on_close, tracer, server_histograms),
+        )
+    else:
+        _wrap(cls, "on_finish", partial(_on_finish, tracer, server_histograms))
     return True
 
 
@@ -377,8 +386,11 @@ def unpatch_handler_class(cls):
         return
 
     unwrap(cls, "prepare")
-    unwrap(cls, "on_finish")
     unwrap(cls, "log_exception")
+    if issubclass(cls, tornado.websocket.WebSocketHandler):
+        unwrap(cls, "on_close")
+    else:
+        unwrap(cls, "on_finish")
     delattr(cls, _OTEL_PATCHED_KEY)
 
 
@@ -409,13 +421,21 @@ def _prepare(
 
 
 def _on_finish(tracer, server_histograms, func, handler, args, kwargs):
-    response = func(*args, **kwargs)
-
-    _record_on_finish_metrics(server_histograms, handler)
+    try:
+        return func(*args, **kwargs)
+    finally:
+        _record_on_finish_metrics(server_histograms, handler)
+        _finish_span(tracer, handler)
 
-    _finish_span(tracer, handler)
 
-    return response
+def _websockethandler_on_close(
+    tracer, server_histograms, func, handler, args, kwargs
+):
+    try:
+        func()
+    finally:
+        _record_on_finish_metrics(server_histograms, handler)
+        _finish_span(tracer, handler)
 
 
 def _log_exception(tracer, server_histograms, func, handler, args, kwargs):

instrumentation/opentelemetry-instrumentation-tornado/tests/test_instrumentation.py

@@ -13,8 +13,10 @@
 # limitations under the License.
 
 
+import asyncio
 from unittest.mock import Mock, patch
 
+import tornado.websocket
 from http_server_mock import HttpServerMock
 from tornado.httpclient import HTTPClientError
 from tornado.testing import AsyncHTTPTestCase
@@ -454,6 +456,53 @@ def test_handler_on_finish(self):
 
         self.assertEqual(auditor.kind, SpanKind.INTERNAL)
 
+    @tornado.testing.gen_test()
+    async def test_websockethandler(self):
+        ws_client = await tornado.websocket.websocket_connect(
+            f"ws://127.0.0.1:{self.get_http_port()}/echo_socket"
+        )
+
+        await ws_client.write_message("world")
+        resp = await ws_client.read_message()
+        self.assertEqual(resp, "hello world")
+
+        ws_client.close()
+        await asyncio.sleep(0.5)
+
+        spans = self.sorted_spans(self.memory_exporter.get_finished_spans())
+        self.assertEqual(len(spans), 3)
+        close_span, msg_span, req_span = spans
+
+        self.assertEqual(req_span.name, "GET /echo_socket")
+        self.assertEqual(req_span.context.trace_id, msg_span.context.trace_id)
+        self.assertIsNone(req_span.parent)
+        self.assertEqual(req_span.kind, SpanKind.SERVER)
+        self.assertSpanHasAttributes(
+            req_span,
+            {
+                HTTP_METHOD: "GET",
+                HTTP_SCHEME: "http",
+                HTTP_HOST: f"127.0.0.1:{self.get_http_port()}",
+                HTTP_TARGET: "/echo_socket",
+                HTTP_CLIENT_IP: "127.0.0.1",
+                HTTP_STATUS_CODE: 101,
+                "tornado.handler": "tests.tornado_test_app.EchoWebSocketHandler",
+            },
+        )
+
+        self.assertEqual(msg_span.name, "audit_message")
+        self.assertFalse(msg_span.context.is_remote)
+        self.assertEqual(msg_span.kind, SpanKind.INTERNAL)
+        self.assertEqual(msg_span.parent.span_id, req_span.context.span_id)
+
+        self.assertEqual(close_span.name, "audit_on_close")
+        self.assertFalse(close_span.context.is_remote)
+        self.assertEqual(close_span.parent.span_id, req_span.context.span_id)
+        self.assertEqual(
+            close_span.context.trace_id, msg_span.context.trace_id
+        )
+        self.assertEqual(close_span.kind, SpanKind.INTERNAL)
+
     def test_exclude_lists(self):
         def test_excluded(path):
             self.fetch(path)

instrumentation/opentelemetry-instrumentation-tornado/tests/tornado_test_app.py

@@ -2,6 +2,7 @@
 import time
 
 import tornado.web
+import tornado.websocket
 from tornado import gen
 
 
@@ -110,6 +111,16 @@ def get(self):
         raise tornado.web.HTTPError(403)
 
 
+class EchoWebSocketHandler(tornado.websocket.WebSocketHandler):
+    async def on_message(self, message):
+        with self.application.tracer.start_as_current_span("audit_message"):
+            self.write_message(f"hello {message}")
+
+    def on_close(self):
+        with self.application.tracer.start_as_current_span("audit_on_close"):
+            time.sleep(0.05)
+
+
 def make_app(tracer):
     app = tornado.web.Application(
         [
@@ -122,6 +133,7 @@ def make_app(tracer):
             (r"/ping", HealthCheckHandler),
             (r"/test_custom_response_headers", CustomResponseHeaderHandler),
             (r"/raise_403", RaiseHTTPErrorHandler),
+            (r"/echo_socket", EchoWebSocketHandler),
         ]
     )
     app.tracer = tracer