use onEnding to add baggage keys that are added while the span is active

Author: zeitlinger

baggage-processor/src/main/java/io/opentelemetry/contrib/baggage/processor/BaggageSpanProcessor.java

@@ -9,14 +9,14 @@
 import io.opentelemetry.context.Context;
 import io.opentelemetry.sdk.trace.ReadWriteSpan;
 import io.opentelemetry.sdk.trace.ReadableSpan;
-import io.opentelemetry.sdk.trace.SpanProcessor;
+import io.opentelemetry.sdk.trace.internal.ExtendedSpanProcessor;
 import java.util.function.Predicate;
 
 /**
  * This span processor copies attributes stored in {@link Baggage} into each newly created {@link
  * io.opentelemetry.api.trace.Span}.
  */
-public class BaggageSpanProcessor implements SpanProcessor {
+public class BaggageSpanProcessor implements ExtendedSpanProcessor {
   private final Predicate<String> baggageKeyPredicate;
 
   /** use {@link #allowAllBaggageKeys()} instead */
@@ -39,19 +39,11 @@ public static BaggageSpanProcessor allowAllBaggageKeys() {
   }
 
   @Override
-  public void onStart(Context parentContext, ReadWriteSpan span) {
-    Baggage.fromContext(parentContext)
-        .forEach(
-            (s, baggageEntry) -> {
-              if (baggageKeyPredicate.test(s)) {
-                span.setAttribute(s, baggageEntry.getValue());
-              }
-            });
-  }
+  public void onStart(Context parentContext, ReadWriteSpan span) {}
 
   @Override
   public boolean isStartRequired() {
-    return true;
+    return false;
   }
 
   @Override
@@ -61,4 +53,20 @@ public void onEnd(ReadableSpan span) {}
   public boolean isEndRequired() {
     return false;
   }
+
+  @Override
+  public void onEnding(ReadWriteSpan span) {
+    Baggage.current()
+        .forEach(
+            (s, baggageEntry) -> {
+              if (baggageKeyPredicate.test(s)) {
+                span.setAttribute(s, baggageEntry.getValue());
+              }
+            });
+  }
+
+  @Override
+  public boolean isOnEndingRequired() {
+    return true;
+  }
 }

baggage-processor/src/test/java/io/opentelemetry/contrib/baggage/processor/BaggageProcessorCustomizerTest.java

@@ -160,7 +160,7 @@ public void test_baggageSpanProcessor_adds_attributes_to_spans(@Mock ReadWriteSp
     try (BaggageSpanProcessor processor =
         BaggageProcessorCustomizer.createBaggageSpanProcessor(Collections.singletonList("*"))) {
       try (Scope ignore = Baggage.current().toBuilder().put("key", "value").build().makeCurrent()) {
-        processor.onStart(Context.current(), span);
+        processor.onEnding(span);
         verify(span).setAttribute("key", "value");
       }
     }
@@ -177,7 +177,7 @@ public void test_baggageSpanProcessor_adds_attributes_to_spans_when_key_filter_m
               .put("other", "value")
               .build()
               .makeCurrent()) {
-        processor.onStart(Context.current(), span);
+        processor.onEnding(span);
         verify(span).setAttribute("key", "value");
         verify(span, Mockito.never()).setAttribute("other", "value");
       }

baggage-processor/src/test/java/io/opentelemetry/contrib/baggage/processor/BaggageSpanProcessorTest.java

@@ -6,7 +6,6 @@
 package io.opentelemetry.contrib.baggage.processor;
 
 import io.opentelemetry.api.baggage.Baggage;
-import io.opentelemetry.context.Context;
 import io.opentelemetry.context.Scope;
 import io.opentelemetry.sdk.trace.ReadWriteSpan;
 import java.util.regex.Pattern;
@@ -23,7 +22,7 @@ public class BaggageSpanProcessorTest {
   public void test_baggageSpanProcessor_adds_attributes_to_spans(@Mock ReadWriteSpan span) {
     try (BaggageSpanProcessor processor = BaggageSpanProcessor.allowAllBaggageKeys()) {
       try (Scope ignore = Baggage.current().toBuilder().put("key", "value").build().makeCurrent()) {
-        processor.onStart(Context.current(), span);
+        processor.onEnding(span);
         Mockito.verify(span).setAttribute("key", "value");
       }
     }
@@ -39,7 +38,7 @@ public void test_baggageSpanProcessor_adds_attributes_to_spans_when_key_filter_m
               .put("other", "value")
               .build()
               .makeCurrent()) {
-        processor.onStart(Context.current(), span);
+        processor.onEnding(span);
         Mockito.verify(span).setAttribute("key", "value");
         Mockito.verify(span, Mockito.never()).setAttribute("other", "value");
       }
@@ -58,7 +57,7 @@ public void test_baggageSpanProcessor_adds_attributes_to_spans_when_key_filter_m
               .put("other", "value")
               .build()
               .makeCurrent()) {
-        processor.onStart(Context.current(), span);
+        processor.onEnding(span);
         Mockito.verify(span).setAttribute("key", "value");
         Mockito.verify(span, Mockito.never()).setAttribute("other", "value");
       }