Link RabbitMQ receive span with the producer span (#6808)

Similar to #6804, but for RabbitMQ. Also changed the span kind of the receive span to `CONSUMER`, to match the spec.
open-telemetry · Oct 6, 2022 · ab0c875 · ab0c875
1 parent 6f6af66
commit ab0c875
Show file tree

Hide file tree

Showing 4 changed files with 156 additions and 92 deletions.
diff --git a/instrumentation/rabbitmq-2.7/javaagent/build.gradle.kts b/instrumentation/rabbitmq-2.7/javaagent/build.gradle.kts
@@ -32,5 +32,7 @@ dependencies {
 tasks.withType<Test>().configureEach {
   // TODO run tests both with and without experimental span attributes
   jvmArgs("-Dotel.instrumentation.rabbitmq.experimental-span-attributes=true")
+  jvmArgs("-Dotel.instrumentation.messaging.experimental.receive-telemetry.enabled=true")
+
   usesService(gradle.sharedServices.registrations["testcontainersBuildService"].getService())
 }
diff --git a/...t/src/main/java/io/opentelemetry/javaagent/instrumentation/rabbitmq/RabbitSingletons.java b/...t/src/main/java/io/opentelemetry/javaagent/instrumentation/rabbitmq/RabbitSingletons.java
@@ -18,6 +18,7 @@
 import io.opentelemetry.instrumentation.api.instrumenter.messaging.MessagingAttributesExtractor;
 import io.opentelemetry.instrumentation.api.instrumenter.messaging.MessagingAttributesGetter;
 import io.opentelemetry.instrumentation.api.instrumenter.net.NetClientAttributesExtractor;
+import io.opentelemetry.instrumentation.api.internal.PropagatorBasedSpanLinksExtractor;
 import io.opentelemetry.javaagent.bootstrap.internal.ExperimentalConfig;
 import io.opentelemetry.javaagent.bootstrap.internal.InstrumentationConfig;
 import java.util.ArrayList;
@@ -29,18 +30,15 @@ public final class RabbitSingletons {
       InstrumentationConfig.get()
           .getBoolean("otel.instrumentation.rabbitmq.experimental-span-attributes", false);
   private static final String instrumentationName = "io.opentelemetry.rabbitmq-2.7";
-  private static final Instrumenter<ChannelAndMethod, Void> channelInstrumenter;
-  private static final Instrumenter<ReceiveRequest, GetResponse> receiveInstrumenter;
-  private static final Instrumenter<DeliveryRequest, Void> deliverInstrumenter;
+  private static final Instrumenter<ChannelAndMethod, Void> channelInstrumenter =
+      createChannelInstrumenter();
+  private static final Instrumenter<ReceiveRequest, GetResponse> receiveInstrumenter =
+      createReceiveInstrumenter();
+  private static final Instrumenter<DeliveryRequest, Void> deliverInstrumenter =
+      createDeliverInstrumenter();
   static final ContextKey<RabbitChannelAndMethodHolder> CHANNEL_AND_METHOD_CONTEXT_KEY =
       ContextKey.named("opentelemetry-rabbitmq-channel-and-method-context-key");
 
-  static {
-    channelInstrumenter = createChannelInstrumenter();
-    receiveInstrumenter = createReceiveInstrumenter();
-    deliverInstrumenter = createDeliverInstrumenter();
-  }
-
   public static Instrumenter<ChannelAndMethod, Void> channelInstrumenter() {
     return channelInstrumenter;
   }
@@ -82,7 +80,12 @@ private static Instrumenter<ReceiveRequest, GetResponse> createReceiveInstrument
     return Instrumenter.<ReceiveRequest, GetResponse>builder(
             GlobalOpenTelemetry.get(), instrumentationName, ReceiveRequest::spanName)
         .addAttributesExtractors(extractors)
-        .buildInstrumenter(SpanKindExtractor.alwaysClient());
+        .setEnabled(ExperimentalConfig.get().messagingReceiveInstrumentationEnabled())
+        .addSpanLinksExtractor(
+            new PropagatorBasedSpanLinksExtractor<>(
+                GlobalOpenTelemetry.getPropagators().getTextMapPropagator(),
+                ReceiveRequestTextMapGetter.INSTANCE))
+        .buildInstrumenter(SpanKindExtractor.alwaysConsumer());
   }
 
   private static Instrumenter<DeliveryRequest, Void> createDeliverInstrumenter() {

diff --git a/...java/io/opentelemetry/javaagent/instrumentation/rabbitmq/ReceiveRequestTextMapGetter.java b/...java/io/opentelemetry/javaagent/instrumentation/rabbitmq/ReceiveRequestTextMapGetter.java
@@ -0,0 +1,40 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.rabbitmq;
+
+import com.rabbitmq.client.AMQP;
+import com.rabbitmq.client.GetResponse;
+import io.opentelemetry.context.propagation.TextMapGetter;
+import java.util.Collections;
+import java.util.Map;
+import java.util.Optional;
+import javax.annotation.Nullable;
+
+enum ReceiveRequestTextMapGetter implements TextMapGetter<ReceiveRequest> {
+  INSTANCE;
+
+  @Override
+  public Iterable<String> keys(ReceiveRequest carrier) {
+    return Optional.of(carrier)
+        .map(ReceiveRequest::getResponse)
+        .map(GetResponse::getProps)
+        .map(AMQP.BasicProperties::getHeaders)
+        .map(Map::keySet)
+        .orElse(Collections.emptySet());
+  }
+
+  @Nullable
+  @Override
+  public String get(@Nullable ReceiveRequest carrier, String key) {
+    return Optional.ofNullable(carrier)
+        .map(ReceiveRequest::getResponse)
+        .map(GetResponse::getProps)
+        .map(AMQP.BasicProperties::getHeaders)
+        .map(headers -> headers.get(key))
+        .map(Object::toString)
+        .orElse(null);
+  }
+}