From 0e63443215abc3ab0b08295275f3a63b94413f8a Mon Sep 17 00:00:00 2001 From: Kelvin Jin Date: Fri, 21 Sep 2018 17:02:49 -0700 Subject: [PATCH] feat: support context propagation in bluebird --- src/config.ts | 1 + src/plugins/plugin-bluebird.ts | 46 ++++++++ src/plugins/types/index.d.ts | 2 + test/fixtures/plugin-fixtures.json | 5 + test/plugins/test-cls-bluebird.ts | 169 +++++++++++++++++++++++++++++ test/utils.ts | 7 +- tsconfig.json | 2 + 7 files changed, 231 insertions(+), 1 deletion(-) create mode 100644 src/plugins/plugin-bluebird.ts create mode 100644 test/plugins/test-cls-bluebird.ts diff --git a/src/config.ts b/src/config.ts index 58ae1abad..5f0ea52b7 100644 --- a/src/config.ts +++ b/src/config.ts @@ -255,6 +255,7 @@ 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'), diff --git a/src/plugins/plugin-bluebird.ts b/src/plugins/plugin-bluebird.ts new file mode 100644 index 000000000..d42c1af59 --- /dev/null +++ b/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(this: bluebird_3, ...args: any[]) { + return thenFn.apply(this, [ + wrapIfFunction(args[0]), wrapIfFunction(args[1]), ...args.slice(2) + ]); + }; + }); + // tslint:enable:no-any + } +} as PluginTypes.Monkeypatch]; + +export = plugin; diff --git a/src/plugins/types/index.d.ts b/src/plugins/types/index.d.ts index 36d4e3add..870ab2c43 100644 --- a/src/plugins/types/index.d.ts +++ b/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, diff --git a/test/fixtures/plugin-fixtures.json b/test/fixtures/plugin-fixtures.json index 87f077da4..6e8193c14 100644 --- a/test/fixtures/plugin-fixtures.json +++ b/test/fixtures/plugin-fixtures.json @@ -1,4 +1,9 @@ { + "bluebird3": { + "dependencies": { + "bluebird": "^3.5.2" + } + }, "connect3": { "dependencies": { "connect": "^3.5.0" diff --git a/test/plugins/test-cls-bluebird.ts b/test/plugins/test-cls-bluebird.ts new file mode 100644 index 000000000..c7cadd138 --- /dev/null +++ b/test/plugins/test-cls-bluebird.ts @@ -0,0 +1,169 @@ +/** + * 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 assert from 'assert'; + +import {bluebird_3 as BluebirdPromise, bluebird_3} from '../../src/plugins/types'; +import {Trace} from '../../src/trace'; +import * as traceTestModule from '../trace'; + +/** + * Describes a test case. + */ +interface TestCase { + /** + * Description of a test case; included in string argument to it(). + */ + description: string; + /** + * Creates and returns a new Promise. + */ + makePromise: () => BluebirdPromise; + /** + * Given a Promise and a callback, calls the callback some time after the + * Promise has been resolved or rejected. + */ + thenFn: (promise: BluebirdPromise, 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 = + (makePromise: () => BluebirdPromise, + thenFn: (promise: BluebirdPromise, cb: () => void) => void): + Promise<[Trace, Trace]> => new Promise((resolve, reject) => { + const tracer = traceTestModule.get(); + let p: BluebirdPromise; + 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: Array> = [ + { + description: 'immediate resolve + child from then callback', + makePromise: () => new BPromise(res => res()), + thenFn: (p, cb) => p.then(cb) + } as TestCase, + { + description: 'bound, immediate resolve + child from then callback', + makePromise: () => new BPromise(res => res()).bind({}), + thenFn: (p, cb) => p.then(cb) + } as TestCase, + { + description: 'immediate resolve + child from spread callback', + makePromise: () => new BPromise(res => res([])), + thenFn: (p, cb) => p.spread(cb) + } as TestCase, + { + description: 'immediate rejection + child from then callback', + makePromise: () => new BPromise((res, rej) => rej()), + thenFn: (p, cb) => p.then(null, cb) + } as TestCase, + { + description: 'immediate rejection + child from catch callback', + makePromise: () => new BPromise((res, rej) => rej()), + thenFn: (p, cb) => p.catch(cb) + } as TestCase, + { + description: 'immediate rejection + child from error callback', + makePromise: () => new BPromise((res, rej) => rej(new BPromise.OperationalError())), + thenFn: (p, cb) => p.error(cb) + } as TestCase, + { + description: 'immediate rejection + child from finally callback', + makePromise: () => new BPromise((res, rej) => rej()), + thenFn: (p, cb) => p.catch(() => {}).finally(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: '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 + ]; + + testCases.forEach(testCase => { + 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'); + }); + }); +}); diff --git a/test/utils.ts b/test/utils.ts index 762974ba3..527ce3f8c 100644 --- a/test/utils.ts +++ b/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( diff --git a/tsconfig.json b/tsconfig.json index e297aad69..5a4f7131b 100644 --- a/tsconfig.json +++ b/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",