Review suggestions

apm-agent-core/src/main/java/co/elastic/apm/agent/impl/ElasticApmTracer.java

@@ -222,7 +222,7 @@ public <C> Transaction startChildTransaction(@Nullable C headerCarrier, TextHead
         if (!coreConfiguration.isActive()) {
             transaction = noopTransaction();
         } else {
-            transaction = createTransaction().start(TraceContext.getFromTraceContextTextHeaders(), headerCarrier,
+            transaction = createTransaction().start(TraceContext.<C>getFromTraceContextTextHeaders(), headerCarrier,
                 textHeadersGetter, epochMicros, sampler, initiatingClassLoader);
         }
         afterTransactionStart(initiatingClassLoader, transaction);
@@ -260,7 +260,7 @@ public <C> Transaction startChildTransaction(@Nullable C headerCarrier, BinaryHe
         if (!coreConfiguration.isActive()) {
             transaction = noopTransaction();
         } else {
-            transaction = createTransaction().start(TraceContext.getFromTraceContextBinaryHeaders(), headerCarrier,
+            transaction = createTransaction().start(TraceContext.<C>getFromTraceContextBinaryHeaders(), headerCarrier,
                 binaryHeadersGetter, epochMicros, sampler, initiatingClassLoader);
         }
         afterTransactionStart(initiatingClassLoader, transaction);

apm-agent-core/src/main/java/co/elastic/apm/agent/impl/transaction/AbstractHeaderGetter.java

@@ -0,0 +1,35 @@
+/*-
+ * #%L
+ * Elastic APM Java agent
+ * %%
+ * Copyright (C) 2018 - 2020 Elastic and contributors
+ * %%
+ * Licensed to Elasticsearch B.V. under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch B.V. licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ * 
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ * 
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ * #L%
+ */
+package co.elastic.apm.agent.impl.transaction;
+
+public abstract class AbstractHeaderGetter<T, C> implements HeaderGetter<T, C> {
+    @Override
+    public <S> void forEach(String headerName, C carrier, S state, HeaderConsumer<T, S> consumer) {
+        T firstHeader = getFirstHeader(headerName, carrier);
+        if (firstHeader != null) {
+            consumer.accept(firstHeader, state);
+        }
+    }
+}

apm-agent-core/src/main/java/co/elastic/apm/agent/impl/transaction/HeaderGetter.java

@@ -31,6 +31,23 @@ public interface HeaderGetter<T, C> {
     @Nullable
     T getFirstHeader(String headerName, C carrier);
 
-    @Nullable
-    Iterable<T> getHeaders(String headerName, C carrier);
+    /**
+     * Calls the consumer for each header value with the given key
+     * until all entries have been processed or the action throws an exception.
+     * <p>
+     * The third parameter lets callers pass in a stateful object to be modified with header values,
+     * so the {@link HeaderConsumer} implementation itself can be stateless and potentially reusable.
+     * </p>
+     *
+     * @param headerName the name of the header
+     * @param carrier    the object containing the headers
+     * @param state      the object to be passed as the second parameter to each invocation on the specified consumer
+     * @param consumer   the action to be performed for each header value
+     * @param <S>        the type of the state object
+     */
+    <S> void forEach(String headerName, C carrier, S state, HeaderConsumer<T, S> consumer);
+
+    interface HeaderConsumer<T, S> {
+        void accept(T headerValue, S state);
+    }
 }

apm-agent-core/src/main/java/co/elastic/apm/agent/impl/transaction/TraceContext.java

@@ -104,38 +104,36 @@ public boolean asChildOf(TraceContext child, TraceContextHolder<?> parent) {
             return true;
         }
     };
-    private static final ChildContextCreatorTwoArg<Object, TextHeaderGetter<?>> FROM_TRACE_CONTEXT_TEXT_HEADERS =
-        new ChildContextCreatorTwoArg<Object, TextHeaderGetter<?>>() {
+    private static final ChildContextCreatorTwoArg FROM_TRACE_CONTEXT_TEXT_HEADERS =
+        new ChildContextCreatorTwoArg<Object, TextHeaderGetter<Object>>() {
             @Override
-            public boolean asChildOf(TraceContext child, @Nullable Object carrier, TextHeaderGetter<?> traceContextHeaderGetter) {
+            public boolean asChildOf(TraceContext child, @Nullable Object carrier, TextHeaderGetter<Object> traceContextHeaderGetter) {
                 if (carrier == null) {
                     return false;
                 }
                 // This cast is required, otherwise the compiler can't guarantee that the carrier type is allowed by the
                 // TextHeaderGetter (which have a runtime inferred type of ?). The caller must ensure anyway that the
                 // carrier type is the one expected in the provided TextHeaderGetter, we cannot enforce it through generics
                 // if we want to use this single ChildContextCreatorTwoArg instance for all types.
-                //noinspection unchecked
-                String traceparent = ((TextHeaderGetter<Object>) traceContextHeaderGetter).getFirstHeader(TRACE_PARENT_TEXTUAL_HEADER_NAME, carrier);
+                String traceparent = traceContextHeaderGetter.getFirstHeader(TRACE_PARENT_TEXTUAL_HEADER_NAME, carrier);
                 if (traceparent != null) {
                     return child.asChildOf(traceparent);
                 }
                 return false;
             }
         };
-    private static final ChildContextCreatorTwoArg<Object, BinaryHeaderGetter<?>> FROM_TRACE_CONTEXT_BINARY_HEADERS =
-        new ChildContextCreatorTwoArg<Object, BinaryHeaderGetter<?>>() {
+    private static final ChildContextCreatorTwoArg FROM_TRACE_CONTEXT_BINARY_HEADERS =
+        new ChildContextCreatorTwoArg<Object, BinaryHeaderGetter<Object>>() {
             @Override
-            public boolean asChildOf(TraceContext child, @Nullable Object carrier, BinaryHeaderGetter<?> traceContextHeaderGetter) {
+            public boolean asChildOf(TraceContext child, @Nullable Object carrier, BinaryHeaderGetter<Object> traceContextHeaderGetter) {
                 if (carrier == null) {
                     return false;
                 }
                 // This cast is required, otherwise the compiler can't guarantee that the carrier type is allowed by the
                 // BinaryHeaderGetter (which have a runtime inferred type of ?). The caller must ensure anyway that the
                 // carrier type is the one expected in the provided BinaryHeaderGetter, we cannot enforce it through generics
                 // if we want to use this single ChildContextCreatorTwoArg instance for all types.
-                //noinspection unchecked
-                byte[] traceparent = ((BinaryHeaderGetter<Object>) traceContextHeaderGetter).getFirstHeader(TRACE_PARENT_BINARY_HEADER_NAME, carrier);
+                byte[] traceparent = traceContextHeaderGetter.getFirstHeader(TRACE_PARENT_BINARY_HEADER_NAME, carrier);
                 if (traceparent != null) {
                     return child.asChildOf(traceparent);
                 }
@@ -225,12 +223,12 @@ public static TraceContext with128BitId(ElasticApmTracer tracer) {
         return new TraceContext(tracer, Id.new128BitId());
     }
 
-    public static ChildContextCreatorTwoArg<Object, TextHeaderGetter<?>> getFromTraceContextTextHeaders() {
-        return FROM_TRACE_CONTEXT_TEXT_HEADERS;
+    public static <C> ChildContextCreatorTwoArg<C, TextHeaderGetter<C>> getFromTraceContextTextHeaders() {
+        return (ChildContextCreatorTwoArg<C, TextHeaderGetter<C>>) FROM_TRACE_CONTEXT_TEXT_HEADERS;
     }
 
-    public static ChildContextCreatorTwoArg<Object, BinaryHeaderGetter<?>> getFromTraceContextBinaryHeaders() {
-        return FROM_TRACE_CONTEXT_BINARY_HEADERS;
+    public static <C> ChildContextCreatorTwoArg<C, BinaryHeaderGetter<C>> getFromTraceContextBinaryHeaders() {
+        return (ChildContextCreatorTwoArg<C, BinaryHeaderGetter<C>>) FROM_TRACE_CONTEXT_BINARY_HEADERS;
     }
 
     public static ChildContextCreator<ElasticApmTracer> fromActive() {

apm-agent-core/src/test/java/co/elastic/apm/agent/impl/BinaryHeaderMapAccessor.java

@@ -30,7 +30,6 @@
 
 import javax.annotation.Nullable;
 import java.util.HashMap;
-import java.util.List;
 import java.util.Map;
 
 public class BinaryHeaderMapAccessor implements BinaryHeaderGetter<Map<String, byte[]>>, BinaryHeaderSetter<Map<String, byte[]>> {
@@ -48,10 +47,12 @@ public byte[] getFirstHeader(String headerName, Map<String, byte[]> headerMap) {
         return headerMap.get(headerName);
     }
 
-    @Nullable
     @Override
-    public Iterable<byte[]> getHeaders(String headerName, Map<String, byte[]> headerMap) {
-        return List.of(headerMap.get(headerName));
+    public <S> void forEach(String headerName, Map<String, byte[]> carrier, S state, HeaderConsumer<byte[], S> consumer) {
+        byte[] headerValue = carrier.get(headerName);
+        if (headerValue != null) {
+            consumer.accept(headerValue, state);
+        }
     }
 
     @Nullable

apm-agent-core/src/test/java/co/elastic/apm/agent/impl/TextHeaderMapAccessor.java

@@ -24,14 +24,14 @@
  */
 package co.elastic.apm.agent.impl;
 
+import co.elastic.apm.agent.impl.transaction.AbstractHeaderGetter;
 import co.elastic.apm.agent.impl.transaction.TextHeaderGetter;
 import co.elastic.apm.agent.impl.transaction.TextHeaderSetter;
 
 import javax.annotation.Nullable;
-import java.util.List;
 import java.util.Map;
 
-public class TextHeaderMapAccessor implements TextHeaderGetter<Map<String, String>>, TextHeaderSetter<Map<String, String>> {
+public class TextHeaderMapAccessor extends AbstractHeaderGetter<String, Map<String, String>> implements TextHeaderGetter<Map<String, String>>, TextHeaderSetter<Map<String, String>> {
 
     public static final TextHeaderMapAccessor INSTANCE = new TextHeaderMapAccessor();
 
@@ -44,16 +44,6 @@ public String getFirstHeader(String headerName, Map<String, String> headerMap) {
         return headerMap.get(headerName);
     }
 
-    @Nullable
-    @Override
-    public Iterable<String> getHeaders(String headerName, Map<String, String> headerMap) {
-        String value = headerMap.get(headerName);
-        if (value != null) {
-            return List.of(value);
-        }
-        return null;
-    }
-
     @Override
     public void setHeader(String headerName, String headerValue, Map<String, String> headerMap) {
         headerMap.put(headerName, headerValue);

apm-agent-core/src/test/java/co/elastic/apm/agent/impl/transaction/TraceContextTest.java

@@ -24,8 +24,8 @@
  */
 package co.elastic.apm.agent.impl.transaction;
 
-import co.elastic.apm.agent.impl.ElasticApmTracer;
 import co.elastic.apm.agent.impl.BinaryHeaderMapAccessor;
+import co.elastic.apm.agent.impl.ElasticApmTracer;
 import co.elastic.apm.agent.impl.TextHeaderMapAccessor;
 import co.elastic.apm.agent.impl.sampling.ConstantSampler;
 import co.elastic.apm.agent.util.HexUtils;
@@ -55,7 +55,7 @@ class TraceContextTest {
     private void mixTextAndBinaryParsingAndContextPropagation(String flagsValue, boolean isSampled) {
         Map<String, String> textHeaderMap = Map.of(TraceContext.TRACE_PARENT_TEXTUAL_HEADER_NAME, "00-0af7651916cd43dd8448eb211c80319c-b9c7c989f97918e1-" + flagsValue);
         final TraceContext child = TraceContext.with64BitId(mock(ElasticApmTracer.class));
-        assertThat(TraceContext.getFromTraceContextTextHeaders().asChildOf(child, textHeaderMap, TextHeaderMapAccessor.INSTANCE)).isTrue();
+        assertThat(TraceContext.<Map<String, String>>getFromTraceContextTextHeaders().asChildOf(child, textHeaderMap, TextHeaderMapAccessor.INSTANCE)).isTrue();
         assertThat(child.getTraceContext().getTraceId().toString()).isEqualTo("0af7651916cd43dd8448eb211c80319c");
         assertThat(child.getTraceContext().getParentId().toString()).isEqualTo("b9c7c989f97918e1");
         assertThat(child.getTraceContext().getId()).isNotEqualTo(child.getTraceContext().getParentId());
@@ -65,7 +65,7 @@ private void mixTextAndBinaryParsingAndContextPropagation(String flagsValue, boo
         final TraceContext grandchild1 = TraceContext.with64BitId(mock(ElasticApmTracer.class));
         final Map<String, byte[]> binaryHeaderMap = new HashMap<>();
         assertThat(child.setOutgoingTraceContextHeaders(binaryHeaderMap, BinaryHeaderMapAccessor.INSTANCE)).isTrue();
-        assertThat(TraceContext.getFromTraceContextBinaryHeaders().asChildOf(grandchild1, binaryHeaderMap, BinaryHeaderMapAccessor.INSTANCE)).isTrue();
+        assertThat(TraceContext.<Map<String, byte[]>>getFromTraceContextBinaryHeaders().asChildOf(grandchild1, binaryHeaderMap, BinaryHeaderMapAccessor.INSTANCE)).isTrue();
         assertThat(grandchild1.getTraceContext().getTraceId().toString()).isEqualTo("0af7651916cd43dd8448eb211c80319c");
         assertThat(grandchild1.getTraceContext().getParentId().toString()).isEqualTo(child.getTraceContext().getId().toString());
         assertThat(grandchild1.getTraceContext().getId()).isNotEqualTo(child.getTraceContext().getId());
@@ -100,7 +100,7 @@ void parseFromTraceParentHeaderUnsupportedFlag() {
     void testBinaryHeaderSizeEnforcement() {
         final Map<String, String> headerMap = Map.of(TraceContext.TRACE_PARENT_TEXTUAL_HEADER_NAME, "00-0af7651916cd43dd8448eb211c80319c-b9c7c989f97918e1-01");
         final TraceContext child = TraceContext.with64BitId(mock(ElasticApmTracer.class));
-        assertThat(TraceContext.getFromTraceContextTextHeaders().asChildOf(child, headerMap, TextHeaderMapAccessor.INSTANCE)).isTrue();
+        assertThat(TraceContext.<Map<String, String>>getFromTraceContextTextHeaders().asChildOf(child, headerMap, TextHeaderMapAccessor.INSTANCE)).isTrue();
         final byte[] outgoingBinaryHeader = new byte[TraceContext.BINARY_FORMAT_EXPECTED_LENGTH - 1];
         assertThat(child.setOutgoingTraceContextHeaders(new HashMap<>(), new BinaryHeaderSetter<Map<String, byte[]>>() {
             @Override
@@ -120,7 +120,7 @@ public void setHeader(String headerName, byte[] headerValue, Map<String, byte[]>
     void testBinaryHeaderCaching() {
         final Map<String, String> headerMap = Map.of(TraceContext.TRACE_PARENT_TEXTUAL_HEADER_NAME, "00-0af7651916cd43dd8448eb211c80319c-b9c7c989f97918e1-01");
         final TraceContext child = TraceContext.with64BitId(mock(ElasticApmTracer.class));
-        assertThat(TraceContext.getFromTraceContextTextHeaders().asChildOf(child, headerMap, TextHeaderMapAccessor.INSTANCE)).isTrue();
+        assertThat(TraceContext.<Map<String, String>>getFromTraceContextTextHeaders().asChildOf(child, headerMap, TextHeaderMapAccessor.INSTANCE)).isTrue();
         HashMap<String, byte[]> binaryHeaderMap = new HashMap<>();
         assertThat(child.setOutgoingTraceContextHeaders(binaryHeaderMap, BinaryHeaderMapAccessor.INSTANCE)).isTrue();
         byte[] outgoingHeader = binaryHeaderMap.get(TraceContext.TRACE_PARENT_BINARY_HEADER_NAME);
@@ -132,7 +132,7 @@ void testBinaryHeaderCaching() {
     void testBinaryHeader_CachingDisabled() {
         final Map<String, String> headerMap = Map.of(TraceContext.TRACE_PARENT_TEXTUAL_HEADER_NAME, "00-0af7651916cd43dd8448eb211c80319c-b9c7c989f97918e1-01");
         final TraceContext child = TraceContext.with64BitId(mock(ElasticApmTracer.class));
-        assertThat(TraceContext.getFromTraceContextTextHeaders().asChildOf(child, headerMap, TextHeaderMapAccessor.INSTANCE)).isTrue();
+        assertThat(TraceContext.<Map<String, String>>getFromTraceContextTextHeaders().asChildOf(child, headerMap, TextHeaderMapAccessor.INSTANCE)).isTrue();
         BinaryHeaderSetter<Map<String, byte[]>> headerSetter = new BinaryHeaderSetter<>() {
             @Override
             public byte[] getFixedLengthByteArray(String headerName, int length) {

apm-agent-plugins/apm-apache-httpclient-plugin/src/main/java/co/elastic/apm/agent/httpclient/ApacheHttpClientInstrumentation.java

@@ -25,7 +25,6 @@
 package co.elastic.apm.agent.httpclient;
 
 import co.elastic.apm.agent.http.client.HttpClientHelper;
-import co.elastic.apm.agent.httpclient.helper.RequestHeaderAccessor;
 import co.elastic.apm.agent.impl.ElasticApmTracer;
 import co.elastic.apm.agent.impl.transaction.Span;
 import co.elastic.apm.agent.impl.transaction.TextHeaderGetter;
@@ -76,7 +75,7 @@ private static void onBeforeExecute(@Advice.Argument(0) HttpRoute route,
             if (span != null) {
                 span.activate();
                 if (headerSetter != null) {
-                    span.getTraceContext().setOutgoingTraceContextHeaders(request, new RequestHeaderAccessor());
+                    span.getTraceContext().setOutgoingTraceContextHeaders(request, headerSetter);
                 }
             } else if (headerGetter != null && !TraceContext.containsTraceContextTextHeaders(request, headerGetter)
                 && headerSetter != null && parent != null) {

apm-agent-plugins/apm-apache-httpclient-plugin/src/main/java/co/elastic/apm/agent/httpclient/helper/RequestHeaderAccessor.java

@@ -30,8 +30,6 @@
 import org.apache.http.HttpRequest;
 
 import javax.annotation.Nullable;
-import java.util.ArrayList;
-import java.util.List;
 
 public class RequestHeaderAccessor implements TextHeaderGetter<HttpRequest>, TextHeaderSetter<HttpRequest> {
     @Override
@@ -49,17 +47,13 @@ public String getFirstHeader(String headerName, HttpRequest request) {
         return null;
     }
 
-    @Nullable
     @Override
-    public Iterable<String> getHeaders(String headerName, HttpRequest request) {
-        Header[] headers = request.getHeaders(headerName);
-        if (headers == null) {
-            return null;
-        }
-        List<String> ret = new ArrayList<>();
-        for (Header header : headers) {
-            ret.add(header.getValue());
+    public <S> void forEach(String headerName, HttpRequest carrier, S state, HeaderConsumer<String, S> consumer) {
+        Header[] headers = carrier.getHeaders(headerName);
+        if (headers != null) {
+            for (Header header : headers) {
+                consumer.accept(header.getValue(), state);
+            }
         }
-        return ret;
     }
 }

apm-agent-plugins/apm-api-plugin/src/main/java/co/elastic/apm/agent/plugin/api/AbstractSpanInstrumentation.java

@@ -27,7 +27,6 @@
 import co.elastic.apm.agent.bci.VisibleForAdvice;
 import co.elastic.apm.agent.impl.transaction.AbstractSpan;
 import co.elastic.apm.agent.impl.transaction.Span;
-import co.elastic.apm.agent.impl.transaction.TraceContext;
 import co.elastic.apm.agent.impl.transaction.TraceContextHolder;
 import co.elastic.apm.agent.impl.transaction.Transaction;
 import net.bytebuddy.asm.Advice;
@@ -291,8 +290,7 @@ public static void injectTraceHeaders(@Advice.FieldValue(value = "span", typing
                                               @Advice.Argument(0) MethodHandle addHeaderMethodHandle,
                                               @Advice.Argument(1) @Nullable Object headerInjector) throws Throwable {
             if (headerInjector != null) {
-                HeaderInjectorBridge.instance().setAddHeaderMethodHandle(addHeaderMethodHandle);
-                context.getTraceContext().setOutgoingTraceContextHeaders(headerInjector, HeaderInjectorBridge.instance());
+                context.getTraceContext().setOutgoingTraceContextHeaders(headerInjector, HeaderInjectorBridge.get(addHeaderMethodHandle));
             }
         }
     }