'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 }