Repository: honeycombio/beeline-nodejs
Pull request #226 (2020-07-18 01:01:35)

propagate trace-level fields unchanged, reserving prefixing behavior for api-use only
Let me know what you think of this, figured I'd take a quick stab while I was thinking of it. For unmarshaling trace headers it really makes more sense to consider anything in the "context" field of a beeline header to be a trace level field and propagate it untouched. In the api, however, let's keep the prefixing so we can identify fields added at beeline-instrumented runtime.

~~I didn't want to add an additional arg and logic to the existing `api.addTraceContext()`~~ after feedback we're totally just going to add another arg to contain any trace level fields parsed from a header, this will enable trace-level fields to pass through untouched, ~~and a new api function seemed unnecessary, so I added a new traceUtil that adds fields to the payload without prefixing. This also allows the api (and functions it wraps) to remain totally untouched.~~

---

`traceContext` -> `spanContext` change is for clarity. A span context is more in line with what we're referring to when we talk about a parsed header. Trace context continues to make sense though for fetching the async hook. 
  - `api.getTraceContext()` fetches the current active trace from the execution (that's staying the same) and now takes an argument for trace-level fields, 
  - ~~`traceUtil.getTraceContext()` --> renamed `traceUtil.getSpanContext()` is calling a parse function on a header. I think it's gotten confusing for these to have the same name~~
  - update on ^ where it's non-breaking, traceContext has been renamed to spanContext when span context is a more accurate term. This is specifically when we are referring to an object derived from a trace header. trace context is more appropriate than custom context in many places, but as this is more of a breaking change, we're mostly sitting on that for now, with comments to help clarify what's going on there

katiebayes: jumping back into this now, sorry for the delay
katiebayes: > kind of annoyed that a bug fix for not prefixing custom trace fields with `app.` got rolled into this giant bikeshed.

if it helps, this is also going to be very helpful for cleaner propagation of trace-level context, so it was this one or the next one honestly.
katiebayes: Ok, I think I covered everything noted. Let me know what you think of this. I opted to swap the ordering of the dataset and traceContext since dataset is never default-populated.
katiebayes: Updated from `main`

Patch:
diff --git a/docs/API.md b/docs/API.md
index 8cb6923..9789398 100644
--- a/docs/API.md
+++ b/docs/API.md
@@ -27,7 +27,7 @@ In general you're going to be creating spans more often than traces. If you're u
 #### startTrace()
 
 ```javascript
-beeline.startTrace(metadataContext[, withTraceId, withParentSpanId])
+beeline.startTrace(metadataContext[, withTraceId, withParentSpanId, traceContext, withDataset])
 ```
 
 Starts a new local trace and initializes the async context propagation machinery. You _must_ have an active trace for the rest of the API to do anything. If you call `startSpan` when you aren't currently in an trace, an `Error` will be thrown. The instrumentations (which must operate in both trace/non-trace environments) handle this by checking `beeline.traceActive()` and only creating spans if they're within a trace.
@@ -37,7 +37,7 @@ root span is installed as the current span.
 
 Returns a reference to the trace.
 
-For an explanation of `withTraceId` and `withParentSpanId`, see [Interprocess trace propagation](#interprocess-trace-propagation) below.
+For an explanation of `withTraceId`, `withParentSpanId`, and `traceContext` see [Interprocess trace propagation](#interprocess-trace-propagation) below.
 
 example:
 
@@ -240,7 +240,9 @@ console.log(traceContext); // => 1;trace_id=weofijwoeifj,parent_id=owefjoweifj,c
 beeline.unmarshalTraceContext(traceContext);
 ```
 
-Returns an object containing the properties `traceId` and `parentSpanId`, which are the two optional parameters with `startTrace()` above.
+Returns an object containing the properties `traceId`, `parentSpanId`, `dataset`, and `customContext` which are the four optional parameters with `startTrace()` above.
+
+`customContext` will include any custom fields propagated from other services, such as with `beeline.addTraceContext()`.
 
 example:
 
@@ -252,6 +254,14 @@ let { traceId, parentSpanId } = beeline.unmarshalTraceContext(
 let trace = startTrace({ name }, traceId, parentSpanId);
 ```
 
+```javascript
+let { traceId, parentSpanId, dataset, customContext } = beeline.unmarshalTraceContext(
+  req.header[beeline.TRACE_HTTP_HEADER]
+);
+
+let trace = startTrace({ name }, traceId, parentSpanId, dataset, customContext);
+```
+
 #### TRACE_HTTP_HEADER
 
 The HTTP header that the beeline uses both for sending and receiving trace context. The value is `"X-Honeycomb-Trace"`.
diff --git a/lib/api/index.js b/lib/api/index.js
index dbafb0b..cd8b0dd 100644
--- a/lib/api/index.js
+++ b/lib/api/index.js
@@ -76,9 +76,16 @@ module.exports = {
     marshalTraceContext: w3c.marshalTraceContext,
   },
 
-  startTrace(metadataContext, withTraceId, withParentSpanId, withDataset) {
-    return apiImpl.startTrace(metadataContext, withTraceId, withParentSpanId, withDataset);
+  startTrace(fields, withTraceId, withParentSpanId, withDataset, propagatedContext = {}) {
+    return apiImpl.startTrace(
+      fields,
+      withTraceId,
+      withParentSpanId,
+      withDataset,
+      propagatedContext
+    );
   },
+
   finishTrace(trace) {
     return apiImpl.finishTrace(trace);
   },
diff --git a/lib/api/libhoney.js b/lib/api/libhoney.js
index 1a58453..a95558d 100644
--- a/lib/api/libhoney.js
+++ b/lib/api/libhoney.js
@@ -102,18 +102,18 @@ module.exports = class LibhoneyEventAPI {
     });
   }
 
-  startTrace(metadataContext, traceId, parentSpanId, dataset) {
+  // fields and propagatedContext are both objects
+  startTrace(fields, traceId, parentSpanId, dataset, propagatedContext) {
     const id = traceId || generateTraceId();
 
     // initialize an async hook
     tracker.setTracked({
       id,
-      traceContext: {},
+      traceContext: propagatedContext,
       stack: [],
       dataset: dataset || this.defaultDataset,
     });
-    // and create a new span to populate the first span in the stack array
-    return this.startSpan(metadataContext, undefined, parentSpanId);
+    return this.startSpan(fields, undefined, parentSpanId);
   }
 
   finishTrace(span) {
@@ -143,9 +143,6 @@ module.exports = class LibhoneyEventAPI {
       parentId = context.parentId;
     }
 
-    // create new Span object
-    // carry over whatever fields are passed in the first arg
-    // and add the trace id and span id
     const span = new Span(
       Object.assign({}, metadataContext, {
         [schema.TRACE_ID]: context.id,
diff --git a/lib/instrumentation/express.js b/lib/instrumentation/express.js
index 6628795..959ead6 100644
--- a/lib/instrumentation/express.js
+++ b/lib/instrumentation/express.js
@@ -83,16 +83,26 @@ const getMagicMiddleware = ({ userContext, traceIdSource, parentIdSource, packag
   res,
   next
 ) => {
-  let traceContext = traceUtil.getTraceContext(traceIdSource, req);
-  let parentTraceId = traceUtil.getParentSourceId(parentIdSource, req);
-  if (parentTraceId) {
-    traceContext.parentSpanId = parentTraceId;
+  // parse incoming trace headers into a span context
+  let spanContext = traceUtil.getTraceContext(traceIdSource, req);
+
+  // parentSpanId refers to the span id of the parent span
+  let parentSpanId = traceUtil.getParentSourceId(parentIdSource, req);
+  if (parentSpanId) {
+    spanContext.parentSpanId = parentSpanId;
+  }
+
+  // propagate trace level fields as-is, if there are any
+  let propagatedContext = {};
+  if (spanContext.customContext) {
+    propagatedContext = spanContext.customContext;
   }
+
   let fields = {
     [schema.EVENT_TYPE]: "express",
     [schema.PACKAGE_VERSION]: packageVersion,
     [schema.TRACE_SPAN_NAME]: "request",
-    [schema.TRACE_ID_SOURCE]: traceContext.source,
+    [schema.TRACE_ID_SOURCE]: spanContext.source,
     "request.host": req.hostname,
     "request.original_url": req.originalUrl,
     "request.remote_addr": req.ip,
@@ -115,15 +125,12 @@ const getMagicMiddleware = ({ userContext, traceIdSource, parentIdSource, packag
 
   let rootSpan = api.startTrace(
     fields,
-    traceContext.traceId,
-    traceContext.parentSpanId,
-    traceContext.dataset
+    spanContext.traceId,
+    spanContext.parentSpanId,
+    spanContext.dataset,
+    propagatedContext
   );
 
-  if (traceContext.customContext) {
-    api.addTraceContext(traceContext.customContext);
-  }
-
   if (!rootSpan) {
     // sampler has decided that we shouldn't trace this request
     next();
diff --git a/lib/instrumentation/fastify.js b/lib/instrumentation/fastify.js
index ba7674b..3cfa801 100644
--- a/lib/instrumentation/fastify.js
+++ b/lib/instrumentation/fastify.js
@@ -70,16 +70,29 @@ const instrumentFastify = function(fastify, opts = {}) {
 
     app.addHook("onRequest", (request, reply, next) => {
       // start our trace handling
-      let traceContext = traceUtil.getTraceContext(traceIdSource, request);
-      let parentTraceId = traceUtil.getParentSourceId(parentIdSource, request);
-      if (parentTraceId) {
-        traceContext.parentSpanId = parentTraceId;
+
+      // parse incoming trace headers into a span context
+      let spanContext = traceUtil.getTraceContext(traceIdSource, request);
+
+      // parentSpanId refers to the span id of the parent span
+      let parentSpanId = traceUtil.getParentSourceId(parentIdSource, request);
+      if (parentSpanId) {
+        spanContext.parentSpanId = parentSpanId;
+      }
+
+      // if the parsed header contains trace level fields
+      // propagate them without prefixing or changing them
+      // includes custom context from a parent span
+      let propagatedContext = {};
+      if (spanContext.customContext) {
+        propagatedContext = spanContext.customContext;
       }
+
       let fields = {
         [schema.EVENT_TYPE]: "fastify",
         [schema.PACKAGE_VERSION]: opts.packageVersion,
         [schema.TRACE_SPAN_NAME]: "request",
-        [schema.TRACE_ID_SOURCE]: traceContext.source,
+        [schema.TRACE_ID_SOURCE]: spanContext.source,
         "request.host": request.hostname,
         "request.original_url": request.req.originalUrl,
         "request.remote_addr": request.ip,
@@ -97,15 +110,12 @@ const instrumentFastify = function(fastify, opts = {}) {
       }
       let rootSpan = api.startTrace(
         fields,
-        traceContext.traceId,
-        traceContext.parentSpanId,
-        traceContext.dataset
+        spanContext.traceId,
+        spanContext.parentSpanId,
+        spanContext.dataset,
+        propagatedContext
       );
 
-      if (traceContext.customContext) {
-        api.addTraceContext(traceContext.customContext);
-      }
-
       if (!rootSpan) {
         // sampler has decided that we shouldn't trace this request
         next();
diff --git a/lib/instrumentation/trace-util.js b/lib/instrumentation/trace-util.js
index a264567..cc4a694 100644
--- a/lib/instrumentation/trace-util.js
+++ b/lib/instrumentation/trace-util.js
@@ -31,6 +31,8 @@ function _includeSource(context, key) {
   });
 }
 
+// parse a trace header and return object used to populate args to startTrace
+// deprecated: in the next major version this should get renamed to getSpanContext
 exports.getTraceContext = (traceIdSource, req) => {
   const { honeycomb, aws } = api;
   if (typeof traceIdSource === "undefined" || typeof traceIdSource === "string") {
diff --git a/lib/propagation/honeycomb.js b/lib/propagation/honeycomb.js
index 5b4d90a..45ddeae 100644
--- a/lib/propagation/honeycomb.js
+++ b/lib/propagation/honeycomb.js
@@ -97,10 +97,10 @@ function unmarshalTraceContextv1(payload) {
     return;
   }
 
-  let customContext;
+  let traceContext;
   if (contextb64) {
     try {
-      customContext = JSON.parse(Buffer.from(contextb64, "base64").toString("ascii"));
+      traceContext = JSON.parse(Buffer.from(contextb64, "base64").toString("ascii"));
     } catch (e) {
       debug("couldn't decode context", e);
       return;
@@ -110,7 +110,10 @@ function unmarshalTraceContextv1(payload) {
   return {
     traceId,
     parentSpanId,
-    customContext,
+
+    // trace level fields go into customContext
+    // these will propagate unchanged
+    customContext: traceContext,
     dataset,
   };
 }
diff --git a/lib/schema.js b/lib/schema.js
index 937951c..fa17d7a 100644
--- a/lib/schema.js
+++ b/lib/schema.js
@@ -16,4 +16,5 @@ exports.TRACE_SPAN_ID = "trace.span_id";
 exports.TRACE_SERVICE_NAME = "service_name";
 exports.TRACE_SPAN_NAME = "name";
 
+// custom context fields will have this prefix
 exports.customContext = key => `app.${key}`;


Test patch:
diff --git a/lib/api/index.test.js b/lib/api/index.test.js
index 8a40b4b..d7f94ca 100644
--- a/lib/api/index.test.js
+++ b/lib/api/index.test.js
@@ -561,10 +561,10 @@ test("distributed tracing linkage is correct", () => {
       [schema.TRACE_SPAN_NAME]: "sub",
     },
     ctx.traceId,
-    ctx.parentSpanId
+    ctx.parentSpanId,
+    ctx.dataset,
+    ctx.customContext
   );
-  const subContext = tracker.getTracked();
-  subContext.traceContext = ctx.customContext;
 
   let subserviceSubspan = api.startSpan({
     [schema.TRACE_SPAN_NAME]: "subspan",
@@ -598,6 +598,11 @@ test("distributed tracing linkage is correct", () => {
   expect(subTraceEventData[schema.TRACE_PARENT_ID]).toBe(rootEventData[schema.TRACE_SPAN_ID]);
   expect(rootEventData[schema.TRACE_PARENT_ID]).toBeUndefined();
 
+  // make sure dataset propagates correctly
+  expect(honey.transmission.events[0].dataset).toBe("nodejs");
+  expect(honey.transmission.events[1].dataset).toBe("nodejs");
+  expect(honey.transmission.events[2].dataset).toBe("nodejs");
+
   // we should have 3 span ids
   let s = new Set();
   s.add(subSpanEventData[schema.TRACE_SPAN_ID]);
@@ -683,3 +688,38 @@ describe("presend hook", () => {
     expect(api._apiForTesting().honey.transmission.events).toHaveLength(0);
   });
 });
+
+test("traceContext, no prefix on incoming trace context, dataset behaves as expected", () => {
+  const honey = api._apiForTesting().honey;
+
+  let rootSpan = api.startTrace(
+    {
+      [schema.TRACE_SPAN_NAME]: "root",
+    },
+    undefined,
+    undefined,
+    "test_dataset",
+    {
+      traceLevel: "field",
+    }
+  );
+
+  api.customContext.add("custom", "value");
+
+  api.finishTrace(rootSpan);
+
+  const postData = JSON.parse(honey.transmission.events[0].postData);
+
+  // no prefix on incoming trace level context
+  expect(postData.traceLevel).toBe("field");
+
+  // prefix on api-added custom context
+  expect(postData["app.custom"]).toBe("value");
+
+  // trace id and span id gets assigned as expected
+  expect(postData["trace.trace_id"]).toBeDefined();
+  expect(postData["trace.span_id"]).toBeDefined();
+
+  // dataset goes through as expected
+  expect(honey.transmission.events[0].dataset).toBe("test_dataset");
+});
diff --git a/lib/propagation/propagation.test.js b/lib/propagation/propagation.test.js
index 14bd01f..aa38009 100644
--- a/lib/propagation/propagation.test.js
+++ b/lib/propagation/propagation.test.js
@@ -205,7 +205,6 @@ cases(
 describe("roundtrip", () => {
   test("works", () => {
     let contextStr = propagation.marshalTraceContext(testContext);
-
     expect(propagation.unmarshalTraceContext(contextStr)).toEqual({
       traceId: "abcdef123456",
       parentSpanId: "0102030405",