Handle backpressure

This commit is contained in:
Jerry Kurian 2019-08-30 09:33:29 -04:00
parent 2524d51aa7
commit fe0e53147c
4 changed files with 105 additions and 8 deletions

View File

@ -23,7 +23,7 @@
}, },
"scripts": { "scripts": {
"test": "NODE_PATH=src node node_modules/.bin/ava 'tests/*.spec.ts' -e", "test": "NODE_PATH=src node node_modules/.bin/ava 'tests/*.spec.ts' -e",
"test:debug": "NODE_PATH=src node inspect node_modules/ava/profile.ts", "test:debug": "NODE_PATH=src node inspect node_modules/ava/profile.js",
"test:all": "NODE_PATH=src node node_modules/.bin/ava", "test:all": "NODE_PATH=src node node_modules/.bin/ava",
"lint": "tslint -p tsconfig.json", "lint": "tslint -p tsconfig.json",
"validate:tslint": "tslint-config-prettier-check ./tslint.json", "validate:tslint": "tslint-config-prettier-check ./tslint.json",

View File

@ -14,10 +14,15 @@ export function demux(
} }
class Demux extends Writable { class Demux extends Writable {
private keyMap: { private streamsByKey: {
[key: string]: NodeJS.WritableStream | NodeJS.ReadWriteStream; [key: string]: {
stream: NodeJS.WritableStream | NodeJS.ReadWriteStream;
writable: boolean;
};
}; };
private demuxer: (chunk: any) => string; private demuxer: (chunk: any) => string;
private isWritable: boolean;
private nonWritableStreams: Array<string>;
private construct: ( private construct: (
destKey?: string, destKey?: string,
) => NodeJS.WritableStream | NodeJS.ReadWriteStream; ) => NodeJS.WritableStream | NodeJS.ReadWriteStream;
@ -36,16 +41,51 @@ class Demux extends Writable {
} }
this.demuxer = demuxBy.keyBy || ((chunk: any) => chunk[demuxBy.key!]); this.demuxer = demuxBy.keyBy || ((chunk: any) => chunk[demuxBy.key!]);
this.construct = construct; this.construct = construct;
this.keyMap = {}; this.streamsByKey = {};
this.isWritable = true;
} }
public _write(chunk: any, encoding: string, cb: any) { public _write(chunk: any, encoding: string, cb: any) {
const destKey = this.demuxer(chunk); const destKey = this.demuxer(chunk);
if (this.keyMap[destKey] === undefined) { if (this.streamsByKey[destKey] === undefined) {
this.keyMap[destKey] = this.construct(destKey).on("error", e => { this.streamsByKey[destKey] = {
this.emit("error", e); stream: this.construct(destKey),
writable: true,
};
}
// Throttle when one stream is not writable anymore
// Set writable to false
// keep state of all the streams, if one is not writable demux shouldnt be writable
// Small optimization is to keep writing until you get a following event to the unwritable destination
let res = false;
if (this.isWritable && this.streamsByKey[destKey].writable) {
res = this.streamsByKey[destKey].stream.write(chunk, encoding, cb);
} else if (this.isWritable) {
this.isWritable = false;
// Buffer chunk?
return this.isWritable;
}
/* If write above returns false and the stream written to was writable previously, we need to make demux
* non-writable and update state to know the stream is nonWritable.
* If write returns true and the stream was previously not writable, we need to update which streams
* are non writable and determine if it is safe for demux to become writable (all streams are writable)
*/
if (!res) {
this.streamsByKey[destKey].writable = false;
this.nonWritableStreams.push(destKey);
this.isWritable = false;
this.streamsByKey[destKey].stream.once("drain", () => {
this.streamsByKey[destKey].writable = true;
this.nonWritableStreams = this.nonWritableStreams.filter(
key => key !== destKey,
);
this.isWritable = this.nonWritableStreams.length === 0;
}); });
} }
return this.keyMap[destKey].write(chunk, encoding, cb);
return this.writable;
} }
} }

View File

@ -14,6 +14,7 @@ export function map<T, R>(
writableObjectMode: true, writableObjectMode: true,
}, },
): Transform { ): Transform {
// remove try catch
return new Transform({ return new Transform({
...options, ...options,
async transform(chunk: T, encoding, callback) { async transform(chunk: T, encoding, callback) {
@ -22,6 +23,7 @@ export function map<T, R>(
this.push(mapped); this.push(mapped);
callback(); callback();
} catch (err) { } catch (err) {
console.log("caught error", err.message);
callback(err); callback(err);
} }
}, },

View File

@ -107,3 +107,58 @@ test.cb("should spread per key using keyBy", t => {
input.forEach(event => demuxed.write(event)); input.forEach(event => demuxed.write(event));
demuxed.end(); demuxed.end();
}); });
test.cb("should emit errors", t => {
t.plan(2);
const input = [
{ key: "a", val: 1 },
{ key: "b", val: 2 },
{ key: "a", val: 3 },
{ key: "a", val: 4 },
];
const results = [
{ key: "a", val: 2 },
{ key: "b", val: 3 },
{ key: "a", val: 4 },
{ key: "a", val: 5 },
];
const destinationStreamKeys = [];
const sink = new Writable({
objectMode: true,
write(chunk, enc, cb) {
expect(results).to.deep.include(chunk);
expect(input).to.not.deep.include(chunk);
t.pass();
cb();
},
}).on("unpipe", e => console.log("sink err"));
const construct = (destKey: string) => {
destinationStreamKeys.push(destKey);
const dest = map((chunk: Test) => {
if (chunk.key === "b") {
throw new Error("Caught object with key 'b'");
}
return {
...chunk,
val: chunk.val + 1,
};
}).on("error", e => console.log("got err"));
dest.pipe(sink);
return dest;
};
const demuxed = demux(
construct,
{ keyBy: (chunk: any) => chunk.key },
{ objectMode: true },
);
demuxed.on("error", e => {
expect(e.message).to.equal("Caught object with key 'b'");
t.pass();
t.end();
});
input.forEach(event => demuxed.write(event));
demuxed.end();
});