fix server listener set span in context

interop-testing/build.gradle

@@ -13,6 +13,7 @@ dependencies {
     implementation project(path: ':grpc-alts', configuration: 'shadow'),
             project(':grpc-auth'),
             project(':grpc-census'),
+            project(':grpc-opentelemetry'),
             project(':grpc-gcp-csm-observability'),
             project(':grpc-netty'),
             project(':grpc-okhttp'),

interop-testing/src/test/java/io/grpc/testing/integration/OpenTelemetryContextFilterTest.java

@@ -0,0 +1,131 @@
+/*
+ * Copyright 2024 The gRPC Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.grpc.testing.integration;
+
+import static org.junit.Assert.assertEquals;
+
+import io.grpc.ForwardingServerCallListener;
+import io.grpc.InsecureServerCredentials;
+import io.grpc.ManagedChannelBuilder;
+import io.grpc.Metadata;
+import io.grpc.ServerBuilder;
+import io.grpc.ServerCall;
+import io.grpc.ServerCallHandler;
+import io.grpc.ServerInterceptor;
+import io.grpc.netty.NettyChannelBuilder;
+import io.grpc.netty.NettyServerBuilder;
+import io.grpc.opentelemetry.GrpcOpenTelemetry;
+import io.grpc.testing.integration.Messages.SimpleRequest;
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.Tracer;
+import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.Scope;
+import io.opentelemetry.context.propagation.ContextPropagators;
+import io.opentelemetry.sdk.OpenTelemetrySdk;
+import io.opentelemetry.sdk.trace.SdkTracerProvider;
+import java.util.concurrent.atomic.AtomicReference;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class OpenTelemetryContextFilterTest extends AbstractInteropTest {
+  private final OpenTelemetrySdk openTelemetrySdk = OpenTelemetrySdk.builder()
+      .setTracerProvider(SdkTracerProvider.builder().build())
+      .setPropagators(ContextPropagators.create(W3CTraceContextPropagator.getInstance()))
+      .build();
+  private final Tracer tracer = openTelemetrySdk.getTracer("grpc-java");
+  private final GrpcOpenTelemetry grpcOpenTelemetry = GrpcOpenTelemetry.newBuilder()
+      .sdk(openTelemetrySdk).build();
+  private final AtomicReference<Span> applicationSpan = new AtomicReference<>();
+
+  @Override
+  protected ServerBuilder<?> getServerBuilder() {
+    NettyServerBuilder builder = NettyServerBuilder.forPort(0, InsecureServerCredentials.create())
+        .maxInboundMessageSize(AbstractInteropTest.MAX_MESSAGE_SIZE);
+    builder.intercept(new ServerInterceptor() {
+      @Override
+      public <ReqT, RespT> ServerCall.Listener<ReqT> interceptCall(ServerCall<ReqT, RespT> call,
+          Metadata headers, ServerCallHandler<ReqT, RespT> next) {
+        ServerCall.Listener<ReqT> listener = next.startCall(call, headers);
+        return new ForwardingServerCallListener<ReqT>() {
+          @Override
+          protected ServerCall.Listener<ReqT> delegate() {
+            return listener;
+          }
+
+          @Override
+          public void onMessage(ReqT request) {
+            applicationSpan.set(tracer.spanBuilder("InteropTest.Application.Span").startSpan());
+            delegate().onMessage(request);
+          }
+
+          @Override
+          public void onHalfClose() {
+            if (applicationSpan.get() != null) {
+              applicationSpan.get().end();
+            }
+            delegate().onHalfClose();
+          }
+
+          @Override
+          public void onCancel() {
+            if (applicationSpan.get() != null) {
+              applicationSpan.get().end();
+            }
+            delegate().onCancel();
+          }
+
+          @Override
+          public void onComplete() {
+            if (applicationSpan.get() != null) {
+              applicationSpan.get().end();
+            }
+            delegate().onComplete();
+          }
+        };
+      }
+    });
+    grpcOpenTelemetry.configureServerBuilder(builder);
+    return builder;
+  }
+
+  @Override
+  protected boolean metricsExpected() {
+    return false;
+  }
+
+  @Override
+  protected ManagedChannelBuilder<?> createChannelBuilder() {
+    NettyChannelBuilder builder = NettyChannelBuilder.forAddress(getListenAddress())
+        .maxInboundMessageSize(AbstractInteropTest.MAX_MESSAGE_SIZE)
+        .usePlaintext();
+    grpcOpenTelemetry.configureChannelBuilder(builder);
+    return builder;
+  }
+
+  @Test
+  public void otelSpanContextPropagation() {
+    Span parentSpan = tracer.spanBuilder("Test.interopTest").startSpan();
+    try (Scope scope = Context.current().with(parentSpan).makeCurrent()) {
+      blockingStub.unaryCall(SimpleRequest.getDefaultInstance());
+    }
+    assertEquals(parentSpan.getSpanContext().getTraceId(),
+        applicationSpan.get().getSpanContext().getTraceId());
+  }
+}

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

@@ -68,6 +68,7 @@ public Stopwatch get() {
   private final boolean disableDefault;
   private final OpenTelemetryMetricsResource resource;
   private final OpenTelemetryMetricsModule openTelemetryMetricsModule;
+  private final OpenTelemetryTracingModule openTelemetryTracingModule;
   private final List<String> optionalLabels;
   private final MetricSink sink;
 
@@ -88,6 +89,7 @@ private GrpcOpenTelemetry(Builder builder) {
     this.optionalLabels = ImmutableList.copyOf(builder.optionalLabels);
     this.openTelemetryMetricsModule = new OpenTelemetryMetricsModule(
         STOPWATCH_SUPPLIER, resource, optionalLabels, builder.plugins);
+    this.openTelemetryTracingModule = new OpenTelemetryTracingModule(openTelemetrySdk);
     this.sink = new OpenTelemetryMetricSink(meter, enableMetrics, disableDefault, optionalLabels);
   }
 
@@ -152,6 +154,7 @@ public void configureChannelBuilder(ManagedChannelBuilder<?> builder) {
     InternalManagedChannelBuilder.addMetricSink(builder, sink);
     InternalManagedChannelBuilder.interceptWithTarget(
         builder, openTelemetryMetricsModule::getClientInterceptor);
+    builder.intercept(openTelemetryTracingModule.getClientInterceptor());
   }
 
   /**
@@ -161,6 +164,9 @@ public void configureChannelBuilder(ManagedChannelBuilder<?> builder) {
    */
   public void configureServerBuilder(ServerBuilder<?> serverBuilder) {
     serverBuilder.addStreamTracerFactory(openTelemetryMetricsModule.getServerTracerFactory());
+    serverBuilder.addStreamTracerFactory(
+        openTelemetryTracingModule.getServerTracerFactory());
+    serverBuilder.intercept(openTelemetryTracingModule.getServerSpanPropagationInterceptor());
   }
 
   @VisibleForTesting

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

@@ -28,6 +28,7 @@
 import io.grpc.ClientStreamTracer;
 import io.grpc.ForwardingClientCall.SimpleForwardingClientCall;
 import io.grpc.ForwardingClientCallListener.SimpleForwardingClientCallListener;
+import io.grpc.ForwardingServerCallListener;
 import io.grpc.Metadata;
 import io.grpc.MethodDescriptor;
 import io.grpc.ServerCall;
@@ -315,9 +316,59 @@ final class TracingServerSpanPropagationInterceptor implements ServerInterceptor
     public <ReqT, RespT> ServerCall.Listener<ReqT> interceptCall(ServerCall<ReqT, RespT> call,
         Metadata headers, ServerCallHandler<ReqT, RespT> next) {
       Span span = otelSpan.get(io.grpc.Context.current());
-      try (Scope scope = Context.current().with(span).makeCurrent()) {
+      if (span == null) {
+        logger.log(Level.FINE, "Server span not found. ServerTracerFactory for server "
+            + "tracing must be set.");
         return next.startCall(call, headers);
       }
+      try (Scope scope = Context.current().with(span).makeCurrent()) {
+        return new AttachSpanServerCallListener<>(next.startCall(call, headers), span);
+      }
+    }
+  }
+
+  private static class AttachSpanServerCallListener<ReqT> extends
+      ForwardingServerCallListener.SimpleForwardingServerCallListener<ReqT> {
+    private final Span span;
+
+    protected AttachSpanServerCallListener(ServerCall.Listener<ReqT> delegate, Span span) {
+      super(delegate);
+      this.span = checkNotNull(span, "span");
+    }
+
+    @Override
+    public void onMessage(ReqT message) {
+      try (Scope scope = Context.current().with(span).makeCurrent()) {
+        delegate().onMessage(message);
+      }
+    }
+
+    @Override
+    public void onHalfClose() {
+      try (Scope scope = Context.current().with(span).makeCurrent()) {
+        delegate().onHalfClose();
+      }
+    }
+
+    @Override
+    public void onCancel() {
+      try (Scope scope = Context.current().with(span).makeCurrent()) {
+        delegate().onCancel();
+      }
+    }
+
+    @Override
+    public void onComplete() {
+      try (Scope scope = Context.current().with(span).makeCurrent()) {
+        delegate().onComplete();
+      }
+    }
+
+    @Override
+    public void onReady() {
+      try (Scope scope = Context.current().with(span).makeCurrent()) {
+        delegate().onReady();
+      }
     }
   }