feat!: Add baggage key predicate func to baggage span processor (#990)

* update test to add multiple baggage items

* add keyfilter constructor parameter

* rework into regex

* add regex example and test

* move new setup & assertions to #on_start test block

Co-authored-by: Mike Goldsmith <[email protected]>

* consistent start_with; no "starts" to confuse readers

Co-authored-by: Mike Goldsmith <[email protected]>

* spec the new new method

Co-authored-by: Mike Goldsmith <[email protected]>

* update doc comments with new usage

* update README with new usage

* bonus: convenience docker compose service for testing baggage processor

* update description to indicate new baggage entries are not added to active span

* fix whitespace

---------

Co-authored-by: Robb Kidd <[email protected]>

Author: MikeGoldsmith

docker-compose.yml

@@ -162,6 +162,12 @@ services:
     command: ./start_server.sh
     working_dir: /app/instrumentation/sinatra/example
 
+  processor-baggage-test:
+    <<: *base
+    working_dir: /app/processor/baggage
+    command: |
+      bash -c "bundle install && rake"
+
   mongo:
     image: mongo:4.4
     expose:

processor/baggage/README.md

@@ -2,7 +2,7 @@
 
 This is an OpenTelemetry [span processor](https://opentelemetry.io/docs/specs/otel/trace/sdk/#span-processor) that reads key/values stored in [Baggage](https://opentelemetry.io/docs/specs/otel/baggage/api/) in the starting span's parent context and adds them as attributes to the span.
 
-Keys and values added to Baggage will appear on all subsequent child spans for a trace within this service *and* will be propagated to external services via propagation headers.
+Keys and values added to Baggage will appear on all subsequent child spans, not the current active span, for a trace within this service *and* will be propagated to external services via propagation headers.
 If the external services also have a Baggage span processor, the keys and values will appear in those child spans as well.
 
 ⚠️ Waning ⚠️
@@ -31,7 +31,7 @@ To install the instrumentation, add the gem to your Gemfile:
 gem 'opentelemetry-processor-baggage'
 ```
 
-Then add the processor to an SDK's configuration:
+Then configure the span processor to copy all baggage entries:
 
 ```ruby
 require 'rubygems'
@@ -40,8 +40,11 @@ require 'bundler/setup'
 Bundler.require
 
 OpenTelemetry::SDK.configure do |c|
-  # Add the BaggageSpanProcessor to the collection of span processors
-  c.add_span_processor(OpenTelemetry::Processor::Baggage::BaggageSpanProcessor.new)
+  # Add the BaggageSpanProcessor to the collection of span processors and
+  # copy all baggage entries
+  c.add_span_processor(OpenTelemetry::Processor::Baggage::BaggageSpanProcessor.new(
+    OpenTelemetry::Processor::Baggage::ALLOW_ALL_BAGGAGE_KEYS
+  ))
 
   # Because the span processor list is no longer empty, the SDK will not use the
   # values in OTEL_TRACES_EXPORTER to instantiate exporters.
@@ -57,6 +60,27 @@ OpenTelemetry::SDK.configure do |c|
 end
 ```
 
+Alternatively, you can provide a custom baggage key predicate to select which baggage keys you want to copy.
+
+For example, to only copy baggage entries that start with `myapp.`:
+
+```ruby
+OUR_BAGGAGE_KEY_PREFIX = 'myapp.'.freeze
+OpenTelemetry::Processor::Baggage::BaggageSpanProcessor.new(
+  # a constant here improves performance
+  ->(baggage_key) { baggage_key.start_with?(OUR_BAGGAGE_KEY_PREFIX) }
+)
+```
+
+For example, to only copy baggage entries that match `myapp.`, `myapp1.` and `myapp42.`:
+
+```ruby
+OUR_BAGGAGE_KEY_MATCHER = /\Amyapp\d*\./
+OpenTelemetry::Processor::Baggage::BaggageSpanProcessor.new(
+  ->(baggage_key) { OUR_BAGGAGE_KEY_MATCHER.match?(baggage_key) }
+)
+```
+
 ## How can I get involved?
 
 The `opentelemetry-processor-baggage` gem source is [on github][repo-github], along with related gems including `opentelemetry-api` and `opentelemetry-sdk`.

processor/baggage/lib/opentelemetry/processor/baggage/baggage_span_processor.rb

@@ -10,24 +10,36 @@
 module OpenTelemetry
   module Processor
     module Baggage
+      # A baggage key predicate that allows all keys to be added to the span as attributes.
+      ALLOW_ALL_BAGGAGE_KEYS = ->(_) { true }
+
       # The BaggageSpanProcessor reads key/values stored in Baggage in the
-      # starting span's parent context and adds them as attributes to the span.
+      # starting span's parent context and adds them as attributes to the span,
+      # if a key matches a provided predicate lambda.
+      #
+      # Keys and values added to Baggage will appear on all subsequent child spans,
+      # not the current active span, for a trace within this service *and* will be
+      # propagated to external services via propagation headers.
       #
-      # Keys and values added to Baggage will appear on all subsequent child spans
-      # for a trace within this service *and* will be propagated to external services
-      # via propagation headers. If the external services also have a Baggage span
-      # processor, the keys and values will appear in those child spans as well.
+      # If the external services also have a Baggage span processor, the keys and
+      # values will appear in those child spans as well.
       #
       # ⚠️
       # To repeat: a consequence of adding data to Baggage is that the keys and
       # values will appear in all outgoing HTTP headers from the application.
       # Do not put sensitive information in Baggage.
       # ⚠️
       #
-      # @example
+      # @example Adding the BaggageSpanProcessor to the SDK, only add attributes for keys that start with 'myapp.'
+      #   OUR_BAGGAGE_KEY_PREFIX = 'myapp.'.freeze
+      #
       #   OpenTelemetry::SDK.configure do |c|
       #     # Add the BaggageSpanProcessor to the collection of span processors
-      #     c.add_span_processor(OpenTelemetry::Processor::Baggage::BaggageSpanProcessor.new)
+      #     c.add_span_processor(
+      #       OpenTelemetry::Processor::Baggage::BaggageSpanProcessor.new(
+      #         ->(key) { key.start_with?(OUR_BAGGAGE_KEY_PREFIX) } # a constant here improves performance
+      #       )
+      #     )
       #
       #     # Because the span processor list is no longer empty, the SDK will not use the
       #     # values in OTEL_TRACES_EXPORTER to instantiate exporters.
@@ -41,7 +53,31 @@ module Baggage
       #       )
       #     )
       #   end
+      #
+      # @example Allow all Baggage keys to be added to the span as attributes
+      #   OpenTelemetry::Processor::Baggage::BaggageSpanProcessor.new(
+      #     # This processor provides a convenience predicate that allows all keys to be added as attributes.
+      #     OpenTelemetry::Processor::Baggage::ALLOW_ALL_BAGGAGE_KEYS
+      #   )
       class BaggageSpanProcessor < OpenTelemetry::SDK::Trace::SpanProcessor
+        # Create a new BaggageSpanProcessor that reads Baggage keys and values from the parent context
+        # and adds them as attributes to the span.
+        #
+        # @param [lambda] baggage_key_predicate A lambda that takes a baggage key [String] and returns true if
+        #  the key should be added to the span as an attribute, false otherwise.
+        #
+        # @example Only add attributes for keys that start with a specific prefix
+        #   OUR_BAGGAGE_KEY_PREFIX = 'myapp.'.freeze
+        #   OpenTelemetry::Processor::Baggage::BaggageSpanProcessor.new(
+        #     ->(key) { key.start_with?(OUR_BAGGAGE_KEY_PREFIX) } # a constant here improves performance
+        #   )
+        def initialize(baggage_key_predicate)
+          raise ArgumentError, 'baggage_key_predicate must respond to :call (lambda/Proc)' unless baggage_key_predicate.respond_to?(:call)
+
+          @baggage_key_predicate = baggage_key_predicate
+          super()
+        end
+
         # Called when a `Span` is started, adds Baggage keys/values to the span as attributes.
         #
         # @param [Span] span the `Span` that just started, expected to conform
@@ -51,7 +87,11 @@ class BaggageSpanProcessor < OpenTelemetry::SDK::Trace::SpanProcessor
         def on_start(span, parent_context)
           return unless span.respond_to?(:add_attributes) && parent_context.is_a?(::OpenTelemetry::Context)
 
-          span.add_attributes(::OpenTelemetry::Baggage.values(context: parent_context))
+          span.add_attributes(
+            ::OpenTelemetry::Baggage
+              .values(context: parent_context)
+              .select { |k, _v| @baggage_key_predicate.call(k) }
+          )
         end
 
         # Called when a Span is ended, does nothing.

processor/baggage/test/opentelemetry/processor/baggage/baggage_span_processor_test.rb

@@ -11,8 +11,10 @@
 TEST_EXPORTER = OpenTelemetry::SDK::Trace::Export::InMemorySpanExporter.new
 
 OpenTelemetry::SDK.configure do |c|
-  # the baggage processor getting wired in for testing
-  c.add_span_processor OpenTelemetry::Processor::Baggage::BaggageSpanProcessor.new
+  # the baggage processor getting wired in for integration testing
+  c.add_span_processor OpenTelemetry::Processor::Baggage::BaggageSpanProcessor.new(
+    OpenTelemetry::Processor::Baggage::ALLOW_ALL_BAGGAGE_KEYS
+  )
 
   # use a simple processor and in-memory export for testing sent spans
   c.add_span_processor(
@@ -24,17 +26,68 @@
 end
 
 describe OpenTelemetry::Processor::Baggage::BaggageSpanProcessor do
-  let(:processor) { OpenTelemetry::Processor::Baggage::BaggageSpanProcessor.new }
+  let(:processor) do
+    OpenTelemetry::Processor::Baggage::BaggageSpanProcessor.new(
+      OpenTelemetry::Processor::Baggage::ALLOW_ALL_BAGGAGE_KEYS
+    )
+  end
   let(:span) { Minitest::Mock.new }
-  let(:context_with_baggage) { OpenTelemetry::Baggage.set_value('a_key', 'a_value') }
+  let(:context_with_baggage) do
+    OpenTelemetry::Baggage.build(context: OpenTelemetry::Context.empty) do |baggage|
+      baggage.set_value('a_key', 'a_value')
+      baggage.set_value('b_key', 'b_value')
+    end
+  end
+
+  describe '#new' do
+    it 'requires a callable baggage_key_predicate' do
+      _(-> { OpenTelemetry::Processor::Baggage::BaggageSpanProcessor.new }).must_raise(ArgumentError)
+      err = _(-> { OpenTelemetry::Processor::Baggage::BaggageSpanProcessor.new(:not_a_callable) }).must_raise(ArgumentError)
+      _(err.message).must_match(/must respond to :call/)
+    end
+  end
 
   describe '#on_start' do
-    it 'adds current baggage keys/values as attributes when a span starts' do
-      span.expect(:add_attributes, span, [{ 'a_key' => 'a_value' }])
+    describe 'with the ALLOW_ALL_BAGGAGE_KEYS predicate' do
+      it 'adds current baggage keys/values as attributes when a span starts' do
+        span.expect(:add_attributes, span, [{ 'a_key' => 'a_value', 'b_key' => 'b_value' }])
+
+        processor.on_start(span, context_with_baggage)
+
+        span.verify
+      end
+    end
+
+    describe 'with a start_with? key predicate' do
+      let(:start_with_processor) do
+        OpenTelemetry::Processor::Baggage::BaggageSpanProcessor.new(
+          ->(baggage_key) { baggage_key.start_with?('a') }
+        )
+      end
+
+      it 'only adds attributes that pass the keyfilter' do
+        span.expect(:add_attributes, span, [{ 'a_key' => 'a_value' }])
+
+        start_with_processor.on_start(span, context_with_baggage)
+
+        span.verify
+      end
+    end
+
+    describe 'with a regex key predicate' do
+      let(:regex_processor) do
+        OpenTelemetry::Processor::Baggage::BaggageSpanProcessor.new(
+          ->(baggage_key) { baggage_key.match?(/^b_ke.+/) }
+        )
+      end
+
+      it 'only adds attributes that pass the keyfilter' do
+        span.expect(:add_attributes, span, [{ 'b_key' => 'b_value' }])
 
-      processor.on_start(span, context_with_baggage)
+        regex_processor.on_start(span, context_with_baggage)
 
-      span.verify
+        span.verify
+      end
     end
 
     it 'does not blow up when given nil context' do
@@ -89,7 +142,7 @@
 
       _(exporter.finished_spans.size).must_equal(1)
       _(exporter.finished_spans.first.name).must_equal('integration test span')
-      _(exporter.finished_spans.first.attributes).must_equal('a_key' => 'a_value')
+      _(exporter.finished_spans.first.attributes).must_equal('a_key' => 'a_value', 'b_key' => 'b_value')
     end
   end
 end