remove useless allocation in TraceState.copyOf() (#1508)

* use full header to minimize allocation

Author: SylvainJuge

CHANGELOG.asciidoc

@@ -28,6 +28,7 @@ endif::[]
 
 [float]
 ===== Bug fixes
+* Fix small memory allocation regression introduced with tracestate header {pull}1508[#1508]
 
 [float]
 ===== Refactors

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

@@ -25,6 +25,7 @@
 package co.elastic.apm.agent.impl.sampling;
 
 import co.elastic.apm.agent.impl.transaction.Id;
+import co.elastic.apm.agent.impl.transaction.TraceState;
 
 /**
  * This is a implementation of {@link Sampler} which always returns the same sampling decision.
@@ -36,12 +37,13 @@ public class ConstantSampler implements Sampler {
 
  private final boolean decision;
  private final double rate;
- private final String rateString;
+
+ private final String traceStateHeader;
 
  private ConstantSampler(boolean decision) {
  this.decision = decision;
- rate = decision ? 1.0d : 0.0d;
- rateString = Double.toString(rate);
+ this.rate = decision ? 1.0d : 0.0d;
+ this.traceStateHeader = TraceState.getHeaderValue(rate);
  }
 
  public static Sampler of(boolean decision) {
@@ -63,7 +65,7 @@ public double getSampleRate() {
  }
 
  @Override
- public String getSampleRateString() {
- return rateString;
+ public String getTraceStateHeader() {
+ return traceStateHeader;
  }
 }

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

@@ -56,13 +56,13 @@ public class ProbabilitySampler implements Sampler {
  private final double sampleRate;
 
  // Because header value only contains sampling rate, we can cache it here
- private final String rateString;
+ private final String traceStateHeader;
 
  private ProbabilitySampler(double samplingRate) {
- higherBound = (long) (Long.MAX_VALUE * samplingRate);
- lowerBound = -higherBound;
+ this.higherBound = (long) (Long.MAX_VALUE * samplingRate);
+ this.lowerBound = -higherBound;
  this.sampleRate = samplingRate;
- rateString = Double.toString(samplingRate);
+ this.traceStateHeader = TraceState.getHeaderValue(samplingRate);
  }
 
  public static Sampler of(double samplingRate) {
@@ -87,7 +87,7 @@ public double getSampleRate() {
  }
 
  @Override
- public String getSampleRateString() {
- return rateString;
+ public String getTraceStateHeader() {
+ return traceStateHeader;
  }
 }

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

@@ -57,8 +57,14 @@ public interface Sampler {
  */
  double getSampleRate();
 
+
  /**
- * @return current sample rate as constant string
+ * @return sample rate as (constant) header for context propagation
+ *
+ * <p>
+ * While the {@code tracestate} header is not related to sampler itself, putting this here allows to reuse the same
+ * {@link String} instance as long as the sample rate does not change to minimize allocation
+ * </p>
  */
- String getSampleRateString();
+ String getTraceStateHeader();
 }

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

@@ -427,7 +427,7 @@ public void asRootSpan(Sampler sampler) {
  transactionId.copyFrom(id);
  if (sampler.isSampled(traceId)) {
  flags = FLAG_RECORDED;
- traceState.set(sampler.getSampleRate(), sampler.getSampleRateString());
+ traceState.set(sampler.getSampleRate(), sampler.getTraceStateHeader());
  }
  clock.init();
  onMutation();

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

@@ -42,11 +42,9 @@ public class TraceState implements Recyclable {
 
  private int sizeLimit;
 
- private final StringBuilder header;
-
  private final StringBuilder rewriteBuffer;
 
- private final List<CharSequence> tracestate;
+ private final List<String> tracestate;
 
  /**
  * sample rate, {@link Double#NaN} if unknown or not set
@@ -58,20 +56,17 @@ public TraceState() {
  sizeLimit = DEFAULT_SIZE_LIMIT;
  tracestate = new ArrayList<>(1);
  rewriteBuffer = new StringBuilder();
- header = new StringBuilder(FULL_PREFIX.length());
  }
 
  public void copyFrom(TraceState other) {
  sampleRate = other.sampleRate;
  sizeLimit = other.sizeLimit;
  tracestate.clear();
- // copy and make sure we have the immutable variant
  for (int i = 0; i < other.tracestate.size(); i++) {
- tracestate.add(other.tracestate.get(i).toString());
+ //noinspection UseBulkOperation
+ tracestate.add(other.tracestate.get(i));
  }
  rewriteBuffer.setLength(0);
- header.setLength(0);
- header.append(other.header);
  }
 
  public void addTextHeader(String headerValue) {
@@ -125,23 +120,32 @@ public void addTextHeader(String headerValue) {
  /**
  * Sets value for trace state. Provided rate and string value are assumed to be correct and consistent
  *
- * @param rate sample rate
- * @param rateString rate written as a string, used to minimize allocation
+ * @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
  */
- public void set(double rate, String rateString) {
+ 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");
  }
+
  sampleRate = rate;
- header.setLength(0);
- header.append(FULL_PREFIX);
- header.append(rateString);
- tracestate.add(header);
+ tracestate.add(headerValue);
+ }
+
+ /**
+ * Generates the header value for the provided sample rate. As this method allocates a new string on each
+ * call, we should minimize the number of calls, ideally once for a given sample rate value.
+ *
+ * @param sampleRate sample rate
+ * @return header value
+ */
+ public static String getHeaderValue(double sampleRate) {
+ return FULL_PREFIX + sampleRate;
  }
 
  /**
@@ -169,7 +173,7 @@ public void resetState() {
  }
 
  public void setSizeLimit(int limit) {
- if(!tracestate.isEmpty()) {
+ if (!tracestate.isEmpty()) {
  throw new IllegalStateException("can't change size limit once headers have been added");
  }
  this.sizeLimit = limit;
@@ -189,27 +193,24 @@ private TextTracestateAppender() {
  }
 
  @Nullable
- public String join(List<? extends CharSequence> tracestate, int tracestateSizeLimit) {
- Object singleEntry = tracestate.size() != 1 ? null : tracestate.get(0);
- if (singleEntry instanceof String) {
- String entry = (String) singleEntry;
- if (entry.length() <= tracestateSizeLimit) {
- return entry;
- }
+ public String join(List<String> tracestate, int tracestateSizeLimit) {
+ String singleEntry = tracestate.size() != 1 ? null : tracestate.get(0);
+ if (singleEntry != null && singleEntry.length() <= tracestateSizeLimit) {
+ return singleEntry;
  }
 
  StringBuilder buffer = getTracestateBuffer();
  //noinspection ForLoopReplaceableByForEach
  for (int i = 0, size = tracestate.size(); i < size; i++) {
- CharSequence value = tracestate.get(i);
+ String value = tracestate.get(i);
  if (value != null) { // ignore null entries to allow removing entries without resizing collection
  appendTracestateHeaderValue(value, buffer, tracestateSizeLimit);
  }
  }
  return buffer.length() == 0 ? null : buffer.toString();
  }
 
- void appendTracestateHeaderValue(CharSequence headerValue, StringBuilder tracestateBuffer, int tracestateSizeLimit) {
+ void appendTracestateHeaderValue(String headerValue, StringBuilder tracestateBuffer, int tracestateSizeLimit) {
  int requiredLength = headerValue.length();
  boolean needsComma = tracestateBuffer.length() > 0;
  if (needsComma) {

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

@@ -53,12 +53,12 @@ void headerCaching(double rate) {
  // will indirectly test ConstantSampler as we delegate to it for 0 and 1 values
  sampler = ProbabilitySampler.of(rate);
 
- String rateString = sampler.getSampleRateString();
- assertThat(rateString).isEqualTo(Double.toString(rate));
+ String rateHeader = sampler.getTraceStateHeader();
+ assertThat(rateHeader).isEqualTo(TraceState.getHeaderValue(rate));
 
- assertThat(rateString)
- .describedAs("rate as string should return same instance on each call")
- .isSameAs(sampler.getSampleRateString());
+ assertThat(rateHeader)
+ .describedAs("sample rate header should return same instance on each call")
+ .isSameAs(sampler.getTraceStateHeader());
  }
 
  @Test

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

@@ -494,7 +494,7 @@ void testSetSampled() {
  assertThat(traceContext.getSampleRate()).isNaN();
 
  // sampled with sample rate
- traceContext.getTraceState().set(0.5d, Double.toString(0.5d));
+ traceContext.getTraceState().set(0.5d, TraceState.getHeaderValue(0.5d));
 
  assertThat(traceContext.isSampled()).isTrue();
  assertThat(traceContext.getSampleRate()).isEqualTo(0.5d);
@@ -518,7 +518,7 @@ private TraceContext createRootSpan(double sampleRate){
  Sampler sampler = mock(Sampler.class);
  when(sampler.isSampled(any(Id.class))).thenReturn(true);
  when(sampler.getSampleRate()).thenReturn(sampleRate);
- when(sampler.getSampleRateString()).thenReturn(Double.toString(sampleRate));
+ when(sampler.getTraceStateHeader()).thenReturn(TraceState.getHeaderValue(sampleRate));
 
  traceContext.asRootSpan(sampler);
  return traceContext;

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

@@ -75,7 +75,7 @@ void createEmpty() {
  "17|one=two_three=four,five=six,seven=eight_nine=ten,eleven-twelve|one=two,nine=ten", // just fits
  })
  void sizeLimit(int limit, String headers, @Nullable String expected) {
- Function<String,String> replaceSpaces = s -> s == null ? null : s.replace('X', ' ');
+ Function<String, String> replaceSpaces = s -> s == null ? null : s.replace('X', ' ');
 
  traceState.setSizeLimit(limit);
  for (String h : headers.split("_")) {
@@ -89,9 +89,14 @@ void sizeLimit(int limit, String headers, @Nullable String expected) {
 
  @Test
  void addSampleRate() {
- traceState.set(0.5d, Double.toString(0.5d));
+ String headerValue = TraceState.getHeaderValue(0.5d);
+ assertThat(headerValue).isEqualTo("es=s:0.5");
+ traceState.set(0.5d, headerValue);
  assertThat(traceState.getSampleRate()).isEqualTo(0.5d);
- assertThat(traceState.toTextHeader()).isEqualTo("es=s:0.5");
+
+ assertThat(traceState.toTextHeader())
+ .describedAs("should reuse the same string without allocating a new one")
+ .isSameAs(headerValue);
  }
 
  @Test
@@ -123,15 +128,16 @@ void sampleRateAddedToHeaders() {
  traceState.addTextHeader("aa=1|2|3");
  traceState.addTextHeader("bb=4|5|6");
 
- traceState.set(0.444d, Double.toString(0.444d));
+ traceState.set(0.444d, TraceState.getHeaderValue(0.444d));
  assertThat(traceState.getSampleRate()).isEqualTo(0.444d);
  assertThat(traceState.toTextHeader()).isEqualTo("aa=1|2|3,bb=4|5|6,es=s:0.444");
  }
 
  @Test
  void buildCopy() {
  TraceState other = new TraceState();
- other.set(0.2d, Double.toString(0.2d));
+
+ other.set(0.2d, TraceState.getHeaderValue(0.2d));
  other.addTextHeader("aa=1_2");
 
  traceState.copyFrom(other);
@@ -162,7 +168,7 @@ void unknownKeysAreIgnored(String header, String rewrittenHeader) {
  "es=s:aa"
  })
  void invalidValuesIgnored(String header) {
- traceState.addTextHeader( header);
+ traceState.addTextHeader(header);
  assertThat(traceState.getSampleRate()).isNaN();
  }
 
@@ -177,15 +183,4 @@ void appliesRoundingOnUpstreamHeader(String headerRate, Double expectedRate) {
  assertThat(traceState.toTextHeader()).isEqualTo("es=s:" + expectedRate);
  }
 
- @Test
- void useProvidedValueForTextHeader() {
- double rate = 0.43d;
- String headerValue = Double.toString(rate + 0.01d);
- traceState.set(rate, headerValue);
- assertThat(traceState.getSampleRate()).isEqualTo(rate);
-
- String textHeader = traceState.toTextHeader();
- assertThat(textHeader).isEqualTo("es=s:0.44"); // not exactly the same on purpose
- }
-
 }