manage caches in SkRateLimiterService

This commit is contained in:
Hazelnoot 2025-10-01 11:30:35 -04:00
parent 338dad9a87
commit b32f3b5019
2 changed files with 32 additions and 32 deletions

View file

@ -5,14 +5,14 @@
import { Inject, Injectable } from '@nestjs/common'; import { Inject, Injectable } from '@nestjs/common';
import Redis from 'ioredis'; import Redis from 'ioredis';
import type { TimeService } from '@/core/TimeService.js'; import type { MiUser } from '@/models/_.js';
import type { EnvService } from '@/core/EnvService.js'; import { TimeService } from '@/core/TimeService.js';
import { EnvService } from '@/core/EnvService.js';
import { BucketRateLimit, LegacyRateLimit, LimitInfo, RateLimit, hasMinLimit, isLegacyRateLimit, Keyed, hasMaxLimit, disabledLimitInfo, MaxLegacyLimit, MinLegacyLimit } from '@/misc/rate-limit-utils.js'; import { BucketRateLimit, LegacyRateLimit, LimitInfo, RateLimit, hasMinLimit, isLegacyRateLimit, Keyed, hasMaxLimit, disabledLimitInfo, MaxLegacyLimit, MinLegacyLimit } from '@/misc/rate-limit-utils.js';
import { RoleService } from '@/core/RoleService.js';
import { CacheManagementService, type ManagedMemoryKVCache } from '@/core/CacheManagementService.js';
import { ConflictError } from '@/misc/errors/ConflictError.js'; import { ConflictError } from '@/misc/errors/ConflictError.js';
import { DI } from '@/di-symbols.js'; import { DI } from '@/di-symbols.js';
import { MemoryKVCache } from '@/misc/cache.js';
import type { MiUser } from '@/models/_.js';
import type { RoleService } from '@/core/RoleService.js';
// Sentinel value used for caching the default role template. // Sentinel value used for caching the default role template.
// Required because MemoryKVCache doesn't support null keys. // Required because MemoryKVCache doesn't support null keys.
@ -30,26 +30,24 @@ interface ParsedLimit {
@Injectable() @Injectable()
export class SkRateLimiterService { export class SkRateLimiterService {
// 1-minute cache interval private readonly factorCache: ManagedMemoryKVCache<number>;
private readonly factorCache = new MemoryKVCache<number>(1000 * 60); private readonly lockoutCache: ManagedMemoryKVCache<number>;
// 10-second cache interval
private readonly lockoutCache = new MemoryKVCache<number>(1000 * 10);
private readonly requestCounts = new Map<string, number>(); private readonly requestCounts = new Map<string, number>();
private readonly disabled: boolean; private readonly disabled: boolean;
constructor( constructor(
@Inject('TimeService')
private readonly timeService: TimeService,
@Inject(DI.redisForRateLimit) @Inject(DI.redisForRateLimit)
private readonly redisClient: Redis.Redis, private readonly redisClient: Redis.Redis,
@Inject('RoleService')
private readonly roleService: RoleService, private readonly roleService: RoleService,
private readonly timeService: TimeService,
@Inject('EnvService')
envService: EnvService, envService: EnvService,
cacheManagementService: CacheManagementService,
) { ) {
this.factorCache = cacheManagementService.createMemoryKVCache<number>(1000 * 60); // 1m
this.lockoutCache = cacheManagementService.createMemoryKVCache<number>(1000 * 10); // 10s
this.disabled = envService.env.NODE_ENV === 'test'; this.disabled = envService.env.NODE_ENV === 'test';
} }

View file

@ -3,30 +3,27 @@
* SPDX-License-Identifier: AGPL-3.0-only * SPDX-License-Identifier: AGPL-3.0-only
*/ */
import { GodOfTimeService } from '../../../misc/GodOfTimeService.js';
import type Redis from 'ioredis'; import type Redis from 'ioredis';
import type { MiUser } from '@/models/User.js'; import type { MiUser } from '@/models/User.js';
import type { RolePolicies, RoleService } from '@/core/RoleService.js'; import type { RolePolicies, RoleService } from '@/core/RoleService.js';
import type { Config } from '@/config.js';
import { SkRateLimiterService } from '@/server/SkRateLimiterService.js'; import { SkRateLimiterService } from '@/server/SkRateLimiterService.js';
import { BucketRateLimit, Keyed, LegacyRateLimit } from '@/misc/rate-limit-utils.js'; import { BucketRateLimit, Keyed, LegacyRateLimit } from '@/misc/rate-limit-utils.js';
import { CacheManagementService } from '@/core/CacheManagementService.js';
/* eslint-disable @typescript-eslint/no-non-null-assertion */ import { InternalEventService } from '@/core/InternalEventService.js';
describe(SkRateLimiterService, () => { describe(SkRateLimiterService, () => {
let mockTimeService: { now: number, date: Date } = null!; let mockTimeService: GodOfTimeService;
let mockRedis: Array<(command: [string, ...unknown[]]) => [Error | null, unknown] | null> = null!; let mockRedis: Array<(command: [string, ...unknown[]]) => [Error | null, unknown] | null>;
let mockRedisExec: (batch: [string, ...unknown[]][]) => Promise<[Error | null, unknown][] | null> = null!; let mockRedisExec: (batch: [string, ...unknown[]][]) => Promise<[Error | null, unknown][] | null>;
let mockEnvironment: Record<string, string | undefined> = null!; let mockEnvironment: Record<string, string | undefined>;
let serviceUnderTest: () => SkRateLimiterService = null!; let serviceUnderTest: () => SkRateLimiterService;
let mockDefaultUserPolicies: Partial<RolePolicies> = null!; let mockDefaultUserPolicies: Partial<RolePolicies>;
let mockUserPolicies: Record<string, Partial<RolePolicies>> = null!; let mockUserPolicies: Record<string, Partial<RolePolicies>>;
beforeEach(() => { beforeEach(() => {
mockTimeService = { mockTimeService = new GodOfTimeService();
now: 0,
get date() {
return new Date(mockTimeService.now);
},
};
function callMockRedis(command: [string, ...unknown[]]) { function callMockRedis(command: [string, ...unknown[]]) {
const handlerResults = mockRedis.map(handler => handler(command)); const handlerResults = mockRedis.map(handler => handler(command));
@ -62,9 +59,10 @@ describe(SkRateLimiterService, () => {
}, },
}; };
}, },
reset() { reset: () => Promise.resolve(),
return Promise.resolve(); publish: () => Promise.resolve(),
}, on() {},
off() {},
} as unknown as Redis.Redis; } as unknown as Redis.Redis;
mockEnvironment = Object.create(process.env); mockEnvironment = Object.create(process.env);
@ -82,9 +80,13 @@ describe(SkRateLimiterService, () => {
}, },
} as unknown as RoleService; } as unknown as RoleService;
const fakeConfig = { host: 'example.com' } as unknown as Config;
const internalEventService = new InternalEventService(mockRedisClient, mockRedisClient, fakeConfig);
const cacheManagementService = new CacheManagementService(mockRedisClient, mockTimeService, internalEventService);
let service: SkRateLimiterService | undefined = undefined; let service: SkRateLimiterService | undefined = undefined;
serviceUnderTest = () => { serviceUnderTest = () => {
return service ??= new SkRateLimiterService(mockTimeService, mockRedisClient, mockRoleService, mockEnvService); return service ??= new SkRateLimiterService(mockRedisClient, mockRoleService, mockTimeService, mockEnvService, cacheManagementService);
}; };
}); });