OTEL Attribute Decorator

attr.decorator.ts

import { ATTR_METADATA_KEY } from './constants';
import { AttrKey, AttrMetadata } from './types';

/**
 * Sets a custom attribute key on a parameter so that it can be picked up
 * by the Span decorator and mapped to the args values.
 */
export function Attr(key: AttrKey): ParameterDecorator {
  return function (
    target: NonNullable<unknown>,
    propertyKey: string | symbol | undefined, // method name
    parameterIndex: number,
  ) {
    if (!propertyKey) {
      console.log('propertyKey is undefined');
      return;
    }

    const existingMetadata: AttrMetadata[] =
      Reflect.getMetadata(ATTR_METADATA_KEY, target, propertyKey) || [];

    const metadata: AttrMetadata = {
      type: 'attr',
      key,
      parameterIndex,
    };

    const newMetadata = [...existingMetadata, metadata];

    Reflect.defineMetadata(ATTR_METADATA_KEY, newMetadata, target, propertyKey);
  };
}

constants.ts

export const ATTR_METADATA_KEY = Symbol('span:attr');

example.service.ts

@Injectable()
export class ExampleService {

  @Span()
  async whatchamacallit(
    @Attr('company.id') companyId: number,
    @Attr('data') data: SomeDataDto,
  ) {
    /* Your awesome code in here */
  }
}

span.decorator.ts

import { SpanOptions, trace } from '@opentelemetry/api';
import {
  copyMetadataFromFunctionToFunction,
  recordException,
  setAttribute,
  setParameterAttributes,
} from './utils';

interface ExtendedSpanOptions extends SpanOptions {
  name?: string;
  ignoreOutcomeAttr?: boolean;
}

/**
 * A modified version of the Span decorator from https://github.com/pragmaticivan/nestjs-otel
 */
export function Span(options: ExtendedSpanOptions = {}) {
  return (
    target: NonNullable<unknown>,
    propertyKey: string,
    propertyDescriptor: PropertyDescriptor,
  ) => {
    const originalFunction = propertyDescriptor.value;
    const wrappedFunction = function PropertyDescriptor(
      this: NonNullable<unknown>,
      ...args: unknown[]
    ) {
      const tracer = trace.getTracer('default');
      const spanName =
        options.name || `${target.constructor.name}.${propertyKey}`;

      return tracer.startActiveSpan(spanName, options, (span) => {
        setParameterAttributes(span, target, propertyKey, args);

        if (originalFunction.constructor.name === 'AsyncFunction') {
          return originalFunction
            .apply(this, args)
            .then((result: unknown) => {
              if (!options.ignoreOutcomeAttr) {
                setAttribute(propertyKey, span, 'outcome', result);
              }
              return result;
            })
            .catch((error: unknown) => {
              recordException(span, error);
              // Throw error to propagate it further
              throw error;
            })
            .finally(() => {
              span.end();
            });
        }

        try {
          const result = originalFunction.apply(this, args);
          if (!options.ignoreOutcomeAttr) {
            setAttribute(propertyKey, span, 'outcome', result);
          }
          return result;
        } catch (error) {
          recordException(span, error);
          // Throw error to propagate it further
          throw error;
        } finally {
          span.end();
        }
      });
    };
    // eslint-disable-next-line no-param-reassign
    propertyDescriptor.value = wrappedFunction;

    copyMetadataFromFunctionToFunction(originalFunction, wrappedFunction);
  };
}

types.ts

export interface AttrMetadata {
  type: 'attr';
  key: AttrKey;
  parameterIndex: number;
}

export type AttrKey =
  /* ---------------------- Frequent ---------------------- */
  | 'outcome'
  | 'company.id'
  /* ----------------------- PubSub ----------------------- */
  | 'pubsub.event'
  /* ------------------------ Redis ----------------------- */
  | 'redis.key'
  | 'redis.value'
  | 'redis.ttl.seconds'
  | 'redis.score'
  | 'redis.min.score'
  | 'redis.max.score'
  /* ----------------------- General ---------------------- */
  | 'str'
  | 'data'
  | 'reason';

utils.ts

import { Span, SpanStatusCode } from '@opentelemetry/api';
import { AttrKey, AttrMetadata } from './types';
import { ATTR_METADATA_KEY } from './constants';

export function recordException(span: Span, error: unknown) {
  if (error instanceof Error) {
    span.recordException(error);
    span.setStatus({ code: SpanStatusCode.ERROR, message: error.message });
  }
}

export function getAttrs(
  target: NonNullable<unknown>,
  propertyKey: string,
): AttrMetadata[] {
  try {
    const metadata: AttrMetadata[] =
      Reflect.getMetadata(ATTR_METADATA_KEY, target, propertyKey) || [];
    const filteredMetadata = metadata.filter((item) => item.type === 'attr');
    return filteredMetadata;
  } catch (error) {
    console.log({ error });
    return [];
  }
}

export function setAttribute(
  propertyKey: string,
  span: Span,
  key: AttrKey,
  value: unknown,
) {
  const serialized = (() => {
    try {
      if (
        typeof value === 'boolean' ||
        typeof value === 'number' ||
        typeof value === 'string'
      ) {
        return value;
      } else {
        // circular objects will fail to stringify
        return JSON.stringify(value) || 'null';
      }
    } catch (e) {
      const error =
        typeof e === 'object' && e !== null && 'message' in e ? e.message : e;
      console.error(
        `Failed to stringify value for span ${propertyKey}. Error ${error}`,
      );
      return 'un-json-stringify-able';
    }
  })();

  try {
    span.setAttribute(key, serialized);
  } catch (error) {
    console.error(error);
  }
}

export function setParameterAttributes(
  span: Span,
  targetClass: NonNullable<unknown>,
  propertyKey: string,
  args: unknown[],
) {
  try {
    const attrs = getAttrs(targetClass, propertyKey);
    attrs.forEach((attr) => {
      const value = args?.[attr.parameterIndex];
      setAttribute(propertyKey, span, attr.key, value);
    });
  } catch (error) {
    console.error(error);
  }
}