fix(redis): use new tracer after setTracerProvider

blumamir · blumamir · commit 64c1b7a99b62 · 2025-05-29T06:40:54.000+07:00
diff --git a/plugins/node/opentelemetry-instrumentation-redis/src/instrumentation.ts b/plugins/node/opentelemetry-instrumentation-redis/src/instrumentation.ts
@@ -18,16 +18,28 @@ import {
   isWrapped,
   InstrumentationBase,
   InstrumentationNodeModuleDefinition,
+  safeExecuteInTheMiddle,
 } from '@opentelemetry/instrumentation';
 import {
+    endSpan,
   getTracedCreateClient,
   getTracedCreateStreamTrace,
-  getTracedInternalSendCommand,
 } from './utils';
-import { RedisInstrumentationConfig } from './types';
+import { RedisCommand, RedisInstrumentationConfig } from './types';
 /** @knipignore */
 import { PACKAGE_NAME, PACKAGE_VERSION } from './version';
-
+import { RedisPluginClientTypes } from './internal-types';
+import { SpanKind, context, trace } from '@opentelemetry/api';
+import {
+    DBSYSTEMVALUES_REDIS,
+    SEMATTRS_DB_CONNECTION_STRING,
+    SEMATTRS_DB_STATEMENT,
+    SEMATTRS_DB_SYSTEM,
+    SEMATTRS_NET_PEER_NAME,
+    SEMATTRS_NET_PEER_PORT,
+  } from '@opentelemetry/semantic-conventions';
+import { defaultDbStatementSerializer } from '@opentelemetry/redis-common';
+  
 const DEFAULT_CONFIG: RedisInstrumentationConfig = {
   requireParentSpan: false,
 };
@@ -96,28 +108,110 @@ export class RedisInstrumentation extends InstrumentationBase<RedisInstrumentati
       ),
     ];
   }
+
   /**
    * Patch internal_send_command(...) to trace requests
    */
   private _getPatchInternalSendCommand() {
-    const tracer = this.tracer;
-    const config = this.getConfig();
+    const instrumentation = this;
     return function internal_send_command(original: Function) {
-      return getTracedInternalSendCommand(tracer, original, config);
+        return function internal_send_command_trace(
+            this: RedisPluginClientTypes,
+            cmd?: RedisCommand
+          ) {
+            // Versions of redis (2.4+) use a single options object
+            // instead of named arguments
+            if (arguments.length !== 1 || typeof cmd !== 'object') {
+              // We don't know how to trace this call, so don't start/stop a span
+              return original.apply(this, arguments);
+            }
+
+            const config = instrumentation.getConfig();
+        
+            const hasNoParentSpan = trace.getSpan(context.active()) === undefined;
+            if (config.requireParentSpan === true && hasNoParentSpan) {
+              return original.apply(this, arguments);
+            }
+        
+            const dbStatementSerializer =
+              config?.dbStatementSerializer || defaultDbStatementSerializer;
+            const span = instrumentation.tracer.startSpan(
+              `${RedisInstrumentation.COMPONENT}-${cmd.command}`,
+              {
+                kind: SpanKind.CLIENT,
+                attributes: {
+                  [SEMATTRS_DB_SYSTEM]: DBSYSTEMVALUES_REDIS,
+                  [SEMATTRS_DB_STATEMENT]: dbStatementSerializer(cmd.command, cmd.args),
+                },
+              }
+            );
+        
+            // Set attributes for not explicitly typed RedisPluginClientTypes
+            if (this.connection_options) {
+              span.setAttributes({
+                [SEMATTRS_NET_PEER_NAME]: this.connection_options.host,
+                [SEMATTRS_NET_PEER_PORT]: this.connection_options.port,
+              });
+            }
+            if (this.address) {
+              span.setAttribute(
+                SEMATTRS_DB_CONNECTION_STRING,
+                `redis://${this.address}`
+              );
+            }
+        
+            const originalCallback = arguments[0].callback;
+            if (originalCallback) {
+              const originalContext = context.active();
+              (arguments[0] as RedisCommand).callback = function callback<T>(
+                this: unknown,
+                err: Error | null,
+                reply: T
+              ) {
+                if (config?.responseHook) {
+                  const responseHook = config.responseHook;
+                  safeExecuteInTheMiddle(
+                    () => {
+                      responseHook(span, cmd.command, cmd.args, reply);
+                    },
+                    err => {
+                      if (err) {
+                        instrumentation._diag.error('Error executing responseHook', err);
+                      }
+                    },
+                    true
+                  );
+                }
+        
+                endSpan(span, err);
+                return context.with(
+                  originalContext,
+                  originalCallback,
+                  this,
+                  ...arguments
+                );
+              };
+            }
+            try {
+              // Span will be ended in callback
+              return original.apply(this, arguments);
+            } catch (rethrow: any) {
+              endSpan(span, rethrow);
+              throw rethrow; // rethrow after ending span
+            }
+          };
     };
   }
 
   private _getPatchCreateClient() {
-    const tracer = this.tracer;
     return function createClient(original: Function) {
-      return getTracedCreateClient(tracer, original);
+      return getTracedCreateClient(original);
     };
   }
 
   private _getPatchCreateStream() {
-    const tracer = this.tracer;
     return function createReadStream(original: Function) {
-      return getTracedCreateStreamTrace(tracer, original);
+      return getTracedCreateStreamTrace(original);
     };
   }
 }
