fix datadog tracer scope issues

src/context.js

@@ -140,6 +140,12 @@ class Context {
 		else if (opts.parentCtx != null && opts.parentCtx.requestID != null)
 			ctx.requestID = opts.parentCtx.requestID;
 
+		// TraceID
+		if (opts.traceID != null)
+			ctx.traceID = opts.traceID;
+		else if (opts.parentCtx != null && opts.parentCtx.traceID != null)
+			ctx.traceID = opts.parentCtx.traceID;
+
 		// Meta
 		if (opts.parentCtx != null && opts.parentCtx.meta != null)
 			ctx.meta = Object.assign({}, opts.parentCtx.meta || {}, opts.meta || {});
@@ -169,6 +175,7 @@ class Context {
 		if (opts.parentSpan != null) {
 			ctx.parentID = opts.parentSpan.id;
 			ctx.requestID = opts.parentSpan.traceID;
+			ctx.traceID = opts.parentSpan.traceID;
 			ctx.tracing = opts.parentSpan.sampled;
 		}
 
@@ -198,6 +205,7 @@ class Context {
 		newCtx.meta = this.meta;
 		newCtx.locals = this.locals;
 		newCtx.requestID = this.requestID;
+		newCtx.traceID = this.traceID;
 		newCtx.tracing = this.tracing;
 		newCtx.span = this.span;
 		newCtx.needAck = this.needAck;
@@ -459,6 +467,7 @@ class Context {
 			"meta",
 			//"locals",
 			"requestID",
+			"traceID",
 			"tracing",
 			"span",
 			"needAck",

src/middlewares/tracing.js

@@ -23,6 +23,7 @@ module.exports = function TracingMiddleware(broker) {
 			return function tracingLocalActionMiddleware(ctx) {
 
 				ctx.requestID = ctx.requestID || tracer.getCurrentTraceID();
+				ctx.traceID = ctx.traceID || tracer.getCurrentTraceID();
 				ctx.parentID = ctx.parentID || tracer.getActiveSpanID();
 
 				const tags = {
@@ -84,17 +85,18 @@ module.exports = function TracingMiddleware(broker) {
 				const span = ctx.startSpan(spanName, {
 					id: ctx.id,
 					type: "action",
-					traceID: ctx.requestID,
+					traceID: ctx.traceID || ctx.requestID,
 					parentID: ctx.parentID,
 					service: ctx.service,
 					sampled: ctx.tracing,
+					autoActivate: false,
 					tags
 				});
 
 				ctx.tracing = span.sampled;
 
 				// Call the handler
-				return handler(ctx).then(res => {
+				return tracer.activate(span, () => handler(ctx)).then(res => {
 					const tags = {
 						fromCache: ctx.cachedResult
 					};
@@ -142,6 +144,7 @@ module.exports = function TracingMiddleware(broker) {
 			return function tracingLocalEventMiddleware(ctx) {
 
 				ctx.requestID = ctx.requestID || tracer.getCurrentTraceID();
+				ctx.traceID = ctx.traceID || tracer.getCurrentTraceID();
 				ctx.parentID = ctx.parentID || tracer.getActiveSpanID();
 
 				const tags = {
@@ -202,17 +205,18 @@ module.exports = function TracingMiddleware(broker) {
 				const span = ctx.startSpan(spanName, {
 					id: ctx.id,
 					type: "event",
-					traceID: ctx.requestID,
+					traceID: ctx.traceID || ctx.requestID,
 					parentID: ctx.parentID,
 					service,
 					sampled: ctx.tracing,
+					autoActivate: false,
 					tags
 				});
 
 				ctx.tracing = span.sampled;
 
 				// Call the handler
-				return handler.apply(service, arguments).then(() => {
+				return tracer.activate(span, () => handler.apply(service, arguments)).then(() => {
 					ctx.finishSpan(span);
 				}).catch(err => {
 					span.setError(err);

src/tracing/exporters/datadog.js

@@ -1,8 +1,8 @@
 "use strict";
 
+const { executionAsyncResource } = require("async_hooks");
 const _ 					= require("lodash");
 const BaseTraceExporter 	= require("./base");
-const asyncHooks			= require("async_hooks");
 const { isFunction } 		= require("../../utils");
 
 /*
@@ -11,6 +11,7 @@ const { isFunction } 		= require("../../utils");
 
 let DatadogSpanContext;
 let DatadogID;
+let Reference;
 
 /**
  * Datadog Trace Exporter with 'dd-trace'.
@@ -51,6 +52,7 @@ class DatadogTraceExporter extends BaseTraceExporter {
 			const ddTrace = require("dd-trace");
 			DatadogSpanContext = require("dd-trace/packages/dd-trace/src/opentracing/span_context");
 			DatadogID = require("dd-trace/packages/dd-trace/src/id");
+			Reference = require("opentracing").Reference;
 			if (!this.ddTracer) {
 				this.ddTracer = ddTrace.init(_.defaultsDeep(this.opts.tracerOptions, {
 					url: this.opts.agentUrl
@@ -68,6 +70,12 @@ class DatadogTraceExporter extends BaseTraceExporter {
 
 		this.ddScope = this.ddTracer.scope();
 
+		// Determine which scope implementation is being used by dd tracer.
+		// dd-trace-js defaults to async_resource for node >=14.5 || ^12.19.0, but async_hooks
+		// and async_local_storage can be forced with the `scope` init option
+		this.usesAsyncResource = this.ddScope._ddResourceStore !== undefined;
+		this.usesAsyncLocalStorage = this.ddScope._storage !== undefined;
+
 		const oldGetCurrentTraceID = this.tracer.getCurrentTraceID.bind(this.tracer);
 		this.tracer.getCurrentTraceID = () => {
 			const traceID = oldGetCurrentTraceID();
@@ -121,18 +129,18 @@ class DatadogTraceExporter extends BaseTraceExporter {
 
 		const serviceName = span.service ? span.service.fullName : null;
 
-		let parentCtx;
+		let reference;
 		if (span.parentID) {
-			parentCtx = new DatadogSpanContext({
+			const parentCtx = new DatadogSpanContext({
 				traceId: this.convertID(span.traceID),
-				spanId: this.convertID(span.parentID),
-				parentId: this.convertID(span.parentID)
+				spanId: this.convertID(span.parentID)
 			});
+			reference = new Reference(span.type === "event" ? "follows_from" : "child_of", parentCtx);
 		}
 
 		const ddSpan = this.ddTracer.startSpan(span.name, {
 			startTime: span.startTime,
-			childOf: parentCtx,
+			references: reference ? [reference] : [],
 			tags: this.flattenTags(_.defaultsDeep({}, span.tags, {
 				span: {
 					kind: "server",
@@ -152,20 +160,59 @@ class DatadogTraceExporter extends BaseTraceExporter {
 		sc._traceId = this.convertID(span.traceID);
 		sc._spanId = this.convertID(span.id);
 
-		// Activate span in Datadog tracer
-		const asyncId = asyncHooks.executionAsyncId();
-		this.ddScope._spans = this.ddScope._spans || {};
-		const oldSpan = this.ddScope._spans[asyncId];
+		let deactivate;
+
+		// If the caller doesn't plan to use the preferred `activate` method,
+		// we need to do our best to activate the span in the dd-trace internals.
+		// In certain situations, this has issues with active spans leaking outside their async scope.
+		if (span.autoActivate) {
+			// Save current active span to restore later
+			const previousSpan = this.ddScope.active();
+			const asyncResource = executionAsyncResource();
+
+			// Activate new span in Datadog tracer
+			if (this.usesAsyncResource) {
+				this.ddScope._enter(ddSpan, asyncResource);
+			} else if (this.usesAsyncLocalStorage) {
+				this.ddScope._storage.enterWith(ddSpan);
+			} else {
+				this.ddScope._enter(ddSpan);
+			}
 
-		this.ddScope._spans[asyncId] = ddSpan;
+			// Prepare deactivation function for later
+			deactivate = () => {
+				if (this.usesAsyncResource) {
+					this.ddScope._exit(asyncResource);
+					const currentResource = executionAsyncResource();
+					if (currentResource !== asyncResource) {
+						this.ddScope._exit(currentResource);
+						this.ddScope._enter(previousSpan || null, currentResource);
+					}
+				} else if (this.usesAsyncLocalStorage) {
+					this.ddScope._storage.enterWith(previousSpan || null);
+				} else {
+					this.ddScope._exit(previousSpan || null);
+				}
+			};
+		}
 
 		span.meta.datadog = {
 			span: ddSpan,
-			asyncId,
-			oldSpan
+			deactivate
 		};
 	}
 
+	/**
+	 * Activate a span in the scope of a function.
+	 *
+	 * @param {Span} span
+	 * @param {Function} fn
+	 * @returns
+	 */
+	activate(span, fn) {
+		return this.ddScope.activate(span.meta.datadog.span, fn);
+	}
+
 	/**
 	 * Span is finished.
 	 *
@@ -186,13 +233,12 @@ class DatadogTraceExporter extends BaseTraceExporter {
 
 		this.addLogs(ddSpan, span.logs);
 
-		ddSpan.finish(span.finishTime);
-
-		if (item.oldSpan) {
-			this.ddScope._spans[item.asyncId] = item.oldSpan;
-		} else {
-			delete this.ddScope._spans[item.asyncId];
+		// Deactivate the span in the context if needed
+		if (item.deactivate) {
+			item.deactivate();
 		}
+
+		ddSpan.finish(span.finishTime);
 	}
 
 	/**
@@ -278,9 +324,11 @@ class DatadogTraceExporter extends BaseTraceExporter {
 	 */
 	convertID(id) {
 		if (id) {
+			// Parse guid as a hex string
 			if (id.indexOf("-") !== -1)
-				return DatadogID(id.replace(/-/g, "").substring(0,16));
-			return DatadogID(id);
+				return DatadogID(id.replace(/-/g, "").substring(0,16), 16);
+			// Otherwise, parse as a decimal string because this is likely an ID from dd-trace
+			return DatadogID(id, 10);
 		}
 
 		return null;

src/tracing/span.js

@@ -54,6 +54,7 @@ class Span {
 
 		this.priority = this.opts.priority != null ? this.opts.priority : 5;
 		this.sampled = this.opts.sampled != null ? this.opts.sampled : this.tracer.shouldSample(this);
+		defProp(this, "autoActivate", this.opts.autoActivate != null ? this.opts.autoActivate : true);
 
 		this.startTime = null;
 		this.startTicks = null;

src/tracing/tracer.js

@@ -207,6 +207,27 @@ class Tracer {
 		return span;
 	}
 
+	/**
+	 * Call a function with the provided span activated in the exporters' context
+	 *
+	 * @param {Span} span
+	 * @param {Function} fn
+	 */
+	activate(span, fn) {
+		if (this.exporter) {
+			const handler = this.exporter.reduce((next, exporter) => {
+				if (typeof exporter.activate === "function") {
+					return exporter.activate.bind(exporter, span, next);
+				}
+				return next;
+			}, fn);
+
+			return handler();
+		} else {
+			return fn();
+		}
+	}
+
 	/**
 	 * Invoke Exporter method.
 	 *

src/transit.js

@@ -384,6 +384,7 @@ class Transit {
 		ctx.level = payload.level;
 		ctx.tracing = !!payload.tracing;
 		ctx.parentID = payload.parentID;
+		ctx.traceID = payload.traceID;
 		ctx.requestID = payload.requestID;
 		ctx.caller = payload.caller;
 		ctx.nodeID = payload.sender;
@@ -422,6 +423,7 @@ class Transit {
 			ctx.id = payload.id;
 			ctx.setParams(pass ? pass : payload.params, this.broker.options.contextParamsCloning);
 			ctx.parentID = payload.parentID;
+			ctx.traceID = payload.traceID;
 			ctx.requestID = payload.requestID;
 			ctx.caller = payload.caller;
 			ctx.meta = payload.meta || {};
@@ -736,6 +738,7 @@ class Transit {
 			level: ctx.level,
 			tracing: ctx.tracing,
 			parentID: ctx.parentID,
+			traceID: ctx.traceID,
 			requestID: ctx.requestID,
 			caller: ctx.caller,
 			stream: isStream,
@@ -855,6 +858,7 @@ class Transit {
 			level: ctx.level,
 			tracing: ctx.tracing,
 			parentID: ctx.parentID,
+			traceID: ctx.traceID,
 			requestID: ctx.requestID,
 			caller: ctx.caller,
 			needAck: ctx.needAck