feat:setTraceAttributes api (#36)

* feat: set trace attributes api

* style: cleanups

* docs: add comment

* docs: README

* style: cleanups

* docs: add changeset

* feat: hook custom attributes (console instrumentation)

* style: move context class

* feat: inject trace attributes to logger message (winston + pino)

* docs: add changeset

* feedback: docs

* perf: setup hard limit on spans per trace + total trace ids

* fix: check trace ids limit

* fix: logger meta should override global trace attributes

.changeset/poor-dingos-judge.md

@@ -0,0 +1,5 @@
+---
+'@hyperdx/node-opentelemetry': patch
+---
+
+feat: setTraceAttributes user-facing API

.changeset/rude-rocks-matter.md

@@ -0,0 +1,5 @@
+---
+'@hyperdx/node-logger': patch
+---
+
+feat: add getCustomMeta option (pino + winston)

.changeset/violet-cats-marry.md

@@ -0,0 +1,5 @@
+---
+'@hyperdx/node-opentelemetry': patch
+---
+
+feat: add trace attributes to console + logger

packages/node-logger/src/pino.ts

@@ -15,7 +15,9 @@ const PINO_LEVELS = {
   60: 'fatal',
 };
 
-export type HyperDXPinoOptions = LoggerOptions;
+export type HyperDXPinoOptions = LoggerOptions & {
+  getCustomMeta?: () => Record<string, any>;
+};
 
 export default (opts: HyperDXPinoOptions) => {
   try {
@@ -27,7 +29,10 @@ export default (opts: HyperDXPinoOptions) => {
         for await (const obj of source) {
           const { level, message, meta } = parsePinoLog(obj);
           hdx('Sending log to HyperDX');
-          logger.postMessage(PINO_LEVELS[level], message, meta);
+          logger.postMessage(PINO_LEVELS[level], message, {
+            ...opts.getCustomMeta?.(),
+            ...meta,
+          });
           hdx('Log sent to HyperDX');
         }
       },

packages/node-logger/src/winston.ts

@@ -7,11 +7,14 @@ import type { LoggerOptions } from './logger';
 
 export type HyperDXWinstonOptions = LoggerOptions & {
   maxLevel?: string;
+  getCustomMeta?: () => Record<string, any>;
 };
 
 export default class HyperDXWinston extends Transport {
   private readonly logger: Logger;
 
+  private readonly getCustomMeta: () => Record<string, any>;
+
   constructor({
     apiKey,
     baseUrl,
@@ -20,9 +23,11 @@ export default class HyperDXWinston extends Transport {
     sendIntervalMs,
     service,
     timeout,
+    getCustomMeta,
   }: HyperDXWinstonOptions) {
     hdx('Initializing HyperDX winston transport...');
     super({ level: maxLevel ?? 'info' });
+    this.getCustomMeta = getCustomMeta;
     this.logger = new Logger({
       apiKey,
       baseUrl,
@@ -45,7 +50,10 @@ export default class HyperDXWinston extends Transport {
 
     const { level, message, meta } = parseWinstonLog(info);
     hdx('Sending log to HyperDX');
-    this.logger.postMessage(level, message, meta);
+    this.logger.postMessage(level, message, {
+      ...this.getCustomMeta?.(),
+      ...meta,
+    });
     hdx('Log sent to HyperDX');
     callback();
   }

packages/node-opentelemetry/README.md

@@ -113,7 +113,30 @@ And run your application with the following command (example using `ts-node`):
 npx ts-node -r './instrument.ts' index.ts
 ```
 
-### Advanced Instrumentation Configuration (Optional)
+### (Optional) Attach User Information or Metadata
+
+To easily tag all events related to a given attribute or identifier (ex. user id or email), you can call the `setTraceAttributes` function which will tag every log/span associated with the current trace after the call with the declared attributes. It's recommended to call this function as early as possible within a given request/trace (ex. as early in an Express middleware stack as possible).
+
+This is a convenient way to ensure all logs/spans are automatically tagged with the right identifiers to be searched on later, instead of needing to manually tagging and propagating identifiers yourself.
+
+`userId`, `userEmail`, `userName`, and `teamName` will populate the sessions UI with the corresponding values, but can be omitted. Any other additional values can be specified and used to search for events.
+
+```ts
+import { setTraceAttributes } from '@hyperdx/node-opentelemetry';
+
+app.use((req, res, next) => {
+  // Get user information from the request...
+
+  // Attach user information to the current trace
+  setTraceAttributes({
+    userId,
+    userEmail,
+  });
+  next();
+});
+```
+
+### (Optional) Advanced Instrumentation Configuration
 
 #### Capture Console Logs
 

packages/node-opentelemetry/dummy.js

@@ -11,6 +11,7 @@ const {
   getPinoTransport,
   init,
 } = require('./build/src/logger');
+const { setTraceAttributes } = require('./build/src');
 
 const logger = winston.createLogger({
   level: 'info',
@@ -60,6 +61,25 @@ async function sendGetRequest() {
   });
 }
 
+function generateRandomString(length) {
+  const characters =
+    'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789';
+  let result = '';
+  for (let i = 0; i < length; i++) {
+    const randomIndex = Math.floor(Math.random() * characters.length);
+    result += characters.charAt(randomIndex);
+  }
+  return result;
+}
+
+// set custom trace attributes
+app.use((req, res, next) => {
+  setTraceAttributes({
+    userId: generateRandomString(8),
+  });
+  next();
+});
+
 app.use(express.json());
 
 app.post('/dump', (req, res) => {

packages/node-opentelemetry/src/context.ts

@@ -0,0 +1,114 @@
+import { Attributes, context, trace } from '@opentelemetry/api';
+import { Span } from '@opentelemetry/sdk-trace-base';
+import { suppressTracing } from '@opentelemetry/core';
+
+import hdx from './debug';
+
+const HDX_CONTEXT_MAX_SPANS_PER_TRACE = 100;
+const HDX_CONTEXT_MAX_TRACKED_TRACE_IDS = 50000; // ~ 500 MB of memory (10 spans / trace, 1k per span)
+const HDX_CONTEXT_REFRESHER_INTERVAL = 10000;
+const HDX_CONTEXT_TRACE_ATTRIBUTES_EXPIRATION = 5 * 60 * 1000;
+
+class HyperDXContext {
+  private readonly _traceAttributes = new Map<string, Attributes>();
+  private readonly _traceMap = new Map<
+    string,
+    {
+      spans: Span[];
+      lastUpdateAt: number;
+    }
+  >();
+
+  constructor() {
+    // expires after 5 minutes
+    setInterval(() => {
+      hdx(`Running _traceMap expiration check`);
+      const now = Date.now();
+      for (const [traceId, data] of this._traceMap.entries()) {
+        if (now - data.lastUpdateAt > HDX_CONTEXT_TRACE_ATTRIBUTES_EXPIRATION) {
+          this._traceMap.delete(traceId);
+          this._traceAttributes.delete(traceId);
+          hdx(`Deleted traceId ${traceId} from _traceMap and _traceAttributes`);
+        }
+      }
+    }, HDX_CONTEXT_REFRESHER_INTERVAL);
+  }
+
+  private _getActiveSpanTraceId = (): string | undefined => {
+    const activeSpan = trace.getActiveSpan();
+    if (!activeSpan) {
+      return undefined;
+    }
+    return activeSpan.spanContext().traceId;
+  };
+
+  private _setTraceAttributesForAllSpans = (traceId: string): void => {
+    const attributes = this.getTraceAttributes(traceId);
+    if (!attributes) {
+      return;
+    }
+    const traceData = this._traceMap.get(traceId);
+    if (traceData) {
+      for (const span of traceData.spans) {
+        if (!span.ended) {
+          span.setAttributes(attributes);
+        }
+      }
+    }
+  };
+
+  addTraceSpan = (traceId: string, span: Span): void => {
+    // prevent downstream exporter calls from generating spans
+    context.with(suppressTracing(context.active()), () => {
+      hdx(`Adding traceId ${traceId} to _traceMap`);
+      const traceData = this._traceMap.get(traceId);
+      if (!traceData) {
+        if (this._traceMap.size >= HDX_CONTEXT_MAX_TRACKED_TRACE_IDS) {
+          hdx(
+            `Exceeded max tracked trace ids: ${HDX_CONTEXT_MAX_TRACKED_TRACE_IDS}`,
+          );
+          return;
+        }
+        this._traceMap.set(traceId, {
+          spans: [span],
+          lastUpdateAt: Date.now(),
+        });
+      } else {
+        if (traceData.spans.length >= HDX_CONTEXT_MAX_SPANS_PER_TRACE) {
+          hdx(
+            `Exceeded max spans per trace: ${HDX_CONTEXT_MAX_SPANS_PER_TRACE}`,
+          );
+          return;
+        }
+        traceData.spans.push(span);
+        traceData.lastUpdateAt = Date.now();
+      }
+      this._setTraceAttributesForAllSpans(traceId);
+    });
+  };
+
+  getTraceAttributes = (traceId: string): Attributes | undefined => {
+    return this._traceAttributes.get(traceId);
+  };
+
+  // user facing API
+  setTraceAttributes = (attributes: Attributes): void => {
+    // prevent downstream exporter calls from generating spans
+    context.with(suppressTracing(context.active()), () => {
+      const currentActiveSpanTraceId = this._getActiveSpanTraceId();
+      if (!currentActiveSpanTraceId) {
+        return;
+      }
+      if (this._traceAttributes.size >= HDX_CONTEXT_MAX_TRACKED_TRACE_IDS) {
+        hdx(
+          `Exceeded max tracked trace ids: ${HDX_CONTEXT_MAX_TRACKED_TRACE_IDS}`,
+        );
+        return;
+      }
+      this._traceAttributes.set(currentActiveSpanTraceId, attributes);
+      this._setTraceAttributesForAllSpans(currentActiveSpanTraceId);
+    });
+  };
+}
+
+export const hyperDXGlobalContext = new HyperDXContext();

packages/node-opentelemetry/src/instrumentations/console.ts

@@ -1,13 +1,14 @@
 import * as shimmer from 'shimmer';
 import _ from 'lodash';
-import api from '@opentelemetry/api';
+import opentelemetry from '@opentelemetry/api';
 import {
   Logger,
   LoggerOptions,
   parseWinstonLog,
 } from '@hyperdx/node-logger/build/src/logger';
 
 import hdx from '../debug';
+import { hyperDXGlobalContext } from '../context';
 
 export const _parseConsoleArgs = (args: any[]) => {
   const stringifiedArgs = [];
@@ -53,13 +54,19 @@ export default class HyperDXConsoleInstrumentation {
         level,
       });
 
-      const currentSpan = api.trace.getSpan(api.context.active());
+      const currentActiveSpan = opentelemetry.trace.getActiveSpan();
+      const traceId = currentActiveSpan?.spanContext().traceId;
+      const attributes = traceId
+        ? hyperDXGlobalContext.getTraceAttributes(traceId)
+        : {};
 
       this._logger.postMessage(parsedLog.level, parsedLog.message, {
         ...parsedLog.meta,
+        // attach custom attributes
+        ...attributes,
         // attached traceId and spanId,
-        trace_id: currentSpan?.spanContext().traceId,
-        span_id: currentSpan?.spanContext().spanId,
+        trace_id: traceId,
+        span_id: currentActiveSpan?.spanContext().spanId,
       });
     } catch (e) {
       hdx(`error in _patchConsole: ${e}`);

packages/node-opentelemetry/src/logger.ts

@@ -1,7 +1,10 @@
+import opentelemetry from '@opentelemetry/api';
+
 // @ts-ignore
 import { HyperDXWinston } from '@hyperdx/node-logger';
 
 import hdx from './debug';
+import { hyperDXGlobalContext } from './context';
 
 const env = process.env;
 
@@ -19,6 +22,12 @@ type WinstonTransportOptions = {
 
 type PinotTransportOptions = WinstonTransportOptions;
 
+const getCustomMeta = () => {
+  const currentActiveSpan = opentelemetry.trace.getActiveSpan();
+  const traceId = currentActiveSpan?.spanContext().traceId;
+  return traceId ? hyperDXGlobalContext.getTraceAttributes(traceId) : {};
+};
+
 export const getWinsonTransport = (
   maxLevel = 'info',
   options: WinstonTransportOptions = {},
@@ -28,6 +37,7 @@ export const getWinsonTransport = (
     apiKey: HYPERDX_API_KEY,
     maxLevel,
     service: SERVICE_NAME,
+    getCustomMeta,
     ...options,
   });
 };
@@ -40,6 +50,7 @@ export const getPinoTransport = (
   options: {
     apiKey: HYPERDX_API_KEY,
     service: SERVICE_NAME,
+    getCustomMeta,
     ...options,
   },
   level: maxLevel,

packages/node-opentelemetry/src/otel.ts

@@ -10,7 +10,9 @@ import hdx, {
   LOG_PREFIX as _LOG_PREFIX,
 } from './debug';
 import HyperDXConsoleInstrumentation from './instrumentations/console';
+import HyperDXSpanProcessor from './spanProcessor';
 import { getHyperDXHTTPInstrumentationConfig } from './instrumentations/http';
+import { hyperDXGlobalContext } from './context';
 
 const LOG_PREFIX = `⚠️  ${_LOG_PREFIX}`;
 
@@ -22,6 +24,8 @@ type SDKConfig = {
   advancedNetworkCapture?: boolean;
 };
 
+export const setTraceAttributes = hyperDXGlobalContext.setTraceAttributes;
+
 export const initSDK = (config: SDKConfig) => {
   // enable otel debug mode if HDX_DEBUG_MODE_ENABLED is set
   if (HDX_DEBUG_MODE_ENABLED) {
@@ -53,10 +57,12 @@ export const initSDK = (config: SDKConfig) => {
   });
 
   const sdk = new NodeSDK({
-    traceExporter: new OTLPTraceExporter({
-      timeoutMillis: 60000,
-    }),
     // metricReader: metricReader,
+    spanProcessor: new HyperDXSpanProcessor(
+      new OTLPTraceExporter({
+        timeoutMillis: 60000,
+      }),
+    ) as any,
     instrumentations: [
       getNodeAutoInstrumentations({
         '@opentelemetry/instrumentation-http': config.advancedNetworkCapture

packages/node-opentelemetry/src/spanProcessor.ts

@@ -0,0 +1,11 @@
+import { BatchSpanProcessor, Span } from '@opentelemetry/sdk-trace-base';
+import { Context } from '@opentelemetry/api';
+
+import { hyperDXGlobalContext } from './context';
+
+export default class HyperDXSpanProcessor extends BatchSpanProcessor {
+  onStart(_span: Span, _parentContext: Context): void {
+    const traceId = _span.spanContext().traceId;
+    hyperDXGlobalContext.addTraceSpan(traceId, _span);
+  }
+}