fix(ws): avoid retaining connection span (#7469)

Run ws socket setup in a store scope without the connection span so
internal ws handlers don't capture and keep it alive.

Author: BridgeAR

packages/datadog-instrumentations/src/ws.js

@@ -13,6 +13,7 @@ const producerCh = tracingChannel('ws:send')
 const receiverCh = tracingChannel('ws:receive')
 const closeCh = tracingChannel('ws:close')
 const emitCh = channel('tracing:ws:server:connect:emit')
+const setSocketCh = channel('tracing:ws:server:connect:setSocket')
 // TODO: Add a error channel / handle error events properly.
 
 /**
@@ -181,12 +182,33 @@ addHook({
   return ws
 })
 
+/**
+ * Prevent internal event handlers (data, close, etc.) registered by the ws library to
+ * capture the connection span in their async context. Otherwise, the
+ * finished connection span is retained for the entire lifetime of the WebSocket
+ * (via ACF -> handle -> WeakMap).
+ *
+ * @param {Function} setSocket
+ * @returns {(...args: unknown[]) => unknown}
+ */
+function wrapSetSocket (setSocket) {
+  return function wrappedSetSocket (...args) {
+    if (!setSocketCh.hasSubscribers) {
+      return setSocket.apply(this, args)
+    }
+    return setSocketCh.runStores({}, () => {
+      return setSocket.apply(this, args)
+    })
+  }
+}
+
 addHook({
   name: 'ws',
   file: 'lib/websocket.js',
   versions: ['>=8.0.0'],
 }, moduleExports => {
   const ws = /** @type {WebSocketClass} */ (moduleExports)
+  shimmer.wrap(ws.prototype, 'setSocket', wrapSetSocket)
   shimmer.wrap(ws.prototype, 'send', wrapSend)
   shimmer.wrap(ws.prototype, 'close', wrapClose)
 

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

@@ -17,6 +17,14 @@ class WSServerPlugin extends TracingPlugin {
   static get type () { return 'websocket' }
   static get kind () { return 'request' }
 
+  constructor (...args) {
+    super(...args)
+
+    // Bind the setSocket channel so internal ws event handlers (data, close)
+    // don't capture their async context.
+    this.addBind('tracing:ws:server:connect:setSocket', () => {})
+  }
+
   bindStart (ctx) {
     const req = ctx.req
 

packages/datadog-plugin-ws/test/index.spec.js

@@ -3,9 +3,11 @@
 const assert = require('node:assert')
 const { once } = require('node:events')
 
+const dc = require('dc-polyfill')
 const { after, afterEach, before, beforeEach, describe, it } = require('mocha')
 
 const agent = require('../../dd-trace/test/plugins/agent')
+const { storage } = require('../../datadog-core')
 const { withVersions } = require('../../dd-trace/test/setup/mocha')
 const { assertObjectContains } = require('../../../integration-tests/helpers')
 
@@ -65,6 +67,29 @@ describe('Plugin', () => {
           agent.close({ ritmReset: false, wipe: true })
         })
 
+        it('should not retain the connection span during socket setup', async () => {
+          const setSocketCh = dc.channel('tracing:ws:server:connect:setSocket')
+          let resolve
+          const promise = new Promise((_resolve) => {
+            resolve = _resolve
+          })
+
+          const handler = () => {
+            resolve(storage('legacy').getStore())
+            setSocketCh.unsubscribe(handler)
+          }
+          setSocketCh.subscribe(handler)
+
+          // Trigger setSocket
+          const newClient = new WebSocket(`ws://localhost:${clientPort}/test`)
+          newClient.on('open', () => newClient.close())
+
+          const store = await promise
+
+          assert.strictEqual(store?.span, undefined,
+            'connection span should not be in the store during setSocket')
+        })
+
         it('should do automatic instrumentation and remove broken handler', () => {
           wsServer.on('connection', (ws) => {
             connectionReceived = true