Repository: honeycombio/beeline-nodejs (js)

Issue:
Trace concurrent actions
Beeline assumes that, within a trace, all actions are sequential, and keeps an internal stack of spans to generate parent-child relationships. This assumption breaks down with e.g. concurrent database queries. Is it possible to trace those events and preserve the correct span containment?

Pull request:
Support async/concurrent spans
Fixes #93 

Add `startAsyncSpan` to the API and only to the libhoney implementation for the time being (i.e. no tests).  It doesn't work the same as `startSpan`, unfortunately.  This is due to my understanding of how async context gets propagated by `async-hooks` - I may be wrong, so I'll keep digging there.

the signature is `startAsyncSpan(metadataContext, spanFn)`.  `metadataContext` is the same as with `startSpan`, and `spanFn` looks similar to `withSpan` but it's not.  it's called with a single argument - the asyncSpan itself.

From my motivating example in #93, here's `slowSpan` when using normal spans:

```
function slowSpan() {
  let asyncSpan = beeline.startSpan({ name: "slow" }); // should be startAsyncSpan here
  return new Promise(resolve => {
    setTimeout(() => {
      beeline.withSpan({ name: "slow-child" }, () => {});
      beeline.finishSpan(asyncSpan);
      resolve();
    }, 5000);
  });
}
```

and here's the converted form:

```
function slowSpan() {
  return beeline.startAsyncSpan(
    { name: "slow" },
    asyncSpan =>
      new Promise(resolve => {
        console.log("in slow promise");
        setTimeout(() => {
          beeline.withSpan({ name: "slow-child" }, () => {});
          beeline.finishSpan(asyncSpan);
          resolve();
        }, 5000);
      })
  );
}
```

The reason for passing the `asyncSpan` into this function is that the `finishSpan` call _must_ happen inside the async context created and installed around this inner function.  we create a new tracking context in `startAsyncSpan`, give it the same trace id, and install the async span as the root span in the new context's stack.  doing it this way means that other non-async child spans will be pushed onto the correct stacks and get the correct parent span id.


