dd-trace 6.9.0 → 6.10.0

package/packages/datadog-plugin-kafkajs/src/consumer.js

@@ -1,9 +1,8 @@
 'use strict'
 
 const dc = require('dc-polyfill')
-const { getMessageSize } = require('../../dd-trace/src/datastreams')
 const ConsumerPlugin = require('../../dd-trace/src/plugins/consumer')
-const { convertToTextMap } = require('./utils')
+const { convertToTextMap, getKafkaMessageSize } = require('./utils')
 const afterStartCh = dc.channel('dd-trace:kafkajs:consumer:afterStart')
 const beforeFinishCh = dc.channel('dd-trace:kafkajs:consumer:beforeFinish')
 
@@ -69,7 +68,7 @@ class KafkajsConsumerPlugin extends ConsumerPlugin {
     if (!headers) return
 
     const { span } = ctx.currentStore
-    const payloadSize = getMessageSize(message)
+    const payloadSize = getKafkaMessageSize(message)
     this.tracer.decodeDataStreamsContext(headers)
     const edgeTags = ['direction:in', `group:${groupId}`, `topic:${topic}`, 'type:kafka']
     if (clusterId) {

package/packages/datadog-plugin-kafkajs/src/producer.js

@@ -1,7 +1,8 @@
 'use strict'
 
 const ProducerPlugin = require('../../dd-trace/src/plugins/producer')
-const { DsmPathwayCodec, getMessageSize } = require('../../dd-trace/src/datastreams')
+const { DsmPathwayCodec } = require('../../dd-trace/src/datastreams')
+const { getKafkaMessageSize } = require('./utils')
 
 const BOOTSTRAP_SERVERS_KEY = 'messaging.kafka.bootstrap.servers'
 const MESSAGING_DESTINATION_KEY = 'messaging.destination.name'
@@ -77,7 +78,7 @@ class KafkajsProducerPlugin extends ProducerPlugin {
 
     for (const message of messages) {
       if (message !== null && typeof message === 'object') {
-        const payloadSize = getMessageSize(message)
+        const payloadSize = getKafkaMessageSize(message)
         const edgeTags = ['direction:out', `topic:${topic}`, 'type:kafka']
 
         if (clusterId) {

package/packages/datadog-plugin-kafkajs/src/utils.js

@@ -1,27 +1,85 @@
 'use strict'
 
+const { getSizeOrZero } = require('../../dd-trace/src/datastreams')
+
+/** @typedef {string | Buffer | null | undefined} KafkaHeaderValue */
+/** @typedef {Record<string, KafkaHeaderValue | KafkaHeaderValue[]>} KafkaHeaderMap */
+/** @typedef {Array<Record<string, KafkaHeaderValue>>} NativeHeaderList */
+
+/** @typedef {Record<string, string | string[]>} TextMap */
+
+/**
+ * @param {TextMap} textMap
+ * @param {string} key
+ * @param {KafkaHeaderValue} value
+ */
+function addHeaderValue (textMap, key, value) {
+  if (value === null || value === undefined) return
+
+  const existing = textMap[key]
+  if (existing === undefined) {
+    textMap[key] = value.toString()
+  } else if (typeof existing === 'string') {
+    textMap[key] = [existing, value.toString()]
+  } else {
+    existing.push(value.toString())
+  }
+}
+
+/** @param {KafkaHeaderMap | NativeHeaderList | null | undefined} bufferMap */
 function convertToTextMap (bufferMap) {
   if (!bufferMap) return null
 
-  // rdKafka returns an array of header maps
+  // A `toString` or `__proto__` wire key would read an inherited value back as a repeat.
+  /** @type {TextMap} */
+  const textMap = Object.create(null)
+
+  // rdKafka reports one single-key record per wire field.
   if (Array.isArray(bufferMap)) {
-    const headers = {}
     for (const headerMap of bufferMap) {
       for (const key of Object.keys(headerMap)) {
-        headers[key] = headerMap[key].toString()
+        addHeaderValue(textMap, key, headerMap[key])
       }
     }
-    return headers
+    return textMap
   }
 
-  const textMap = {}
   for (const key of Object.keys(bufferMap)) {
-    if (bufferMap[key] === null || bufferMap[key] === undefined) continue
-    textMap[key] = bufferMap[key].toString()
+    const value = bufferMap[key]
+    if (Array.isArray(value)) {
+      for (const repeated of value) {
+        addHeaderValue(textMap, key, repeated)
+      }
+    } else {
+      addHeaderValue(textMap, key, value)
+    }
   }
   return textMap
 }
 
+/**
+ * A KafkaJS header value array becomes one wire record per element, so its key
+ * counts once per value. A native header array already holds one record per
+ * entry, where the array indices are not wire bytes.
+ *
+ * @param {{ key?: unknown, value?: unknown,
+ *   headers?: KafkaHeaderMap | NativeHeaderList | null }} message
+ */
+function getKafkaMessageSize (message) {
+  const { key, value, headers } = message
+  let size = getSizeOrZero(key) + getSizeOrZero(value)
+  if (headers === undefined || headers === null) return size
+  if (Array.isArray(headers)) return size + getSizeOrZero(headers)
+
+  for (const headerKey of Object.keys(headers)) {
+    const headerValue = headers[headerKey]
+    const keySize = Buffer.byteLength(headerKey, 'utf8')
+    size += (Array.isArray(headerValue) ? keySize * headerValue.length : keySize) + getSizeOrZero(headerValue)
+  }
+  return size
+}
+
 module.exports = {
   convertToTextMap,
+  getKafkaMessageSize,
 }