fix(downloads): share throttle budget globally

This commit is contained in:
Sucukdeluxe
2026-08-12 03:00:36 +02:00
parent 40e4e630cf
commit 86920b5410
6 changed files with 226 additions and 34 deletions
@@ -52,8 +52,16 @@ describe('download policy integration contract', () => {
const end = source.indexOf('const outputFinished = output.finished', start);
const section = source.slice(start, end);
expect(section).toContain('createTokenBucketTransform');
expect(section).toContain('createDownloadThrottleTransform()');
expect(section).toContain("const args = [...streamlinkCmd.prefixArgs, url, getStreamlinkStreamArg(), '--stdout'];");
expect(section).not.toMatch(/args\.push\([^\n]*(?:bandwidth|rate-limit|max-rate|throttle)/i);
});
it('routes queue and clip stdout through one app-wide token bucket budget', () => {
const source = readFileSync(join(process.cwd(), 'src', 'main.ts'), 'utf8');
expect(source).toContain('const downloadThrottleBudget = createTokenBucketBudget(null);');
expect(source).toContain('function createDownloadThrottleTransform(): Transform | undefined');
expect(source.match(/createDownloadThrottleTransform\(\)/g)).toHaveLength(3);
});
});
+36 -1
View File
@@ -1,6 +1,6 @@
import { PassThrough } from 'node:stream';
import { describe, expect, it } from 'vitest';
import { createTokenBucketTransform, type TokenBucketClock } from './token-bucket-transform';
import { createTokenBucketBudget, createTokenBucketTransform, type TokenBucketClock } from './token-bucket-transform';
class ManualClock implements TokenBucketClock {
private nextTimerId = 0;
@@ -75,4 +75,39 @@ describe('app-side token bucket transform', () => {
expect(clock.timerCount).toBe(0);
expect(Buffer.concat(output).toString()).toBe('a');
});
it('shares one byte budget across concurrent output transforms', () => {
const clock = new ManualClock();
const budget = createTokenBucketBudget(2, clock);
const first = createTokenBucketTransform(2, clock, budget);
const second = createTokenBucketTransform(2, clock, budget);
const firstOutput: Buffer[] = [];
const secondOutput: Buffer[] = [];
first.on('data', (chunk: Buffer) => firstOutput.push(Buffer.from(chunk)));
second.on('data', (chunk: Buffer) => secondOutput.push(Buffer.from(chunk)));
first.write(Buffer.from('ab'));
second.write(Buffer.from('cd'));
expect(Buffer.concat(firstOutput).toString()).toBe('ab');
expect(Buffer.concat(secondOutput).toString()).toBe('');
expect(clock.timerCount).toBe(1);
clock.advance(1_000);
expect(Buffer.concat(secondOutput).toString()).toBe('cd');
});
it('seeds an app-wide budget when throttling is enabled after startup', () => {
const clock = new ManualClock();
const budget = createTokenBucketBudget(null, clock);
budget.setMaxBytesPerSecond(2);
const transform = createTokenBucketTransform(2, clock, budget);
const output: Buffer[] = [];
transform.on('data', (chunk: Buffer) => output.push(Buffer.from(chunk)));
transform.write(Buffer.from('ab'));
expect(Buffer.concat(output).toString()).toBe('ab');
});
});
+120 -26
View File
@@ -6,53 +6,147 @@ export interface TokenBucketClock {
clearTimeout(handle: ReturnType<typeof setTimeout>): void;
}
export interface TokenBucketBudget {
reserve(bytes: number, release: () => void): () => void;
setMaxBytesPerSecond(maxBytesPerSecond: number | null): void;
}
interface TokenBucketReservation {
bytes: number;
release: () => void;
cancelled: boolean;
}
const systemClock: TokenBucketClock = {
now: () => Date.now(),
setTimeout: (callback, delayMs) => setTimeout(callback, delayMs),
clearTimeout: (handle) => clearTimeout(handle),
};
class TokenBucketTransform extends Transform {
function assertRate(maxBytesPerSecond: number): void {
if (!Number.isSafeInteger(maxBytesPerSecond) || maxBytesPerSecond <= 0) throw new RangeError('maxBytesPerSecond must be a positive safe integer');
}
class SharedTokenBucketBudget implements TokenBucketBudget {
private availableBytes: number;
private lastRefillAt: number;
private timer: ReturnType<typeof setTimeout> | null = null;
private draining = false;
private readonly reservations: TokenBucketReservation[] = [];
constructor(private readonly maxBytesPerSecond: number, private readonly clock: TokenBucketClock) {
super();
this.availableBytes = maxBytesPerSecond;
constructor(private maxBytesPerSecond: number | null, private readonly clock: TokenBucketClock) {
if (maxBytesPerSecond !== null) assertRate(maxBytesPerSecond);
this.availableBytes = maxBytesPerSecond ?? 0;
this.lastRefillAt = clock.now();
}
reserve(bytes: number, release: () => void): () => void {
const reservation: TokenBucketReservation = { bytes, release, cancelled: false };
this.reservations.push(reservation);
this.drain();
return () => {
if (reservation.cancelled) return;
reservation.cancelled = true;
this.drain();
};
}
setMaxBytesPerSecond(maxBytesPerSecond: number | null): void {
if (maxBytesPerSecond !== null) assertRate(maxBytesPerSecond);
if (this.maxBytesPerSecond === maxBytesPerSecond) return;
const wasUnlimited = this.maxBytesPerSecond === null;
this.maxBytesPerSecond = maxBytesPerSecond;
this.availableBytes = maxBytesPerSecond === null ? 0 : wasUnlimited ? maxBytesPerSecond : Math.min(this.availableBytes, maxBytesPerSecond);
this.lastRefillAt = this.clock.now();
this.drain();
}
private refill(capacity: number): void {
if (this.maxBytesPerSecond === null) return;
const now = this.clock.now();
const elapsed = Math.max(0, now - this.lastRefillAt);
this.availableBytes = Math.min(capacity, this.availableBytes + (elapsed * this.maxBytesPerSecond) / 1000);
this.lastRefillAt = now;
}
private clearTimer(): void {
if (!this.timer) return;
this.clock.clearTimeout(this.timer);
this.timer = null;
}
private removeCancelledReservations(): void {
while (this.reservations[0]?.cancelled) this.reservations.shift();
}
private drain(): void {
if (this.draining) return;
this.draining = true;
try {
this.clearTimer();
while (true) {
this.removeCancelledReservations();
const reservation = this.reservations[0];
if (!reservation) return;
if (this.maxBytesPerSecond === null) {
this.reservations.shift();
reservation.release();
continue;
}
const capacity = Math.max(this.maxBytesPerSecond, reservation.bytes);
this.refill(capacity);
if (this.availableBytes >= reservation.bytes) {
this.availableBytes -= reservation.bytes;
this.reservations.shift();
reservation.release();
continue;
}
const delayMs = Math.max(1, Math.ceil(((reservation.bytes - this.availableBytes) * 1000) / this.maxBytesPerSecond));
this.timer = this.clock.setTimeout(() => {
this.timer = null;
this.drain();
}, delayMs);
return;
}
} finally {
this.draining = false;
}
}
}
class TokenBucketTransform extends Transform {
private cancelReservation: (() => void) | null = null;
constructor(private readonly budget: TokenBucketBudget) {
super();
}
override _transform(chunk: Buffer, _encoding: BufferEncoding, callback: (error?: Error | null) => void): void {
const output = Buffer.from(chunk);
const capacity = Math.max(this.maxBytesPerSecond, output.length);
const release = (): void => {
this.timer = null;
this.cancelReservation = this.budget.reserve(output.length, () => {
this.cancelReservation = null;
if (this.destroyed) return;
const now = this.clock.now();
const elapsed = Math.max(0, now - this.lastRefillAt);
this.availableBytes = Math.min(capacity, this.availableBytes + (elapsed * this.maxBytesPerSecond) / 1000);
this.lastRefillAt = now;
if (this.availableBytes >= output.length) {
this.availableBytes -= output.length;
this.push(output);
callback();
return;
}
const delayMs = Math.max(1, Math.ceil(((output.length - this.availableBytes) * 1000) / this.maxBytesPerSecond));
this.timer = this.clock.setTimeout(release, delayMs);
};
release();
this.push(output);
callback();
});
}
override _destroy(error: Error | null, callback: (error: Error | null) => void): void {
if (this.timer) this.clock.clearTimeout(this.timer);
this.timer = null;
this.cancelReservation?.();
this.cancelReservation = null;
callback(error);
}
}
export function createTokenBucketTransform(maxBytesPerSecond: number, clock: TokenBucketClock = systemClock): Transform {
if (!Number.isSafeInteger(maxBytesPerSecond) || maxBytesPerSecond <= 0) throw new RangeError('maxBytesPerSecond must be a positive safe integer');
return new TokenBucketTransform(maxBytesPerSecond, clock);
export function createTokenBucketBudget(maxBytesPerSecond: number | null, clock: TokenBucketClock = systemClock): TokenBucketBudget {
return new SharedTokenBucketBudget(maxBytesPerSecond, clock);
}
export function createTokenBucketTransform(
maxBytesPerSecond: number,
clock: TokenBucketClock = systemClock,
budget: TokenBucketBudget = createTokenBucketBudget(maxBytesPerSecond, clock),
): Transform {
assertRate(maxBytesPerSecond);
return new TokenBucketTransform(budget);
}