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(
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);
this.streamsByKey[destKey] = newPipeline;
const newPipelines = this.construct(destKey, chunk);
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);
let totalEnded = 0;
const pipelines: Array<DemuxStreams> = [].concat.apply(
[],
Object.values(this.streamsByKey),
);
const flushPromises: Array<Promise<void>> = [];
pipelines.forEach(pipeline => {
flushPromises.push(
new Promise(resolve => {
pipeline.once("end", () => {
totalEnded++;
if (pipelines.length === totalEnded) {
this.push(null);
this.emit("end");
}
resolve();
});
}),
);
});
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);
}