Add BaggageRestrictionManager

jaeger-core/src/main/java/com/uber/jaeger/Span.java

@@ -22,6 +22,7 @@
 
 package com.uber.jaeger;
 
+import com.uber.jaeger.baggage.SanitizedBaggage;
 import io.opentracing.tag.Tags;
 import java.util.ArrayList;
 import java.util.Collections;
@@ -117,21 +118,35 @@ public List<LogData> getLogs() {
   @Override
   public Span setBaggageItem(String key, String value) {
     synchronized (this) {
-      // TODO emit a metric whenever baggage is updated
+      if (key == null || value == null) {
+        return this;
+      }
+      SanitizedBaggage baggage = this.tracer.getBaggageRestrictionManager().sanitizeBaggage(key, value);
+      if (!baggage.isValid()) {
+        this.tracer.getMetrics().baggageUpdateFailure.inc(1);
+        return this;
+      }
+      if (baggage.isTruncated()) {
+        this.tracer.getMetrics().baggageTruncate.inc(1);
+      }
       String prevItem = this.getBaggageItem(key);
-      this.context = this.context.withBaggageItem(key, value);
+      this.context = this.context.withBaggageItem(key, baggage.getSanitizedValue());
+      this.tracer.getMetrics().baggageUpdateSuccess.inc(1);
       if (context.isSampled()) {
         Map<String, String> fields = new HashMap<String, String>();
         fields.put("event", "baggage");
         fields.put("key", key);
-        fields.put("value", value);
+        fields.put("value", baggage.getSanitizedValue());
         if (prevItem != null) {
           fields.put("override", "true");
         }
+        if (baggage.isTruncated()) {
+          fields.put("truncated", "true");
+        }
         return this.log(fields);
       }
+      return this;
     }
-    return this;
   }
 
   @Override

jaeger-core/src/main/java/com/uber/jaeger/Tracer.java

@@ -22,7 +22,8 @@
 
 package com.uber.jaeger;
 
-import com.uber.jaeger.Constants;
+import com.uber.jaeger.baggage.BaggageRestrictionManager;
+import com.uber.jaeger.baggage.DefaultBaggageRestrictionManager;
 import com.uber.jaeger.exceptions.UnsupportedFormatException;
 import com.uber.jaeger.metrics.Metrics;
 import com.uber.jaeger.metrics.NullStatsReporter;
@@ -74,6 +75,7 @@ public class Tracer implements io.opentracing.Tracer {
   private final Map<String, ?> tags;
   private final boolean zipkinSharedRpcSpan;
   private final ActiveSpanSource activeSpanSource;
+  private final BaggageRestrictionManager baggageRestrictionManager;
 
   private Tracer(
       String serviceName,
@@ -84,7 +86,8 @@ private Tracer(
       Metrics metrics,
       Map<String, Object> tags,
       boolean zipkinSharedRpcSpan,
-      ActiveSpanSource activeSpanSource) {
+      ActiveSpanSource activeSpanSource,
+      BaggageRestrictionManager baggageRestrictionManager) {
     this.serviceName = serviceName;
     this.reporter = reporter;
     this.sampler = sampler;
@@ -93,6 +96,7 @@ private Tracer(
     this.metrics = metrics;
     this.zipkinSharedRpcSpan = zipkinSharedRpcSpan;
     this.activeSpanSource = activeSpanSource;
+    this.baggageRestrictionManager = baggageRestrictionManager;
 
     this.version = loadVersion();
 
@@ -140,6 +144,10 @@ Reporter getReporter() {
     return reporter;
   }
 
+  BaggageRestrictionManager getBaggageRestrictionManager() {
+    return baggageRestrictionManager;
+  }
+
   void reportSpan(Span span) {
     reporter.report(span);
     metrics.spansFinished.inc(1);
@@ -448,6 +456,7 @@ public static final class Builder {
     private Map<String, Object> tags = new HashMap<String, Object>();
     private boolean zipkinSharedRpcSpan;
     private ActiveSpanSource activeSpanSource = new ThreadLocalActiveSpanSource();
+    private BaggageRestrictionManager baggageRestrictionManager = new DefaultBaggageRestrictionManager();
 
     public Builder(String serviceName, Reporter reporter, Sampler sampler) {
       if (serviceName == null || serviceName.trim().length() == 0) {
@@ -524,9 +533,14 @@ public Builder withTags(Map<String, String> tags) {
       return this;
     }
 
+    public Builder withBaggageRestrictionManager(BaggageRestrictionManager baggageRestrictionManager) {
+      this.baggageRestrictionManager = baggageRestrictionManager;
+      return this;
+    }
+
     public Tracer build() {
       return new Tracer(this.serviceName, reporter, sampler, registry, clock, metrics, tags,
-          zipkinSharedRpcSpan, activeSpanSource);
+          zipkinSharedRpcSpan, activeSpanSource, baggageRestrictionManager);
     }
   }
 

jaeger-core/src/main/java/com/uber/jaeger/baggage/BaggageRestrictionManager.java

@@ -0,0 +1,36 @@
+/*
+ * Copyright (c) 2017, Uber Technologies, Inc
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+
+package com.uber.jaeger.baggage;
+
+public abstract class BaggageRestrictionManager {
+  static final int DEFAULT_MAX_VALUE_LENGTH = 2048;
+
+  public abstract SanitizedBaggage sanitizeBaggage(String key, String value);
+
+  SanitizedBaggage truncateBaggage(String value, int maxValueLength) {
+    if (value.length() > maxValueLength) {
+      return SanitizedBaggage.of(true, value.substring(0, maxValueLength), true);
+    }
+    return SanitizedBaggage.of(true, value, false);
+  }
+}

jaeger-core/src/main/java/com/uber/jaeger/baggage/DefaultBaggageRestrictionManager.java

@@ -0,0 +1,41 @@
+/*
+ * Copyright (c) 2017, Uber Technologies, Inc
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+
+package com.uber.jaeger.baggage;
+
+public class DefaultBaggageRestrictionManager extends BaggageRestrictionManager {
+
+  private final int maxValueLength;
+
+  public DefaultBaggageRestrictionManager() {
+    this(DEFAULT_MAX_VALUE_LENGTH);
+  }
+
+  public DefaultBaggageRestrictionManager(int maxValueLength) {
+    this.maxValueLength = maxValueLength;
+  }
+
+  @Override
+  public SanitizedBaggage sanitizeBaggage(String key, String value) {
+    return truncateBaggage(value, maxValueLength);
+  }
+}

jaeger-core/src/main/java/com/uber/jaeger/baggage/SanitizedBaggage.java

@@ -0,0 +1,32 @@
+/*
+ * Copyright (c) 2017, Uber Technologies, Inc
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+
+package com.uber.jaeger.baggage;
+
+import lombok.Value;
+
+@Value(staticConstructor = "of")
+public class SanitizedBaggage {
+  final boolean valid;
+  final String sanitizedValue;
+  final boolean truncated;
+}

jaeger-core/src/main/java/com/uber/jaeger/metrics/Metrics.java

@@ -211,4 +211,22 @@ public static Metrics fromStatsReporter(StatsReporter reporter) {
   )
   // Number of times the Sampler failed to parse retrieved sampling strategy
   public Counter samplerParsingFailure;
+
+  @Metric(
+      name = "baggage-update",
+      tags = {@Tag(key = "result", value = "ok")}
+  )
+  // Number of times baggage was successfully written or updated on spans
+  public Counter baggageUpdateSuccess;
+
+  @Metric(
+      name = "baggage-update",
+      tags = {@Tag(key = "result", value = "err")}
+  )
+  // Number of times baggage failed to write or update on spans
+  public Counter baggageUpdateFailure;
+
+  @Metric(name = "baggage-truncate")
+  // Number of times baggage was truncated as per baggage restrictions
+  public Counter baggageTruncate;
 }

jaeger-core/src/test/java/com/uber/jaeger/SpanTest.java

@@ -28,6 +28,9 @@
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
+import com.uber.jaeger.baggage.BaggageRestrictionManager;
+import com.uber.jaeger.baggage.DefaultBaggageRestrictionManager;
+import com.uber.jaeger.baggage.SanitizedBaggage;
 import com.uber.jaeger.metrics.InMemoryStatsReporter;
 import com.uber.jaeger.reporters.InMemoryReporter;
 import com.uber.jaeger.samplers.ConstSampler;
@@ -58,6 +61,7 @@ public void setUp() throws Exception {
         new Tracer.Builder("SamplerTest", reporter, new ConstSampler(true))
             .withStatsReporter(metricsReporter)
             .withClock(clock)
+            .withBaggageRestrictionManager(new DefaultBaggageRestrictionManager(15))
             .build();
     span = (Span) tracer.buildSpan("some-operation").startManual();
   }
@@ -80,7 +84,48 @@ public void testSetAndGetBaggageItem() {
     assertEquals(expected, span.getBaggageItem(key));
 
     // Ensure the baggage was logged
-    this.assertBaggageLogs(span, key, expected, false);
+    this.assertBaggageLogs(span, key, expected, false, false);
+
+    assertEquals(
+        1L, metricsReporter.counters.get("jaeger.baggage-update.result=ok").longValue());
+  }
+
+  @Test
+  public void testInvalidBaggageKey() {
+    String key = "some.BAGGAGE";
+    String value = "luggage";
+    BaggageRestrictionManager mgr = mock(BaggageRestrictionManager.class);
+
+    tracer =
+        new Tracer.Builder("SamplerTest", reporter, new ConstSampler(true))
+            .withStatsReporter(metricsReporter)
+            .withClock(clock)
+            .withBaggageRestrictionManager(mgr)
+            .build();
+    span = (Span) tracer.buildSpan("some-operation").startManual();
+
+    when(mgr.sanitizeBaggage(key, value)).thenReturn(SanitizedBaggage.of(false, null, false));
+
+    span.setBaggageItem(key, value);
+    assertNull(span.getBaggageItem(key));
+
+    assertEquals(
+        1L, metricsReporter.counters.get("jaeger.baggage-update.result=err").longValue());
+  }
+
+  @Test
+  public void testTruncateBaggage() {
+    String value = "01234567890123456789";
+    String expected = "012345678901234";
+    String key = "some.BAGGAGE";
+    span.setBaggageItem(key, value);
+    assertEquals(expected, span.getBaggageItem(key));
+
+    // Ensure the baggage was logged
+    this.assertBaggageLogs(span, key, expected, false, true);
+
+    assertEquals(
+        1L, metricsReporter.counters.get("jaeger.baggage-truncate").longValue());
   }
 
   @Test
@@ -295,14 +340,14 @@ public void testSpanDetectsSamplingPriorityLessThanZero() {
   public void testBaggageOneReference() {
     io.opentracing.Span parent = tracer.buildSpan("foo").startManual();
     parent.setBaggageItem("foo", "bar");
-    this.assertBaggageLogs(parent, "foo", "bar", false);
+    this.assertBaggageLogs(parent, "foo", "bar", false, false);
 
     io.opentracing.Span child = tracer.buildSpan("foo")
         .asChildOf(parent)
         .startManual();
 
     child.setBaggageItem("a", "a");
-    this.assertBaggageLogs(child, "a", "a", false);
+    this.assertBaggageLogs(child, "a", "a", false, false);
 
     assertNull(parent.getBaggageItem("a"));
     assertEquals("a", child.getBaggageItem("a"));
@@ -313,20 +358,20 @@ public void testBaggageOneReference() {
   public void testBaggageMultipleReferences() {
     io.opentracing.Span parent1 = tracer.buildSpan("foo").startManual();
     parent1.setBaggageItem("foo", "bar");
-    this.assertBaggageLogs(parent1, "foo", "bar", false);
+    this.assertBaggageLogs(parent1, "foo", "bar", false, false);
     io.opentracing.Span parent2 = tracer.buildSpan("foo").startManual();
     parent2.setBaggageItem("foo2", "bar");
-    this.assertBaggageLogs(parent2, "foo2", "bar", false);
+    this.assertBaggageLogs(parent2, "foo2", "bar", false, false);
 
     io.opentracing.Span child = tracer.buildSpan("foo")
         .asChildOf(parent1)
         .addReference(References.FOLLOWS_FROM, parent2.context())
         .startManual();
 
     child.setBaggageItem("a", "a");
-    this.assertBaggageLogs(child, "a", "a", false);
+    this.assertBaggageLogs(child, "a", "a", false, false);
     child.setBaggageItem("foo2", "b");
-    this.assertBaggageLogs(child, "foo2", "b", true);
+    this.assertBaggageLogs(child, "foo2", "b", true, false);
 
     assertNull(parent1.getBaggageItem("a"));
     assertNull(parent2.getBaggageItem("a"));
@@ -350,17 +395,25 @@ public void testImmutableBaggage() {
     assertFalse(baggageIter.hasNext());
   }
 
-  private void assertBaggageLogs(io.opentracing.Span span, String key, String value, boolean override) {
+  private void assertBaggageLogs(
+      io.opentracing.Span span,
+      String key,
+      String value,
+      boolean override,
+      boolean truncate
+  ) {
     Span sp = (Span)span;
     List<LogData> logs = sp.getLogs();
     assertEquals(false, logs.isEmpty());
     Map<String, ?> fields = logs.get(logs.size() - 1).getFields();
-    assertEquals(override ? 4 : 3, fields.size());
     assertEquals("baggage", fields.get("event"));
     assertEquals(key, fields.get("key"));
     assertEquals(value, fields.get("value"));
     if (override) {
       assertEquals("true", fields.get("override"));
     }
+    if (truncate) {
+      assertEquals("true", fields.get("truncated"));
+    }
   }
 }