Add BaggageRestrictionManager

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

@@ -22,7 +22,7 @@
 
 package com.uber.jaeger;
 
-import com.uber.jaeger.baggage.BaggageValidity;
+import com.uber.jaeger.baggage.BaggageSetter;
 import io.opentracing.tag.Tags;
 import java.util.ArrayList;
 import java.util.Collections;
@@ -117,12 +117,12 @@ public List<LogData> getLogs() {
 
   @Override
   public Span setBaggageItem(String key, String value) {
+    if (key == null || value == null) {
+      return this;
+    }
+    BaggageSetter baggageValidity = this.getTracer().getBaggageRestrictionManager().getBaggageSetter(key);
     synchronized (this) {
-      if (key == null || value == null) {
-        return this;
-      }
-      BaggageValidity baggageValidity = this.getTracer().getBaggageRestrictionManager().isBaggageValid(key, value);
-      this.context = baggageValidity.sanitizeBaggage(this, key, value);
+      this.context = baggageValidity.setBaggage(this, key, value);
       return this;
     }
   }

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

@@ -450,13 +450,14 @@ public static final class Builder {
     private final Sampler sampler;
     private final Reporter reporter;
     private final PropagationRegistry registry = new PropagationRegistry();
-    private Metrics metrics;
+    private Metrics metrics = new Metrics(new StatsFactoryImpl(new NullStatsReporter()));
     private String serviceName;
     private Clock clock = new SystemClock();
     private Map<String, Object> tags = new HashMap<String, Object>();
     private boolean zipkinSharedRpcSpan;
     private ActiveSpanSource activeSpanSource = new ThreadLocalActiveSpanSource();
-    private BaggageRestrictionManager baggageRestrictionManager = new DefaultBaggageRestrictionManager();
+    private BaggageRestrictionManager baggageRestrictionManager =
+        new DefaultBaggageRestrictionManager(this.metrics);
 
     public Builder(String serviceName, Reporter reporter, Sampler sampler) {
       if (serviceName == null || serviceName.trim().length() == 0) {
@@ -465,7 +466,6 @@ public Builder(String serviceName, Reporter reporter, Sampler sampler) {
       this.serviceName = serviceName;
       this.reporter = reporter;
       this.sampler = sampler;
-      this.metrics = new Metrics(new StatsFactoryImpl(new NullStatsReporter()));
 
       TextMapCodec textMapCodec = new TextMapCodec(false);
       this.registerInjector(Format.Builtin.TEXT_MAP, textMapCodec);

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

@@ -25,5 +25,5 @@
 public abstract class BaggageRestrictionManager {
   static final int DEFAULT_MAX_VALUE_LENGTH = 2048;
 
-  public abstract BaggageValidity isBaggageValid(String key, String value);
+  public abstract BaggageSetter getBaggageSetter(String key);
 }

jaeger-core/src/main/java/com/uber/jaeger/baggage/BaggageValidity.java → jaeger-core/src/main/java/com/uber/jaeger/baggage/BaggageSetter.java

@@ -25,18 +25,19 @@
 import com.uber.jaeger.Span;
 import com.uber.jaeger.SpanContext;
 import com.uber.jaeger.metrics.Metrics;
-import lombok.Value;
 
 import java.util.HashMap;
 import java.util.Map;
 
+import lombok.Value;
+
 @Value(staticConstructor = "of")
-public class BaggageValidity {
+public class BaggageSetter {
   final boolean valid;
   final int maxValueLength;
+  final Metrics metrics;
 
-  public SpanContext sanitizeBaggage(Span span, String key, String value) {
-    Metrics metrics = span.getTracer().getMetrics();
+  public SpanContext setBaggage(Span span, String key, String value) {
     if (!this.isValid()) {
       metrics.baggageUpdateFailure.inc(1);
       this.logFields(span, key, value, null, false, true);

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

@@ -22,20 +22,21 @@
 
 package com.uber.jaeger.baggage;
 
-public class DefaultBaggageRestrictionManager extends BaggageRestrictionManager {
+import com.uber.jaeger.metrics.Metrics;
 
-  private final BaggageValidity baggageValidity;
+public class DefaultBaggageRestrictionManager extends BaggageRestrictionManager {
+  private final BaggageSetter baggageSetter;
 
-  public DefaultBaggageRestrictionManager() {
-    this(DEFAULT_MAX_VALUE_LENGTH);
+  public DefaultBaggageRestrictionManager(Metrics metrics) {
+    this(metrics, DEFAULT_MAX_VALUE_LENGTH);
   }
 
-  public DefaultBaggageRestrictionManager(int maxValueLength) {
-    baggageValidity = BaggageValidity.of(true, maxValueLength);
+  public DefaultBaggageRestrictionManager(Metrics metrics, int maxValueLength) {
+    baggageSetter = BaggageSetter.of(true, maxValueLength, metrics);
   }
 
   @Override
-  public BaggageValidity isBaggageValid(String key, String value) {
-    return baggageValidity;
+  public BaggageSetter getBaggageSetter(String key) {
+    return baggageSetter;
   }
 }

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

@@ -29,9 +29,11 @@
 import static org.mockito.Mockito.when;
 
 import com.uber.jaeger.baggage.BaggageRestrictionManager;
+import com.uber.jaeger.baggage.BaggageSetter;
 import com.uber.jaeger.baggage.DefaultBaggageRestrictionManager;
-import com.uber.jaeger.baggage.BaggageValidity;
 import com.uber.jaeger.metrics.InMemoryStatsReporter;
+import com.uber.jaeger.metrics.Metrics;
+import com.uber.jaeger.metrics.StatsFactoryImpl;
 import com.uber.jaeger.reporters.InMemoryReporter;
 import com.uber.jaeger.samplers.ConstSampler;
 import com.uber.jaeger.utils.Clock;
@@ -51,17 +53,19 @@ public class SpanTest {
   private Tracer tracer;
   private Span span;
   private InMemoryStatsReporter metricsReporter;
+  private Metrics metrics;
 
   @Before
   public void setUp() throws Exception {
     metricsReporter = new InMemoryStatsReporter();
     reporter = new InMemoryReporter();
     clock = mock(Clock.class);
+    metrics = new Metrics(new StatsFactoryImpl(metricsReporter));
     tracer =
         new Tracer.Builder("SamplerTest", reporter, new ConstSampler(true))
             .withStatsReporter(metricsReporter)
             .withClock(clock)
-            .withBaggageRestrictionManager(new DefaultBaggageRestrictionManager(15))
+            .withBaggageRestrictionManager(new DefaultBaggageRestrictionManager(metrics, 15))
             .build();
     span = (Span) tracer.buildSpan("some-operation").startManual();
   }
@@ -78,22 +82,24 @@ public void testSpanMetrics() {
 
   @Test
   public void testSetAndGetBaggageItem() {
-    String expected = "expected";
+    String value = "01234567890123456789";
+    String expected = "012345678901234";
     String key = "some.BAGGAGE";
-    span.setBaggageItem(key, expected);
+    span.setBaggageItem(key, value);
     assertEquals(expected, span.getBaggageItem(key));
 
     // Ensure the baggage was logged
-    this.assertBaggageLogs(span, key, expected, false, false);
+    this.assertBaggageLogs(span, key, expected, false, true);
 
+    assertEquals(
+        1L, metricsReporter.counters.get("jaeger.baggage-truncate").longValue());
     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 =
@@ -104,30 +110,15 @@ public void testInvalidBaggageKey() {
             .build();
     span = (Span) tracer.buildSpan("some-operation").startManual();
 
-    when(mgr.isBaggageValid(key, value)).thenReturn(BaggageValidity.of(false, 0));
+    when(mgr.getBaggageSetter(key)).thenReturn(BaggageSetter.of(false, 0, metrics));
 
-    span.setBaggageItem(key, value);
+    span.setBaggageItem(key, "luggage");
     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
   public void testSetBooleanTag() {
     Boolean expected = true;

jaeger-core/src/test/java/com/uber/jaeger/baggage/BaggageRestrictionManagerTest.java

@@ -24,18 +24,21 @@
 
 import static org.junit.Assert.assertEquals;
 
+import com.uber.jaeger.metrics.Metrics;
+import com.uber.jaeger.metrics.NullStatsReporter;
+import com.uber.jaeger.metrics.StatsFactoryImpl;
 import org.junit.Test;
 
 public class BaggageRestrictionManagerTest {
 
+
   @Test
   public void testIsBaggageValid() {
-    final String BAGGAGE_KEY = "key";
-    final String BAGGAGE_VALUE = "value";
-    final DefaultBaggageRestrictionManager undertest = new DefaultBaggageRestrictionManager();
+    final Metrics nullMetrics = new Metrics(new StatsFactoryImpl(new NullStatsReporter()));
+    final DefaultBaggageRestrictionManager undertest = new DefaultBaggageRestrictionManager(nullMetrics);
 
-    BaggageValidity actual = undertest.isBaggageValid(BAGGAGE_KEY, BAGGAGE_VALUE);
-    BaggageValidity expected = BaggageValidity.of(true, 2048);
+    BaggageSetter actual = undertest.getBaggageSetter("key");
+    BaggageSetter expected = BaggageSetter.of(true, 2048, nullMetrics);
     assertEquals(expected, actual);
   }
 }

jaeger-core/src/test/java/com/uber/jaeger/baggage/BaggageSetterTest.java

@@ -0,0 +1,116 @@
+/*
+ * 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 static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
+
+import com.uber.jaeger.LogData;
+import com.uber.jaeger.Span;
+import com.uber.jaeger.SpanContext;
+import com.uber.jaeger.Tracer;
+import com.uber.jaeger.metrics.InMemoryStatsReporter;
+import com.uber.jaeger.metrics.Metrics;
+import com.uber.jaeger.metrics.StatsFactoryImpl;
+import com.uber.jaeger.reporters.InMemoryReporter;
+import com.uber.jaeger.samplers.ConstSampler;
+
+import java.util.List;
+import java.util.Map;
+
+import org.junit.Before;
+import org.junit.Test;
+
+public class BaggageSetterTest {
+
+  private InMemoryReporter reporter;
+  private Tracer tracer;
+  private Span span;
+  private InMemoryStatsReporter metricsReporter;
+  private Metrics metrics;
+
+  private static final String KEY = "key";
+
+  @Before
+  public void setUp() throws Exception {
+    metricsReporter = new InMemoryStatsReporter();
+    reporter = new InMemoryReporter();
+    metrics = new Metrics(new StatsFactoryImpl(metricsReporter));
+    tracer =
+        new Tracer.Builder("SamplerTest", reporter, new ConstSampler(true))
+            .withStatsReporter(metricsReporter)
+            .withBaggageRestrictionManager(new DefaultBaggageRestrictionManager(metrics, 15))
+            .build();
+    span = (Span) tracer.buildSpan("some-operation").startManual();
+  }
+
+  @Test
+  public void testInvalidBaggage() {
+    BaggageSetter setter = BaggageSetter.of(false, 0, metrics);
+    final String value = "value";
+    SpanContext ctx = setter.setBaggage(span, KEY, value);
+
+    assertBaggageLogs(span, KEY, value, false, true);
+    assertNull(ctx.getBaggageItem(KEY));
+
+    assertEquals(
+        1L, metricsReporter.counters.get("jaeger.baggage-update.result=err").longValue());
+  }
+
+  @Test
+  public void testTruncatedBaggage() {
+    BaggageSetter setter = BaggageSetter.of(true, 5, metrics);
+    final String value = "0123456789";
+    final String expected = "01234";
+    SpanContext ctx = setter.setBaggage(span, KEY, value);
+
+    assertBaggageLogs(span, KEY, expected, true, false);
+    assertEquals(expected, ctx.getBaggageItem(KEY));
+
+    assertEquals(
+        1L, metricsReporter.counters.get("jaeger.baggage-truncate").longValue());
+    assertEquals(
+        1L, metricsReporter.counters.get("jaeger.baggage-update.result=ok").longValue());
+  }
+
+  private void assertBaggageLogs(
+      Span span,
+      String key,
+      String value,
+      boolean truncate,
+      boolean invalid
+  ) {
+    List<LogData> logs = span.getLogs();
+    assertEquals(false, logs.isEmpty());
+    Map<String, ?> fields = logs.get(logs.size() - 1).getFields();
+    assertEquals("baggage", fields.get("event"));
+    assertEquals(key, fields.get("key"));
+    assertEquals(value, fields.get("value"));
+    if (truncate) {
+      assertEquals("true", fields.get("truncated"));
+    }
+    if (invalid) {
+      assertEquals("true", fields.get("invalid"));
+    }
+  }
+}