dd-trace/packages/datadog-plugin-kafkajs/src/producer.js

Datadog APM tracing client for JavaScript

'use strict'

const ProducerPlugin = require('../../dd-trace/src/plugins/producer')
const { DsmPathwayCodec, getMessageSize } = require('../../dd-trace/src/datastreams')

const BOOTSTRAP_SERVERS_KEY = 'messaging.kafka.bootstrap.servers'
const MESSAGING_DESTINATION_KEY = 'messaging.destination.name'

class KafkajsProducerPlugin extends ProducerPlugin {
  static id = 'kafkajs'
  static operation = 'produce'
  static peerServicePrecursors = [BOOTSTRAP_SERVERS_KEY]

  constructor () {
    super(...arguments)
    this.addSub(`apm:${this.constructor.id}:produce:commit`, message => this.commit(message))
  }

  /**
   * Transform individual commit details sent by kafkajs' event reporter
   * into actionable backlog items for DSM
   *
   * @typedef {object} ProducerBacklog
   * @property {number} type
   * @property {string} topic
   * @property {number} partition
   * @property {number} offset
   */
  /**
   *
   * @typedef {object} ProducerResponseItem
   * @property {string} topic
   * @property {number} partition
   * @property {import('kafkajs/utils/long').Long} [offset]
   * @property {import('kafkajs/utils/long').Long} [baseOffset]
   *
   * @param {ProducerResponseItem} response
   * @returns {ProducerBacklog}
   */
  transformProduceResponse (response, clusterId) {
    // In produce protocol >=v3, the offset key changes from `offset` to `baseOffset`
    const { topicName: topic, partition, offset, baseOffset } = response
    const offsetAsLong = offset || baseOffset
    const backlog = {
      type: 'kafka_produce',
      partition,
      offset: offsetAsLong ? Number(offsetAsLong) : undefined,
      topic,
    }
    if (clusterId) {
      backlog.kafka_cluster_id = clusterId
    }
    return backlog
  }

  /**
   *
   * @param {{ result: ProducerResponseItem[] }} ctx
   * @returns {void}
   */
  commit (ctx) {
    const commitList = ctx.result
    const clusterId = ctx.clusterId

    if (!this.config.dsmEnabled) return
    if (!commitList || !Array.isArray(commitList)) return
    for (const rawCommit of commitList) {
      const commit = this.transformProduceResponse(rawCommit, clusterId)
      this.tracer.setOffset(commit)
    }
  }

  start (ctx) {
    if (!this.config.dsmEnabled) return
    const { topic, messages, clusterId, disableHeaderInjection, currentStore: { span } } = ctx

    for (const message of messages) {
      if (message !== null && typeof message === 'object') {
        const payloadSize = getMessageSize(message)
        const edgeTags = ['direction:out', `topic:${topic}`, 'type:kafka']

        if (clusterId) {
          edgeTags.push(`kafka_cluster_id:${clusterId}`)
        }

        const dataStreamsContext = this.tracer.setCheckpoint(edgeTags, span, payloadSize)
        if (!disableHeaderInjection) {
          DsmPathwayCodec.encode(dataStreamsContext, message.headers)
        }
      }
    }
  }

  finish (ctx) {
    const span = ctx?.currentStore?.span
    const result = ctx?.result
    if (span && Array.isArray(result) && result.length > 0) {
      // The broker response is one entry per (topic, partition). Each entry
      // carries a `baseOffset` — the offset assigned to the first record sent
      // to that partition. We don't know per-partition record counts from the
      // response, only the starting offset.
      const offsets = []
      for (const entry of result) {
        const offsetAsLong = entry.offset ?? entry.baseOffset
        if (entry.partition === undefined || offsetAsLong === undefined) continue
        // Kafka offsets are 64-bit; coercing to Number loses precision past
        // 2^53. Keep them as strings so the tag matches the exact offset on
        // long-lived/high-throughput topics.
        offsets.push({ partition: entry.partition, start_offset: String(offsetAsLong) })
      }
      if (offsets.length > 0) {
        offsets.sort((a, b) => a.partition - b.partition)
        span.setTag('kafka.messages.offsets', JSON.stringify(offsets))
      }
      // Single-message send: the one entry's partition/offset describes the
      // exact record. Also expose them as flat tags for easy filtering.
      if (offsets.length === 1 && ctx.messages?.length === 1) {
        span.setTag('kafka.partition', offsets[0].partition)
        // Set as a string meta tag (not a metric) to preserve full 64-bit precision.
        span.setTag('kafka.message.offset', offsets[0].start_offset)
      }
    }
    super.finish(ctx)
  }

  bindStart (ctx) {
    const { topic, messages, bootstrapServers, clusterId, disableHeaderInjection } = ctx
    const span = this.startSpan({
      resource: topic,
      meta: {
        component: this.constructor.id,
        'kafka.topic': topic,
        'kafka.cluster_id': clusterId,
        [MESSAGING_DESTINATION_KEY]: topic,
      },
      metrics: {
        'kafka.batch_size': messages.length,
      },
    }, ctx)
    if (bootstrapServers) {
      span.setTag(BOOTSTRAP_SERVERS_KEY, bootstrapServers)
    }
    for (const message of messages) {
      // message headers are not supported for kafka broker versions <0.11
      if (message !== null && typeof message === 'object' && !disableHeaderInjection) {
        message.headers ??= {}
        this.tracer.inject(span, 'text_map', message.headers)
      }
    }

    return ctx.currentStore
  }
}

module.exports = KafkajsProducerPlugin