diff --git a/opentelemetry/src/main/java/io/grpc/opentelemetry/OpenTelemetryTracingModule.java b/opentelemetry/src/main/java/io/grpc/opentelemetry/OpenTelemetryTracingModule.java index 8c42a189ac2..997d10e3e09 100644 --- a/opentelemetry/src/main/java/io/grpc/opentelemetry/OpenTelemetryTracingModule.java +++ b/opentelemetry/src/main/java/io/grpc/opentelemetry/OpenTelemetryTracingModule.java @@ -38,6 +38,7 @@ import io.grpc.ServerStreamTracer; import io.grpc.opentelemetry.internal.OpenTelemetryConstants; import io.opentelemetry.api.OpenTelemetry; +import io.opentelemetry.api.baggage.Baggage; import io.opentelemetry.api.common.AttributesBuilder; import io.opentelemetry.api.trace.Span; import io.opentelemetry.api.trace.StatusCode; @@ -58,6 +59,11 @@ final class OpenTelemetryTracingModule { @VisibleForTesting final io.grpc.Context.Key otelSpan = io.grpc.Context.key("opentelemetry-span-key"); + + @VisibleForTesting + final io.grpc.Context.Key baggageKey = + io.grpc.Context.key("opentelemetry-baggage-key"); + @Nullable private static final AtomicIntegerFieldUpdater callEndedUpdater; @Nullable @@ -330,6 +336,8 @@ public ServerCall.Listener interceptCall(ServerCall(next.startCall(call, headers), serverCallContext); } diff --git a/opentelemetry/src/test/java/io/grpc/opentelemetry/OpenTelemetryTracingModuleTest.java b/opentelemetry/src/test/java/io/grpc/opentelemetry/OpenTelemetryTracingModuleTest.java index bca6be94b9f..389454878dd 100644 --- a/opentelemetry/src/test/java/io/grpc/opentelemetry/OpenTelemetryTracingModuleTest.java +++ b/opentelemetry/src/test/java/io/grpc/opentelemetry/OpenTelemetryTracingModuleTest.java @@ -59,6 +59,7 @@ import io.grpc.testing.GrpcCleanupRule; import io.grpc.testing.GrpcServerRule; import io.opentelemetry.api.OpenTelemetry; +import io.opentelemetry.api.baggage.Baggage; import io.opentelemetry.api.trace.Span; import io.opentelemetry.api.trace.SpanBuilder; import io.opentelemetry.api.trace.SpanId; @@ -750,6 +751,48 @@ public void onComplete() { } } + @Test + public void serverSpanPropagationInterceptor_propagatesBaggage() { + // 1. Arrange: Setup the module, interceptor, and mocks. + OpenTelemetryTracingModule tracingModule = new OpenTelemetryTracingModule( + openTelemetryRule.getOpenTelemetry()); + ServerInterceptor interceptor = tracingModule.getServerSpanPropagationInterceptor(); + + @SuppressWarnings("unchecked") + ServerCallHandler handler = mock(ServerCallHandler.class); + ServerCall call = new NoopServerCall<>(); + Metadata metadata = new Metadata(); + final AtomicReference capturedBaggage = new AtomicReference<>(); + + // Mock the handler to capture the Baggage from the current context when it's called. + when(handler.startCall(any(), any())).thenAnswer(invocation -> { + capturedBaggage.set(io.opentelemetry.api.baggage.Baggage.current()); + return mockServerCallListener; + }); + + // Create a test Span and Baggage to be propagated. + Span parentSpan = tracerRule.spanBuilder("parent-span").startSpan(); + io.opentelemetry.api.baggage.Baggage testBaggage = + io.opentelemetry.api.baggage.Baggage.builder().put("testKey", "testValue").build(); + + // Attach the Span and Baggage to the gRPC context. + io.grpc.Context grpcContext = io.grpc.Context.current() + .withValue(tracingModule.otelSpan, parentSpan) + .withValue(tracingModule.baggageKey, testBaggage); + + io.grpc.Context previous = grpcContext.attach(); + try { + // 2. Act: Call the interceptor. + interceptor.interceptCall(call, metadata, handler); + } finally { + grpcContext.detach(previous); + } + + // 3. Assert: Verify the handler was called and the correct Baggage was propagated. + verify(handler).startCall(same(call), same(metadata)); + assertEquals(testBaggage, capturedBaggage.get()); + } + @Test public void generateTraceSpanName() { assertEquals(