feat(tracing): implement protobufjs DSM schema support (#4701)

* add protobufjs schemas support for DSM

.github/workflows/plugins.yml

@@ -803,6 +803,15 @@ jobs:
       - uses: actions/checkout@v4
       - uses: ./.github/actions/plugins/test
 
+  protobufjs:
+    runs-on: ubuntu-latest
+    env:
+      PLUGINS: protobufjs
+      DD_DATA_STREAMS_ENABLED: true
+    steps:
+      - uses: actions/checkout@v4
+      - uses: ./.github/actions/plugins/test-and-upstream
+
   q:
     runs-on: ubuntu-latest
     env:

docs/API.md

@@ -87,6 +87,7 @@ tracer.use('pg', {
 <h5 id="pg"></h5>
 <h5 id="pg-tags"></h5>
 <h5 id="pg-config"></h5>
+<h5 id="protobufjs"></h5>
 <h5 id="redis"></h5>
 <h5 id="redis-tags"></h5>
 <h5 id="redis-config"></h5>
@@ -142,6 +143,7 @@ tracer.use('pg', {
 * [pg](./interfaces/export_.plugins.pg.html)
 * [promise](./interfaces/export_.plugins.promise.html)
 * [promise-js](./interfaces/export_.plugins.promise_js.html)
+* [protobufjs](./interfaces/export_.plugins.protobufjs.html)
 * [q](./interfaces/export_.plugins.q.html)
 * [redis](./interfaces/export_.plugins.redis.html)
 * [restify](./interfaces/export_.plugins.restify.html)

docs/add-redirects.sh

@@ -55,6 +55,7 @@ declare -a plugins=(
   "pg"
   "promise"
   "promise_js"
+  "protobufjs"
   "q"
   "redis"
   "restify"

docs/test.ts

@@ -362,6 +362,7 @@ tracer.use('playwright');
 tracer.use('pg');
 tracer.use('pg', { service: params => `${params.host}-${params.database}` });
 tracer.use('pino');
+tracer.use('protobufjs');
 tracer.use('redis');
 tracer.use('redis', redisOptions);
 tracer.use('restify');

index.d.ts

@@ -190,6 +190,7 @@ interface Plugins {
   "playwright": tracer.plugins.playwright;
   "pg": tracer.plugins.pg;
   "pino": tracer.plugins.pino;
+  "protobufjs": tracer.plugins.protobufjs;
   "redis": tracer.plugins.redis;
   "restify": tracer.plugins.restify;
   "rhea": tracer.plugins.rhea;
@@ -1731,6 +1732,12 @@ declare namespace tracer {
      * on the tracer.
      */
     interface pino extends Integration {}
+
+    /**
+     * This plugin automatically patches the [protobufjs](https://protobufjs.github.io/protobuf.js/)
+     * to collect protobuf message schemas when Datastreams Monitoring is enabled.
+     */
+    interface protobufjs extends Integration {}
 
     /**
      * This plugin automatically instruments the

packages/datadog-instrumentations/src/helpers/hooks.js

@@ -100,6 +100,7 @@ module.exports = {
   playwright: () => require('../playwright'),
   'promise-js': () => require('../promise-js'),
   promise: () => require('../promise'),
+  protobufjs: () => require('../protobufjs'),
   q: () => require('../q'),
   qs: () => require('../qs'),
   redis: () => require('../redis'),

packages/datadog-instrumentations/src/protobufjs.js

@@ -0,0 +1,127 @@
+const shimmer = require('../../datadog-shimmer')
+const { addHook } = require('./helpers/instrument')
+
+const dc = require('dc-polyfill')
+const serializeChannel = dc.channel('apm:protobufjs:serialize-start')
+const deserializeChannel = dc.channel('apm:protobufjs:deserialize-end')
+
+function wrapSerialization (messageClass) {
+  if (messageClass?.encode) {
+    shimmer.wrap(messageClass, 'encode', original => function () {
+      if (!serializeChannel.hasSubscribers) {
+        return original.apply(this, arguments)
+      }
+      serializeChannel.publish({ messageClass: this })
+      return original.apply(this, arguments)
+    })
+  }
+}
+
+function wrapDeserialization (messageClass) {
+  if (messageClass?.decode) {
+    shimmer.wrap(messageClass, 'decode', original => function () {
+      if (!deserializeChannel.hasSubscribers) {
+        return original.apply(this, arguments)
+      }
+      const result = original.apply(this, arguments)
+      deserializeChannel.publish({ messageClass: result })
+      return result
+    })
+  }
+}
+
+function wrapSetup (messageClass) {
+  if (messageClass?.setup) {
+    shimmer.wrap(messageClass, 'setup', original => function () {
+      const result = original.apply(this, arguments)
+
+      wrapSerialization(messageClass)
+      wrapDeserialization(messageClass)
+
+      return result
+    })
+  }
+}
+
+function wrapProtobufClasses (root) {
+  if (!root) {
+    return
+  }
+
+  if (root.decode) {
+    wrapSetup(root)
+  }
+
+  if (root.nestedArray) {
+    for (const subRoot of root.nestedArray) {
+      wrapProtobufClasses(subRoot)
+    }
+  }
+}
+
+function wrapReflection (protobuf) {
+  const reflectionMethods = [
+    {
+      target: protobuf.Root,
+      name: 'fromJSON'
+    },
+    {
+      target: protobuf.Type.prototype,
+      name: 'fromObject'
+    }
+  ]
+
+  reflectionMethods.forEach(method => {
+    shimmer.wrap(method.target, method.name, original => function () {
+      const result = original.apply(this, arguments)
+      if (result.nested) {
+        for (const type in result.nested) {
+          wrapSetup(result.nested[type])
+        }
+      }
+      if (result.$type) {
+        wrapSetup(result.$type)
+      }
+      return result
+    })
+  })
+}
+
+function isPromise (obj) {
+  return !!obj && (typeof obj === 'object' || typeof obj === 'function') && typeof obj.then === 'function'
+}
+
+addHook({
+  name: 'protobufjs',
+  versions: ['>=6.8.0']
+}, protobuf => {
+  shimmer.wrap(protobuf.Root.prototype, 'load', original => function () {
+    const result = original.apply(this, arguments)
+    if (isPromise(result)) {
+      return result.then(root => {
+        wrapProtobufClasses(root)
+        return root
+      })
+    } else {
+      // If result is not a promise, directly wrap the protobuf classes
+      wrapProtobufClasses(result)
+      return result
+    }
+  })
+
+  shimmer.wrap(protobuf.Root.prototype, 'loadSync', original => function () {
+    const root = original.apply(this, arguments)
+    wrapProtobufClasses(root)
+    return root
+  })
+
+  shimmer.wrap(protobuf, 'Type', Original => function () {
+    const typeInstance = new Original(...arguments)
+    wrapSetup(typeInstance)
+    return typeInstance
+  })
+
+  wrapReflection(protobuf)
+
+  return protobuf
+})

packages/datadog-plugin-protobufjs/src/index.js

@@ -0,0 +1,14 @@
+const SchemaPlugin = require('../../dd-trace/src/plugins/schema')
+const SchemaExtractor = require('./schema_iterator')
+
+class ProtobufjsPlugin extends SchemaPlugin {
+  static get id () {
+    return 'protobufjs'
+  }
+
+  static get schemaExtractor () {
+    return SchemaExtractor
+  }
+}
+
+module.exports = ProtobufjsPlugin