feat: add redis 4 connect span (#1125)

* feat: trace redis connect * test: fix typo in test name * fix: remove code duplication, record exception, fix span name * refactor: don't create unnecessary promises Co-authored-by: Rauno Viskus <Rauno56@users.noreply.github.com>
open-telemetry · Sep 13, 2022 · dbf37fb · dbf37fb
1 parent 82b8a84
commit dbf37fb
Show file tree

Hide file tree

Showing 3 changed files with 115 additions and 14 deletions.
diff --git a/plugins/node/opentelemetry-instrumentation-redis-4/src/instrumentation.ts b/plugins/node/opentelemetry-instrumentation-redis-4/src/instrumentation.ts
@@ -28,13 +28,10 @@ import {
   InstrumentationNodeModuleDefinition,
   InstrumentationNodeModuleFile,
 } from '@opentelemetry/instrumentation';
-import { defaultDbStatementSerializer } from './utils';
+import { defaultDbStatementSerializer, getClientAttributes } from './utils';
 import { RedisInstrumentationConfig } from './types';
 import { VERSION } from './version';
-import {
-  DbSystemValues,
-  SemanticAttributes,
-} from '@opentelemetry/semantic-conventions';
+import { SemanticAttributes } from '@opentelemetry/semantic-conventions';
 
 const OTEL_OPEN_SPANS = Symbol(
   'opentelemetry.instruemntation.redis.open_spans'
@@ -186,6 +183,12 @@ export class RedisInstrumentation extends InstrumentationBase<any> {
           this._getPatchRedisClientSendCommand()
         );
 
+        this._wrap(
+          redisClientPrototype,
+          'connect',
+          this._getPatchedClientConnect()
+        );
+
         return moduleExports;
       },
       (moduleExports: any) => {
@@ -332,6 +335,44 @@ export class RedisInstrumentation extends InstrumentationBase<any> {
     };
   }
 
+  private _getPatchedClientConnect() {
+    const plugin = this;
+    return function connectWrapper(original: Function) {
+      return function patchedConnect(this: any): Promise<void> {
+        const options = this.options;
+
+        const attributes = getClientAttributes(options);
+
+        const span = plugin.tracer.startSpan(
+          `${RedisInstrumentation.COMPONENT}-connect`,
+          {
+            kind: SpanKind.CLIENT,
+            attributes,
+          }
+        );
+
+        const res = context.with(trace.setSpan(context.active(), span), () => {
+          return original.apply(this);
+        });
+
+        return res
+          .then((result: unknown) => {
+            span.end();
+            return result;
+          })
+          .catch((error: Error) => {
+            span.recordException(error);
+            span.setStatus({
+              code: SpanStatusCode.ERROR,
+              message: error.message,
+            });
+            span.end();
+            return Promise.reject(error);
+          });
+      };
+    };
+  }
+
   private _traceClientCommand(
     origFunction: Function,
     origThis: any,
@@ -351,12 +392,7 @@ export class RedisInstrumentation extends InstrumentationBase<any> {
     const dbStatementSerializer =
       this._config?.dbStatementSerializer || defaultDbStatementSerializer;
 
-    const attributes = {
-      [SemanticAttributes.DB_SYSTEM]: DbSystemValues.REDIS,
-      [SemanticAttributes.NET_PEER_NAME]: clientOptions?.socket?.host,
-      [SemanticAttributes.NET_PEER_PORT]: clientOptions?.socket?.port,
-      [SemanticAttributes.DB_CONNECTION_STRING]: clientOptions?.url,
-    };
+    const attributes = getClientAttributes(clientOptions);
 
     try {
       const dbStatement = dbStatementSerializer(commandName, commandArgs);
@@ -377,7 +413,9 @@ export class RedisInstrumentation extends InstrumentationBase<any> {
       }
     );
 
-    const res = origFunction.apply(origThis, origArguments);
+    const res = context.with(trace.setSpan(context.active(), span), () => {
+      return origFunction.apply(origThis, origArguments);
+    });
     if (typeof res?.then === 'function') {
       res.then(
         (redisRes: unknown) => {

diff --git a/plugins/node/opentelemetry-instrumentation-redis-4/src/utils.ts b/plugins/node/opentelemetry-instrumentation-redis-4/src/utils.ts
@@ -14,6 +14,19 @@
  * limitations under the License.
  */
 import { DbStatementSerializer } from './types';
+import {
+  DbSystemValues,
+  SemanticAttributes,
+} from '@opentelemetry/semantic-conventions';
 
 export const defaultDbStatementSerializer: DbStatementSerializer = cmdName =>
   cmdName;
+
+export function getClientAttributes(options: any) {
+  return {
+    [SemanticAttributes.DB_SYSTEM]: DbSystemValues.REDIS,
+    [SemanticAttributes.NET_PEER_NAME]: options?.socket?.host,
+    [SemanticAttributes.NET_PEER_PORT]: options?.socket?.port,
+    [SemanticAttributes.DB_CONNECTION_STRING]: options?.url,
+  };
+}
diff --git a/plugins/node/opentelemetry-instrumentation-redis-4/test/redis.test.ts b/plugins/node/opentelemetry-instrumentation-redis-4/test/redis.test.ts
@@ -42,7 +42,7 @@ import {
 } from '@opentelemetry/api';
 import { SemanticAttributes } from '@opentelemetry/semantic-conventions';
 import { RedisResponseCustomAttributeFunction } from '../src/types';
-import { hrTimeToMilliseconds } from '@opentelemetry/core';
+import { hrTimeToMilliseconds, suppressTracing } from '@opentelemetry/core';
 
 describe('redis@^4.0.0', () => {
   before(function () {
@@ -71,7 +71,9 @@ describe('redis@^4.0.0', () => {
     client = createClient({
       url: redisTestUrl,
     });
-    await client.connect();
+    context.with(suppressTracing(context.active()), async () => {
+      await client.connect();
+    });
   });
 
   afterEach(async () => {
@@ -183,6 +185,54 @@ describe('redis@^4.0.0', () => {
     });
   });
 
+  describe('client connect', () => {
+    it('produces a span', async () => {
+      const newClient = createClient({
+        url: redisTestUrl,
+      });
+
+      after(async () => {
+        await newClient.disconnect();
+      });
+
+      await newClient.connect();
+
+      const [span] = getTestSpans();
+
+      assert.strictEqual(span.name, 'redis-connect');
+
+      assert.strictEqual(
+        span.attributes[SemanticAttributes.DB_SYSTEM],
+        'redis'
+      );
+      assert.strictEqual(
+        span.attributes[SemanticAttributes.NET_PEER_NAME],
+        redisTestConfig.host
+      );
+      assert.strictEqual(
+        span.attributes[SemanticAttributes.NET_PEER_PORT],
+        redisTestConfig.port
+      );
+      assert.strictEqual(
+        span.attributes[SemanticAttributes.DB_CONNECTION_STRING],
+        redisTestUrl
+      );
+    });
+
+    it('sets error status on connection failure', async () => {
+      const newClient = createClient({
+        url: `redis://${redisTestConfig.host}:${redisTestConfig.port + 1}`,
+      });
+
+      await assert.rejects(newClient.connect());
+
+      const [span] = getTestSpans();
+
+      assert.strictEqual(span.name, 'redis-connect');
+      assert.strictEqual(span.status.code, SpanStatusCode.ERROR);
+    });
+  });
+
   describe('multi (transactions) commands', () => {
     it('multi commands', async () => {
       await client.set('another-key', 'another-value');