Merge pull request #71 from breedx-splk/add_propagators

Allow user to customize propagators
open-telemetry · Sep 13, 2023 · 71b213a · 71b213a
2 parents 29a171e + e208d04
commit 71b213a
Show file tree

Hide file tree

Showing 2 changed files with 86 additions and 0 deletions.
diff --git a/instrumentation/src/main/java/io/opentelemetry/android/OpenTelemetryRumBuilder.java b/instrumentation/src/main/java/io/opentelemetry/android/OpenTelemetryRumBuilder.java
@@ -5,8 +5,14 @@
 
 package io.opentelemetry.android;
 
+import static java.util.Objects.requireNonNull;
+
 import android.app.Application;
 import io.opentelemetry.android.instrumentation.InstrumentedApplication;
+import io.opentelemetry.api.baggage.propagation.W3CBaggagePropagator;
+import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator;
+import io.opentelemetry.context.propagation.ContextPropagators;
+import io.opentelemetry.context.propagation.TextMapPropagator;
 import io.opentelemetry.sdk.OpenTelemetrySdk;
 import io.opentelemetry.sdk.logs.SdkLoggerProvider;
 import io.opentelemetry.sdk.logs.SdkLoggerProviderBuilder;
@@ -16,9 +22,12 @@
 import io.opentelemetry.sdk.trace.SdkTracerProvider;
 import io.opentelemetry.sdk.trace.SdkTracerProviderBuilder;
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.function.BiFunction;
 import java.util.function.Consumer;
+import java.util.function.Function;
 
 /**
  * A builder of {@link OpenTelemetryRum}. It enabled configuring the OpenTelemetry SDK and disabling
@@ -39,8 +48,18 @@ public final class OpenTelemetryRumBuilder {
             loggerProviderCustomizers = new ArrayList<>();
     private final List<Consumer<InstrumentedApplication>> instrumentationInstallers =
             new ArrayList<>();
+
+    private Function<? super TextMapPropagator, ? extends TextMapPropagator> propagatorCustomizer =
+            (a) -> a;
+
     private Resource resource;
 
+    private static TextMapPropagator buildDefaultPropagator() {
+        Map<Class<? extends TextMapPropagator>, TextMapPropagator> result = new HashMap<>();
+        return TextMapPropagator.composite(
+                W3CTraceContextPropagator.getInstance(), W3CBaggagePropagator.getInstance());
+    }
+
     OpenTelemetryRumBuilder(Application application) {
         this.application = application;
         SessionIdTimeoutHandler timeoutHandler = new SessionIdTimeoutHandler();
@@ -142,6 +161,27 @@ public OpenTelemetryRumBuilder addInstrumentation(
         return this;
     }
 
+    /**
+     * Adds a {@link Function} to invoke with the default {@link TextMapPropagator} to allow
+     * customization. The return value of the {@link BiFunction} will replace the passed-in
+     * argument. To add new propagators, use {@code TextMapPropagator.composite()} with the existing
+     * propagator passed to your function.
+     *
+     * <p>Multiple calls will execute the customizers in order.
+     */
+    public OpenTelemetryRumBuilder addPropagatorCustomizer(
+            Function<? super TextMapPropagator, ? extends TextMapPropagator> propagatorCustomizer) {
+        requireNonNull(propagatorCustomizer, "propagatorCustomizer");
+        Function<? super TextMapPropagator, ? extends TextMapPropagator> existing =
+                this.propagatorCustomizer;
+        this.propagatorCustomizer =
+                propagator -> {
+                    TextMapPropagator result = existing.apply(propagator);
+                    return propagatorCustomizer.apply(result);
+                };
+        return this;
+    }
+
     public SessionId getSessionId() {
         return sessionId;
     }
@@ -161,6 +201,7 @@ public OpenTelemetryRum build() {
                         .setTracerProvider(buildTracerProvider(sessionId, application))
                         .setMeterProvider(buildMeterProvider(application))
                         .setLoggerProvider(buildLoggerProvider(application))
+                        .setPropagators(buildFinalPropagators())
                         .build();
 
         SdkPreconfiguredRumBuilder delegate =
@@ -200,4 +241,9 @@ private SdkLoggerProvider buildLoggerProvider(Application application) {
         }
         return loggerProviderBuilder.build();
     }
+
+    private ContextPropagators buildFinalPropagators() {
+        TextMapPropagator defaultPropagator = buildDefaultPropagator();
+        return ContextPropagators.create(propagatorCustomizer.apply(defaultPropagator));
+    }
 }
diff --git a/instrumentation/src/test/java/io/opentelemetry/android/OpenTelemetryRumBuilderTest.java b/instrumentation/src/test/java/io/opentelemetry/android/OpenTelemetryRumBuilderTest.java
@@ -9,11 +9,16 @@
 import static io.opentelemetry.sdk.testing.assertj.OpenTelemetryAssertions.assertThat;
 import static io.opentelemetry.sdk.testing.assertj.OpenTelemetryAssertions.equalTo;
 import static org.mockito.ArgumentMatchers.isA;
+import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
 import android.app.Activity;
 import android.app.Application;
 import io.opentelemetry.android.instrumentation.ApplicationStateListener;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.propagation.TextMapGetter;
+import io.opentelemetry.context.propagation.TextMapPropagator;
 import io.opentelemetry.sdk.resources.Resource;
 import io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter;
 import io.opentelemetry.sdk.trace.data.SpanData;
@@ -92,4 +97,39 @@ void shouldInstallInstrumentation() {
         activityCallbacksCaptor.getValue().onActivityStopped(activity);
         verify(listener).onApplicationBackgrounded();
     }
+
+    @Test
+    void canAddPropagator() {
+        Context context = Context.root();
+        Object carrier = new Object();
+
+        Context expected = mock(Context.class);
+        TextMapGetter<? super Object> getter = mock(TextMapGetter.class);
+        TextMapPropagator customPropagator = mock(TextMapPropagator.class);
+
+        when(customPropagator.extract(context, carrier, getter)).thenReturn(expected);
+
+        OpenTelemetryRum rum =
+                OpenTelemetryRum.builder(application)
+                        .addPropagatorCustomizer(x -> customPropagator)
+                        .build();
+        Context result =
+                rum.getOpenTelemetry()
+                        .getPropagators()
+                        .getTextMapPropagator()
+                        .extract(context, carrier, getter);
+        assertThat(result).isSameAs(expected);
+    }
+
+    @Test
+    void canSetPropagator() {
+        TextMapPropagator customPropagator = mock(TextMapPropagator.class);
+
+        OpenTelemetryRum rum =
+                OpenTelemetryRum.builder(application)
+                        .addPropagatorCustomizer(x -> customPropagator)
+                        .build();
+        TextMapPropagator result = rum.getOpenTelemetry().getPropagators().getTextMapPropagator();
+        assertThat(result).isSameAs(customPropagator);
+    }
 }