From cb1857fd7d20aceb0fc0ccee39a32b881cbd3b5f Mon Sep 17 00:00:00 2001
From: Diana Huang <dkh@edx.org>
Date: Mon, 29 Jul 2024 14:13:17 -0400
Subject: [PATCH] feat: Create manual spans for monitoring backends.

---
 CHANGELOG.rst                                 |  6 +++++
 edx_django_utils/__init__.py                  |  2 +-
 edx_django_utils/monitoring/__init__.py       |  3 ++-
 .../monitoring/internal/backends.py           | 25 +++++++++++++++++++
 edx_django_utils/monitoring/internal/utils.py | 14 +++++++++++
 5 files changed, 48 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index f5b6c49c..69bc569d 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -11,6 +11,12 @@ Change Log
 
 .. There should always be an "Unreleased" section for changes pending release.
 
+[5.15.0] - 2024-07-29
+---------------------
+Added
+~~~~~
+* Added ``start_trace`` to ``monitoring.utils`` so that manual spans in our monitoring backends can be created.
+
 [5.14.2] - 2024-05-31
 ---------------------
 Fixed
diff --git a/edx_django_utils/__init__.py b/edx_django_utils/__init__.py
index 825ad372..d2b27a88 100644
--- a/edx_django_utils/__init__.py
+++ b/edx_django_utils/__init__.py
@@ -2,7 +2,7 @@
 EdX utilities for Django Application development..
 """
 
-__version__ = "5.14.2"
+__version__ = "5.15.0"
 
 default_app_config = (
     "edx_django_utils.apps.EdxDjangoUtilsConfig"
diff --git a/edx_django_utils/monitoring/__init__.py b/edx_django_utils/monitoring/__init__.py
index 696ba167..a01337c1 100644
--- a/edx_django_utils/monitoring/__init__.py
+++ b/edx_django_utils/monitoring/__init__.py
@@ -29,7 +29,8 @@
     increment,
     record_exception,
     set_custom_attribute,
-    set_custom_attributes_for_course_key
+    set_custom_attributes_for_course_key,
+    start_trace
 )
 # "set_custom_metric*" methods are deprecated
 from .utils import set_custom_metric, set_custom_metrics_for_course_key
diff --git a/edx_django_utils/monitoring/internal/backends.py b/edx_django_utils/monitoring/internal/backends.py
index c982a3fb..f605d4a8 100644
--- a/edx_django_utils/monitoring/internal/backends.py
+++ b/edx_django_utils/monitoring/internal/backends.py
@@ -50,6 +50,13 @@ def record_exception(self):
         Record the exception that is currently being handled.
         """
 
+    @abstractmethod
+    def trace(self, name, params=None):
+        """
+        Create a manual trace with the given name. The underlying implementation is generally
+        assumed to be a contextmanager, and the level above calls it as such.
+        """
+
 
 class NewRelicBackend(TelemetryBackend):
     """
@@ -77,6 +84,10 @@ def record_exception(self):
         # https://docs.newrelic.com/docs/apm/agents/python-agent/python-agent-api/recordexception-python-agent-api/
         newrelic.agent.record_exception()
 
+    def trace(self, name, params=None):
+        # To pass along parameters to New Relic, we can add functionality here. Currently, this is not implemented.
+        return newrelic.agent.function_trace(name)
+
 
 class OpenTelemetryBackend(TelemetryBackend):
     """
@@ -89,6 +100,7 @@ def __init__(self):
         # If import fails, the backend won't be used.
         from opentelemetry import trace
         self.otel_trace = trace
+        self.otel_tracer = trace.get_tracer(__name__)
 
     def set_attribute(self, key, value):
         # Sets the value on the current span, not necessarily the root
@@ -98,6 +110,10 @@ def set_attribute(self, key, value):
     def record_exception(self):
         self.otel_trace.get_current_span().record_exception(sys.exc_info()[1])
 
+    def trace(self, name, params=None):
+        # To pass along parameters to OpenTelemetry, we can add functionality here. Currently, this is not implemented.
+        return self.otel_tracer.start_as_current_span(name)
+
 
 class DatadogBackend(TelemetryBackend):
     """
@@ -119,6 +135,15 @@ def record_exception(self):
         if span := self.dd_tracer.current_span():
             span.set_traceback()
 
+    def trace(self, name, params=None):
+        resource = None
+        service = None
+        if params is not None:
+            resource = params.get('dd-resource')
+            service = params.get('dd-service')
+
+        return self.dd_tracer.start_span(name, resource=resource, service=service, activate=True)
+
 
 # We're using an lru_cache instead of assigning the result to a variable on
 # module load. With the default settings (pointing to a TelemetryBackend
diff --git a/edx_django_utils/monitoring/internal/utils.py b/edx_django_utils/monitoring/internal/utils.py
index 2dc2cfb1..70b3e137 100644
--- a/edx_django_utils/monitoring/internal/utils.py
+++ b/edx_django_utils/monitoring/internal/utils.py
@@ -17,6 +17,8 @@
 At this time, the custom monitoring will only be reported to New Relic.
 
 """
+from contextlib import contextmanager, ExitStack
+
 from .backends import configured_backends
 from .middleware import CachedCustomMonitoringMiddleware
 
@@ -92,6 +94,18 @@ def record_exception():
         backend.record_exception()
 
 
+@contextmanager
+def start_trace(name, params=None):
+    """
+    A context manager that starts spans for all configured backends.
+    """
+
+    # ExitStack handles the underlying context managers.
+    with ExitStack() as stack:
+        for backend in configured_backends():
+            yield stack.enter_context(backend.trace(name, params=params))
+
+
 def background_task(*args, **kwargs):
     """
     Handles monitoring for background tasks that are not passed in through the web server like