Crysmags/ws integration

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,84 @@
+'use strict'
+
+const {
+  addHook
+} = require('./helpers/instrument')
+const shimmer = require('../../datadog-shimmer')
+
+const tracingChannel = require('dc-polyfill').tracingChannel
+const serverCh = tracingChannel('ws:server:connect')
+const producerCh = tracingChannel('ws:send')
+
+function createWrapRequest (ws, options) {
+  return function wrapRequest (request) {
+    return function (headers) {
+      if (!serverCh.start.hasSubscribers) return request.apply(this, arguments)
+
+      const ctx = { headers, ws, options }
+
+      return serverCh.tracePromise(() => request.call(this), ctx)
+    }
+  }
+}
+function createWrapEmit (ctx) {
+  return function wrapEmit (emit) {
+    return function (title, headers, req) {
+      ctx.resStatus = headers[0].split(' ')[1]
+
+      return serverCh.asyncStart.runStores(ctx, () => {
+        try {
+          return emit.apply(this, arguments)
+        } finally {
+          serverCh.asyncEnd.publish(ctx)
+        }
+      })
+    }
+  }
+}
+
+function wrapHandleUpgrade (handleUpgrade) {
+  return function () {
+    if (!serverCh.start.hasSubscribers) return handleUpgrade.apply(this, arguments)
+
+    const [req, socket, head, cb] = arguments
+    const ctx = { req }
+
+    return serverCh.tracePromise(() => {
+      shimmer.wrap(this, 'emit', createWrapEmit(ctx))
+      handleUpgrade.call(this, req, socket, head, cb)
+    }, ctx)
+  }
+}
+
+function wrapHandleSend (send) {
+  return function () {
+    if (!serverCh.start.hasSubscribers) return send.apply(this, arguments)
+
+    // console.log('argu,ent', arguments)
+    const [data, options, cb] = arguments
+    const ctx = { data }
+
+    return producerCh.tracePromise(() => {
+      send.call(this, data, options, cb)
+    }, ctx)
+  }
+}
+
+addHook({
+  name: 'ws',
+  file: 'lib/websocket-server.js'
+}, ws => {
+  shimmer.wrap(ws.prototype, 'handleUpgrade', wrapHandleUpgrade)
+
+  return ws
+})
+
+addHook({
+  name: 'ws',
+  file: 'lib/websocket.js'
+}, ws => {
+  // console.log('ws prototyoe', ws.prototype.send)
+  shimmer.wrap(ws.prototype, 'send', wrapHandleSend)
+
+  return ws
+})

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

@@ -0,0 +1,22 @@
+'use strict'
+
+const CompositePlugin = require('../../dd-trace/src/plugins/composite')
+
+const WSServerPlugin = require('./server')
+const WSProducerPlugin = require('./producer')
+
+class WSPlugin extends CompositePlugin {
+  static get id () { return 'websocket' }
+  static get plugins () {
+    return {
+      server: WSServerPlugin,
+      producer: WSProducerPlugin
+    }
+  }
+
+  configure (config) {
+    return super.configure(config)
+  }
+}
+
+module.exports = WSPlugin

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

@@ -0,0 +1,52 @@
+'use strict'
+
+const TracingPlugin = require('../../dd-trace/src/plugins/tracing.js')
+const tags = require('../../../ext/tags.js')
+const { storage } = require('../../datadog-core')
+
+class WSProducerPlugin extends TracingPlugin {
+  static get id () { return 'websocket' }
+  static get prefix () { return 'tracing:ws:send' }
+  static get type () { return 'websocket' }
+  static get kind () { return 'producer' }
+
+  bindStart (ctx) {
+    // const store = storage('legacy').getStore()
+    // const childOf = store ? store.span : null
+    // console.log('store', store, childOf)
+    const span = this.startSpan(this.operationName(), {
+      meta: {
+        service: this.serviceName({ pluginConfig: this.config }),
+        // 'resource.name': 'websocket ' + ,
+        'span.type': 'websocket',
+        'span.kind': 'producer'
+
+      }
+
+    }, true)
+    // console.log('ctx', span.operationName())
+    // span.addLink(ctx)
+    // console.log(span)
+    ctx.span = span
+    ctx.currentStore = { span }
+
+    return ctx.currentStore
+  }
+
+  end (ctx) {
+    const store = storage('legacy').getStore()
+    const childOf = store ? store.span : null
+    console.log('store in producer', store, childOf)
+
+    ctx.span.addLink(ctx.span._spanContext)
+    // console.log(ctx.span)
+    // console.log('in the end', ctx)
+    // ctx.req.res = ctx.resStatus
+
+    // ctx.span.setTag(HTTP_STATUS_CODE, ctx.req.res)
+    // if (!ctx.span) return
+    ctx.span.finish()
+  }
+}
+
+module.exports = WSProducerPlugin

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

