feat: contribute xray lambda propagator (#1464)

Co-authored-by: Kayla Reopelle <[email protected]>

Author: yiyuan-he

propagator/xray/lib/opentelemetry/propagator/xray.rb

@@ -12,6 +12,7 @@
 # See the documentation for the `opentelemetry-api` gem for details.
 
 require_relative 'xray/text_map_propagator'
+require_relative 'xray/lambda_text_map_propagator'
 
 module OpenTelemetry
   # Namespace for OpenTelemetry propagator extension libraries
@@ -22,7 +23,8 @@ module XRay
 
       DEBUG_CONTEXT_KEY = Context.create_key('xray-debug-key')
       TEXT_MAP_PROPAGATOR = TextMapPropagator.new
-      private_constant :DEBUG_CONTEXT_KEY, :TEXT_MAP_PROPAGATOR
+      LAMBDA_TEXT_MAP_PROPAGATOR = LambdaTextMapPropagator.new
+      private_constant :DEBUG_CONTEXT_KEY, :TEXT_MAP_PROPAGATOR, :LAMBDA_TEXT_MAP_PROPAGATOR
 
       # @api private
       # Returns a new context with the xray debug flag enabled
@@ -41,6 +43,12 @@ def debug?(context)
       def text_map_propagator
         TEXT_MAP_PROPAGATOR
       end
+
+      # Returns a text map propagator that propagates context in the XRay
+      # format with special handling for Lambda environment.
+      def lambda_text_map_propagator
+        LAMBDA_TEXT_MAP_PROPAGATOR
+      end
     end
   end
 end

propagator/xray/lib/opentelemetry/propagator/xray/lambda_text_map_propagator.rb

@@ -0,0 +1,61 @@
+# frozen_string_literal: true
+
+# Copyright OpenTelemetry Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+# OpenTelemetry is an open source observability framework, providing a
+# general-purpose API, SDK, and related tools required for the instrumentation
+# of cloud-native software, frameworks, and libraries.
+#
+# The OpenTelemetry module provides global accessors for telemetry objects.
+# See the documentation for the `opentelemetry-api` gem for details.
+module OpenTelemetry
+  # Namespace for OpenTelemetry propagator extension libraries
+  module Propagator
+    # Namespace for OpenTelemetry XRay propagation
+    module XRay
+      # Implementation of AWS X-Ray Trace Header propagation with special handling for
+      # Lambda's _X_AMZN_TRACE_ID environment variable
+      class LambdaTextMapPropagator < TextMapPropagator
+        AWS_TRACE_HEADER_ENV_KEY = '_X_AMZN_TRACE_ID'
+
+        # Extract trace context from the supplied carrier or from Lambda environment variable
+        # If extraction fails, the original context will be returned
+        #
+        # @param [Carrier] carrier The carrier to get the header from
+        # @param [optional Context] context Context to be updated with the trace context
+        #   extracted from the carrier. Defaults to +Context.current+.
+        # @param [optional Getter] getter If the optional getter is provided, it
+        #   will be used to read the header from the carrier, otherwise the default
+        #   text map getter will be used.
+        #
+        # @return [Context] context updated with extracted baggage, or the original context
+        #   if extraction fails
+        def extract(carrier, context: Context.current, getter: Context::Propagation.text_map_getter)
+          # Check if the original input context already has a valid span
+          span_context = Trace.current_span(context).context
+          # If original context is valid, just return it - do not extract from carrier
+          return context if span_context.valid?
+
+          # First try to extract from the carrier using the standard X-Ray propagator
+          xray_context = super
+
+          # Check if we successfully extracted a context from the carrier
+          span_context = Trace.current_span(xray_context).context
+          return xray_context if span_context.valid?
+
+          # If not, check for the Lambda environment variable
+          trace_header = ENV.fetch(AWS_TRACE_HEADER_ENV_KEY, nil)
+          return xray_context unless trace_header
+
+          # Create a carrier with the trace header and extract from it
+          env_carrier = { XRAY_CONTEXT_KEY => trace_header }
+          super(env_carrier, context: xray_context, getter: getter)
+        rescue OpenTelemetry::Error
+          context
+        end
+      end
+    end
+  end
+end

propagator/xray/test/lambda_text_map_propagator_test.rb

@@ -0,0 +1,159 @@
+# frozen_string_literal: true
+
+# Copyright OpenTelemetry Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+require 'test_helper'
+
+describe OpenTelemetry::Propagator::XRay::LambdaTextMapPropagator do
+  Span = OpenTelemetry::Trace::Span
+  SpanContext = OpenTelemetry::Trace::SpanContext
+  TraceFlags = OpenTelemetry::Trace::TraceFlags
+
+  let(:propagator) { OpenTelemetry::Propagator::XRay::LambdaTextMapPropagator.new }
+
+  describe('#extract') do
+    it 'returns the original context when no headers or env vars exist' do
+      parent_context = OpenTelemetry::Context.empty
+      carrier = {}
+
+      # Ensure environment variable is not set
+      original_env = ENV.fetch(OpenTelemetry::Propagator::XRay::LambdaTextMapPropagator::AWS_TRACE_HEADER_ENV_KEY, nil)
+      ENV.delete(OpenTelemetry::Propagator::XRay::LambdaTextMapPropagator::AWS_TRACE_HEADER_ENV_KEY)
+
+      context = propagator.extract(carrier, context: parent_context)
+      extracted_context = OpenTelemetry::Trace.current_span(context).context
+
+      _(extracted_context.valid?).must_equal(false)
+      _(context).must_equal(parent_context)
+
+      # Restore original env value if it existed
+      ENV[OpenTelemetry::Propagator::XRay::LambdaTextMapPropagator::AWS_TRACE_HEADER_ENV_KEY] = original_env if original_env
+    end
+
+    it 'returns existing context when valid and no env var exists' do
+      # Create a valid context
+      valid_context = create_context(
+        trace_id: '80f198eae56343ba864fe8b2a57d3eff',
+        span_id: 'e457b5a2e4d86bd1',
+        trace_flags: TraceFlags::SAMPLED
+      )
+
+      carrier = {}
+
+      # Ensure environment variable is not set
+      original_env = ENV.fetch(OpenTelemetry::Propagator::XRay::LambdaTextMapPropagator::AWS_TRACE_HEADER_ENV_KEY, nil)
+      ENV.delete(OpenTelemetry::Propagator::XRay::LambdaTextMapPropagator::AWS_TRACE_HEADER_ENV_KEY)
+
+      context = propagator.extract(carrier, context: valid_context)
+      extracted_context = OpenTelemetry::Trace.current_span(context).context
+
+      _(extracted_context.hex_trace_id).must_equal('80f198eae56343ba864fe8b2a57d3eff')
+      _(extracted_context.hex_span_id).must_equal('e457b5a2e4d86bd1')
+      _(extracted_context.trace_flags).must_equal(TraceFlags::SAMPLED)
+      _(context).must_equal(valid_context)
+
+      # Restore original env value if it existed
+      ENV[OpenTelemetry::Propagator::XRay::LambdaTextMapPropagator::AWS_TRACE_HEADER_ENV_KEY] = original_env if original_env
+    end
+
+    it 'extracts context from environment variable when no headers exist' do
+      parent_context = OpenTelemetry::Context.empty
+      carrier = {}
+
+      # Set environment variable with trace information
+      original_env = ENV.fetch(OpenTelemetry::Propagator::XRay::LambdaTextMapPropagator::AWS_TRACE_HEADER_ENV_KEY, nil)
+      ENV[OpenTelemetry::Propagator::XRay::LambdaTextMapPropagator::AWS_TRACE_HEADER_ENV_KEY] = 'Root=1-80f198ea-e56343ba864fe8b2a57d3eff;Parent=e457b5a2e4d86bd1;Sampled=1'
+
+      context = propagator.extract(carrier, context: parent_context)
+      extracted_context = OpenTelemetry::Trace.current_span(context).context
+
+      _(extracted_context.hex_trace_id).must_equal('80f198eae56343ba864fe8b2a57d3eff')
+      _(extracted_context.hex_span_id).must_equal('e457b5a2e4d86bd1')
+      _(extracted_context.trace_flags).must_equal(TraceFlags::SAMPLED)
+      _(extracted_context).must_be(:remote?)
+
+      # Restore original env value if it existed
+      ENV[OpenTelemetry::Propagator::XRay::LambdaTextMapPropagator::AWS_TRACE_HEADER_ENV_KEY] = original_env if original_env
+    end
+
+    it 'prioritizes header over environment variable' do
+      parent_context = OpenTelemetry::Context.empty
+      carrier = { 'X-Amzn-Trace-Id' => 'Root=1-90f198ea-f56343ba964fe8b2a67d3eff;Parent=f457b5a2e4d86bd2;Sampled=1' }
+
+      # Set environment variable with different trace information
+      original_env = ENV.fetch(OpenTelemetry::Propagator::XRay::LambdaTextMapPropagator::AWS_TRACE_HEADER_ENV_KEY, nil)
+      ENV[OpenTelemetry::Propagator::XRay::LambdaTextMapPropagator::AWS_TRACE_HEADER_ENV_KEY] = 'Root=1-80f198ea-e56343ba864fe8b2a57d3eff;Parent=e457b5a2e4d86bd1;Sampled=1'
+
+      context = propagator.extract(carrier, context: parent_context)
+      extracted_context = OpenTelemetry::Trace.current_span(context).context
+
+      # Should use the header, not the environment variable
+      _(extracted_context.hex_trace_id).must_equal('90f198eaf56343ba964fe8b2a67d3eff')
+      _(extracted_context.hex_span_id).must_equal('f457b5a2e4d86bd2')
+      _(extracted_context.trace_flags).must_equal(TraceFlags::SAMPLED)
+      _(extracted_context).must_be(:remote?)
+
+      # Restore original env value if it existed
+      ENV[OpenTelemetry::Propagator::XRay::LambdaTextMapPropagator::AWS_TRACE_HEADER_ENV_KEY] = original_env if original_env
+    end
+
+    it 'handles malformed environment variable gracefully' do
+      parent_context = OpenTelemetry::Context.empty
+      carrier = {}
+
+      # Set environment variable with malformed trace information
+      original_env = ENV.fetch(OpenTelemetry::Propagator::XRay::LambdaTextMapPropagator::AWS_TRACE_HEADER_ENV_KEY, nil)
+      ENV[OpenTelemetry::Propagator::XRay::LambdaTextMapPropagator::AWS_TRACE_HEADER_ENV_KEY] = 'NotAValidTraceHeader'
+
+      context = propagator.extract(carrier, context: parent_context)
+
+      # Should return the original context since the env var is invalid
+      _(context).must_equal(parent_context)
+
+      # Restore original env value if it existed
+      ENV[OpenTelemetry::Propagator::XRay::LambdaTextMapPropagator::AWS_TRACE_HEADER_ENV_KEY] = original_env if original_env
+    end
+
+    it 'uses existing context when valid, even if headers exist' do
+      # Create a valid context
+      valid_context = create_context(
+        trace_id: '80f198eae56343ba864fe8b2a57d3eff',
+        span_id: 'e457b5a2e4d86bd1',
+        trace_flags: TraceFlags::SAMPLED
+      )
+
+      # Create a carrier with header
+      carrier = { 'X-Amzn-Trace-Id' => 'Root=1-90f198ea-f56343ba964fe8b2a67d3eff;Parent=f457b5a2e4d86bd2;Sampled=1' }
+
+      context = propagator.extract(carrier, context: valid_context)
+      extracted_context = OpenTelemetry::Trace.current_span(context).context
+
+      # Should use the existing context, not the header
+      _(extracted_context.hex_trace_id).must_equal('80f198eae56343ba864fe8b2a57d3eff')
+      _(extracted_context.hex_span_id).must_equal('e457b5a2e4d86bd1')
+      _(extracted_context.trace_flags).must_equal(TraceFlags::SAMPLED)
+      _(context).must_equal(valid_context)
+    end
+  end
+
+  # Helper method to create a context with specified values
+  def create_context(trace_id:, span_id:, trace_flags: TraceFlags::DEFAULT, xray_debug: false)
+    span_context = SpanContext.new(
+      trace_id: Array(trace_id).pack('H*'),
+      span_id: Array(span_id).pack('H*'),
+      trace_flags: trace_flags,
+      remote: false
+    )
+
+    span = Span.new(
+      span_context: span_context
+    )
+
+    context = OpenTelemetry::Trace.context_with_span(span)
+    return context unless xray_debug
+
+    OpenTelemetry::Propagator::XRay.send(:context_with_debug, context)
+  end
+end

propagator/xray/test/xray_test.rb

@@ -14,5 +14,10 @@
         OpenTelemetry::Propagator::XRay::TextMapPropagator
       )
     end
+
+    it 'provides access to lambda text map propagator' do
+      propagator = OpenTelemetry::Propagator::XRay.lambda_text_map_propagator
+      _(propagator).must_be_instance_of(OpenTelemetry::Propagator::XRay::LambdaTextMapPropagator)
+    end
   end
 end