dd-trace/packages/dd-trace/src/profiling/profiler.js

Datadog APM tracing client for JavaScript

'use strict'

const { EventEmitter } = require('events')
const dc = require('dc-polyfill')
const crashtracker = require('../crashtracking')
const log = require('../log')
const { Config } = require('./config')
const { snapshotKinds } = require('./constants')
const { threadNamePrefix } = require('./profilers/shared')
const { isWebServerSpan, endpointNameFromTags, getStartedSpans } = require('./webspan-utils')

const profileSubmittedChannel = dc.channel('datadog:profiling:profile-submitted')
const spanFinishedChannel = dc.channel('dd-trace:span:finish')

function findWebSpan (startedSpans, spanId) {
  for (let i = startedSpans.length; --i >= 0;) {
    const ispan = startedSpans[i]
    const context = ispan.context()
    if (context._spanId === spanId) {
      if (isWebServerSpan(context._tags)) {
        return true
      }
      spanId = context._parentId
    }
  }
  return false
}

const MISSING_SOURCE_MAPS_TOKEN = 'dd:has-missing-map-files'

function profileHasMissingSourceMaps (profile) {
  const strings = profile?.stringTable?.strings
  if (!strings) return false
  return profile.comment?.some(idx => strings[idx] === MISSING_SOURCE_MAPS_TOKEN) ?? false
}

function processInfo (infos, info, type) {
  if (Object.keys(info).length > 0) {
    infos[type] = info
  }
}

class Profiler extends EventEmitter {
  #compressionFn
  #compressionFnInitialized = false
  #compressionOptions
  #config
  #customLabelKeys = new Set()
  #enabled = false
  #endpointCounts = new Map()
  #lastStart
  #logger
  #profileSeq = 0
  #spanFinishListener
  #timer

  constructor () {
    super()
    this._timeoutInterval = undefined
  }

  get serverless () { return false }

  get flushInterval () {
    return this.#config?.flushInterval
  }

  get enabled () {
    return this.#enabled
  }

