dd-trace/packages/datadog-plugin-aws-sdk/src/services/sqs.js

Datadog APM tracing client for JavaScript

'use strict'

const log = require('../../../dd-trace/src/log')
const BaseAwsSdkPlugin = require('../base')
const { DsmPathwayCodec, getHeadersSize } = require('../../../dd-trace/src/datastreams')
const { extractQueueMetadata, isEmpty } = require('../util')

class Sqs extends BaseAwsSdkPlugin {
  static id = 'sqs'
  static peerServicePrecursors = ['queuename']
  static isPayloadReporter = true

  constructor (...args) {
    super(...args)
    //
    // TODO(bengl) Find a way to create the response span tags without this WeakMap being populated
    // in the base class
    this.requestTags = new WeakMap()

    this.addBind('apm:aws:response:start:sqs', ctx => {
      const { request, response } = ctx
      const contextExtraction = this.responseExtract(request.params, request.operation, response)

      let store = this._parentMap.get(request)
      let span
      let parsedMessageAttributes
      let parsedFirstBody
      let firstBodyChecked = false
      if (contextExtraction !== undefined) {
        parsedFirstBody = contextExtraction.parsedBody
        firstBodyChecked = contextExtraction.bodyChecked === true
        if (contextExtraction.datadogContext !== undefined) {
          ctx.needsFinish = true
          const options = {
            childOf: contextExtraction.datadogContext,
            meta: {
              ...this.requestTags.get(request),
              'span.kind': 'server',
            },
            integrationName: 'aws-sdk',
          }
          parsedMessageAttributes = contextExtraction.parsedAttributes
          span = this.startSpan('aws.response', options, ctx)
          store = ctx.currentStore
        }
      }

      // Extract DSM context after, as we might not have a parent-child but may have a DSM context.
      this.responseExtractDSMContext(
        request.operation, request.params, response, span ?? null,
        { parsedAttributes: parsedMessageAttributes, parsedFirstBody, firstBodyChecked }
      )

      return store
    })

    this.addSub('apm:aws:response:finish:sqs', ctx => {
      if (!ctx.needsFinish) return
      this.finish(ctx)
    })
  }

  operationFromRequest (request) {
    switch (request.operation) {
      case 'receiveMessage':
        return this.operationName({
          type: 'messaging',
          kind: 'consumer',
        })
      case 'sendMessage':
      case 'sendMessageBatch':
        return this.operationName({
          type: 'messaging',
          kind: 'producer',
        })
    }

    return this.operationName({
      id: 'aws',
      type: 'web',
      kind: 'client',
      awsService: 'sqs',
    })
  }

  isEnabled (request) {
    // TODO(bengl) Figure out a way to make separate plugins for consumer and producer so that
    // config can be isolated to `.configure()` instead of this whole isEnabled() thing.
    const config = this.config
    switch (request.operation) {
      case 'receiveMessage':
        return config.consumer !== false
      case 'sendMessage':
      case 'sendMessageBatch':
        return config.producer !== false
      default:
        return true
    }
  }

  generateTags (params, operation, response) {
    if (!params || (!params.QueueName && !params.QueueUrl)) return

    const queueMetadata = extractQueueMetadata(params.QueueUrl)
    const queueName = queueMetadata?.queueName || params.QueueName

    const tags = {
      'resource.name': `${operation} ${params.QueueName || params.QueueUrl}`,
      'aws.sqs.queue_name': params.QueueName || params.QueueUrl,
      'messaging.system': 'aws_sqs',
      queuename: queueName,
    }

    if (queueMetadata?.arn) {
      tags['cloud.resource_id'] = queueMetadata.arn
    }

    switch (operation) {
      case 'receiveMessage':
        tags['span.type'] = 'worker'
        tags['span.kind'] = 'consumer'
        break
      case 'sendMessage':
      case 'sendMessageBatch':
        tags['span.kind'] = 'producer'
        break
    }

    return tags
  }

  responseExtract (params, operation, response) {
    if (operation !== 'receiveMessage') return
    if (params.MaxNumberOfMessages && params.MaxNumberOfMessages !== 1) return
    if (!response || !response.Messages || !response.Messages[0]) return

    let message = response.Messages[0]
    let parsedBody

    if (message.Body) {
      try {
        parsedBody = JSON.parse(message.Body)
      } catch {
        // SQS to SQS
      }
      // SNS to SQS
      if (parsedBody?.Type === 'Notification') {
        message = parsedBody
      }
    }

    if (!message.MessageAttributes || !message.MessageAttributes._datadog) {
      return { parsedBody, bodyChecked: true }
    }

    const datadogAttribute = message.MessageAttributes._datadog

    const parsedAttributes = this.parseDatadogAttributes(datadogAttribute)
    if (parsedAttributes) {
      return {
        datadogContext: this.tracer.extract('text_map', parsedAttributes),
        parsedAttributes,
        parsedBody,
        bodyChecked: true,
      }
    }

    return { parsedBody, bodyChecked: true }
  }