diff --git a/plugins/node/opentelemetry-instrumentation-redis/src/utils.ts b/plugins/node/opentelemetry-instrumentation-redis/src/utils.ts
@@ -17,29 +17,12 @@
 import type * as redisTypes from 'redis';
 import {
   context,
-  Tracer,
-  SpanKind,
   Span,
   SpanStatusCode,
-  trace,
-  diag,
 } from '@opentelemetry/api';
-import { RedisCommand, RedisInstrumentationConfig } from './types';
 import { EventEmitter } from 'events';
-import { RedisInstrumentation } from './';
-import {
-  DBSYSTEMVALUES_REDIS,
-  SEMATTRS_DB_CONNECTION_STRING,
-  SEMATTRS_DB_STATEMENT,
-  SEMATTRS_DB_SYSTEM,
-  SEMATTRS_NET_PEER_NAME,
-  SEMATTRS_NET_PEER_PORT,
-} from '@opentelemetry/semantic-conventions';
-import { safeExecuteInTheMiddle } from '@opentelemetry/instrumentation';
-import { RedisPluginClientTypes } from './internal-types';
-import { defaultDbStatementSerializer } from '@opentelemetry/redis-common';
 
-const endSpan = (span: Span, err?: Error | null) => {
+export const endSpan = (span: Span, err?: Error | null) => {
   if (err) {
     span.setStatus({
       code: SpanStatusCode.ERROR,
@@ -49,15 +32,14 @@ const endSpan = (span: Span, err?: Error | null) => {
   span.end();
 };
 
-export const getTracedCreateClient = (tracer: Tracer, original: Function) => {
+export const getTracedCreateClient = (original: Function) => {
   return function createClientTrace(this: redisTypes.RedisClient) {
     const client: redisTypes.RedisClient = original.apply(this, arguments);
     return context.bind(context.active(), client);
   };
 };
 
 export const getTracedCreateStreamTrace = (
-  tracer: Tracer,
   original: Function
 ) => {
   return function create_stream_trace(this: redisTypes.RedisClient) {
@@ -75,93 +57,3 @@ export const getTracedCreateStreamTrace = (
     return original.apply(this, arguments);
   };
 };
-
-export const getTracedInternalSendCommand = (
-  tracer: Tracer,
-  original: Function,
-  config?: RedisInstrumentationConfig
-) => {
-  return function internal_send_command_trace(
-    this: RedisPluginClientTypes,
-    cmd?: RedisCommand
-  ) {
-    // New versions of redis (2.4+) use a single options object
-    // instead of named arguments
-    if (arguments.length !== 1 || typeof cmd !== 'object') {
-      // We don't know how to trace this call, so don't start/stop a span
-      return original.apply(this, arguments);
-    }
-
-    const hasNoParentSpan = trace.getSpan(context.active()) === undefined;
-    if (config?.requireParentSpan === true && hasNoParentSpan) {
-      return original.apply(this, arguments);
-    }
-
-    const dbStatementSerializer =
-      config?.dbStatementSerializer || defaultDbStatementSerializer;
-    const span = tracer.startSpan(
-      `${RedisInstrumentation.COMPONENT}-${cmd.command}`,
-      {
-        kind: SpanKind.CLIENT,
-        attributes: {
-          [SEMATTRS_DB_SYSTEM]: DBSYSTEMVALUES_REDIS,
-          [SEMATTRS_DB_STATEMENT]: dbStatementSerializer(cmd.command, cmd.args),
-        },
-      }
-    );
-
-    // Set attributes for not explicitly typed RedisPluginClientTypes
-    if (this.connection_options) {
-      span.setAttributes({
-        [SEMATTRS_NET_PEER_NAME]: this.connection_options.host,
-        [SEMATTRS_NET_PEER_PORT]: this.connection_options.port,
-      });
-    }
-    if (this.address) {
-      span.setAttribute(
-        SEMATTRS_DB_CONNECTION_STRING,
-        `redis://${this.address}`
-      );
-    }
-
-    const originalCallback = arguments[0].callback;
-    if (originalCallback) {
-      const originalContext = context.active();
-      (arguments[0] as RedisCommand).callback = function callback<T>(
-        this: unknown,
-        err: Error | null,
-        reply: T
-      ) {
-        if (config?.responseHook) {
-          const responseHook = config.responseHook;
-          safeExecuteInTheMiddle(
-            () => {
-              responseHook(span, cmd.command, cmd.args, reply);
-            },
-            err => {
-              if (err) {
-                diag.error('Error executing responseHook', err);
-              }
-            },
-            true
-          );
-        }
-
-        endSpan(span, err);
-        return context.with(
-          originalContext,
-          originalCallback,
-          this,
-          ...arguments
-        );
-      };
-    }
-    try {
-      // Span will be ended in callback
-      return original.apply(this, arguments);
-    } catch (rethrow: any) {
-      endSpan(span, rethrow);
-      throw rethrow; // rethrow after ending span
-    }
-  };
-};
diff --git a/plugins/node/opentelemetry-instrumentation-redis/test/redis.test.ts b/plugins/node/opentelemetry-instrumentation-redis/test/redis.test.ts
@@ -80,6 +80,9 @@ describe('redis@2.x', () => {
   beforeEach(() => {
     contextManager = new AsyncHooksContextManager().enable();
     context.setGlobalContextManager(contextManager);
+    // set the default tracer provider before each test
+    // specific ones can override it to assert certain things
+    instrumentation.setTracerProvider(provider);
   });
 
   afterEach(() => {
@@ -141,30 +144,30 @@ describe('redis@2.x', () => {
       expectedDbStatement: string;
       method: (cb: redisTypes.Callback<unknown>) => unknown;
     }> = [
-      {
-        description: 'insert',
-        command: 'hset',
-        args: ['hash', 'random', 'random'],
-        expectedDbStatement: 'hash random [1 other arguments]',
-        method: (cb: redisTypes.Callback<number>) =>
-          client.hset('hash', 'random', 'random', cb),
-      },
-      {
-        description: 'get',
-        command: 'get',
-        args: ['test'],
-        expectedDbStatement: 'test',
-        method: (cb: redisTypes.Callback<string | null>) =>
-          client.get('test', cb),
-      },
-      {
-        description: 'delete',
-        command: 'del',
-        args: ['test'],
-        expectedDbStatement: 'test',
-        method: (cb: redisTypes.Callback<number>) => client.del('test', cb),
-      },
-    ];
+        {
+          description: 'insert',
+          command: 'hset',
+          args: ['hash', 'random', 'random'],
+          expectedDbStatement: 'hash random [1 other arguments]',
+          method: (cb: redisTypes.Callback<number>) =>
+            client.hset('hash', 'random', 'random', cb),
+        },
+        {
+          description: 'get',
+          command: 'get',
+          args: ['test'],
+          expectedDbStatement: 'test',
+          method: (cb: redisTypes.Callback<string | null>) =>
+            client.get('test', cb),
+        },
+        {
+          description: 'delete',
+          command: 'del',
+          args: ['test'],
+          expectedDbStatement: 'test',
+          method: (cb: redisTypes.Callback<number>) => client.del('test', cb),
+        },
+      ];
 
     before(done => {
       client = redis.createClient(URL);
@@ -389,5 +392,34 @@ describe('redis@2.x', () => {
         });
       });
     });
+
+    describe('setTracerProvider', () => {
+      before(() => {
+        instrumentation.disable();
+        instrumentation.setConfig({});
+        instrumentation.enable();
+      });
+
+      it('should use new tracer provider after setTracerProvider is called', done => {
+        const testSpecificMemoryExporter = new InMemorySpanExporter();
+        const spanProcessor = new SimpleSpanProcessor(testSpecificMemoryExporter);
+        const tracerProvider = new NodeTracerProvider({
+          spanProcessors: [spanProcessor],
+        });
+
+        // key point of this test, setting new tracer provider and making sure
+        // new spans use it.
+        instrumentation.setTracerProvider(tracerProvider);
+
+        client.set('foo', 'bar-value-from-test', (err) => {
+          assert.ifError(err);
+          // assert that the span was exported by the new tracer provider
+          // which is using the test specific span processor
+          const spans = testSpecificMemoryExporter.getFinishedSpans();
+          assert.strictEqual(spans.length, 1);
+          done();
+        })
+      });
+    });
   });
 });
