Add uvicorn instrumentation. Closes #4

Author: newrelic-opensource

newrelic/config.py

@@ -2351,6 +2351,10 @@ def _process_module_builtin_defaults():
             'newrelic.hooks.external_urllib3',
             'instrument_urllib3_connection')
 
+    _process_module_definition('uvicorn.config',
+            'newrelic.hooks.adapter_uvicorn',
+            'instrument_uvicorn_config')
+
     _process_module_definition('sanic.app',
             'newrelic.hooks.framework_sanic',
             'instrument_sanic_app')

newrelic/core/environment.py

@@ -153,12 +153,25 @@ def environment_settings():
             dispatcher.append(('Dispatcher', 'gunicorn (gevent)'))
         elif 'gunicorn.workers.geventlet' in sys.modules:
             dispatcher.append(('Dispatcher', 'gunicorn (eventlet)'))
+        elif 'uvicorn.workers' in sys.modules:
+            dispatcher.append(('Dispatcher', 'gunicorn (uvicorn)'))
+            uvicorn = sys.modules.get('uvicorn')
+            if hasattr(uvicorn, '__version__'):
+                dispatcher.append(('Worker Version', uvicorn.__version__))
         else:
             dispatcher.append(('Dispatcher', 'gunicorn'))
+
         gunicorn = sys.modules['gunicorn']
         if hasattr(gunicorn, '__version__'):
             dispatcher.append(('Dispatcher Version', gunicorn.__version__))
 
+    if not dispatcher and 'uvicorn' in sys.modules:
+        dispatcher.append(('Dispatcher', 'uvicorn'))
+        uvicorn = sys.modules['uvicorn']
+
+        if hasattr(uvicorn, '__version__'):
+            dispatcher.append(('Dispatcher Version', uvicorn.__version__))
+
     if not dispatcher and 'tornado' in sys.modules:
         dispatcher.append(('Dispatcher', 'tornado'))
         tornado = sys.modules['tornado']

newrelic/hooks/adapter_uvicorn.py

@@ -0,0 +1,36 @@
+# 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 newrelic.api.asgi_application import ASGIApplicationWrapper
+
+
+@property
+def loaded_app(self):
+    return self._nr_loaded_app
+
+
+@loaded_app.setter
+def loaded_app(self, value):
+    # Wrap only the first loaded app
+    if (
+        not getattr(self, "_nr_loaded_app", None)
+        and value
+        and getattr(self, "interface", "") != "wsgi"
+    ):
+        value = ASGIApplicationWrapper(value)
+    self._nr_loaded_app = value
+
+
+def instrument_uvicorn_config(module):
+    module.Config.loaded_app = loaded_app

tests/agent_unittests/test_environment.py

@@ -12,20 +12,99 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import pytest
+import sys
 from newrelic.core.environment import environment_settings
 
 
+def module(version):
+    class Module(object):
+        pass
+
+    if version:
+        Module.__version__ = version
+
+    return Module
+
+
 def test_plugin_list():
     # Let's pretend we fired an import hook
     import newrelic.hooks.adapter_gunicorn
 
     environment_info = environment_settings()
 
     for key, plugin_list in environment_info:
-        if key == 'Plugin List':
+        if key == "Plugin List":
             break
     else:
         assert False, "'Plugin List' not found"
 
     # Check that bogus plugins don't get reported
-    assert 'newrelic.hooks.newrelic' not in plugin_list
+    assert "newrelic.hooks.newrelic" not in plugin_list
+
+
+@pytest.mark.parametrize(
+    "loaded_modules,dispatcher,dispatcher_version,worker_version",
+    (
+        ({"uvicorn": module("4.5.6")}, "uvicorn", "4.5.6", None),
+        (
+            {
+                "gunicorn": module("1.2.3"),
+                "uvicorn": module("4.5.6"),
+                "uvicorn.workers": object(),
+            },
+            "gunicorn (uvicorn)",
+            "1.2.3",
+            "4.5.6",
+        ),
+        ({"uvicorn": object()}, "uvicorn", None, None),
+        (
+            {
+                "gunicorn": object(),
+                "uvicorn": module("4.5.6"),
+                "uvicorn.workers": object(),
+            },
+            "gunicorn (uvicorn)",
+            None,
+            "4.5.6",
+        ),
+        (
+            {"gunicorn": module("1.2.3"), "uvicorn": None, "uvicorn.workers": object()},
+            "gunicorn (uvicorn)",
+            "1.2.3",
+            None,
+        ),
+        (
+            {"gunicorn": object(), "uvicorn": object(), "uvicorn.workers": object()},
+            "gunicorn (uvicorn)",
+            None,
+            None,
+        ),
+    ),
+)
+def test_uvicorn_dispatcher(
+    monkeypatch, loaded_modules, dispatcher, dispatcher_version, worker_version
+):
+    # Let's pretend we load some modules
+    for name, module in loaded_modules.items():
+        monkeypatch.setitem(sys.modules, name, module)
+
+    environment_info = environment_settings()
+
+    actual_dispatcher = None
+    actual_dispatcher_version = None
+    actual_worker_version = None
+    for key, value in environment_info:
+        if key == "Dispatcher":
+            assert actual_dispatcher is None
+            actual_dispatcher = value
+        elif key == "Dispatcher Version":
+            assert actual_dispatcher_version is None
+            actual_dispatcher_version = value
+        elif key == "Worker Version":
+            assert actual_worker_version is None
+            actual_worker_version = value
+
+    assert actual_dispatcher == dispatcher
+    assert actual_dispatcher_version == dispatcher_version
+    assert actual_worker_version == worker_version