Repository: honeycombio/beeline-nodejs
Pull request #146 (2019-06-21 00:36:08)

dataset propagation, fix parent span ids on async spans
The go beeline will put a dataset k/v in the marshaled trace context.  If it's there, use it.  This will result in distributed traces being sent entirely to the same dataset (chosen by the root service.)

The async span problem is a little difficult to explain: The node beeline creates async spans by essentially forking/execing the trace context, putting the new span as the root span of the new trace context.  everything else is copied over, so the trace id remains the same, etc.  The parent id is also grabbed so the root async span will be shown as a child of the current span at the time the async span was created.

The problem comes in when that root async span finishes, and another span starts.  The stack is now empty, so there's no parentId to use.  The fix is to also store the parent id on the new trace context, and use it from there if the stack is empty.

Lastly, I noticed that in the express instrumentation when unmarshaling trace context, we were using the customContext api to add the context from the calling service.  This results in `app.` being prepended, so we were seeing fields like `app.app.user_id`, `app.app.whatever`.


Patch:
diff --git a/lib/api/index.js b/lib/api/index.js
index 209d216..c1e8b29 100644
--- a/lib/api/index.js
+++ b/lib/api/index.js
@@ -45,14 +45,14 @@ module.exports = {
   marshalTraceContext: propagation.marshalTraceContext,
   unmarshalTraceContext: propagation.unmarshalTraceContext,
 
-  startTrace(metadataContext, withTraceId, withParentSpanId) {
-    return apiImpl.startTrace(metadataContext, withTraceId, withParentSpanId);
+  startTrace(metadataContext, withTraceId, withParentSpanId, withDataset) {
+    return apiImpl.startTrace(metadataContext, withTraceId, withParentSpanId, withDataset);
   },
   finishTrace(trace) {
     return apiImpl.finishTrace(trace);
   },
-  withTrace(metadataContext, fn, withTraceId, withParentSpanId) {
-    const trace = apiImpl.startTrace(metadataContext, withTraceId, withParentSpanId);
+  withTrace(metadataContext, fn, withTraceId, withParentSpanId, withDataset) {
+    const trace = apiImpl.startTrace(metadataContext, withTraceId, withParentSpanId, withDataset);
     try {
       return fn();
     } finally {
diff --git a/lib/api/libhoney.js b/lib/api/libhoney.js
index a65e87e..ddb911c 100644
--- a/lib/api/libhoney.js
+++ b/lib/api/libhoney.js
@@ -69,13 +69,15 @@ module.exports = class LibhoneyEventAPI {
       debug(`using proxy ${proxy}`);
     }
 
+    this.defaultDataset = opts.dataset || process.env["HONEYCOMB_DATASET"] || defaultName;
+
     this.honey = new libhoney(
       Object.assign(
         {
           apiHost,
           proxy,
+          dataset: this.defaultDataset,
           writeKey: process.env["HONEYCOMB_WRITEKEY"],
-          dataset: process.env["HONEYCOMB_DATASET"] || defaultName,
           userAgentAddition: `honeycomb-beeline/${pkg.version}`,
         },
         getFilteredOptions(opts)
@@ -87,13 +89,14 @@ module.exports = class LibhoneyEventAPI {
     });
   }
 
-  startTrace(metadataContext, traceId, parentSpanId) {
+  startTrace(metadataContext, traceId, parentSpanId, dataset) {
     const id = traceId || uuidv4();
 
     tracker.setTracked({
       id,
       customContext: {},
       stack: [],
+      dataset: dataset || this.defaultDataset,
     });
     return this.startSpan(metadataContext, undefined, parentSpanId);
   }
@@ -111,6 +114,10 @@ module.exports = class LibhoneyEventAPI {
       }
       parentId = context.stack[context.stack.length - 1][schema.TRACE_SPAN_ID];
     }
+    if (!parentId) {
+      parentId = context.parentId;
+    }
+
     const eventPayload = Object.assign({}, metadataContext, {
       [schema.TRACE_ID]: context.id,
       [schema.TRACE_SPAN_ID]: spanId,
@@ -131,6 +138,9 @@ module.exports = class LibhoneyEventAPI {
     if (parentContext.stack.length > 0) {
       parentId = parentContext.stack[parentContext.stack.length - 1][schema.TRACE_SPAN_ID];
     }
+    if (!parentId) {
+      parentId = parentContext.parentId;
+    }
 
     const eventPayload = Object.assign({}, metadataContext, {
       [schema.TRACE_ID]: parentContext.id,
@@ -144,7 +154,9 @@ module.exports = class LibhoneyEventAPI {
 
     let newContext = {
       id: parentContext.id,
+      parentId: parentId,
       customContext: parentContext.customContext,
+      dataset: parentContext.dataset,
       stack: [eventPayload],
     };
 
@@ -211,6 +223,7 @@ module.exports = class LibhoneyEventAPI {
       const active_instrumentations = instrumentation.activeInstrumentations();
       const active_instrumentation_count = active_instrumentations.length;
       const ev = this.honey.newEvent();
+      ev.dataset = context.dataset;
       ev.timestamp = new Date(startTime);
       ev.add(payload);
       ev.add(context.customContext);
diff --git a/lib/instrumentation/express.js b/lib/instrumentation/express.js
index 1c34d3a..b35cd52 100644
--- a/lib/instrumentation/express.js
+++ b/lib/instrumentation/express.js
@@ -191,13 +191,12 @@ const getMagicMiddleware = ({ userContext, traceIdSource, packageVersion }) => (
       "request.xhr": req.xhr,
     },
     traceContext.traceId,
-    traceContext.parentSpanId
+    traceContext.parentSpanId,
+    traceContext.dataset
   );
 
   if (traceContext.customContext) {
-    Object.keys(traceContext.customContext).forEach(k =>
-      api.customContext.add(k, traceContext.customContext[k])
-    );
+    api.addContext(traceContext.customContext);
   }
 
   if (!trace) {
diff --git a/lib/propagation.js b/lib/propagation.js
index cfc8bae..a182131 100644
--- a/lib/propagation.js
+++ b/lib/propagation.js
@@ -18,9 +18,10 @@ exports.VERSION = VERSION;
 // PAYLOAD is a list of comma-separated params (k=v pairs), with no spaces.  recognized
 // keys + value types:
 //
-//  trace_id=${traceId}    - traceId is an opaque ascii string which shall not include ','
-//  parent_id=${spanId}    - spanId is an opaque ascii string which shall not include ','
-//  context=${contextBlob} - contextBlob is a base64 encoded json object.
+//  trace_id=${traceId}     - traceId is an opaque ascii string which shall not include ','
+//  parent_id=${spanId}     - spanId is an opaque ascii string which shall not include ','
+//  dataset=${datasetId}   - datasetId is the slug for the honeycomb dataset to which downstream spans should be sent; shall not include ','
+//  context=${contextBlob}  - contextBlob is a base64 encoded json object.
 //
 // ex: X-Honeycomb-Trace: 1;trace_id=weofijwoeifj,parent_id=owefjoweifj,context=SGVsbG8gV29ybGQ=
 
@@ -29,8 +30,14 @@ function marshalTraceContextv1(traceContext) {
   let traceId = traceContext.id;
   let spanId = traceContext.stack[traceContext.stack.length - 1][schema.TRACE_SPAN_ID];
   let contextToSend = traceContext.customContext || {};
+  let dataset = traceContext.dataset;
 
-  return `${VERSION};trace_id=${traceId},parent_id=${spanId},context=${Buffer.from(
+  let datasetClause = "";
+  if (dataset) {
+    datasetClause = `dataset=${encodeURIComponent(dataset)},`;
+  }
+
+  return `${VERSION};trace_id=${traceId},parent_id=${spanId},${datasetClause}context=${Buffer.from(
     JSON.stringify(contextToSend)
   ).toString("base64")}`;
 }
@@ -49,7 +56,7 @@ function unmarshalTraceContext(contextStr) {
 function unmarshalTraceContextv1(payload) {
   let clauses = payload.split(",");
 
-  let traceId, parentSpanId, contextb64;
+  let traceId, parentSpanId, dataset, contextb64;
 
   clauses.forEach(cl => {
     let [k, v] = cl.split("=");
@@ -60,6 +67,9 @@ function unmarshalTraceContextv1(payload) {
       case "parent_id":
         parentSpanId = v;
         break;
+      case "dataset":
+        dataset = decodeURIComponent(v);
+        break;
       case "context":
         contextb64 = v;
         break;
@@ -89,5 +99,5 @@ function unmarshalTraceContextv1(payload) {
     }
   }
 
-  return { traceId, parentSpanId, customContext };
+  return { traceId, parentSpanId, customContext, dataset };
 }


Test patch:
diff --git a/lib/propagation.test.js b/lib/propagation.test.js
index d8367ed..21300c4 100644
--- a/lib/propagation.test.js
+++ b/lib/propagation.test.js
@@ -6,6 +6,7 @@ const propagation = require("./propagation"),
 // this context structure is the same as what the libhoney event api implementation generate.
 let testContext = {
   id: "abcdef123456",
+  dataset: "testDataset",
   stack: [{ [schema.TRACE_SPAN_ID]: "0102030405" }],
   customContext: {
     userID: 1,
@@ -72,6 +73,7 @@ describe("roundtrip", () => {
     expect(propagation.unmarshalTraceContext(contextStr)).toEqual({
       traceId: "abcdef123456",
       parentSpanId: "0102030405",
+      dataset: "testDataset",
       customContext: {
         userID: 1,
         errorMsg: "failed to sign on",