  /**
   * Declares the set of custom label keys that will be used with
   * {@link runWithLabels}. This is used for profile upload metadata and
   * for pprof serialization optimization (low-cardinality deduplication).
   *
   * @param {Iterable<string>} keys - Custom label key names
   */
  setCustomLabelKeys (keys) {
    this.#customLabelKeys.clear()
    for (const key of keys) {
      this.#customLabelKeys.add(key)
    }
    if (this.#config) {
      for (const profiler of this.#config.profilers) {
        profiler.setCustomLabelKeys?.(this.#customLabelKeys)
      }
    }
  }

  /**
   * Runs a function with custom profiling labels attached to wall profiler samples.
   *
   * @param {Record<string, string | number>} labels - Custom labels to attach
   * @param {function(): T} fn - Function to execute with the labels
   * @returns {T} The return value of fn
   * @template T
   */
  runWithLabels (labels, fn) {
    if (!this.#enabled || !this.#config) {
      return fn()
    }
    for (const profiler of this.#config.profilers) {
      if (profiler.runWithLabels) {
        return profiler.runWithLabels(labels, fn)
      }
    }
    return fn()
  }

  #getCompressionFn () {
    if (!this.#compressionFnInitialized) {
      this.#compressionFnInitialized = true
      try {
        const { promisify } = require('util')
        const zlib = require('zlib')
        const { method, level: clevel } = this.#config.uploadCompression
        switch (method) {
          case 'gzip':
            this.#compressionFn = promisify(zlib.gzip)
            if (clevel !== undefined) {
              this.#compressionOptions = {
                level: clevel,
              }
            }
            break
          case 'zstd':
            // eslint-disable-next-line n/no-unsupported-features/node-builtins
            if (typeof zlib.zstdCompress === 'function') {
              // eslint-disable-next-line n/no-unsupported-features/node-builtins
              this.#compressionFn = promisify(zlib.zstdCompress)
              if (clevel !== undefined) {
                this.#compressionOptions = {
                  params: {
                    // eslint-disable-next-line n/no-unsupported-features/node-builtins
                    [zlib.constants.ZSTD_c_compressionLevel]: clevel,
                  },
                }
              }
            } else {
              const zstdCompress = require('@datadog/libdatadog').load('datadog-js-zstd').zstd_compress
              const level = clevel ?? 0 // 0 is zstd default compression level
              this.#compressionFn = (buffer) => Promise.resolve(Buffer.from(zstdCompress(buffer, level)))
            }
            break
        }
      } catch (error) {
        log.error(error)
      }
    }
    return this.#compressionFn
  }

  /**
   * @param {import('../config/config-base')} options - Tracer configuration
   */
  start (options) {
    if (this.enabled) return true
    this.#enabled = true

    const config = this.#config = new Config(options)
    this.#logger = config.logger

    this._setInterval()
    // Log errors if the source map finder fails, but don't prevent the rest
    // of the profiler from running without source maps.
    let mapper
    const { setLogger, SourceMapper } = require('@datadog/pprof')
    setLogger(config.logger)

    if (config.sourceMap) {
      mapper = new SourceMapper(config.debugSourceMaps)
      mapper.loadDirectory(process.cwd())
        .then(() => {
          if (config.debugSourceMaps) {
            const count = mapper.infoMap.size
            this.#logger.debug(() => {
              return count === 0
                ? 'Found no source maps'
                : `Found source maps for following files: [${[...mapper.infoMap.keys()].join(', ')}]`
            })
          }
        })
        .catch((error) => {
          log.error(error)
        })
    }

    try {
      const start = new Date()
      const nearOOMCallback = this.#nearOOMExport.bind(this)
      for (const profiler of config.profilers) {
        // TODO: move this out of Profiler when restoring sourcemap support
        profiler.start({
          mapper,
          nearOOMCallback,
        })
        this.#logger.debug(`Started ${profiler.type} profiler in ${threadNamePrefix} thread`)
      }

      if (config.endpointCollectionEnabled) {
        this.#spanFinishListener = this.#onSpanFinish.bind(this)
        spanFinishedChannel.subscribe(this.#spanFinishListener)
      }

      this._capture(this._timeoutInterval, start)
    } catch (error) {
      log.error(error)
      this.#stop()
      return false
    }

    return true
  }

  #nearOOMExport (profileType, encodedProfile, info) {
    const start = this.#lastStart
    const end = new Date()
    const infos = this.#createInitialInfos()
    processInfo(infos, info, profileType)
    this.#submit({
      [profileType]: encodedProfile,
    }, infos, start, end, snapshotKinds.ON_OUT_OF_MEMORY)
  }

  _setInterval () {
    this._timeoutInterval = this.#config.flushInterval
  }

  stop () {
    if (!this.enabled) return

    // collect and export current profiles
    // once collect returns, profilers can be safely stopped
    this._collect(snapshotKinds.ON_SHUTDOWN, false)
    this.#stop()
  }

  #stop () {
    if (!this.enabled) return

    this.#enabled = false

    if (this.#spanFinishListener !== undefined) {
      spanFinishedChannel.unsubscribe(this.#spanFinishListener)
      this.#spanFinishListener = undefined
    }

    for (const profiler of this.#config.profilers) {
      profiler.stop()
      this.#logger.debug(`Stopped ${profiler.type} profiler in ${threadNamePrefix} thread`)
    }

    clearTimeout(this.#timer)
    this.#timer = undefined
  }

  _capture (timeout, start) {
    if (!this.enabled) return
    this.#lastStart = start
    if (!this.#timer || timeout !== this._timeoutInterval) {
      this.#timer = setTimeout(() => this._collect(snapshotKinds.PERIODIC), timeout)
      this.#timer.unref()
    } else {
      this.#timer.refresh()
    }
  }

  #onSpanFinish (span) {
    const context = span.context()
    const tags = context._tags
    if (!isWebServerSpan(tags)) return

    const endpointName = endpointNameFromTags(tags)
    if (!endpointName) return

    // Make sure this is the outermost web span, just in case so we don't overcount
    if (findWebSpan(getStartedSpans(context), context._parentId)) return

    let counter = this.#endpointCounts.get(endpointName)
    if (counter === undefined) {
      counter = { count: 1 }
      this.#endpointCounts.set(endpointName, counter)
    } else {
      counter.count++
    }
  }

  #createInitialInfos () {
    return {
      serverless: this.serverless,
      settings: this.#config.systemInfoReport,
      hasMissingSourceMaps: false,
    }
  }

  async _collect (snapshotKind, restart = true) {
    if (!this.enabled) return

    try {
      if (this.#config.profilers.length === 0) {
        throw new Error('No profile types configured.')
      }

      const startDate = this.#lastStart
      const endDate = new Date()
      const profiles = []

      crashtracker.withProfilerSerializing(() => {
        // collect profiles synchronously so that profilers can be safely stopped asynchronously
        for (const profiler of this.#config.profilers) {
          const info = profiler.getInfo()
          const profile = profiler.profile(restart, startDate, endDate)
          if (!restart) {
            this.#logger.debug(`Stopped ${profiler.type} profiler in ${threadNamePrefix} thread`)
          }
          if (!profile) continue
          profiles.push({ profiler, profile, info })
        }
      })

      if (restart) {
        this._capture(this._timeoutInterval, endDate)
      }

      let hasEncoded = false

      const encodedProfiles = {}
      const infos = this.#createInitialInfos()
      const compressionFn = this.#getCompressionFn()

      // encode and export asynchronously
      await Promise.all(profiles.map(async ({ profiler, profile, info }) => {
        try {
          const encoded = await profiler.encode(profile)
          const compressed = encoded instanceof Buffer && compressionFn !== undefined
            ? await compressionFn(encoded, this.#compressionOptions)
            : encoded
          encodedProfiles[profiler.type] = compressed
          if (profileHasMissingSourceMaps(profile)) {
            infos.hasMissingSourceMaps = true
          }
          processInfo(infos, info, profiler.type)
          this.#logger.debug(() => {
            const profileJson = JSON.stringify(profile, (_, value) => {
              return typeof value === 'bigint' ? value.toString() : value
            })
            return `Collected ${profiler.type} profile: ` + profileJson
          })
          hasEncoded = true
        } catch (error) {
          // If encoding one of the profile types fails, we should still try to
          // encode and submit the other profile types.
          log.error(error)
        }
      }))

      if (hasEncoded) {
        await this.#submit(encodedProfiles, infos, startDate, endDate, snapshotKind)
        profileSubmittedChannel.publish()
        this.#logger.debug('Submitted profiles')
      }
    } catch (error) {
      log.error(error)
      this.#stop()
    }
  }

  #submit (profiles, infos, start, end, snapshotKind) {
    const { tags } = this.#config

    // Flatten endpoint counts
    const endpointCounts = {}
    for (const [endpoint, { count }] of this.#endpointCounts) {
      endpointCounts[endpoint] = count
    }
    this.#endpointCounts.clear()

    tags.snapshot = snapshotKind
    tags.profile_seq = this.#profileSeq++
    const customAttributes = this.#customLabelKeys.size > 0
      ? [...this.#customLabelKeys]
      : undefined
    const exportSpec = { profiles, infos, start, end, tags, endpointCounts, customAttributes }
    const tasks = this.#config.exporters.map(exporter =>
      exporter.export(exportSpec).catch(error => {
        log.warn(error)
      })
    )

    return Promise.all(tasks)
  }
}

class ServerlessProfiler extends Profiler {
  #profiledIntervals = 0
  #interval = 1 // seconds
  #flushAfterIntervals

  constructor () {
    super()
    this.#profiledIntervals = 0
    this.#interval = 1
    this.#flushAfterIntervals = undefined
  }

  get serverless () { return true }

  get profiledIntervals () {
    return this.#profiledIntervals
  }

  _setInterval () {
    this._timeoutInterval = this.#interval * 1000
    this.#flushAfterIntervals = this.flushInterval / 1000
  }

  async _collect (snapshotKind, restart = true) {
    if (this.#profiledIntervals >= this.#flushAfterIntervals || !restart) {
      this.#profiledIntervals = 0
      await super._collect(snapshotKind, restart)
    } else {
      this.#profiledIntervals += 1
      this._capture(this._timeoutInterval, new Date())
      // Don't submit profile until 65 (flushAfterIntervals) intervals have elapsed
    }
  }
}

module.exports = { Profiler, ServerlessProfiler }