dnstrace.js

import dns from 'dns'
import AWSXRay from 'aws-xray-sdk'

function patchDnsLookup() {
    const realLookup = dns.lookup
    dns.lookup = (hostname, ...args) => {
        if (hostname === process.env.AWS_XRAY_DAEMON_ADDRESS) {
            return realLookup(hostname, ...args)
        }
        AWSXRay.captureAsyncFunc('DNS Lookup', (subsegment) => {
            if (!subsegment) {
                // if we aren't in a segment just behave like normal dns.lookup
                return realLookup(hostname, ...args)
            }

            // rewrite callback
            const callback = args[args.length - 1]
            args[args.length - 1] = (err, addr, ...callbackArgs) => {
                if (err) {
                    subsegment.addError(err)
                    subsegment.addErrorFlag()
                }
                subsegment.addMetadata('result', addr)
                // usually args would be [addrType], but since it is
                // undocumented, it is safer to treat them as blackbox
                subsegment.addMetadata('additionalResult', callbackArgs)

                subsegment.close()
                callback(err, addr, ...callbackArgs)
            }

            subsegment.addAnnotation('hostname', hostname)
            realLookup(hostname, ...args)
        })
    }
}