feat: Add OpenTelemetry tracing to the Publisher and Subscriber

google-cloud-pubsub/pom.xml

@@ -100,6 +100,18 @@
       <artifactId>google-http-client</artifactId>
       <scope>runtime</scope>
     </dependency>
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-api</artifactId>
+    </dependency>
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-context</artifactId>
+    </dependency>
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-semconv</artifactId>
+    </dependency>
 
     <!-- Test dependencies -->
     <dependency>
@@ -142,6 +154,21 @@
       <artifactId>opencensus-impl</artifactId>
       <scope>test</scope>
     </dependency>
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-sdk-trace</artifactId>
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-sdk-testing</artifactId>
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>org.assertj</groupId>
+      <artifactId>assertj-core</artifactId>
+      <scope>test</scope>
+    </dependency>
     <!-- Need testing utility classes for generated gRPC clients tests -->
     <dependency>
       <groupId>com.google.api</groupId>
@@ -174,6 +201,7 @@
               <ignoredUnusedDeclaredDependency>com.google.auth:google-auth-library-oauth2-http:jar</ignoredUnusedDeclaredDependency>
               <ignoredUnusedDeclaredDependency>io.opencensus:opencensus-impl</ignoredUnusedDeclaredDependency>
               <ignoredUnusedDeclaredDependency>javax.annotation:javax.annotation-api</ignoredUnusedDeclaredDependency>
+              <ignoredUnusedDeclaredDependency>org.assertj:assertj-core</ignoredUnusedDeclaredDependency>
             </ignoredUnusedDeclaredDependencies>
           </configuration>
         </plugin>

google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/AckRequestData.java

