refactor: Refactor child_process instrumentation to use tracing channel instead of Shim

lib/instrumentation/core/child_process.js

@@ -5,33 +5,112 @@
 
 '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 logger = require('../../logger').child({ component: 'child_process' })
+
+// Create channels based on instrumented methods
+const methods = ['exec', 'execFile']
+const channels = new Map(methods.map((name) => [name, tracingChannel('child_process.' + name)]))
+
+// Patched once per process and never unwrapped; each `initialize()` call
+// just re-points the bound store at the current agent.
+let patched = false
+let currentTracer = null
+let boundStore = null
 
 module.exports = initialize
 
-function initialize(agent, childProcess, moduleName, shim) {
+function initialize(agent, childProcess) {
   if (!childProcess) {
-    shim.logger.debug('Could not find child_process, not instrumenting')
+    logger.debug('Could not find child_process, not instrumenting')
     return false
   }
 
-  const methods = ['exec', 'execFile']
+  patch(childProcess)
+  rebindStore(agent)
+}
 
-  shim.record(childProcess, methods, function recordExec(shim, fn, name) {
-    return new RecorderSpec({ name: 'child_process.' + name, callback: shim.LAST })
-  })
+function patch(childProcess) {
+  if (patched === true) {
+    return
+  }
+  patched = true
 
-  makePromisifyCompatible(shim, childProcess)
-}
+  for (const [methodName, channel] of channels) {
+    const original = childProcess[methodName]
+
+    childProcess[methodName] = function wrappedMethod(...args) {
+      const ctx = currentTracer?.getContext()
+      if (!ctx?.transaction?.isActive()) {
+        return original.apply(this, args)
+      }
+
+      const lastArg = args[args.length - 1]
+      const hasCallback = typeof lastArg === 'function'
+      const data = { methodName, callbackName: hasCallback ? lastArg.name || '<anonymous>' : null }
 
-function makePromisifyCompatible(shim, childProcess) {
-  const originalExec = shim.getOriginal(childProcess.exec)
-  for (const symbol of Object.getOwnPropertySymbols(originalExec)) {
-    childProcess.exec[symbol] = originalExec[symbol]
+      return hasCallback
+        ? channel.traceCallback(original, -1, data, this, ...args)
+        : channel.traceSync(original, data, this, ...args)
+    }
+
+    for (const symbol of Object.getOwnPropertySymbols(original)) {
+      childProcess[methodName][symbol] = original[symbol]
+    }
+  }
+}
+// TODO: This can be repurposed for other core instrumentation
+// refactors over to tracing channel.
+function createSegment(tracer, ctx, name) {
+  const segment = tracer.createSegment({ name, parent: ctx?.segment, transaction: ctx?.transaction })
+  if (segment) {
+    segment.start()
   }
+  return segment
+}
+
+function rebindStore(agent) {
+  const tracer = agent.tracer
+  const store = tracer._contextManager._asyncLocalStorage
+
+  for (const channel of channels.values()) {
+    if (boundStore && boundStore !== store) {
+      channel.start.unbindStore(boundStore)
+      channel.asyncStart.unbindStore(boundStore)
+    }
+
+    channel.start.bindStore(store, (data) => {
+      const ctx = tracer.getContext()
+      const segment = createSegment(tracer, ctx, 'child_process.' + data.methodName)
+      data.ctx = segment ? ctx.enterSegment({ segment }) : ctx
+      return data.ctx
+    })
 
-  const originalExecFile = shim.getOriginal(childProcess.execFile)
-  for (const symbol of Object.getOwnPropertySymbols(originalExecFile)) {
-    childProcess.execFile[symbol] = originalExecFile[symbol]
+    channel.asyncStart.bindStore(store, (data) => {
+      const ctx = data.ctx
+      ctx.segment.touch()
+
+      const callbackSegment = createSegment(tracer, ctx, 'Callback: ' + data.callbackName)
+      data.callbackSegment = callbackSegment
+      return ctx.enterSegment({ segment: callbackSegment })
+    })
   }
+
+  boundStore = store
+  currentTracer = tracer
+}
+
+// `end`/`asyncEnd` run once at module load and don't need
+// agent.tracer (segment is already created within `data`),
+// while `start`/`asyncStart` are handled inside rebindStore().
+for (const channel of channels.values()) {
+  channel.subscribe({
+    end(data) {
+      data.ctx?.segment?.touch()
+    },
+    asyncEnd(data) {
+      data.callbackSegment?.touch()
+    }
+  })
 }

test/unit/instrumentation/core/child_process.test.js

@@ -0,0 +1,33 @@
+/*
+ * Copyright 2020 New Relic Corporation. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+'use strict'
+
+const test = require('node:test')
+const assert = require('node:assert')
+const sinon = require('sinon')
+const helper = require('#testlib/agent_helper.js')
+const logger = require('#agentlib/logger.js')
+const childProcessInstrumentation = require('#agentlib/instrumentation/core/child_process.js')
+
+test('child_process instrumentation', async (t) => {
+  const agent = helper.loadMockedAgent()
+  t.after(() => {
+    helper.unloadAgent(agent)
+  })
+
+  await t.test('should log and return false when child_process is not available', (t) => {
+    const proto = Object.getPrototypeOf(logger.child({}))
+    const debugStub = sinon.stub(proto, 'debug')
+    t.after(() => {
+      debugStub.restore()
+    })
+
+    const result = childProcessInstrumentation(agent, null)
+
+    assert.equal(result, false)
+    assert.ok(debugStub.calledWith('Could not find child_process, not instrumenting'))
+  })
+})