diff --git a/apm-agent-core/src/main/java/co/elastic/apm/agent/impl/ElasticApmTracer.java b/apm-agent-core/src/main/java/co/elastic/apm/agent/impl/ElasticApmTracer.java index fd044d64f2..1bf1e2209a 100644 --- a/apm-agent-core/src/main/java/co/elastic/apm/agent/impl/ElasticApmTracer.java +++ b/apm-agent-core/src/main/java/co/elastic/apm/agent/impl/ElasticApmTracer.java @@ -222,7 +222,7 @@ public Transaction startChildTransaction(@Nullable C headerCarrier, TextHead if (!coreConfiguration.isActive()) { transaction = noopTransaction(); } else { - transaction = createTransaction().start(TraceContext.getFromTraceContextTextHeaders(), headerCarrier, + transaction = createTransaction().start(TraceContext.getFromTraceContextTextHeaders(), headerCarrier, textHeadersGetter, epochMicros, sampler, initiatingClassLoader); } afterTransactionStart(initiatingClassLoader, transaction); @@ -260,7 +260,7 @@ public Transaction startChildTransaction(@Nullable C headerCarrier, BinaryHe if (!coreConfiguration.isActive()) { transaction = noopTransaction(); } else { - transaction = createTransaction().start(TraceContext.getFromTraceContextBinaryHeaders(), headerCarrier, + transaction = createTransaction().start(TraceContext.getFromTraceContextBinaryHeaders(), headerCarrier, binaryHeadersGetter, epochMicros, sampler, initiatingClassLoader); } afterTransactionStart(initiatingClassLoader, transaction); diff --git a/apm-agent-core/src/main/java/co/elastic/apm/agent/impl/transaction/AbstractHeaderGetter.java b/apm-agent-core/src/main/java/co/elastic/apm/agent/impl/transaction/AbstractHeaderGetter.java new file mode 100644 index 0000000000..6e1f2bcdb2 --- /dev/null +++ b/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 implements HeaderGetter { + @Override + public void forEach(String headerName, C carrier, S state, HeaderConsumer consumer) { + T firstHeader = getFirstHeader(headerName, carrier); + if (firstHeader != null) { + consumer.accept(firstHeader, state); + } + } +} diff --git a/apm-agent-core/src/main/java/co/elastic/apm/agent/impl/transaction/HeaderGetter.java b/apm-agent-core/src/main/java/co/elastic/apm/agent/impl/transaction/HeaderGetter.java index bd809c1d0f..f1f27e19e2 100644 --- a/apm-agent-core/src/main/java/co/elastic/apm/agent/impl/transaction/HeaderGetter.java +++ b/apm-agent-core/src/main/java/co/elastic/apm/agent/impl/transaction/HeaderGetter.java @@ -31,6 +31,23 @@ public interface HeaderGetter { @Nullable T getFirstHeader(String headerName, C carrier); - @Nullable - Iterable 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. + *

+ * 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. + *

+ * + * @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 the type of the state object + */ + void forEach(String headerName, C carrier, S state, HeaderConsumer consumer); + + interface HeaderConsumer { + void accept(T headerValue, S state); + } } diff --git a/apm-agent-core/src/main/java/co/elastic/apm/agent/impl/transaction/TraceContext.java b/apm-agent-core/src/main/java/co/elastic/apm/agent/impl/transaction/TraceContext.java index ccd0299de0..f2a051c540 100644 --- a/apm-agent-core/src/main/java/co/elastic/apm/agent/impl/transaction/TraceContext.java +++ b/apm-agent-core/src/main/java/co/elastic/apm/agent/impl/transaction/TraceContext.java @@ -104,10 +104,10 @@ public boolean asChildOf(TraceContext child, TraceContextHolder parent) { return true; } }; - private static final ChildContextCreatorTwoArg> FROM_TRACE_CONTEXT_TEXT_HEADERS = - new ChildContextCreatorTwoArg>() { + private static final ChildContextCreatorTwoArg FROM_TRACE_CONTEXT_TEXT_HEADERS = + new ChildContextCreatorTwoArg>() { @Override - public boolean asChildOf(TraceContext child, @Nullable Object carrier, TextHeaderGetter traceContextHeaderGetter) { + public boolean asChildOf(TraceContext child, @Nullable Object carrier, TextHeaderGetter traceContextHeaderGetter) { if (carrier == null) { return false; } @@ -115,18 +115,17 @@ public boolean asChildOf(TraceContext child, @Nullable Object carrier, TextHeade // 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) 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> FROM_TRACE_CONTEXT_BINARY_HEADERS = - new ChildContextCreatorTwoArg>() { + private static final ChildContextCreatorTwoArg FROM_TRACE_CONTEXT_BINARY_HEADERS = + new ChildContextCreatorTwoArg>() { @Override - public boolean asChildOf(TraceContext child, @Nullable Object carrier, BinaryHeaderGetter traceContextHeaderGetter) { + public boolean asChildOf(TraceContext child, @Nullable Object carrier, BinaryHeaderGetter traceContextHeaderGetter) { if (carrier == null) { return false; } @@ -134,8 +133,7 @@ public boolean asChildOf(TraceContext child, @Nullable Object carrier, BinaryHea // 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) 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> getFromTraceContextTextHeaders() { - return FROM_TRACE_CONTEXT_TEXT_HEADERS; + public static ChildContextCreatorTwoArg> getFromTraceContextTextHeaders() { + return (ChildContextCreatorTwoArg>) FROM_TRACE_CONTEXT_TEXT_HEADERS; } - public static ChildContextCreatorTwoArg> getFromTraceContextBinaryHeaders() { - return FROM_TRACE_CONTEXT_BINARY_HEADERS; + public static ChildContextCreatorTwoArg> getFromTraceContextBinaryHeaders() { + return (ChildContextCreatorTwoArg>) FROM_TRACE_CONTEXT_BINARY_HEADERS; } public static ChildContextCreator fromActive() { diff --git a/apm-agent-core/src/test/java/co/elastic/apm/agent/impl/BinaryHeaderMapAccessor.java b/apm-agent-core/src/test/java/co/elastic/apm/agent/impl/BinaryHeaderMapAccessor.java index 9f4c83f50d..036c3ef4f8 100644 --- a/apm-agent-core/src/test/java/co/elastic/apm/agent/impl/BinaryHeaderMapAccessor.java +++ b/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>, BinaryHeaderSetter> { @@ -48,10 +47,12 @@ public byte[] getFirstHeader(String headerName, Map headerMap) { return headerMap.get(headerName); } - @Nullable @Override - public Iterable getHeaders(String headerName, Map headerMap) { - return List.of(headerMap.get(headerName)); + public void forEach(String headerName, Map carrier, S state, HeaderConsumer consumer) { + byte[] headerValue = carrier.get(headerName); + if (headerValue != null) { + consumer.accept(headerValue, state); + } } @Nullable diff --git a/apm-agent-core/src/test/java/co/elastic/apm/agent/impl/TextHeaderMapAccessor.java b/apm-agent-core/src/test/java/co/elastic/apm/agent/impl/TextHeaderMapAccessor.java index 4c6c91b741..cdd7902568 100644 --- a/apm-agent-core/src/test/java/co/elastic/apm/agent/impl/TextHeaderMapAccessor.java +++ b/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>, TextHeaderSetter> { +public class TextHeaderMapAccessor extends AbstractHeaderGetter> implements TextHeaderGetter>, TextHeaderSetter> { public static final TextHeaderMapAccessor INSTANCE = new TextHeaderMapAccessor(); @@ -44,16 +44,6 @@ public String getFirstHeader(String headerName, Map headerMap) { return headerMap.get(headerName); } - @Nullable - @Override - public Iterable getHeaders(String headerName, Map headerMap) { - String value = headerMap.get(headerName); - if (value != null) { - return List.of(value); - } - return null; - } - @Override public void setHeader(String headerName, String headerValue, Map headerMap) { headerMap.put(headerName, headerValue); diff --git a/apm-agent-core/src/test/java/co/elastic/apm/agent/impl/transaction/TraceContextTest.java b/apm-agent-core/src/test/java/co/elastic/apm/agent/impl/transaction/TraceContextTest.java index bb563f78bf..64addf7780 100644 --- a/apm-agent-core/src/test/java/co/elastic/apm/agent/impl/transaction/TraceContextTest.java +++ b/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 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.>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 binaryHeaderMap = new HashMap<>(); assertThat(child.setOutgoingTraceContextHeaders(binaryHeaderMap, BinaryHeaderMapAccessor.INSTANCE)).isTrue(); - assertThat(TraceContext.getFromTraceContextBinaryHeaders().asChildOf(grandchild1, binaryHeaderMap, BinaryHeaderMapAccessor.INSTANCE)).isTrue(); + assertThat(TraceContext.>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 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.>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>() { @Override @@ -120,7 +120,7 @@ public void setHeader(String headerName, byte[] headerValue, Map void testBinaryHeaderCaching() { final Map 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.>getFromTraceContextTextHeaders().asChildOf(child, headerMap, TextHeaderMapAccessor.INSTANCE)).isTrue(); HashMap 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 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.>getFromTraceContextTextHeaders().asChildOf(child, headerMap, TextHeaderMapAccessor.INSTANCE)).isTrue(); BinaryHeaderSetter> headerSetter = new BinaryHeaderSetter<>() { @Override public byte[] getFixedLengthByteArray(String headerName, int length) { diff --git a/apm-agent-plugins/apm-apache-httpclient-plugin/src/main/java/co/elastic/apm/agent/httpclient/ApacheHttpClientInstrumentation.java b/apm-agent-plugins/apm-apache-httpclient-plugin/src/main/java/co/elastic/apm/agent/httpclient/ApacheHttpClientInstrumentation.java index ccce8ba71b..2bf2079dde 100644 --- a/apm-agent-plugins/apm-apache-httpclient-plugin/src/main/java/co/elastic/apm/agent/httpclient/ApacheHttpClientInstrumentation.java +++ b/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) { diff --git a/apm-agent-plugins/apm-apache-httpclient-plugin/src/main/java/co/elastic/apm/agent/httpclient/helper/RequestHeaderAccessor.java b/apm-agent-plugins/apm-apache-httpclient-plugin/src/main/java/co/elastic/apm/agent/httpclient/helper/RequestHeaderAccessor.java index eda60b1ad9..04ec912c9a 100644 --- a/apm-agent-plugins/apm-apache-httpclient-plugin/src/main/java/co/elastic/apm/agent/httpclient/helper/RequestHeaderAccessor.java +++ b/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, TextHeaderSetter { @Override @@ -49,17 +47,13 @@ public String getFirstHeader(String headerName, HttpRequest request) { return null; } - @Nullable @Override - public Iterable getHeaders(String headerName, HttpRequest request) { - Header[] headers = request.getHeaders(headerName); - if (headers == null) { - return null; - } - List ret = new ArrayList<>(); - for (Header header : headers) { - ret.add(header.getValue()); + public void forEach(String headerName, HttpRequest carrier, S state, HeaderConsumer consumer) { + Header[] headers = carrier.getHeaders(headerName); + if (headers != null) { + for (Header header : headers) { + consumer.accept(header.getValue(), state); + } } - return ret; } } diff --git a/apm-agent-plugins/apm-api-plugin/src/main/java/co/elastic/apm/agent/plugin/api/AbstractSpanInstrumentation.java b/apm-agent-plugins/apm-api-plugin/src/main/java/co/elastic/apm/agent/plugin/api/AbstractSpanInstrumentation.java index 65e2611615..1f5398f557 100644 --- a/apm-agent-plugins/apm-api-plugin/src/main/java/co/elastic/apm/agent/plugin/api/AbstractSpanInstrumentation.java +++ b/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)); } } } diff --git a/apm-agent-plugins/apm-api-plugin/src/main/java/co/elastic/apm/agent/plugin/api/ElasticApmApiInstrumentation.java b/apm-agent-plugins/apm-api-plugin/src/main/java/co/elastic/apm/agent/plugin/api/ElasticApmApiInstrumentation.java index e74802c6ed..73a720611b 100644 --- a/apm-agent-plugins/apm-api-plugin/src/main/java/co/elastic/apm/agent/plugin/api/ElasticApmApiInstrumentation.java +++ b/apm-agent-plugins/apm-api-plugin/src/main/java/co/elastic/apm/agent/plugin/api/ElasticApmApiInstrumentation.java @@ -86,14 +86,12 @@ private static void doStartTransaction(@Advice.Origin Class clazz, @Advice.Argument(2) MethodHandle getAllHeaders, @Advice.Argument(3) @Nullable Object headersExtractor) { if (tracer != null) { - if (headerExtractor != null) { - HeaderExtractorBridge headerExtractorBridge = HeaderExtractorBridge.instance(); - headerExtractorBridge.setGetHeaderMethodHandle(getFirstHeader); - transaction = tracer.startChildTransaction(headerExtractor, headerExtractorBridge, clazz.getClassLoader()); - } else if (headersExtractor != null) { - HeadersExtractorBridge headersExtractorBridge = HeadersExtractorBridge.instance(); - headersExtractorBridge.setGetHeaderMethodHandle(getAllHeaders); - transaction = tracer.startChildTransaction(headersExtractor, headersExtractorBridge, clazz.getClassLoader()); + if (headersExtractor != null) { + HeadersExtractorBridge headersExtractorBridge = HeadersExtractorBridge.get(getFirstHeader, getAllHeaders); + transaction = tracer.startChildTransaction(new HeadersExtractorBridge.Extractor(headerExtractor, headersExtractor), headersExtractorBridge, clazz.getClassLoader()); + } else if (headerExtractor != null) { + HeaderExtractorBridge headersExtractorBridge = HeaderExtractorBridge.get(getFirstHeader); + transaction = tracer.startChildTransaction(headerExtractor, headersExtractorBridge, clazz.getClassLoader()); } else { transaction = tracer.startRootTransaction(clazz.getClassLoader()); } diff --git a/apm-agent-plugins/apm-api-plugin/src/main/java/co/elastic/apm/agent/plugin/api/HeaderExtractorBridge.java b/apm-agent-plugins/apm-api-plugin/src/main/java/co/elastic/apm/agent/plugin/api/HeaderExtractorBridge.java index dc597d8c45..a4c8a8209c 100644 --- a/apm-agent-plugins/apm-api-plugin/src/main/java/co/elastic/apm/agent/plugin/api/HeaderExtractorBridge.java +++ b/apm-agent-plugins/apm-api-plugin/src/main/java/co/elastic/apm/agent/plugin/api/HeaderExtractorBridge.java @@ -25,6 +25,7 @@ package co.elastic.apm.agent.plugin.api; import co.elastic.apm.agent.bci.VisibleForAdvice; +import co.elastic.apm.agent.impl.transaction.AbstractHeaderGetter; import co.elastic.apm.agent.impl.transaction.TextHeaderGetter; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -33,30 +34,24 @@ import java.lang.invoke.MethodHandle; @VisibleForAdvice -public class HeaderExtractorBridge implements TextHeaderGetter { +public class HeaderExtractorBridge extends AbstractHeaderGetter implements TextHeaderGetter { private static final Logger logger = LoggerFactory.getLogger(HeaderExtractorBridge.class); - private static final HeaderExtractorBridge INSTANCE = new HeaderExtractorBridge(); - - @VisibleForAdvice - public static HeaderExtractorBridge instance() { - return INSTANCE; - } - @Nullable - private MethodHandle getFirstHeaderMethod; + private static HeaderExtractorBridge INSTANCE; + + private final MethodHandle getFirstHeaderMethod; - private HeaderExtractorBridge() { + private HeaderExtractorBridge(MethodHandle getFirstHeaderMethod) { + this.getFirstHeaderMethod = getFirstHeaderMethod; } - @VisibleForAdvice - public void setGetHeaderMethodHandle(MethodHandle getHeaderMethodHandle) { - // No need to make thread-safe - only one methodHandle can be set in practice, so we don't care replacing its - // reference a couple of times on startup. Cheaper than volatile access. - if (this.getFirstHeaderMethod == null) { - this.getFirstHeaderMethod = getHeaderMethodHandle; + public static HeaderExtractorBridge get(MethodHandle getFirstHeaderMethod) { + if (INSTANCE == null) { + INSTANCE = new HeaderExtractorBridge(getFirstHeaderMethod); } + return INSTANCE; } @Nullable @@ -64,21 +59,11 @@ public void setGetHeaderMethodHandle(MethodHandle getHeaderMethodHandle) { public String getFirstHeader(String headerName, Object carrier) { String value = null; try { - if (getFirstHeaderMethod != null) { - value = (String) getFirstHeaderMethod.invoke(carrier, headerName); - } + value = (String) getFirstHeaderMethod.invoke(carrier, headerName); } catch (Throwable throwable) { logger.error("Failed to extract trace context headers", throwable); } return value; } - /** - * Returns null. {@link HeadersExtractorBridge} should be used instead for this functionality - */ - @Nullable - @Override - public Iterable getHeaders(String headerName, Object carrier) { - return null; - } } diff --git a/apm-agent-plugins/apm-api-plugin/src/main/java/co/elastic/apm/agent/plugin/api/HeaderInjectorBridge.java b/apm-agent-plugins/apm-api-plugin/src/main/java/co/elastic/apm/agent/plugin/api/HeaderInjectorBridge.java index d289ed57f8..e459f3aa1c 100644 --- a/apm-agent-plugins/apm-api-plugin/src/main/java/co/elastic/apm/agent/plugin/api/HeaderInjectorBridge.java +++ b/apm-agent-plugins/apm-api-plugin/src/main/java/co/elastic/apm/agent/plugin/api/HeaderInjectorBridge.java @@ -25,7 +25,6 @@ package co.elastic.apm.agent.plugin.api; import co.elastic.apm.agent.bci.VisibleForAdvice; -import co.elastic.apm.agent.impl.transaction.TextHeaderGetter; import co.elastic.apm.agent.impl.transaction.TextHeaderSetter; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -38,34 +37,27 @@ public class HeaderInjectorBridge implements TextHeaderSetter { private static final Logger logger = LoggerFactory.getLogger(HeaderInjectorBridge.class); - private static final HeaderInjectorBridge INSTANCE = new HeaderInjectorBridge(); + @Nullable + private static HeaderInjectorBridge INSTANCE; @VisibleForAdvice - public static HeaderInjectorBridge instance() { + public static HeaderInjectorBridge get(MethodHandle addHeaderMethod) { + if (INSTANCE == null) { + INSTANCE = new HeaderInjectorBridge(addHeaderMethod); + } return INSTANCE; } - @Nullable - private MethodHandle addHeaderMethod; + private final MethodHandle addHeaderMethod; - private HeaderInjectorBridge() { - } - - @VisibleForAdvice - public void setAddHeaderMethodHandle(MethodHandle addMethodHandle) { - // No need to make thread-safe - only one methodHandle can be set in practice, so we don't care replacing its - // reference a couple of times on startup. Cheaper than volatile access. - if (this.addHeaderMethod == null) { - this.addHeaderMethod = addMethodHandle; - } + private HeaderInjectorBridge(MethodHandle addHeaderMethod) { + this.addHeaderMethod = addHeaderMethod; } @Override public void setHeader(String headerName, String headerValue, Object carrier) { try { - if (addHeaderMethod != null) { - addHeaderMethod.invoke(carrier, headerName, headerValue); - } + addHeaderMethod.invoke(carrier, headerName, headerValue); } catch (Throwable throwable) { logger.error("Failed to add trace context headers", throwable); } diff --git a/apm-agent-plugins/apm-api-plugin/src/main/java/co/elastic/apm/agent/plugin/api/HeadersExtractorBridge.java b/apm-agent-plugins/apm-api-plugin/src/main/java/co/elastic/apm/agent/plugin/api/HeadersExtractorBridge.java index d25f9c73fa..07ac5ceeb2 100644 --- a/apm-agent-plugins/apm-api-plugin/src/main/java/co/elastic/apm/agent/plugin/api/HeadersExtractorBridge.java +++ b/apm-agent-plugins/apm-api-plugin/src/main/java/co/elastic/apm/agent/plugin/api/HeadersExtractorBridge.java @@ -34,65 +34,78 @@ import java.util.Iterator; @VisibleForAdvice -public class HeadersExtractorBridge implements TextHeaderGetter { +public class HeadersExtractorBridge implements TextHeaderGetter { private static final Logger logger = LoggerFactory.getLogger(HeadersExtractorBridge.class); + @Nullable + private static HeadersExtractorBridge INSTANCE; - private static final HeadersExtractorBridge INSTANCE = new HeadersExtractorBridge(); + public static class Extractor { + @Nullable + private final Object headerExtractor; + private final Object headersExtractor; - @VisibleForAdvice - public static HeadersExtractorBridge instance() { - return INSTANCE; + public Extractor(@Nullable Object headerExtractor, Object headersExtractor) { + this.headerExtractor = headerExtractor; + this.headersExtractor = headersExtractor; + } } - @Nullable - private MethodHandle getAllHeadersMethod; + private final MethodHandle getFirstHeaderMethod; + private final MethodHandle getAllHeadersMethod; - private HeadersExtractorBridge() { + public static HeadersExtractorBridge get(MethodHandle getFirstHeaderMethod, MethodHandle getAllHeadersMethod) { + if (INSTANCE == null) { + INSTANCE = new HeadersExtractorBridge(getFirstHeaderMethod, getAllHeadersMethod); + } + return INSTANCE; } - @VisibleForAdvice - public void setGetHeaderMethodHandle(MethodHandle getAllHeadersMethodHandle) { - // No need to make thread-safe - only one methodHandle can be set in practice, so we don't care replacing its - // reference a couple of times on startup. Cheaper than volatile access. - if (this.getAllHeadersMethod == null) { - this.getAllHeadersMethod = getAllHeadersMethodHandle; - } + private HeadersExtractorBridge(MethodHandle getFirstHeaderMethod, MethodHandle getAllHeadersMethod) { + this.getFirstHeaderMethod = getFirstHeaderMethod; + this.getAllHeadersMethod = getAllHeadersMethod; } @Nullable @Override - public String getFirstHeader(String headerName, Object carrier) { - String value = null; - try { - if (getAllHeadersMethod != null) { - //noinspection unchecked - final Iterable headersIterable = (Iterable) getAllHeadersMethod.invoke(carrier, headerName); - if (headersIterable != null) { - final Iterator headersIterator = headersIterable.iterator(); - if (headersIterator.hasNext()) { - value = headersIterator.next(); - } + public String getFirstHeader(String headerName, Extractor carrier) { + if (carrier.headerExtractor != null) { + try { + return (String) getFirstHeaderMethod.invoke(carrier.headerExtractor, headerName); + } catch (Throwable throwable) { + logger.error("Failed to extract trace context headers", throwable); + } + } else { + Iterable headers = getValues(headerName, carrier); + if (headers != null) { + Iterator iterator = headers.iterator(); + if (iterator.hasNext()) { + return iterator.next(); } } - } catch (Throwable throwable) { - logger.error("Failed to extract trace context headers", throwable); } - return value; + return null; } - @Nullable @Override - public Iterable getHeaders(String headerName, Object carrier) { - Iterable values = null; - if (getAllHeadersMethod != null) { - try { - //noinspection unchecked - values = (Iterable) getAllHeadersMethod.invoke(carrier, headerName); - } catch (Throwable throwable) { - logger.error("Failed to extract trace context headers", throwable); + public void forEach(String headerName, Extractor carrier, S state, HeaderConsumer consumer) { + Iterable values = getValues(headerName, carrier); + if (values != null) { + for (String value : values) { + consumer.accept(value, state); } } + } + + @Nullable + public Iterable getValues(String headerName, Extractor carrier) { + Iterable values = null; + try { + //noinspection unchecked + values = (Iterable) getAllHeadersMethod.invoke(carrier.headersExtractor, headerName); + } catch (Throwable throwable) { + logger.error("Failed to extract trace context headers", throwable); + } return values; } } diff --git a/apm-agent-plugins/apm-asynchttpclient-plugin/src/main/java/co/elastic/apm/agent/asynchttpclient/AbstractAsyncHttpClientInstrumentation.java b/apm-agent-plugins/apm-asynchttpclient-plugin/src/main/java/co/elastic/apm/agent/asynchttpclient/AbstractAsyncHttpClientInstrumentation.java index 059bf4c6a7..640fd907ed 100644 --- a/apm-agent-plugins/apm-asynchttpclient-plugin/src/main/java/co/elastic/apm/agent/asynchttpclient/AbstractAsyncHttpClientInstrumentation.java +++ b/apm-agent-plugins/apm-asynchttpclient-plugin/src/main/java/co/elastic/apm/agent/asynchttpclient/AbstractAsyncHttpClientInstrumentation.java @@ -76,7 +76,7 @@ public AbstractAsyncHttpClientInstrumentation() { synchronized (AbstractAsyncHandlerInstrumentation.class) { if (headerSetterManager == null) { headerSetterManager = HelperClassManager.ForAnyClassLoader.of(tracer, - "co.elastic.apm.agent.asynchttpclient.helper.RequestHeaderAccessor" + "co.elastic.apm.agent.asynchttpclient.helper.RequestHeaderSetter" ); } } diff --git a/apm-agent-plugins/apm-asynchttpclient-plugin/src/main/java/co/elastic/apm/agent/asynchttpclient/helper/RequestHeaderAccessor.java b/apm-agent-plugins/apm-asynchttpclient-plugin/src/main/java/co/elastic/apm/agent/asynchttpclient/helper/RequestHeaderSetter.java similarity index 71% rename from apm-agent-plugins/apm-asynchttpclient-plugin/src/main/java/co/elastic/apm/agent/asynchttpclient/helper/RequestHeaderAccessor.java rename to apm-agent-plugins/apm-asynchttpclient-plugin/src/main/java/co/elastic/apm/agent/asynchttpclient/helper/RequestHeaderSetter.java index 20cda7a56a..edd644f917 100644 --- a/apm-agent-plugins/apm-asynchttpclient-plugin/src/main/java/co/elastic/apm/agent/asynchttpclient/helper/RequestHeaderAccessor.java +++ b/apm-agent-plugins/apm-asynchttpclient-plugin/src/main/java/co/elastic/apm/agent/asynchttpclient/helper/RequestHeaderSetter.java @@ -25,29 +25,15 @@ package co.elastic.apm.agent.asynchttpclient.helper; import co.elastic.apm.agent.bci.VisibleForAdvice; -import co.elastic.apm.agent.impl.transaction.TextHeaderGetter; import co.elastic.apm.agent.impl.transaction.TextHeaderSetter; import org.asynchttpclient.Request; -import javax.annotation.Nullable; - @VisibleForAdvice -public class RequestHeaderAccessor implements TextHeaderGetter, TextHeaderSetter { +public class RequestHeaderSetter implements TextHeaderSetter { @Override public void setHeader(String headerName, String headerValue, Request request) { request.getHeaders().set(headerName, headerValue); } - @Nullable - @Override - public String getFirstHeader(String headerName, Request request) { - return request.getHeaders().get(headerName); - } - - @Nullable - @Override - public Iterable getHeaders(String headerName, Request request) { - return request.getHeaders().getAll(headerName); - } } diff --git a/apm-agent-plugins/apm-httpclient-core/src/test/java/co/elastic/apm/agent/httpclient/AbstractHttpClientInstrumentationTest.java b/apm-agent-plugins/apm-httpclient-core/src/test/java/co/elastic/apm/agent/httpclient/AbstractHttpClientInstrumentationTest.java index 035c974b04..c51f5c94e3 100644 --- a/apm-agent-plugins/apm-httpclient-core/src/test/java/co/elastic/apm/agent/httpclient/AbstractHttpClientInstrumentationTest.java +++ b/apm-agent-plugins/apm-httpclient-core/src/test/java/co/elastic/apm/agent/httpclient/AbstractHttpClientInstrumentationTest.java @@ -42,7 +42,6 @@ import org.junit.Test; import javax.annotation.Nullable; - import java.util.HashMap; import java.util.List; import java.util.Map; @@ -50,13 +49,10 @@ import static com.github.tomakehurst.wiremock.client.WireMock.aResponse; import static com.github.tomakehurst.wiremock.client.WireMock.any; import static com.github.tomakehurst.wiremock.client.WireMock.anyRequestedFor; -import static com.github.tomakehurst.wiremock.client.WireMock.equalTo; import static com.github.tomakehurst.wiremock.client.WireMock.get; -import static com.github.tomakehurst.wiremock.client.WireMock.getRequestedFor; import static com.github.tomakehurst.wiremock.client.WireMock.seeOther; import static com.github.tomakehurst.wiremock.client.WireMock.urlEqualTo; import static com.github.tomakehurst.wiremock.client.WireMock.urlPathEqualTo; -import static com.github.tomakehurst.wiremock.client.WireMock.verify; import static org.assertj.core.api.Assertions.assertThat; public abstract class AbstractHttpClientInstrumentationTest extends AbstractInstrumentationTest { @@ -229,17 +225,18 @@ public String getFirstHeader(String headerName, LoggedRequest loggedRequest) { return loggedRequest.getHeader(headerName); } - @Nullable @Override - public Iterable getHeaders(String headerName, LoggedRequest loggedRequest) { + public void forEach(String headerName, LoggedRequest loggedRequest, S state, HeaderConsumer consumer) { HttpHeaders headers = loggedRequest.getHeaders(); if (headers != null) { HttpHeader header = headers.getHeader(headerName); if (header != null) { - return header.values(); + List values = header.values(); + for (int i = 0, size = values.size(); i < size; i++) { + consumer.accept(values.get(i), state); + } } } - return null; } } } diff --git a/apm-agent-plugins/apm-jms-plugin/src/main/java/co/elastic/apm/agent/jms/JmsInstrumentationHelperImpl.java b/apm-agent-plugins/apm-jms-plugin/src/main/java/co/elastic/apm/agent/jms/JmsInstrumentationHelperImpl.java index 6ed56f0b00..b9f509ee0a 100644 --- a/apm-agent-plugins/apm-jms-plugin/src/main/java/co/elastic/apm/agent/jms/JmsInstrumentationHelperImpl.java +++ b/apm-agent-plugins/apm-jms-plugin/src/main/java/co/elastic/apm/agent/jms/JmsInstrumentationHelperImpl.java @@ -127,7 +127,7 @@ public Transaction startJmsTransaction(Message parentMessage, Class instrumen @Override public void makeChildOf(Transaction childTransaction, Message parentMessage) { - TraceContext.getFromTraceContextTextHeaders().asChildOf(childTransaction.getTraceContext(), parentMessage, JmsMessagePropertyAccessor.instance()); + TraceContext.getFromTraceContextTextHeaders().asChildOf(childTransaction.getTraceContext(), parentMessage, JmsMessagePropertyAccessor.instance()); } @VisibleForAdvice diff --git a/apm-agent-plugins/apm-jms-plugin/src/main/java/co/elastic/apm/agent/jms/JmsMessagePropertyAccessor.java b/apm-agent-plugins/apm-jms-plugin/src/main/java/co/elastic/apm/agent/jms/JmsMessagePropertyAccessor.java index 9a3954d1e9..6ee6a187ad 100644 --- a/apm-agent-plugins/apm-jms-plugin/src/main/java/co/elastic/apm/agent/jms/JmsMessagePropertyAccessor.java +++ b/apm-agent-plugins/apm-jms-plugin/src/main/java/co/elastic/apm/agent/jms/JmsMessagePropertyAccessor.java @@ -24,6 +24,7 @@ */ package co.elastic.apm.agent.jms; +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 co.elastic.apm.agent.impl.transaction.TraceContext; @@ -37,7 +38,7 @@ import static co.elastic.apm.agent.jms.JmsInstrumentationHelper.JMS_TRACE_PARENT_PROPERTY; -public class JmsMessagePropertyAccessor implements TextHeaderGetter, TextHeaderSetter { +public class JmsMessagePropertyAccessor extends AbstractHeaderGetter implements TextHeaderGetter, TextHeaderSetter { private static final Logger logger = LoggerFactory.getLogger(JmsMessagePropertyAccessor.class); @@ -72,13 +73,6 @@ private String jmsifyHeaderName(String headerName) { return headerName; } - @Nullable - @Override - public Iterable getHeaders(String headerName, Message message) { - // not supported for JMS - return null; - } - @Override public void setHeader(String headerName, String headerValue, Message message) { headerName = jmsifyHeaderName(headerName); diff --git a/apm-agent-plugins/apm-kafka-plugin/apm-kafka-headers-plugin/src/main/java/co/elastic/apm/agent/kafka/BaseKafkaHeadersInstrumentation.java b/apm-agent-plugins/apm-kafka-plugin/apm-kafka-headers-plugin/src/main/java/co/elastic/apm/agent/kafka/BaseKafkaHeadersInstrumentation.java index 4d51810e41..01ba0633ce 100644 --- a/apm-agent-plugins/apm-kafka-plugin/apm-kafka-headers-plugin/src/main/java/co/elastic/apm/agent/kafka/BaseKafkaHeadersInstrumentation.java +++ b/apm-agent-plugins/apm-kafka-plugin/apm-kafka-headers-plugin/src/main/java/co/elastic/apm/agent/kafka/BaseKafkaHeadersInstrumentation.java @@ -52,8 +52,7 @@ private synchronized static void init(ElasticApmTracer tracer) { "co.elastic.apm.agent.kafka.helper.ConsumerRecordsIterableWrapper", "co.elastic.apm.agent.kafka.helper.ConsumerRecordsListWrapper", "co.elastic.apm.agent.kafka.helper.ElasticHeaderImpl", - "co.elastic.apm.agent.kafka.helper.KafkaRecordHeaderAccessor", - "co.elastic.apm.agent.kafka.helper.KafkaRecordHeaderAccessor$HeaderValuesIterator"); + "co.elastic.apm.agent.kafka.helper.KafkaRecordHeaderAccessor"); } } diff --git a/apm-agent-plugins/apm-kafka-plugin/apm-kafka-headers-plugin/src/main/java/co/elastic/apm/agent/kafka/helper/KafkaRecordHeaderAccessor.java b/apm-agent-plugins/apm-kafka-plugin/apm-kafka-headers-plugin/src/main/java/co/elastic/apm/agent/kafka/helper/KafkaRecordHeaderAccessor.java index a47d565146..a8a4a4a7c0 100644 --- a/apm-agent-plugins/apm-kafka-plugin/apm-kafka-headers-plugin/src/main/java/co/elastic/apm/agent/kafka/helper/KafkaRecordHeaderAccessor.java +++ b/apm-agent-plugins/apm-kafka-plugin/apm-kafka-headers-plugin/src/main/java/co/elastic/apm/agent/kafka/helper/KafkaRecordHeaderAccessor.java @@ -35,20 +35,16 @@ import javax.annotation.Nullable; import java.util.HashMap; -import java.util.Iterator; import java.util.Map; -import java.util.NoSuchElementException; @SuppressWarnings("rawtypes") class KafkaRecordHeaderAccessor implements BinaryHeaderGetter, BinaryHeaderSetter, - HeaderRemover, Iterable { + HeaderRemover { public static final Logger logger = LoggerFactory.getLogger(KafkaRecordHeaderAccessor.class); private static final KafkaRecordHeaderAccessor INSTANCE = new KafkaRecordHeaderAccessor(); - private static final HeaderValuesIterator headerIteratorWrapper = new HeaderValuesIterator(); - private static final ThreadLocal> threadLocalHeaderMap = new ThreadLocal<>(); public static KafkaRecordHeaderAccessor instance() { @@ -65,16 +61,11 @@ public byte[] getFirstHeader(String headerName, ConsumerRecord record) { return null; } - @Nullable - @Override - public Iterable getHeaders(String headerName, ConsumerRecord record) { - headerIteratorWrapper.headerIterator = record.headers().headers(headerName).iterator(); - return this; - } - @Override - public Iterator iterator() { - return headerIteratorWrapper; + public void forEach(String headerName, ConsumerRecord carrier, S state, HeaderConsumer consumer) { + for (Header header : carrier.headers().headers(headerName)) { + consumer.accept(header.value(), state); + } } @Override @@ -114,24 +105,4 @@ public void remove(String headerName, ProducerRecord carrier) { carrier.headers().remove(headerName); } - private static final class HeaderValuesIterator implements Iterator { - @Nullable - private Iterator
headerIterator; - - @Override - public boolean hasNext() { - if (headerIterator != null) { - return headerIterator.hasNext(); - } - return false; - } - - @Override - public byte[] next() { - if (headerIterator != null) { - return headerIterator.next().value(); - } - throw new NoSuchElementException("Kafka record header iterator is not initialized"); - } - } } diff --git a/apm-agent-plugins/apm-okhttp-plugin/src/main/java/co/elastic/apm/agent/okhttp/AbstractOkHttp3ClientInstrumentation.java b/apm-agent-plugins/apm-okhttp-plugin/src/main/java/co/elastic/apm/agent/okhttp/AbstractOkHttp3ClientInstrumentation.java index 62fdc9a968..b315d1ec1c 100644 --- a/apm-agent-plugins/apm-okhttp-plugin/src/main/java/co/elastic/apm/agent/okhttp/AbstractOkHttp3ClientInstrumentation.java +++ b/apm-agent-plugins/apm-okhttp-plugin/src/main/java/co/elastic/apm/agent/okhttp/AbstractOkHttp3ClientInstrumentation.java @@ -28,7 +28,6 @@ import co.elastic.apm.agent.bci.HelperClassManager; import co.elastic.apm.agent.bci.VisibleForAdvice; import co.elastic.apm.agent.impl.ElasticApmTracer; -import co.elastic.apm.agent.impl.transaction.TextHeaderGetter; import co.elastic.apm.agent.impl.transaction.TextHeaderSetter; import okhttp3.Request; @@ -41,17 +40,11 @@ public abstract class AbstractOkHttp3ClientInstrumentation extends ElasticApmIns // We can refer OkHttp types thanks to type erasure @VisibleForAdvice @Nullable - public static HelperClassManager> headerGetterHelperManager; - @VisibleForAdvice - @Nullable public static HelperClassManager> headerSetterHelperManager; public AbstractOkHttp3ClientInstrumentation(ElasticApmTracer tracer) { - headerGetterHelperManager = HelperClassManager.ForAnyClassLoader.of(tracer, - "co.elastic.apm.agent.okhttp.OkHttp3RequestHeaderAccessor" - ); headerSetterHelperManager = HelperClassManager.ForAnyClassLoader.of(tracer, - "co.elastic.apm.agent.okhttp.OkHttp3RequestHeaderAccessor" + "co.elastic.apm.agent.okhttp.OkHttp3RequestHeaderSetter" ); } diff --git a/apm-agent-plugins/apm-okhttp-plugin/src/main/java/co/elastic/apm/agent/okhttp/AbstractOkHttpClientInstrumentation.java b/apm-agent-plugins/apm-okhttp-plugin/src/main/java/co/elastic/apm/agent/okhttp/AbstractOkHttpClientInstrumentation.java index fc5a1df096..1362321841 100644 --- a/apm-agent-plugins/apm-okhttp-plugin/src/main/java/co/elastic/apm/agent/okhttp/AbstractOkHttpClientInstrumentation.java +++ b/apm-agent-plugins/apm-okhttp-plugin/src/main/java/co/elastic/apm/agent/okhttp/AbstractOkHttpClientInstrumentation.java @@ -28,7 +28,6 @@ import co.elastic.apm.agent.bci.HelperClassManager; import co.elastic.apm.agent.bci.VisibleForAdvice; import co.elastic.apm.agent.impl.ElasticApmTracer; -import co.elastic.apm.agent.impl.transaction.TextHeaderGetter; import co.elastic.apm.agent.impl.transaction.TextHeaderSetter; import com.squareup.okhttp.Request; @@ -41,17 +40,11 @@ public abstract class AbstractOkHttpClientInstrumentation extends ElasticApmInst // We can refer OkHttp types thanks to type erasure @VisibleForAdvice @Nullable - public static HelperClassManager> headerGetterHelperManager; - @VisibleForAdvice - @Nullable public static HelperClassManager> headerSetterHelperManager; public AbstractOkHttpClientInstrumentation(ElasticApmTracer tracer) { - headerGetterHelperManager = HelperClassManager.ForAnyClassLoader.of(tracer, - "co.elastic.apm.agent.okhttp.OkHttpRequestHeaderAccessor" - ); headerSetterHelperManager = HelperClassManager.ForAnyClassLoader.of(tracer, - "co.elastic.apm.agent.okhttp.OkHttpRequestHeaderAccessor" + "co.elastic.apm.agent.okhttp.OkHttpRequestHeaderSetter" ); } diff --git a/apm-agent-plugins/apm-okhttp-plugin/src/main/java/co/elastic/apm/agent/okhttp/OkHttp3RequestHeaderAccessor.java b/apm-agent-plugins/apm-okhttp-plugin/src/main/java/co/elastic/apm/agent/okhttp/OkHttp3RequestHeaderSetter.java similarity index 70% rename from apm-agent-plugins/apm-okhttp-plugin/src/main/java/co/elastic/apm/agent/okhttp/OkHttp3RequestHeaderAccessor.java rename to apm-agent-plugins/apm-okhttp-plugin/src/main/java/co/elastic/apm/agent/okhttp/OkHttp3RequestHeaderSetter.java index be24fb7dd5..ef74770ab0 100644 --- a/apm-agent-plugins/apm-okhttp-plugin/src/main/java/co/elastic/apm/agent/okhttp/OkHttp3RequestHeaderAccessor.java +++ b/apm-agent-plugins/apm-okhttp-plugin/src/main/java/co/elastic/apm/agent/okhttp/OkHttp3RequestHeaderSetter.java @@ -24,25 +24,10 @@ */ package co.elastic.apm.agent.okhttp; -import co.elastic.apm.agent.impl.transaction.TextHeaderGetter; import co.elastic.apm.agent.impl.transaction.TextHeaderSetter; import okhttp3.Request; -import javax.annotation.Nullable; - -public class OkHttp3RequestHeaderAccessor implements TextHeaderSetter, TextHeaderGetter { - - @Nullable - @Override - public String getFirstHeader(String headerName, Request request) { - return request.header(headerName); - } - - @Nullable - @Override - public Iterable getHeaders(String headerName, Request request) { - return request.headers(headerName); - } +public class OkHttp3RequestHeaderSetter implements TextHeaderSetter { @Override public void setHeader(String headerName, String headerValue, Request.Builder requestBuilder) { diff --git a/apm-agent-plugins/apm-okhttp-plugin/src/main/java/co/elastic/apm/agent/okhttp/OkHttpRequestHeaderAccessor.java b/apm-agent-plugins/apm-okhttp-plugin/src/main/java/co/elastic/apm/agent/okhttp/OkHttpRequestHeaderSetter.java similarity index 70% rename from apm-agent-plugins/apm-okhttp-plugin/src/main/java/co/elastic/apm/agent/okhttp/OkHttpRequestHeaderAccessor.java rename to apm-agent-plugins/apm-okhttp-plugin/src/main/java/co/elastic/apm/agent/okhttp/OkHttpRequestHeaderSetter.java index d3ee86e3c7..11088bf307 100644 --- a/apm-agent-plugins/apm-okhttp-plugin/src/main/java/co/elastic/apm/agent/okhttp/OkHttpRequestHeaderAccessor.java +++ b/apm-agent-plugins/apm-okhttp-plugin/src/main/java/co/elastic/apm/agent/okhttp/OkHttpRequestHeaderSetter.java @@ -24,25 +24,10 @@ */ package co.elastic.apm.agent.okhttp; -import co.elastic.apm.agent.impl.transaction.TextHeaderGetter; import co.elastic.apm.agent.impl.transaction.TextHeaderSetter; import com.squareup.okhttp.Request; -import javax.annotation.Nullable; - -public class OkHttpRequestHeaderAccessor implements TextHeaderSetter, TextHeaderGetter { - - @Nullable - @Override - public String getFirstHeader(String headerName, Request request) { - return request.header(headerName); - } - - @Nullable - @Override - public Iterable getHeaders(String headerName, Request request) { - return request.headers(headerName); - } +public class OkHttpRequestHeaderSetter implements TextHeaderSetter { @Override public void setHeader(String headerName, String headerValue, Request.Builder requestBuilder) { diff --git a/apm-agent-plugins/apm-opentracing-plugin/src/main/java/co/elastic/apm/agent/opentracing/impl/ExternalSpanContextInstrumentation.java b/apm-agent-plugins/apm-opentracing-plugin/src/main/java/co/elastic/apm/agent/opentracing/impl/ExternalSpanContextInstrumentation.java index 0015cd8d51..5b72b19ee1 100644 --- a/apm-agent-plugins/apm-opentracing-plugin/src/main/java/co/elastic/apm/agent/opentracing/impl/ExternalSpanContextInstrumentation.java +++ b/apm-agent-plugins/apm-opentracing-plugin/src/main/java/co/elastic/apm/agent/opentracing/impl/ExternalSpanContextInstrumentation.java @@ -35,7 +35,6 @@ import org.slf4j.LoggerFactory; import javax.annotation.Nullable; - import java.util.Map; import static net.bytebuddy.matcher.ElementMatchers.named; @@ -105,7 +104,7 @@ public static void toSpanId(@Advice.FieldValue(value = "textMap", typing = Assig public static TraceContext parseTextMap(Iterable> textMap) { if (tracer != null) { TraceContext childTraceContext = TraceContext.with64BitId(tracer); - if (TraceContext.getFromTraceContextTextHeaders().asChildOf(childTraceContext, textMap, OpenTracingTextMapBridge.instance())) { + if (TraceContext.>>getFromTraceContextTextHeaders().asChildOf(childTraceContext, textMap, OpenTracingTextMapBridge.instance())) { return childTraceContext; } } diff --git a/apm-agent-plugins/apm-opentracing-plugin/src/main/java/co/elastic/apm/agent/opentracing/impl/OpenTracingTextMapBridge.java b/apm-agent-plugins/apm-opentracing-plugin/src/main/java/co/elastic/apm/agent/opentracing/impl/OpenTracingTextMapBridge.java index 4c5902a785..c085ffe94e 100644 --- a/apm-agent-plugins/apm-opentracing-plugin/src/main/java/co/elastic/apm/agent/opentracing/impl/OpenTracingTextMapBridge.java +++ b/apm-agent-plugins/apm-opentracing-plugin/src/main/java/co/elastic/apm/agent/opentracing/impl/OpenTracingTextMapBridge.java @@ -55,11 +55,13 @@ public String getFirstHeader(String headerName, Iterable getHeaders(String headerName, Iterable> textMap) { - // Not supported for OpenTracing TextMap - return null; + public void forEach(String headerName, Iterable> carrier, S state, HeaderConsumer consumer) { + for (Map.Entry entry : carrier) { + if (entry.getKey().equalsIgnoreCase(headerName)) { + consumer.accept(entry.getValue(), state); + } + } } @Override diff --git a/apm-agent-plugins/apm-servlet-plugin/src/main/java/co/elastic/apm/agent/servlet/ServletInstrumentation.java b/apm-agent-plugins/apm-servlet-plugin/src/main/java/co/elastic/apm/agent/servlet/ServletInstrumentation.java index 4938c11036..3293349160 100644 --- a/apm-agent-plugins/apm-servlet-plugin/src/main/java/co/elastic/apm/agent/servlet/ServletInstrumentation.java +++ b/apm-agent-plugins/apm-servlet-plugin/src/main/java/co/elastic/apm/agent/servlet/ServletInstrumentation.java @@ -69,8 +69,7 @@ public ServletInstrumentation(ElasticApmTracer tracer) { ServletApiAdvice.init(tracer); servletTransactionCreationHelperManager = HelperClassManager.ForSingleClassLoader.of(tracer, "co.elastic.apm.agent.servlet.helper.ServletTransactionCreationHelperImpl", - "co.elastic.apm.agent.servlet.helper.RequestHeaderGetter", - "co.elastic.apm.agent.servlet.helper.RequestHeaderGetter$HeaderValuesIterator" + "co.elastic.apm.agent.servlet.helper.RequestHeaderGetter" ); } diff --git a/apm-agent-plugins/apm-servlet-plugin/src/main/java/co/elastic/apm/agent/servlet/helper/RequestHeaderGetter.java b/apm-agent-plugins/apm-servlet-plugin/src/main/java/co/elastic/apm/agent/servlet/helper/RequestHeaderGetter.java index bb7c3449ed..32a5156fc7 100644 --- a/apm-agent-plugins/apm-servlet-plugin/src/main/java/co/elastic/apm/agent/servlet/helper/RequestHeaderGetter.java +++ b/apm-agent-plugins/apm-servlet-plugin/src/main/java/co/elastic/apm/agent/servlet/helper/RequestHeaderGetter.java @@ -29,26 +29,13 @@ import javax.annotation.Nullable; import javax.servlet.http.HttpServletRequest; import java.util.Enumeration; -import java.util.Iterator; -import java.util.NoSuchElementException; -class RequestHeaderGetter implements TextHeaderGetter, Iterable { +class RequestHeaderGetter implements TextHeaderGetter { - // Caching in ThreadLocal in order to reuse the Enumeration-based iterator - private static final ThreadLocal threadLocalRequestHeaderGetter = new ThreadLocal<>(); + private static final RequestHeaderGetter INSTANCE = new RequestHeaderGetter(); static RequestHeaderGetter getInstance() { - RequestHeaderGetter requestHeaderGetter = threadLocalRequestHeaderGetter.get(); - if (requestHeaderGetter == null) { - requestHeaderGetter = new RequestHeaderGetter(); - threadLocalRequestHeaderGetter.set(requestHeaderGetter); - } - return requestHeaderGetter; - } - - private final HeaderValuesIterator headerValuesIterator = new HeaderValuesIterator(); - - private RequestHeaderGetter() { + return INSTANCE; } @Nullable @@ -57,37 +44,12 @@ public String getFirstHeader(String headerName, HttpServletRequest request) { return request.getHeader(headerName); } - @Nullable - @Override - public Iterable getHeaders(String headerName, HttpServletRequest request) { - headerValuesIterator.headerValues = request.getHeaders(headerName); - return this; - } - @Override - public Iterator iterator() { - return headerValuesIterator; - } - - private static final class HeaderValuesIterator implements Iterator { - - @Nullable - private Enumeration headerValues; - - @Override - public boolean hasNext() { - if (headerValues != null) { - return headerValues.hasMoreElements(); - } - return false; - } - - @Override - public String next() { - if (headerValues != null) { - return headerValues.nextElement(); - } - throw new NoSuchElementException("Header values Enumeration is not initialized"); + public void forEach(String headerName, HttpServletRequest carrier, S state, HeaderConsumer consumer) { + Enumeration headers = carrier.getHeaders(headerName); + while (headers.hasMoreElements()) { + consumer.accept(headers.nextElement(), state); } } + } diff --git a/apm-agent-plugins/apm-urlconnection-plugin/src/main/java/co/elastic/apm/agent/urlconnection/UrlConnectionPropertyAccessor.java b/apm-agent-plugins/apm-urlconnection-plugin/src/main/java/co/elastic/apm/agent/urlconnection/UrlConnectionPropertyAccessor.java index 74cce038af..37cb413f18 100644 --- a/apm-agent-plugins/apm-urlconnection-plugin/src/main/java/co/elastic/apm/agent/urlconnection/UrlConnectionPropertyAccessor.java +++ b/apm-agent-plugins/apm-urlconnection-plugin/src/main/java/co/elastic/apm/agent/urlconnection/UrlConnectionPropertyAccessor.java @@ -29,6 +29,7 @@ import javax.annotation.Nullable; import java.net.HttpURLConnection; +import java.util.List; public class UrlConnectionPropertyAccessor implements TextHeaderSetter, TextHeaderGetter { @@ -49,9 +50,13 @@ public String getFirstHeader(String headerName, HttpURLConnection urlConnection) return urlConnection.getRequestProperty(headerName); } - @Nullable @Override - public Iterable getHeaders(String headerName, HttpURLConnection urlConnection) { - return urlConnection.getRequestProperties().get(headerName); + public void forEach(String headerName, HttpURLConnection carrier, S state, HeaderConsumer consumer) { + List values = carrier.getRequestProperties().get(headerName); + if (values != null) { + for (int i = 0, size = values.size(); i < size; i++) { + consumer.accept(values.get(i), state); + } + } } } diff --git a/apm-opentracing/src/test/java/co/elastic/apm/opentracing/OpenTracingBridgeTest.java b/apm-opentracing/src/test/java/co/elastic/apm/opentracing/OpenTracingBridgeTest.java index c7e8502f2b..990d7e59c9 100644 --- a/apm-opentracing/src/test/java/co/elastic/apm/opentracing/OpenTracingBridgeTest.java +++ b/apm-opentracing/src/test/java/co/elastic/apm/opentracing/OpenTracingBridgeTest.java @@ -34,7 +34,9 @@ import io.opentracing.SpanContext; import io.opentracing.Tracer; import io.opentracing.log.Fields; -import io.opentracing.propagation.*; +import io.opentracing.propagation.Format; +import io.opentracing.propagation.TextMap; +import io.opentracing.propagation.TextMapAdapter; import io.opentracing.tag.Tags; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; @@ -457,7 +459,7 @@ void testInjectExtract() { final HashMap map = new HashMap<>(); apmTracer.inject(otSpan.context(), Format.Builtin.TEXT_MAP, new TextMapAdapter(map)); final TraceContext injectedContext = TraceContext.with64BitId(tracer); - assertThat(TraceContext.getFromTraceContextTextHeaders().asChildOf(injectedContext, map, TextHeaderMapAccessor.INSTANCE)).isTrue(); + assertThat(TraceContext.>getFromTraceContextTextHeaders().asChildOf(injectedContext, map, TextHeaderMapAccessor.INSTANCE)).isTrue(); assertThat(injectedContext.getTraceId().toString()).isEqualTo(traceId); assertThat(injectedContext.getParentId()).isEqualTo(transaction.getTraceContext().getId()); assertThat(injectedContext.isSampled()).isTrue();