baggage propagation

Author: AgraVator

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<Span> otelSpan = io.grpc.Context.key("opentelemetry-span-key");
+
+  @VisibleForTesting
+  final io.grpc.Context.Key<Baggage> baggageKey =
+      io.grpc.Context.key("opentelemetry-baggage-key");
+
   @Nullable
   private static final AtomicIntegerFieldUpdater<CallAttemptsTracerFactory> callEndedUpdater;
   @Nullable
@@ -330,6 +336,8 @@ public <ReqT, RespT> ServerCall.Listener<ReqT> interceptCall(ServerCall<ReqT, Re
         return next.startCall(call, headers);
       }
       Context serverCallContext = Context.current().with(span);
+      Baggage baggage = baggageKey.get(io.grpc.Context.current());
+      serverCallContext = serverCallContext.with(baggage);
       try (Scope scope = serverCallContext.makeCurrent()) {
         return new ContextServerCallListener<>(next.startCall(call, headers), serverCallContext);
       }

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<String, String> handler = mock(ServerCallHandler.class);
+    ServerCall<String, String> call = new NoopServerCall<>();
+    Metadata metadata = new Metadata();
+    final AtomicReference<Baggage> 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(