Patch:
diff --git a/docs/API.md b/docs/API.md
index f2c43e6..1418334 100644
--- a/docs/API.md
+++ b/docs/API.md
@@ -123,7 +123,36 @@ let span = beeline.startSpan({
 });
 fs.writeFile(filePath, fileContents, err => {
   beeline.customContext.add("fileError", err.toString());
-  beeline.finishSpan(trace);
+  beeline.finishSpan(span);
+});
+```
+
+#### startAsyncSpan
+
+_(unstable API - possibly changing in the future)_
+
+```javascript
+beeline.startAsyncSpan(metadataContext, spanFn);
+```
+
+Starts a new async span within an existing trace. This new span is added as a child of the current span, and recorded as the current
+span within the context of calling `spanFn`. Outside of `spanFn` the current span is unchanged. The newly created span is passed as
+the only argument to `spanFn` (to be used in `finishSpan` below.)
+
+Async spans are the norm within the beeline-packaged instrumentation, and you should use them for all asynchronous operations.
+
+example:
+
+```javascript
+beeline.startAsyncSpan({
+    task: "writing a file",
+    filePath,
+}, span => {
+       fs.writeFile(filePath, fileContents, err => {
+           beeline.customContext.add("fileError", err.toString());
+           beeline.finishSpan(span);
+       });
+   }
 });
 ```
 
diff --git a/lib/api/index.js b/lib/api/index.js
index 988710c..9413b51 100644
--- a/lib/api/index.js
+++ b/lib/api/index.js
@@ -75,6 +75,10 @@ module.exports = {
     }
   },
 
+  startAsyncSpan(metadataContext, spanFn) {
+    return apiImpl.startAsyncSpan(metadataContext, spanFn);
+  },
+
   startTimer(name) {
     return {
       name,
diff --git a/lib/api/libhoney.js b/lib/api/libhoney.js
index 4f51523..fd246d8 100644
--- a/lib/api/libhoney.js
+++ b/lib/api/libhoney.js
@@ -109,7 +109,35 @@ module.exports = class LibhoneyEventAPI {
     return eventPayload;
   }
 
+  startAsyncSpan(metadataContext, spanFn) {
+    let parentId;
+    let spanId = uuidv4();
+    let context = tracker.getTracked();
+    if (context.stack.length > 0) {
+      parentId = context.stack[context.stack.length - 1][schema.TRACE_SPAN_ID];
+    }
+
+    const eventPayload = Object.assign({}, metadataContext, {
+      [schema.TRACE_ID]: context.id,
+      [schema.TRACE_SPAN_ID]: spanId,
+      startTime: Date.now(),
+      startTimeHR: process.hrtime(),
+    });
+    if (parentId) {
+      eventPayload[schema.TRACE_PARENT_ID] = parentId;
+    }
+
+    let newContext = {
+      id: context.id,
+      customContext: {},
+      stack: [eventPayload],
+    };
+
+    return tracker.callWithContext(() => spanFn(eventPayload), newContext);
+  }
+
   finishSpan(ev, rollup) {
+    debug(`finishing span ${JSON.stringify(ev)}`);
     const context = tracker.getTracked();
     if (!context) {
       // valid, since we can end up in our instrumentation outside of requests we're tracking
diff --git a/lib/api/mock.js b/lib/api/mock.js
index 511e657..4f09f67 100644
--- a/lib/api/mock.js
+++ b/lib/api/mock.js
@@ -39,6 +39,33 @@ module.exports = class MockEventAPI {
     context.stack.push(eventPayload);
     return eventPayload;
   }
+  startAsyncSpan(metadataContext, spanFn) {
+    let parentId;
+    let context = tracker.getTracked();
+    let spanId = context.spanId + 1000;
+    if (context.stack.length > 0) {
+      parentId = context.stack[context.stack.length - 1][schema.TRACE_SPAN_ID];
+    }
+
+    const eventPayload = Object.assign({}, metadataContext, {
+      [schema.TRACE_ID]: context.id,
+      [schema.TRACE_SPAN_ID]: ++spanId,
+      startTime: Date.now(),
+      startTimeHR: process.hrtime(),
+    });
+    if (parentId) {
+      eventPayload[schema.TRACE_PARENT_ID] = parentId;
+    }
+
+    let newContext = {
+      id: context.id,
+      spanId,
+      stack: [eventPayload],
+    };
+
+    return tracker.callWithContext(() => spanFn(eventPayload), newContext);
+  }
+
   finishSpan(ev, _rollup) {
     Object.assign(ev, this.customContext, {
       [schema.DURATION_MS]: 0,
diff --git a/lib/async_tracker.js b/lib/async_tracker.js
index ae573cc..6ed4087 100644
--- a/lib/async_tracker.js
+++ b/lib/async_tracker.js
@@ -51,6 +51,18 @@ class AsyncTracker {
     };
   }
 
+  // XXX(toshok) this feels wrong, but maybe not?
+  callWithContext(fn, context) {
+    let previousContext = this.getTracked();
+
+    tracker.setTracked(context);
+    try {
+      return fn();
+    } finally {
+      tracker.setTracked(previousContext);
+    }
+  }
+
   // below is the portion of the async_hooks api we need.  they shouldn't be called directly
   // from user code.  They also aren't async safe - if any async code is added to them (like console.log)
   // we'll blow the stack.
diff --git a/lib/instrumentation/child_process.js b/lib/instrumentation/child_process.js
index 5044427..e283242 100644
--- a/lib/instrumentation/child_process.js
+++ b/lib/instrumentation/child_process.js
@@ -10,27 +10,31 @@ function wrapExecLike(name, packageVersion) {
         return original.apply(this, arguments);
       }
 
-      let span = api.startSpan({
-        [schema.EVENT_TYPE]: "child_process",
-        [schema.PACKAGE_VERSION]: packageVersion,
-        [schema.TRACE_SPAN_NAME]: name,
-        "exec.file": file,
-      });
-      if (Array.isArray(args)) {
-        api.addContext({ "exec.args": args });
-      }
+      api.startAsyncSpan(
+        {
+          [schema.EVENT_TYPE]: "child_process",
+          [schema.PACKAGE_VERSION]: packageVersion,
+          [schema.TRACE_SPAN_NAME]: name,
+          "exec.file": file,
+        },
+        span => {
+          if (Array.isArray(args)) {
+            api.addContext({ "exec.args": args });
+          }
 
-      let argsArray = Array.from(arguments);
-      argsArray.slice(1).forEach((arg, idx) => {
-        if (typeof arg !== "function") {
-          return;
+          let argsArray = Array.from(arguments);
+          argsArray.slice(1).forEach((arg, idx) => {
+            if (typeof arg !== "function") {
+              return;
+            }
+            argsArray[idx + 1] = function(error, stdout, stderr) {
+              api.finishSpan(span, name);
+              arg(error, stdout, stderr);
+            };
+          });
+          return original.apply(this, argsArray);
         }
-        argsArray[idx + 1] = function(error, stdout, stderr) {
-          api.finishSpan(span, name);
-          arg(error, stdout, stderr);
-        };
-      });
-      return original.apply(this, argsArray);
+      );
     };
   };
 }
diff --git a/lib/instrumentation/http.js b/lib/instrumentation/http.js
index 8565657..f89bc27 100644
--- a/lib/instrumentation/http.js
+++ b/lib/instrumentation/http.js
@@ -37,31 +37,34 @@ let instrumentHTTP = (http, opts = {}) => {
         options.hostname = "localhost";
       }
 
-      let span = api.startSpan({
-        [schema.EVENT_TYPE]: "http",
-        [schema.PACKAGE_VERSION]: opts.packageVersion,
-        [schema.TRACE_SPAN_NAME]: options.method || "GET",
-        url: url.format(options),
-      });
+      return api.startAsyncSpan(
+        {
+          [schema.EVENT_TYPE]: "http",
+          [schema.PACKAGE_VERSION]: opts.packageVersion,
+          [schema.TRACE_SPAN_NAME]: options.method || "GET",
+          url: url.format(options),
+        },
+        span => {
+          let wrapped_cb = function(res) {
+            api.finishSpan(span, "request");
+            if (cb) {
+              return cb.apply(this, [res]);
+            }
+          };
 
-      let wrapped_cb = function(res) {
-        api.finishSpan(span, "request");
-        if (cb) {
-          return cb.apply(this, [res]);
-        }
-      };
-
-      // shallow clone options.headers, adding our tracing headers in
-      options.headers = Object.assign({}, options.headers, {
-        [api.TRACE_HTTP_HEADER]: api.marshalTraceContext(api.getTraceContext()),
-      });
+          // shallow clone options.headers, adding our tracing headers in
+          options.headers = Object.assign({}, options.headers, {
+            [api.TRACE_HTTP_HEADER]: api.marshalTraceContext(api.getTraceContext()),
+          });
 
-      if (cb) {
-        return original.apply(this, [options, wrapped_cb]);
-      } else {
-        // the user didn't specify a callback, add it as a "response" handler ourselves
-        return original.apply(this, [options]).on("response", wrapped_cb);
-      }
+          if (cb) {
+            return original.apply(this, [options, wrapped_cb]);
+          } else {
+            // the user didn't specify a callback, add it as a "response" handler ourselves
+            return original.apply(this, [options]).on("response", wrapped_cb);
+          }
+        }
+      );
     };
   });
 
diff --git a/lib/instrumentation/https.js b/lib/instrumentation/https.js
index 5f8002f..0e8515c 100644
--- a/lib/instrumentation/https.js
+++ b/lib/instrumentation/https.js
@@ -37,31 +37,34 @@ let instrumentHTTPS = (https, opts = {}) => {
         options.hostname = "localhost";
       }
 
-      let span = api.startSpan({
-        [schema.EVENT_TYPE]: "https",
-        [schema.PACKAGE_VERSION]: opts.packageVersion,
-        [schema.TRACE_SPAN_NAME]: options.method || "GET",
-        url: url.format(options),
-      });
+      return api.startAsyncSpan(
+        {
+          [schema.EVENT_TYPE]: "https",
+          [schema.PACKAGE_VERSION]: opts.packageVersion,
+          [schema.TRACE_SPAN_NAME]: options.method || "GET",
+          url: url.format(options),
+        },
+        span => {
+          let wrapped_cb = function(res) {
+            api.finishSpan(span, "request");
+            if (cb) {
+              return cb.apply(this, [res]);
+            }
+          };
 
-      let wrapped_cb = function(res) {
-        api.finishSpan(span, "request");
-        if (cb) {
-          return cb.apply(this, [res]);
-        }
-      };
-
-      // shallow clone options.headers, adding our tracing headers in
-      options.headers = Object.assign({}, options.headers, {
-        [api.TRACE_HTTP_HEADER]: api.marshalTraceContext(api.getTraceContext()),
-      });
+          // shallow clone options.headers, adding our tracing headers in
+          options.headers = Object.assign({}, options.headers, {
+            [api.TRACE_HTTP_HEADER]: api.marshalTraceContext(api.getTraceContext()),
+          });
 
-      if (cb) {
-        return original.apply(this, [options, wrapped_cb]);
-      } else {
-        // the user didn't specify a callback, add it as a "response" handler ourselves
-        return original.apply(this, [options]).on("response", wrapped_cb);
-      }
+          if (cb) {
+            return original.apply(this, [options, wrapped_cb]);
+          } else {
+            // the user didn't specify a callback, add it as a "response" handler ourselves
+            return original.apply(this, [options]).on("response", wrapped_cb);
+          }
+        }
+      );
     };
   });
 
diff --git a/lib/instrumentation/mongodb.js b/lib/instrumentation/mongodb.js
index 1cb4206..2cc9548 100644
--- a/lib/instrumentation/mongodb.js
+++ b/lib/instrumentation/mongodb.js
@@ -41,46 +41,49 @@ function instrumentMongodb(mongodb, opts = {}) {
 
         let eventName = `collection.${name}`;
 
-        let span = api.startSpan({
-          [schema.EVENT_TYPE]: "mongodb",
-          [schema.PACKAGE_VERSION]: opts.packageVersion,
-          [schema.TRACE_SPAN_NAME]: eventName,
-        });
-
-        api.addContext(prefixKeys("db.options", options));
-        if (additionalContext) {
-          api.addContext(prefixKeys("db", additionalContext(...populatedArgs)));
-        }
+        api.startAsyncSpan(
+          {
+            [schema.EVENT_TYPE]: "mongodb",
+            [schema.PACKAGE_VERSION]: opts.packageVersion,
+            [schema.TRACE_SPAN_NAME]: eventName,
+          },
+          span => {
+            api.addContext(prefixKeys("db.options", options));
+            if (additionalContext) {
+              api.addContext(prefixKeys("db", additionalContext(...populatedArgs)));
+            }
 
-        if (callback) {
-          let wrapped_cb = api.bindFunctionToTrace(function(...cb_args) {
-            api.finishSpan(span, eventName);
-            return callback(...cb_args);
-          });
+            if (callback) {
+              let wrapped_cb = api.bindFunctionToTrace(function(...cb_args) {
+                api.finishSpan(span, eventName);
+                return callback(...cb_args);
+              });
 
-          return original.apply(this, [...populatedArgs, wrapped_cb]);
-        }
+              return original.apply(this, [...populatedArgs, wrapped_cb]);
+            }
 
-        let rv = original.apply(this, populatedArgs);
-        if (typeof rv.then === "function") {
-          // it's a promise?
-          return rv.then(
-            () => {
-              api.finishSpan(span, eventName);
-            },
-            err => {
-              if (err) {
-                api.addContext({ error: err.toString() });
-              }
-              api.finishSpan(span, eventName);
+            let rv = original.apply(this, populatedArgs);
+            if (typeof rv.then === "function") {
+              // it's a promise?
+              return rv.then(
+                () => {
+                  api.finishSpan(span, eventName);
+                },
+                err => {
+                  if (err) {
+                    api.addContext({ error: err.toString() });
+                  }
+                  api.finishSpan(span, eventName);
+                }
+              );
             }
-          );
-        }
 
-        // otherwise it's a cursor.  we don't do a good job of instrumenting
-        // through cursor usage, but regardless this span is done.
-        api.finishSpan(span, eventName);
-        return rv;
+            // otherwise it's a cursor.  we don't do a good job of instrumenting
+            // through cursor usage, but regardless this span is done.
+            api.finishSpan(span, eventName);
+            return rv;
+          }
+        );
       };
     });
   }
diff --git a/lib/instrumentation/mysql2.js b/lib/instrumentation/mysql2.js
index 55a3e18..fb9c3b3 100644
--- a/lib/instrumentation/mysql2.js
+++ b/lib/instrumentation/mysql2.js
@@ -13,20 +13,23 @@ let instrumentConnection = function(conn, packageVersion) {
         return original.apply(this, args);
       }
 
-      let span = api.startSpan({
-        [schema.EVENT_TYPE]: "mysql2",
-        [schema.PACKAGE_VERSION]: packageVersion,
-        [schema.TRACE_SPAN_NAME]: "query",
-        "db.query": args[0].sql,
-      });
+      api.startAsyncSpan(
+        {
+          [schema.EVENT_TYPE]: "mysql2",
+          [schema.PACKAGE_VERSION]: packageVersion,
+          [schema.TRACE_SPAN_NAME]: "query",
+          "db.query": args[0].sql,
+        },
+        span => {
+          let cb = args[args.length - 1];
+          let wrapped_cb = function(...cb_args) {
+            api.finishSpan(span, "query");
+            return cb(...cb_args);
+          };
 
-      let cb = args[args.length - 1];
-      let wrapped_cb = function(...cb_args) {
-        api.finishSpan(span, "query");
-        return cb(...cb_args);
-      };
-
-      return original.apply(this, args.slice(0, -1).concat(wrapped_cb));
+          return original.apply(this, args.slice(0, -1).concat(wrapped_cb));
+        }
+      );
     };
   });
 
@@ -40,22 +43,25 @@ let instrumentConnection = function(conn, packageVersion) {
         return original.apply(this, args);
       }
 
-      let span = api.startSpan({
-        [schema.EVENT_TYPE]: "mysql2",
-        [schema.PACKAGE_VERSION]: packageVersion,
-        [schema.TRACE_SPAN_NAME]: "query",
-        "db.query": args[0].sql,
-      });
-
-      let cb = args[args.length - 1];
-      // XXX(toshok) this bindFunction shouldn't be necessary, but we aren't
-      // finishing the event properly without it.
-      let wrapped_cb = api.bindFunctionToTrace(function(...cb_args) {
-        api.finishSpan(span, "query");
-        return cb(...cb_args);
-      });
+      api.startAsyncSpan(
+        {
+          [schema.EVENT_TYPE]: "mysql2",
+          [schema.PACKAGE_VERSION]: packageVersion,
+          [schema.TRACE_SPAN_NAME]: "query",
+          "db.query": args[0].sql,
+        },
+        span => {
+          let cb = args[args.length - 1];
+          // XXX(toshok) this bindFunction shouldn't be necessary, but we aren't
+          // finishing the event properly without it.
+          let wrapped_cb = api.bindFunctionToTrace(function(...cb_args) {
+            api.finishSpan(span, "query");
+            return cb(...cb_args);
+          });
 
-      return original.apply(this, args.slice(0, -1).concat(wrapped_cb));
+          return original.apply(this, args.slice(0, -1).concat(wrapped_cb));
+        }
+      );
     };
   });
   return conn;
diff --git a/package.json b/package.json
index c79eb32..d270e2f 100644
--- a/package.json
+++ b/package.json
@@ -48,6 +48,7 @@
   },
   "jest": {
     "verbose": true,
-    "timers": "real"
+    "timers": "real",
+    "testURL": "http://localhost/"
   }
 }


Test patch:
diff --git a/lib/instrumentation/child_process.test.js b/lib/instrumentation/child_process.test.js
index 4e53986..7a6190f 100644
--- a/lib/instrumentation/child_process.test.js
+++ b/lib/instrumentation/child_process.test.js
@@ -5,18 +5,20 @@ const child_process = require("child_process"),
   api = require("../api"),
   schema = require("../schema");
 
+function newMockContext() {
+  return { id: 0, spanId: 50000, stack: [] };
+}
+
 instrumentChildProcess(child_process);
 
 beforeAll(() => api.configure({ impl: "mock" }));
 
 test("exec 'echo hi' works", done => {
-  let context = { stack: [] };
-  tracker.setTracked(context);
+  tracker.setTracked(newMockContext());
 
   child_process.exec("echo hi", (error, stdout) => {
     expect(error).toBeNull();
     expect(stdout).toBe("hi\n");
-    expect(tracker.getTracked()).toBe(context);
     expect(api._apiForTesting().sentEvents.length).toBe(1);
 
     let ev = api._apiForTesting().sentEvents[0];
@@ -29,13 +31,11 @@ test("exec 'echo hi' works", done => {
 });
 
 test("execFile '/bin/echo hi' works", done => {
-  let context = { stack: [] };
-  tracker.setTracked(context);
+  tracker.setTracked(newMockContext());
 
   child_process.execFile("echo", ["hi"], (error, stdout) => {
     expect(error).toBeNull();
     expect(stdout).toBe("hi\n");
-    expect(tracker.getTracked()).toBe(context);
     expect(api._apiForTesting().sentEvents.length).toBe(1);
 
     let ev = api._apiForTesting().sentEvents[0];
diff --git a/lib/instrumentation/http.test.js b/lib/instrumentation/http.test.js
index 91726af..0d9366f 100644
--- a/lib/instrumentation/http.test.js
+++ b/lib/instrumentation/http.test.js
@@ -38,7 +38,7 @@ test("url as a string", done => {
 
   http.get("http://localhost:9009", res => {
     expect(res.headers[api.TRACE_HTTP_HEADER.toLowerCase()]).toBe(
-      "1;trace_id=0,parent_id=50001,context=e30="
+      "1;trace_id=0,parent_id=51001,context=e30="
     );
     expect(api._apiForTesting().sentEvents).toMatchObject([
       {
@@ -61,7 +61,7 @@ test("url as options", done => {
     },
     res => {
       expect(res.headers[api.TRACE_HTTP_HEADER.toLowerCase()]).toBe(
-        "1;trace_id=0,parent_id=50001,context=e30="
+        "1;trace_id=0,parent_id=51001,context=e30="
       );
       expect(api._apiForTesting().sentEvents).toEqual([
         expect.objectContaining({