Update demux

This commit is contained in:
Jerry Kurian 2019-09-12 09:08:49 -04:00
parent 65c36a8f22
commit 586f618e95
2 changed files with 128 additions and 73 deletions

View File

@ -24,7 +24,7 @@ const eventsTarget = {
export function demux( export function demux(
construct: () => NodeJS.WritableStream | NodeJS.ReadWriteStream, construct: () => NodeJS.WritableStream | NodeJS.ReadWriteStream,
demuxBy: { key?: string; keyBy?: (chunk: any) => string }, demuxBy: string | ((chunk: any) => string),
options?: WritableOptions, options?: WritableOptions,
): Writable { ): Writable {
return new Demux(construct, demuxBy, options); return new Demux(construct, demuxBy, options);
@ -42,19 +42,17 @@ class Demux extends Writable {
construct: ( construct: (
destKey?: string, destKey?: string,
) => NodeJS.WritableStream | NodeJS.ReadWriteStream, ) => NodeJS.WritableStream | NodeJS.ReadWriteStream,
demuxBy: { key?: string; keyBy?: (chunk: any) => string }, demuxBy: string | ((chunk: any) => string),
options?: WritableOptions, options?: WritableOptions,
) { ) {
super(options); super(options);
if (demuxBy.keyBy === undefined && demuxBy.key === undefined) { this.demuxer =
throw new Error("keyBy or key must be provided in second argument"); typeof demuxBy === "string" ? chunk => chunk[demuxBy] : demuxBy;
}
this.demuxer = demuxBy.keyBy || ((chunk: any) => chunk[demuxBy.key!]);
this.construct = construct; this.construct = construct;
this.streamsByKey = {}; this.streamsByKey = {};
} }
public async _write(chunk: any, encoding: any, cb: any) { public _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) {
this.streamsByKey[destKey] = this.construct(destKey); this.streamsByKey[destKey] = this.construct(destKey);

View File

@ -30,7 +30,7 @@ test.cb("demux() constructor should be called once per key", t => {
return dest; return dest;
}); });
const demuxed = demux(construct, { key: "key" }, { objectMode: true }); const demuxed = demux(construct, "key", { objectMode: true });
demuxed.on("finish", () => { demuxed.on("finish", () => {
expect(construct.withArgs("a").callCount).to.equal(1); expect(construct.withArgs("a").callCount).to.equal(1);
@ -65,7 +65,7 @@ test.cb("demux() should send input through correct pipeline", t => {
return dest; return dest;
}; };
const demuxed = demux(construct, { key: "key" }, { objectMode: true }); const demuxed = demux(construct, "key", { objectMode: true });
demuxed.on("finish", () => { demuxed.on("finish", () => {
pipelineSpies["a"].getCalls().forEach(call => { pipelineSpies["a"].getCalls().forEach(call => {
@ -107,11 +107,7 @@ test.cb("demux() constructor should be called once per key using keyBy", t => {
return dest; return dest;
}); });
const demuxed = demux( const demuxed = demux(construct, item => item.key, { objectMode: true });
construct,
{ keyBy: item => item.key },
{ objectMode: true },
);
demuxed.on("finish", () => { demuxed.on("finish", () => {
expect(construct.withArgs("a").callCount).to.equal(1); expect(construct.withArgs("a").callCount).to.equal(1);
@ -146,11 +142,7 @@ test.cb("demux() should send input through correct pipeline using keyBy", t => {
return dest; return dest;
}; };
const demuxed = demux( const demuxed = demux(construct, item => item.key, { objectMode: true });
construct,
{ keyBy: item => item.key },
{ objectMode: true },
);
demuxed.on("finish", () => { demuxed.on("finish", () => {
pipelineSpies["a"].getCalls().forEach(call => { pipelineSpies["a"].getCalls().forEach(call => {
@ -211,14 +203,10 @@ test("demux() write should return false after if it has >= highWaterMark items b
return first; return first;
}; };
const _demux = demux( const _demux = demux(construct, "key", {
construct,
{ key: "key" },
{
objectMode: true, objectMode: true,
highWaterMark, highWaterMark,
}, });
);
_demux.on("error", err => { _demux.on("error", err => {
reject(); reject();
@ -278,14 +266,10 @@ test("demux() should emit one drain event after slowProcessorSpeed * highWaterMa
}); });
return first; return first;
}; };
const _demux = demux( const _demux = demux(construct, "key", {
construct,
{ key: "key" },
{
objectMode: true, objectMode: true,
highWaterMark, highWaterMark,
}, });
);
_demux.on("error", err => { _demux.on("error", err => {
reject(); reject();
}); });
@ -345,14 +329,10 @@ test("demux() should emit one drain event when writing 6 items with highWaterMar
}); });
return first; return first;
}; };
const _demux = demux( const _demux = demux(construct, "key", {
construct,
{ key: "key" },
{
objectMode: true, objectMode: true,
highWaterMark: 5, highWaterMark: 5,
}, });
);
_demux.on("error", err => { _demux.on("error", err => {
reject(); reject();
@ -375,9 +355,9 @@ test("demux() should emit one drain event when writing 6 items with highWaterMar
}); });
test.cb( test.cb(
"demux() should emit drain event when second stream is bottleneck", "demux() should emit drain event when third stream is bottleneck",
t => { t => {
t.plan(6); t.plan(8);
const slowProcessorSpeed = 100; const slowProcessorSpeed = 100;
const highWaterMark = 5; const highWaterMark = 5;
interface Chunk { interface Chunk {
@ -417,20 +397,15 @@ test.cb(
first.pipe(second).pipe(sink); first.pipe(second).pipe(sink);
return first; return first;
}; };
const _demux = demux( const _demux = demux(construct, () => "a", {
construct,
{ key: "key" },
{
objectMode: true, objectMode: true,
highWaterMark, highWaterMark,
}, });
);
_demux.on("error", err => { _demux.on("error", err => {
t.end(err); t.end(err);
}); });
// This event should be received after at least 3 * slowProcessorSpeed (two are read immediately by first) // This event should be received after at least 3 * slowProcessorSpeed (two are read immediately by first)
// @TODO Verify this is correct behaviour
_demux.on("drain", () => { _demux.on("drain", () => {
expect(_demux._writableState.length).to.be.equal(0); expect(_demux._writableState.length).to.be.equal(0);
expect(performance.now() - start).to.be.greaterThan( expect(performance.now() - start).to.be.greaterThan(
@ -441,10 +416,100 @@ test.cb(
const input = [ const input = [
{ key: "a", mapped: [] }, { key: "a", mapped: [] },
{ key: "b", mapped: [] },
{ key: "c", mapped: [] },
{ key: "d", mapped: [] },
{ key: "e", mapped: [] },
{ key: "f", mapped: [] },
{ key: "g", mapped: [] },
];
let pendingReads = input.length;
const start = performance.now();
input.forEach(item => {
_demux.write(item);
});
},
);
test.cb(
"demux() should emit drain event when second stream is bottleneck",
t => {
t.plan(8);
const slowProcessorSpeed = 100;
const highWaterMark = 5;
interface Chunk {
key: string;
mapped: number[];
}
const sink = new Writable({
objectMode: true,
write(chunk, encoding, cb) {
expect(chunk.mapped).to.deep.equal([1, 2, 3]);
t.pass();
pendingReads--;
if (pendingReads === 0) {
t.end();
}
cb();
},
});
const construct = (destKey: string) => {
const first = map(
(chunk: Chunk) => {
chunk.mapped.push(1);
return chunk;
},
{ objectMode: true, highWaterMark: 1 },
);
const second = map(
(chunk: Chunk) => {
chunk.mapped.push(2);
return chunk;
},
{ objectMode: true, highWaterMark: 1 },
);
const third = map(
async (chunk: Chunk) => {
await sleep(slowProcessorSpeed);
chunk.mapped.push(3);
return chunk;
},
{ objectMode: true, highWaterMark: 1 },
);
first
.pipe(second)
.pipe(third)
.pipe(sink);
return first;
};
const _demux = demux(construct, () => "a", {
objectMode: true,
highWaterMark,
});
_demux.on("error", err => {
t.end(err);
});
// This event should be received after at least 3 * slowProcessorSpeed (two are read immediately by first)
_demux.on("drain", () => {
expect(_demux._writableState.length).to.be.equal(0);
expect(performance.now() - start).to.be.greaterThan(
slowProcessorSpeed * (input.length - 4),
);
t.pass();
});
const input = [
{ key: "a", mapped: [] }, { key: "a", mapped: [] },
{ key: "a", mapped: [] }, { key: "b", mapped: [] },
{ key: "a", mapped: [] }, { key: "c", mapped: [] },
{ key: "a", mapped: [] }, { key: "d", mapped: [] },
{ key: "e", mapped: [] },
{ key: "f", mapped: [] },
{ key: "g", mapped: [] },
]; ];
let pendingReads = input.length; let pendingReads = input.length;
@ -486,14 +551,10 @@ test("demux() should be blocked by slowest pipeline", t => {
}); });
return first; return first;
}; };
const _demux = demux( const _demux = demux(construct, "key", {
construct,
{ key: "key" },
{
objectMode: true, objectMode: true,
highWaterMark: 1, highWaterMark: 1,
}, });
);
_demux.on("error", err => { _demux.on("error", err => {
reject(err); reject(err);
}); });
@ -567,14 +628,10 @@ test("demux() should emit drain event when second stream in pipeline is bottlene
return first; return first;
}; };
const _demux = demux( const _demux = demux(construct, "key", {
construct,
{ key: "key" },
{
objectMode: true, objectMode: true,
highWaterMark, highWaterMark,
}, });
);
_demux.on("error", err => { _demux.on("error", err => {
reject(); reject();
}); });