ws plugin handles handshake

Author: crysmags

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

@@ -132,5 +132,6 @@ module.exports = {
   vm: () => require('../vm'),
   when: () => require('../when'),
   winston: () => require('../winston'),
-  workerpool: () => require('../mocha')
+  workerpool: () => require('../mocha'),
+  ws: () => require('../ws')
 }

packages/datadog-instrumentations/src/ws.js

@@ -0,0 +1,45 @@
+'use strict'
+
+const {
+  addHook,
+  channel,
+  AsyncResource
+} = require('./helpers/instrument')
+const shimmer = require('../../datadog-shimmer')
+
+const tracingChannel = require('dc-polyfill').tracingChannel
+const ch = tracingChannel('ws:client:connect')
+
+function createWrapRequest (ws, options) {
+  return function wrapRequest (request) {
+    return function (headers) {
+      if (!ch.start.hasSubscribers) return request.apply(this, arguments)
+
+      const ctx = { headers, ws, options }
+
+      return ch.tracePromise(() => request.call(this), ctx)
+    }
+  }
+}
+
+function wrapHandleUpgrade (handleUpgrade) {
+  return function () {
+    if (!ch.start.hasSubscribers) return handleUpgrade.apply(this, arguments)
+
+    const [req, socket, head, cb] = arguments
+    const ctx = { req }
+
+    return ch.tracePromise(() => {
+      handleUpgrade.call(this, req, socket, head, cb)
+    }, ctx)
+  }
+}
+
+addHook({
+  name: 'ws',
+  file: 'lib/websocket-server.js'
+}, ws => {
+  shimmer.wrap(ws.prototype, 'handleUpgrade', wrapHandleUpgrade)
+
+  return ws
+})

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

@@ -0,0 +1,56 @@
+'use strict'
+
+const TracingPlugin = require('../../dd-trace/src/plugins/tracing.js')
+
+class WSPlugin extends TracingPlugin {
+  static get id () { return 'ws' }
+  static get prefix () { return 'tracing:ws:client:connect' }
+  static get type () { return 'websocket' }
+  static get kind () { return 'consumer' }
+
+  bindStart (message) {
+    const req = message.req
+
+    const options = {}
+    const headers = Object.entries(req.headers)
+    options.headers = Object.fromEntries(headers)
+    options.method = req.method
+
+    message.args = { options }
+
+    const span = this.startSpan(this.operationName(), {
+      meta: {
+        service: this.serviceName({ pluginConfig: this.config }),
+        'span.type': 'ws',
+        'http.upgraded': 'websocket'
+
+      }
+
+    }, true)
+    message.span = span
+    // message.parentStore = store
+    message.currentStore = { span }
+
+    return message.currentStore
+  }
+
+  asyncStart (ctx) {
+    console.log('ctx', ctx)
+    ctx?.currentStore?.span.finish()
+    return ctx.parentStore
+  }
+  // finish ({ req, res, span }) {
+  //   console.log('arguments', arguments)
+  //   if (!span) return
+  //   console.log('span')
+  //   span.finish()
+  // }
+
+  // asyncEnd (message) {
+  //   console.log('async end')
+  //   message.res = message.result
+  //   return this.finish(message)
+  // }
+}
+
+module.exports = WSPlugin

packages/dd-trace/src/plugins/index.js

@@ -97,5 +97,6 @@ module.exports = {
   get sharedb () { return require('../../../datadog-plugin-sharedb/src') },
   get tedious () { return require('../../../datadog-plugin-tedious/src') },
   get undici () { return require('../../../datadog-plugin-undici/src') },
-  get winston () { return require('../../../datadog-plugin-winston/src') }
+  get winston () { return require('../../../datadog-plugin-winston/src') },
+  get ws () { return require('../../../datadog-plugin-ws/src') }
 }

packages/dd-trace/src/service-naming/schemas/v0/index.js

@@ -4,5 +4,6 @@ const storage = require('./storage')
 const graphql = require('./graphql')
 const web = require('./web')
 const serverless = require('./serverless')
+const websocket = require('./websocket')
 
-module.exports = new SchemaDefinition({ messaging, storage, web, graphql, serverless })
+module.exports = new SchemaDefinition({ messaging, storage, web, graphql, serverless, websocket })

packages/dd-trace/src/service-naming/schemas/v0/websocket.js

@@ -0,0 +1,18 @@
+const { identityService } = require('../util')
+
+const websocket = {
+  consumer: {
+    ws: {
+      opName: () => 'ws.request',
+      serviceName: identityService
+    }
+  },
+  producer: {
+    ws: {
+      opName: () => 'ws.request',
+      serviceName: identityService
+    }
+  }
+}
+
+module.exports = websocket

packages/dd-trace/src/service-naming/schemas/v1/index.js

@@ -4,5 +4,6 @@ const storage = require('./storage')
 const graphql = require('./graphql')
 const web = require('./web')
 const serverless = require('./serverless')
+const websocket = require('./websocket')
 
-module.exports = new SchemaDefinition({ messaging, storage, web, graphql, serverless })
+module.exports = new SchemaDefinition({ messaging, storage, web, graphql, serverless, websocket })

packages/dd-trace/src/service-naming/schemas/v1/websocket.js

@@ -0,0 +1,18 @@
+const { identityService } = require('../util')
+
+const websocket = {
+  consumer: {
+    ws: {
+      opName: () => 'ws.request',
+      serviceName: identityService
+    }
+  },
+  producer: {
+    ws: {
+      opName: () => 'ws.request',
+      serviceName: identityService
+    }
+  }
+}
+
+module.exports = websocket