Avoid tracestate error with sample rate provided through headers (#1808)

* allow to set identical sample rate multiple times * change check method name * make vertx version test more robust * ignore 2cnd value + log in debug * fail on double set + warn on double header * fix tests * early skip if sample rate is already set * revert micrometer plugin hack * update changelog * fix javadoc

Author: SylvainJuge

CHANGELOG.asciidoc

@@ -35,6 +35,7 @@ events. With the new `OVERRIDE` option, non-file logs can be ECS-reformatted aut
 ===== Bug fixes
 * Fix another error related to instrumentation plugins loading on Windows - {pull}1785[#1785]
 * Load Spring AMQP plugin- {pull}1784[#1784]
+* Avoid `IllegalStateException` when multiple `tracestate` headers are used - {pull}1808[#1808]
 
 [float]
 ===== Refactors

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

@@ -26,13 +26,17 @@
 
 import co.elastic.apm.agent.configuration.converter.RoundedDoubleConverter;
 import co.elastic.apm.agent.objectpool.Recyclable;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 import javax.annotation.Nullable;
 import java.util.ArrayList;
 import java.util.List;
 
 public class TraceState implements Recyclable {
 
+ private static final Logger log = LoggerFactory.getLogger(TraceState.class);
+
  private static final int DEFAULT_SIZE_LIMIT = 4096;
 
  private static final RoundedDoubleConverter DOUBLE_CONVERTER = RoundedDoubleConverter.withDefaultPrecision();
@@ -72,68 +76,109 @@ public void copyFrom(TraceState other) {
  rewriteBuffer.setLength(0);
  }
 
-
  public void addTextHeader(String headerValue) {
- int elasticEntryStartIndex = headerValue.indexOf(VENDOR_PREFIX);
-
- if (elasticEntryStartIndex >= 0) {
- // parsing (and maybe fixing) current tracestate required
- int entriesStart = headerValue.indexOf(SAMPLE_RATE_PREFIX, elasticEntryStartIndex);
- if (entriesStart >= 0) {
- int valueStart = entriesStart + 2;
- int valueEnd = valueStart;
- if (valueEnd < headerValue.length()) {
- char c = headerValue.charAt(valueEnd);
- while (valueEnd < headerValue.length() && c != VENDOR_SEPARATOR && c != ENTRY_SEPARATOR) {
- c = headerValue.charAt(valueEnd++);
- }
- if (valueEnd < headerValue.length()) {
- // end due to separator char that needs to be trimmed
- valueEnd--;
- }
- }
- double value;
- try {
- value = Double.parseDouble(headerValue.substring(valueStart, valueEnd));
- if (0 <= value && value <= 1.0) {
- // ensure proper rounding of sample rate to minimize storage
- // even if configuration should not allow this, any upstream value might require rounding
- double rounded = DOUBLE_CONVERTER.round(value);
- if (rounded != value) {
- // value needs to be re-written due to rounding
-
- rewriteBuffer.setLength(0);
- rewriteBuffer.append(headerValue, 0, valueStart);
- rewriteBuffer.append(rounded);
- rewriteBuffer.append(headerValue, valueEnd, headerValue.length());
- // we don't minimize allocation as re-writing should be an exception
- headerValue = rewriteBuffer.toString();
- }
- sampleRate = rounded;
- }
- } catch (NumberFormatException e) {
- // silently ignored
+ int vendorStart = headerValue.indexOf(VENDOR_PREFIX);
+
+ if (vendorStart < 0) {
+ // no ES entry
+ tracestate.add(headerValue);
+ return;
+ }
+
+ int vendorEnd = headerValue.indexOf(VENDOR_SEPARATOR, vendorStart);
+ if (vendorEnd < 0) {
+ vendorEnd = headerValue.length();
+ }
+
+ int sampleRateStart = headerValue.indexOf(SAMPLE_RATE_PREFIX, vendorStart);
+ if (sampleRateStart < 0) {
+ // no sample rate, rewrite
+ log.warn("invalid header, no sample rate {}", headerValue);
+ headerValue = rewriteRemoveInvalidHeader(headerValue, vendorStart, vendorEnd);
+ if (headerValue.length() > 0) {
+ tracestate.add(headerValue);
+ }
+ return;
+ }
+
+ int valueStart = sampleRateStart + 2;
+ int valueEnd = valueStart;
+ if (valueEnd < headerValue.length()) {
+ char c = headerValue.charAt(valueEnd);
+ while (valueEnd < headerValue.length() && c != VENDOR_SEPARATOR && c != ENTRY_SEPARATOR) {
+ c = headerValue.charAt(valueEnd++);
+ }
+ if (valueEnd < headerValue.length()) {
+ // end due to separator char that needs to be trimmed
+ valueEnd--;
+ }
+ }
+
+ String headerValueString = headerValue.substring(valueStart, valueEnd);
+ double value = Double.NaN;
+ try {
+ value = Double.parseDouble(headerValueString);
+ } catch (NumberFormatException e) {
+ // silently ignored
+ }
+
+ if (Double.isNaN(value) || value < 0 || value > 1) {
+ log.warn("invalid sample rate header {}", headerValueString);
+ headerValue = rewriteRemoveInvalidHeader(headerValue, vendorStart, vendorEnd);
+ } else {
+ if (!Double.isNaN(sampleRate)) {
+ log.warn("sample rate already set to {}, trying to set it to {} through header will be ignored", sampleRate, value);
+ headerValue = rewriteRemoveInvalidHeader(headerValue, vendorStart, vendorEnd);
+ } else {
+ // ensure proper rounding of sample rate to minimize storage
+ // even if configuration should not allow this, any upstream value might require rounding
+ double rounded = DOUBLE_CONVERTER.round(value);
+ if (rounded != value) {
+ // value needs to be re-written due to rounding
+ headerValue = rewriteRoundedHeader(headerValue, valueStart, valueEnd, rounded);
+ sampleRate = rounded;
+ } else {
+ sampleRate = value;
  }
  }
  }
 
- tracestate.add(headerValue);
+ if (!headerValue.isEmpty()) {
+ tracestate.add(headerValue);
+ }
+
+ }
+
+ private String rewriteRoundedHeader(String fullHeader, int valueStart, int valueEnd, double rounded) {
+ // we don't minimize allocation as re-writing should be an exception
+ rewriteBuffer.setLength(0);
+ rewriteBuffer.append(fullHeader, 0, valueStart);
+ rewriteBuffer.append(DOUBLE_CONVERTER.toString(rounded));
+ rewriteBuffer.append(fullHeader, valueEnd, fullHeader.length());
+ return rewriteBuffer.toString();
+ }
+
+ private String rewriteRemoveInvalidHeader(String fullHeader, int start, int end) {
+ rewriteBuffer.setLength(0);
+ if (start > 0) {
+ rewriteBuffer.append(fullHeader, 0, start);
+ }
+ if (end < fullHeader.length()) {
+ rewriteBuffer.append(fullHeader, end + 1, fullHeader.length());
+ }
+ return rewriteBuffer.toString();
  }
 
  /**
  * Sets value for trace state. Provided rate and string value are assumed to be correct and consistent
  *
  * @param rate sample rate
  * @param headerValue header value, as provided by a call to {@link #getHeaderValue(double)}
- * @throws IllegalStateException if sample rate has already been set
- * @throws IllegalArgumentException if rate has an invalid value
+ * @throws IllegalStateException if sample rate has already been set
  */
  public void set(double rate, String headerValue) {
  if (!Double.isNaN(sampleRate)) {
- // sample rate is set either explicitly from this method (for root transactions)
- // or through upstream header, thus there is no need to change after. This allows to only
- // write/rewrite headers once
- throw new IllegalStateException("sample rate has already been set from headers");
+ throw new IllegalStateException(String.format("sample rate already set to %f, trying to set it to %f", sampleRate, rate));
  }
 
  sampleRate = rate;
@@ -148,7 +193,10 @@ public void set(double rate, String headerValue) {
  * @return header value
  */
  public static String getHeaderValue(double sampleRate) {
- return FULL_PREFIX + sampleRate;
+ if (Double.isNaN(sampleRate) || sampleRate < 0 || sampleRate > 1) {
+ throw new IllegalArgumentException("invalid sample rate " + sampleRate);
+ }
+ return FULL_PREFIX + DOUBLE_CONVERTER.toString(DOUBLE_CONVERTER.round(sampleRate));
  }
 
  /**

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

@@ -548,19 +548,19 @@ void testNonTracedChildSpanWithoutTraceState() {
 
  @ParameterizedTest
  @CsvSource(delimiter = '|', value = {
- // invalid tracestate values: just assume no sample rate is provided
- "|NaN",
- "es=|NaN",
- "es=s|NaN",
- "es=s:|NaN",
- "es=s:a|NaN",
+ // invalid tracestate values: no sample rate and header fixed
+ "|NaN|",
+ "es=|NaN|",
+ "es=s|NaN|",
+ "es=s:|NaN|",
+ "es=s:a|NaN|",
  // valid tracestate values with sample rate
- "es=s:1|1",
- "es=s:0.42|0.42",
+ "es=s:1|1|es=s:1",
+ "es=s:0.42|0.42|es=s:0.42",
  // other vendors entries
- "a=123,es=s:0.42|0.42",
+ "a=123,es=s:0.42|0.42|a=123,es=s:0.42",
  })
- void checkExpectedSampleRate(@Nullable String traceState, double expectedRate) {
+ void checkExpectedSampleRate(@Nullable String traceState, double expectedRate, @Nullable String expectedHeader) {
  Map<String, String> headers = new HashMap<>();
  headers.put(TraceContext.W3C_TRACE_PARENT_TEXTUAL_HEADER_NAME, "00-0af7651916cd43dd8448eb211c80319c-b9c7c989f97918e1-01");
  if (null != traceState) {
@@ -577,8 +577,7 @@ void checkExpectedSampleRate(@Nullable String traceState, double expectedRate) {
  .isEqualTo(Double.valueOf(expectedRate));
 
  assertThat(child.getTraceState().toTextHeader())
- .describedAs("tracestate should be kept as-is")
- .isEqualTo(traceState);
+ .isEqualTo(expectedHeader);
 
  }