@drip_sync/drip
Version:
Scalable incremental sync for MongoDB aggregation pipelines
75 lines (74 loc) • 2.22 kB
JavaScript
;
Object.defineProperty(exports, "__esModule", { value: true });
exports.streamSquashMerge = streamSquashMerge;
exports.streamAppend = streamAppend;
exports.streamTake = streamTake;
async function* streamSquashMerge(ss, lt) {
try {
const ress = await Promise.all(ss.map((s) => s.next()));
while (true) {
let state = {
idxs: undefined,
};
for (const [idx, res] of ress.entries()) {
if (res.done)
continue;
if (typeof state.idxs === "undefined" ||
lt(res.value, state.smallest)) {
state = {
idxs: [idx],
smallest: res.value,
};
}
else if (!lt(state.smallest, res.value)) {
state = {
idxs: [...state.idxs, idx],
smallest: state.smallest,
};
}
}
if (typeof state.idxs === "undefined") {
// All streams are done
break;
}
// Yield the smallest element, giving priority
// to the instance coming from earlier streams
// if there are multiple minimums
yield state.smallest;
// Advance the relevant stream(s)
const news = await Promise.all(state.idxs.map((idx) => ss[idx].next()));
for (const [i, idx] of state.idxs.entries()) {
ress[idx] = news[i];
}
}
}
finally {
// Close the streams
await Promise.all(ss.map((s) => s.return()));
}
}
async function* streamAppend(ss) {
try {
for (const s of ss) {
yield* s;
}
}
finally {
// Close the streams
await Promise.all(ss.map((s) => s.return()));
}
}
async function* streamTake(limit, s) {
if (limit == 0) {
// Close the stream
await s.return();
return;
}
let cnt = 0;
for await (const elem of s) {
yield elem;
cnt++;
if (cnt >= limit)
break;
}
}