Spring Integration library instrumentation

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

@@ -55,7 +55,9 @@ public static <REQUEST, RESPONSE> InstrumenterBuilder<REQUEST, RESPONSE> newBuil
   private final SpanNameExtractor<? super REQUEST> spanNameExtractor;
   private final SpanKindExtractor<? super REQUEST> spanKindExtractor;
   private final SpanStatusExtractor<? super REQUEST, ? super RESPONSE> spanStatusExtractor;
-  private final List<? extends AttributesExtractor<? super REQUEST, ? super RESPONSE>> extractors;
+  private final List<? extends AttributesExtractor<? super REQUEST, ? super RESPONSE>>
+      attributesExtractors;
+  private final List<? extends SpanLinkExtractor<? super REQUEST>> spanLinkExtractors;
   private final List<? extends RequestListener> requestListeners;
   private final ErrorCauseExtractor errorCauseExtractor;
   @Nullable private final StartTimeExtractor<REQUEST> startTimeExtractor;
@@ -68,7 +70,8 @@ public static <REQUEST, RESPONSE> InstrumenterBuilder<REQUEST, RESPONSE> newBuil
     this.spanNameExtractor = builder.spanNameExtractor;
     this.spanKindExtractor = builder.spanKindExtractor;
     this.spanStatusExtractor = builder.spanStatusExtractor;
-    this.extractors = new ArrayList<>(builder.attributesExtractors);
+    this.attributesExtractors = new ArrayList<>(builder.attributesExtractors);
+    this.spanLinkExtractors = new ArrayList<>(builder.spanLinkExtractors);
     this.requestListeners = new ArrayList<>(builder.requestListeners);
     this.errorCauseExtractor = builder.errorCauseExtractor;
     this.startTimeExtractor = builder.startTimeExtractor;
@@ -119,8 +122,12 @@ public Context start(Context parentContext, REQUEST request) {
       spanBuilder.setStartTimestamp(startTimeExtractor.extract(request));
     }
 
+    for (SpanLinkExtractor<? super REQUEST> extractor : spanLinkExtractors) {
+      spanBuilder.addLink(extractor.extract(parentContext, request));
+    }
+
     UnsafeAttributes attributesBuilder = new UnsafeAttributes();
