Add span status customizer (#15288)

Author: laurit

instrumentation-api-incubator/src/main/java/io/opentelemetry/instrumentation/api/incubator/instrumenter/InstrumenterCustomizer.java

@@ -11,6 +11,7 @@
 import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
 import io.opentelemetry.instrumentation.api.instrumenter.OperationMetrics;
 import io.opentelemetry.instrumentation.api.instrumenter.SpanNameExtractor;
+import io.opentelemetry.instrumentation.api.instrumenter.SpanStatusExtractor;
 import io.opentelemetry.instrumentation.api.semconv.http.HttpClientAttributesExtractor;
 import java.util.function.Function;
 import java.util.function.UnaryOperator;
@@ -106,6 +107,16 @@ default InstrumenterCustomizer setSpanNameExtractor(
   InstrumenterCustomizer setSpanNameExtractor(
       UnaryOperator<SpanNameExtractor<?>> spanNameExtractor);
 
+  /**
+   * Sets a transformer function that will modify the {@link SpanStatusExtractor}. This allows
+   * customizing how span statuses are generated for the instrumented operations.
+   *
+   * @param spanStatusExtractor function that transforms the original span status extractor
+   * @return this InstrumenterCustomizer for method chaining
+   */
+  InstrumenterCustomizer setSpanStatusExtractor(
+      UnaryOperator<SpanStatusExtractor<?, ?>> spanStatusExtractor);
+
   /** Types of instrumentation. */
   enum InstrumentationType {
     HTTP_CLIENT,

instrumentation-api-incubator/src/main/java/io/opentelemetry/instrumentation/api/incubator/instrumenter/internal/InstrumenterCustomizerImpl.java

@@ -10,6 +10,7 @@
 import io.opentelemetry.instrumentation.api.instrumenter.ContextCustomizer;
 import io.opentelemetry.instrumentation.api.instrumenter.OperationMetrics;
 import io.opentelemetry.instrumentation.api.instrumenter.SpanNameExtractor;
+import io.opentelemetry.instrumentation.api.instrumenter.SpanStatusExtractor;
 import io.opentelemetry.instrumentation.api.internal.InternalInstrumenterCustomizer;
 import io.opentelemetry.instrumentation.api.internal.SpanKey;
 import java.util.HashMap;
@@ -87,4 +88,11 @@ public InstrumenterCustomizer setSpanNameExtractor(
     customizer.setSpanNameExtractor(spanNameExtractor);
     return this;
   }
+
+  @Override
+  public InstrumenterCustomizer setSpanStatusExtractor(
+      UnaryOperator<SpanStatusExtractor<?, ?>> spanStatusExtractor) {
+    customizer.setSpanStatusExtractor(spanStatusExtractor);
+    return this;
+  }
 }

instrumentation-api/src/main/java/io/opentelemetry/instrumentation/api/instrumenter/InstrumenterBuilder.java

@@ -442,6 +442,14 @@ public void setSpanNameExtractor(
               builder.spanNameExtractor =
                   spanNameExtractorTransformer.apply(builder.spanNameExtractor);
             }
+
+            @Override
+            public void setSpanStatusExtractor(
+                UnaryOperator<SpanStatusExtractor<? super REQUEST, ? super RESPONSE>>
+                    spanStatusExtractorTransformer) {
+              builder.spanStatusExtractor =
+                  spanStatusExtractorTransformer.apply(builder.spanStatusExtractor);
+            }
           });
     }
   }

instrumentation-api/src/main/java/io/opentelemetry/instrumentation/api/internal/InternalInstrumenterCustomizer.java

@@ -9,6 +9,7 @@
 import io.opentelemetry.instrumentation.api.instrumenter.ContextCustomizer;
 import io.opentelemetry.instrumentation.api.instrumenter.OperationMetrics;
 import io.opentelemetry.instrumentation.api.instrumenter.SpanNameExtractor;
+import io.opentelemetry.instrumentation.api.instrumenter.SpanStatusExtractor;
 import java.util.function.UnaryOperator;
 
 /**
@@ -32,4 +33,8 @@ void addAttributesExtractors(
 
   void setSpanNameExtractor(
       UnaryOperator<SpanNameExtractor<? super REQUEST>> spanNameExtractorTransformer);
+
+  void setSpanStatusExtractor(
+      UnaryOperator<SpanStatusExtractor<? super REQUEST, ? super RESPONSE>>
+          spanStatusExtractorTransformer);
 }

instrumentation-api/src/test/java/io/opentelemetry/instrumentation/api/internal/InstrumentationCustomizerTest.java

@@ -21,6 +21,7 @@
 import io.opentelemetry.api.trace.SpanContext;
 import io.opentelemetry.api.trace.SpanId;
 import io.opentelemetry.api.trace.SpanKind;
+import io.opentelemetry.api.trace.StatusCode;
 import io.opentelemetry.context.Context;
 import io.opentelemetry.instrumentation.api.incubator.instrumenter.InstrumenterCustomizer.InstrumentationType;
 import io.opentelemetry.instrumentation.api.incubator.instrumenter.InstrumenterCustomizerProvider;
@@ -359,6 +360,50 @@ void testSetSpanNameExtractor() {
                             .hasAttributes(Attributes.empty())));
   }
 
+  @Test
+  void testSetSpanStatusExtractor() {
+    AtomicBoolean customizerCalled = new AtomicBoolean();
+    setCustomizer(
+        customizer -> {
+          customizerCalled.set(true);
+          customizer.setSpanStatusExtractor(
+              unused ->
+                  (spanStatusBuilder, request, response, error) ->
+                      spanStatusBuilder.setStatus(StatusCode.OK));
+        });
+
+    Instrumenter<Map<String, String>, Map<String, String>> instrumenter =
+        Instrumenter.<Map<String, String>, Map<String, String>>builder(
+                otelTesting.getOpenTelemetry(), "test", unused -> "span")
+            .setSpanStatusExtractor(
+                (spanStatusBuilder, request, response, error) ->
+                    spanStatusBuilder.setStatus(StatusCode.ERROR))
+            .buildInstrumenter();
+
+    assertThat(customizerCalled).isTrue();
+
+    Context context = instrumenter.start(Context.root(), REQUEST);
+    SpanContext spanContext = Span.fromContext(context).getSpanContext();
+    assertThat(spanContext.isValid()).isTrue();
+
+    instrumenter.end(context, REQUEST, RESPONSE, null);
+
+    otelTesting
+        .assertTraces()
+        .hasTracesSatisfyingExactly(
+            trace ->
+                trace.hasSpansSatisfyingExactly(
+                    span ->
+                        span.hasName("span")
+                            .hasKind(SpanKind.INTERNAL)
+                            .hasInstrumentationScopeInfo(InstrumentationScopeInfo.create("test"))
+                            .hasTraceId(spanContext.getTraceId())
+                            .hasSpanId(spanContext.getSpanId())
+                            .hasParentSpanId(SpanId.getInvalid())
+                            .hasStatus(StatusData.ok())
+                            .hasAttributes(Attributes.empty())));
+  }
+
   static class AttributesExtractor1
       implements AttributesExtractor<Map<String, String>, Map<String, String>> {