  parseDatadogAttributes (attributes) {
    try {
      if (attributes.StringValue) {
        const textMap = attributes.StringValue
        return JSON.parse(textMap)
      } else if (attributes.Type === 'Binary' || attributes.DataType === 'Binary') {
        const buffer = Buffer.from(attributes.Value ?? attributes.BinaryValue, 'base64')
        return JSON.parse(buffer)
      }
    } catch (error) {
      log.error('Sqs error parsing DD attributes', error)
    }
  }

  responseExtractDSMContext (operation, params, response, span, kwargs = {}) {
    let { parsedAttributes } = kwargs
    const { parsedFirstBody, firstBodyChecked } = kwargs
    if (!this.config.dsmEnabled) return
    if (operation !== 'receiveMessage') return
    if (!response || !response.Messages || !response.Messages[0]) return

    // Only attribute payloadSize to the span when there is a single message.
    span = response.Messages.length > 1 ? null : span

    // QueueUrl is the same for the whole receive batch.
    const queue = params.QueueUrl.slice(params.QueueUrl.lastIndexOf('/') + 1)

    for (let i = 0; i < response.Messages.length; i++) {
      let message = response.Messages[i]
      if (!parsedAttributes) {
        let body
        // responseExtract already parsed message[0]; reuse that result instead of re-parsing.
        if (i === 0 && firstBodyChecked) {
          body = parsedFirstBody
        } else if (message.Body) {
          try {
            body = JSON.parse(message.Body)
          } catch {
            // SQS to SQS
          }
        }
        // SNS to SQS
        if (body?.Type === 'Notification') {
          message = body
        }
        if (message.MessageAttributes && message.MessageAttributes._datadog) {
          parsedAttributes = this.parseDatadogAttributes(message.MessageAttributes._datadog)
        }
      }
      const payloadSize = getHeadersSize({
        Body: message.Body,
        MessageAttributes: message.MessageAttributes,
      })
      if (parsedAttributes) {
        this.tracer.decodeDataStreamsContext(parsedAttributes)
      }
      this.tracer
        .setCheckpoint(['direction:in', `topic:${queue}`, 'type:sqs'], span, payloadSize)
    }
  }

  requestInject (span, request) {
    const { operation, params } = request

    if (!params) return

    switch (operation) {
      case 'sendMessage':
        this.injectToMessage(span, params, params.QueueUrl, true)
        break
      case 'sendMessageBatch':
        for (let i = 0; i < params.Entries.length; i++) {
          this.injectToMessage(
            span,
            params.Entries[i],
            params.QueueUrl,
            i === 0 || (this.config.batchPropagationEnabled)
          )
        }
        break
      case 'receiveMessage':
        if (!params.MessageAttributeNames) {
          params.MessageAttributeNames = ['_datadog']
        } else if (
          !params.MessageAttributeNames.includes('_datadog') &&
          !params.MessageAttributeNames.includes('.*') &&
          !params.MessageAttributeNames.includes('All')
        ) {
          params.MessageAttributeNames.push('_datadog')
        }
        break
    }
  }

  injectToMessage (span, params, queueUrl, injectTraceContext) {
    if (!params) {
      params = {}
    }
    if (!params.MessageAttributes) {
      params.MessageAttributes = {}
    } else if (Object.keys(params.MessageAttributes).length >= 10) { // SQS quota
      // TODO: add test when the test suite is fixed
      return
    }

    const ddInfo = {}
    // For now we only inject to the first message; batches may change later.
    if (injectTraceContext) {
      this.tracer.inject(span, 'text_map', ddInfo)
    }

    if (this.config.dsmEnabled) {
      // Attach `_datadog` before measuring so the DSM payload size metric
      // matches the on-wire payload, then update with the encoded context.
      params.MessageAttributes._datadog = {
        DataType: 'String',
        StringValue: JSON.stringify(ddInfo),
      }
      const dataStreamsContext = this.setDSMCheckpoint(span, params, queueUrl)
      if (dataStreamsContext) {
        DsmPathwayCodec.encode(dataStreamsContext, ddInfo)
        params.MessageAttributes._datadog.StringValue = JSON.stringify(ddInfo)
      } else if (isEmpty(ddInfo)) {
        delete params.MessageAttributes._datadog
      }
      return
    }

    if (isEmpty(ddInfo)) return

    params.MessageAttributes._datadog = {
      DataType: 'String',
      StringValue: JSON.stringify(ddInfo),
    }
  }

  setDSMCheckpoint (span, params, queueUrl) {
    const payloadSize = getHeadersSize({
      Body: params.MessageBody,
      MessageAttributes: params.MessageAttributes,
    })
    const queue = queueUrl.slice(queueUrl.lastIndexOf('/') + 1)
    return this.tracer
      .setCheckpoint(['direction:out', `topic:${queue}`, 'type:sqs'], span, payloadSize)
  }
}

module.exports = Sqs