support fractional rate limit scaling

This commit is contained in:
Hazelnoot 2024-12-08 08:32:05 -05:00
parent afb026ebea
commit 2781f53d6b
2 changed files with 18 additions and 21 deletions

View file

@ -173,35 +173,31 @@ export class SkRateLimiterService {
private async limitMin(limit: LegacyRateLimit & { minInterval: number }, actor: string, factor: number): Promise<LimitInfo | null> { private async limitMin(limit: LegacyRateLimit & { minInterval: number }, actor: string, factor: number): Promise<LimitInfo | null> {
const counter = await this.getLimitCounter(limit, actor, 'min'); const counter = await this.getLimitCounter(limit, actor, 'min');
const maxCalls = Math.max(Math.ceil(factor), 1); const minInterval = Math.max(Math.ceil(limit.minInterval / factor), 0);
// Update expiration // Update expiration
if (counter.c >= maxCalls) { if (counter.c > 0) {
const isCleared = this.timeService.now - counter.t >= limit.minInterval; const isCleared = this.timeService.now - counter.t >= minInterval;
if (isCleared) { if (isCleared) {
counter.c = 0; counter.c = 0;
} }
} }
const blocked = counter.c >= maxCalls; const blocked = counter.c > 0;
if (!blocked) { if (!blocked) {
counter.c++; counter.c++;
counter.t = this.timeService.now; counter.t = this.timeService.now;
} }
// Calculate limit status // Calculate limit status
const remaining = Math.max(maxCalls - counter.c, 0); const resetMs = Math.max(Math.ceil(minInterval - (this.timeService.now - counter.t)), 0);
const fullResetMs = Math.max(Math.ceil(limit.minInterval - (this.timeService.now - counter.t)), 0); const resetSec = Math.ceil(resetMs / 1000);
const fullResetSec = Math.ceil(fullResetMs / 1000); const limitInfo: LimitInfo = { blocked, remaining: 0, resetSec, resetMs, fullResetSec: resetSec, fullResetMs: resetMs };
const resetMs = remaining < 1 ? fullResetMs : 0;
const resetSec = remaining < 1 ? fullResetSec : 0;
const limitInfo: LimitInfo = { blocked, remaining, resetSec, resetMs, fullResetSec, fullResetMs,
};
// Update the limit counter, but not if blocked // Update the limit counter, but not if blocked
if (!blocked) { if (!blocked) {
// Don't await, or we will slow down the API. // Don't await, or we will slow down the API.
this.setLimitCounter(limit, actor, counter, fullResetSec, 'min') this.setLimitCounter(limit, actor, counter, resetSec, 'min')
.catch(err => this.logger.error(`Failed to update limit ${limit.key}:min for ${actor}:`, err)); .catch(err => this.logger.error(`Failed to update limit ${limit.key}:min for ${actor}:`, err));
} }
@ -210,9 +206,9 @@ export class SkRateLimiterService {
private async limitBucket(limit: RateLimit, actor: string, factor: number): Promise<LimitInfo> { private async limitBucket(limit: RateLimit, actor: string, factor: number): Promise<LimitInfo> {
const counter = await this.getLimitCounter(limit, actor, 'bucket'); const counter = await this.getLimitCounter(limit, actor, 'bucket');
const bucketSize = Math.max(Math.ceil(limit.size * factor), 1);
const dripRate = (limit.dripRate ?? 1000); const dripRate = (limit.dripRate ?? 1000);
const dripSize = (limit.dripSize ?? 1); const dripSize = (limit.dripSize ?? 1);
const bucketSize = (limit.size * factor);
// Update drips // Update drips
if (counter.c > 0) { if (counter.c > 0) {

View file

@ -424,14 +424,13 @@ describe(SkRateLimiterService, () => {
expect(info.blocked).toBeFalsy(); expect(info.blocked).toBeFalsy();
}); });
it('should scale limit by factor', async () => { it('should scale interval by factor', async () => {
minCounter = { c: 1, t: 0 }; minCounter = { c: 1, t: 0 };
mockTimeService.now += 500;
const i1 = await serviceUnderTest().limit(limit, actor, 2); // 1 + 1 = 2 const info = await serviceUnderTest().limit(limit, actor, 2);
const i2 = await serviceUnderTest().limit(limit, actor, 2); // 2 + 1 = 3
expect(i1.blocked).toBeFalsy(); expect(info.blocked).toBeFalsy();
expect(i2.blocked).toBeTruthy();
}); });
it('should set key expiration', async () => { it('should set key expiration', async () => {
@ -662,12 +661,14 @@ describe(SkRateLimiterService, () => {
expect(info.blocked).toBeFalsy(); expect(info.blocked).toBeFalsy();
}); });
it('should scale limit by factor', async () => { it('should scale limit and interval by factor', async () => {
minCounter = { c: 5, t: 0 }; counter = { c: 5, t: 0 };
minCounter = { c: 1, t: 0 };
mockTimeService.now += 500;
const info = await serviceUnderTest().limit(limit, actor, 2); const info = await serviceUnderTest().limit(limit, actor, 2);
expect(info.blocked).toBeTruthy(); expect(info.blocked).toBeFalsy();
}); });
it('should set key expiration', async () => { it('should set key expiration', async () => {