[Extensions] Update BaggageActivityProcessor to use baggage key predi…

…cate (#1816)

Co-authored-by: Piotr Kiełkowicz <pkiekowicz@splunk.com>
Co-authored-by: Cijo Thomas <cithomas@microsoft.com>
Co-authored-by: Tyler Helmuth <12352919+TylerHelmuth@users.noreply.github.com>

src/OpenTelemetry.Extensions/.publicApi/PublicAPI.Unshipped.txt

@@ -8,7 +8,10 @@ OpenTelemetry.Logs.LogToActivityEventConversionOptions.ScopeConverter.get -> Sys
 OpenTelemetry.Logs.LogToActivityEventConversionOptions.ScopeConverter.set -> void
 OpenTelemetry.Logs.LogToActivityEventConversionOptions.StateConverter.get -> System.Action<System.Diagnostics.ActivityTagsCollection!, System.Collections.Generic.IReadOnlyList<System.Collections.Generic.KeyValuePair<string!, object?>>!>!
 OpenTelemetry.Logs.LogToActivityEventConversionOptions.StateConverter.set -> void
+OpenTelemetry.Trace.BaggageActivityProcessor
 OpenTelemetry.Trace.TracerProviderBuilderExtensions
+override OpenTelemetry.Trace.BaggageActivityProcessor.OnStart(System.Diagnostics.Activity! data) -> void
 static Microsoft.Extensions.Logging.OpenTelemetryLoggingExtensions.AttachLogsToActivityEvent(this OpenTelemetry.Logs.OpenTelemetryLoggerOptions! loggerOptions, System.Action<OpenTelemetry.Logs.LogToActivityEventConversionOptions!>? configure = null) -> OpenTelemetry.Logs.OpenTelemetryLoggerOptions!
+static OpenTelemetry.Trace.BaggageActivityProcessor.AllowAllBaggageKeys.get -> System.Predicate<string!>!
 static OpenTelemetry.Trace.TracerProviderBuilderExtensions.AddAutoFlushActivityProcessor(this OpenTelemetry.Trace.TracerProviderBuilder! builder, System.Func<System.Diagnostics.Activity!, bool>! predicate, int timeoutMilliseconds = 10000) -> OpenTelemetry.Trace.TracerProviderBuilder!
-static OpenTelemetry.Trace.TracerProviderBuilderExtensions.AddBaggageActivityProcessor(this OpenTelemetry.Trace.TracerProviderBuilder! builder) -> OpenTelemetry.Trace.TracerProviderBuilder!
+static OpenTelemetry.Trace.TracerProviderBuilderExtensions.AddBaggageActivityProcessor(this OpenTelemetry.Trace.TracerProviderBuilder! builder, System.Predicate<string!>! baggageKeyPredicate) -> OpenTelemetry.Trace.TracerProviderBuilder!

src/OpenTelemetry.Extensions/CHANGELOG.md

@@ -2,6 +2,9 @@
 
 ## Unreleased
 
+* Update BaggageActivityProcessor to require baggage key predicate.
+  ([#1816](https://github.com/open-telemetry/opentelemetry-dotnet-contrib/pull/1816))
+
 * Updated OpenTelemetry core component version(s) to `1.9.0`.
   ([#1888](https://github.com/open-telemetry/opentelemetry-dotnet-contrib/pull/1888))
 

src/OpenTelemetry.Extensions/Internal/OpenTelemetryExtensionsEventSource.cs

@@ -43,4 +43,10 @@ public void LogRecordFilterException(string? categoryName, string exception)
     {
         this.WriteEvent(2, categoryName, exception);
     }
+
+    [Event(3, Message = "Baggage key predicate threw exeption when trying to add baggage entry with key '{0}'. Baggage entry will not be added to the activity. Exception: '{1}'", Level = EventLevel.Warning)]
+    public void BaggageKeyPredicateException(string baggageKey, string exception)
+    {
+        this.WriteEvent(3, baggageKey, exception);
+    }
 }

src/OpenTelemetry.Extensions/README.md

@@ -46,10 +46,36 @@ and adds the baggage keys and values to the `Activity` as tags (attributes) on s
 
 Add this activity processor to a tracer provider.
 
-Example of adding BaggageActivityProcessor to `TracerProvider`:
+For example, to add all baggage entries to new activities:
 
 ```cs
 var tracerProvider = Sdk.CreateTracerProviderBuilder()
-    .AddBaggageActivityProcessor()
+    .AddBaggageActivityProcessor(BaggageActivityProcessor.AllowAllBaggageKeys)
     .Build();
 ```
+
+Alternatively, you can select which baggage keys you want to copy using a
+custom predicate function.
+
+For example, to only copy baggage entries where the key start with `my-key`
+using a custom function:
+
+```cs
+var tracerProvider = Sdk.CreateTracerProviderBuilder()
+  .AddBaggageActivityProcessor((baggageKey) => baggageKey.StartWith("my-key", System.StringComparison.Ordinal))
+  .Build();
+```
+
+For example, to only copy baggage entries where the key matches the regular
+expression `^my-key`:
+
+```cs
+var baggageKeyRegex = new Regex("^mykey", RegexOptions.Compiled);
+var tracerProvider = Sdk.CreateTracerProviderBuilder()
+  .AddBaggageActivityProcessor((baggageKey) => baggageKeyRegex.IsMatch(baggageKey))
+  .Build();
+```
+
+Warning: The baggage key predicate is executed for every baggage entry for each
+started activity.
+Do not use slow or intensive operations.

src/OpenTelemetry.Extensions/Trace/BaggageActivityProcessor.cs

@@ -6,16 +6,35 @@
 namespace OpenTelemetry.Trace;
 
 /// <summary>
-/// Activity processor that adds <see cref="Baggage"/> fields to every new span.
+/// Activity processor that adds <see cref="Baggage"/> fields to every new activity.
 /// </summary>
-internal sealed class BaggageActivityProcessor : BaseProcessor<Activity>
+public sealed class BaggageActivityProcessor : BaseProcessor<Activity>
 {
+    private readonly Predicate<string> baggageKeyPredicate;
+
+    /// <summary>
+    /// Initializes a new instance of the <see cref="BaggageActivityProcessor"/> class.
+    /// </summary>
+    /// <param name="baggageKeyPredicate">Predicate to determine which baggage keys should be added to the activity.</param>
+    internal BaggageActivityProcessor(Predicate<string> baggageKeyPredicate)
+    {
+        this.baggageKeyPredicate = baggageKeyPredicate ?? throw new ArgumentNullException(nameof(baggageKeyPredicate));
+    }
+
+    /// <summary>
+    /// Gets a baggage key predicate that returns <c>true</c> for all baggage keys.
+    /// </summary>
+    public static Predicate<string> AllowAllBaggageKeys => (_) => true;
+
     /// <inheritdoc />
-    public override void OnStart(Activity activity)
+    public override void OnStart(Activity data)
     {
         foreach (var entry in Baggage.Current)
         {
-            activity.SetTag(entry.Key, entry.Value);
+            if (this.baggageKeyPredicate(entry.Key))
+            {
+                data?.SetTag(entry.Key, entry.Value);
+            }
         }
     }
 }

src/OpenTelemetry.Extensions/TracerProviderBuilderExtensions.cs

@@ -51,14 +51,26 @@ public static TracerProviderBuilder AddAutoFlushActivityProcessor(
     /// Adds the <see cref="BaggageActivityProcessor"/> to the <see cref="TracerProviderBuilder"/>.
     /// </summary>
     /// <param name="builder"><see cref="TracerProviderBuilder"/> to add the <see cref="BaggageActivityProcessor"/> to.</param>
+    /// <param name="baggageKeyPredicate">Predicate to determine which baggage keys should be added to the activity.</param>
     /// <returns>The instance of <see cref="TracerProviderBuilder"/> to chain the calls.</returns>
     public static TracerProviderBuilder AddBaggageActivityProcessor(
-        this TracerProviderBuilder builder)
+        this TracerProviderBuilder builder,
+        Predicate<string> baggageKeyPredicate)
     {
         Guard.ThrowIfNull(builder);
+        Guard.ThrowIfNull(baggageKeyPredicate);
 
-#pragma warning disable CA2000 // Dispose objects before losing scope
-        return builder.AddProcessor(new BaggageActivityProcessor());
-#pragma warning restore CA2000 // Dispose objects before losing scope
+        return builder.AddProcessor(b => new BaggageActivityProcessor(baggageKey =>
+        {
+            try
+            {
+                return baggageKeyPredicate(baggageKey);
+            }
+            catch (Exception exception)
+            {
+                OpenTelemetryExtensionsEventSource.Log.BaggageKeyPredicateException(baggageKey, exception.Message);
+                return false;
+            }
+        }));
     }
 }

test/OpenTelemetry.Extensions.Tests/Trace/BaggageActivityProcessorTests.cs

@@ -0,0 +1,139 @@
+// Copyright The OpenTelemetry Authors
+// SPDX-License-Identifier: Apache-2.0
+
+using System.Diagnostics;
+using System.Runtime.CompilerServices;
+using System.Text.RegularExpressions;
+using OpenTelemetry.Trace;
+using Xunit;
+
+namespace OpenTelemetry.Extensions.Tests.Trace;
+
+public class BaggageActivityProcessorTests
+{
+    [Fact]
+    public void BaggageActivityProcessor_CanAddAllowAllBaggageKeysPredicate()
+    {
+        var sourceName = GetTestMethodName();
+
+        using var provider = Sdk.CreateTracerProviderBuilder()
+            .AddBaggageActivityProcessor(BaggageActivityProcessor.AllowAllBaggageKeys)
+            .AddSource(sourceName)
+            .Build();
+
+        Baggage.SetBaggage("key", "value");
+        Baggage.SetBaggage("other_key", "other_value");
+
+        using var source = new ActivitySource(sourceName);
+        using var activity = source.StartActivity("name", ActivityKind.Server);
+        Assert.NotNull(activity);
+        activity.Stop();
+
+        Assert.Contains(activity.Tags, kv => kv.Key == "key" && kv.Value == "value");
+        Assert.Contains(activity.Tags, kv => kv.Key == "other_key" && kv.Value == "other_value");
+    }
+
+    [Fact]
+    public void BaggageActivityProcessor_CanUseCustomPredicate()
+    {
+        var sourceName = GetTestMethodName();
+
+        using var provider = Sdk.CreateTracerProviderBuilder()
+            .AddBaggageActivityProcessor((baggageKey) => baggageKey.StartsWith("key", StringComparison.Ordinal))
+            .AddSource(sourceName)
+            .Build();
+
+        Baggage.SetBaggage("key", "value");
+        Baggage.SetBaggage("other_key", "other_value");
+
+        using var source = new ActivitySource(sourceName);
+        using var activity = source.StartActivity("name", ActivityKind.Client);
+        Assert.NotNull(activity);
+        activity.Stop();
+
+        Assert.Contains(activity.Tags, kv => kv.Key == "key" && kv.Value == "value");
+        Assert.DoesNotContain(activity.Tags, kv => kv.Key == "other_key" && kv.Value == "other_value");
+    }
+
+    [Fact]
+    public void BaggageActivityProcessor_CanUseRegex()
+    {
+        var sourceName = GetTestMethodName();
+
+        var regex = new Regex("^mykey", RegexOptions.Compiled);
+        using var provider = Sdk.CreateTracerProviderBuilder()
+            .AddBaggageActivityProcessor((baggageKey) => regex.IsMatch(baggageKey))
+            .AddSource(sourceName)
+            .Build();
+
+        Baggage.SetBaggage("mykey", "value");
+        Baggage.SetBaggage("other_key", "other_value");
+
+        using var source = new ActivitySource(sourceName);
+        using var activity = source.StartActivity("name", ActivityKind.Client);
+        Assert.NotNull(activity);
+        activity.Stop();
+
+        Assert.Contains(activity.Tags, kv => kv.Key == "mykey" && kv.Value == "value");
+        Assert.DoesNotContain(activity.Tags, kv => kv.Key == "other_key" && kv.Value == "other_value");
+    }
+
+    [Fact]
+    public void BaggageActivityProcessor_PredicateThrows_DoesNothing()
+    {
+        var sourceName = GetTestMethodName();
+
+        using var provider = Sdk.CreateTracerProviderBuilder()
+            .AddBaggageActivityProcessor(_ => throw new Exception("Predicate throws an exception."))
+            .AddSource(sourceName)
+            .Build();
+
+        Baggage.SetBaggage("key", "value");
+
+        using var source = new ActivitySource(sourceName);
+        using var activity = source.StartActivity("name", ActivityKind.Server);
+        Assert.NotNull(activity);
+        activity.Stop();
+
+        Assert.DoesNotContain(activity.Tags, kv => kv.Key == "key" && kv.Value == "value");
+    }
+
+    [Fact]
+    public void BaggageActivityProcessor_PredicateThrows_OnlyDropsEntriesThatThrow()
+    {
+        var sourceName = GetTestMethodName();
+
+        // First call to predicate should not throw, second call should.
+        using var provider = Sdk.CreateTracerProviderBuilder()
+            .AddBaggageActivityProcessor(key =>
+            {
+                if (key == "key")
+                {
+                    throw new Exception("Predicate throws an exception.");
+                }
+
+                return true;
+            })
+            .AddSource(sourceName)
+            .Build();
+
+        Baggage.SetBaggage("key", "value");
+        Baggage.SetBaggage("other_key", "other_value");
+        Baggage.SetBaggage("another_key", "another_value");
+
+        using var source = new ActivitySource(sourceName);
+        using var activity = source.StartActivity("name", ActivityKind.Server);
+        Assert.NotNull(activity);
+        activity.Stop();
+
+        // Only keys that do not throw should be added.
+        Assert.DoesNotContain(activity.Tags, kv => kv.Key == "key" && kv.Value == "value");
+        Assert.Contains(activity.Tags, kv => kv.Key == "other_key" && kv.Value == "other_value");
+        Assert.Contains(activity.Tags, kv => kv.Key == "another_key" && kv.Value == "another_value");
+    }
+
+    private static string GetTestMethodName([CallerMemberName] string callingMethodName = "")
+    {
+        return callingMethodName;
+    }
+}