Extends TextMapGetter with GetAll() method, implement usage in W3CBag…

…gagePropagator (#6852)

Co-authored-by: Jack Berg <[email protected]>

api/all/src/main/java/io/opentelemetry/api/baggage/propagation/W3CBaggagePropagator.java

@@ -16,7 +16,9 @@
 import io.opentelemetry.context.propagation.TextMapGetter;
 import io.opentelemetry.context.propagation.TextMapPropagator;
 import io.opentelemetry.context.propagation.TextMapSetter;
+import io.opentelemetry.context.propagation.internal.ExtendedTextMapGetter;
 import java.util.Collection;
+import java.util.Iterator;
 import java.util.List;
 import javax.annotation.Nullable;
 
@@ -95,6 +97,14 @@ public <C> Context extract(Context context, @Nullable C carrier, TextMapGetter<C
       return context;
     }
 
+    if (getter instanceof ExtendedTextMapGetter) {
+      return extractMulti(context, carrier, (ExtendedTextMapGetter<C>) getter);
+    }
+    return extractSingle(context, carrier, getter);
+  }
+
+  private static <C> Context extractSingle(
+      Context context, @Nullable C carrier, TextMapGetter<C> getter) {
     String baggageHeader = getter.get(carrier, FIELD);
     if (baggageHeader == null) {
       return context;
@@ -112,6 +122,33 @@ public <C> Context extract(Context context, @Nullable C carrier, TextMapGetter<C
     return context.with(baggageBuilder.build());
   }
 
+  private static <C> Context extractMulti(
+      Context context, @Nullable C carrier, ExtendedTextMapGetter<C> getter) {
+    Iterator<String> baggageHeaders = getter.getAll(carrier, FIELD);
+    if (baggageHeaders == null) {
+      return context;
+    }
+
+    boolean extracted = false;
+    BaggageBuilder baggageBuilder = Baggage.builder();
+
+    while (baggageHeaders.hasNext()) {
+      String header = baggageHeaders.next();
+      if (header.isEmpty()) {
+        continue;
+      }
+
+      try {
+        extractEntries(header, baggageBuilder);
+        extracted = true;
+      } catch (RuntimeException expected) {
+        // invalid baggage header, continue
+      }
+    }
+
+    return extracted ? context.with(baggageBuilder.build()) : context;
+  }
+
   private static void extractEntries(String baggageHeader, BaggageBuilder baggageBuilder) {
     new Parser(baggageHeader).parseInto(baggageBuilder);
   }

api/all/src/test/java/io/opentelemetry/api/baggage/propagation/W3CBaggagePropagatorTest.java

@@ -8,14 +8,18 @@
 import static java.util.Collections.singletonMap;
 import static org.assertj.core.api.Assertions.assertThat;
 
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import io.opentelemetry.api.baggage.Baggage;
 import io.opentelemetry.api.baggage.BaggageEntryMetadata;
 import io.opentelemetry.context.Context;
 import io.opentelemetry.context.propagation.TextMapGetter;
+import io.opentelemetry.context.propagation.internal.ExtendedTextMapGetter;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.Iterator;
 import java.util.LinkedHashMap;
+import java.util.List;
 import java.util.Map;
 import javax.annotation.Nullable;
 import org.junit.jupiter.api.Test;
@@ -36,6 +40,28 @@ public String get(Map<String, String> carrier, String key) {
         }
       };
 
+  private static final ExtendedTextMapGetter<Map<String, List<String>>> multiGetter =
+      new ExtendedTextMapGetter<Map<String, List<String>>>() {
+        @Override
+        public Iterable<String> keys(Map<String, List<String>> carrier) {
+          return carrier.keySet();
+        }
+
+        @Nullable
+        @Override
+        public String get(Map<String, List<String>> carrier, String key) {
+          return carrier.getOrDefault(key, Collections.emptyList()).stream()
+              .findFirst()
+              .orElse(null);
+        }
+
+        @Override
+        public Iterator<String> getAll(Map<String, List<String>> carrier, String key) {
+          List<String> values = carrier.get(key);
+          return values == null ? Collections.emptyIterator() : values.iterator();
+        }
+      };
+
   @Test
   void fields() {
     assertThat(W3CBaggagePropagator.getInstance().fields()).containsExactly("baggage");
@@ -421,6 +447,101 @@ void extract_nullGetter() {
         .isSameAs(context);
   }
 
+  @Test
+  void extract_multiple_headers() {
+    W3CBaggagePropagator propagator = W3CBaggagePropagator.getInstance();
+
+    Context result =
+        propagator.extract(
+            Context.root(),
+            ImmutableMap.of("baggage", ImmutableList.of("k1=v1", "k2=v2")),
+            multiGetter);
+
+    Baggage expectedBaggage = Baggage.builder().put("k1", "v1").put("k2", "v2").build();
+    assertThat(Baggage.fromContext(result)).isEqualTo(expectedBaggage);
+  }
+
+  @Test
+  void extract_multiple_headers_duplicate_key() {
+    W3CBaggagePropagator propagator = W3CBaggagePropagator.getInstance();
+
+    Context result =
+        propagator.extract(
+            Context.root(),
+            ImmutableMap.of("baggage", ImmutableList.of("k1=v1", "k1=v2")),
+            multiGetter);
+
+    Baggage expectedBaggage = Baggage.builder().put("k1", "v2").build();
+    assertThat(Baggage.fromContext(result)).isEqualTo(expectedBaggage);
+  }
+
+  @Test
+  void extract_multiple_headers_mixed_duplicates_non_duplicates() {
+    W3CBaggagePropagator propagator = W3CBaggagePropagator.getInstance();
+
+    Context result =
+        propagator.extract(
+            Context.root(),
+            ImmutableMap.of("baggage", ImmutableList.of("k1=v1,k2=v0", "k2=v2,k3=v3")),
+            multiGetter);
+
+    Baggage expectedBaggage =
+        Baggage.builder().put("k1", "v1").put("k2", "v2").put("k3", "v3").build();
+    assertThat(Baggage.fromContext(result)).isEqualTo(expectedBaggage);
+  }
+
+  @Test
+  void extract_multiple_headers_all_empty() {
+    W3CBaggagePropagator propagator = W3CBaggagePropagator.getInstance();
+
+    Context result =
+        propagator.extract(
+            Context.root(), ImmutableMap.of("baggage", ImmutableList.of("", "")), multiGetter);
+
+    Baggage expectedBaggage = Baggage.builder().build();
+    assertThat(Baggage.fromContext(result)).isEqualTo(expectedBaggage);
+  }
+
+  @Test
+  void extract_multiple_headers_some_empty() {
+    W3CBaggagePropagator propagator = W3CBaggagePropagator.getInstance();
+
+    Context result =
+        propagator.extract(
+            Context.root(), ImmutableMap.of("baggage", ImmutableList.of("", "k=v")), multiGetter);
+
+    Baggage expectedBaggage = Baggage.builder().put("k", "v").build();
+    assertThat(Baggage.fromContext(result)).isEqualTo(expectedBaggage);
+  }
+
+  @Test
+  void extract_multiple_headers_all_invalid() {
+    W3CBaggagePropagator propagator = W3CBaggagePropagator.getInstance();
+
+    Context result =
+        propagator.extract(
+            Context.root(),
+            ImmutableMap.of("baggage", ImmutableList.of("!@#$%^", "key=va%lue")),
+            multiGetter);
+
+    Baggage expectedBaggage = Baggage.builder().build();
+    assertThat(Baggage.fromContext(result)).isEqualTo(expectedBaggage);
+  }
+
+  @Test
+  void extract_multiple_headers_some_invalid() {
+    W3CBaggagePropagator propagator = W3CBaggagePropagator.getInstance();
+
+    Context result =
+        propagator.extract(
+            Context.root(),
+            ImmutableMap.of("baggage", ImmutableList.of("k1=v1", "key=va%lue", "k2=v2")),
+            multiGetter);
+
+    Baggage expectedBaggage = Baggage.builder().put("k1", "v1").put("k2", "v2").build();
+    assertThat(Baggage.fromContext(result)).isEqualTo(expectedBaggage);
+  }
+
   @Test
   void inject_noBaggage() {
     W3CBaggagePropagator propagator = W3CBaggagePropagator.getInstance();

context/src/main/java/io/opentelemetry/context/propagation/internal/ExtendedTextMapGetter.java

@@ -0,0 +1,44 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.context.propagation.internal;
+
+import io.opentelemetry.context.propagation.TextMapGetter;
+import java.util.Collections;
+import java.util.Iterator;
+import javax.annotation.Nullable;
+
+/**
+ * Extends {@link TextMapGetter} to return possibly multiple values for a given key.
+ *
+ * <p>This class is internal and experimental. Its APIs are unstable and can change at any time. Its
+ * APIs (or a version of them) may be promoted to the public stable API in the future, but no
+ * guarantees are made.
+ *
+ * @param <C> carrier of propagation fields, such as an http request.
+ */
+public interface ExtendedTextMapGetter<C> extends TextMapGetter<C> {
+  /**
+   * If implemented, returns all values for a given {@code key} in order, or returns an empty list.
+   *
+   * <p>The default method returns the first value of the given propagation {@code key} as a
+   * singleton list, or returns an empty list.
+   *
+   * <p>This class is internal and is hence not for public use. Its APIs are unstable and can change
+   * at any time.
+   *
+   * @param carrier carrier of propagation fields, such as an http request.
+   * @param key the key of the field.
+   * @return all values for a given {@code key} in order, or returns an empty list. Default method
+   *     wraps {@code get()} as an {@link Iterator}.
+   */
+  default Iterator<String> getAll(@Nullable C carrier, String key) {
+    String first = get(carrier, key);
+    if (first == null) {
+      return Collections.emptyIterator();
+    }
+    return Collections.singleton(first).iterator();
+  }
+}

context/src/test/java/io/opentelemetry/context/propagation/internal/ExtendedTextMapGetterTest.java

@@ -0,0 +1,70 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.context.propagation.internal;
+
+import static org.assertj.core.api.AssertionsForClassTypes.assertThat;
+
+import com.google.common.collect.ImmutableList;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.Iterator;
+import java.util.List;
+import javax.annotation.Nullable;
+import org.junit.jupiter.api.Test;
+
+class ExtendedTextMapGetterTest {
+
+  final ExtendedTextMapGetter<Void> nullGet =
+      new ExtendedTextMapGetter<Void>() {
+        @Override
+        public Iterable<String> keys(Void carrier) {
+          return ImmutableList.of("key");
+        }
+
+        @Nullable
+        @Override
+        public String get(@Nullable Void carrier, String key) {
+          return null;
+        }
+      };
+
+  final ExtendedTextMapGetter<Void> nonNullGet =
+      new ExtendedTextMapGetter<Void>() {
+        @Override
+        public Iterable<String> keys(Void carrier) {
+          return ImmutableList.of("key");
+        }
+
+        @Override
+        public String get(@Nullable Void carrier, String key) {
+          return "123";
+        }
+      };
+
+  @Test
+  void extendedTextMapGetterdefaultMethod_returnsEmpty() {
+    Iterator<String> result = nullGet.getAll(null, "key");
+    assertThat(result).isNotNull();
+    List<String> values = iterToList(result);
+    assertThat(values).isEqualTo(Collections.emptyList());
+  }
+
+  @Test
+  void extendedTextMapGetterdefaultMethod_returnsSingleVal() {
+    Iterator<String> result = nonNullGet.getAll(null, "key");
+    assertThat(result).isNotNull();
+    List<String> values = iterToList(result);
+    assertThat(values).isEqualTo(Collections.singletonList("123"));
+  }
+
+  private static <T> List<T> iterToList(Iterator<T> iter) {
+    List<T> list = new ArrayList<>();
+    while (iter.hasNext()) {
+      list.add(iter.next());
+    }
+    return list;
+  }
+}