Demux handles arrays by key

This commit is contained in:
Jerry Kurian 2020-05-08 15:58:31 -04:00
parent 71b03678ba
commit 58a95a91d0

View File

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