fix: ensure correct run context for 'ws' instrumentation

Refs: #2430

examples/trace-ws.js

@@ -0,0 +1,46 @@
+#!/usr/bin/env node --unhandled-rejections=strict
+
+// A small example showing Elastic APM tracing of the 'ws' package.
+//
+// Currently Elastic APM will create a span for `ws.send()` calls.
+// For tracing spans to be created, there must be an active transaction.
+// Typically, a transaction is automatically started for incoming HTTP
+// requests to a Node.js server. However, because this script is not running
+// an HTTP server, we manually start transactions in our websocket server and
+// client. More details at:
+// https://www.elastic.co/guide/en/apm/agent/nodejs/current/custom-transactions.html
+
+const apm = require('../').start({ // elastic-apm-node
+  serviceName: 'example-trace-ws'
+})
+
+var WebSocket = require('ws')
+var PORT = 4567
+
+// Server
+const wss = new WebSocket.Server({ port: PORT })
+wss.on('connection', function (ws) {
+  ws.on('message', function (message) {
+    console.log('server on "message": %j', message)
+    const tserver = apm.startTransaction('tserver')
+    ws.send('pong')
+    tserver.end()
+  })
+})
+
+// Client
+const ws = new WebSocket('ws://localhost:' + PORT)
+ws.on('open', function () {
+  const tclient = apm.startTransaction('tclient')
+  console.log('client: send "ping"')
+  ws.send('ping', function () {
+    console.log('client: "ping" has been sent')
+  })
+  console.log('client: send "ring"')
+  ws.send('ring')
+  tclient.end()
+})
+ws.on('message', function (message) {
+  console.log('client on "message": %j', message)
+  wss.close()
+})

lib/instrumentation/modules/ws.js

@@ -11,35 +11,38 @@ module.exports = function (ws, agent, { version, enabled }) {
     return ws
   }
 
+  const ins = agent._instrumentation
+
   agent.logger.debug('shimming ws.prototype.send function')
   shimmer.wrap(ws.prototype, 'send', wrapSend)
 
   return ws
 
   function wrapSend (orig) {
     return function wrappedSend () {
-      var span = agent.startSpan('Send WebSocket Message', 'websocket', 'send')
-      var id = span && span.transaction.id
-
-      agent.logger.debug('intercepted call to ws.prototype.send %o', { id: id })
-
-      if (!span) return orig.apply(this, arguments)
+      agent.logger.debug('intercepted call to ws.prototype.send')
+      const span = ins.createSpan('Send WebSocket Message', 'websocket', 'send')
+      if (!span) {
+        return orig.apply(this, arguments)
+      }
 
-      var args = [].slice.call(arguments)
-      var cb = args[args.length - 1]
+      const args = Array.prototype.slice.call(arguments)
+      let cb = args[args.length - 1]
+      const onDone = function () {
+        span.end()
+        if (cb) {
+          cb.apply(this, arguments)
+        }
+      }
       if (typeof cb === 'function') {
-        args[args.length - 1] = done
+        args[args.length - 1] = ins.bindFunction(onDone)
       } else {
         cb = null
-        args.push(done)
+        args.push(ins.bindFunction(onDone))
       }
 
-      return orig.apply(this, args)
-
-      function done () {
-        span.end()
-        if (cb) cb.apply(this, arguments)
-      }
+      const spanRunContext = ins.currRunContext().enterSpan(span)
+      return ins.withRunContext(spanRunContext, orig, this, ...args)
     }
   }
 }

test/instrumentation/modules/ws.test.js

@@ -1,8 +1,7 @@
 'use strict'
 
 var agent = require('../../..').start({
-  serviceName: 'test',
-  secretToken: 'test',
+  serviceName: 'test-ws',
   captureExceptions: false,
   metricsInterval: 0,
   centralConfig: false
@@ -32,8 +31,10 @@ test('ws.send', function (t) {
   ws.on('open', function () {
     agent.startTransaction('foo', 'websocket')
     ws.send('ping', function () {
+      t.ok(agent.currentSpan === null, 'websocket span should not be the currentSpan in user callback')
       agent.endTransaction()
     })
+    t.ok(agent.currentSpan === null, 'websocket span should not spill into user code')
   })
 
   ws.on('message', function (message) {