fix: rack event baggage handling

instrumentation/rack/lib/opentelemetry/instrumentation/rack/middlewares/event_handler.rb

@@ -42,7 +42,7 @@ module Middlewares
         class EventHandler
           include ::Rack::Events::Abstract
 
-          TOKENS_KEY = 'otel.context.tokens'
+          TOKEN_KEY = 'otel.context.token'
           GOOD_HTTP_STATUSES = (100..499)
 
           # Creates a server span for this current request using the incoming parent context
@@ -56,7 +56,9 @@ def on_start(request, _)
 
             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[TOKEN_KEY] = OpenTelemetry::Context.attach(rack_ctx)
           rescue StandardError => e
             OpenTelemetry.handle_error(exception: e)
           end
@@ -108,7 +110,7 @@ def on_finish(request, response)
           rescue StandardError => e
             OpenTelemetry.handle_error(exception: e)
           ensure
-            detach_contexts(request)
+            detach_context(request)
           end
 
           private
@@ -191,11 +193,11 @@ def request_span_attributes(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[TOKEN_KEY]
+
+            OpenTelemetry::Trace.current_span.finish
+            OpenTelemetry::Context.detach(request.env[TOKEN_KEY])
           rescue StandardError => e
             OpenTelemetry.handle_error(exception: e)
           end
@@ -244,15 +246,6 @@ 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),

instrumentation/rack/test/opentelemetry/instrumentation/rack/middlewares/event_handler_test.rb

@@ -84,6 +84,25 @@
       _(proxy_event).must_be_nil
     end
 
+    describe 'when baggage is set' do
+      let(:headers) do
+        Hash(
+          'baggage' => 'foo=123'
+        )
+      end
+
+      let(:service) do
+        lambda do |_env|
+          _(OpenTelemetry::Baggage.raw_entries['foo'].value).must_equal('123')
+          [200, { 'Content-Type' => 'text/plain' }, response_body]
+        end
+      end
+
+      it 'sets baggage in the request context' do
+        _(rack_span.name).must_equal 'HTTP GET'
+      end
+    end
+
     describe 'when a query is passed in' do
       let(:uri) { '/endpoint?query=true' }