Wait for drain when write returns false in demux

This commit is contained in:
Jerry Kurian 2019-09-10 18:13:13 -04:00
parent ee3d9b9ded
commit 9d280b1662
2 changed files with 62 additions and 143 deletions

View File

@ -62,7 +62,7 @@ class Demux extends Writable {
} }
// Throttles when one stream is not writable // Throttles when one stream is not writable
public _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) {
this.streamsByKey[destKey] = { this.streamsByKey[destKey] = {
@ -70,21 +70,12 @@ class Demux extends Writable {
writable: true, writable: true,
}; };
} }
let res = false; if (!this.streamsByKey[destKey].stream.write(chunk, encoding, cb)) {
if (this.streamsByKey[destKey].writable && this.isWritable) { await new Promise((resolve, reject) => {
res = this.streamsByKey[destKey].stream.write(chunk, encoding, cb);
}
if (!res && this.isWritable) {
this.isWritable = false;
this.streamsByKey[destKey].writable = false;
this.nonWritableStreams.push(destKey);
this.streamsByKey[destKey].stream.once("drain", () => { this.streamsByKey[destKey].stream.once("drain", () => {
this.nonWritableStreams.filter(key => key !== destKey); resolve();
this.isWritable = this.nonWritableStreams.length === 0;
this.streamsByKey[destKey].stream.write(chunk, encoding, cb);
if (this.isWritable) {
this.emit("drain"); this.emit("drain");
} });
}); });
} }
} }

View File

