Demux handles arrays by key

src/functions/demux.ts

@@ -9,7 +9,10 @@ export interface DemuxOptions extends DuplexOptions {
 }
 
 export function demux(
-    construct: (destKey?: string, chunk?: any) => DemuxStreams,
+    construct: (
+        destKey?: string,
+        chunk?: any,
+    ) => DemuxStreams | Array<DemuxStreams>,
     demuxBy: string | ((chunk: any) => string),
     options?: DemuxOptions,
 ): Duplex {
@@ -18,21 +21,27 @@ export function demux(
 
 class Demux extends Duplex {
     private streamsByKey: {
-        [key: string]: DemuxStreams;
+        [key: string]: Array<DemuxStreams>;
     };
     private demuxer: (chunk: any) => string;
-    private construct: (destKey?: string, chunk?: any) => DemuxStreams;
+    private construct: (destKey?: string, chunk?: any) => Array<DemuxStreams>;
     private remultiplex: boolean;
     private transform: Transform;
     constructor(
-        construct: (destKey?: string) => DemuxStreams,
+        construct: (
+            destKey?: string,
+            chunk?: any,
+        ) => DemuxStreams | Array<DemuxStreams>,
         demuxBy: string | ((chunk: any) => string),
         options: DemuxOptions = {},
     ) {
         super(options);
         this.demuxer =
             typeof demuxBy === "string" ? chunk => chunk[demuxBy] : demuxBy;
-        this.construct = construct;
+        this.construct = (destKey: string, chunk?: any) => {
+            const pipeline = construct(destKey, chunk);
+            return Array.isArray(pipeline) ? pipeline : [pipeline];
+        };
         this.remultiplex =
             options.remultiplex === undefined ? true : options.remultiplex;
         this.streamsByKey = {};
@@ -53,43 +62,66 @@ class Demux extends Duplex {
     public async _write(chunk: any, encoding: any, cb: any) {
         const destKey = this.demuxer(chunk);
         if (this.streamsByKey[destKey] === undefined) {
-            const newPipeline = await this.construct(destKey, chunk);
+            const newPipelines = this.construct(destKey, chunk);
-            this.streamsByKey[destKey] = newPipeline;
+            this.streamsByKey[destKey] = newPipelines;
+
+            newPipelines.forEach(newPipeline => {
                 if (this.remultiplex && isReadable(newPipeline)) {
-                (newPipeline as NodeJS.ReadWriteStream).pipe(this.transform);
+                    (newPipeline as NodeJS.ReadWriteStream).pipe(
+                        this.transform,
+                    );
                 } else if (this.remultiplex) {
                     console.error(
                         `Pipeline construct for ${destKey} does not implement readable interface`,
                     );
                 }
-        }
-
-        if (!this.streamsByKey[destKey].write(chunk, encoding)) {
-            this.streamsByKey[destKey].once("drain", () => {
-                cb();
             });
-        } else {
-            cb();
         }
+        const pipelines = this.streamsByKey[destKey];
+        let pendingDrains: Array<Promise<any>> = [];
+
+        pipelines.forEach(pipeline => {
+            if (!pipeline.write(chunk, encoding)) {
+                pendingDrains.push(
+                    new Promise(resolve => {
+                        pipeline.once("drain", () => {
+                            resolve();
+                        });
+                    }),
+                );
+            }
+        });
+        await Promise.all(pendingDrains);
+        cb();
     }
 
     public _flush() {
-        const pipelines = Object.values(this.streamsByKey);
+        const pipelines: Array<DemuxStreams> = [].concat.apply(
-        let totalEnded = 0;
+            [],
+            Object.values(this.streamsByKey),
+        );
+        const flushPromises: Array<Promise<void>> = [];
         pipelines.forEach(pipeline => {
+            flushPromises.push(
+                new Promise(resolve => {
                     pipeline.once("end", () => {
-                totalEnded++;
+                        resolve();
-                if (pipelines.length === totalEnded) {
-                    this.push(null);
-                    this.emit("end");
-                }
                     });
+                }),
+            );
         });
         pipelines.forEach(pipeline => pipeline.end());
+        Promise.all(flushPromises).then(() => {
+            this.push(null);
+            this.emit("end");
+        });
     }
 
     public _destroy(error: any, cb: (error?: any) => void) {
-        const pipelines = Object.values(this.streamsByKey);
+        const pipelines: Array<DemuxStreams> = [].concat.apply(
+            [],
+            Object.values(this.streamsByKey),
+        );
         pipelines.forEach(p => (p as any).destroy());
         cb(error);
     }