index.js 5.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223
  1. import {on, once} from 'node:events';
  2. import {PassThrough as PassThroughStream} 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. if (streams.length === 0) {
  22. endStream(passThroughStream);
  23. }
  24. return passThroughStream;
  25. }
  26. const getHighWaterMark = (streams, objectMode) => {
  27. if (streams.length === 0) {
  28. // @todo Use `node:stream` `getDefaultHighWaterMark(objectMode)` in next major release
  29. return 16_384;
  30. }
  31. const highWaterMarks = streams
  32. .filter(({readableObjectMode}) => readableObjectMode === objectMode)
  33. .map(({readableHighWaterMark}) => readableHighWaterMark);
  34. return Math.max(...highWaterMarks);
  35. };
  36. class MergedStream extends PassThroughStream {
  37. #streams = new Set([]);
  38. #ended = new Set([]);
  39. #aborted = new Set([]);
  40. #onFinished;
  41. add(stream) {
  42. validateStream(stream);
  43. if (this.#streams.has(stream)) {
  44. return;
  45. }
  46. this.#streams.add(stream);
  47. this.#onFinished ??= onMergedStreamFinished(this, this.#streams);
  48. endWhenStreamsDone({
  49. passThroughStream: this,
  50. stream,
  51. streams: this.#streams,
  52. ended: this.#ended,
  53. aborted: this.#aborted,
  54. onFinished: this.#onFinished,
  55. });
  56. stream.pipe(this, {end: false});
  57. }
  58. remove(stream) {
  59. validateStream(stream);
  60. if (!this.#streams.has(stream)) {
  61. return false;
  62. }
  63. stream.unpipe(this);
  64. return true;
  65. }
  66. }
  67. const onMergedStreamFinished = async (passThroughStream, streams) => {
  68. updateMaxListeners(passThroughStream, PASSTHROUGH_LISTENERS_COUNT);
  69. const controller = new AbortController();
  70. try {
  71. await Promise.race([
  72. onMergedStreamEnd(passThroughStream, controller),
  73. onInputStreamsUnpipe(passThroughStream, streams, controller),
  74. ]);
  75. } finally {
  76. controller.abort();
  77. updateMaxListeners(passThroughStream, -PASSTHROUGH_LISTENERS_COUNT);
  78. }
  79. };
  80. const onMergedStreamEnd = async (passThroughStream, {signal}) => {
  81. await finished(passThroughStream, {signal, cleanup: true});
  82. };
  83. const onInputStreamsUnpipe = async (passThroughStream, streams, {signal}) => {
  84. for await (const [unpipedStream] of on(passThroughStream, 'unpipe', {signal})) {
  85. if (streams.has(unpipedStream)) {
  86. unpipedStream.emit(unpipeEvent);
  87. }
  88. }
  89. };
  90. const validateStream = stream => {
  91. if (typeof stream?.pipe !== 'function') {
  92. throw new TypeError(`Expected a readable stream, got: \`${typeof stream}\`.`);
  93. }
  94. };
  95. const endWhenStreamsDone = async ({passThroughStream, stream, streams, ended, aborted, onFinished}) => {
  96. updateMaxListeners(passThroughStream, PASSTHROUGH_LISTENERS_PER_STREAM);
  97. const controller = new AbortController();
  98. try {
  99. await Promise.race([
  100. afterMergedStreamFinished(onFinished, stream),
  101. onInputStreamEnd({passThroughStream, stream, streams, ended, aborted, controller}),
  102. onInputStreamUnpipe({stream, streams, ended, aborted, controller}),
  103. ]);
  104. } finally {
  105. controller.abort();
  106. updateMaxListeners(passThroughStream, -PASSTHROUGH_LISTENERS_PER_STREAM);
  107. }
  108. if (streams.size === ended.size + aborted.size) {
  109. if (ended.size === 0 && aborted.size > 0) {
  110. abortStream(passThroughStream);
  111. } else {
  112. endStream(passThroughStream);
  113. }
  114. }
  115. };
  116. // This is the error thrown by `finished()` on `stream.destroy()`
  117. const isAbortError = error => error?.code === 'ERR_STREAM_PREMATURE_CLOSE';
  118. const afterMergedStreamFinished = async (onFinished, stream) => {
  119. try {
  120. await onFinished;
  121. abortStream(stream);
  122. } catch (error) {
  123. if (isAbortError(error)) {
  124. abortStream(stream);
  125. } else {
  126. errorStream(stream, error);
  127. }
  128. }
  129. };
  130. const onInputStreamEnd = async ({passThroughStream, stream, streams, ended, aborted, controller: {signal}}) => {
  131. try {
  132. await finished(stream, {signal, cleanup: true, readable: true, writable: false});
  133. if (streams.has(stream)) {
  134. ended.add(stream);
  135. }
  136. } catch (error) {
  137. if (signal.aborted || !streams.has(stream)) {
  138. return;
  139. }
  140. if (isAbortError(error)) {
  141. aborted.add(stream);
  142. } else {
  143. errorStream(passThroughStream, error);
  144. }
  145. }
  146. };
  147. const onInputStreamUnpipe = async ({stream, streams, ended, aborted, controller: {signal}}) => {
  148. await once(stream, unpipeEvent, {signal});
  149. streams.delete(stream);
  150. ended.delete(stream);
  151. aborted.delete(stream);
  152. };
  153. const unpipeEvent = Symbol('unpipe');
  154. const endStream = stream => {
  155. if (stream.writable) {
  156. stream.end();
  157. }
  158. };
  159. const abortStream = stream => {
  160. if (stream.readable || stream.writable) {
  161. stream.destroy();
  162. }
  163. };
  164. // `stream.destroy(error)` crashes the process with `uncaughtException` if no `error` event listener exists on `stream`.
  165. // We take care of error handling on user behalf, so we do not want this to happen.
  166. const errorStream = (stream, error) => {
  167. if (!stream.destroyed) {
  168. stream.once('error', noop);
  169. stream.destroy(error);
  170. }
  171. };
  172. const noop = () => {};
  173. const updateMaxListeners = (passThroughStream, increment) => {
  174. const maxListeners = passThroughStream.getMaxListeners();
  175. if (maxListeners !== 0 && maxListeners !== Number.POSITIVE_INFINITY) {
  176. passThroughStream.setMaxListeners(maxListeners + increment);
  177. }
  178. };
  179. // Number of times `passThroughStream.on()` is called regardless of streams:
  180. // - once due to `finished(passThroughStream)`
  181. // - once due to `on(passThroughStream)`
  182. const PASSTHROUGH_LISTENERS_COUNT = 2;
  183. // Number of times `passThroughStream.on()` is called per stream:
  184. // - once due to `stream.pipe(passThroughStream)`
  185. const PASSTHROUGH_LISTENERS_PER_STREAM = 1;