Starlette Instrumentation Updates (#254)

* Fix starlette background task bug Co-authored-by: Uma Annamalai <[email protected]> * Formatting Co-authored-by: Uma Annamalai <[email protected]>
newrelic · Jun 23, 2021 · 208344e · 208344e
1 parent b891adc
commit 208344e
Show file tree

Hide file tree

Showing 7 changed files with 365 additions and 109 deletions.
diff --git a/newrelic/hooks/framework_starlette.py b/newrelic/hooks/framework_starlette.py
@@ -14,15 +14,18 @@
 
 from newrelic.api.asgi_application import wrap_asgi_application
 from newrelic.api.background_task import BackgroundTaskWrapper
-from newrelic.api.time_trace import current_trace
-from newrelic.api.function_trace import FunctionTraceWrapper, wrap_function_trace
+from newrelic.api.function_trace import FunctionTraceWrapper
+from newrelic.api.time_trace import current_trace, notice_error
+from newrelic.api.transaction import current_transaction
+from newrelic.common.coroutine import is_coroutine_function
 from newrelic.common.object_names import callable_name
-from newrelic.common.object_wrapper import wrap_function_wrapper, function_wrapper, FunctionWrapper
-from newrelic.core.trace_cache import trace_cache
-from newrelic.api.time_trace import notice_error
+from newrelic.common.object_wrapper import (
+    FunctionWrapper,
+    function_wrapper,
+    wrap_function_wrapper,
+)
 from newrelic.core.config import should_ignore_error
-from newrelic.common.coroutine import is_coroutine_function
-from newrelic.api.transaction import current_transaction
+from newrelic.core.trace_cache import trace_cache
 
 
 def framework_details():
@@ -93,10 +96,19 @@ def wrap_request(wrapped, instance, args, kwargs):
 def wrap_background_method(wrapped, instance, args, kwargs):
     func = getattr(instance, "func", None)
     if func:
-        instance.func = BackgroundTaskWrapper(func)
+        instance.func = wrap_background_task(func)
     return wrapped(*args, **kwargs)
 
 
+@function_wrapper
+def wrap_background_task(wrapped, instance, args, kwargs):
+    transaction = current_transaction(active_only=False)
+    if not transaction:
+        return BackgroundTaskWrapper(wrapped)(*args, **kwargs)
+    else:
+        return FunctionTraceWrapper(wrapped)(*args, **kwargs)
+
+
 async def middleware_wrapper(wrapped, instance, args, kwargs):
     transaction = current_transaction()
     if transaction:
@@ -158,7 +170,9 @@ async def wrap_exception_handler_async(coro, exc):
 
 def wrap_exception_handler(wrapped, instance, args, kwargs):
     if is_coroutine_function(wrapped):
-        return wrap_exception_handler_async(FunctionTraceWrapper(wrapped)(*args, **kwargs), bind_exc(*args, **kwargs))
+        return wrap_exception_handler_async(
+            FunctionTraceWrapper(wrapped)(*args, **kwargs), bind_exc(*args, **kwargs)
+        )
     else:
         with RequestContext(bind_request(*args, **kwargs)):
             response = FunctionTraceWrapper(wrapped)(*args, **kwargs)
@@ -168,14 +182,16 @@ def wrap_exception_handler(wrapped, instance, args, kwargs):
 
 def wrap_server_error_handler(wrapped, instance, args, kwargs):
     result = wrapped(*args, **kwargs)
-    handler = getattr(instance, 'handler', None)
+    handler = getattr(instance, "handler", None)
     if handler:
         instance.handler = FunctionWrapper(handler, wrap_exception_handler)
     return result
 
 
 def wrap_add_exception_handler(wrapped, instance, args, kwargs):
-    exc_class_or_status_code, handler, args, kwargs = bind_add_exception_handler(*args, **kwargs)
+    exc_class_or_status_code, handler, args, kwargs = bind_add_exception_handler(
+        *args, **kwargs
+    )
     handler = FunctionWrapper(handler, wrap_exception_handler)
     return wrapped(exc_class_or_status_code, handler, *args, **kwargs)
 
@@ -207,25 +223,36 @@ def instrument_starlette_requests(module):
 
 
 def instrument_starlette_middleware_errors(module):
-    wrap_function_wrapper(module, "ServerErrorMiddleware.__call__", error_middleware_wrapper)
+    wrap_function_wrapper(
+        module, "ServerErrorMiddleware.__call__", error_middleware_wrapper
+    )
 
-    wrap_function_wrapper(module, "ServerErrorMiddleware.__init__", wrap_server_error_handler)
+    wrap_function_wrapper(
+        module, "ServerErrorMiddleware.__init__", wrap_server_error_handler
+    )
 
-    wrap_function_wrapper(module, "ServerErrorMiddleware.error_response", wrap_exception_handler)
+    wrap_function_wrapper(
+        module, "ServerErrorMiddleware.error_response", wrap_exception_handler
+    )
 
-    wrap_function_wrapper(module, "ServerErrorMiddleware.debug_response", wrap_exception_handler)
+    wrap_function_wrapper(
+        module, "ServerErrorMiddleware.debug_response", wrap_exception_handler
+    )
 
 
 def instrument_starlette_exceptions(module):
-    wrap_function_wrapper(module, "ExceptionMiddleware.__call__", error_middleware_wrapper)
+    wrap_function_wrapper(
+        module, "ExceptionMiddleware.__call__", error_middleware_wrapper
+    )
 
-    wrap_function_wrapper(module, "ExceptionMiddleware.http_exception",
-        wrap_exception_handler)
+    wrap_function_wrapper(
+        module, "ExceptionMiddleware.http_exception", wrap_exception_handler
+    )
 
-    wrap_function_wrapper(module, "ExceptionMiddleware.add_exception_handler",
-        wrap_add_exception_handler)
+    wrap_function_wrapper(
+        module, "ExceptionMiddleware.add_exception_handler", wrap_add_exception_handler
+    )
 
 
 def instrument_starlette_background_task(module):
     wrap_function_wrapper(module, "BackgroundTask.__call__", wrap_background_method)
-
diff --git a/...ramework_starlette/_target_application.py → .../framework_starlette/_test_application.py b/...ramework_starlette/_target_application.py → .../framework_starlette/_test_application.py
@@ -14,12 +14,13 @@
 
 from starlette.applications import Starlette
 from starlette.background import BackgroundTasks
+from starlette.exceptions import HTTPException
 from starlette.responses import PlainTextResponse
 from starlette.routing import Route
-from starlette.exceptions import HTTPException
 from testing_support.asgi_testing import AsgiTest
-from newrelic.api.transaction import current_transaction
+
 from newrelic.api.function_trace import FunctionTrace
+from newrelic.api.transaction import current_transaction
 from newrelic.common.object_names import callable_name
 
 try:
@@ -101,6 +102,7 @@ async def bg_task_async():
 def bg_task_non_async():
     pass
 
+
 routes = [
     Route("/index", index),
     Route("/418", teapot),
@@ -130,7 +132,11 @@ async def middleware_decorator(request, call_next):
 # Generating target applications
 app_name_map = {
     "no_error_handler": (True, False, {}),
-    "async_error_handler_no_middleware": (False, False, {Exception: async_error_handler}),
+    "async_error_handler_no_middleware": (
+        False,
+        False,
+        {Exception: async_error_handler},
+    ),
     "non_async_error_handler_no_middleware": (False, False, {}),
     "no_middleware": (False, False, {}),
     "debug_no_middleware": (False, True, {}),
@@ -145,12 +151,21 @@ async def middleware_decorator(request, call_next):
 
     # Instantiate app
     if not middleware_on:
-        app = Starlette(debug=debug, routes=routes, exception_handlers=exception_handlers)
+        app = Starlette(
+            debug=debug, routes=routes, exception_handlers=exception_handlers
+        )
     else:
         if Middleware:
-            app = Starlette(debug=debug, routes=routes, middleware=[Middleware(middleware_factory)], exception_handlers=exception_handlers)
+            app = Starlette(
+                debug=debug,
+                routes=routes,
+                middleware=[Middleware(middleware_factory)],
+                exception_handlers=exception_handlers,
+            )
         else:
-            app = Starlette(debug=debug, routes=routes, exception_handlers=exception_handlers)
+            app = Starlette(
+                debug=debug, routes=routes, exception_handlers=exception_handlers
+            )
             # in earlier versions of starlette, middleware is not a legal argument on the Starlette application class
             # In order to keep the counts the same, we add the middleware twice using the add_middleware interface
             app.add_middleware(middleware_factory)

diff --git a/tests/framework_starlette/_test_bg_tasks.py b/tests/framework_starlette/_test_bg_tasks.py
@@ -0,0 +1,76 @@
+# Copyright 2010 New Relic, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from starlette.applications import Starlette
+from starlette.background import BackgroundTasks
+from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.responses import PlainTextResponse
+from starlette.routing import Route
+from testing_support.asgi_testing import AsgiTest
+
+
+class ASGIStyleMiddleware:
+    def __init__(self, app):
+        self.app = app
+
+    async def __call__(self, scope, receive, send):
+        response = await self.app(scope, receive, send)
+        return response
+
+
+class BaseHTTPStyleMiddleware(BaseHTTPMiddleware):
+    async def dispatch(self, request, call_next):
+        # simple middleware that does absolutely nothing
+        response = await call_next(request)
+        return response
+
+
+async def run_async_bg_task(request):
+    tasks = BackgroundTasks()
+    tasks.add_task(async_bg_task)
+    return PlainTextResponse("Hello, world!", background=tasks)
+
+
+async def run_sync_bg_task(request):
+    tasks = BackgroundTasks()
+    tasks.add_task(sync_bg_task)
+    return PlainTextResponse("Hello, world!", background=tasks)
+
+
+async def async_bg_task():
+    pass
+
+
+async def sync_bg_task():
+    pass
+
+
+routes = [
+    Route("/async", run_async_bg_task),
+    Route("/sync", run_sync_bg_task),
+]
+
+# Generating target applications
+target_application = {}
+
+app = Starlette(routes=routes)
+app.add_middleware(ASGIStyleMiddleware)
+target_application["asgi"] = AsgiTest(app)
+
+app = Starlette(routes=routes)
+app.add_middleware(BaseHTTPStyleMiddleware)
+target_application["basehttp"] = AsgiTest(app)
+
+app = Starlette(routes=routes)
+target_application["none"] = AsgiTest(app)
diff --git a/tests/framework_starlette/conftest.py b/tests/framework_starlette/conftest.py
@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import pytest
-
 from testing_support.fixtures import (
     code_coverage_fixture,
     collector_agent_registration_fixture,
@@ -38,10 +36,3 @@
     app_name="Python Agent Test (framework_starlette)",
     default_settings=_default_settings,
 )
-
-
-@pytest.fixture(scope="session")
-def target_application():
-    import _target_application
-
-    return _target_application.target_application