79 lines
2.6 KiB
TypeScript
79 lines
2.6 KiB
TypeScript
import { PassThrough } from 'node:stream';
|
|
import { describe, expect, it } from 'vitest';
|
|
import { createTokenBucketTransform, type TokenBucketClock } from './token-bucket-transform';
|
|
|
|
class ManualClock implements TokenBucketClock {
|
|
private nextTimerId = 0;
|
|
private readonly timers = new Map<number, { dueAt: number; callback: () => void }>();
|
|
nowMs = 0;
|
|
|
|
now(): number {
|
|
return this.nowMs;
|
|
}
|
|
|
|
setTimeout(callback: () => void, delayMs: number): ReturnType<typeof setTimeout> {
|
|
const id = ++this.nextTimerId;
|
|
this.timers.set(id, { dueAt: this.nowMs + delayMs, callback });
|
|
return id as unknown as ReturnType<typeof setTimeout>;
|
|
}
|
|
|
|
clearTimeout(handle: ReturnType<typeof setTimeout>): void {
|
|
this.timers.delete(handle as unknown as number);
|
|
}
|
|
|
|
advance(ms: number): void {
|
|
this.nowMs += ms;
|
|
while (true) {
|
|
const due = [...this.timers.entries()]
|
|
.filter(([, timer]) => timer.dueAt <= this.nowMs)
|
|
.sort(([, left], [, right]) => left.dueAt - right.dueAt)[0];
|
|
if (!due) return;
|
|
this.timers.delete(due[0]);
|
|
due[1].callback();
|
|
}
|
|
}
|
|
|
|
get timerCount(): number {
|
|
return this.timers.size;
|
|
}
|
|
}
|
|
|
|
describe('app-side token bucket transform', () => {
|
|
it('backpressures stdout after its initial bucket without changing the source bytes', () => {
|
|
const clock = new ManualClock();
|
|
const transform = createTokenBucketTransform(2, clock);
|
|
const output: Buffer[] = [];
|
|
transform.on('data', (chunk: Buffer) => output.push(Buffer.from(chunk)));
|
|
|
|
transform.write(Buffer.from('ab'));
|
|
transform.write(Buffer.from('cd'));
|
|
|
|
expect(Buffer.concat(output).toString()).toBe('ab');
|
|
expect(clock.timerCount).toBe(1);
|
|
|
|
clock.advance(999);
|
|
expect(Buffer.concat(output).toString()).toBe('ab');
|
|
|
|
clock.advance(1);
|
|
expect(Buffer.concat(output).toString()).toBe('abcd');
|
|
});
|
|
|
|
it('cancels a pending throttle timer when the output stream is destroyed', () => {
|
|
const clock = new ManualClock();
|
|
const source = new PassThrough();
|
|
const transform = createTokenBucketTransform(1, clock);
|
|
const output: Buffer[] = [];
|
|
source.pipe(transform).on('data', (chunk: Buffer) => output.push(Buffer.from(chunk)));
|
|
|
|
source.write(Buffer.from('a'));
|
|
source.write(Buffer.from('b'));
|
|
expect(clock.timerCount).toBe(1);
|
|
|
|
transform.destroy();
|
|
clock.advance(10_000);
|
|
|
|
expect(clock.timerCount).toBe(0);
|
|
expect(Buffer.concat(output).toString()).toBe('a');
|
|
});
|
|
});
|