Fix open-telemetry#3364

Signed-off-by: Francesco Guardiani <[email protected]>

context/src/main/java/io/opentelemetry/context/propagation/CompositeTextMapPropagatorBuilder.java

@@ -0,0 +1,69 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.context.propagation;
+
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * A builder for configuring an {@link TextMapPropagator} specifying which propagators should be used for extracting the context and which should be used for injecting the context.
+ */
+public final class CompositeTextMapPropagatorBuilder {
+
+  private final List<TextMapPropagator> extractors;
+  private final List<TextMapPropagator> injectors;
+
+  /**
+   * Package protected to disallow direct initialization.
+   *
+   * @see TextMapPropagator#builder()
+   */
+  CompositeTextMapPropagatorBuilder() {
+    this.extractors = new ArrayList<>();
+    this.injectors = new ArrayList<>();
+  }
+
+  /**
+   * Add a {@link TextMapPropagator} to be used only to extract the context.
+   */
+  public CompositeTextMapPropagatorBuilder extractor(TextMapPropagator propagator) {
+    this.extractors.add(propagator);
+    return this;
+  }
+
+  /**
+   * Add a {@link TextMapPropagator} to be used only to inject the context.
+   */
+  public CompositeTextMapPropagatorBuilder injector(TextMapPropagator propagator) {
+    this.injectors.add(propagator);
+    return this;
+  }
+
+  /**
+   * Add a {@link TextMapPropagator} to be used both to extract and inject the context.
+   */
+  public CompositeTextMapPropagatorBuilder propagator(TextMapPropagator propagator) {
+    this.injectors.add(propagator);
+    this.extractors.add(propagator);
+    return this;
+  }
+
+  /**
+   * Returns the built {@link TextMapPropagator}
+   *
+   * @see CompositeTextMapPropagatorBuilder
+   */
+  public TextMapPropagator build() {
+    if (this.injectors.isEmpty()) {
+      this.injectors.add(TextMapPropagator.noop());
+    }
+    if (this.extractors.isEmpty()) {
+      this.extractors.add(TextMapPropagator.noop());
+    }
+
+    return new MultiTextMapPropagator(this.extractors, this.injectors);
+  }
+}

context/src/main/java/io/opentelemetry/context/propagation/MultiTextMapPropagator.java

@@ -16,17 +16,29 @@
 import javax.annotation.Nullable;
 
 final class MultiTextMapPropagator implements TextMapPropagator {
-  private final TextMapPropagator[] textPropagators;
+  private final TextMapPropagator[] extractors;
+  private final TextMapPropagator[] injectors;
   private final Collection<String> allFields;
 
   MultiTextMapPropagator(TextMapPropagator... textPropagators) {
     this(Arrays.asList(textPropagators));
   }
 
   MultiTextMapPropagator(List<TextMapPropagator> textPropagators) {
-    this.textPropagators = new TextMapPropagator[textPropagators.size()];
-    textPropagators.toArray(this.textPropagators);
-    this.allFields = Collections.unmodifiableList(getAllFields(this.textPropagators));
+    this.extractors = new TextMapPropagator[textPropagators.size()];
+    textPropagators.toArray(this.extractors);
+    this.injectors = this.extractors; // No need to copy
+    this.allFields = Collections.unmodifiableList(getAllFields(this.injectors));
+  }
+
+  MultiTextMapPropagator(List<TextMapPropagator> extractors, List<TextMapPropagator> injectors) {
+    this.extractors = new TextMapPropagator[extractors.size()];
+    extractors.toArray(this.extractors);
+
+    this.injectors = new TextMapPropagator[injectors.size()];
+    injectors.toArray(this.injectors);
+
+    this.allFields = Collections.unmodifiableList(getAllFields(this.injectors));
   }
 
   @Override
@@ -48,7 +60,7 @@ public <C> void inject(Context context, @Nullable C carrier, TextMapSetter<C> se
     if (context == null || setter == null) {
       return;
     }
-    for (TextMapPropagator textPropagator : textPropagators) {
+    for (TextMapPropagator textPropagator : injectors) {
       textPropagator.inject(context, carrier, setter);
     }
   }
@@ -61,7 +73,7 @@ public <C> Context extract(Context context, @Nullable C carrier, TextMapGetter<C
     if (getter == null) {
       return context;
     }
-    for (TextMapPropagator textPropagator : textPropagators) {
+    for (TextMapPropagator textPropagator : extractors) {
       context = textPropagator.extract(context, carrier, getter);
     }
     return context;

context/src/main/java/io/opentelemetry/context/propagation/TextMapPropagator.java

@@ -43,6 +43,13 @@
 @ThreadSafe
 public interface TextMapPropagator {
 
+  /**
+   * Returns a composite text map propagator builder that allows to specify specific propagators as extractors and specific propagators as injectors.
+   */
+  static CompositeTextMapPropagatorBuilder builder() {
+    return new CompositeTextMapPropagatorBuilder();
+  }
+
   /**
    * Returns a {@link TextMapPropagator} which simply delegates injection and extraction to the
    * provided propagators.

context/src/test/java/io/opentelemetry/context/propagation/MultiTextMapPropagatorTest.java

@@ -8,6 +8,7 @@
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
@@ -168,4 +169,32 @@ void inject_nullSetter() {
     new MultiTextMapPropagator(propagator1, propagator2).inject(context, carrier, null);
     assertThat(carrier).isEmpty();
   }
+
+  @Test
+  void different_inject_and_extract() {
+    Map<String, String> carrier = new HashMap<>();
+
+    TextMapPropagator prop = TextMapPropagator.builder()
+        .injector(propagator1)
+        .extractor(propagator2)
+        .propagator(propagator3)
+        .build();
+
+    Context context1 = mock(Context.class);
+    Context context2 = mock(Context.class);
+    Context expectedContext = mock(Context.class);
+
+    when(propagator2.extract(context1, carrier, getter)).thenReturn(context2);
+    when(propagator3.extract(context2, carrier, getter)).thenReturn(expectedContext);
+
+    assertThat(prop.extract(context1, carrier, getter)).isEqualTo(expectedContext);
+
+    Context context = mock(Context.class);
+    TextMapSetter<Map<String, String>> setter = Map::put;
+
+    prop.inject(context, carrier, setter);
+    verify(propagator1).inject(context, carrier, setter);
+    verify(propagator2, never()).inject(context, carrier, setter);
+    verify(propagator3).inject(context, carrier, setter);
+  }
 }