event_handler.rb in opentelemetry-instrumentation-rack-0.24.5

- old
+ new

@@ -40,25 +40,30 @@
         # @see Rack::Events
         # @see OpenTelemetry::Instrumentation::Rack.current_span
         class EventHandler
           include ::Rack::Events::Abstract
 
-          TOKENS_KEY = 'otel.context.tokens'
+          OTEL_TOKEN_AND_SPAN = 'otel.rack.token_and_span'
           GOOD_HTTP_STATUSES = (100..499)
 
           # Creates a server span for this current request using the incoming parent context
           # and registers them as the {current_span}
           #
           # @param [Rack::Request] The current HTTP request
           # @param [Rack::Response] This is nil in practice
           # @return [void]
           def on_start(request, _)
-            return if untraced_request?(request.env)
+            parent_context = if untraced_request?(request.env)
+                               extract_remote_context(request, OpenTelemetry::Common::Utilities.untraced)
+                             else
+                               extract_remote_context(request)
+                             end
 
-            parent_context = extract_remote_context(request)
             span = create_span(parent_context, request)
-            request.env[TOKENS_KEY] = register_current_span(span)
+            span_ctx = OpenTelemetry::Trace.context_with_span(span, parent_context: parent_context)
+            rack_ctx = OpenTelemetry::Instrumentation::Rack.context_with_span(span, parent_context: span_ctx)
+            request.env[OTEL_TOKEN_AND_SPAN] = [OpenTelemetry::Context.attach(rack_ctx), span]
           rescue StandardError => e
             OpenTelemetry.handle_error(exception: e)
           end
 
           # Optionally adds debugging response headers injected from {response_propagators}
@@ -106,11 +111,11 @@
 
             add_response_attributes(span, response) if response
           rescue StandardError => e
             OpenTelemetry.handle_error(exception: e)
           ensure
-            detach_contexts(request)
+            detach_context(request)
           end
 
           private
 
           EMPTY_HASH = {}.freeze
@@ -169,13 +174,14 @@
             else
               "HTTP #{request.request_method}"
             end
           end
 
-          def extract_remote_context(request)
+          def extract_remote_context(request, context = Context.current)
             OpenTelemetry.propagation.extract(
               request.env,
+              context: context,
               getter: OpenTelemetry::Common::Propagation.rack_env_getter
             )
           end
 
           def request_span_attributes(env)
@@ -189,15 +195,16 @@
             attributes['http.user_agent'] = env['HTTP_USER_AGENT'] if env['HTTP_USER_AGENT']
             attributes.merge!(extract_request_headers(env))
             attributes
           end
 
-          def detach_contexts(request)
-            request.env[TOKENS_KEY]&.reverse_each do |token|
-              OpenTelemetry::Context.detach(token)
-              OpenTelemetry::Trace.current_span.finish
-            end
+          def detach_context(request)
+            return nil unless request.env[OTEL_TOKEN_AND_SPAN]
+
+            token, span = request.env[OTEL_TOKEN_AND_SPAN]
+            span.finish
+            OpenTelemetry::Context.detach(token)
           rescue StandardError => e
             OpenTelemetry.handle_error(exception: e)
           end
 
           def add_response_attributes(span, response)
@@ -240,18 +247,9 @@
             OpenTelemetry::Instrumentation::Rack::Instrumentation.instance.tracer
           end
 
           def config
             OpenTelemetry::Instrumentation::Rack::Instrumentation.instance.config
-          end
-
-          def register_current_span(span)
-            ctx = OpenTelemetry::Trace.context_with_span(span)
-            rack_ctx = OpenTelemetry::Instrumentation::Rack.context_with_span(span, parent_context: ctx)
-
-            contexts = [ctx, rack_ctx]
-            contexts.compact!
-            contexts.map { |context| OpenTelemetry::Context.attach(context) }
           end
 
           def create_span(parent_context, request)
             span = tracer.start_span(
               create_request_span_name(request),