From 7fe57cd40d0a9b9a58c1174c0a5721ed2621b1ac Mon Sep 17 00:00:00 2001 From: Gregor Zeitlinger Date: Mon, 6 Jan 2025 12:16:40 +0100 Subject: [PATCH] use onEnding to add baggage keys that are added while the span is active --- .../processor/BaggageSpanProcessor.java | 32 ++++++++++++------- .../BaggageProcessorCustomizerTest.java | 4 +-- .../processor/BaggageSpanProcessorTest.java | 7 ++-- 3 files changed, 25 insertions(+), 18 deletions(-) diff --git a/baggage-processor/src/main/java/io/opentelemetry/contrib/baggage/processor/BaggageSpanProcessor.java b/baggage-processor/src/main/java/io/opentelemetry/contrib/baggage/processor/BaggageSpanProcessor.java index dbc1bd92c..ac5baf0d6 100644 --- a/baggage-processor/src/main/java/io/opentelemetry/contrib/baggage/processor/BaggageSpanProcessor.java +++ b/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 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; + } } diff --git a/baggage-processor/src/test/java/io/opentelemetry/contrib/baggage/processor/BaggageProcessorCustomizerTest.java b/baggage-processor/src/test/java/io/opentelemetry/contrib/baggage/processor/BaggageProcessorCustomizerTest.java index 9dcb9a4a7..eb2c05f35 100644 --- a/baggage-processor/src/test/java/io/opentelemetry/contrib/baggage/processor/BaggageProcessorCustomizerTest.java +++ b/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"); } diff --git a/baggage-processor/src/test/java/io/opentelemetry/contrib/baggage/processor/BaggageSpanProcessorTest.java b/baggage-processor/src/test/java/io/opentelemetry/contrib/baggage/processor/BaggageSpanProcessorTest.java index ca1180dcc..b1e70ab6a 100644 --- a/baggage-processor/src/test/java/io/opentelemetry/contrib/baggage/processor/BaggageSpanProcessorTest.java +++ b/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"); }