[Service Bus] Tracing for send API

sdk/servicebus/service-bus/review/service-bus.api.md

@@ -15,6 +15,8 @@ import { PageSettings } from '@azure/core-paging';
 import { PipelineOptions } from '@azure/core-http';
 import { RetryOptions } from '@azure/core-amqp';
 import { ServiceClient } from '@azure/core-http';
+import { Span } from '@opentelemetry/api';
+import { SpanContext } from '@opentelemetry/api';
 import { TokenCredential } from '@azure/core-amqp';
 import { TokenType } from '@azure/core-amqp';
 import { UserAgentOptions } from '@azure/core-http';
@@ -226,6 +228,11 @@ export { OperationOptions }
 // @public
 export type OperationOptionsBase = Pick<OperationOptions, "abortSignal" | "tracingOptions">;
 
+// @public
+export interface ParentSpanOptions {
+    parentSpan?: Span | SpanContext | null;
+}
+
 // @public
 export interface PeekMessagesOptions extends OperationOptionsBase {
     fromSequenceNumber?: Long;
@@ -395,8 +402,10 @@ export interface ServiceBusMessageBatch {
     // @internal
     _generateMessage(): Buffer;
     readonly maxSizeInBytes: number;
+    // @internal
+    readonly _messageSpanContexts: SpanContext[];
     readonly sizeInBytes: number;
-    tryAdd(message: ServiceBusMessage): boolean;
+    tryAdd(message: ServiceBusMessage, options?: TryAddOptions): boolean;
 }
 
 // @public
@@ -577,6 +586,10 @@ export interface TopicRuntimeProperties {
 export interface TopicRuntimePropertiesResponse extends TopicRuntimeProperties, Response {
 }
 
+// @public
+export interface TryAddOptions extends ParentSpanOptions {
+}
+
 export { WebSocketImpl }
 
 export { WebSocketOptions }

sdk/servicebus/service-bus/src/diagnostics/instrumentServiceBusMessage.ts

@@ -0,0 +1,56 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+
+import { extractSpanContextFromTraceParentHeader, getTraceParentHeader } from "@azure/core-tracing";
+import { Span, SpanContext } from "@opentelemetry/api";
+import { ServiceBusMessage } from "../serviceBusMessage";
+
+/**
+ * @ignore
+ */
+export const TRACEPARENT_PROPERTY = "Diagnostic-Id";
+
+/**
+ * Populates the `ServiceBusMessage` with `SpanContext` info to support trace propagation.
+ * Creates and returns a copy of the passed in `ServiceBusMessage` unless the `ServiceBusMessage`
+ * has already been instrumented.
+ * @param message The `ServiceBusMessage` to instrument.
+ * @param span The `Span` containing the context to propagate tracing information.
+ * @ignore
+ * @internal
+ */
+export function instrumentServiceBusMessage(
+  message: ServiceBusMessage,
+  span: Span
+): ServiceBusMessage {
+  if (message.properties && message.properties[TRACEPARENT_PROPERTY]) {
+    return message;
+  }
+
+  // create a copy so the original isn't modified
+  message = { ...message, properties: { ...message.properties } };
+
+  const traceParent = getTraceParentHeader(span.context());
+  if (traceParent) {
+    message.properties![TRACEPARENT_PROPERTY] = traceParent;
+  }
+
+  return message;
+}
+
+/**
+ * Extracts the `SpanContext` from an `ServiceBusMessage` if the context exists.
+ * @param message An individual `ServiceBusMessage` object.
+ * @internal
+ * @ignore
+ */
+export function extractSpanContextFromServiceBusMessage(
+  message: ServiceBusMessage
+): SpanContext | undefined {
+  if (!message.properties || !message.properties[TRACEPARENT_PROPERTY]) {
+    return;
+  }
+
+  const diagnosticId = message.properties[TRACEPARENT_PROPERTY] as string;
+  return extractSpanContextFromTraceParentHeader(diagnosticId);
+}

sdk/servicebus/service-bus/src/diagnostics/messageSpan.ts

@@ -0,0 +1,27 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+
+import { ConnectionConfig } from "@azure/core-amqp";
+import { getTracer } from "@azure/core-tracing";
+import { Span, SpanContext, SpanKind } from "@opentelemetry/api";
+
+/**
+ * @internal
+ * @ignore
+ */
+export function createMessageSpan(
+  parentSpan?: Span | SpanContext | null,
+  config?: Pick<ConnectionConfig, "entityPath" | "host">
+): Span {
+  const tracer = getTracer();
+  const span = tracer.startSpan("Azure.ServiceBus.message", {
+    kind: SpanKind.PRODUCER,
+    parent: parentSpan
+  });
+  span.setAttribute("az.namespace", "Microsoft.ServiceBus");
+  if (config) {
+    span.setAttribute("message_bus.destination", config.entityPath);
+    span.setAttribute("peer.address", config.host);
+  }
+  return span;
+}

sdk/servicebus/service-bus/src/index.ts

@@ -31,7 +31,11 @@ export {
   SubQueue,
   SubscribeOptions
 } from "./models";
-export { OperationOptionsBase } from "./modelsToBeSharedWithEventHubs";
+export {
+  OperationOptionsBase,
+  TryAddOptions,
+  ParentSpanOptions
+} from "./modelsToBeSharedWithEventHubs";
 export { ServiceBusReceiver } from "./receivers/receiver";
 export { ServiceBusSessionReceiver } from "./receivers/sessionReceiver";
 export { ServiceBusSender } from "./sender";

sdk/servicebus/service-bus/src/modelsToBeSharedWithEventHubs.ts

@@ -5,6 +5,7 @@
 
 import { Span, SpanContext } from "@opentelemetry/api";
 import { OperationOptions } from "@azure/core-http";
+import { OperationTracingOptions } from "@azure/core-tracing";
 
 /**
  * NOTE: This type is intended to mirror the relevant fields and structure from @azure/core-http OperationOptions
@@ -18,7 +19,22 @@ export type OperationOptionsBase = Pick<OperationOptions, "abortSignal" | "traci
  * @ignore
  */
 export function getParentSpan(
-  options: Pick<OperationOptionsBase, "tracingOptions">
+  options?: OperationTracingOptions
 ): Span | SpanContext | null | undefined {
-  return options.tracingOptions?.spanOptions?.parent;
+  return options?.spanOptions?.parent;
 }
+
+/**
+ * The set of options to manually propagate `Span` context for distributed tracing.
+ * - `parentSpan` : The `Span` or `SpanContext` for the operation to use as a `parent` when creating its own span.
+ */
+export interface ParentSpanOptions {
+  /**
+   * The `Span` or `SpanContext` to use as the `parent` of any spans created while calling operations that make a request to the service.
+   */
+  parentSpan?: Span | SpanContext | null;
+}
+/**
+ * Options to configure the behavior of the `tryAdd` method on the `ServiceBusMessageBatch` class.
+ */
+export interface TryAddOptions extends ParentSpanOptions {}

sdk/servicebus/service-bus/src/sender.ts

@@ -21,7 +21,9 @@ import {
   RetryOptions,
   retry
 } from "@azure/core-amqp";
-import { OperationOptionsBase } from "./modelsToBeSharedWithEventHubs";
+import { getParentSpan, OperationOptionsBase } from "./modelsToBeSharedWithEventHubs";
+import { CanonicalCode, Link, Span, SpanContext, SpanKind } from "@opentelemetry/api";
+import { getTracer } from "@azure/core-tracing";
 
 /**
  * A Sender can be used to send messages, schedule messages to be sent at a later time
@@ -179,14 +181,19 @@ export class ServiceBusSenderImpl implements ServiceBusSender {
     const invalidTypeErrMsg =
       "Provided value for 'messages' must be of type ServiceBusMessage, ServiceBusMessageBatch or an array of type ServiceBusMessage.";
 
+    // link message span contexts
+    let spanContextsToLink: SpanContext[] = [];
+    if (isServiceBusMessage(messages)) {
+      messages = [messages];
+    }
+    let batch: ServiceBusMessageBatch;
     if (Array.isArray(messages)) {
-      const batch = await this.createBatch(options);
-
-      for (const message of messages) {
+      batch = await this.createBatch(options);
+      for (let message of messages) {
         if (!isServiceBusMessage(message)) {
           throw new TypeError(invalidTypeErrMsg);
         }
-        if (!batch.tryAdd(message)) {
+        if (!batch.tryAdd(message, { parentSpan: getParentSpan(options?.tracingOptions) })) {
           // this is too big - throw an error
           const error = new MessagingError(
             "Messages were too big to fit in a single batch. Remove some messages and try again or create your own batch using createBatch(), which gives more fine-grained control."
@@ -195,14 +202,31 @@ export class ServiceBusSenderImpl implements ServiceBusSender {
           throw error;
         }
       }
-
-      return this._sender.sendBatch(batch, options);
     } else if (isServiceBusMessageBatch(messages)) {
-      return this._sender.sendBatch(messages, options);
-    } else if (isServiceBusMessage(messages)) {
-      return this._sender.send(messages, options);
+      spanContextsToLink = messages._messageSpanContexts;
+      batch = messages;
+    } else {
+      throw new TypeError(invalidTypeErrMsg);
+    }
+
+    const sendSpan = this._createSendSpan(
+      getParentSpan(options?.tracingOptions),
+      spanContextsToLink
+    );
+
+    try {
+      const result = await this._sender.sendBatch(batch, options);
+      sendSpan.setStatus({ code: CanonicalCode.OK });
+      return result;
+    } catch (error) {
+      sendSpan.setStatus({
+        code: CanonicalCode.UNKNOWN,
+        message: error.message
+      });
+      throw error;
+    } finally {
+      sendSpan.end();
     }
-    throw new TypeError(invalidTypeErrMsg);
   }
 
   async createBatch(options?: CreateBatchOptions): Promise<ServiceBusMessageBatch> {
@@ -323,6 +347,29 @@ export class ServiceBusSenderImpl implements ServiceBusSender {
       throw err;
     }
   }
+
+  private _createSendSpan(
+    parentSpan?: Span | SpanContext | null,
+    spanContextsToLink: SpanContext[] = []
+  ): Span {
+    const links: Link[] = spanContextsToLink.map((context) => {
+      return {
+        context
+      };
+    });
+    const tracer = getTracer();
+    const span = tracer.startSpan("Azure.ServiceBus.send", {
+      kind: SpanKind.CLIENT,
+      parent: parentSpan,
+      links
+    });
+
+    span.setAttribute("az.namespace", "Microsoft.ServiceBus");
+    span.setAttribute("message_bus.destination", this.entityPath);
+    span.setAttribute("peer.address", this._context.config.host);
+
+    return span;
+  }
 }
 
 /**

sdk/servicebus/service-bus/src/serviceBusMessageBatch.ts

@@ -10,6 +10,13 @@ import {
   message as RheaMessageUtil
 } from "rhea-promise";
 import { AmqpMessage } from "@azure/core-amqp";
+import { SpanContext } from "@opentelemetry/api";
+import {
+  instrumentServiceBusMessage,
+  TRACEPARENT_PROPERTY
+} from "./diagnostics/instrumentServiceBusMessage";
+import { createMessageSpan } from "./diagnostics/messageSpan";
+import { TryAddOptions } from "./modelsToBeSharedWithEventHubs";
 
 /**
  * @internal
@@ -65,7 +72,7 @@ export interface ServiceBusMessageBatch {
    * @param message  An individual service bus message.
    * @returns A boolean value indicating if the message has been added to the batch or not.
    */
-  tryAdd(message: ServiceBusMessage): boolean;
+  tryAdd(message: ServiceBusMessage, options?: TryAddOptions): boolean;
 
   /**
    * The AMQP message containing encoded events that were added to the batch.
@@ -77,6 +84,14 @@ export interface ServiceBusMessageBatch {
    * @ignore
    */
   _generateMessage(): Buffer;
+
+  /**
+   * Gets the "message" span contexts that were created when adding events to the batch.
+   * Used internally by the `sendBatch()` method to set up the right spans in traces if tracing is enabled.
+   * @internal
+   * @ignore
+   */
+  readonly _messageSpanContexts: SpanContext[];
 }
 
 /**
@@ -95,6 +110,10 @@ export class ServiceBusMessageBatchImpl implements ServiceBusMessageBatch {
    * @property Encoded amqp messages.
    */
   private _encodedMessages: Buffer[] = [];
+  /**
+   * List of 'message' span contexts.
+   */
+  private _spanContexts: SpanContext[] = [];
   /**
    * ServiceBusMessageBatch should not be constructed using `new ServiceBusMessageBatch()`
    * Use the `createBatch()` method on your `Sender` instead.
@@ -132,6 +151,15 @@ export class ServiceBusMessageBatchImpl implements ServiceBusMessageBatch {
     return this._encodedMessages.length;
   }
 
+  /**
+   * Gets the "message" span contexts that were created when adding messages to the batch.
+   * @internal
+   * @ignore
+   */
+  get _messageSpanContexts(): SpanContext[] {
+    return this._spanContexts;
+  }
+
   /**
    * Generates an AMQP message that contains the provided encoded messages and annotations.
    *
@@ -210,12 +238,24 @@ export class ServiceBusMessageBatchImpl implements ServiceBusMessageBatch {
    * @param message  An individual service bus message.
    * @returns A boolean value indicating if the message has been added to the batch or not.
    */
-  public tryAdd(message: ServiceBusMessage): boolean {
+  public tryAdd(message: ServiceBusMessage, options: TryAddOptions = {}): boolean {
     throwTypeErrorIfParameterMissing(this._context.connectionId, "message", message);
     if (!isServiceBusMessage(message)) {
       throw new TypeError("Provided value for 'message' must be of type ServiceBusMessage.");
     }
 
+    // check if the event has already been instrumented
+    const previouslyInstrumented = Boolean(
+      message.properties && message.properties[TRACEPARENT_PROPERTY]
+    );
+    let spanContext: SpanContext | undefined;
+    if (!previouslyInstrumented) {
+      const messageSpan = createMessageSpan(options?.parentSpan, this._context.config);
+      message = instrumentServiceBusMessage(message, messageSpan);
+      spanContext = messageSpan.context();
+      messageSpan.end();
+    }
+
     // Convert ServiceBusMessage to AmqpMessage.
     const amqpMessage = toAmqpMessage(message);
     amqpMessage.body = this._context.dataTransformer.encode(message.body);
@@ -261,6 +301,9 @@ export class ServiceBusMessageBatchImpl implements ServiceBusMessageBatch {
 
     // The message will fit in the batch, so it is now safe to store it.
     this._encodedMessages.push(encodedMessage);
+    if (spanContext) {
+      this._spanContexts.push(spanContext);
+    }
 
     this._sizeInBytes = currentSize;
     return true;