[Service Bus] Tracing for send API

sdk/servicebus/service-bus/CHANGELOG.md

@@ -10,6 +10,8 @@
   [PR 11250](https://github.com/Azure/azure-sdk-for-js/pull/11250)
 - "properties" in the correlation rule filter now supports `Date`.
   [PR 11117](https://github.com/Azure/azure-sdk-for-js/pull/11117)
+- `sendMessages` method on the sender and `tryAdd` method to add messages to a batch now support tracing.
+  [PR 11651](https://github.com/Azure/azure-sdk-for-js/pull/11651)
 - Message locks can be auto-renewed in all receive methods (receiver.receiveMessages, receiver.subcribe
   and receiver.getMessageIterator). This can be configured in options when calling `ServiceBusClient.createReceiver()`.
   [PR 11658](https://github.com/Azure/azure-sdk-for-js/pull/11658)

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';
@@ -395,8 +397,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 +581,11 @@ export interface TopicRuntimeProperties {
 export interface TopicRuntimePropertiesResponse extends TopicRuntimeProperties, Response {
 }
 
+// @public
+export interface TryAddOptions {
+    parentSpan?: Span | SpanContext | null;
+}
+
 export { WebSocketImpl }
 
 export { WebSocketOptions }

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

@@ -329,29 +329,14 @@ export namespace ConnectionContext {
         await delay(Constants.connectionReconnectDelay);
 
         const detachCalls: Promise<void>[] = [];
-
-        // Call onDetached() on sender so that it can gracefully shutdown
-        for (const senderName of Object.keys(connectionContext.senders)) {
-          const sender = connectionContext.senders[senderName];
-          if (sender) {
-            logger.verbose(
-              "[%s] calling detached on sender '%s'.",
-              connectionContext.connection.id,
-              sender.name
-            );
-            detachCalls.push(
-              sender.onDetached().catch((err) => {
-                logError(
-                  err,
-                  "[%s] An error occurred while calling onDetached() the sender '%s': %O.",
-                  connectionContext.connection.id,
-                  sender.name,
-                  err
-                );
-              })
-            );
-          }
-        }
+        // Neither we do recovery for the sender, nor we cleanup
+        // No recovery:
+        //   Because we don't want to keep the sender active all the time
+        //   and the "next" send call would bear the burden of creating the link
+        // No cleanup:
+        //   "Closing the link" cleanup would step over new link initializations
+        //   and can possibly clear the link once created, hence we do the cleanup
+        //   at the time of new link creation
 
         // Call onDetached() on receivers so that batching receivers it can gracefully close any ongoing batch operation
         // and streaming receivers can decide whether to reconnect or not.

sdk/servicebus/service-bus/src/core/linkEntity.ts

@@ -262,6 +262,12 @@ export abstract class LinkEntity<LinkT extends Receiver | AwaitableSender | Requ
       this.checkIfConnectionReady();
 
       logger.verbose(`${this._logPrefix} Creating with options %O`, options);
+      if (this.link) {
+        // For the cases such as connection is disconnected, onAmqpClose and onSessionClose events,
+        // we don't clear the link for sender. To compensate for the above,
+        // we are calling close on the link to cleanup before creating the new one.
+        await this.link.close();
+      }
       this._link = await this.createRheaLink(options);
       checkAborted();
 

sdk/servicebus/service-bus/src/core/messageSender.ts

@@ -110,14 +110,6 @@ export class MessageSender extends LinkEntity<AwaitableSender> {
         `${this.logPrefix} 'sender_close' event occurred. The associated error is: %O`,
         senderError
       );
-
-      await this.onDetached().catch((err) => {
-        logError(
-          err,
-          `${this.logPrefix} error when closing sender after 'sender_close' event: %O`,
-          err
-        );
-      });
     };
 
     this._onSessionClose = async (context: EventContext) => {
@@ -128,14 +120,6 @@ export class MessageSender extends LinkEntity<AwaitableSender> {
         `${this.logPrefix} 'session_close' event occurred. The associated error is: %O`,
         sessionError
       );
-
-      await this.onDetached().catch((err) => {
-        logError(
-          err,
-          `${this.logPrefix} error when closing sender after 'session_close' event: %O`,
-          err
-        );
-      });
     };
   }
 
@@ -333,17 +317,6 @@ export class MessageSender extends LinkEntity<AwaitableSender> {
     }
   }
 
-  /**
-   * Closes the rhea link.
-   * To be called when connection is disconnected, onAmqpClose and onSessionClose events.
-   * @returns {Promise<void>} Promise<void>.
-   */
-  async onDetached(): Promise<void> {
-    // Clears the token renewal timer. Closes the link and its session if they are open.
-    // Removes the link and its session if they are present in rhea's cache.
-    await this.closeLink();
-  }
-
   /**
    * Determines whether the AMQP sender link is open. If open then returns true else returns false.
    * @return {boolean} boolean

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,7 @@ export {
   SubQueue,
   SubscribeOptions
 } from "./models";
-export { OperationOptionsBase } from "./modelsToBeSharedWithEventHubs";
+export { OperationOptionsBase, TryAddOptions } from "./modelsToBeSharedWithEventHubs";
 export { ServiceBusReceiver } from "./receivers/receiver";
 export { ServiceBusSessionReceiver } from "./receivers/sessionReceiver";
 export { ServiceBusSender } from "./sender";

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

@@ -3,8 +3,9 @@
 
 // TODO: this code is a straight-copy from EventHubs. Need to merge.
 
-import { Span, SpanContext } from "@opentelemetry/api";
+import { Link, Span, SpanContext, SpanKind } from "@opentelemetry/api";
 import { OperationOptions } from "@azure/core-http";
+import { getTracer, 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,41 @@ 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;
+}
+
+export function createSendSpan(
+  parentSpan?: Span | SpanContext | null,
+  spanContextsToLink: SpanContext[] = [],
+  entityPath?: string,
+  host?: string
+): 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", entityPath);
+  span.setAttribute("peer.address", host);
+
+  return span;
+}
+/**
+ * The set of options to manually propagate `Span` context for distributed tracing.
+ */
+export interface TryAddOptions {
+  /**
+   * 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;
 }

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

@@ -21,7 +21,12 @@ import {
   RetryOptions,
   retry
 } from "@azure/core-amqp";
-import { OperationOptionsBase } from "./modelsToBeSharedWithEventHubs";
+import {
+  createSendSpan,
+  getParentSpan,
+  OperationOptionsBase
+} from "./modelsToBeSharedWithEventHubs";
+import { CanonicalCode, SpanContext } from "@opentelemetry/api";
 
 /**
  * A Sender can be used to send messages, schedule messages to be sent at a later time
@@ -179,14 +184,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);
-
+      batch = await this.createBatch(options);
       for (const 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 +205,33 @@ 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 = createSendSpan(
+      getParentSpan(options?.tracingOptions),
+      spanContextsToLink,
+      this.entityPath,
+      this._context.config.host
+    );
+
+    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> {