refactor: Refactored dns instrumentation to subscribe to events emitted

lib/instrumentation/core/dns.js

@@ -4,28 +4,50 @@
  */
 
 'use strict'
-
-const { RecorderSpec } = require('../../../lib/shim/specs')
+// eslint-disable-next-line n/no-unsupported-features/node-builtins
+const { tracingChannel } = require('node:diagnostics_channel')
+const { INSTRUMENTED_METHODS, RESOLVE_METHODS } = require('#agentlib/subscribers/dns/constants.js')
+const shimmer = require('#agentlib/shimmer.js')
 
 module.exports = initialize
 
-function initialize(agent, dns, moduleName, shim) {
-  const methods = [
-    'lookup',
-    'resolve',
-    'resolve4',
-    'resolve6',
-    'resolveCname',
-    'resolveMx',
-    'resolveNaptr',
-    'resolveNs',
-    'resolvePtr',
-    'resolveSrv',
-    'resolveTxt',
-    'reverse'
-  ]
+function initialize(_agent, dns) {
+  shimmer.wrapMethod(dns, 'dns', INSTRUMENTED_METHODS, function wrapMethod(original, method) {
+    const channel = tracingChannel(`nr:dns:${method}`)
+    return function wrappedMethod(...args) {
+      const callback = args.at(-1)
+      const callbackName = callback?.name || '<anonymous>'
+      const data = { name: `dns.${method}`, callbackName }
+      return channel.traceCallback(original, -1, data, this, ...args)
+    }
+  })
 
-  shim.record(dns, methods, function recordDnsMethod(shim, fn, name) {
-    return new RecorderSpec({ name: 'dns.' + name, callback: shim.LAST })
+  // All code below can be added separately. it adds the following:
+  //  * instrumentation on resolve methods from `dns.Resolver.prototype`
+  //  * instrumentation on `dns.promises.<method>`
+  //  * instrumentation on resolve methods from `dns.promises.Resolver.prototype`
+  shimmer.wrapMethod(dns.Resolver.prototype, 'dns', RESOLVE_METHODS, function wrapMethod(original, method) {
+    const channel = tracingChannel(`nr:dns:${method}`)
+    return function wrappedMethod(...args) {
+      const callback = args.at(-1)
+      const callbackName = callback?.name || '<anonymous>'
+      const data = { name: `dns.${method}`, callbackName }
+      return channel.traceCallback(original, -1, data, this, ...args)
+    }
+  })
+
+  shimmer.wrapMethod(dns.promises, 'dns', INSTRUMENTED_METHODS, function wrapMethod(original, method) {
+    const channel = tracingChannel(`nr:dns:promises:${method}`)
+    return function wrappedMethod(...args) {
+      const data = { name: `dns.${method}` }
+      return channel.tracePromise(original, data, this, ...args)
+    }
+  })
+  shimmer.wrapMethod(dns.promises.Resolver.prototype, 'dns', RESOLVE_METHODS, function wrapMethod(original, method) {
+    const channel = tracingChannel(`nr:dns:promises:${method}`)
+    return function wrappedMethod(...args) {
+      const data = { name: `dns.${method}` }
+      return channel.tracePromise(original, data, this, ...args)
+    }
   })
 }

lib/subscriber-configs.js

@@ -17,6 +17,7 @@ const subscribers = {
   ...require('./subscribers/bunyan/config'),
   ...require('./subscribers/cassandra-driver/config'),
   ...require('./subscribers/connect/config'),
+  ...require('./subscribers/dns/config'),
   ...require('./subscribers/elasticsearch/config'),
   ...require('./subscribers/express/config'),
   ...require('./subscribers/fastify/config'),

lib/subscribers/dns/config.js

@@ -0,0 +1,16 @@
+/*
+ * Copyright 2026 New Relic Corporation. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+'use strict'
+module.exports = {
+  dns: [
+    {
+      path: './dns/index', instrumentations: []
+    },
+    {
+      path: './dns/promises', instrumentations: []
+    }
+  ]
+}

lib/subscribers/dns/constants.js

@@ -0,0 +1,31 @@
+/*
+ * Copyright 2026 New Relic Corporation. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+'use strict'
+const RESOLVE_METHODS = [
+  'resolve',
+  'resolve4',
+  'resolve6',
+  'resolveCname',
+  'resolveMx',
+  'resolveNaptr',
+  'resolveNs',
+  'resolvePtr',
+  'resolveSrv',
+  'resolveTxt',
+  // missing
+  // resolveAny
+  // resolveCaa
+  // resolveNaptr
+  // resolveSoa
+  // resolveTlsa
+]
+const INSTRUMENTED_METHODS = [
+  'lookup',
+  'reverse',
+  ...RESOLVE_METHODS
+]
+
+module.exports = { INSTRUMENTED_METHODS, RESOLVE_METHODS }

lib/subscribers/dns/index.js

@@ -0,0 +1,101 @@
+/*
+ * Copyright 2026 New Relic Corporation. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+'use strict'
+const { INSTRUMENTED_METHODS } = require('./constants')
+// eslint-disable-next-line n/no-unsupported-features/node-builtins
+const { tracingChannel } = require('node:diagnostics_channel')
+
+class DnsSubscriber {
+  constructor({ agent, logger }) {
+    this.agent = agent
+    this.config = agent.config
+    this.logger = logger.child({ component: 'dns-subscriber' })
+    this.prefix = 'nr:dns'
+    this.channels = this.buildChannels()
+    this.store = agent.tracer._contextManager._asyncLocalStorage
+    this.boundEnd = this.#end.bind(this)
+    this.boundAsyncEnd = this.#asyncEnd.bind(this)
+  }
+
+  buildChannels() {
+    return INSTRUMENTED_METHODS.map((method) => tracingChannel(`${this.prefix}:${method}`))
+  }
+
+  get enabled() {
+    return this.config.instrumentation.dns.enabled === true
+  }
+
+  createSegment({ name, ctx }) {
+    const parent = ctx.segment
+
+    const segment = this.agent.tracer.createSegment({
+      name,
+      parent,
+      transaction: ctx.transaction
+    })
+
+    if (segment) {
+      segment.start()
+      return ctx.enterSegment({ segment })
+    }
+    return ctx
+  }
+
+  enable() {
+    for (const channel of this.channels) {
+      channel.start.bindStore(this.store, (data) => {
+        const ctx = this.agent.tracer.getContext()
+        if (!ctx?.transaction?.isActive()) {
+          return ctx
+        }
+
+        return this.createSegment({ name: data.name, ctx })
+      })
+
+      channel.asyncStart.bindStore(this.store, (data) => {
+        const ctx = this.agent.tracer.getContext()
+        if (!ctx?.transaction?.isActive()) {
+          return ctx
+        }
+
+        const name = `Callback: ${data.callbackName}`
+        const newCtx = this.createSegment({ name, ctx })
+        return newCtx
+      })
+    }
+  }
+
+  #end() {
+    const ctx = this.agent.tracer.getContext()
+    ctx?.segment?.touch()
+  }
+
+  #asyncEnd(data) {
+    const ctx = this.agent.tracer.getContext()
+    ctx?.segment?.touch()
+  }
+
+  subscribe() {
+    for (const channel of this.channels) {
+      channel.subscribe({ end: this.boundEnd, asyncEnd: this.boundAsyncEnd })
+    }
+  }
+
+  unsubscribe() {
+    for (const channel of this.channels) {
+      channel.unsubscribe({ end: this.boundEnd, asyncEnd: this.boundAsyncEnd })
+    }
+  }
+
+  disable() {
+    for (const channel of this.channels) {
+      channel.start.unbindStore(this.store)
+      channel.asyncStart.unbindStore(this.store)
+    }
+  }
+}
+
+module.exports = DnsSubscriber

lib/subscribers/dns/promises.js

@@ -0,0 +1,83 @@
+/*
+ * Copyright 2026 New Relic Corporation. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+'use strict'
+const { INSTRUMENTED_METHODS } = require('./constants')
+// eslint-disable-next-line n/no-unsupported-features/node-builtins
+const { tracingChannel } = require('node:diagnostics_channel')
+
+class DnsPromisesSubscriber {
+  constructor({ agent, logger }) {
+    this.agent = agent
+    this.config = agent.config
+    this.logger = logger.child({ component: 'dns-promies-subscriber' })
+    this.prefix = 'nr:dns:promises'
+    this.channels = this.buildChannels()
+    this.store = agent.tracer._contextManager._asyncLocalStorage
+    this.boundAsyncEnd = this.#asyncEnd.bind(this)
+  }
+
+  buildChannels() {
+    return INSTRUMENTED_METHODS.map((method) => tracingChannel(`${this.prefix}:${method}`))
+  }
+
+  get enabled() {
+    return this.config.instrumentation.dns.enabled === true
+  }
+
+  createSegment({ name, ctx }) {
+    const parent = ctx.segment
+
+    const segment = this.agent.tracer.createSegment({
+      name,
+      parent,
+      transaction: ctx.transaction
+    })
+
+    if (segment) {
+      segment.start()
+      return ctx.enterSegment({ segment })
+    }
+    return ctx
+  }
+
+  enable() {
+    for (const channel of this.channels) {
+      channel.start.bindStore(this.store, (data) => {
+        const ctx = this.agent.tracer.getContext()
+        if (!ctx?.transaction?.isActive()) {
+          return ctx
+        }
+
+        return this.createSegment({ name: data.name, ctx })
+      })
+    }
+  }
+
+  #asyncEnd(data) {
+    const ctx = this.agent.tracer.getContext()
+    ctx?.segment?.touch()
+  }
+
+  subscribe() {
+    for (const channel of this.channels) {
+      channel.subscribe({ asyncEnd: this.boundAsyncEnd })
+    }
+  }
+
+  unsubscribe() {
+    for (const channel of this.channels) {
+      channel.unsubscribe({ asyncEnd: this.boundAsyncEnd })
+    }
+  }
+
+  disable() {
+    for (const channel of this.channels) {
+      channel.start.unbindStore(this.store)
+    }
+  }
+}
+
+module.exports = DnsPromisesSubscriber

test/integration/core/dns-utils.js

@@ -0,0 +1,80 @@
+/*
+ * Copyright 2026 New Relic Corporation. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+'use strict'
+
+/**
+ * Mock most methods so we can control the results.
+ * Wrap calling the callback in a `setImmediate` so it will
+ * properly emit both `end` and `asyncEnd` within the tracing channel
+ * @param {object} params to function
+ * @param {object} params.dns dns package
+ * @param {object} params.sandbox sinon sandbox
+ */
+module.exports = function mockDns({ dns, sandbox }) {
+  sandbox.stub(dns, 'reverse').callsFake((_addr, cb) => {
+    setImmediate(() => {
+      cb(null, ['localhost'])
+    })
+  })
+  sandbox.stub(dns.promises, 'reverse').callsFake(async() => Promise.resolve(['localhost']))
+
+  sandbox.stub(dns, 'resolve').callsFake((_, cb) => {
+    setImmediate(() => {
+      cb(null, ['127.0.0.1'])
+    })
+  })
+  sandbox.stub(dns.Resolver.prototype, 'resolve').callsFake((_, cb) => {
+    setImmediate(() => {
+      cb(null, ['127.0.0.1'])
+    })
+  })
+  sandbox.stub(dns.promises, 'resolve').callsFake(async () => Promise.resolve(['127.0.0.1']))
+  sandbox.stub(dns.promises.Resolver.prototype, 'resolve').callsFake(async () => Promise.resolve(['127.0.0.1']))
+  sandbox.stub(dns, 'resolve4').callsFake((_, cb) => {
+    setImmediate(() => {
+      cb(null, ['127.0.0.1'])
+    })
+  })
+  sandbox.stub(dns.promises, 'resolve4').callsFake(async () => Promise.resolve(['127.0.0.1']))
+  sandbox.stub(dns, 'resolve6').callsFake((_, cb) => {
+    setImmediate(() => {
+      cb(null, ['::1'])
+    })
+  })
+  sandbox.stub(dns.promises, 'resolve6').callsFake(async () => Promise.resolve(['::1']))
+  const error = Error('boom')
+  error.code = 'ENODATA'
+  sandbox.stub(dns, 'resolveCname').callsFake((_, cb) => {
+    setImmediate(() => {
+      cb(error)
+    })
+  })
+  sandbox.stub(dns.promises, 'resolveCname').callsFake(async () => Promise.reject(error))
+  sandbox.stub(dns, 'resolveMx').callsFake((_, cb) => {
+    setImmediate(() => {
+      cb(null, ['127.0.0.1'])
+    })
+  })
+  sandbox.stub(dns.promises, 'resolveMx').callsFake(async () => Promise.resolve(['127.0.0.1']))
+  sandbox.stub(dns, 'resolveNs').callsFake((_, cb) => {
+    setImmediate(() => {
+      cb(null, ['a.iana-servers.net', 'b.iana-servers.net'])
+    })
+  })
+  sandbox.stub(dns.promises, 'resolveNs').callsFake(async () => Promise.resolve(['a.iana-servers.net', 'b.iana-servers.net']))
+  sandbox.stub(dns, 'resolveTxt').callsFake((_, cb) => {
+    setImmediate(() => {
+      cb(null, ['one', 'two', 'three'])
+    })
+  })
+  sandbox.stub(dns.promises, 'resolveTxt').callsFake(async () => Promise.resolve(['one', 'two', 'three']))
+  sandbox.stub(dns, 'resolveSrv').callsFake((_, cb) => {
+    setImmediate(() => {
+      cb(error)
+    })
+  })
+  sandbox.stub(dns.promises, 'resolveSrv').callsFake(async () => Promise.reject(error))
+}