Repository: honeycombio/beeline-nodejs
Pull request #60 (2018-07-08 17:11:52)

[express] in debug mode, wrap middleware such that we can detect context losses
Pretty simple change considering how scary it might appear:

if the express instrumentation's `debugWrapMiddleware` is true (normal value == `debug.enabled`, but overridable for testing), record the call site that added the middleware and wrap it in one additional layer to detect if we've lost async context.  If async context is lost, a message will be printed to the console (using `debug`) containing the call site, asking them to file an issue.

The context should continue to propagate as normal to each middleware, but there will be holes in custom context / columns from within a middleware that loses async context.


Patch:
diff --git a/lib/instrumentation/express.js b/lib/instrumentation/express.js
index 1f69d3c..11c08c1 100644
--- a/lib/instrumentation/express.js
+++ b/lib/instrumentation/express.js
@@ -1,13 +1,18 @@
 /* eslint-env node */
 const onHeaders = require("on-headers"),
+  shimmer = require("shimmer"),
+  flatten = require("array-flatten"),
   tracker = require("../async_tracker"),
   schema = require("../schema"),
   event = require("../event_api"),
   propagation = require("../propagation"),
+  util = require("../util"),
   path = require("path"),
   pkg = require(path.join(__dirname, "..", "..", "package.json")),
   debug = require("debug")(`${pkg.name}:express`);
 
