openTelemetry: propagate baggage

Author: AgraVator

opentelemetry/build.gradle

@@ -15,11 +15,15 @@ dependencies {
             libraries.auto.value.annotations
 
     testImplementation project(':grpc-testing'),
+            project(':grpc-testing-proto'),
             project(':grpc-inprocess'),
+            project(':grpc-stub'),
+            project(':grpc-protobuf'),
             testFixtures(project(':grpc-core')),
             testFixtures(project(':grpc-api')),
             libraries.opentelemetry.sdk.testing,
             libraries.assertj.core // opentelemetry.sdk.testing uses compileOnly for assertj
+            libraries.protobuf.java
 
     annotationProcessor libraries.auto.value
 

opentelemetry/src/main/java/io/grpc/opentelemetry/GrpcOpenTelemetry.java

@@ -179,12 +179,14 @@ public void configureChannelBuilder(ManagedChannelBuilder<?> builder) {
    * @param serverBuilder the server builder to configure
    */
   public void configureServerBuilder(ServerBuilder<?> serverBuilder) {
-    serverBuilder.addStreamTracerFactory(openTelemetryMetricsModule.getServerTracerFactory());
+    /* To ensure baggage propagation to metrics, we need the tracing
+    tracers to be initialised before metrics */
     if (ENABLE_OTEL_TRACING) {
       serverBuilder.addStreamTracerFactory(
           openTelemetryTracingModule.getServerTracerFactory());
       serverBuilder.intercept(openTelemetryTracingModule.getServerSpanPropagationInterceptor());
     }
+    serverBuilder.addStreamTracerFactory(openTelemetryMetricsModule.getServerTracerFactory());
   }
 
   @VisibleForTesting

opentelemetry/src/main/java/io/grpc/opentelemetry/OpenTelemetryMetricsModule.java

@@ -18,6 +18,7 @@
 
 import static com.google.common.base.Preconditions.checkNotNull;
 import static io.grpc.opentelemetry.internal.OpenTelemetryConstants.BACKEND_SERVICE_KEY;
+import static io.grpc.opentelemetry.internal.OpenTelemetryConstants.BAGGAGE_KEY;
 import static io.grpc.opentelemetry.internal.OpenTelemetryConstants.LOCALITY_KEY;
 import static io.grpc.opentelemetry.internal.OpenTelemetryConstants.METHOD_KEY;
 import static io.grpc.opentelemetry.internal.OpenTelemetryConstants.STATUS_KEY;
@@ -44,7 +45,9 @@
 import io.grpc.Status;
 import io.grpc.Status.Code;
 import io.grpc.StreamTracer;
+import io.opentelemetry.api.baggage.Baggage;
 import io.opentelemetry.api.common.AttributesBuilder;
+import io.opentelemetry.context.Context;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
@@ -94,8 +97,8 @@ final class OpenTelemetryMetricsModule {
   private final ImmutableList<OpenTelemetryPlugin> plugins;
 
   OpenTelemetryMetricsModule(Supplier<Stopwatch> stopwatchSupplier,
-      OpenTelemetryMetricsResource resource, Collection<String> optionalLabels,
-      List<OpenTelemetryPlugin> plugins) {
+                             OpenTelemetryMetricsResource resource,
+                             Collection<String> optionalLabels, List<OpenTelemetryPlugin> plugins) {
     this.resource = checkNotNull(resource, "resource");
     this.stopwatchSupplier = checkNotNull(stopwatchSupplier, "stopwatchSupplier");
     this.localityEnabled = optionalLabels.contains(LOCALITY_KEY.getKey());
@@ -128,6 +131,14 @@ static String recordMethodName(String fullMethodName, boolean isGeneratedMethod)
     return isGeneratedMethod ? fullMethodName : "other";
   }
 
+  private static Context otelContextWithBaggage() {
+    Baggage baggage = BAGGAGE_KEY.get(io.grpc.Context.current());
+    if (baggage == null) {
+      return Context.current();
+    }
+    return Context.current().with(baggage);
+  }
+
   private static final class ClientTracer extends ClientStreamTracer {
     @Nullable private static final AtomicLongFieldUpdater<ClientTracer> outboundWireSizeUpdater;
     @Nullable private static final AtomicLongFieldUpdater<ClientTracer> inboundWireSizeUpdater;
@@ -243,6 +254,7 @@ public void streamClosed(Status status) {
     }
 
     void recordFinishedAttempt() {
+      Context otelContext = otelContextWithBaggage();
       AttributesBuilder builder = io.opentelemetry.api.common.Attributes.builder()
           .put(METHOD_KEY, fullMethodName)
           .put(TARGET_KEY, target)
@@ -268,15 +280,15 @@ void recordFinishedAttempt() {
 
       if (module.resource.clientAttemptDurationCounter() != null ) {
         module.resource.clientAttemptDurationCounter()
-            .record(attemptNanos * SECONDS_PER_NANO, attribute);
+            .record(attemptNanos * SECONDS_PER_NANO, attribute, otelContext);
       }
       if (module.resource.clientTotalSentCompressedMessageSizeCounter() != null) {
         module.resource.clientTotalSentCompressedMessageSizeCounter()
-            .record(outboundWireSize, attribute);
+            .record(outboundWireSize, attribute, otelContext);
       }
       if (module.resource.clientTotalReceivedCompressedMessageSizeCounter() != null) {
         module.resource.clientTotalReceivedCompressedMessageSizeCounter()
-            .record(inboundWireSize, attribute);
+            .record(inboundWireSize, attribute, otelContext);
       }
     }
   }
@@ -408,6 +420,7 @@ void callEnded(Status status) {
     }
 
     void recordFinishedCall() {
+      Context otelContext = otelContextWithBaggage();
       if (attemptsPerCall.get() == 0) {
         ClientTracer tracer = newClientTracer(null);
         tracer.attemptNanos = attemptDelayStopwatch.elapsed(TimeUnit.NANOSECONDS);
@@ -429,15 +442,17 @@ void recordFinishedCall() {
             callLatencyNanos * SECONDS_PER_NANO,
             baseAttributes.toBuilder()
                 .put(STATUS_KEY, status.getCode().toString())
-                .build()
+                .build(),
+            otelContext
         );
       }
 
       // Retry counts
       if (module.resource.clientCallRetriesCounter() != null) {
         long retriesPerCall = Math.max(attemptsPerCall.get() - 1, 0);
         if (retriesPerCall > 0) {
-          module.resource.clientCallRetriesCounter().record(retriesPerCall, baseAttributes);
+          module.resource.clientCallRetriesCounter()
+              .record(retriesPerCall, baseAttributes, otelContext);
         }
       }
 
@@ -446,24 +461,25 @@ void recordFinishedCall() {
         long hedges = hedgedAttemptsPerCall.get();
         if (hedges > 0) {
           module.resource.clientCallHedgesCounter()
-              .record(hedges, baseAttributes);
+              .record(hedges, baseAttributes, otelContext);
         }
       }
 
       // Transparent Retry counts
       if (module.resource.clientCallTransparentRetriesCounter() != null) {
         long transparentRetries = transparentRetriesPerCall.get();
         if (transparentRetries > 0) {
-          module.resource.clientCallTransparentRetriesCounter().record(
-              transparentRetries, baseAttributes);
+          module.resource.clientCallTransparentRetriesCounter()
+              .record(transparentRetries, baseAttributes, otelContext);
         }
       }
 
       // Retry delay
       if (module.resource.clientCallRetryDelayCounter() != null) {
         module.resource.clientCallRetryDelayCounter().record(
             retryDelayNanos * SECONDS_PER_NANO,
-            baseAttributes
+            baseAttributes,
+            otelContext
         );
       }
     }
@@ -562,6 +578,7 @@ public void inboundWireSize(long bytes) {
      */
     @Override
     public void streamClosed(Status status) {
+      Context otelContext = otelContextWithBaggage();
       if (streamClosedUpdater != null) {
         if (streamClosedUpdater.getAndSet(this, 1) != 0) {
           return;
@@ -584,15 +601,15 @@ public void streamClosed(Status status) {
 
       if (module.resource.serverCallDurationCounter() != null) {
         module.resource.serverCallDurationCounter()
-            .record(elapsedTimeNanos * SECONDS_PER_NANO, attributes);
+            .record(elapsedTimeNanos * SECONDS_PER_NANO, attributes, otelContext);
       }
       if (module.resource.serverTotalSentCompressedMessageSizeCounter() != null) {
         module.resource.serverTotalSentCompressedMessageSizeCounter()
-            .record(outboundWireSize, attributes);
+            .record(outboundWireSize, attributes, otelContext);
       }
       if (module.resource.serverTotalReceivedCompressedMessageSizeCounter() != null) {
         module.resource.serverTotalReceivedCompressedMessageSizeCounter()
-            .record(inboundWireSize, attributes);
+            .record(inboundWireSize, attributes, otelContext);
       }
     }
   }

opentelemetry/src/main/java/io/grpc/opentelemetry/OpenTelemetryTracingModule.java

@@ -19,6 +19,7 @@
 import static com.google.common.base.Preconditions.checkNotNull;
 import static io.grpc.ClientStreamTracer.NAME_RESOLUTION_DELAYED;
 import static io.grpc.internal.GrpcUtil.IMPLEMENTATION_VERSION;
+import static io.grpc.opentelemetry.internal.OpenTelemetryConstants.BAGGAGE_KEY;
 
 import com.google.common.annotations.VisibleForTesting;
 import io.grpc.Attributes;
@@ -60,10 +61,6 @@ 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
@@ -248,13 +245,15 @@ private final class ServerTracer extends ServerStreamTracer {
     private final Span span;
     volatile int streamClosed;
     private int seqNo;
+    private Baggage baggage;
 
-    ServerTracer(String fullMethodName, @Nullable Span remoteSpan) {
+    ServerTracer(String fullMethodName, @Nullable Span remoteSpan, Baggage baggage) {
       checkNotNull(fullMethodName, "fullMethodName");
       this.span =
           otelTracer.spanBuilder(generateTraceSpanName(true, fullMethodName))
               .setParent(remoteSpan == null ? null : Context.current().with(remoteSpan))
               .startSpan();
+      this.baggage = baggage;
     }
 
     /**
@@ -280,7 +279,9 @@ public void streamClosed(io.grpc.Status status) {
 
     @Override
     public io.grpc.Context filterContext(io.grpc.Context context) {
-      return context.withValue(otelSpan, span);
+      return context
+          .withValue(otelSpan, span)
+          .withValue(BAGGAGE_KEY, baggage);
     }
 
     @Override
@@ -320,7 +321,8 @@ public ServerStreamTracer newServerStreamTracer(String fullMethodName, Metadata
       if (remoteSpan == Span.getInvalid()) {
         remoteSpan = null;
       }
-      return new ServerTracer(fullMethodName, remoteSpan);
+      Baggage baggage = Baggage.fromContext(context);
+      return new ServerTracer(fullMethodName, remoteSpan, baggage);
     }
   }
 
@@ -335,9 +337,12 @@ public <ReqT, RespT> ServerCall.Listener<ReqT> interceptCall(ServerCall<ReqT, Re
             + "tracing must be set.");
         return next.startCall(call, headers);
       }
-      Context serverCallContext = Context.current().with(span);
-      Baggage baggage = baggageKey.get(io.grpc.Context.current());
-      serverCallContext = serverCallContext.with(baggage);
+      Context serverCallContext = Context.current();
+      serverCallContext = serverCallContext.with(span);
+      Baggage baggage = BAGGAGE_KEY.get(io.grpc.Context.current());
+      if (baggage != null) {
+        serverCallContext = serverCallContext.with(baggage);
+      }
       try (Scope scope = serverCallContext.makeCurrent()) {
         return new ContextServerCallListener<>(next.startCall(call, headers), serverCallContext);
       }

opentelemetry/src/main/java/io/grpc/opentelemetry/internal/OpenTelemetryConstants.java

@@ -16,7 +16,9 @@
 
 package io.grpc.opentelemetry.internal;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableList;
+import io.opentelemetry.api.baggage.Baggage;
 import io.opentelemetry.api.common.AttributeKey;
 import java.util.List;
 
@@ -42,6 +44,10 @@ public final class OpenTelemetryConstants {
   public static final AttributeKey<String> SECURITY_LEVEL_KEY =
       AttributeKey.stringKey("grpc.security_level");
 
+  @VisibleForTesting
+  public static final io.grpc.Context.Key<Baggage> BAGGAGE_KEY =
+      io.grpc.Context.key("opentelemetry-baggage-key");
+
   public static final List<Double> LATENCY_BUCKETS =
       ImmutableList.of(
           0d,     0.00001d, 0.00005d, 0.0001d, 0.0003d, 0.0006d, 0.0008d, 0.001d, 0.002d,