Handle baggage null values better (#308)

* This closes #256 by ignoring attempts to set a new baggage value of null, but also removing existing baggage when their value has been updated as null Signed-off-by: Babak Mozaffari <[email protected]> * Using span.setBaggage API in unit test, since it has additional null checks Signed-off-by: Babak Mozaffari <[email protected]>
jaegertracing · Feb 1, 2018 · c0409b9 · c0409b9
1 parent 8dabf7a
commit c0409b9
Show file tree

Hide file tree

Showing 4 changed files with 34 additions and 3 deletions.
diff --git a/jaeger-core/src/main/java/com/uber/jaeger/Span.java b/jaeger-core/src/main/java/com/uber/jaeger/Span.java
@@ -118,7 +118,8 @@ public List<LogData> getLogs() {
 
   @Override
   public Span setBaggageItem(String key, String value) {
-    if (key == null || value == null) {
+    if (key == null || (value == null && context.getBaggageItem(key) == null)) {
+      //Ignore attempts to add new baggage items with null values, they're not accessible anyway
       return this;
     }
     synchronized (this) {

diff --git a/jaeger-core/src/main/java/com/uber/jaeger/SpanContext.java b/jaeger-core/src/main/java/com/uber/jaeger/SpanContext.java
@@ -125,7 +125,11 @@ public static SpanContext contextFromString(String value)
 
   public SpanContext withBaggageItem(String key, String val) {
     Map<String, String> newBaggage = new HashMap<String, String>(this.baggage);
-    newBaggage.put(key, val);
+    if (val == null) {
+      newBaggage.remove(key);
+    } else {
+      newBaggage.put(key, val);
+    }
     return new SpanContext(traceId, spanId, parentId, flags, newBaggage, debugId);
   }
 

diff --git a/jaeger-core/src/main/java/com/uber/jaeger/baggage/BaggageSetter.java b/jaeger-core/src/main/java/com/uber/jaeger/baggage/BaggageSetter.java
@@ -58,7 +58,7 @@ public SpanContext setBaggage(Span span, String key, String value) {
       logFields(span, key, value, prevItem, truncated, restriction.isKeyAllowed());
       return span.context();
     }
-    if (value.length() > restriction.getMaxValueLength()) {
+    if (value != null && value.length() > restriction.getMaxValueLength()) {
       truncated = true;
       value = value.substring(0, restriction.getMaxValueLength());
       metrics.baggageTruncate.inc(1);

diff --git a/jaeger-core/src/test/java/com/uber/jaeger/baggage/BaggageSetterTest.java b/jaeger-core/src/test/java/com/uber/jaeger/baggage/BaggageSetterTest.java
@@ -124,6 +124,32 @@ public void testUnsampledSpan() {
     assertNull(span.getLogs());
   }
 
+  @Test
+  public void testBaggageNullValueTolerated() {
+    when(mgr.getRestriction(SERVICE, KEY)).thenReturn(Restriction.of(true, 5));
+    final String value = null;
+    SpanContext ctx = setter.setBaggage(span, KEY, value);
+
+    assertBaggageLogs(span, KEY, null, false, false, false);
+    assertNull(ctx.getBaggageItem(KEY));
+  }
+
+  @Test
+  public void testBaggageNullRemoveValue() {
+    when(mgr.getRestriction(SERVICE, KEY)).thenReturn(Restriction.of(true, 5));
+    final String value = "value";
+    Span originalSpan = span.setBaggageItem(KEY, value);
+    assertEquals(value, originalSpan.getBaggageItem(KEY));
+    Span child = (Span) tracer.buildSpan("some-operation").asChildOf(originalSpan).startManual();
+    child = child.setBaggageItem(KEY, null);
+
+    assertBaggageLogs(child, KEY, null, false, true, false);
+    assertNull(child.getBaggageItem(KEY));
+
+    assertEquals(
+            2L, metricsReporter.counters.get("jaeger.baggage-update.result=ok").longValue());
+  }
+
   private void assertBaggageLogs(
       Span span,
       String key,