@ -172,6 +172,7 @@ test.cb("demux() should send input through correct pipeline using keyBy", t => {
demuxed.end(); demuxed.end();
}); });
// Probably needs to be removed
test.cb("should emit errors", t => { test.cb("should emit errors", t => {
t.plan(2); t.plan(2);
let index = 0; let index = 0;
@ -209,7 +210,7 @@ test.cb("should emit errors", t => {
_chunk.visited.push(index); _chunk.visited.push(index);
index++; index++;
return _chunk; return _chunk;
}).on("error", () => {}); }).on("error", () => {}); // Otherwise ava complains
dest.pipe(sink); dest.pipe(sink);
return dest; return dest;
@ -226,18 +227,26 @@ test.cb("should emit errors", t => {
t.end(); t.end();
}); });
input.forEach(event => demuxed.write(event)); input.forEach(event => demuxed.write(event));
demuxed.end();
}); });
test("demux() should emit drain event ~rate * highWaterMark ms for every write that causes backpressure", t => { test("demux() when write returns false, drain event should be emitted after at least slowProcessorSpeed * highWaterMark", t => {
return new Promise(async (resolve, reject) => {
t.plan(7); t.plan(7);
interface Chunk { interface Chunk {
key: string; key: string;
mapped: number[]; mapped: number[];
} }
const input: Chunk[] = [
{ key: "a", mapped: [] },
{ key: "a", mapped: [] },
{ key: "a", mapped: [] },
{ key: "a", mapped: [] },
{ key: "a", mapped: [] },
{ key: "a", mapped: [] },
];
let pendingReads = input.length;
const highWaterMark = 5; const highWaterMark = 5;
const _rate = 25; const slowProcessorSpeed = 25;
return new Promise(async (resolve, reject) => {
const sink = new Writable({ const sink = new Writable({
objectMode: true, objectMode: true,
write(chunk, encoding, cb) { write(chunk, encoding, cb) {
@ -251,7 +260,7 @@ test("demux() should emit drain event ~rate * highWaterMark ms for every write t
}); });
const construct = (destKey: string) => { const construct = (destKey: string) => {
const first = map(async (chunk: Chunk) => { const first = map(async (chunk: Chunk) => {
await sleep(_rate); await sleep(slowProcessorSpeed);
chunk.mapped.push(1); chunk.mapped.push(1);
return chunk; return chunk;
}); });
@ -278,40 +287,45 @@ test("demux() should emit drain event ~rate * highWaterMark ms for every write t
_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(_rate); expect(performance.now() - start).to.be.greaterThan(
slowProcessorSpeed * highWaterMark,
);
t.pass(); t.pass();
}); });
const input = [ let start = null;
{ key: "a", mapped: [] },
{ key: "a", mapped: [] },
{ key: "a", mapped: [] },
{ key: "a", mapped: [] },
{ key: "a", mapped: [] },
{ key: "a", mapped: [] },
];
let pendingReads = input.length;
let start = performance.now();
for (const item of input) { for (const item of input) {
const res = _demux.write(item); const res = _demux.write(item);
expect(_demux._writableState.length).to.be.at.most(highWaterMark); expect(_demux._writableState.length).to.be.at.most(highWaterMark);
if (!res) { if (!res) {
start = performance.now(); start = performance.now();
await sleep(100); await new Promise((resolv, rej) => {
_demux.once("drain", () => {
resolv();
});
});
} }
} }
}); });
}); });
test("demux() should emit one drain event when writing 6 items with highWaterMark of 5", t => { test("demux() should emit one drain event when writing 6 items with highWaterMark of 5", t => {
t.plan(7);
const highWaterMark = 5;
return new Promise(async (resolve, reject) => { return new Promise(async (resolve, reject) => {
t.plan(7);
interface Chunk { interface Chunk {
key: string; key: string;
mapped: number[]; mapped: number[];
} }
const highWaterMark = 5;
const input = [
{ key: "a", mapped: [] },
{ key: "a", mapped: [] },
{ key: "a", mapped: [] },
{ key: "a", mapped: [] },
{ key: "a", mapped: [] },
{ key: "a", mapped: [] },
];
let pendingReads = input.length;
const sink = new Writable({ const sink = new Writable({
objectMode: true, objectMode: true,
write(chunk, encoding, cb) { write(chunk, encoding, cb) {
@ -354,110 +368,26 @@ test("demux() should emit one drain event when writing 6 items with highWaterMar
t.pass(); t.pass();
}); });
const input = [
{ key: "a", mapped: [] },
{ key: "a", mapped: [] },
{ key: "a", mapped: [] },
{ key: "a", mapped: [] },
{ key: "a", mapped: [] },
{ key: "a", mapped: [] },
];
let pendingReads = input.length;
for (const item of input) { for (const item of input) {
const res = _demux.write(item); const res = _demux.write(item);
expect(_demux._writableState.length).to.be.at.most(highWaterMark); expect(_demux._writableState.length).to.be.at.most(highWaterMark);
if (!res) { if (!res) {
await sleep(10); await new Promise((_resolve, _reject) => {
_demux.once("drain", () => {
_resolve();
});
});
} }
} }
}); });
}); });
test.cb(
"demux() should emit drain event after 500 ms when writing 5 items that take 100ms to process with a highWaterMark of 5 ",
t => {
t.plan(6);
const _rate = 100;
const highWaterMark = 5;
interface Chunk {
key: string;
mapped: number[];
}
const sink = new Writable({
objectMode: true,
write(chunk, encoding, cb) {
t.pass();
cb();
if (pendingReads === 0) {
t.end();
}
},
});
const construct = (destKey: string) => {
const first = map(
async (chunk: Chunk) => {
chunk.mapped.push(1);
await sleep(_rate);
return chunk;
},
{ objectMode: true },
);
const second = map(
(chunk: Chunk) => {
pendingReads--;
chunk.mapped.push(2);
return chunk;
},
{ objectMode: true, highWaterMark: 1 },
);
first.pipe(second).pipe(sink);
return first;
};
const _demux = demux(
construct,
{ key: "key" },
{
objectMode: true,
highWaterMark,
},
);
_demux.on("error", err => {
t.end(err);
});
_demux.on("drain", () => {
expect(_demux._writableState.length).to.be.equal(0);
expect(performance.now() - start).to.be.greaterThan(
_rate * input.length,
);
t.pass();
});
const input = [
{ key: "a", mapped: [] },
{ key: "a", mapped: [] },
{ key: "a", mapped: [] },
{ key: "a", mapped: [] },
{ key: "a", mapped: [] },
];
let pendingReads = input.length;
input.forEach(item => {
_demux.write(item);
});
const start = performance.now();
},
);
test.cb( test.cb(
"demux() should emit drain event immediately when second stream is bottleneck", "demux() should emit drain event immediately when second stream is bottleneck",
t => { t => {
t.plan(6); t.plan(6);
const highWaterMark = 5; const highWaterMark = 5;
const _rate = 200; const slowProcessorSpeed = 200;
interface Chunk { interface Chunk {
key: string; key: string;
mapped: number[]; mapped: number[];
@ -484,7 +414,7 @@ test.cb(
const second = map( const second = map(
async (chunk: Chunk) => { async (chunk: Chunk) => {
pendingReads--; pendingReads--;
await sleep(_rate); await sleep(slowProcessorSpeed);
chunk.mapped.push(2); chunk.mapped.push(2);
expect(second._writableState.length).to.be.equal(1); expect(second._writableState.length).to.be.equal(1);
expect(first._readableState.length).to.equal(pendingReads); expect(first._readableState.length).to.equal(pendingReads);
@ -510,7 +440,9 @@ test.cb(
_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.lessThan(_rate); expect(performance.now() - start).to.be.lessThan(
slowProcessorSpeed,
);
t.pass(); t.pass();
}); });
@ -530,7 +462,7 @@ test.cb(
}, },
); );
test("demux() should only emit drain event when all streams are writable", t => { test.only("demux() should only emit drain event when all streams are writable", t => {
t.plan(1); t.plan(1);
const highWaterMark = 2; const highWaterMark = 2;
interface Chunk { interface Chunk {
@ -558,12 +490,12 @@ test("demux() should only emit drain event when all streams are writable", t =>
chunk.mapped.push(1); chunk.mapped.push(1);
return chunk; return chunk;
}, },
{ objectMode: true }, { objectMode: true, highWaterMark: 1 },
); );
const second = map( const second = map(
async (chunk: Chunk) => { async (chunk: Chunk) => {
await sleep(50); await sleep(2000);
chunk.mapped.push(2); chunk.mapped.push(2);
return chunk; return chunk;
}, },
@ -578,11 +510,10 @@ test("demux() should only emit drain event when all streams are writable", t =>
{ key: "key" }, { key: "key" },
{ {
objectMode: true, objectMode: true,
highWaterMark,
}, },
); );
_demux.on("error", err => { _demux.on("error", err => {
reject(); reject(err);
}); });
const input = [ const input = [
@ -590,16 +521,13 @@ test("demux() should only emit drain event when all streams are writable", t =>
{ key: "a", mapped: [] }, { key: "a", mapped: [] },
{ key: "c", mapped: [] }, { key: "c", mapped: [] },
{ key: "c", mapped: [] }, { key: "c", mapped: [] },
{ key: "b", mapped: [] }, // should only be recieved after a becomes writable { key: "b", mapped: [] }, // should only be recieved after a and c become writable
]; ];
let pendingReads = input.length; let pendingReads = input.length;
const start = performance.now(); const start = performance.now();
for (const item of input) { for (const item of input) {
const res = _demux.write(item); console.log("DEMUX", _demux.write(item));
if (!res) {
await sleep(50);
}
} }
}); });
}); });