Fixing it so we check to see if a created span is actually being recorded before we add it to the links for the overall span. If not, then we don't bother since it'll be the empty/null span and won't actually record anything of interest. (Azure#14040)

richardpark-msft · web-flow · commit d28b4ae11a00 · 2021-03-01T15:43:43.000-08:00
When we were instrumenting spans we wouldn't pay attention to whether we were currently recording or not, which could lead to us creating a bunch of empty/null spans and trying to add them as links to the span for our batches. Rather than do that I just check (when we instrument) if the span we created is recording and if not don't add it. Fixes Azure#13049
diff --git a/sdk/servicebus/service-bus/src/diagnostics/instrumentServiceBusMessage.ts b/sdk/servicebus/service-bus/src/diagnostics/instrumentServiceBusMessage.ts
@@ -1,7 +1,7 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT license.
 
-import { extractSpanContextFromTraceParentHeader, getTraceParentHeader } from "@azure/core-tracing";
+import { extractSpanContextFromTraceParentHeader } from "@azure/core-tracing";
 import { CanonicalCode, Link, Span, SpanContext, SpanKind } from "@opentelemetry/api";
 import { ConnectionContext } from "../connectionContext";
 import { OperationOptionsBase } from "../modelsToBeSharedWithEventHubs";
@@ -14,34 +14,6 @@ import { createServiceBusSpan } from "./tracing";
  */
 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.
- * @hidden
- * @internal
- */
-export function instrumentServiceBusMessage(
-  message: ServiceBusMessage,
-  span: Span
-): ServiceBusMessage {
-  if (message.applicationProperties && message.applicationProperties[TRACEPARENT_PROPERTY]) {
-    return message;
-  }
-
-  // create a copy so the original isn't modified
-  message = { ...message, applicationProperties: { ...message.applicationProperties } };
-
-  const traceParent = getTraceParentHeader(span.context());
-  if (traceParent) {
-    message.applicationProperties![TRACEPARENT_PROPERTY] = traceParent;
-  }
-
-  return message;
-}
-
 /**
  * Extracts the `SpanContext` from an `ServiceBusMessage` if the context exists.
  * @param message - An individual `ServiceBusMessage` object.
diff --git a/sdk/servicebus/service-bus/src/diagnostics/tracing.ts b/sdk/servicebus/service-bus/src/diagnostics/tracing.ts
@@ -2,8 +2,14 @@
 // Licensed under the MIT license.
 
 import { OperationOptions, RestError } from "@azure/core-http";
-import { CanonicalCode, Span, SpanKind } from "@opentelemetry/api";
-import { createSpanFunction, SpanOptions } from "@azure/core-tracing";
+import { CanonicalCode, Span, SpanContext, SpanKind } from "@opentelemetry/api";
+import {
+  createSpanFunction,
+  extractSpanContextFromTraceParentHeader,
+  getTraceParentHeader,
+  SpanOptions
+} from "@azure/core-tracing";
+import { ServiceBusMessage } from "../serviceBusMessage";
 
 /**
  * Creates a span using the global tracer.
@@ -75,3 +81,101 @@ export function createServiceBusSpan(
     updatedOptions
   };
 }
+
+/**
+ * @internal
+ */
+export const TRACEPARENT_PROPERTY = "Diagnostic-Id";
+
+/**
+ * @hidden
+ */
+export interface InstrumentableMessage {
+  /**
+   * The application specific properties which can be
+   * used for custom message metadata.
+   */
+  applicationProperties?: { [key: string]: number | boolean | string | Date };
+}
+
+/**
+ * Instruments an AMQP message with a proper `Diagnostic-Id` for tracing.
+ *
+ * @hidden
+ */
+export function instrumentMessage<T extends InstrumentableMessage>(
+  message: T,
+  options: OperationOptions,
+  entityPath: string,
+  host: string
+): {
+  /**
+   * If instrumentation was done, a copy of the message with
+   * message.applicationProperties['Diagnostic-Id'] filled
+   * out appropriately.
+   */
+  message: T;
+
+  /**
+   * A valid SpanContext if this message should be linked to a parent span, or undefined otherwise.
+   */
+  spanContext: SpanContext | undefined;
+} {
+  const previouslyInstrumented = Boolean(
+    message.applicationProperties && message.applicationProperties[TRACEPARENT_PROPERTY]
+  );
+
+  if (previouslyInstrumented) {
+    return {
+      message,
+      spanContext: undefined
+    };
+  }
+
+  const { span: messageSpan } = createMessageSpan(options, entityPath, host);
+
+  try {
+    if (!messageSpan.isRecording()) {
+      return {
+        message,
+        spanContext: undefined
+      };
+    }
+
+    const traceParent = getTraceParentHeader(messageSpan.context());
+
+    if (traceParent) {
+      // create a copy so the original isn't modified
+      message = {
+        ...message,
+        applicationProperties: {
+          ...message.applicationProperties,
+          [TRACEPARENT_PROPERTY]: traceParent
+        }
+      };
+    }
+
+    return {
+      message,
+      spanContext: messageSpan.context()
+    };
+  } finally {
+    messageSpan.end();
+  }
+}
+
+/**
+ * Extracts the `SpanContext` from an `ServiceBusMessage` if the context exists.
+ * @param message - An individual `ServiceBusMessage` object.
+ * @internal
+ */
+export function extractSpanContextFromServiceBusMessage(
+  message: ServiceBusMessage
+): SpanContext | undefined {
+  if (!message.applicationProperties || !message.applicationProperties[TRACEPARENT_PROPERTY]) {
+    return;
+  }
+
+  const diagnosticId = message.applicationProperties[TRACEPARENT_PROPERTY] as string;
+  return extractSpanContextFromTraceParentHeader(diagnosticId);
+}
diff --git a/sdk/servicebus/service-bus/src/serviceBusMessageBatch.ts b/sdk/servicebus/service-bus/src/serviceBusMessageBatch.ts
@@ -15,11 +15,7 @@ import {
   Message as RheaMessage
 } from "rhea-promise";
 import { SpanContext } from "@opentelemetry/api";