@@ -0,0 +1,87 @@
+'use strict'
+
+const TracingPlugin = require('../../dd-trace/src/plugins/tracing.js')
+const tags = require('../../../ext/tags.js')
+const { storage } = require('../../datadog-core')
+
+const HTTP_STATUS_CODE = tags.HTTP_STATUS_CODE
+const HTTP_REQUEST_HEADERS = tags.HTTP_REQUEST_HEADERS
+const HTTP_RESPONSE_HEADERS = tags.HTTP_RESPONSE_HEADERS
+
+class WSServerPlugin extends TracingPlugin {
+  static get id () { return 'websocket' }
+  static get prefix () { return 'tracing:ws:server:connect' }
+  static get type () { return 'websocket' }
+  static get kind () { return 'consumer' }
+
+  bindStart (ctx) {
+    const { http = {} } = ctx
+    const req = ctx.req
+
+    const options = {}
+    const headers = Object.entries(req.headers)
+    options.headers = Object.fromEntries(headers)
+    options.method = req.method
+
+    const agent = options.agent || options._defaultAgent || http.globalAgent || {}
+    const protocol = options.protocol || agent.protocol || 'http:'
+    const hostname = options.hostname || options.host || 'localhost'
+    const host = options.port ? `${hostname}:${options.port}` : hostname
+    const pathname = options.path || options.pathname
+    const path = pathname ? pathname.split(/[?#]/)[0] : '/'
+    const uri = `${protocol}//${host}${path}`
+
+    ctx.args = { options }
+
+    const span = this.startSpan(this.operationName(), {
+      meta: {
+        service: this.serviceName({ pluginConfig: this.config }),
+        'span.type': 'ws',
+        'http.upgraded': 'websocket',
+        'http.method': options.method,
+        'http.url': uri,
+        'resource.name': options.method,
+        'span.kind': 'server'
+
+      }
+
+    }, true)
+    ctx.span = span
+    ctx.currentStore = { span }
+
+    return ctx.currentStore
+  }
+
+  // asyncStart (ctx) {
+
+  //   ctx?.currentStore?.span.finish()
+
+  //   ctx.res = ctx.resStatus
+
+  //   ctx.currentStore.span.setTag(HTTP_STATUS_CODE, ctx.res)
+
+  //   return ctx.currentStore
+  // }
+
+  end (ctx) {
+    const store = storage('legacy').getStore()
+    const childOf = store ? store.span : null
+    console.log('store in server', store, childOf)
+    ctx.req.res = ctx.resStatus
+
+    ctx.span.setTag(HTTP_STATUS_CODE, ctx.req.res)
+    if (!ctx.span) return
+    ctx.span.finish()
+  }
+
+  // asyncEnd (ctx) {
+  //   ctx.res = ctx.resStatus
+
+  //   ctx.currentStore.span.setTag(HTTP_STATUS_CODE, ctx.res)
+  //   ctx.span.setTag(HTTP_STATUS_CODE, ctx.res)
+
+  //   return this.finish(ctx)
+  // }
+}
+
+module.exports = WSServerPlugin

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: {
+    websocket: {
+      opName: () => 'websocket.request',
+      serviceName: identityService
+    }
+  },
+  producer: {
+    websocket: {
+      opName: () => 'websocket.send',
+      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: {
+    websocket: {
+      opName: () => 'websocket.request',
+      serviceName: identityService
+    }
+  },
+  producer: {
+    websocket: {
+      opName: () => 'websocket.send',
+      serviceName: identityService
+    }
+  }
+}
+
+module.exports = websocket

packages/dd-trace/src/supported-configurations.json

@@ -397,7 +397,7 @@
     "DD_TRACE_WHEN_ENABLED": ["A"],
     "DD_TRACE_WINSTON_ENABLED": ["A"],
     "DD_TRACE_WORKERPOOL_ENABLED": ["A"],
-    "DD_TRACE_WS_ENABLED": ["A"],
+    "DD_TRACE_WEBSOCKET_ENABLED": ["A"],
     "DD_TRACE_X_DATADOG_TAGS_MAX_LENGTH": ["A"],
     "DD_TRACING_ENABLED": ["A"],
     "DD_VERSION": ["A"],