-    for (AttributesExtractor<? super REQUEST, ? super RESPONSE> extractor : extractors) {
+    for (AttributesExtractor<? super REQUEST, ? super RESPONSE> extractor : attributesExtractors) {
       extractor.onStart(attributesBuilder, request);
     }
     Attributes attributes = attributesBuilder;
@@ -154,7 +161,7 @@ public void end(Context context, REQUEST request, RESPONSE response, @Nullable T
     Span span = Span.fromContext(context);
 
     UnsafeAttributes attributesBuilder = new UnsafeAttributes();
-    for (AttributesExtractor<? super REQUEST, ? super RESPONSE> extractor : extractors) {
+    for (AttributesExtractor<? super REQUEST, ? super RESPONSE> extractor : attributesExtractors) {
       extractor.onEnd(attributesBuilder, request, response);
     }
     Attributes attributes = attributesBuilder;

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

@@ -33,6 +33,7 @@ public final class InstrumenterBuilder<REQUEST, RESPONSE> {
 
   final List<AttributesExtractor<? super REQUEST, ? super RESPONSE>> attributesExtractors =
       new ArrayList<>();
+  final List<SpanLinkExtractor<? super REQUEST>> spanLinkExtractors = new ArrayList<>();
   final List<RequestListener> requestListeners = new ArrayList<>();
 
   SpanKindExtractor<? super REQUEST> spanKindExtractor = SpanKindExtractor.alwaysInternal();
@@ -83,6 +84,13 @@ public InstrumenterBuilder<REQUEST, RESPONSE> addAttributesExtractors(
     return addAttributesExtractors(Arrays.asList(attributesExtractors));
   }
 
+  /** Adds a {@link SpanLinkExtractor} to extract span link from requests. */
+  public InstrumenterBuilder<REQUEST, RESPONSE> addSpanLinkExtractor(
+      SpanLinkExtractor<REQUEST> spanLinkExtractor) {
+    spanLinkExtractors.add(spanLinkExtractor);
+    return this;
+  }
+
   /** Adds a {@link RequestMetrics} whose metrics will be recorded for request start and stop. */
   @UnstableApi
   public InstrumenterBuilder<REQUEST, RESPONSE> addRequestMetrics(RequestMetrics factory) {
@@ -126,8 +134,7 @@ public Instrumenter<REQUEST, RESPONSE> newClientInstrumenter(TextMapSetter<REQUE
    * requests.
    */
   public Instrumenter<REQUEST, RESPONSE> newServerInstrumenter(TextMapGetter<REQUEST> getter) {
-    return newInstrumenter(
-        InstrumenterConstructor.propagatingFromUpstream(getter), SpanKindExtractor.alwaysServer());
+    return newUpstreamPropagatingInstrumenter(SpanKindExtractor.alwaysServer(), getter);
   }
 
   /**
@@ -145,9 +152,17 @@ public Instrumenter<REQUEST, RESPONSE> newProducerInstrumenter(TextMapSetter<REQ
    * requests.
    */
   public Instrumenter<REQUEST, RESPONSE> newConsumerInstrumenter(TextMapGetter<REQUEST> getter) {
+    return newUpstreamPropagatingInstrumenter(SpanKindExtractor.alwaysConsumer(), getter);
+  }
+
+  /**
+   * Returns a new {@link Instrumenter} which will create spans with kind determined by the passed
+   * {@code spanKindExtractor} and extract context from requests.
+   */
+  public Instrumenter<REQUEST, RESPONSE> newUpstreamPropagatingInstrumenter(
+      SpanKindExtractor<REQUEST> spanKindExtractor, TextMapGetter<REQUEST> getter) {
     return newInstrumenter(
-        InstrumenterConstructor.propagatingFromUpstream(getter),
-        SpanKindExtractor.alwaysConsumer());
+        InstrumenterConstructor.propagatingFromUpstream(getter), spanKindExtractor);
   }
 
   /**

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

@@ -0,0 +1,28 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.api.instrumenter;
+
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.SpanContext;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.propagation.ContextPropagators;
+import io.opentelemetry.context.propagation.TextMapGetter;
+
+final class PropagatorBasedSpanLinkExtractor<REQUEST> implements SpanLinkExtractor<REQUEST> {
+  private final ContextPropagators propagators;
+  private final TextMapGetter<REQUEST> getter;
+
+  PropagatorBasedSpanLinkExtractor(ContextPropagators propagators, TextMapGetter<REQUEST> getter) {
+    this.propagators = propagators;
+    this.getter = getter;
+  }
+
+  @Override
+  public SpanContext extract(Context parentContext, REQUEST request) {
+    Context extracted = propagators.getTextMapPropagator().extract(parentContext, request, getter);
+    return Span.fromContext(extracted).getSpanContext();
+  }
+}

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

@@ -0,0 +1,30 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.api.instrumenter;
+
+import io.opentelemetry.api.trace.SpanContext;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.propagation.ContextPropagators;
+import io.opentelemetry.context.propagation.TextMapGetter;
+
+/** Extractor of a span link for a request. */
+@FunctionalInterface
+public interface SpanLinkExtractor<REQUEST> {
+  /**
+   * Extract a {@link SpanContext} that should be linked to the newly created span. Returning {@code
+   * SpanContext.getInvalid()} will not add any link to the span.
+   */
+  SpanContext extract(Context parentContext, REQUEST request);
+
+  /**
+   * Returns a new {@link SpanLinkExtractor} that will extract a {@link SpanContext} from the
+   * request using configured {@code propagators}.
+   */
+  static <REQUEST> SpanLinkExtractor<REQUEST> fromUpstreamRequest(
+      ContextPropagators propagators, TextMapGetter<REQUEST> getter) {
+    return new PropagatorBasedSpanLinkExtractor<>(propagators, getter);
+  }
+}

instrumentation-api/src/test/java/io/opentelemetry/instrumentation/api/instrumenter/InstrumenterTest.java

@@ -15,13 +15,15 @@
 import io.opentelemetry.api.trace.SpanId;
 import io.opentelemetry.api.trace.SpanKind;
 import io.opentelemetry.api.trace.TraceFlags;
+import io.opentelemetry.api.trace.TraceId;
 import io.opentelemetry.api.trace.TraceState;
 import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator;
 import io.opentelemetry.context.Context;
 import io.opentelemetry.context.propagation.TextMapGetter;
 import io.opentelemetry.instrumentation.api.tracer.ServerSpan;
 import io.opentelemetry.sdk.common.InstrumentationLibraryInfo;
 import io.opentelemetry.sdk.testing.junit5.OpenTelemetryExtension;
+import io.opentelemetry.sdk.trace.data.LinkData;
 import io.opentelemetry.sdk.trace.data.StatusData;
 import java.time.Instant;
 import java.util.Collections;
@@ -33,14 +35,18 @@
 import org.junit.jupiter.api.extension.RegisterExtension;
 
 class InstrumenterTest {
+  private static final String LINK_TRACE_ID = TraceId.fromLongs(0, 42);
+  private static final String LINK_SPAN_ID = SpanId.fromLong(123);
 
   private static final Map<String, String> REQUEST =
       Collections.unmodifiableMap(
           Stream.of(
                   entry("req1", "req1_value"),
                   entry("req2", "req2_value"),
                   entry("req2_2", "req2_2_value"),
-                  entry("req3", "req3_value"))
+                  entry("req3", "req3_value"),
+                  entry("linkTraceId", LINK_TRACE_ID),
+                  entry("linkSpanId", LINK_SPAN_ID))
               .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)));
 
   private static final Map<String, String> RESPONSE =
@@ -86,6 +92,17 @@ protected void onEnd(
     }
   }
 
+  static class LinkExtractor implements SpanLinkExtractor<Map<String, String>> {
+    @Override
+    public SpanContext extract(Context parentContext, Map<String, String> request) {
+      return SpanContext.create(
+          request.get("linkTraceId"),
+          request.get("linkSpanId"),
+          TraceFlags.getSampled(),
+          TraceState.getDefault());
+    }
+  }
+
   static class MapGetter implements TextMapGetter<Map<String, String>> {
 
     @Override
@@ -108,6 +125,7 @@ void server() {
         Instrumenter.<Map<String, String>, Map<String, String>>newBuilder(
                 otelTesting.getOpenTelemetry(), "test", unused -> "span")
             .addAttributesExtractors(new AttributesExtractor1(), new AttributesExtractor2())
+            .addSpanLinkExtractor(new LinkExtractor())
             .newServerInstrumenter(new MapGetter());
 
     Context context = instrumenter.start(Context.root(), REQUEST);
@@ -132,6 +150,7 @@ void server() {
                             .hasSpanId(spanContext.getSpanId())
                             .hasParentSpanId(SpanId.getInvalid())
                             .hasStatus(StatusData.unset())
+                            .hasLinks(expectedSpanLink())
                             .hasAttributesSatisfying(
                                 attributes ->
                                     assertThat(attributes)
@@ -215,6 +234,7 @@ void client() {
         Instrumenter.<Map<String, String>, Map<String, String>>newBuilder(
                 otelTesting.getOpenTelemetry(), "test", unused -> "span")
             .addAttributesExtractors(new AttributesExtractor1(), new AttributesExtractor2())
+            .addSpanLinkExtractor(new LinkExtractor())
             .newClientInstrumenter(Map::put);
 
     Map<String, String> request = new HashMap<>(REQUEST);
@@ -240,6 +260,7 @@ void client() {
                             .hasSpanId(spanContext.getSpanId())
                             .hasParentSpanId(SpanId.getInvalid())
                             .hasStatus(StatusData.unset())
+                            .hasLinks(expectedSpanLink())
                             .hasAttributesSatisfying(
                                 attributes ->
                                     assertThat(attributes)
@@ -340,4 +361,32 @@ void shouldStartSpanWithGivenStartTime() {
                 trace.hasSpansSatisfyingExactly(
                     span -> span.hasName("test span").startsAt(startTime).endsAt(endTime)));
   }
+
+  @Test
+  void shouldNotAddInvalidLink() {
+    // given
+    Instrumenter<String, String> instrumenter =
+        Instrumenter.<String, String>newBuilder(
+                otelTesting.getOpenTelemetry(), "test", request -> "test span")
+            .addSpanLinkExtractor((parentContext, request) -> SpanContext.getInvalid())
+            .newInstrumenter();
+
+    // when
+    Context context = instrumenter.start(Context.root(), "request");
+    instrumenter.end(context, "request", "response", null);
+
+    // then
+    otelTesting
+        .assertTraces()
+        .hasTracesSatisfyingExactly(
+            trace ->
+                trace.hasSpansSatisfyingExactly(
+                    span -> span.hasName("test span").hasTotalRecordedLinks(0)));
+  }
+
+  private static LinkData expectedSpanLink() {
+    return LinkData.create(
+        SpanContext.create(
+            LINK_TRACE_ID, LINK_SPAN_ID, TraceFlags.getSampled(), TraceState.getDefault()));
+  }
 }

instrumentation-api/src/test/java/io/opentelemetry/instrumentation/api/instrumenter/PropagatorBasedSpanLinkExtractorTest.java

@@ -0,0 +1,61 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.api.instrumenter;
+
+import static java.util.Collections.singletonMap;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+import io.opentelemetry.api.trace.SpanContext;
+import io.opentelemetry.api.trace.SpanId;
+import io.opentelemetry.api.trace.TraceFlags;
+import io.opentelemetry.api.trace.TraceId;
+import io.opentelemetry.api.trace.TraceState;
+import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.propagation.ContextPropagators;
+import io.opentelemetry.context.propagation.TextMapGetter;
+import java.util.Map;
+import org.junit.jupiter.api.Test;
+
+class PropagatorBasedSpanLinkExtractorTest {
+  private static final String TRACE_ID = TraceId.fromLongs(0, 123);
+  private static final String SPAN_ID = SpanId.fromLong(456);
+
+  @Test
+  void shouldExtractSpanLink() {
+    // given
+    ContextPropagators propagators =
+        ContextPropagators.create(W3CTraceContextPropagator.getInstance());
+
+    SpanLinkExtractor<Map<String, String>> underTest =
+        SpanLinkExtractor.fromUpstreamRequest(propagators, new MapGetter());
+
+    Map<String, String> request =
+        singletonMap("traceparent", String.format("00-%s-%s-01", TRACE_ID, SPAN_ID));
+
+    // when
+    SpanContext link = underTest.extract(Context.root(), request);
+
+    // then
+    assertEquals(
+        SpanContext.createFromRemoteParent(
+            TRACE_ID, SPAN_ID, TraceFlags.getSampled(), TraceState.getDefault()),
+        link);
+  }
+
+  static final class MapGetter implements TextMapGetter<Map<String, String>> {
+
+    @Override
+    public Iterable<String> keys(Map<String, String> carrier) {
+      return carrier.keySet();
+    }
+
+    @Override
+    public String get(Map<String, String> carrier, String key) {
+      return carrier.get(key);
+    }
+  }
+}

instrumentation/spring/spring-integration-4.1/library/spring-integration-4.1-library.gradle

@@ -0,0 +1,15 @@
+apply from: "$rootDir/gradle/instrumentation-library.gradle"
+
+dependencies {
+  compileOnly "com.google.auto.value:auto-value-annotations"
+  annotationProcessor "com.google.auto.value:auto-value"
+
+  library 'org.springframework.integration:spring-integration-core:4.1.0.RELEASE'
+
+  testLibrary "org.springframework.boot:spring-boot-starter-test:1.5.22.RELEASE"
+  testLibrary "org.springframework.boot:spring-boot-starter:1.5.22.RELEASE"
+}
+
+test {
+  systemProperty "testLatestDeps", testLatestDeps
+}

instrumentation/spring/spring-integration-4.1/library/src/main/java/io/opentelemetry/instrumentation/spring/integration/ContextAndScope.java

@@ -0,0 +1,30 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.spring.integration;
+
+import com.google.auto.value.AutoValue;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.Scope;
+
+@AutoValue
+abstract class ContextAndScope {
+
+  abstract Context getContext();
+
+  abstract Scope getScope();
+
+  void close() {
+    getScope().close();
+  }
+
+  static ContextAndScope makeCurrent(Context context) {
+    return create(context, context.makeCurrent());
+  }
+
+  static ContextAndScope create(Context context, Scope scope) {
+    return new AutoValue_ContextAndScope(context, scope);
+  }
+}