From 3b47e26c3a8d97dd9d0383a15ba91e710bda3413 Mon Sep 17 00:00:00 2001
From: Naseem <naseemkullah@gmail.com>
Date: Wed, 18 Dec 2019 23:34:05 -0500
Subject: [PATCH] feat: trace connection

---
 .../src/ioredis.ts                            | 17 ++++++++++-
 .../opentelemetry-plugin-ioredis/src/utils.ts | 30 +++++++++++++++++++
 .../test/ioredis.test.ts                      | 15 +++++-----
 3 files changed, 54 insertions(+), 8 deletions(-)

diff --git a/packages/opentelemetry-plugin-ioredis/src/ioredis.ts b/packages/opentelemetry-plugin-ioredis/src/ioredis.ts
index adcf69d7bb..4fedf82c53 100644
--- a/packages/opentelemetry-plugin-ioredis/src/ioredis.ts
+++ b/packages/opentelemetry-plugin-ioredis/src/ioredis.ts
@@ -17,7 +17,7 @@
 import { BasePlugin } from '@opentelemetry/core';
 import * as ioredisTypes from 'ioredis';
 import * as shimmer from 'shimmer';
-import { traceSendCommand } from './utils';
+import { traceConnection, traceSendCommand } from './utils';
 
 export class IORedisPlugin extends BasePlugin<typeof ioredisTypes> {
   static readonly COMPONENT = 'ioredis';
@@ -36,12 +36,20 @@ export class IORedisPlugin extends BasePlugin<typeof ioredisTypes> {
       this._patchSendCommand()
     );
 
+    this._logger.debug('patching ioredis.prototype.connect');
+    shimmer.wrap(
+      this._moduleExports.prototype,
+      'connect',
+      this._patchConnection()
+    );
+
     return this._moduleExports.prototype;
   }
 
   protected unpatch(): void {
     if (this._moduleExports) {
       shimmer.unwrap(this._moduleExports.prototype, 'sendCommand');
+      shimmer.unwrap(this._moduleExports.prototype, 'connect');
     }
   }
 
@@ -54,6 +62,13 @@ export class IORedisPlugin extends BasePlugin<typeof ioredisTypes> {
       return traceSendCommand(tracer, original);
     };
   }
+
+  private _patchConnection() {
+    const tracer = this._tracer;
+    return (original: Function) => {
+      return traceConnection(tracer, original);
+    };
+  }
 }
 
 export const plugin = new IORedisPlugin(IORedisPlugin.COMPONENT);
diff --git a/packages/opentelemetry-plugin-ioredis/src/utils.ts b/packages/opentelemetry-plugin-ioredis/src/utils.ts
index c10a1589a2..8dbae4478b 100644
--- a/packages/opentelemetry-plugin-ioredis/src/utils.ts
+++ b/packages/opentelemetry-plugin-ioredis/src/utils.ts
@@ -32,6 +32,36 @@ const endSpan = (span: Span, err: NodeJS.ErrnoException | null | undefined) => {
   span.end();
 };
 
+export const traceConnection = (tracer: Tracer, original: Function) => {
+  return function(this: ioredisTypes.Redis & IORedisPluginClientTypes) {
+    const parentSpan = tracer.getCurrentSpan();
+    const span = tracer.startSpan('connect', {
+      kind: SpanKind.CLIENT,
+      parent: parentSpan,
+      attributes: {
+        [AttributeNames.COMPONENT]: IORedisPlugin.COMPONENT,
+        [AttributeNames.DB_TYPE]: IORedisPlugin.DB_TYPE,
+        [AttributeNames.DB_STATEMENT]: 'connect',
+      },
+    });
+    const { host, port } = this.options;
+
+    span.setAttributes({
+      [AttributeNames.PEER_HOSTNAME]: host,
+      [AttributeNames.PEER_PORT]: port,
+      [AttributeNames.PEER_ADDRESS]: `redis://${host}:${port}`,
+    });
+    try {
+      const client = original.apply(this, arguments);
+      endSpan(span, null);
+      return client;
+    } catch (error) {
+      endSpan(span, error);
+      throw error;
+    }
+  };
+};
+
 export const traceSendCommand = (tracer: Tracer, original: Function) => {
   return function(
     this: ioredisTypes.Redis & IORedisPluginClientTypes,
diff --git a/packages/opentelemetry-plugin-ioredis/test/ioredis.test.ts b/packages/opentelemetry-plugin-ioredis/test/ioredis.test.ts
index 9cbcdf99d1..5ba9aebac7 100644
--- a/packages/opentelemetry-plugin-ioredis/test/ioredis.test.ts
+++ b/packages/opentelemetry-plugin-ioredis/test/ioredis.test.ts
@@ -89,14 +89,15 @@ describe('ioredis', () => {
       let client: ioredisTypes.Redis;
       const attributes = {
         ...DEFAULT_ATTRIBUTES,
-        [AttributeNames.DB_STATEMENT]: 'info',
+        [AttributeNames.DB_STATEMENT]: 'connect',
       };
       const readyHandler = () => {
         const endedSpans = memoryExporter.getFinishedSpans();
 
         assert.strictEqual(tracer.getCurrentSpan(), span);
-        assert.strictEqual(endedSpans.length, 1);
-        assert.strictEqual(endedSpans[0].name, `info`);
+        assert.strictEqual(endedSpans.length, 2);
+        assert.strictEqual(endedSpans[0].name, `connect`);
+        assert.strictEqual(endedSpans[1].name, `info`);
         assertionUtils.assertPropagation(endedSpans[0], span);
 
         assertionUtils.assertSpan(
@@ -107,12 +108,12 @@ describe('ioredis', () => {
           okStatus
         );
         span.end();
-        assert.strictEqual(endedSpans.length, 2);
-        assert.strictEqual(endedSpans[1].name, `test span`);
+        assert.strictEqual(endedSpans.length, 3);
+        assert.strictEqual(endedSpans[2].name, `test span`);
 
         client.quit(done);
-        assert.strictEqual(endedSpans.length, 3);
-        assert.strictEqual(endedSpans[2].name, `quit`);
+        assert.strictEqual(endedSpans.length, 4);
+        assert.strictEqual(endedSpans[3].name, `quit`);
       };
       const errorHandler = (err: Error) => {
         assert.ifError(err);