fix(agent): Fallback methods for loading traced plugins from layers

Author: mrickard

src/index.js

@@ -1,9 +1,8 @@
 import Perf from 'performance-node';
-
 import pkg from '../package.json'; // eslint-disable-line import/extensions
 import { addToReport, addTraceData } from './addToReport';
 
-const load = plugin => {
+const loadPlugin = plugin => {
   /*eslint-disable camelcase, no-undef*/
   if (typeof __non_webpack_require__ === 'function') {
     return __non_webpack_require__(`./plugins/${plugin}`);
@@ -144,7 +143,7 @@ class TracePlugin {
 
       if (context.config[conf] && context.config[conf].enabled) {
         // getting plugin; allows this to be loaded only if enabled.
-        await load(`${k}`).then(mod => {
+        await loadPlugin(`${k}`).then(mod => {
           plugins[k].wrap = mod.wrap;
           plugins[k].unwrap = mod.unwrap;
           context[namespace] = {
@@ -153,7 +152,7 @@ class TracePlugin {
             data: {},
             config: context.config[conf]
           };
-          plugins[k].wrap(context[namespace]);
+          return plugins[k].wrap(context[namespace]);
         });
       }
     });

src/loadHelper.js

@@ -0,0 +1,52 @@
+import { debuglog } from 'util';
+
+const debug = debuglog('@iopipe:trace:loadHelper');
+
+// The module being traced might not be in the lambda NODE_PATH,
+// particularly if IOpipe has been installed with lambda layers.
+// So here's an attempt at a fallback.
+
+const deriveTargetPath = manualPath => {
+  const path = manualPath ? manualPath : process.env.LAMBDA_TASK_ROOT;
+  if (!path) {
+    return false;
+  }
+  const targetPath = `${path}/node_modules`;
+  process.env.IOPIPE_TARGET_PATH = targetPath;
+  return targetPath;
+};
+
+const appendToPath = manualPath => {
+  if (!process.env || !process.env.NODE_PATH) {
+    return false;
+  }
+  const targetPath = deriveTargetPath(manualPath);
+  const pathArray = process.env.NODE_PATH.split(':');
+  if (!targetPath && pathArray.indexOf(targetPath) === -1) {
+    process.env.NODE_PATH = `${process.env.NODE_PATH}:${targetPath}`;
+  }
+  return targetPath;
+};
+
+appendToPath();
+
+const loadModuleForTracing = async (module, path) => {
+  const targetPath = path ? path : process.env.IOPIPE_TARGET_PATH;
+  let loadedModule;
+  try {
+    loadedModule = await require(module);
+  } catch (e) {
+    debug(`Unable to load ${module} from require; trying TASK_ROOT path.`, e);
+    try {
+      if (targetPath) {
+        loadedModule = await require(`${targetPath}/${module}`);
+      }
+    } catch (err) {
+      debug(`Unable to load ${module} from path ${targetPath}`, err);
+    }
+  }
+
+  return loadedModule;
+};
+
+export default loadModuleForTracing;

src/plugins/https.js

@@ -9,7 +9,7 @@ import pickBy from 'lodash.pickby';
 import isArray from 'isarray';
 import { flatten } from 'flat';
 
-const debug = debuglog('@iopipe/trace');
+const debug = debuglog('@iopipe:trace:https');
 
 /*eslint-disable babel/no-invalid-this*/
 

src/plugins/ioredis.js

@@ -1,10 +1,27 @@
 import { debuglog } from 'util';
 import shimmer from 'shimmer';
-import Redis from 'ioredis';
 import Perf from 'performance-node';
 import uuid from 'uuid/v4';
+import loadModuleForTracing from '../loadHelper';
 
-const debug = debuglog('@iopipe/trace');
+const debug = debuglog('@iopipe:trace:ioredis');
+
+let Redis, RedisTarget, RedisCmdTarget;
+
+const loadModule = () =>
+  loadModuleForTracing('ioredis')
+    .then(module => {
+      Redis = module;
+      RedisTarget = Redis && Redis.prototype;
+      if (Redis.Command && Redis.Command.prototype) {
+        RedisCmdTarget = Redis.Command && Redis.Command.prototype;
+      }
+      return module;
+    })
+    .catch(e => {
+      debug('Not loading ioredis', e);
+      return false;
+    });
 
 /*eslint-disable babel/no-invalid-this*/
 /*eslint-disable func-name-matching */
@@ -28,23 +45,24 @@ const filterRequest = (command, context) => {
   };
 };
 
-function wrap({ timeline, data = {} } = {}) {
+async function wrap({ timeline, data = {} } = {}) {
+  await loadModule();
+  if (!Redis) {
+    debug('ioredis plugin not accessible from trace plugin. Skipping.');
+    return false;
+  }
   if (!(timeline instanceof Perf)) {
     debug(
       'Timeline passed to plugins/ioredis.wrap not an instance of performance-node. Skipping.'
     );
     return false;
   }
 
-  if (!Redis.__iopipeShimmer) {
+  if (Redis && !Redis.__iopipeShimmer) {
     if (process.env.IOPIPE_TRACE_IOREDIS_INITPROMISE) {
-      shimmer.wrap(
-        Redis.Command && Redis.Command.prototype,
-        'initPromise',
-        wrapPromise
-      );
+      shimmer.wrap(RedisCmdTarget, 'initPromise', wrapPromise);
     }
-    shimmer.wrap(Redis && Redis.prototype, 'sendCommand', wrapSendCommand);
+    shimmer.wrap(RedisTarget, 'sendCommand', wrapSendCommand);
     Redis.__iopipeShimmer = true;
   }
 
@@ -128,11 +146,18 @@ function wrap({ timeline, data = {} } = {}) {
 }
 
 function unwrap() {
+  if (!Redis) {
+    debug(
+      'ioredis plugin not accessible from trace plugin. Nothing to unwrap.'
+    );
+    return false;
+  }
   if (process.env.IOPIPE_TRACE_IOREDIS_INITPROMISE) {
-    shimmer.unwrap(Redis.Command && Redis.Command.prototype, 'initPromise');
+    shimmer.unwrap(RedisCmdTarget, 'initPromise');
   }
-  shimmer.unwrap(Redis && Redis.prototype, 'sendCommand');
+  shimmer.unwrap(RedisTarget, 'sendCommand');
   delete Redis.__iopipeShimmer;
+  return true;
 }
 
 export { unwrap, wrap };

src/plugins/ioredis.test.js

@@ -64,8 +64,8 @@ xtest('Redis works as normal if wrap is not called', done => {
     });
 });
 
-test('Bails if timeline is not instance of performance-node', () => {
-  const bool = wrap({ timeline: [] });
+test('Bails if timeline is not instance of performance-node', async () => {
+  const bool = await wrap({ timeline: [] });
   expect(bool).toBe(false);
 });
 
@@ -128,7 +128,7 @@ xdescribe('Wrapping Redis', () => {
         db: 1
       });
 
-      expect(redis.sendCommand.__wrapped).toBeDefined();
+      expect(redis.__wrapped).toBeDefined();
 
       const expectedStr = 'wrapped ioredis, async/await';
 
@@ -152,7 +152,7 @@ xdescribe('Wrapping Redis', () => {
 
       redis = new Redis({ host: '127.0.0.1', connectionName: 'Test 2', db: 1 });
 
-      expect(redis.sendCommand.__wrapped).toBeDefined();
+      expect(redis.__wrapped).toBeDefined();
 
       const expectedStr = 'wrapped ioredis, promise syntax';
 
@@ -185,7 +185,7 @@ xdescribe('Wrapping Redis', () => {
 
       redis = new Redis({ host: 'localhost', connectionName: 'Test 3', db: 1 });
 
-      expect(redis.sendCommand.__wrapped).toBeDefined();
+      expect(redis.__wrapped).toBeDefined();
 
       const expectedStr = 'wrapped ioredis, callback syntax';
       redis.set('testString', expectedStr);

src/plugins/mongodb.js

@@ -1,9 +1,42 @@
 import { debuglog } from 'util';
 import shimmer from 'shimmer';
-import { MongoClient, Server, Cursor, Collection } from 'mongodb';
 import Perf from 'performance-node';
 import uuid from 'uuid/v4';
 import get from 'lodash/get';
+import loadModuleForTracing from '../loadHelper';
+
+const debug = debuglog('@iopipe:trace:mongodb');
+
+let MongoClient,
+  Server,
+  Cursor,
+  Collection,
+  clientTarget,
+  collectionTarget,
+  serverTarget,
+  cursorTarget;
+
+const loadModule = async () => {
+  const mod = await loadModuleForTracing('mongodb')
+    .then(module => {
+      MongoClient = module.MongoClient;
+      Server = module.Server;
+      Cursor = module.Cursor;
+      Collection = module.Collection;
+
+      clientTarget = MongoClient && MongoClient.prototype;
+      collectionTarget = Collection && Collection.prototype;
+      serverTarget = Server && Server.prototype;
+      cursorTarget = Cursor && Cursor.prototype;
+
+      return module;
+    })
+    .catch(e => {
+      debug('Not loading mongodb', e);
+      return null;
+    });
+  return mod;
+};
 
 const dbType = 'mongodb';
 const serverOps = ['command', 'insert', 'update', 'remove'];
@@ -23,13 +56,6 @@ const collectionOps = [
 const cursorOps = ['next', 'filter', 'sort', 'hint', 'toArray'];
 const clientOps = ['connect', 'close', 'db'];
 
-const clientTarget = MongoClient && MongoClient.prototype;
-const collectionTarget = Collection && Collection.prototype;
-const serverTarget = Server && Server.prototype;
-const cursorTarget = Cursor && Cursor.prototype;
-
-const debug = debuglog('@iopipe/trace');
-
 /*eslint-disable babel/no-invalid-this*/
 /*eslint-disable func-name-matching */
 /*eslint-disable prefer-rest-params */
@@ -146,7 +172,13 @@ const filterRequest = (params, context) => {
   };
 };
 
-function wrap({ timeline, data = {} } = {}) {
+async function wrap({ timeline, data = {} } = {}) {
+  await loadModule();
+
+  if (!clientTarget) {
+    debug('mongodb plugin not accessible from trace plugin. Skipping.');
+    return false;
+  }
   if (!(timeline instanceof Perf)) {
     debug(
       'Timeline passed to plugins/mongodb.wrap not an instance of performance-node. Skipping.'
@@ -217,6 +249,13 @@ function wrap({ timeline, data = {} } = {}) {
 }
 
 function unwrap() {
+  if (!clientTarget) {
+    debug(
+      'mongodb plugin not accessible from trace plugin. Nothing to unwrap.'
+    );
+    return false;
+  }
+
   if (serverTarget.__iopipeShimmer) {
     shimmer.massUnwrap(serverTarget, serverOps);
     delete serverTarget.__iopipeShimmer;
@@ -233,6 +272,7 @@ function unwrap() {
     shimmer.massUnwrap(clientTarget, clientOps); // mass just seems to hang and not complete
     delete clientTarget.__iopipeShimmer;
   }
+  return true;
 }
 
 export { unwrap, wrap };