index.js 6.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265
  1. import {on, once} from 'node:events';
  2. import {PassThrough as PassThroughStream, getDefaultHighWaterMark} from 'node:stream';
  3. import {finished} from 'node:stream/promises';
  4. export default function mergeStreams(streams) {
  5. if (!Array.isArray(streams)) {
  6. throw new TypeError(`Expected an array, got \`${typeof streams}\`.`);
  7. }
  8. for (const stream of streams) {
  9. validateStream(stream);
  10. }
  11. const objectMode = streams.some(({readableObjectMode}) => readableObjectMode);
  12. const highWaterMark = getHighWaterMark(streams, objectMode);
  13. const passThroughStream = new MergedStream({
  14. objectMode,
  15. writableHighWaterMark: highWaterMark,
  16. readableHighWaterMark: highWaterMark,
  17. });
  18. for (const stream of streams) {
  19. passThroughStream.add(stream);
  20. }
  21. return passThroughStream;
  22. }
  23. const getHighWaterMark = (streams, objectMode) => {
  24. if (streams.length === 0) {
  25. return getDefaultHighWaterMark(objectMode);
  26. }
  27. const highWaterMarks = streams
  28. .filter(({readableObjectMode}) => readableObjectMode === objectMode)
  29. .map(({readableHighWaterMark}) => readableHighWaterMark);
  30. return Math.max(...highWaterMarks);
  31. };
  32. class MergedStream extends PassThroughStream {
  33. #streams = new Set([]);
  34. #ended = new Set([]);
  35. #aborted = new Set([]);
  36. #onFinished;
  37. #unpipeEvent = Symbol('unpipe');
  38. #streamPromises = new WeakMap();
  39. add(stream) {
  40. validateStream(stream);
  41. if (this.#streams.has(stream)) {
  42. return;
  43. }
  44. this.#streams.add(stream);
  45. this.#onFinished ??= onMergedStreamFinished(this, this.#streams, this.#unpipeEvent);
  46. const streamPromise = endWhenStreamsDone({
  47. passThroughStream: this,
  48. stream,
  49. streams: this.#streams,
  50. ended: this.#ended,
  51. aborted: this.#aborted,
  52. onFinished: this.#onFinished,
  53. unpipeEvent: this.#unpipeEvent,
  54. });
  55. this.#streamPromises.set(stream, streamPromise);
  56. stream.pipe(this, {end: false});
  57. }
  58. async remove(stream) {
  59. validateStream(stream);
  60. if (!this.#streams.has(stream)) {
  61. return false;
  62. }
  63. const streamPromise = this.#streamPromises.get(stream);
  64. if (streamPromise === undefined) {
  65. return false;
  66. }
  67. this.#streamPromises.delete(stream);
  68. stream.unpipe(this);
  69. await streamPromise;
  70. return true;
  71. }
  72. }
  73. const onMergedStreamFinished = async (passThroughStream, streams, unpipeEvent) => {
  74. updateMaxListeners(passThroughStream, PASSTHROUGH_LISTENERS_COUNT);
  75. const controller = new AbortController();
  76. try {
  77. await Promise.race([
  78. onMergedStreamEnd(passThroughStream, controller),
  79. onInputStreamsUnpipe(passThroughStream, streams, unpipeEvent, controller),
  80. ]);
  81. } finally {
  82. controller.abort();
  83. updateMaxListeners(passThroughStream, -PASSTHROUGH_LISTENERS_COUNT);
  84. }
  85. };
  86. const onMergedStreamEnd = async (passThroughStream, {signal}) => {
  87. try {
  88. await finished(passThroughStream, {signal, cleanup: true});
  89. } catch (error) {
  90. errorOrAbortStream(passThroughStream, error);
  91. throw error;
  92. }
  93. };
  94. const onInputStreamsUnpipe = async (passThroughStream, streams, unpipeEvent, {signal}) => {
  95. for await (const [unpipedStream] of on(passThroughStream, 'unpipe', {signal})) {
  96. if (streams.has(unpipedStream)) {
  97. unpipedStream.emit(unpipeEvent);
  98. }
  99. }
  100. };
  101. const validateStream = stream => {
  102. if (typeof stream?.pipe !== 'function') {
  103. throw new TypeError(`Expected a readable stream, got: \`${typeof stream}\`.`);
  104. }
  105. };
  106. const endWhenStreamsDone = async ({passThroughStream, stream, streams, ended, aborted, onFinished, unpipeEvent}) => {
  107. updateMaxListeners(passThroughStream, PASSTHROUGH_LISTENERS_PER_STREAM);
  108. const controller = new AbortController();
  109. try {
  110. await Promise.race([
  111. afterMergedStreamFinished(onFinished, stream, controller),
  112. onInputStreamEnd({
  113. passThroughStream,
  114. stream,
  115. streams,
  116. ended,
  117. aborted,
  118. controller,
  119. }),
  120. onInputStreamUnpipe({
  121. stream,
  122. streams,
  123. ended,
  124. aborted,
  125. unpipeEvent,
  126. controller,
  127. }),
  128. ]);
  129. } finally {
  130. controller.abort();
  131. updateMaxListeners(passThroughStream, -PASSTHROUGH_LISTENERS_PER_STREAM);
  132. }
  133. if (streams.size > 0 && streams.size === ended.size + aborted.size) {
  134. if (ended.size === 0 && aborted.size > 0) {
  135. abortStream(passThroughStream);
  136. } else {
  137. endStream(passThroughStream);
  138. }
  139. }
  140. };
  141. const afterMergedStreamFinished = async (onFinished, stream, {signal}) => {
  142. try {
  143. await onFinished;
  144. if (!signal.aborted) {
  145. abortStream(stream);
  146. }
  147. } catch (error) {
  148. if (!signal.aborted) {
  149. errorOrAbortStream(stream, error);
  150. }
  151. }
  152. };
  153. const onInputStreamEnd = async ({passThroughStream, stream, streams, ended, aborted, controller: {signal}}) => {
  154. try {
  155. await finished(stream, {
  156. signal,
  157. cleanup: true,
  158. readable: true,
  159. writable: false,
  160. });
  161. if (streams.has(stream)) {
  162. ended.add(stream);
  163. }
  164. } catch (error) {
  165. if (signal.aborted || !streams.has(stream)) {
  166. return;
  167. }
  168. if (isAbortError(error)) {
  169. aborted.add(stream);
  170. } else {
  171. errorStream(passThroughStream, error);
  172. }
  173. }
  174. };
  175. const onInputStreamUnpipe = async ({stream, streams, ended, aborted, unpipeEvent, controller: {signal}}) => {
  176. await once(stream, unpipeEvent, {signal});
  177. if (!stream.readable) {
  178. return once(signal, 'abort', {signal});
  179. }
  180. streams.delete(stream);
  181. ended.delete(stream);
  182. aborted.delete(stream);
  183. };
  184. const endStream = stream => {
  185. if (stream.writable) {
  186. stream.end();
  187. }
  188. };
  189. const errorOrAbortStream = (stream, error) => {
  190. if (isAbortError(error)) {
  191. abortStream(stream);
  192. } else {
  193. errorStream(stream, error);
  194. }
  195. };
  196. // This is the error thrown by `finished()` on `stream.destroy()`
  197. const isAbortError = error => error?.code === 'ERR_STREAM_PREMATURE_CLOSE';
  198. const abortStream = stream => {
  199. if (stream.readable || stream.writable) {
  200. stream.destroy();
  201. }
  202. };
  203. // `stream.destroy(error)` crashes the process with `uncaughtException` if no `error` event listener exists on `stream`.
  204. // We take care of error handling on user behalf, so we do not want this to happen.
  205. const errorStream = (stream, error) => {
  206. if (!stream.destroyed) {
  207. stream.once('error', noop);
  208. stream.destroy(error);
  209. }
  210. };
  211. const noop = () => {};
  212. const updateMaxListeners = (passThroughStream, increment) => {
  213. const maxListeners = passThroughStream.getMaxListeners();
  214. if (maxListeners !== 0 && maxListeners !== Number.POSITIVE_INFINITY) {
  215. passThroughStream.setMaxListeners(maxListeners + increment);
  216. }
  217. };
  218. // Number of times `passThroughStream.on()` is called regardless of streams:
  219. // - once due to `finished(passThroughStream)`
  220. // - once due to `on(passThroughStream)`
  221. const PASSTHROUGH_LISTENERS_COUNT = 2;
  222. // Number of times `passThroughStream.on()` is called per stream:
  223. // - once due to `stream.pipe(passThroughStream)`
  224. const PASSTHROUGH_LISTENERS_PER_STREAM = 1;