Hybrid agent context management

newrelic/api/opentelemetry.py

@@ -17,6 +17,10 @@
 from contextlib import contextmanager
 
 from opentelemetry import trace as otel_api_trace
+from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
+from opentelemetry.baggage.propagation import W3CBaggagePropagator
+from opentelemetry.propagators.composite import CompositePropagator
+from opentelemetry.propagate import set_global_textmap
 
 from newrelic.api.application import application_instance
 from newrelic.api.background_task import BackgroundTask
@@ -26,29 +30,69 @@
 from newrelic.api.message_trace import MessageTrace
 from newrelic.api.message_transaction import MessageTransaction
 from newrelic.api.time_trace import current_trace, notice_error
-from newrelic.api.transaction import Sentinel, current_transaction
+from newrelic.api.transaction import Sentinel, current_transaction, accept_distributed_trace_headers, insert_distributed_trace_headers
 from newrelic.api.web_transaction import WebTransaction
+
 from newrelic.core.otlp_utils import create_resource
 
 _logger = logging.getLogger(__name__)
 
 
+class NRTraceContextPropagator(TraceContextTextMapPropagator):
+    LIST_OF_TRACEPARENT_KEYS = ("traceparent", "HTTP_TRACEPARENT")
+    LIST_OF_TRACESTATE_KEYS = ("tracestate", "HTTP_TRACESTATE")
+    HEADER_KEY_MAPPING = dict((LIST_OF_TRACEPARENT_KEYS, LIST_OF_TRACESTATE_KEYS, ("newrelic", "HTTP_NEWRELIC")))
+
+    def extract(self, carrier, context=None, getter=None):
+        # If we are passing into New Relic, traceparent 
+        # and/or tracestate's keys also need to be NR compatible.
+        nr_headers = {lowercase_name: carrier.get(lowercase_name, carrier.get(http_name, "")) for lowercase_name, http_name in self.HEADER_KEY_MAPPING.items()}
+        accept_distributed_trace_headers(nr_headers)
+
+        return super().extract(carrier=carrier, context=context, getter=getter)
+
+
+    def inject(self, carrier, context=None, setter=None):
+        transaction = current_transaction()
+        # Only insert headers if we have not done so already this transaction
+        # Distributed Trace State will have the following states:
+        #   0 if not set
+        #   1 if already accepted
+        #   2 if inserted but not accepted
+
+        if transaction and not transaction._distributed_trace_state:
+            if isinstance(carrier, dict):
+                nr_headers = list(carrier.items())
+                insert_distributed_trace_headers(nr_headers)
+            elif isinstance(carrier, list):
+                insert_distributed_trace_headers(carrier)
+            else:
+                raise TypeError("Unsupported carrier type")
+
+            return super().inject(carrier=carrier, context=context, setter=setter)
+
+        elif not transaction:
+            return super().inject(carrier=carrier, context=context, setter=setter)
+
+        else:
+            # Do NOT call inject in this case.  Transaction has already received
+            # and/or received and inserted distributed trace headers.
+            pass
+
+
+# Context and Context Propagator Setup
+otel_context_propagator = CompositePropagator(
+    propagators=[
+        NRTraceContextPropagator(),
+        W3CBaggagePropagator(),
+    ]
+)
+set_global_textmap(otel_context_propagator)
+
 # ----------------------------------------------
 # Custom OTel Spans and Traces
 # ----------------------------------------------
 
