basic implementation

lib/shimmer.js

@@ -14,6 +14,7 @@ const shims = require('./shim')
 const { Hook } = require('@newrelic/ritm')
 const IitmHook = require('import-in-the-middle')
 const { nrEsmProxy } = require('./symbols')
+const moduleIds = new (require('./util/idgen'))()
 let pkgsToHook = []
 
 const NAMES = require('./metrics/names')
@@ -412,9 +413,11 @@ const shimmer = (module.exports = {
       }
     }
 
-    opts[symbols.instrumented] = new Set()
-    opts[symbols.instrumentedErrored] = new Set()
-    shimmer.registeredInstrumentations[opts.moduleName].push({ ...opts })
+    const instrumentation = {
+      ...opts,
+      instrumentationId: moduleIds.idFor(opts.moduleName)
+    }
+    shimmer.registeredInstrumentations[opts.moduleName].push(instrumentation)
   },
 
   registeredInstrumentations: Object.create(null),
@@ -467,23 +470,20 @@ const shimmer = (module.exports = {
    * only if every hook succeeded.
    *
    * @param {string} moduleName name of registered instrumentation
+   * @param {object} nodule the module that is being instrumented
    * @returns {boolean} if all instrumentation hooks ran for a given version
    */
-  isInstrumented(moduleName) {
-    const pkgVersion = shimmer.getPackageVersion(moduleName)
-    const instrumentations = shimmer.registeredInstrumentations[moduleName]
-    const didInstrument = instrumentations.filter((instrumentation) =>
-      instrumentation[symbols.instrumented].has(pkgVersion)
-    )
-    return didInstrument.length === instrumentations.length
+  isInstrumented(moduleName, nodule) {
+    const instrumentations = shimmer.registeredInstrumentations[moduleName] ?? []
+    return nodule?.[symbols.instrumented]?.length === instrumentations.length
   },
 
   instrumentPostLoad(agent, module, moduleName, resolvedName, returnModule = false) {
     const result = _postLoad(agent, module, moduleName, resolvedName)
     // This is to not break the public API
     // previously it would just call instrumentation
     // and not check the result
-    return returnModule ? result : shimmer.isInstrumented(moduleName)
+    return returnModule ? result : shimmer.isInstrumented(moduleName, module)
   },
 
   /**
@@ -512,6 +512,8 @@ function applyDebugState(shim, nodule, inEsm) {
       instrumented.push({
         [symbols.unwrap]: function unwrapNodule() {
           delete nodule[symbols.shim]
+          delete nodule[symbols.instrumented]
+          delete nodule[symbols.instrumentedErrored]
         }
       })
       nodule[symbols.shim] = shim
@@ -536,9 +538,11 @@ function instrumentPostLoad(agent, nodule, moduleName, resolvedName, esmResolver
   const pkgVersion = resolvedName ? shimmer.getPackageVersion(moduleName) : process.version
   const instrumentations = shimmer.registeredInstrumentations[moduleName]
   instrumentations.forEach((instrumentation) => {
-    const isInstrumented = instrumentation[symbols.instrumented].has(pkgVersion)
-    const failedInstrumentation = instrumentation[symbols.instrumentedErrored].has(pkgVersion)
-    if (isInstrumented || failedInstrumentation) {
+    const isInstrumented =
+      nodule[symbols.instrumented]?.includes(instrumentation.instrumentationId) === true
+    const failedInstrumentation =
+      nodule[symbols.instrumentedErrored]?.includes(instrumentation.instrumentationId) === true
+    if (isInstrumented === true || failedInstrumentation === true) {
       const msg = isInstrumented ? 'Already instrumented' : 'Failed to instrument'
       logger.trace(`${msg} ${moduleName}@${pkgVersion}, skipping registering instrumentation`)
       return
@@ -566,7 +570,7 @@ function instrumentPostLoad(agent, nodule, moduleName, resolvedName, esmResolver
     // that occur directly above this.  No reason to attempt to load instrumentation
     // as it does not exist.
     if (instrumentation.type === MODULE_TYPE.TRACKING) {
-      instrumentation[symbols.instrumented].add(pkgVersion)
+      registerHook(nodule, instrumentation)
       return
     }
 
@@ -660,11 +664,12 @@ function loadInstrumentation({
 }) {
   try {
     if (instrumentation.onRequire(shim, resolvedNodule, moduleName) !== false) {
+      instrumentation.instrumentationId = moduleIds.idFor(moduleName)
       nodule = shim.getExport(nodule)
-      instrumentation[symbols.instrumented].add(pkgVersion)
+      registerHook(nodule, instrumentation)
     }
   } catch (instrumentationError) {
-    instrumentation[symbols.instrumentedErrored].add(pkgVersion)
+    registerHook(nodule, instrumentation, symbols.instrumentedErrored)
     if (instrumentation.onError) {
       try {
         instrumentation.onError(instrumentationError)
@@ -786,3 +791,29 @@ function tryGetVersion(shim) {
 
   return shim.pkgVersion
 }
+
+/**
+ * When a module is being instrumented, this function is used to track all of
+ * the hooks being registered for that module. As an example, the core agent
+ * may register a set of `onRequire` and `onError` hooks and then the security
+ * agent may also register its own set of `onRequire` and `onError` hooks. We
+ * need to keep track of all such hooks, and this utility function allows us
+ * to do that without replicating this code in all of the places we end up
+ * doing the registration (which is, unfortunately, in more than one place).
+ *
+ * @param {object} nodule The module being instrumented.
+ * @param {object} instrumentation The instrumentation configuration. It has
+ * properties that represent the hooks being registered, in addition to a
+ * unique identifier that identifies the instrumentation configuration defining
+ * those hooks.
+ * @param {Symbol} [sym] The symbol to use when updating
+ * the module's registered hooks. Possible values are `symbols.instrumented`
+ * and `symbols.instrumentedErrored`.
+ */
+function registerHook(nodule, instrumentation, sym = symbols.instrumented) {
+  if (Array.isArray(nodule[sym]) === true) {
+    nodule[sym].push(instrumentation.instrumentationId)
+  } else {
+    nodule[sym] = [instrumentation.instrumentationId]
+  }
+}

test/integration/agent/reinstrument.tap.js

@@ -0,0 +1,41 @@
+/*
+ * Copyright 2024 New Relic Corporation. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+'use strict'
+
+const helper = require('../../lib/agent_helper')
+const agent = helper.instrumentMockedAgent({})
+agent.start(() => {
+  const api = helper.getAgentApi()
+
+  const didWeDoIt = []
+  api.instrument('test-logger', (shim, mod) => {
+    shim.wrap(mod.prototype, 'info', (shim, fn) => {
+      return function wrappedInfo() {
+        didWeDoIt.push('yes')
+        console.log('wrapped info')
+        console.log(shim.isWrapped(mod.prototype.info))
+        return fn.apply(this, arguments)
+      }
+    })
+  })
+
+  const { Writable } = require('stream')
+  const Person = require('./testdata/index')
+
+  const lines = []
+  const dest = new Writable({
+    write(chunk, _, cb) {
+      lines.push(chunk.toString())
+      cb()
+    }
+  })
+  const person = new Person(dest)
+  console.log(person.isHuman)
+  console.log(lines)
+
+  // This should be true if we have solved the problem correctly.
+  console.log(didWeDoIt.length === 2)
+})

test/integration/agent/testdata/.gitignore

@@ -0,0 +1 @@
+!node_modules

test/integration/agent/testdata/index.js

@@ -0,0 +1,11 @@
+/*
+ * Copyright 2024 New Relic Corporation. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+'use strict'
+
+// eslint-disable-next-line node/no-extraneous-require
+const Person = require('person')
+
+module.exports = Person

test/integration/module-loading/module-loading.tap.js

@@ -69,8 +69,8 @@ tap.test('should instrument multiple versions of the same package', function (t)
   const pkg2 = require(CUSTOM_MODULE_PATH_SUB)
   t.ok(pkg1[symbols.shim], 'should wrap first package')
   t.ok(pkg2[symbols.shim], 'should wrap sub package of same name, different version')
-  t.ok(instrumentation[symbols.instrumented].has('3.0.0'))
-  t.ok(instrumentation[symbols.instrumented].has('1.0.0'))
+  t.equal(pkg1[symbols.instrumented][0] > -1, true)
+  t.equal(pkg2[symbols.instrumented][0] > -1, true)
 })
 
 tap.test('should only log supportability metric for tracking type instrumentation', function (t) {
@@ -87,12 +87,12 @@ tap.test('should only log supportability metric for tracking type instrumentatio
   const PKG_VERSION = `${FEATURES.INSTRUMENTATION.ON_REQUIRE}/knex/Version/1`
 
   // eslint-disable-next-line node/no-extraneous-require
-  require('knex')
+  const knex = require('knex')
   const knexOnRequiredMetric = agent.metrics._metrics.unscoped[PKG]
   t.equal(knexOnRequiredMetric.callCount, 1, `should record ${PKG}`)
   const knexVersionMetric = agent.metrics._metrics.unscoped[PKG_VERSION]
   t.equal(knexVersionMetric.callCount, 1, `should record ${PKG_VERSION}`)
-  t.ok(shimmer.isInstrumented('knex'), 'should mark tracking modules as instrumented')
+  t.ok(shimmer.isInstrumented('knex', knex), 'should mark tracking modules as instrumented')
   t.end()
 })
 

test/integration/shimmer/module-load.tap.js

@@ -51,7 +51,9 @@ tap.test('Test Module Instrumentation Loading', (t) => {
 
     t.ok(module, 'loaded module')
     t.equal(module(), 'hello world', 'module behaves as expected')
-    t.ok(instrumentation[symbols.instrumented].has(process.version))
+    t.equal(module[symbols.instrumented].length, 1)
+    t.equal(module[symbols.instrumented][0], 1)
+
     t.end()
   })
 
@@ -84,7 +86,7 @@ tap.test('Test Module Instrumentation Loading', (t) => {
 
     t.ok(module, 'loaded module')
     t.equal(module(), 'hello world', 'module behaves as expected')
-    t.ok(instrumentation[symbols.instrumentedErrored].has(process.version))
+    t.strictSame(module[symbols.instrumentedErrored], [0])
     t.end()
   })
 })