Narrow ExtendedSpanBuilder return types for chaining (#6514)

open-telemetry · Jun 24, 2024 · 0aacc55 · 0aacc55
1 parent a09eff6
commit 0aacc55
Show file tree

Hide file tree

Showing 3 changed files with 75 additions and 11 deletions.
diff --git a/api/incubator/src/main/java/io/opentelemetry/api/incubator/trace/ExtendedSpanBuilder.java b/api/incubator/src/main/java/io/opentelemetry/api/incubator/trace/ExtendedSpanBuilder.java
@@ -6,10 +6,17 @@
 package io.opentelemetry.api.incubator.trace;
 
 import io.opentelemetry.api.OpenTelemetry;
+import io.opentelemetry.api.common.AttributeKey;
+import io.opentelemetry.api.common.Attributes;
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.api.trace.SpanBuilder;
+import io.opentelemetry.api.trace.SpanContext;
+import io.opentelemetry.api.trace.SpanKind;
+import io.opentelemetry.context.Context;
 import io.opentelemetry.context.propagation.ContextPropagators;
+import java.time.Instant;
 import java.util.Map;
+import java.util.concurrent.TimeUnit;
 import java.util.function.BiConsumer;
 
 /** Extended {@link SpanBuilder} with experimental APIs. */
@@ -90,4 +97,60 @@ <T, E extends Throwable> T startAndCall(
    */
   <E extends Throwable> void startAndRun(
       SpanRunnable<E> runnable, BiConsumer<Span, Throwable> handleException) throws E;
+
+  /** {@inheritDoc} */
+  @Override
+  ExtendedSpanBuilder setParent(Context context);
+
+  /** {@inheritDoc} */
+  @Override
+  ExtendedSpanBuilder setNoParent();
+
+  /** {@inheritDoc} */
+  @Override
+  ExtendedSpanBuilder addLink(SpanContext spanContext);
+
+  /** {@inheritDoc} */
+  @Override
+  ExtendedSpanBuilder addLink(SpanContext spanContext, Attributes attributes);
+
+  /** {@inheritDoc} */
+  @Override
+  ExtendedSpanBuilder setAttribute(String key, String value);
+
+  /** {@inheritDoc} */
+  @Override
+  ExtendedSpanBuilder setAttribute(String key, long value);
+
+  /** {@inheritDoc} */
+  @Override
+  ExtendedSpanBuilder setAttribute(String key, double value);
+
+  /** {@inheritDoc} */
+  @Override
+  ExtendedSpanBuilder setAttribute(String key, boolean value);
+
+  /** {@inheritDoc} */
+  @Override
+  <T> ExtendedSpanBuilder setAttribute(AttributeKey<T> key, T value);
+
+  /** {@inheritDoc} */
+  @Override
+  default ExtendedSpanBuilder setAllAttributes(Attributes attributes) {
+    return (ExtendedSpanBuilder) SpanBuilder.super.setAllAttributes(attributes);
+  }
+
+  /** {@inheritDoc} */
+  @Override
+  ExtendedSpanBuilder setSpanKind(SpanKind spanKind);
+
+  /** {@inheritDoc} */
+  @Override
+  ExtendedSpanBuilder setStartTimestamp(long startTimestamp, TimeUnit unit);
+
+  /** {@inheritDoc} */
+  @Override
+  default ExtendedSpanBuilder setStartTimestamp(Instant startTimestamp) {
+    return (ExtendedSpanBuilder) SpanBuilder.super.setStartTimestamp(startTimestamp);
+  }
 }
diff --git a/...cubator/src/test/java/io/opentelemetry/api/incubator/trace/ExtendedTraceApiUsageTest.java b/...cubator/src/test/java/io/opentelemetry/api/incubator/trace/ExtendedTraceApiUsageTest.java
@@ -110,6 +110,7 @@ void startAndCallOrRun() {
     // Wrap the resetCheckout method in a span
     String cartId =
         ((ExtendedSpanBuilder) tracer.spanBuilder("reset_checkout_and_return"))
+            .setAttribute("key123", "val456")
             .startAndCall(() -> resetCheckoutAndReturn("abc123", /* throwException= */ false));
     assertThat(cartId).isEqualTo("abc123");
     // ...or runnable variation

diff --git a/sdk/trace/src/main/java/io/opentelemetry/sdk/trace/SdkSpanBuilder.java b/sdk/trace/src/main/java/io/opentelemetry/sdk/trace/SdkSpanBuilder.java
@@ -68,7 +68,7 @@ final class SdkSpanBuilder implements ExtendedSpanBuilder {
   }
 
   @Override
-  public SpanBuilder setParent(Context context) {
+  public ExtendedSpanBuilder setParent(Context context) {
     if (context == null) {
       return this;
     }
@@ -77,13 +77,13 @@ public SpanBuilder setParent(Context context) {
   }
 
   @Override
-  public SpanBuilder setNoParent() {
+  public ExtendedSpanBuilder setNoParent() {
     this.parent = Context.root();
     return this;
   }
 
   @Override
-  public SpanBuilder setSpanKind(SpanKind spanKind) {
+  public ExtendedSpanBuilder setSpanKind(SpanKind spanKind) {
     if (spanKind == null) {
       return this;
     }
@@ -92,7 +92,7 @@ public SpanBuilder setSpanKind(SpanKind spanKind) {
   }
 
   @Override
-  public SpanBuilder addLink(SpanContext spanContext) {
+  public ExtendedSpanBuilder addLink(SpanContext spanContext) {
     if (spanContext == null || !spanContext.isValid()) {
       return this;
     }
@@ -101,7 +101,7 @@ public SpanBuilder addLink(SpanContext spanContext) {
   }
 
   @Override
-  public SpanBuilder addLink(SpanContext spanContext, Attributes attributes) {
+  public ExtendedSpanBuilder addLink(SpanContext spanContext, Attributes attributes) {
     if (spanContext == null || !spanContext.isValid()) {
       return this;
     }
@@ -135,27 +135,27 @@ private void addLink(LinkData link) {
   }
 
   @Override
-  public SpanBuilder setAttribute(String key, String value) {
+  public ExtendedSpanBuilder setAttribute(String key, String value) {
     return setAttribute(stringKey(key), value);
   }
 
   @Override
-  public SpanBuilder setAttribute(String key, long value) {
+  public ExtendedSpanBuilder setAttribute(String key, long value) {
     return setAttribute(longKey(key), value);
   }
 
   @Override
-  public SpanBuilder setAttribute(String key, double value) {
+  public ExtendedSpanBuilder setAttribute(String key, double value) {
     return setAttribute(doubleKey(key), value);
   }
 
   @Override
-  public SpanBuilder setAttribute(String key, boolean value) {
+  public ExtendedSpanBuilder setAttribute(String key, boolean value) {
     return setAttribute(booleanKey(key), value);
   }
 
   @Override
-  public <T> SpanBuilder setAttribute(AttributeKey<T> key, T value) {
+  public <T> ExtendedSpanBuilder setAttribute(AttributeKey<T> key, T value) {
     if (key == null || key.getKey().isEmpty() || value == null) {
       return this;
     }
@@ -164,7 +164,7 @@ public <T> SpanBuilder setAttribute(AttributeKey<T> key, T value) {
   }
 
   @Override
-  public SpanBuilder setStartTimestamp(long startTimestamp, TimeUnit unit) {
+  public ExtendedSpanBuilder setStartTimestamp(long startTimestamp, TimeUnit unit) {
     if (startTimestamp < 0 || unit == null) {
       return this;
     }