@@ -22,10 +22,12 @@
 public class AckRequestData {
   private final String ackId;
   private final Optional<SettableApiFuture<AckResponse>> messageFuture;
+  private PubsubMessageWrapper messageWrapper;
 
   protected AckRequestData(Builder builder) {
     this.ackId = builder.ackId;
     this.messageFuture = builder.messageFuture;
+    this.messageWrapper = builder.messageWrapper;
   }
 
   public String getAckId() {
@@ -36,6 +38,17 @@ public SettableApiFuture<AckResponse> getMessageFutureIfExists() {
     return this.messageFuture.orElse(null);
   }
 
+  /**
+   * Returns an empty PubsubMessageWrapper with OpenTelemetry tracing disabled. This allows methods
+   * that use this method to be unit tested.
+   */
+  public PubsubMessageWrapper getMessageWrapper() {
+    if (this.messageWrapper == null) {
+      return PubsubMessageWrapper.newBuilder(null, null, false).build();
+    }
+    return messageWrapper;
+  }
+
   public AckRequestData setResponse(AckResponse ackResponse, boolean setResponseOnSuccess) {
     if (this.messageFuture.isPresent() && !this.messageFuture.get().isDone()) {
       switch (ackResponse) {
@@ -68,6 +81,7 @@ public static Builder newBuilder(String ackId) {
   protected static final class Builder {
     private final String ackId;
     private Optional<SettableApiFuture<AckResponse>> messageFuture = Optional.empty();
+    private PubsubMessageWrapper messageWrapper;
 
     protected Builder(String ackId) {
       this.ackId = ackId;
@@ -78,6 +92,11 @@ public Builder setMessageFuture(SettableApiFuture<AckResponse> messageFuture) {
       return this;
     }
 
+    public Builder setMessageWrapper(PubsubMessageWrapper messageWrapper) {
+      this.messageWrapper = messageWrapper;
+      return this;
+    }
+
     public AckRequestData build() {
       return new AckRequestData(this);
     }

google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/MessageDispatcher.java

@@ -28,6 +28,7 @@
 import com.google.common.util.concurrent.MoreExecutors;
 import com.google.pubsub.v1.PubsubMessage;
 import com.google.pubsub.v1.ReceivedMessage;
+import io.opentelemetry.api.trace.Tracer;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.Iterator;
@@ -104,6 +105,10 @@ class MessageDispatcher {
   // To keep track of number of seconds the receiver takes to process messages.
   private final Distribution ackLatencyDistribution;
 
+  private final String subscriptionName;
+  private final boolean enableOpenTelemetryTracing;
+  private final Tracer tracer;
+
   /** Internal representation of a reply to a Pubsub message, to be sent back to the service. */
   public enum AckReply {
     ACK,
@@ -157,6 +162,7 @@ public void onFailure(Throwable t) {
           t);
       this.ackRequestData.setResponse(AckResponse.OTHER, false);
       pendingNacks.add(this.ackRequestData);
+      this.ackRequestData.getMessageWrapper().endSubscribeProcessSpan("nack");
       forget();
     }
 
@@ -169,9 +175,11 @@ public void onSuccess(AckReply reply) {
           ackLatencyDistribution.record(
               Ints.saturatedCast(
                   (long) Math.ceil((clock.millisTime() - receivedTimeMillis) / 1000D)));
+          this.ackRequestData.getMessageWrapper().endSubscribeProcessSpan("ack");
           break;
         case NACK:
           pendingNacks.add(this.ackRequestData);
+          this.ackRequestData.getMessageWrapper().endSubscribeProcessSpan("nack");
           break;
         default:
           throw new IllegalArgumentException(String.format("AckReply: %s not supported", reply));
@@ -217,6 +225,10 @@ private MessageDispatcher(Builder builder) {
     jobLock = new ReentrantLock();
     messagesWaiter = new Waiter();
     sequentialExecutor = new SequentialExecutorService.AutoExecutor(builder.executor);
+
+    subscriptionName = builder.subscriptionName;
+    enableOpenTelemetryTracing = builder.enableOpenTelemetryTracing;
+    tracer = builder.tracer;
   }
 
   private boolean shouldSetMessageFuture() {
@@ -351,13 +363,15 @@ void setMessageOrderingEnabled(boolean messageOrderingEnabled) {
   }
 
   private static class OutstandingMessage {
-    private final ReceivedMessage receivedMessage;
     private final AckHandler ackHandler;
 
-    private OutstandingMessage(ReceivedMessage receivedMessage, AckHandler ackHandler) {
-      this.receivedMessage = receivedMessage;
+    private OutstandingMessage(AckHandler ackHandler) {
       this.ackHandler = ackHandler;
     }
+
+    public PubsubMessageWrapper messageWrapper() {
+      return this.ackHandler.ackRequestData.getMessageWrapper();
+    }
   }
 
   private static class ReceiptCompleteData {
@@ -390,10 +404,21 @@ void processReceivedMessages(List<ReceivedMessage> messages) {
       if (shouldSetMessageFuture()) {
         builder.setMessageFuture(SettableApiFuture.create());
       }
+      PubsubMessageWrapper messageWrapper =
+          PubsubMessageWrapper.newBuilder(
+                  message.getMessage(),
+                  subscriptionName,
+                  message.getAckId(),
+                  message.getDeliveryAttempt(),
+                  enableOpenTelemetryTracing)
+              .build();
+      builder.setMessageWrapper(messageWrapper);
+      messageWrapper.startSubscriberSpan(tracer, this.exactlyOnceDeliveryEnabled.get());
+
       AckRequestData ackRequestData = builder.build();
       AckHandler ackHandler =
           new AckHandler(ackRequestData, message.getMessage().getSerializedSize(), totalExpiration);
-      OutstandingMessage outstandingMessage = new OutstandingMessage(message, ackHandler);
+      OutstandingMessage outstandingMessage = new OutstandingMessage(ackHandler);
 
       if (this.exactlyOnceDeliveryEnabled.get()) {
         // For exactly once deliveries we don't add to outstanding batch because we first
@@ -457,30 +482,39 @@ private void processBatch(List<OutstandingMessage> batch) {
     for (OutstandingMessage message : batch) {
       // This is a blocking flow controller.  We have already incremented messagesWaiter, so
       // shutdown will block on processing of all these messages anyway.
+      message.messageWrapper().startSubscribeConcurrencyControlSpan(tracer);
       try {
-        flowController.reserve(1, message.receivedMessage.getMessage().getSerializedSize());
+        flowController.reserve(1, message.messageWrapper().getPubsubMessage().getSerializedSize());
+        message.messageWrapper().endSubscribeConcurrencyControlSpan();
       } catch (FlowControlException unexpectedException) {
         // This should be a blocking flow controller and never throw an exception.
+        message.messageWrapper().setSubscribeConcurrencyControlSpanException(unexpectedException);
         throw new IllegalStateException("Flow control unexpected exception", unexpectedException);
       }
-      processOutstandingMessage(addDeliveryInfoCount(message.receivedMessage), message.ackHandler);
+      addDeliveryInfoCount(message.messageWrapper());
+      processOutstandingMessage(message.ackHandler);
     }
   }
 
-  private PubsubMessage addDeliveryInfoCount(ReceivedMessage receivedMessage) {
-    PubsubMessage originalMessage = receivedMessage.getMessage();
-    int deliveryAttempt = receivedMessage.getDeliveryAttempt();
+  private void addDeliveryInfoCount(PubsubMessageWrapper messageWrapper) {
+    PubsubMessage originalMessage = messageWrapper.getPubsubMessage();
+    int deliveryAttempt = messageWrapper.getDeliveryAttempt();
     // Delivery Attempt will be set to 0 if DeadLetterPolicy is not set on the subscription. In
     // this case, do not populate the PubsubMessage with the delivery attempt attribute.
     if (deliveryAttempt > 0) {
-      return PubsubMessage.newBuilder(originalMessage)
-          .putAttributes("googclient_deliveryattempt", Integer.toString(deliveryAttempt))
-          .build();
+      messageWrapper.setPubsubMessage(
+          PubsubMessage.newBuilder(originalMessage)
+              .putAttributes("googclient_deliveryattempt", Integer.toString(deliveryAttempt))
+              .build());
     }
-    return originalMessage;
   }
 
-  private void processOutstandingMessage(final PubsubMessage message, final AckHandler ackHandler) {
+  private void processOutstandingMessage(final AckHandler ackHandler) {
+    // Get the PubsubMessageWrapper and the PubsubMessage it wraps that are stored withing the
+    // AckHandler object.
+    PubsubMessageWrapper messageWrapper = ackHandler.ackRequestData.getMessageWrapper();
+    PubsubMessage message = messageWrapper.getPubsubMessage();
+
     // This future is for internal bookkeeping to be sent to the StreamingSubscriberConnection
     // use below in the consumers
     SettableApiFuture<AckReply> ackReplySettableApiFuture = SettableApiFuture.create();
@@ -499,8 +533,10 @@ public void run() {
                 // so it was probably sent to someone else. Don't work on it.
                 // Don't nack it either, because we'd be nacking someone else's message.
                 ackHandler.forget();
+                messageWrapper.setSubscriberSpanExpirationResult();
                 return;
               }
+              messageWrapper.startSubscribeProcessSpan(tracer);
               if (shouldSetMessageFuture()) {
                 // This is the message future that is propagated to the user
                 SettableApiFuture<AckResponse> messageFuture =
@@ -521,7 +557,9 @@ public void run() {
     if (!messageOrderingEnabled.get() || message.getOrderingKey().isEmpty()) {
       executor.execute(deliverMessageTask);
     } else {
+      messageWrapper.startSubscribeSchedulerSpan(tracer);
       sequentialExecutor.submit(message.getOrderingKey(), deliverMessageTask);
+      messageWrapper.endSubscribeSchedulerSpan();
     }
   }
 
@@ -607,8 +645,10 @@ void processOutstandingOperations() {
     List<AckRequestData> ackRequestDataReceipts = new ArrayList<AckRequestData>();
     pendingReceipts.drainTo(ackRequestDataReceipts);
     if (!ackRequestDataReceipts.isEmpty()) {
-      modackRequestData.add(
-          new ModackRequestData(this.getMessageDeadlineSeconds(), ackRequestDataReceipts));
+      ModackRequestData receiptModack =
+          new ModackRequestData(this.getMessageDeadlineSeconds(), ackRequestDataReceipts);
+      receiptModack.setIsReceiptModack(true);
+      modackRequestData.add(receiptModack);
     }
     logger.log(Level.FINER, "Sending {0} receipts", ackRequestDataReceipts.size());
 
@@ -645,6 +685,10 @@ public static final class Builder {
     private ScheduledExecutorService systemExecutor;
     private ApiClock clock;
 
+    private String subscriptionName;
+    private boolean enableOpenTelemetryTracing;
+    private Tracer tracer;
+
     protected Builder(MessageReceiver receiver) {
       this.receiver = receiver;
     }
@@ -715,6 +759,21 @@ public Builder setApiClock(ApiClock clock) {
       return this;
     }
 
+    public Builder setSubscriptionName(String subscriptionName) {
+      this.subscriptionName = subscriptionName;
+      return this;
+    }
+
+    public Builder setEnableOpenTelemetryTracing(boolean enableOpenTelemetryTracing) {
+      this.enableOpenTelemetryTracing = enableOpenTelemetryTracing;
+      return this;
+    }
+
+    public Builder setTracer(Tracer tracer) {
+      this.tracer = tracer;
+      return this;
+    }
+
     public MessageDispatcher build() {
       return new MessageDispatcher(this);
     }

google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/ModackRequestData.java

@@ -21,6 +21,7 @@
 class ModackRequestData {
   private final int deadlineExtensionSeconds;
   private List<AckRequestData> ackRequestData;
+  private boolean isReceiptModack;
 
   ModackRequestData(int deadlineExtensionSeconds) {
     this.deadlineExtensionSeconds = deadlineExtensionSeconds;
@@ -45,8 +46,17 @@ public List<AckRequestData> getAckRequestData() {
     return ackRequestData;
   }
 
+  public boolean getIsReceiptModack() {
+    return isReceiptModack;
+  }
+
   public ModackRequestData addAckRequestData(AckRequestData ackRequestData) {
     this.ackRequestData.add(ackRequestData);
     return this;
   }
+
+  public ModackRequestData setIsReceiptModack(boolean isReceiptModack) {
+    this.isReceiptModack = isReceiptModack;
+    return this;
+  }
 }