-# TracerProvider: we can think of this as the agent instance.  Only one can exist
-# SpanProcessor: we can think of this as an application.  In NR, we can have multiple applications
-#   though right now, we can only do SpanProcessor and SynchronousMultiSpanProcessor
-# Tracer: we can think of this as the transaction.
-# Span: we can think of this as the trace.
-# Links functionality has now been enabled but not implemented yet.  Links are relationships
-#   between spans, but lateral in hierarchy.  In NR we only have parent-child relationships.
-#   We may want to preserve this information with a custom attribute.  We can also add this
-#   as a new attribute in a trace, but it will still not be seen in the UI other than a trace
-#   attribute.
-
-
 class Span(otel_api_trace.Span):
     def __init__(
         self,
@@ -141,6 +185,9 @@
         self.nr_trace.__enter__()
 
     def _sampled(self):
+        # NOTE: This logic is using the old logic from before
+        # the various samplers had been implemented.
+        #
         # Uses NR to determine if the trace is sampled
         #
         # transaction.sampled can be `None`, `True`, `False`.
@@ -154,6 +201,12 @@
         # The primary reason for this behavior is because Otel expects to
         # only be able to record information like events and attributes
         # when `is_recording()` == `True`
+        # TODO: Provided that the trace has not already ended, 
+        # configure based on sampler configuration.
+        #   sampler==always_on => return True
+        #   sampler==always_off => return False
+        #   sampler in (default, adaptive, trace_id_ratio_based) 
+        #       => return (if remote parent, parent._sampled(), else transaction.sampled)
 
         if self.otel_parent:
             return bool(self.otel_parent.trace_flags)
@@ -168,7 +221,17 @@
         if not getattr(self, "nr_trace", False):
             return otel_api_trace.INVALID_SPAN_CONTEXT
 
-        otel_tracestate_headers = None
+        if self.nr_transaction.settings.distributed_tracing.enabled:
+            nr_tracestate_headers = (
+                self.nr_transaction._create_distributed_trace_data()
+            )
+
+            nr_tracestate_headers["sa"] = self._sampled()
+            otel_tracestate_headers = [
+                (key, str(value)) for key, value in nr_tracestate_headers.items()
+            ]
+        else:
+            otel_tracestate_headers = None
 
         return otel_api_trace.SpanContext(
             trace_id=int(self.nr_transaction.trace_id, 16),
@@ -208,7 +271,10 @@
             self.nr_trace.name = self._name
 
     def is_recording(self):
-        return self._sampled() and not (getattr(self.nr_trace, None), "end_time", None)
+        # TODO: Similar to self._sampled, we need to
+        # implement a compatible method now that
+        # samplers have been implemented.
+        return self._sampled() and not (getattr(self.nr_trace, "end_time", None))
 
     def set_status(self, status, description=None):
         # TODO: not implemented yet
@@ -277,6 +343,10 @@
         self.nr_application = application_instance()
         self.attributes = attributes or {}
 
+        if not self.nr_application.active:
+            # Force application registration if not already active
+            self.nr_application.activate()
+
         if not self.nr_application.settings.otel_bridge.enabled:
             return otel_api_trace.INVALID_SPAN
 
@@ -286,7 +356,17 @@
         if parent_span_context is None or not parent_span_context.is_valid:
             parent_span_context = None
 
+        # If parent_span_context exists, we can create traceparent
+        # and tracestate headers
+        _headers = {}
+        if parent_span_context and self.nr_application.settings.distributed_tracing.enabled:
+            parent_span_trace_id = parent_span_context.trace_id
+            parent_span_span_id = parent_span_context.span_id
+            parent_span_trace_flags = parent_span_context.trace_flags
+
+
         # If remote_parent, transaction must be created, regardless of kind type
+        # Make sure we transfer DT headers when we are here, if DT is enabled
         if parent_span_context and parent_span_context.is_remote:
             if kind in (otel_api_trace.SpanKind.SERVER, otel_api_trace.SpanKind.CLIENT):
                 # This is a web request
@@ -296,6 +376,10 @@
                 port = self.attributes.get("net.host.port")
                 request_method = self.attributes.get("http.method")
                 request_path = self.attributes.get("http.route")
+
+                update_sampled_flag = False if headers else True
+                headers = headers if headers else _headers
+
                 transaction = WebTransaction(
                     self.nr_application,
                     name=name,
@@ -306,7 +390,16 @@
                     request_path=request_path,
                     headers=headers,
                 )
-            elif kind in (otel_api_trace.SpanKind.PRODUCER, otel_api_trace.SpanKind.INTERNAL):
+
+                # If headers do not contain the traceparent/tracestate
+                # the sampled flag needs to be updated to that of the
+                # parent span.
+                if update_sampled_flag and parent_span_context:
+                    transaction._sampled = bool(parent_span_trace_flags)
+            elif kind in (
+                otel_api_trace.SpanKind.PRODUCER,
+                otel_api_trace.SpanKind.INTERNAL,
+            ):
                 transaction = BackgroundTask(self.nr_application, name=name)
             elif kind == otel_api_trace.SpanKind.CONSUMER:
                 transaction = MessageTransaction(
@@ -338,6 +431,9 @@
                     request_method = self.attributes.get("http.method")
                     request_path = self.attributes.get("http.route")
 
+                    update_GUID_flag = False if headers else True
+                    headers = headers if headers else _headers
+
                     transaction = WebTransaction(
                         self.nr_application,
                         name=name,
@@ -348,6 +444,14 @@
                         request_path=request_path,
                         headers=headers,
                     )
+
+                    # If headers do not contain the traceparent/tracestate
+                    # the transaction GUID needs to be updated to that of
+                    # the parent span.
+                    if update_GUID_flag and parent_span_context:
+                        guid = parent_span_trace_id >> 64
+                        transaction.guid = f"{guid:x}"
+
                 transaction.__enter__()
             elif kind == otel_api_trace.SpanKind.INTERNAL:
                 if transaction:

newrelic/config.py

@@ -4408,11 +4408,27 @@ def _process_module_builtin_defaults():
 
     # Hybrid Agent Hooks
     _process_module_definition(
-        "opentelemetry.trace", "newrelic.hooks.hybridagent_opentelemetry", "instrument_trace_api"
+        "opentelemetry.context",
+        "newrelic.hooks.hybridagent_opentelemetry",
+        "instrument_context_api",
     )
 
     _process_module_definition(
-        "opentelemetry.instrumentation.utils", "newrelic.hooks.hybridagent_opentelemetry", "instrument_utils"
+        "opentelemetry.instrumentation.propagators",
+        "newrelic.hooks.hybridagent_opentelemetry",
+        "instrument_global_propagators_api",
+    )
+
+    _process_module_definition(
+        "opentelemetry.trace",
+        "newrelic.hooks.hybridagent_opentelemetry",
+        "instrument_trace_api",
+    )
+
+    _process_module_definition(
+        "opentelemetry.instrumentation.utils",
+        "newrelic.hooks.hybridagent_opentelemetry",
+        "instrument_utils",
     )
 
 

newrelic/core/config.py

@@ -1431,7 +1431,9 @@ def default_otlp_host(host):
 _settings.azure_operator.enabled = _environ_as_bool("NEW_RELIC_AZURE_OPERATOR_ENABLED", default=False)
 _settings.package_reporting.enabled = _environ_as_bool("NEW_RELIC_PACKAGE_REPORTING_ENABLED", default=True)
 _settings.ml_insights_events.enabled = _environ_as_bool("NEW_RELIC_ML_INSIGHTS_EVENTS_ENABLED", default=False)
-_settings.otel_bridge.enabled = _environ_as_bool("NEW_RELIC_OTEL_BRIDGE_ENABLED", default=False)
+_settings.otel_bridge.enabled = _environ_as_bool(
+    "NEW_RELIC_OTEL_BRIDGE_ENABLED", default=False
+)
 
 
 def global_settings():