Set custom gRPC client/server span name extractor (#5244)

* Set custom gRPC client/server span name extractor * Fix imports * Fix compilation * Fix style rule * Update to use Functions instead * Fix formatting
open-telemetry · Feb 24, 2022 · bf16564 · bf16564
1 parent 6d58704
commit bf16564
Showing 1 changed file with 42 additions and 3 deletions.
diff --git a/.../library/src/main/java/io/opentelemetry/instrumentation/grpc/v1_6/GrpcTracingBuilder.java b/.../library/src/main/java/io/opentelemetry/instrumentation/grpc/v1_6/GrpcTracingBuilder.java
@@ -12,13 +12,15 @@
 import io.opentelemetry.instrumentation.api.instrumenter.InstrumenterBuilder;
 import io.opentelemetry.instrumentation.api.instrumenter.PeerServiceAttributesExtractor;
 import io.opentelemetry.instrumentation.api.instrumenter.SpanKindExtractor;
+import io.opentelemetry.instrumentation.api.instrumenter.SpanNameExtractor;
 import io.opentelemetry.instrumentation.api.instrumenter.net.NetClientAttributesExtractor;
 import io.opentelemetry.instrumentation.api.instrumenter.net.NetServerAttributesExtractor;
 import io.opentelemetry.instrumentation.grpc.v1_6.internal.GrpcNetClientAttributesGetter;
 import io.opentelemetry.instrumentation.grpc.v1_6.internal.GrpcNetServerAttributesGetter;
 import io.opentelemetry.semconv.trace.attributes.SemanticAttributes;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.function.Function;
 import java.util.stream.Stream;
 import javax.annotation.Nullable;
 
@@ -30,6 +32,14 @@ public final class GrpcTracingBuilder {
   private final OpenTelemetry openTelemetry;
   @Nullable private String peerService;
 
+  @Nullable
+  private Function<SpanNameExtractor<GrpcRequest>, ? extends SpanNameExtractor<? super GrpcRequest>>
+      clientSpanNameExtractorTransformer;
+
+  @Nullable
+  private Function<SpanNameExtractor<GrpcRequest>, ? extends SpanNameExtractor<? super GrpcRequest>>
+      serverSpanNameExtractorTransformer;
+
   private final List<AttributesExtractor<? super GrpcRequest, ? super Status>>
       additionalExtractors = new ArrayList<>();
 
@@ -49,9 +59,26 @@ public GrpcTracingBuilder addAttributeExtractor(
     return this;
   }
 
+  /** Sets custom client {@link SpanNameExtractor} via transform function. */
+  public GrpcTracingBuilder setClientSpanNameExtractor(
+      Function<SpanNameExtractor<GrpcRequest>, ? extends SpanNameExtractor<? super GrpcRequest>>
+          clientSpanNameExtractor) {
+    this.clientSpanNameExtractorTransformer = clientSpanNameExtractor;
+    return this;
+  }
+
+  /** Sets custom server {@link SpanNameExtractor} via transform function. */
+  public GrpcTracingBuilder setServerSpanNameExtractor(
+      Function<SpanNameExtractor<GrpcRequest>, ? extends SpanNameExtractor<? super GrpcRequest>>
+          serverSpanNameExtractor) {
+    this.serverSpanNameExtractorTransformer = serverSpanNameExtractor;
+    return this;
+  }
+
   /** Sets the {@code peer.service} attribute for http client spans. */
-  public void setPeerService(String peerService) {
+  public GrpcTracingBuilder setPeerService(String peerService) {
     this.peerService = peerService;
+    return this;
   }
 
   /**
@@ -67,10 +94,22 @@ public GrpcTracingBuilder setCaptureExperimentalSpanAttributes(
 
   /** Returns a new {@link GrpcTracing} with the settings of this {@link GrpcTracingBuilder}. */
   public GrpcTracing build() {
+    SpanNameExtractor<GrpcRequest> originalSpanNameExtractor = new GrpcSpanNameExtractor();
+
+    SpanNameExtractor<? super GrpcRequest> clientSpanNameExtractor = originalSpanNameExtractor;
+    if (clientSpanNameExtractorTransformer != null) {
+      clientSpanNameExtractor = clientSpanNameExtractorTransformer.apply(originalSpanNameExtractor);
+    }
+
+    SpanNameExtractor<? super GrpcRequest> serverSpanNameExtractor = originalSpanNameExtractor;
+    if (serverSpanNameExtractorTransformer != null) {
+      serverSpanNameExtractor = serverSpanNameExtractorTransformer.apply(originalSpanNameExtractor);
+    }
+
     InstrumenterBuilder<GrpcRequest, Status> clientInstrumenterBuilder =
-        Instrumenter.builder(openTelemetry, INSTRUMENTATION_NAME, new GrpcSpanNameExtractor());
+        Instrumenter.builder(openTelemetry, INSTRUMENTATION_NAME, clientSpanNameExtractor);
     InstrumenterBuilder<GrpcRequest, Status> serverInstrumenterBuilder =
-        Instrumenter.builder(openTelemetry, INSTRUMENTATION_NAME, new GrpcSpanNameExtractor());
+        Instrumenter.builder(openTelemetry, INSTRUMENTATION_NAME, serverSpanNameExtractor);
 
     Stream.of(clientInstrumenterBuilder, serverInstrumenterBuilder)
         .forEach(