feat: support context propagation in bluebird (#872)

src/config.ts

@@ -255,14 +255,14 @@ export const defaultConfig = {
   maximumLabelValueSize: 512,
   plugins: {
     // enable all by default
+    'bluebird': path.join(pluginDirectory, 'plugin-bluebird.js'),
     'connect': path.join(pluginDirectory, 'plugin-connect.js'),
     'express': path.join(pluginDirectory, 'plugin-express.js'),
     'generic-pool': path.join(pluginDirectory, 'plugin-generic-pool.js'),
     'grpc': path.join(pluginDirectory, 'plugin-grpc.js'),
     'hapi': path.join(pluginDirectory, 'plugin-hapi.js'),
     'http': path.join(pluginDirectory, 'plugin-http.js'),
     'http2': path.join(pluginDirectory, 'plugin-http2.js'),
-    'knex': path.join(pluginDirectory, 'plugin-knex.js'),
     'koa': path.join(pluginDirectory, 'plugin-koa.js'),
     'mongodb-core': path.join(pluginDirectory, 'plugin-mongodb-core.js'),
     'mysql': path.join(pluginDirectory, 'plugin-mysql.js'),

src/plugins/plugin-bluebird.ts

@@ -0,0 +1,46 @@
+/**
+ * Copyright 2018 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import * as shimmer from 'shimmer';
+
+import {PluginTypes} from '..';
+
+import {bluebird_3} from './types';
+
+type BluebirdModule = typeof bluebird_3&{prototype: {_then: Function;}};
+
+const plugin: PluginTypes.Plugin = [{
+  // Bluebird is a class.
+  // tslint:disable-next-line:variable-name
+  patch: (Bluebird, tracer) => {
+    // any is a type arg; args are type checked when read directly, otherwise
+    // passed through to a function with the same type signature.
+    // tslint:disable:no-any
+    const wrapIfFunction = (fn: any) =>
+        typeof fn === 'function' ? tracer.wrap(fn) : fn;
+    shimmer.wrap(Bluebird.prototype, '_then', (thenFn: Function) => {
+      // Inherit context from the call site of .then().
+      return function<T>(this: bluebird_3<T>, ...args: any[]) {
+        return thenFn.apply(this, [
+          wrapIfFunction(args[0]), wrapIfFunction(args[1]), ...args.slice(2)
+        ]);
+      };
+    });
+    // tslint:enable:no-any
+  }
+} as PluginTypes.Monkeypatch<BluebirdModule>];
+
+export = plugin;

src/plugins/types/index.d.ts

@@ -27,6 +27,7 @@
  * contain dots.
  */
 
+import * as bluebird_3 from './bluebird_3'; // bluebird@3
 import * as connect_3 from './connect_3'; // connect@3
 import * as express_4 from './express_4'; // express@4
 import * as hapi_16 from './hapi_16'; // hapi@16
@@ -86,6 +87,7 @@ declare namespace pg_6 {
 //---exports---//
 
 export {
+  bluebird_3,
   connect_3,
   express_4,
   hapi_16,

test/fixtures/plugin-fixtures.json

@@ -1,4 +1,9 @@
 {
+  "bluebird3": {
+    "dependencies": {
+      "bluebird": "^3.5.2"
+    }
+  },
   "connect3": {
     "dependencies": {
       "connect": "^3.5.0"

test/plugins/test-cls-bluebird.ts

@@ -0,0 +1,177 @@
+/**
+ * Copyright 2019 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import * as assert from 'assert';
+
+import {bluebird_3 as BluebirdPromise} from '../../src/plugins/types';
+import {Trace} from '../../src/trace';
+import * as traceTestModule from '../trace';
+
+/**
+ * Describes a test case.
+ */
+interface TestCase<T = void> {
+  /**
+   * Description of a test case; included in string argument to it().
+   */
+  description: string;
+  /**
+   * Creates and returns a new Promise.
+   */
+  makePromise: () => BluebirdPromise<T>;
+  /**
+   * Given a Promise and a callback, calls the callback some time after the
+   * Promise has been resolved or rejected.
+   */
+  thenFn: (promise: BluebirdPromise<T>, cb: () => void) => void;
+}
+/**
+ * For a given Promise implementation, create two traces:
+ * 1. Constructs a new Promise and resolves it.
+ * 2. Within a then callback to the above mentioned Promise, construct a child
+ *    span.
+ */
+const getTracesForPromiseImplementation =
+    <T>(makePromise: () => BluebirdPromise<T>,
+        thenFn: (promise: BluebirdPromise<T>, cb: () => void) =>
+            void): Promise<[Trace, Trace]> => {
+      return new Promise((resolve, reject) => {
+        const tracer = traceTestModule.get();
+        let p: BluebirdPromise<T>;
+        const firstSpan = tracer.runInRootSpan({name: 'first'}, span => {
+          p = makePromise();
+          return span;
+        });
+        tracer.runInRootSpan({name: 'second'}, secondSpan => {
+          // Note to maintainers: Do NOT convert this to async/await,
+          // as it changes context propagation behavior.
+          thenFn(p, () => {
+            tracer.createChildSpan().endSpan();
+            secondSpan.endSpan();
+            firstSpan.endSpan();
+            setImmediate(() => {
+              try {
+                const trace1 = traceTestModule.getOneTrace(
+                    trace => trace.spans.some(root => root.name === 'first'));
+                const trace2 = traceTestModule.getOneTrace(
+                    trace => trace.spans.some(root => root.name === 'second'));
+                traceTestModule.clearTraceData();
+                resolve([trace1, trace2]);
+              } catch (e) {
+                traceTestModule.clearTraceData();
+                reject(e);
+              }
+            });
+          });
+        });
+      });
+    };
+
+
+describe('Patch plugin for bluebird', () => {
+  // BPromise is a class.
+  // tslint:disable-next-line:variable-name
+  let BPromise: typeof BluebirdPromise;
+
+  before(() => {
+    traceTestModule.setCLSForTest();
+    traceTestModule.setPluginLoaderForTest();
+    traceTestModule.start();
+    BPromise = require('./fixtures/bluebird3');
+  });
+
+  after(() => {
+    traceTestModule.setCLSForTest(traceTestModule.TestCLS);
+    traceTestModule.setPluginLoaderForTest(traceTestModule.TestPluginLoader);
+  });
+
+  const testCases = [
+    {
+      description: 'immediate resolve + child from then callback',
+      makePromise: () => new BPromise(res => res()),
+      thenFn: (p, cb) => p.then(cb)
+    } as TestCase,
+    {
+      description: 'deferred resolve + child from then callback',
+      makePromise: () => new BPromise(res => setTimeout(res, 0)),
+      thenFn: (p, cb) => p.then(cb)
+    } as TestCase,
+    {
+      description: 'bound, deferred resolve + child from then callback',
+      makePromise: () => new BPromise<void>(res => setTimeout(res, 0)).bind({}),
+      thenFn: (p, cb) => p.then(cb)
+    } as TestCase,
+    {
+      description: 'deferred resolve + child from spread callback',
+      makePromise: () => new BPromise(res => setTimeout(() => res([]), 0)),
+      thenFn: (p, cb) => p.spread(cb)
+    } as TestCase<never[]>,
+    {
+      description: 'deferred rejection + child from then callback',
+      makePromise: () => new BPromise((res, rej) => setTimeout(rej, 0)),
+      thenFn: (p, cb) => p.then(null, cb)
+    } as TestCase,
+    {
+      description: 'deferred rejection + child from catch callback',
+      makePromise: () => new BPromise((res, rej) => setTimeout(rej, 0)),
+      thenFn: (p, cb) => p.catch(cb)
+    } as TestCase,
+    {
+      description: 'deferred rejection + child from error callback',
+      makePromise: () => new BPromise(
+          (res, rej) =>
+              setTimeout(() => rej(new BPromise.OperationalError()), 0)),
+      thenFn: (p, cb) => p.error(cb)
+    } as TestCase,
+    {
+      description: 'deferred rejection + child from finally callback',
+      makePromise: () => new BPromise((res, rej) => setTimeout(rej, 0)),
+      thenFn: (p, cb) => p.catch(() => {}).finally(cb)
+    } as TestCase,
+    {
+      description: 'immediate resolve + child after await',
+      makePromise: () => new BPromise(res => res()),
+      thenFn: async (p, cb) => {
+        await p;
+        cb();
+      }
+    } as TestCase,
+    {
+      description: 'deferred resolve + child after await',
+      makePromise: () => new BPromise(res => setTimeout(res, 0)),
+      thenFn: async (p, cb) => {
+        await p;
+        cb();
+      }
+    } as TestCase
+  ];
+
+  // tslint:disable-next-line:no-any
+  testCases.forEach((testCase: TestCase<any>) => {
+    it(`enables context propagation in the same way as native promises for test case: ${
+           testCase.description}`,
+       async () => {
+         const actual = (await getTracesForPromiseImplementation(
+                             testCase.makePromise, testCase.thenFn))
+                            .map(trace => trace.spans.length)
+                            .join(', ');
+         // In each case, the second trace should have the child span.
+         // The format here is "[numSpansInFirstTrace],
+         // [numSpansInSecondTrace]".
+         assert.strictEqual(actual, '1, 2');
+       });
+  });
+});

test/plugins/test-trace-knex.ts

@@ -218,23 +218,18 @@ describeInterop<typeof knexTypes>('knex', (fixture) => {
                 return span.name === 'mysql-query';
               });
               let expectedCmds;
-              if (parsedVersion.minor === 10 || parsedVersion.minor >= 14) {
+              if (parsedVersion.minor === 10 || parsedVersion.minor >= 12) {
                 expectedCmds = [
                   /^BEGIN/, 'insert into `t` (`k`, `v`) values (?, ?)',
                   'select * from `t`', /^ROLLBACK/, 'select * from `t`'
                 ];
-              } else if (parsedVersion.minor === 11) {
+              } else /*if (parsedVersion.minor === 11)*/ {
                 expectedCmds = [
                   'SELECT 1', 'BEGIN;',
                   'insert into `t` (`k`, `v`) values (?, ?)',
                   'select * from `t`', 'ROLLBACK;', 'SELECT 1',
                   'select * from `t`'
                 ];
-              } else {
-                expectedCmds = [
-                  'insert into `t` (`k`, `v`) values (?, ?)',
-                  'select * from `t`', 'ROLLBACK;', 'select * from `t`'
-                ];
               }
               assert.strictEqual(expectedCmds.length, spans.length);
               for (let i = 0; i < spans.length; i++) {

test/utils.ts

@@ -52,9 +52,14 @@ export function wait(ms: number) {
   return new Promise(resolve => setTimeout(resolve, ms));
 }
 
+// Get the given span's duration in MS.
+export function getDuration(span: TraceSpan) {
+  return Date.parse(span.endTime) - Date.parse(span.startTime);
+}
+
 // Assert that the given span's duration is within the given range.
 export function assertSpanDuration(span: TraceSpan, bounds: [number, number?]) {
-  const spanDuration = Date.parse(span.endTime) - Date.parse(span.startTime);
+  const spanDuration = getDuration(span);
   const lowerBound = bounds[0];
   const upperBound = bounds[1] !== undefined ? bounds[1] : bounds[0];
   assert.ok(

tsconfig.json

@@ -9,6 +9,7 @@
   "include": [
     "src/*.ts",
     "src/cls/*.ts",
+    "src/plugins/plugin-bluebird.ts",
     "src/plugins/plugin-connect.ts",
     "src/plugins/plugin-express.ts",
     "src/plugins/plugin-grpc.ts",
@@ -19,6 +20,7 @@
     "src/plugins/plugin-koa.ts",
     "src/plugins/plugin-pg.ts",
     "src/plugins/plugin-restify.ts",
+    "test/plugins/test-cls-bluebird.ts",
     "test/plugins/test-trace-google-gax.ts",
     "test/plugins/test-trace-http.ts",
     "test/plugins/test-trace-http2.ts",