clean up tests

Author: mabdinur

ddtrace/contrib/internal/tornado/__init__.py

@@ -15,8 +15,7 @@
 
     # create your handlers
     class MainHandler(tornado.web.RequestHandler):
-        @tornado.gen.coroutine
-        def get(self):
+        async def get(self):
             self.write("Hello, world")
 
     # create your application
@@ -33,10 +32,9 @@ def get(self):
 the ``trace()`` method as usual::
 
     class MainHandler(tornado.web.RequestHandler):
-        @tornado.gen.coroutine
-        def get(self):
-            yield self.notify()
-            yield self.blocking_method()
+        async def get(self):
+            await self.notify()
+            await self.blocking_method()
             with tracer.trace('tornado.before_write') as span:
                 # trace more work in the handler
 
@@ -46,17 +44,15 @@ def blocking_method(self):
             # do something expensive
 
         @tracer.wrap('tornado.notify', service='tornado-notification')
-        @tornado.gen.coroutine
-        def notify(self):
+        async def notify(self):
             # do something
 
 If you are overriding the ``on_finish`` or ``log_exception`` methods on a
 ``RequestHandler``, you will need to call the super method to ensure the
 tracer's patched methods are called::
 
     class MainHandler(tornado.web.RequestHandler):
-        @tornado.gen.coroutine
-        def get(self):
+        async def get(self):
             self.write("Hello, world")
 
         def on_finish(self):

ddtrace/contrib/internal/tornado/application.py

@@ -6,9 +6,54 @@
 from ddtrace import config
 from ddtrace._trace.pin import Pin
 from ddtrace.contrib.internal.tornado.constants import CONFIG_KEY
+from ddtrace.internal.logger import get_logger
 from ddtrace.internal.schema import schematize_service_name
 
 
+log = get_logger(__name__)
+
+
+def _wrap_executor(tracer, fn, args, kwargs, span_name, service=None, resource=None, span_type=None):
+    span = tracer.trace(span_name, service=service, resource=resource, span_type=span_type)
+    prev_active = tracer.context_provider.active()
+
+    def _finish_span(future):
+        try:
+            exc = future.exception()
+            if exc is not None:
+                span.set_exc_info(type(exc), exc, exc.__traceback__)
+        except Exception:
+            try:
+                future.result()
+            except Exception:
+                span.set_traceback()
+        finally:
+            tracer.context_provider.activate(prev_active)
+            span.finish()
+
+    try:
+        result = fn(*args, **kwargs)
+        if callable(getattr(result, "add_done_callback", None)):
+            tracer.context_provider.activate(span)
+            result.add_done_callback(_finish_span)
+        else:
+            span.finish()
+        return result
+    except Exception:
+        span.set_traceback()
+        span.finish()
+        raise
+
+
+def _wrapped_sleep(wrapped, instance, args, kwargs):
+    log.warning(
+        "ddtrace does not support tornado.gen.sleep. Its usage can impact span parenting and trace accuracy. "
+        "With Tornado >6.1 being asyncio-based, tornado.gen.sleep is no longer supported. "
+        "Please migrate to asyncio.sleep() instead."
+    )
+    return wrapped(*args, **kwargs)
+
+
 def tracer_config(__init__, app, args, kwargs):
     """
     Wrap Tornado web application so that we can configure services info and

ddtrace/contrib/internal/tornado/constants.py

@@ -5,4 +5,3 @@
 
 CONFIG_KEY = "datadog_trace"
 REQUEST_SPAN_KEY = "__datadog_request_span"
-FUTURE_SPAN_KEY = "__datadog_future_span"

ddtrace/contrib/internal/tornado/patch.py

@@ -5,6 +5,7 @@
 from wrapt import wrap_function_wrapper as _w
 
 from ddtrace import config
+from ddtrace import tracer
 from ddtrace.internal.utils.formats import asbool
 from ddtrace.internal.utils.wrappers import unwrap as _u
 
@@ -30,39 +31,41 @@ def _supported_versions() -> Dict[str, str]:
     return {"tornado": ">=6.1"}
 
 
+DEFAULT_WRAP_EXECUTOR = None
+
+
 def patch():
-    """
-    Tracing function that patches the Tornado web application so that it will be
-    traced using the given ``tracer``.
-    """
-    # patch only once
     if getattr(tornado, "__datadog_patch", False):
         return
     tornado.__datadog_patch = True
+    global DEFAULT_WRAP_EXECUTOR
+    DEFAULT_WRAP_EXECUTOR = getattr(tracer, "_wrap_executor", None)
+    tracer._wrap_executor = application._wrap_executor
 
-    # patch Application to initialize properly our settings and tracer
     _w("tornado.web", "Application.__init__", application.tracer_config)
-
-    # patch RequestHandler to trace all Tornado handlers
     _w("tornado.web", "RequestHandler._execute", handlers.execute)
     _w("tornado.web", "RequestHandler.on_finish", handlers.on_finish)
     _w("tornado.web", "RequestHandler.log_exception", handlers.log_exception)
-
-    # patch Template system
     _w("tornado.template", "Template.generate", template.generate)
+    # wrapt handles lazy imports, so we can wrap even if tornado.gen isn't imported yet
+    try:
+        _w("tornado.gen", "sleep", application._wrapped_sleep)
+    except (ImportError, AttributeError):
+        pass
 
 
 def unpatch():
-    """
-    Remove all tracing functions in a Tornado web application.
-    """
     if not getattr(tornado, "__datadog_patch", False):
         return
     tornado.__datadog_patch = False
 
-    # unpatch Tornado
+    tracer._wrap_executor = DEFAULT_WRAP_EXECUTOR
+
     _u(tornado.web.RequestHandler, "_execute")
     _u(tornado.web.RequestHandler, "on_finish")
     _u(tornado.web.RequestHandler, "log_exception")
     _u(tornado.web.Application, "__init__")
     _u(tornado.template.Template, "generate")
+    # Use getattr to avoid creating a local variable that shadows the module-level import
+    if getattr(tornado, "gen", None) and hasattr(tornado.gen, "sleep"):
+        _u(tornado.gen, "sleep")

releasenotes/notes/fix-tornado-wrap-executor-contextvars.yaml

@@ -0,0 +1,8 @@
+---
+fixes:
+  - |
+    tornado: Fixed ``tracer.wrap()`` to correctly measure span durations for Tornado coroutines and
+    functions that return Futures, ensuring spans complete only after the asynchronous operation finishes.
+    ``tornado.gen.sleep`` is no longer supported for Tornado >=6.1; a warning will be logged if used.
+    Users should migrate to ``asyncio.sleep``.
+