dd-trace 4.18.0 → 5.6.0

package/packages/datadog-plugin-aws-sdk/src/services/kinesis.js

@@ -1,10 +1,68 @@
 'use strict'
+const {
+  getSizeOrZero
+} = require('../../../dd-trace/src/datastreams/processor')
+const { DsmPathwayCodec } = require('../../../dd-trace/src/datastreams/pathway')
 const log = require('../../../dd-trace/src/log')
 const BaseAwsSdkPlugin = require('../base')
+const { storage } = require('../../../datadog-core')
+
 class Kinesis extends BaseAwsSdkPlugin {
   static get id () { return 'kinesis' }
   static get peerServicePrecursors () { return ['streamname'] }
 
+  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.addSub('apm:aws:response:start:kinesis', obj => {
+      const { request, response } = obj
+      const store = storage.getStore()
+      const plugin = this
+
+      // if we have either of these operations, we want to store the streamName param
+      // since it is not typically available during get/put records requests
+      if (request.operation === 'getShardIterator' || request.operation === 'listShards') {
+        this.storeStreamName(request.params, request.operation, store)
+        return
+      }
+
+      if (request.operation === 'getRecords') {
+        let span
+        const responseExtraction = this.responseExtract(request.params, request.operation, response)
+        if (responseExtraction && responseExtraction.maybeChildOf) {
+          obj.needsFinish = true
+          const options = {
+            childOf: responseExtraction.maybeChildOf,
+            tags: Object.assign(
+              {},
+              this.requestTags.get(request) || {},
+              { 'span.kind': 'server' }
+            )
+          }
+          span = plugin.tracer.startSpan('aws.response', options)
+          this.enter(span, store)
+        }
+
+        // get the stream name that should have been stored previously
+        const { streamName } = storage.getStore()
+
+        // extract DSM context after as we might not have a parent-child but may have a DSM context
+        this.responseExtractDSMContext(
+          request.operation, response, span || null, streamName
+        )
+      }
+    })
+
+    this.addSub('apm:aws:response:finish:kinesis', err => {
+      const { span } = storage.getStore()
+      this.finish(span, null, err)
+    })
+  }
+
   generateTags (params, operation, response) {
     if (!params || !params.StreamName) return {}
 
@@ -15,6 +73,55 @@ class Kinesis extends BaseAwsSdkPlugin {
     }
   }
 
+  storeStreamName (params, operation, store) {
+    if (!operation || (operation !== 'getShardIterator' && operation !== 'listShards')) return
+    if (!params || !params.StreamName) return
+
+    const streamName = params.StreamName
+    storage.enterWith({ ...store, streamName })
+  }
+
+  responseExtract (params, operation, response) {
+    if (operation !== 'getRecords') return
+    if (params.Limit && params.Limit !== 1) return
+    if (!response || !response.Records || !response.Records[0]) return
+
+    const record = response.Records[0]
+
+    try {
+      const decodedData = JSON.parse(Buffer.from(record.Data).toString())
+
+      return {
+        maybeChildOf: this.tracer.extract('text_map', decodedData._datadog),
+        parsedAttributes: decodedData._datadog
+      }
+    } catch (e) {
+      log.error(e)
+    }
+  }
+
+  responseExtractDSMContext (operation, response, span, streamName) {
+    if (!this.config.dsmEnabled) return
+    if (operation !== 'getRecords') return
+    if (!response || !response.Records || !response.Records[0]) return
+
+    // we only want to set the payloadSize on the span if we have one message, not repeatedly
+    span = response.Records.length > 1 ? null : span
+
+    response.Records.forEach(record => {
+      const parsedAttributes = JSON.parse(Buffer.from(record.Data).toString())
+
+      if (
+        parsedAttributes?._datadog && streamName && DsmPathwayCodec.contextExists(parsedAttributes._datadog)
+      ) {
+        const payloadSize = getSizeOrZero(record.Data)
+        this.tracer.decodeDataStreamsContext(parsedAttributes._datadog)
+        this.tracer
+          .setCheckpoint(['direction:in', `topic:${streamName}`, 'type:kinesis'], span, payloadSize)
+      }
+    })
+  }
+
   // AWS-SDK will b64 kinesis payloads
   // or will accept an already b64 encoded payload
   // This method handles both
