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,77 @@
 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):
+    HEADER_KEYS = ("traceparent", "tracestate", "newrelic")
+
+    def extract(self, carrier, context=None, getter=None):
+        transaction = current_transaction()
+        # If we are passing into New Relic, traceparent 
+        # and/or tracestate's keys also need to be NR compatible.
+
+        if transaction:
+            nr_headers = {header_key: ("" if not getter.get(carrier, header_key) else getter.get(carrier, header_key)[0]) for header_key in self.HEADER_KEYS}
+            transaction.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
+        # New Relic's Distributed Trace State will have the following states:
+        #   0 (00) if not set:
+        #       Transaction has not inserted any outbound headers nor has
+        #       it accepted any inbound headers (yet).
+        #   1 (01) if already accepted:
+        #       Transaction has accepted inbound headers and is able to 
+        #       insert outbound headers to the next app if needed.
+        #   2 (10) if inserted but not accepted:
+        #       Transaction has inserted outbound headers already.
+        #       Do not insert outbound headers multiple times. This is
+        #       a fundamental difference in OTel vs NR behavior: if
+        #       headers are inserted by OTel multiple times, it will
+        #       propagate the last set of data that was inserted. NR
+        #       will not allow more than one header insertion per
+        #       transaction.
+        #   3 (11) if accepted, then inserted:
+        #       Transaction has accepted inbound headers and has inserted
+        #       outbound headers.
+
+        if not transaction:
+            return super().inject(carrier=carrier, context=context, setter=setter)
+
+        if transaction._distributed_trace_state < 2:
+            nr_headers = []
+            transaction.insert_distributed_trace_headers(nr_headers)
+            for key, value in nr_headers:
+                setter.set(carrier, key, value)
+
+            return super().inject(carrier=carrier, context=context, setter=setter)
+
+        # If distributed_trace_state == 2 or 3, do not inject headers.
+
+
+# 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,
@@ -100,7 +152,7 @@ def __init__(
             _logger.error(
                 "OpenTelemetry span (%s) and NR trace (%s) do not match nor correspond to a remote span. Open Telemetry span will not be reported to New Relic. Please report this problem to New Relic.",
                 self.otel_parent,
-                current_nr_trace,
+                current_nr_trace,   # NR parent trace
             )
             return
 
@@ -141,24 +193,20 @@ def __init__(
         self.nr_trace.__enter__()
 
     def _sampled(self):
-        # Uses NR to determine if the trace is sampled
-        #
-        # transaction.sampled can be `None`, `True`, `False`.
-        # If `None`, this has not been computed by NR which
-        # can also mean the following:
-        # 1. There was not a context passed in that explicitly has sampling disabled.
-        #   This flag would be found in the traceparent or traceparent and tracespan headers.
-        # 2. Transaction was not created where DT headers are accepted during __init__
-        #   Therefore, we will treat a value of `None` as `True` for now.
-        #
-        # 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: Refine this logic to respect the parent
+        # trace/span, even if the parent is from OTel.
+        # priority: except for always_on or always off, 
+        # sampled flag from parent if parent exists takes
+        # precedence over whatever sampling algorithm is
+        # enabled.
 
         if self.otel_parent:
             return bool(self.otel_parent.trace_flags)
+        elif self.nr_transaction:
+            self.nr_transaction._make_sampling_decision()
+            return self.nr_transaction.sampled
         else:
-            return bool(self.nr_transaction and (self.nr_transaction.sampled or (self.nr_transaction.sampled is None)))
+            return False
 
     def _remote(self):
         # Remote span denotes if propagated from a remote parent
@@ -168,14 +216,12 @@ def get_span_context(self):
         if not getattr(self, "nr_trace", False):
             return otel_api_trace.INVALID_SPAN_CONTEXT
 
-        otel_tracestate_headers = None
-
         return otel_api_trace.SpanContext(
             trace_id=int(self.nr_transaction.trace_id, 16),
             span_id=int(self.nr_trace.guid, 16),
             is_remote=self._remote(),
             trace_flags=otel_api_trace.TraceFlags(0x01 if self._sampled() else 0x00),
-            trace_state=otel_api_trace.TraceState(otel_tracestate_headers),
+            trace_state=otel_api_trace.TraceState(),
         )
 
     def set_attribute(self, key, value):
@@ -208,7 +254,13 @@ def update_name(self, name):
             self.nr_trace.name = self._name
 
     def is_recording(self):
-        return self._sampled() and not (getattr(self.nr_trace, None), "end_time", None)
+        # TODO: Refine this logic along with `_sampled()`
+        if getattr(self.nr_trace, "end_time", None):
+            return False
+        if not getattr(self.nr_transaction, "priority", None):
+            self.nr_transaction._make_sampling_decision()
+
+        return self.nr_transaction.priority > 0
 
     def set_status(self, status, description=None):
         # TODO: not implemented yet
@@ -277,6 +329,10 @@ def start_span(
         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 +342,15 @@ def start_span(
         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_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 +360,10 @@ def start_span(
                 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 +374,16 @@ def start_span(
                     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._remote_parent_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 +415,9 @@ def start_span(
                     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 +428,14 @@ def start_span(
                         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

@@ -4429,11 +4429,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():