From fbb2b8f213d641885dee31df07462368d1c2fb78 Mon Sep 17 00:00:00 2001
From: Ran Nozik <ran@gethelios.dev>
Date: Sun, 2 Oct 2022 10:07:13 +0300
Subject: [PATCH] collect pubsub attributes

---
 .../PubsubPublisherInstrumentation.java       |  4 ++-
 .../pubsub/PubsubSingletons.java              | 25 ++++++++-----------
 2 files changed, 13 insertions(+), 16 deletions(-)

diff --git a/instrumentation/pubsub-2.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/pubsub/PubsubPublisherInstrumentation.java b/instrumentation/pubsub-2.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/pubsub/PubsubPublisherInstrumentation.java
index c8f6a4c3cd7b..532252897ff8 100644
--- a/instrumentation/pubsub-2.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/pubsub/PubsubPublisherInstrumentation.java
+++ b/instrumentation/pubsub-2.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/pubsub/PubsubPublisherInstrumentation.java
@@ -9,6 +9,7 @@
 import static net.bytebuddy.matcher.ElementMatchers.isPublic;
 import static net.bytebuddy.matcher.ElementMatchers.named;
 
+import com.google.cloud.pubsub.v1.Publisher;
 import com.google.pubsub.v1.PubsubMessage;
 import io.opentelemetry.context.Context;
 import io.opentelemetry.javaagent.bootstrap.Java8BytecodeBridge;
@@ -35,9 +36,10 @@ public void transform(TypeTransformer typeTransformer) {
   public static class PubsubPublisherAddAttributesAdvice {
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static void onEnterHandle(
+        @Advice.This Publisher publisher,
         @Advice.Argument(value = 0, readOnly = false) PubsubMessage pubsubMessage) {
       Context parentContext = Java8BytecodeBridge.currentContext();
-      startAndInjectSpan(parentContext, pubsubMessage);
+      startAndInjectSpan(parentContext, pubsubMessage, publisher);
     }
   }
 }
diff --git a/instrumentation/pubsub-2.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/pubsub/PubsubSingletons.java b/instrumentation/pubsub-2.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/pubsub/PubsubSingletons.java
index 7073be364c20..4f6aefbe1fa2 100644
--- a/instrumentation/pubsub-2.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/pubsub/PubsubSingletons.java
+++ b/instrumentation/pubsub-2.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/pubsub/PubsubSingletons.java
@@ -5,6 +5,7 @@
 
 package io.opentelemetry.javaagent.instrumentation.pubsub;
 
+import com.google.cloud.pubsub.v1.Publisher;
 import com.google.pubsub.v1.PubsubMessage;
 import io.opentelemetry.api.GlobalOpenTelemetry;
 import io.opentelemetry.api.trace.Span;
@@ -13,6 +14,7 @@
 import io.opentelemetry.instrumentation.api.instrumenter.SpanKindExtractor;
 import io.opentelemetry.instrumentation.api.instrumenter.SpanNameExtractor;
 import io.opentelemetry.javaagent.bootstrap.Java8BytecodeBridge;
+import io.opentelemetry.semconv.trace.attributes.SemanticAttributes;
 import java.lang.reflect.Field;
 import java.util.HashMap;
 import java.util.Map;
@@ -44,13 +46,7 @@ public static Instrumenter<PubsubMessage, Void> publisherInstrumenter() {
   }
 
   private static Instrumenter<PubsubMessage, Void> createPublisherInstrumenter() {
-    SpanNameExtractor publisherSpanNameExtractor =
-        new SpanNameExtractor() {
-          @Override
-          public String extract(Object o) {
-            return publisherSpanName;
-          }
-        };
+    SpanNameExtractor publisherSpanNameExtractor = o -> publisherSpanName;
 
     return Instrumenter.<PubsubMessage, Void>builder(
             GlobalOpenTelemetry.get(), instrumentationName, publisherSpanNameExtractor)
@@ -59,20 +55,15 @@ public String extract(Object o) {
 
   public static Instrumenter<PubsubMessage, Void> createSubscriberInstrumenter() {
 
-    SpanNameExtractor subscriberSpanNameExtractor =
-        new SpanNameExtractor() {
-          @Override
-          public String extract(Object o) {
-            return subscriberSpanName;
-          }
-        };
+    SpanNameExtractor subscriberSpanNameExtractor = o -> subscriberSpanName;
 
     return Instrumenter.<PubsubMessage, Void>builder(
             GlobalOpenTelemetry.get(), instrumentationName, subscriberSpanNameExtractor)
         .newInstrumenter(SpanKindExtractor.alwaysConsumer());
   }
 
-  public static void startAndInjectSpan(Context parentContext, PubsubMessage pubsubMessage) {
+  public static void startAndInjectSpan(
+      Context parentContext, PubsubMessage pubsubMessage, Publisher publisher) {
     if (!publisherInstrumenter().shouldStart(parentContext, pubsubMessage)) {
       return;
     }
@@ -83,6 +74,10 @@ public static void startAndInjectSpan(Context parentContext, PubsubMessage pubsu
     Context context = publisherInstrumenter().start(parentContext, pubsubMessage);
     Span span = Java8BytecodeBridge.spanFromContext(context);
     span.setAttribute(MESSAGE_PAYLOAD_ATTRIBUTE, new String(pubsubMessage.getData().toByteArray()));
+    span.setAttribute(SemanticAttributes.MESSAGING_SYSTEM, "pubsub");
+    span.setAttribute(SemanticAttributes.MESSAGING_DESTINATION_KIND, "topic");
+    span.setAttribute(SemanticAttributes.MESSAGING_DESTINATION, publisher.getTopicNameString());
+
     GlobalOpenTelemetry.get()
         .getPropagators()
         .getTextMapPropagator()