shimCoroutine.js 2.7 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485
  1. 'use strict';
  2. /*
  3. * cls-bluebird
  4. * Function to shim `Promise.coroutine`
  5. *
  6. * Works by binding the `.next()` and `.throw()` methods of generator to CLS context
  7. * at time when coroutine is executed.
  8. *
  9. * In bluebird v3.x, running the coroutine internally calls `.lastly()` if cancellation is enabled.
  10. * To prevent unnecessary binding of the `.lastly()` callback to CLS context, this patch
  11. * temporarily disables the patch on `Promise.prototype.lastly`.
  12. * NB This patch could break if bluebird internals change, but this is covered by the tests.
  13. */
  14. // Modules
  15. var shimmer = require('shimmer');
  16. // Exports
  17. /**
  18. * Patch `Promise.coroutine` or `Promise.spawn` to maintain current CLS context after all `yield` statements.
  19. *
  20. * @param {string} methodName - method name (either 'coroutine' or 'spawn')
  21. * @param {Function} Promise - Bluebird Promise constructor to patch
  22. * @param {Object} ns - CLS namespace to bind callbacks to
  23. * @returns {undefined}
  24. */
  25. module.exports = function(methodName, Promise, ns, v3) {
  26. var lastlyPatched = Promise.prototype.lastly,
  27. lastlyOriginal = Promise.prototype.lastly.__original;
  28. // Patch method
  29. shimmer.wrap(Promise, methodName, function(original) {
  30. return function(generatorFunction, options) {
  31. // NB If `generatorFunction` is not a function, do not alter it.
  32. // Pass value directly to bluebird which will throw an error.
  33. if (typeof generatorFunction === 'function') {
  34. // Create proxy generator function
  35. var generatorFunctionOriginal = generatorFunction;
  36. generatorFunction = function() {
  37. // Create generator from generator function
  38. var generator = generatorFunctionOriginal.apply(this, arguments);
  39. // Bind `.next()`, '.throw()' and `.return()` to current CLS context.
  40. // NB CLS context is from when coroutine is called, not when created.
  41. ['next', 'throw', 'return'].forEach(function(name) {
  42. if (typeof generator[name] === 'function') generator[name] = ns.bind(generator[name]);
  43. });
  44. return generator;
  45. };
  46. }
  47. // Temporarily remove patch from `Promise.prototype.lastly` in bluebird v3
  48. // to avoid unnecessary binding to CLS context.
  49. var self = this;
  50. if (methodName === 'spawn' && v3) {
  51. return tempPatchLastly(function() {
  52. return original.call(self, generatorFunction, options);
  53. });
  54. }
  55. var fn = original.call(this, generatorFunction, options);
  56. if (methodName === 'coroutine' && v3) {
  57. return function() {
  58. var self = this, args = arguments;
  59. return tempPatchLastly(function() {
  60. return fn.apply(self, args);
  61. });
  62. };
  63. }
  64. return fn;
  65. };
  66. });
  67. function tempPatchLastly(fn) {
  68. Promise.prototype.lastly = lastlyOriginal;
  69. var res = fn();
  70. Promise.prototype.lastly = lastlyPatched;
  71. return res;
  72. }
  73. };