Browse Source

refactor: Extract span attribute code to telemetry plugin

Michael Bromley 8 months ago
parent
commit
f63490389d

+ 6 - 5
packages/core/src/common/index.ts

@@ -1,15 +1,16 @@
-export * from './finite-state-machine/finite-state-machine';
-export * from './finite-state-machine/types';
 export * from './async-queue';
 export * from './calculated-decorator';
-export * from './error/errors';
 export * from './error/error-result';
+export * from './error/errors';
 export * from './error/generated-graphql-admin-errors';
+export * from './finite-state-machine/finite-state-machine';
+export * from './finite-state-machine/types';
 export * from './injector';
+export * from './instrument-decorator';
 export * from './permission-definition';
-export * from './ttl-cache';
-export * from './self-refreshing-cache';
 export * from './round-money';
+export * from './self-refreshing-cache';
+export * from './ttl-cache';
 export * from './types/common-types';
 export * from './types/entity-relation-paths';
 export * from './types/injectable-strategy';

+ 15 - 3
packages/core/src/common/instrument-decorator.ts

@@ -1,8 +1,10 @@
 import { getConfig } from '../config';
 
+const INSTRUMENTED_CLASS = Symbol('InstrumentedClassTarget');
+
 export function Instrument(): ClassDecorator {
-    return function (target: Type<any>) {
-        return class extends (target as new (...args: any[]) => any) {
+    return function (target: any) {
+        const InstrumentedClass = class extends (target as new (...args: any[]) => any) {
             constructor(...args: any[]) {
                 // eslint-disable-next-line constructor-super
                 super(...args);
@@ -13,7 +15,7 @@ export function Instrument(): ClassDecorator {
                 }
                 return new Proxy(this, {
                     get: (obj, prop) => {
-                        const original = obj[prop];
+                        const original = obj[prop as string];
                         if (typeof original === 'function') {
                             return function (...methodArgs: any[]) {
                                 const wrappedMethodArgs = {
@@ -30,5 +32,15 @@ export function Instrument(): ClassDecorator {
                 });
             }
         };
+
+        // Set the name property of ProxiedClass to match the target's name
+        Object.defineProperty(InstrumentedClass, 'name', { value: target.name });
+        Object.defineProperty(InstrumentedClass, INSTRUMENTED_CLASS, { value: target });
+
+        return InstrumentedClass;
     };
 }
+
+export function getInstrumentedClassTarget(input: any) {
+    return input[INSTRUMENTED_CLASS];
+}

+ 2 - 2
packages/core/src/config/config.module.ts

@@ -1,6 +1,5 @@
 import { Module, OnApplicationBootstrap, OnApplicationShutdown } from '@nestjs/common';
 import { ModuleRef } from '@nestjs/core';
-import { notNullOrUndefined } from '@vendure/common/lib/shared-utils';
 
 import { ConfigurableOperationDef } from '../common/configurable-operation';
 import { Injector } from '../common/injector';
@@ -112,7 +111,7 @@ export class ConfigModule implements OnApplicationBootstrap, OnApplicationShutdo
         const { healthChecks, errorHandlers } = this.configService.systemOptions;
         const { assetImportStrategy } = this.configService.importExportOptions;
         const { refundProcess: refundProcess } = this.configService.paymentOptions;
-        const { cacheStrategy } = this.configService.systemOptions;
+        const { cacheStrategy, instrumentationStrategy } = this.configService.systemOptions;
         const entityIdStrategy = entityIdStrategyCurrent ?? entityIdStrategyDeprecated;
         return [
             ...adminAuthenticationStrategy,
@@ -155,6 +154,7 @@ export class ConfigModule implements OnApplicationBootstrap, OnApplicationShutdo
             guestCheckoutStrategy,
             ...refundProcess,
             cacheStrategy,
+            ...(instrumentationStrategy ? [instrumentationStrategy] : []),
             ...orderInterceptors,
         ];
     }

+ 0 - 6
packages/core/src/service/services/product.service.ts

@@ -132,10 +132,6 @@ export class ProductService {
         productId: ID,
         relations?: RelationPaths<Product>,
     ): Promise<Translated<Product> | undefined> {
-        const span = getActiveSpan();
-        span?.setAttribute('productId', productId.toString());
-        span?.setAttribute('channelId', ctx.channelId.toString());
-
         const effectiveRelations = relations ?? this.relations.slice();
         if (relations && effectiveRelations.includes('facetValues')) {
             // We need the facet to determine with the FacetValues are public
@@ -149,10 +145,8 @@ export class ProductService {
             },
         });
         if (!product) {
-            span?.setAttribute('found', 'false');
             return;
         }
-        span?.setAttribute('found', 'true');
         return this.translator.translate(product, ctx, ['facetValues', ['facetValues', 'facet']]);
     }
 

+ 23 - 2
packages/telemetry/src/config/otel-instrumentation-strategy.ts

@@ -1,6 +1,8 @@
 import { Span as ApiSpan, SpanStatusCode, trace } from '@opentelemetry/api';
 import { logs } from '@opentelemetry/api-logs';
-import { InstrumentationStrategy, VENDURE_VERSION, WrappedMethodArgs } from '@vendure/core';
+import { Injector, InstrumentationStrategy, VENDURE_VERSION, WrappedMethodArgs } from '@vendure/core';
+
+import { SetAttributeFn, SpanAttributeService } from '../service/span-attribute.service';
 
 export const tracer = trace.getTracer('@vendure/core', VENDURE_VERSION);
 export const otelLogger = logs.getLogger('@vendure/core', VENDURE_VERSION);
@@ -10,13 +12,28 @@ const recordException = (span: ApiSpan, error: any) => {
 };
 
 export class OtelInstrumentationStrategy implements InstrumentationStrategy {
+    private spanAttributeService: SpanAttributeService;
+
+    init(injector: Injector) {
+        this.spanAttributeService = injector.get(SpanAttributeService);
+    }
+
     wrapMethod({ target, methodName, args, applyOriginalFunction }: WrappedMethodArgs) {
         const spanName = `${String(target.name)}.${String(methodName)}`;
 
         return tracer.startActiveSpan(spanName, {}, span => {
+            const hooks = this.spanAttributeService.getHooks(target, methodName);
+            const setAttribute: SetAttributeFn = (key, value) => span.setAttribute(key, value);
+            hooks?.pre?.(args, setAttribute);
             if (applyOriginalFunction.constructor.name === 'AsyncFunction') {
                 return (
                     applyOriginalFunction()
+                        .then((result: any) => {
+                            if (hooks?.post) {
+                                hooks.post(result, setAttribute);
+                            }
+                            return result;
+                        })
                         // @ts-expect-error
                         .catch(error => {
                             recordException(span, error);
@@ -30,7 +47,11 @@ export class OtelInstrumentationStrategy implements InstrumentationStrategy {
             }
 
             try {
-                return applyOriginalFunction();
+                const result = applyOriginalFunction();
+                if (hooks?.post) {
+                    hooks.post(result, setAttribute);
+                }
+                return result;
             } catch (error) {
                 recordException(span, error);
 

+ 46 - 0
packages/telemetry/src/service/span-attribute.service.ts

@@ -0,0 +1,46 @@
+import { Injectable } from '@nestjs/common';
+import { AttributeValue } from '@opentelemetry/api';
+import { Type } from '@vendure/common/lib/shared-types';
+import { getInstrumentedClassTarget } from '@vendure/core';
+
+/**
+ * Extracts only the method names from a class type T
+ */
+export type MethodNames<T> = {
+    // eslint-disable-next-line @typescript-eslint/ban-types
+    [K in keyof T]: T[K] extends Function ? K : never;
+}[keyof T];
+
+export type Unwrap<T> = T extends Promise<infer U> ? U : T;
+
+export type SetAttributeFn = (key: string, value: AttributeValue) => void;
+
+type MethodType<T, K extends keyof T> = T[K] extends (...args: any[]) => any ? T[K] : never;
+
+export interface InstrumentedMethodHooks<T, Method extends MethodNames<T>> {
+    pre?: (args: Parameters<MethodType<T, Method>>, setAttribute: SetAttributeFn) => void;
+    post?: (result: Unwrap<ReturnType<MethodType<T, Method>>>, setAttribute: SetAttributeFn) => void;
+}
+
+export type MethodHooksForType<T> = {
+    [K in MethodNames<T>]?: InstrumentedMethodHooks<T, K>;
+};
+
+@Injectable()
+export class SpanAttributeService {
+    private hooksMap = new Map<any, { [methodName: string]: InstrumentedMethodHooks<any, any> }>();
+
+    registerHooks<T>(target: Type<T>, hooks: MethodHooksForType<T>): void {
+        const actualTarget = getInstrumentedClassTarget(target) ?? target;
+        const existingHooks = this.hooksMap.get(actualTarget);
+        const combinedHooks = {
+            ...existingHooks,
+            ...hooks,
+        };
+        this.hooksMap.set(actualTarget, combinedHooks);
+    }
+
+    getHooks<T>(target: T, methodName: string): InstrumentedMethodHooks<T, any> | undefined {
+        return this.hooksMap.get(target)?.[methodName];
+    }
+}

+ 17 - 2
packages/telemetry/src/telemetry.plugin.ts

@@ -1,12 +1,27 @@
-import { PluginCommonModule, VendurePlugin } from '@vendure/core';
+import { PluginCommonModule, ProductService, VendurePlugin } from '@vendure/core';
 
 import { OtelInstrumentationStrategy } from './config/otel-instrumentation-strategy';
+import { SpanAttributeService } from './service/span-attribute.service';
 
 @VendurePlugin({
     imports: [PluginCommonModule],
+    providers: [SpanAttributeService],
     configuration: config => {
         config.systemOptions.instrumentationStrategy = new OtelInstrumentationStrategy();
         return config;
     },
 })
-export class TelemetryPlugin {}
+export class TelemetryPlugin {
+    constructor(spanAttributeService: SpanAttributeService) {
+        spanAttributeService.registerHooks(ProductService, {
+            findOne: {
+                pre([ctx, productId], setAttribute) {
+                    setAttribute('productId', productId);
+                },
+                post(product, setAttribute) {
+                    setAttribute('found', !!product);
+                },
+            },
+        });
+    }
+}

+ 1 - 1
packages/telemetry/src/types.ts

@@ -1 +1 @@
-export interface TracingPluginOptions {}
+export interface TelemetryPluginOptions {}