@@ -32,40 +139,69 @@ class Kinesis extends BaseAwsSdkPlugin {
   }
 
   requestInject (span, request) {
-    const operation = request.operation
-    if (operation === 'putRecord' || operation === 'putRecords') {
-      if (!request.params) {
+    const { operation, params } = request
+    if (!params) return
+
+    let stream
+    switch (operation) {
+      case 'putRecord':
+        stream = params.StreamArn ? params.StreamArn : (params.StreamName ? params.StreamName : '')
+        this.injectToMessage(span, params, stream, true)
+        break
+      case 'putRecords':
+        stream = params.StreamArn ? params.StreamArn : (params.StreamName ? params.StreamName : '')
+        for (let i = 0; i < params.Records.length; i++) {
+          this.injectToMessage(span, params.Records[i], stream, i === 0)
+        }
+    }
+  }
+
+  injectToMessage (span, params, stream, injectTraceContext) {
+    if (!params) {
+      return
+    }
+
+    let parsedData
+    if (injectTraceContext || this.config.dsmEnabled) {
+      parsedData = this._tryParse(params.Data)
+      if (!parsedData) {
+        log.error('Unable to parse payload, unable to pass trace context or set DSM checkpoint (if enabled)')
         return
       }
+    }
+
+    const ddInfo = {}
+    // for now, we only want to inject to the first message, this may change for batches in the future
+    if (injectTraceContext) { this.tracer.inject(span, 'text_map', ddInfo) }
 
-      const traceData = {}
-      this.tracer.inject(span, 'text_map', traceData)
-      let injectPath
-      if (request.params.Records && request.params.Records.length > 0) {
-        injectPath = request.params.Records[0]
-      } else if (request.params.Data) {
-        injectPath = request.params
-      } else {
-        log.error('No valid payload passed, unable to pass trace context')
+    // set DSM hash if enabled
+    if (this.config.dsmEnabled) {
+      parsedData._datadog = ddInfo
+      const dataStreamsContext = this.setDSMCheckpoint(span, parsedData, stream)
+      DsmPathwayCodec.encode(dataStreamsContext, ddInfo)
+    }
+
+    if (Object.keys(ddInfo).length !== 0) {
+      parsedData._datadog = ddInfo
+      const finalData = Buffer.from(JSON.stringify(parsedData))
+      const byteSize = finalData.length
+      // Kinesis max payload size is 1MB
+      // So we must ensure adding DD context won't go over that (512b is an estimate)
+      if (byteSize >= 1048576) {
+        log.info('Payload size too large to pass context')
         return
       }
-      const parsedData = this._tryParse(injectPath.Data)
-      if (parsedData) {
-        parsedData._datadog = traceData
-        const finalData = Buffer.from(JSON.stringify(parsedData))
-        const byteSize = finalData.length
-        // Kinesis max payload size is 1MB
-        // So we must ensure adding DD context won't go over that (512b is an estimate)
-        if (byteSize >= 1048576) {
-          log.info('Payload size too large to pass context')
-          return
-        }
-        injectPath.Data = finalData
-      } else {
-        log.error('Unable to parse payload, unable to pass trace context')
-      }
+      params.Data = finalData
     }
   }
+
+  setDSMCheckpoint (span, parsedData, stream) {
+    // get payload size of request data
+    const payloadSize = Buffer.from(JSON.stringify(parsedData)).byteLength
+    const dataStreamsContext = this.tracer
+      .setCheckpoint(['direction:out', `topic:${stream}`, 'type:kinesis'], span, payloadSize)
+    return dataStreamsContext
+  }
 }
 
 module.exports = Kinesis

package/packages/datadog-plugin-aws-sdk/src/services/sns.js

@@ -1,4 +1,6 @@
 'use strict'
+const { getHeadersSize } = require('../../../dd-trace/src/datastreams/processor')
+const { DsmPathwayCodec } = require('../../../dd-trace/src/datastreams/pathway')
 const log = require('../../../dd-trace/src/log')
 const BaseAwsSdkPlugin = require('../base')
 
@@ -11,6 +13,7 @@ class Sns extends BaseAwsSdkPlugin {
 
     if (!params.TopicArn && !(response.data && response.data.TopicArn)) return {}
     const TopicArn = params.TopicArn || response.data.TopicArn
+
     // Split the ARN into its parts
     // ex.'arn:aws:sns:us-east-1:123456789012:my-topic'
     const arnParts = TopicArn.split(':')
@@ -52,17 +55,17 @@ class Sns extends BaseAwsSdkPlugin {
 
     switch (operation) {
       case 'publish':
-        this._injectMessageAttributes(span, params)
+        this.injectToMessage(span, params, params.TopicArn, true)
         break
       case 'publishBatch':
-        if (params.PublishBatchRequestEntries && params.PublishBatchRequestEntries.length > 0) {
-          this._injectMessageAttributes(span, params.PublishBatchRequestEntries[0])
+        for (let i = 0; i < params.PublishBatchRequestEntries.length; i++) {
+          this.injectToMessage(span, params.PublishBatchRequestEntries[i], params.TopicArn, i === 0)
         }
         break
     }
   }
 
-  _injectMessageAttributes (span, params) {
+  injectToMessage (span, params, topicArn, injectTraceContext) {
     if (!params.MessageAttributes) {
       params.MessageAttributes = {}
     }
@@ -70,11 +73,46 @@ class Sns extends BaseAwsSdkPlugin {
       log.info('Message attributes full, skipping trace context injection')
       return
     }
+
     const ddInfo = {}
-    this.tracer.inject(span, 'text_map', ddInfo)
-    params.MessageAttributes._datadog = {
-      DataType: 'Binary',
-      BinaryValue: Buffer.from(JSON.stringify(ddInfo)) // BINARY types are automatically base64 encoded
+    // for now, we only want to inject to the first message, this may change for batches in the future
+    if (injectTraceContext) {
+      this.tracer.inject(span, 'text_map', ddInfo)
+      // add ddInfo before checking DSM so we can include DD attributes in payload size
+      params.MessageAttributes._datadog = {
+        DataType: 'Binary',
+        BinaryValue: ddInfo
+      }
+    }
+
+    if (this.config.dsmEnabled) {
+      if (!params.MessageAttributes._datadog) {
+        params.MessageAttributes._datadog = {
+          DataType: 'Binary',
+          BinaryValue: ddInfo
+        }
+      }
+
+      const dataStreamsContext = this.setDSMCheckpoint(span, params, topicArn)
+      DsmPathwayCodec.encode(dataStreamsContext, ddInfo)
+    }
+
+    if (Object.keys(ddInfo).length !== 0) {
+      // BINARY types are automatically base64 encoded
+      params.MessageAttributes._datadog.BinaryValue = Buffer.from(JSON.stringify(ddInfo))
+    } else if (params.MessageAttributes._datadog) {
+      // let's avoid adding any additional information to payload if we failed to inject
+      delete params.MessageAttributes._datadog
+    }
+  }
+
+  setDSMCheckpoint (span, params, topicArn) {
+    // only set a checkpoint if publishing to a topic
+    if (topicArn) {
+      const payloadSize = getHeadersSize(params)
+      const dataStreamsContext = this.tracer
+        .setCheckpoint(['direction:out', `topic:${topicArn}`, 'type:sns'], span, payloadSize)
+      return dataStreamsContext
     }
   }
 }

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

@@ -3,6 +3,8 @@
 const log = require('../../../dd-trace/src/log')
 const BaseAwsSdkPlugin = require('../base')
 const { storage } = require('../../../datadog-core')
+const { getHeadersSize } = require('../../../dd-trace/src/datastreams/processor')
+const { DsmPathwayCodec } = require('../../../dd-trace/src/datastreams/pathway')
 
 class Sqs extends BaseAwsSdkPlugin {
   static get id () { return 'sqs' }
@@ -19,20 +21,27 @@ class Sqs extends BaseAwsSdkPlugin {
       const { request, response } = obj
       const store = storage.getStore()
       const plugin = this
-      const maybeChildOf = this.responseExtract(request.params, request.operation, response)
-      if (maybeChildOf) {
+      const contextExtraction = this.responseExtract(request.params, request.operation, response)
+      let span
+      let parsedMessageAttributes
+      if (contextExtraction && contextExtraction.datadogContext) {
         obj.needsFinish = true
         const options = {
-          childOf: maybeChildOf,
+          childOf: contextExtraction.datadogContext,
           tags: Object.assign(
             {},
             this.requestTags.get(request) || {},
             { 'span.kind': 'server' }
           )
         }
-        const span = plugin.tracer.startSpan('aws.response', options)
+        parsedMessageAttributes = contextExtraction.parsedAttributes
+        span = plugin.tracer.startSpan('aws.response', options)
         this.enter(span, store)
       }
+      // 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, parsedMessageAttributes || null
+      )
     })
 
     this.addSub('apm:aws:response:finish:sqs', err => {
@@ -133,38 +142,136 @@ class Sqs extends BaseAwsSdkPlugin {
 
     const datadogAttribute = message.MessageAttributes._datadog
 
+    const parsedAttributes = this.parseDatadogAttributes(datadogAttribute)
+    if (parsedAttributes) {
+      return {
+        datadogContext: this.tracer.extract('text_map', parsedAttributes),
+        parsedAttributes: parsedAttributes
+      }
+    }
+  }
+
+  parseDatadogAttributes (attributes) {
     try {
-      if (datadogAttribute.StringValue) {
-        const textMap = datadogAttribute.StringValue
-        return this.tracer.extract('text_map', JSON.parse(textMap))
-      } else if (datadogAttribute.Type === 'Binary') {
-        const buffer = Buffer.from(datadogAttribute.Value, 'base64')
-        return this.tracer.extract('text_map', JSON.parse(buffer))
+      if (attributes.StringValue) {
+        const textMap = attributes.StringValue
+        return JSON.parse(textMap)
+      } else if (attributes.Type === 'Binary') {
+        const buffer = Buffer.from(attributes.Value, 'base64')
+        return JSON.parse(buffer)
       }
     } catch (e) {
       log.error(e)
     }
   }
 
-  requestInject (span, request) {
-    const operation = request.operation
-    if (operation === 'sendMessage') {
-      if (!request.params) {
-        request.params = {}
+  responseExtractDSMContext (operation, params, response, span, parsedAttributes) {
+    if (!this.config.dsmEnabled) return
+    if (operation !== 'receiveMessage') return
+    if (!response || !response.Messages || !response.Messages[0]) return
+
+    // we only want to set the payloadSize on the span if we have one message
+    span = response.Messages.length > 1 ? null : span
+
+    response.Messages.forEach(message => {
+      // we may have already parsed the message attributes when extracting trace context
+      if (!parsedAttributes) {
+        if (message.Body) {
+          try {
+            const body = JSON.parse(message.Body)
+
+            // SNS to SQS
+            if (body.Type === 'Notification') {
+              message = body
+            }
+          } catch (e) {
+            // SQS to SQS
+          }
+        }
+        if (message.MessageAttributes && message.MessageAttributes._datadog) {
+          parsedAttributes = this.parseDatadogAttributes(message.MessageAttributes._datadog)
+        }
       }
-      if (!request.params.MessageAttributes) {
-        request.params.MessageAttributes = {}
-      } else if (Object.keys(request.params.MessageAttributes).length >= 10) { // SQS quota
-        // TODO: add test when the test suite is fixed
-        return
+      if (parsedAttributes && DsmPathwayCodec.contextExists(parsedAttributes)) {
+        const payloadSize = getHeadersSize({
+          Body: message.Body,
+          MessageAttributes: message.MessageAttributes
+        })
+        const queue = params.QueueUrl.split('/').pop()
+        this.tracer.decodeDataStreamsContext(parsedAttributes)
+        this.tracer
+          .setCheckpoint(['direction:in', `topic:${queue}`, 'type:sqs'], span, payloadSize)
       }
-      const ddInfo = {}
+    })
+  }
+
+  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)
+        }
+        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 want to inject to the first message, this may change for batches in the future
+    if (injectTraceContext) {
       this.tracer.inject(span, 'text_map', ddInfo)
-      request.params.MessageAttributes._datadog = {
+      params.MessageAttributes._datadog = {
         DataType: 'String',
         StringValue: JSON.stringify(ddInfo)
       }
     }
+
+    if (this.config.dsmEnabled) {
+      if (!params.MessageAttributes._datadog) {
+        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)
+      }
+    }
+
+    if (params.MessageAttributes._datadog && Object.keys(ddInfo).length === 0) {
+      // let's avoid adding any additional information to payload if we failed to inject
+      delete params.MessageAttributes._datadog
+    }
+  }
+
+  setDSMCheckpoint (span, params, queueUrl) {
+    const payloadSize = getHeadersSize({
+      Body: params.MessageBody,
+      MessageAttributes: params.MessageAttributes
+    })
+    const queue = queueUrl.split('/').pop()
+    const dataStreamsContext = this.tracer
+      .setCheckpoint(['direction:out', `topic:${queue}`, 'type:sqs'], span, payloadSize)
+    return dataStreamsContext
   }
 }
 

package/packages/datadog-plugin-child_process/src/index.js

@@ -0,0 +1,91 @@
+'use strict'
+
+const TracingPlugin = require('../../dd-trace/src/plugins/tracing')
+const scrubChildProcessCmd = require('./scrub-cmd-params')
+
+const MAX_ARG_SIZE = 4096 // 4kB
+
+function truncateCommand (cmdFields) {
+  let size = cmdFields[0].length
+  let truncated = false
+  for (let i = 1; i < cmdFields.length; i++) {
+    if (size >= MAX_ARG_SIZE) {
+      truncated = true
+      cmdFields[i] = ''
+      continue
+    }
+
+    const argLen = cmdFields[i].length
+    if (size < MAX_ARG_SIZE && size + argLen > MAX_ARG_SIZE) {
+      cmdFields[i] = cmdFields[i].substring(0, 2)
+      truncated = true
+    }
+
+    size += argLen
+  }
+
+  return truncated
+}
+
+class ChildProcessPlugin extends TracingPlugin {
+  static get id () { return 'child_process' }
+  static get prefix () { return 'tracing:datadog:child_process:execution' }
+
+  get tracer () {
+    return this._tracer
+  }
+
+  start ({ command, shell }) {
+    if (typeof command !== 'string') {
+      return
+    }
+
+    const cmdFields = scrubChildProcessCmd(command)
+    const truncated = truncateCommand(cmdFields)
+    const property = (shell === true) ? 'cmd.shell' : 'cmd.exec'
+
+    const meta = {
+      'component': 'subprocess',
+      [property]: (shell === true) ? cmdFields.join(' ') : JSON.stringify(cmdFields)
+    }
+
+    if (truncated) {
+      meta['cmd.truncated'] = `${truncated}`
+    }
+
+    this.startSpan('command_execution', {
+      service: this.config.service,
+      resource: (shell === true) ? 'sh' : cmdFields[0],
+      type: 'system',
+      meta
+    })
+  }
+
+  end ({ result, error }) {
+    let exitCode
+
+    if (result !== undefined) {
+      exitCode = result?.status || 0
+    } else if (error !== undefined) {
+      exitCode = error?.status || error?.code || 0
+    } else {
+      // TracingChannels call start, end synchronously. Later when the promise is resolved then asyncStart asyncEnd.
+      // Therefore in the case of calling end with neither result nor error means that they will come in the asyncEnd.
+      return
+    }
+
+    this.activeSpan?.setTag('cmd.exit_code', `${exitCode}`)
+    this.activeSpan?.finish()
+  }
+
+  error (error) {
+    this.addError(error)
+  }
+
+  asyncEnd ({ result }) {
+    this.activeSpan?.setTag('cmd.exit_code', `${result}`)
+    this.activeSpan?.finish()
+  }
+}
+
+module.exports = ChildProcessPlugin