-import {
-  instrumentServiceBusMessage,
-  TRACEPARENT_PROPERTY
-} from "./diagnostics/instrumentServiceBusMessage";
-import { createMessageSpan } from "./diagnostics/tracing";
+import { instrumentMessage } from "./diagnostics/tracing";
 import { TryAddOptions } from "./modelsToBeSharedWithEventHubs";
 import { defaultDataTransformer } from "./dataTransformer";
 
@@ -230,32 +226,22 @@ export class ServiceBusMessageBatchImpl implements ServiceBusMessageBatch {
    * **NOTE**: Always remember to check the return value of this method, before calling it again
    * for the next message.
    *
-   * @param message - An individual service bus message.
+   * @param originalMessage - An individual service bus message.
    * @returns A boolean value indicating if the message has been added to the batch or not.
    */
-  public tryAddMessage(message: ServiceBusMessage, options: TryAddOptions = {}): boolean {
-    throwTypeErrorIfParameterMissing(this._context.connectionId, "message", message);
+  public tryAddMessage(originalMessage: ServiceBusMessage, options: TryAddOptions = {}): boolean {
+    throwTypeErrorIfParameterMissing(this._context.connectionId, "message", originalMessage);
     throwIfNotValidServiceBusMessage(
-      message,
+      originalMessage,
       "Provided value for 'message' must be of type ServiceBusMessage."
     );
 
-    // check if the event has already been instrumented
-    const previouslyInstrumented = Boolean(
-      message.applicationProperties && message.applicationProperties[TRACEPARENT_PROPERTY]
+    const { message, spanContext } = instrumentMessage(
+      originalMessage,
+      options,
+      this._context.config.entityPath!,
+      this._context.config.host
     );
-    let spanContext: SpanContext | undefined;
-    if (!previouslyInstrumented) {
-      const { span: messageSpan } = createMessageSpan(
-        options,
-        this._context.config.entityPath!,
-        this._context.config.host
-      );
-
-      message = instrumentServiceBusMessage(message, messageSpan);
-      spanContext = messageSpan.context();
-      messageSpan.end();
-    }
 
     // Convert ServiceBusMessage to AmqpMessage.
     const amqpMessage = toRheaMessage(message);
diff --git a/sdk/servicebus/service-bus/test/internal/unit/tracing.spec.ts b/sdk/servicebus/service-bus/test/internal/unit/tracing.spec.ts
@@ -4,15 +4,10 @@
 import chai from "chai";
 import { getTracer, NoOpSpan, TestSpan, TestTracer } from "@azure/core-tracing";
 import { CanonicalCode, SpanOptions } from "@opentelemetry/api";
-import {
-  ServiceBusMessage,
-  ServiceBusMessageImpl,
-  ServiceBusReceivedMessage
-} from "../../../src/serviceBusMessage";
+import { ServiceBusMessageImpl, ServiceBusReceivedMessage } from "../../../src/serviceBusMessage";
 import {
   createAndEndProcessingSpan,
   createProcessingSpan,
-  instrumentServiceBusMessage,
   TRACEPARENT_PROPERTY
 } from "../../../src/diagnostics/instrumentServiceBusMessage";
 import { OperationOptionsBase, trace } from "../../../src/modelsToBeSharedWithEventHubs";
@@ -26,6 +21,7 @@ import { ServiceBusReceiverImpl } from "../../../src/receivers/receiver";
 import { OnMessage } from "../../../src/core/messageReceiver";
 import { ServiceBusSessionReceiverImpl } from "../../../src/receivers/sessionReceiver";
 import { MessageSession } from "../../../src/session/messageSession";
+import { instrumentMessage } from "../../../src/diagnostics/tracing";
 const should = chai.should();
 const assert = chai.assert;
 
@@ -291,7 +287,7 @@ describe("Tracing tests", () => {
   }
 
   describe("telemetry", () => {
-    const otherProperties = {
+    const receiverProperties = {
       entityPath: "entityPath"
     };
 
@@ -302,7 +298,7 @@ describe("Tracing tests", () => {
     it("basic span properties are set", async () => {
       const fakeParentSpanContext = new NoOpSpan().context();
 
-      createProcessingSpan([], otherProperties, connectionConfig, {
+      createProcessingSpan([], receiverProperties, connectionConfig, {
         tracingOptions: {
           spanOptions: {
             parent: fakeParentSpanContext
@@ -320,46 +316,79 @@ describe("Tracing tests", () => {
 
       attributes!.should.deep.equal({
         "az.namespace": "Microsoft.ServiceBus",
-        "message_bus.destination": otherProperties.entityPath,
+        "message_bus.destination": receiverProperties.entityPath,
         "peer.address": connectionConfig.host
       });
     });
 
-    it("received events are linked to this span using Diagnostic-Id", async () => {
+    it("already instrumented messages are skipped", () => {
+      const alreadyInstrumentedMessage = {
+        body: "hello",
+        enqueuedTimeUtc: new Date(),
+        applicationProperties: {
+          "Diagnostic-Id": "alreadyhasdiagnosticsid"
+        }
+      };
+
+      const { message, spanContext } = instrumentMessage(alreadyInstrumentedMessage, {}, "", "");
+
+      assert.equal(
+        message,
+        alreadyInstrumentedMessage,
+        "Messages that are already instrumented do not get copied"
+      );
+      assert.isUndefined(
+        spanContext,
+        "Messages that are already instrumented do not get a new Span (or SpanContext)"
+      );
+      assert.isUndefined(
+        tracer.spanOptions,
+        "No spans should be created for already instrumented messages"
+      );
+    });
+
+    it("Round trip test - instrument message as it would be for sending and then receive and process them", async () => {
+      // the sender side...
       const requiredMessageProperties = {
         body: "hello",
-        enqueuedTimeUtc: new Date()
+        enqueuedTimeUtc: new Date(),
+        applicationProperties: undefined
       };
 
-      const firstEvent = tracer.startSpan("a");
-      const thirdEvent = tracer.startSpan("c");
+      const originalMessage = { ...requiredMessageProperties };
+
+      const { message, spanContext } = instrumentMessage(originalMessage, {}, "", "");
+      assert.ok(
+        spanContext,
+        "A span context should be created when we instrument a message for the first time"
+      );
+      assert.notEqual(message, originalMessage, "Instrumenting a message should copy it");
+
+      assert.ok(tracer.spanOptions, "A span should be created when we instrumented the messsage");
+      const spanContextFromSender = tracer.span?.context();
+      assert.ok(spanContextFromSender);
+
+      tracer.clearTracingData();
 
-      const receivedMessages: ServiceBusMessage[] = [
-        instrumentServiceBusMessage({ ...requiredMessageProperties }, firstEvent),
-        { applicationProperties: {}, ...requiredMessageProperties }, // no diagnostic ID means it gets skipped
-        instrumentServiceBusMessage({ ...requiredMessageProperties }, thirdEvent)
-      ];
+      // the receiver side...
 
-      createProcessingSpan(
-        (receivedMessages as any) as ServiceBusReceivedMessage[],
-        otherProperties,
+      // messages have been received. Now we'll create a span and link the received spans to it.
+      const processingSpan = createProcessingSpan(
+        ([message] as any) as ServiceBusReceivedMessage[],
+        receiverProperties,
         connectionConfig,
         {}
       );
 
-      // middle event, since it has no trace information, doesn't get included
-      // in the telemetry
-      tracer.spanOptions!.links!.length.should.equal(3 - 1);
-      // the test tracer just hands out a string integer that just gets
-      // incremented
-      tracer.spanOptions!.links![0]!.context.traceId.should.equal(firstEvent.context().traceId);
+      assert.ok(processingSpan);
+
+      const processingSpanOptions = tracer.spanOptions;
+      assert.ok(processingSpanOptions);
+
+      tracer.spanOptions!.links![0]!.context.traceId.should.equal(spanContextFromSender?.traceId);
       (tracer.spanOptions!.links![0]!.attributes!.enqueuedTime as number).should.equal(
         requiredMessageProperties.enqueuedTimeUtc.getTime()
       );
-      tracer.spanOptions!.links![1]!.context.traceId.should.equal(thirdEvent.context().traceId);
-      (tracer.spanOptions!.links![1]!.attributes!.enqueuedTime as number).should.equal(
-        requiredMessageProperties.enqueuedTimeUtc.getTime()
-      );
     });
 
     it("trace - normal", async () => {
@@ -388,12 +417,20 @@ describe("Tracing tests", () => {
 });
 
 class TestTracer2 extends TestTracer {
-  public spanOptions: SpanOptions | undefined;
-  public spanName: string | undefined;
+  spanName: string | undefined;
+  spanOptions: SpanOptions | undefined;
+  span: TestSpan | undefined;
+
+  clearTracingData() {
+    this.spanName = undefined;
+    this.spanOptions = undefined;
+    this.span = undefined;
+  }
 
   startSpan(nameArg: string, optionsArg?: SpanOptions): TestSpan {
     this.spanName = nameArg;
     this.spanOptions = optionsArg;
-    return super.startSpan(nameArg, optionsArg);
+    this.span = super.startSpan(nameArg, optionsArg);
+    return this.span;
   }
 }
