fix: reduce memory allocation and use transparent instrumentation (#7422)

BridgeAR · web-flow · commit c195ee802ffe · 2026-02-05T16:34:50.000+01:00
The instrumentation was noticable by a user so far and hooked into
the public event listener addition / removal. The event handlers
being added could not be detected as not being the ones by the user.

That is improved here by using an internal only instrumentation
that should also reduce overhead and needs less coordination
between the additions and removals.

The plugin is refactored to have less memory retantion. So far,
the whole span was attached and in case long living ws connections
existed with a high load, the spans could add up without ever being
freed. Now that tracking overhead is much lower and the span is
created when needed.
diff --git a/packages/datadog-instrumentations/src/ws.js b/packages/datadog-instrumentations/src/ws.js
@@ -1,6 +1,6 @@
 'use strict'
 
-const tracingChannel = require('dc-polyfill').tracingChannel
+const { tracingChannel } = /** @type {import('node:diagnostics_channel')} */ (require('dc-polyfill'))
 
 const shimmer = require('../../datadog-shimmer')
 const {
@@ -15,7 +15,50 @@ const closeCh = tracingChannel('ws:close')
 const emitCh = channel('tracing:ws:server:connect:emit')
 // TODO: Add a error channel / handle error events properly.
 
-const eventHandlerMap = new WeakMap()
+/**
+ * @typedef {object} WebSocketServerPrototype
+ * @property {(...args: unknown[]) => unknown} handleUpgrade
+ * @property {(...args: unknown[]) => unknown} emit
+ */
+
+/**
+ * @typedef {{ prototype: WebSocketServerPrototype }} WebSocketServerClass
+ */
+
+/**
+ * @typedef {object} WebSocketPrototype
+ * @property {(...args: unknown[]) => unknown} send
+ * @property {(...args: unknown[]) => unknown} close
+ * @property {(...args: unknown[]) => unknown} setSocket
+ */
+
+/**
+ * @typedef {{ prototype: WebSocketPrototype }} WebSocketClass
+ */
+
+/**
+ * @typedef {object} ReceiverPrototype
+ * @property {(eventName: string, listener: (...args: unknown[]) => unknown) => unknown} on
+ * @property {(eventName: string, listener: (...args: unknown[]) => unknown) => unknown} addListener
+ */
+
+/**
+ * @typedef {{ prototype: ReceiverPrototype }} ReceiverClass
+ */
+
+/**
+ * @typedef {string | Buffer | ArrayBuffer | ArrayBufferView | Blob | Buffer[]} WebSocketMessageData
+ */
+
+/**
+ * @typedef {object} WebSocketInstance
+ * @property {(...args: unknown[]) => unknown} emit
+ * @property {(eventName: string) => number} [listenerCount]
+ * @property {{ _socket?: unknown } | undefined} [_sender]
+ * @property {unknown} [_receiver]
+ */
+
+let kWebSocketSymbol
 
 function wrapHandleUpgrade (handleUpgrade) {
   return function () {
@@ -41,7 +84,9 @@ function wrapHandleUpgrade (handleUpgrade) {
 
 function wrapSend (send) {
   return function wrappedSend (...args) {
-    if (!producerCh.start.hasSubscribers) return send.apply(this, arguments)
+    if (!producerCh.start.hasSubscribers) {
+      return send.apply(this, arguments)
+    }
 
     const [data, options, cb] = arguments
 
@@ -55,7 +100,9 @@ function wrapSend (send) {
 
 function createWrapEmit (emit) {
   return function (title, headers, req) {
-    if (!serverCh.start.hasSubscribers || title !== 'headers') return emit.apply(this, arguments)
+    if (!serverCh.start.hasSubscribers || title !== 'headers') {
+      return emit.apply(this, arguments)
+    }
 
     const ctx = { req }
     ctx.req.resStatus = headers[0].split(' ')[1]
@@ -70,35 +117,42 @@ function createWrapEmit (emit) {
   }
 }
 
-function createWrappedHandler (handler) {
-  return shimmer.wrapFunction(handler, originalHandler => function (data, binary) {
-    const byteLength = dataLength(data)
+/**
+ * @param {Function} setSocket
+ * @returns {(...args: unknown[]) => unknown}
+ */
+/**
+ * @param {Function} on
+ * @returns {(...args: unknown[]) => unknown}
+ */
+function wrapReceiverOn (on) {
+  return function wrappedOn (eventName, handler) {
+    if (eventName !== 'message' || typeof handler !== 'function') {
+      return on.apply(this, arguments)
+    }
 
-    const ctx = { data, binary, socket: this._sender?._socket, byteLength }
+    const wrappedHandler = function (data, isBinary) {
+      if (!receiverCh.start.hasSubscribers || !kWebSocketSymbol) {
+        return handler.call(this, data, isBinary)
+      }
 
-    return receiverCh.traceSync(originalHandler, ctx, this, data, binary)
-  })
-}
+      const websocket = /** @type {WebSocketInstance | undefined} */ (this[kWebSocketSymbol])
+      // Avoid receive spans when no one listens to messages.
+      if (websocket && typeof websocket.listenerCount === 'function' && websocket.listenerCount('message') === 0) {
+        return handler.call(this, data, isBinary)
+      }
+      const socket = websocket?._sender?._socket
+      if (!socket) {
+        return handler.call(this, data, isBinary)
+      }
 
-function wrapListener (originalOn) {
-  return function (eventName, handler) {
-    if (eventName === 'message') {
-      // Prevent multiple wrapping of the same handler in case the user adds the listener multiple times
-      const wrappedHandler = eventHandlerMap.get(handler) ?? createWrappedHandler(handler)
-      eventHandlerMap.set(handler, wrappedHandler)
-      return originalOn.call(this, eventName, wrappedHandler)
-    }
-    return originalOn.apply(this, arguments)
-  }
-}
+      const byteLength = dataLength(/** @type {WebSocketMessageData} */ (data))
+      const ctx = { data, binary: isBinary, socket, byteLength }
 
-function removeListener (originalOff) {
-  return function (eventName, handler) {
-    if (eventName === 'message') {
-      const wrappedHandler = eventHandlerMap.get(handler) || handler
-      return originalOff.call(this, eventName, wrappedHandler)
+      return receiverCh.traceSync(handler, ctx, this, data, isBinary)
     }
-    return originalOff.apply(this, arguments)
+
+    return on.call(this, eventName, wrappedHandler)
   }
 }
 
@@ -120,7 +174,8 @@ addHook({
   name: 'ws',
   file: 'lib/websocket-server.js',
   versions: ['>=8.0.0'],
-}, ws => {
+}, moduleExports => {
+  const ws = /** @type {WebSocketServerClass} */ (moduleExports)
   shimmer.wrap(ws.prototype, 'handleUpgrade', wrapHandleUpgrade)
   shimmer.wrap(ws.prototype, 'emit', createWrapEmit)
   return ws
@@ -130,26 +185,58 @@ addHook({
   name: 'ws',
   file: 'lib/websocket.js',
   versions: ['>=8.0.0'],
-}, ws => {
+}, moduleExports => {
+  const ws = /** @type {WebSocketClass} */ (moduleExports)
   shimmer.wrap(ws.prototype, 'send', wrapSend)
   shimmer.wrap(ws.prototype, 'close', wrapClose)
 
-  // TODO: Do not wrap these methods. Instead, add a listener to the websocket instance when one is created.
-  // That way it avoids producing too many spans for the same websocket instance and less user code is impacted.
-  shimmer.wrap(ws.prototype, 'on', wrapListener)
-  shimmer.wrap(ws.prototype, 'addListener', wrapListener)
-  shimmer.wrap(ws.prototype, 'off', removeListener)
-  shimmer.wrap(ws.prototype, 'removeListener', removeListener)
-
   return ws
 })
 
+addHook({
+  name: 'ws',
+  file: 'lib/constants.js',
+  versions: ['>=8.0.0'],
+}, moduleExports => {
+  const constants = /** @type {{ kWebSocket?: symbol }} */ (moduleExports)
+  kWebSocketSymbol = constants.kWebSocket
+  return constants
+})
+
+addHook({
+  name: 'ws',
+  file: 'lib/receiver.js',
+  versions: ['>=8.0.0'],
+}, moduleExports => {
+  const Receiver = /** @type {ReceiverClass} */ (moduleExports)
+  shimmer.wrap(Receiver.prototype, 'on', wrapReceiverOn)
+  shimmer.wrap(Receiver.prototype, 'addListener', wrapReceiverOn)
+  return Receiver
+})
+
+/**
+ * @param {WebSocketMessageData} data
+ * @returns {number}
+ */
 function dataLength (data) {
   if (typeof data === 'string') {
     return Buffer.byteLength(data)
   }
   if (data instanceof Blob) {
     return data.size
   }
-  return data?.length ?? 0
+  if (ArrayBuffer.isView(data)) {
+    return data.byteLength
+  }
+  if (data instanceof ArrayBuffer) {
+    return data.byteLength
+  }
+  let total = 0
+  if (Array.isArray(data)) {
+    const chunks = /** @type {Buffer[]} */ (data)
+    for (const chunk of chunks) {
+      total += chunk.length
+    }
+  }
+  return total
 }
diff --git a/packages/datadog-plugin-ws/src/close.js b/packages/datadog-plugin-ws/src/close.js
@@ -28,7 +28,7 @@ class WSClosePlugin extends TracingPlugin {
     if (!socket?.spanContext) return
 
     const spanKind = isPeerClose ? 'consumer' : 'producer'
-    const spanTags = socket.spanContext.spanTags
+    const spanTags = socket.spanTags
     const path = spanTags['resource.name'].split(' ')[1]
     const service = this.serviceName({ pluginConfig: this.config })
     const span = this.startSpan(this.operationName(), {
@@ -76,14 +76,13 @@ class WSClosePlugin extends TracingPlugin {
       linkAttributes['dd.kind'] = isIncoming ? 'executed_by' : 'resuming'
 
       // Add span pointer for context propagation
-      if (this.config.traceWebsocketMessagesEnabled && ctx.socket.handshakeSpan) {
-        const handshakeSpan = ctx.socket.handshakeSpan
+      if (this.config.traceWebsocketMessagesEnabled && ctx.socket.spanContext) {
+        const handshakeContext = ctx.socket.spanContext
 
         // Only add span pointers if distributed tracing is enabled and handshake has distributed context
-        if (hasDistributedTracingContext(handshakeSpan, ctx.socket)) {
+        if (hasDistributedTracingContext(handshakeContext, ctx.socket)) {
           const counterType = isIncoming ? 'receiveCounter' : 'sendCounter'
           const counter = incrementWebSocketCounter(ctx.socket, counterType)
-          const handshakeContext = handshakeSpan.context()
 
           const ptrHash = buildWebSocketSpanPointerHash(
             handshakeContext._traceId,
diff --git a/packages/datadog-plugin-ws/src/producer.js b/packages/datadog-plugin-ws/src/producer.js
@@ -22,7 +22,7 @@ class WSProducerPlugin extends TracingPlugin {
     const { byteLength, socket, binary } = ctx
     if (!socket.spanContext) return
 
-    const spanTags = socket.spanContext.spanTags
+    const spanTags = socket.spanTags
     const path = spanTags['resource.name'].split(' ')[1]
     const opCode = binary ? 'binary' : 'text'
     const service = this.serviceName({ pluginConfig: this.config })
@@ -61,13 +61,12 @@ class WSProducerPlugin extends TracingPlugin {
       const linkAttributes = { 'dd.kind': 'resuming' }
 
       // Add span pointer for context propagation
-      if (this.config.traceWebsocketMessagesEnabled && ctx.socket.handshakeSpan) {
-        const handshakeSpan = ctx.socket.handshakeSpan
+      if (this.config.traceWebsocketMessagesEnabled && ctx.socket.spanContext) {
+        const handshakeContext = ctx.socket.spanContext
 
         // Only add span pointers if distributed tracing is enabled and handshake has distributed context
-        if (hasDistributedTracingContext(handshakeSpan, ctx.socket)) {
+        if (hasDistributedTracingContext(handshakeContext, ctx.socket)) {
           const counter = incrementWebSocketCounter(ctx.socket, 'sendCounter')
-          const handshakeContext = handshakeSpan.context()
 
           const ptrHash = buildWebSocketSpanPointerHash(
             handshakeContext._traceId,
diff --git a/packages/datadog-plugin-ws/src/receiver.js b/packages/datadog-plugin-ws/src/receiver.js
@@ -27,7 +27,7 @@ class WSReceiverPlugin extends TracingPlugin {
     const { byteLength, socket, binary } = ctx
     if (!socket.spanContext) return
 
-    const spanTags = socket.spanContext.spanTags
+    const spanTags = socket.spanTags
     const path = spanTags['resource.name'].split(' ')[1]
     const opCode = binary ? 'binary' : 'text'
 
@@ -72,13 +72,12 @@ class WSReceiverPlugin extends TracingPlugin {
       const linkAttributes = { 'dd.kind': 'executed_by' }
 
       // Add span pointer for context propagation
-      if (this.config.traceWebsocketMessagesEnabled && ctx.socket.handshakeSpan) {
-        const handshakeSpan = ctx.socket.handshakeSpan
+      if (this.config.traceWebsocketMessagesEnabled && ctx.socket.spanContext) {
+        const handshakeContext = ctx.socket.spanContext
 
         // Only add span pointers if distributed tracing is enabled and handshake has distributed context
-        if (hasDistributedTracingContext(handshakeSpan, ctx.socket)) {
+        if (hasDistributedTracingContext(handshakeContext, ctx.socket)) {
           const counter = incrementWebSocketCounter(ctx.socket, 'receiveCounter')
-          const handshakeContext = handshakeSpan.context()
 
           const ptrHash = buildWebSocketSpanPointerHash(
             handshakeContext._traceId,
diff --git a/packages/datadog-plugin-ws/src/server.js b/packages/datadog-plugin-ws/src/server.js
@@ -2,8 +2,12 @@
 
 const TracingPlugin = require('../../dd-trace/src/plugins/tracing.js')
 const tags = require('../../../ext/tags.js')
-const { FORMAT_HTTP_HEADERS } = require('../../../ext/formats')
-const { initWebSocketMessageCounters } = require('./util')
+const { HTTP_HEADERS } = require('../../../ext/formats')
+const {
+  createWebSocketSpanContext,
+  hasTraceHeaders,
+  initWebSocketMessageCounters,
+} = require('./util')
 
 const HTTP_STATUS_CODE = tags.HTTP_STATUS_CODE
 
@@ -27,11 +31,12 @@ class WSServerPlugin extends TracingPlugin {
     const indexOfParam = url.indexOf('?')
     const route = indexOfParam === -1 ? url : url.slice(0, indexOfParam)
     const uri = `${protocol}//${host}${route}`
+    const resourceName = `${options.method} ${route}`
 
     ctx.args = { options }
 
     // Extract distributed tracing context from request headers
-    const childOf = this.tracer.extract(FORMAT_HTTP_HEADERS, req.headers)
+    const childOf = this.tracer.extract(HTTP_HEADERS, req.headers)
 
     const service = this.serviceName({ pluginConfig: this.config })
     const span = this.startSpan(this.operationName(), {
@@ -42,19 +47,19 @@ class WSServerPlugin extends TracingPlugin {
         'http.upgraded': 'websocket',
         'http.method': options.method,
         'http.url': uri,
-        'resource.name': `${options.method} ${route}`,
+        'resource.name': resourceName,
         'span.kind': 'server',
       },
 
     }, ctx)
     ctx.span = span
 
-    ctx.socket.spanContext = ctx.span._spanContext
-    ctx.socket.spanContext.spanTags = ctx.span._spanContext._tags
-    // Store the handshake span for use in message span pointers
-    ctx.socket.handshakeSpan = ctx.span
-    // Store the request headers for distributed tracing check
-    ctx.socket.requestHeaders = req.headers
+    ctx.socket.spanTags = {
+      'resource.name': resourceName,
+      'service.name': service,
+    }
+    ctx.socket.spanContext = createWebSocketSpanContext(ctx.span._spanContext)
+    ctx.socket.hasTraceHeaders = hasTraceHeaders(req.headers)
 
     // Initialize message counters for span pointers
     initWebSocketMessageCounters(ctx.socket)
diff --git a/packages/datadog-plugin-ws/src/util.js b/packages/datadog-plugin-ws/src/util.js
diff --git a/packages/datadog-plugin-ws/test/index.spec.js b/packages/datadog-plugin-ws/test/index.spec.js