+const slice = Array.prototype.slice;
+
 // returns the header name/value for the first header with a value in the request.
 const getValueFromHeaders = (req, headers) => {
   let value, header;
@@ -99,6 +104,75 @@ const getUserContext = (userContext, req) => {
   return userEventContext;
 };
 
+function wrapNextMiddleware(middleware, debugWrap) {
+  let debugLocation;
+  if (debugWrap) {
+    debugLocation = util.captureStackTrace(5);
+  }
+
+  return function(req, res, next) {
+    let context = tracker.getTracked();
+    if (!context) {
+      return middleware.apply(this, [req, res, next]);
+    }
+
+    let boundWrappedNext = tracker.bindFunction(function(...args) {
+      return next.apply(this, args);
+    });
+
+    let wrappedNext = boundWrappedNext;
+    if (debugWrap) {
+      wrappedNext = function(...args) {
+        let afterContext = tracker.getTracked();
+        if (!afterContext) {
+          event.askForIssue(
+            "we lost our tracking somewhere in the middleware registered:\n" + debugLocation,
+            debug
+          );
+        }
+        return boundWrappedNext.apply(this, args);
+      };
+    }
+
+    return middleware.apply(this, [req, res, wrappedNext]);
+  };
+}
+
+function wrapErrMiddleware(middleware, debugWrap) {
+  let debugLocation;
+  if (debugWrap) {
+    debugLocation = util.captureStackTrace(5);
+  }
+
+  return function(err, req, res, next) {
+    let context = tracker.getTracked();
+    if (!context) {
+      return middleware.apply(this, [err, req, res, next]);
+    }
+
+    let boundWrappedNext = tracker.bindFunction(function(...args) {
+      return next.apply(this, args);
+    });
+
+    let wrappedNext = boundWrappedNext;
+
+    if (debugWrap) {
+      wrappedNext = function(...args) {
+        let afterContext = tracker.getTracked();
+        if (!afterContext) {
+          event.askForIssue(
+            "we lost our tracking somewhere in the middleware registered:\n" + debugLocation,
+            debug
+          );
+        }
+        return boundWrappedNext.apply(this, args);
+      };
+    }
+
+    return middleware.apply(this, [err, req, res, wrappedNext]);
+  };
+}
+
 const getMagicMiddleware = ({ userContext, traceIdSource, packageVersion }) => (req, res, next) => {
   let traceContext = getTraceContext(traceIdSource, req);
   let ev = event.startRequest(
@@ -195,9 +269,63 @@ let instrumentExpress = function(express, opts = {}) {
 
   let packageVersion = opts.packageVersion;
 
+  let debugWrapMiddleware =
+    typeof opts.debugWrapMiddleware !== "undefined" ? opts.debugWrapMiddleware : debug.enabled;
   const wrapper = function() {
     const app = express();
+    // put our middleware in the chain at the start
     app.use(getMagicMiddleware({ userContext, traceIdSource, packageVersion }));
+
+    // and wrap every other middleware's bind/error functions so we can 1) detect context loss, and 2) warn the user about it.
+    shimmer.wrap(express.Router, "use", function instrumentUse(original) {
+      return function(fn) {
+        // copy the argument handling from express, with a couple of modifications:
+        // 1. we don't care about the path here (we just need to locate/modify the functions)
+        // 2. we wrap all middleware functions.
+
+        // BEGIN_COPIED_FROM_EXPRESS
+        var offset = 0;
+        // express has this line, but we don't need it:
+        // var path = "/";
+
+        // default path to '/'
+        // disambiguate app.use([fn])
+        if (typeof fn !== "function") {
+          var arg = fn;
+
+          while (Array.isArray(arg) && arg.length !== 0) {
+            arg = arg[0];
+          }
+
+          // first arg is the path
+          if (typeof arg !== "function") {
+            offset = 1;
+            // express has this line, but we don't need it:
+            // path = fn;
+          }
+        }
+
+        // express has this line, but we need to wrap each of the values,
+        // so add a .map call:
+        //   var fns = flatten(slice.call(arguments, offset));
+        let fns = flatten(slice.call(arguments, offset)).map(function(fn) {
+          if (fn.length > 3) {
+            // args are (err, req, res, next)
+            return wrapErrMiddleware(fn, debugWrapMiddleware);
+          }
+
+          // args are (req, res, next)
+          return wrapNextMiddleware(fn, debugWrapMiddleware);
+        });
+
+        if (fns.length === 0) {
+          throw new TypeError("app.use() requires a middleware function");
+        }
+        // END_COPIED_FROM_EXPRESS
+
+        return original.apply(this, slice.call(arguments, 0, offset).concat(fns));
+      };
+    });
     return app;
   };
   Object.defineProperties(wrapper, Object.getOwnPropertyDescriptors(express));
diff --git a/lib/util.js b/lib/util.js
new file mode 100644
index 0000000..603d723
--- /dev/null
+++ b/lib/util.js
@@ -0,0 +1,19 @@
+/* eslint-env node */
+exports.captureStackTrace = captureStackTrace;
+function captureStackTrace(skipFrames = 0, limitFrames = 10) {
+  let e = new Error();
+
+  // save off what the current (perhaps app-specified) traceTraceLimit is so we can
+  // capture limitFrames frames.
+  let stackTraceLimit = Error.stackTraceLimit;
+  Error.stackTraceLimit = stackTraceLimit < limitFrames ? limitFrames : stackTraceLimit;
+
+  Error.captureStackTrace(e);
+
+  // reinstate previous limit.
+  Error.stackTraceLimit = stackTraceLimit;
+
+  let frames = e.stack.split("\n");
+  // the +1 here to get rid of the `Error\n` line at the top of the stacktrace.
+  return frames.slice(1 + skipFrames).join("\n");
+}
diff --git a/package-lock.json b/package-lock.json
index 0e99037..a1b12bd 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -531,10 +531,9 @@
       "dev": true
     },
     "array-flatten": {
-      "version": "1.1.1",
-      "resolved": "https://registry.npmjs.org/array-flatten/-/array-flatten-1.1.1.tgz",
-      "integrity": "sha1-ml9pkFGx5wczKPKgCJaLZOopVdI=",
-      "dev": true
+      "version": "2.1.1",
+      "resolved": "https://registry.npmjs.org/array-flatten/-/array-flatten-2.1.1.tgz",
+      "integrity": "sha1-Qmu52oQJDBg42BLIFQryCoMx4pY="
     },
     "array-union": {
       "version": "1.0.2",
@@ -2093,6 +2092,12 @@
         "vary": "~1.1.2"
       },
       "dependencies": {
+        "array-flatten": {
+          "version": "1.1.1",
+          "resolved": "https://registry.npmjs.org/array-flatten/-/array-flatten-1.1.1.tgz",
+          "integrity": "sha1-ml9pkFGx5wczKPKgCJaLZOopVdI=",
+          "dev": true
+        },
         "debug": {
           "version": "2.6.9",
           "resolved": "https://registry.npmjs.org/debug/-/debug-2.6.9.tgz",
diff --git a/package.json b/package.json
index b668f5b..a41f98b 100644
--- a/package.json
+++ b/package.json
@@ -37,6 +37,7 @@
     "supertest": "^3.0.0"
   },
   "dependencies": {
+    "array-flatten": "^2.1.1",
     "debug": "^3.1.0",
     "libhoney": "^1.1.0",
     "on-headers": "^1.0.1",


Test patch:
diff --git a/lib/instrumentation/express.test.js b/lib/instrumentation/express.test.js
index cd339d9..5e9d1e1 100644
--- a/lib/instrumentation/express.test.js
+++ b/lib/instrumentation/express.test.js
@@ -1,10 +1,13 @@
 /* eslint-env node, jest */
 const request = require("supertest"),
   path = require("path"),
+  http = require("http"),
+  net = require("net"),
   instrumentExpress = require("./express"),
   cases = require("jest-in-case"),
   schema = require("../schema"),
   event = require("../event_api"),
+  tracker = require("../async_tracker"),
   pkg = require(path.join(__dirname, "..", "..", "package.json"));
 
 describe("userContext", () => {
@@ -244,3 +247,68 @@ describe("trace id callback", () => {
       });
   });
 });
+
+describe("tracking loss detection", () => {
+  const express = instrumentExpress(require("express"), {
+    debugWrapMiddleware: true,
+  });
+
+  let askForIssue;
+  beforeAll(() => {
+    event.configure({ api: "mock" });
+    askForIssue = event.askForIssue;
+  });
+  afterAll(() => {
+    event.resetForTesting();
+    event.askForIssue = askForIssue;
+  });
+
+  let testRequest, testResponse;
+  let app;
+  let mockAsk;
+  beforeEach(() => {
+    app = express();
+    mockAsk = event.askForIssue = jest.fn();
+    testRequest = new http.IncomingMessage(new net.Socket());
+    testRequest.url = "/test/route";
+    testRequest.method = "GET";
+
+    testResponse = new http.ServerResponse(testRequest);
+  });
+  test("we detect it in normal middleware next callback", done => {
+    app.use((req, res, next) => {
+      // we simulate an async operation that loses context by explicitly losing it before calling next.
+      tracker.setTracked(undefined);
+      next();
+    });
+
+    app.use((_req, _res, next) => {
+      expect(mockAsk.mock.calls.length).toBe(1);
+      expect(mockAsk.mock.calls[0][0]).toMatch(/we lost our tracking/);
+      next();
+    });
+
+    app.handle(testRequest, testResponse, done);
+  });
+
+  test("we detect it in error middleware next callback", done => {
+    app.use(() => {
+      // generate an error
+      throw new Error("something bad happened here");
+    });
+
+    app.use((err, _req, _res, next) => {
+      // we simulate an async operation that loses context by explicitly losing it before calling err.
+      tracker.setTracked(undefined);
+      next(err);
+    });
+
+    app.use((err, _req, _res, next) => {
+      expect(mockAsk.mock.calls.length).toBe(1);
+      expect(mockAsk.mock.calls[0][0]).toMatch(/we lost our tracking/);
+      next();
+    });
+
+    app.handle(testRequest, testResponse, done);
+  });
+});