diff --git a/spec/common/misc/throttle.spec.ts b/spec/common/misc/throttle.spec.ts new file mode 100644 index 0000000000..3e5a988341 --- /dev/null +++ b/spec/common/misc/throttle.spec.ts @@ -0,0 +1,110 @@ +import { throttle } from '../../../src/misc/throttle'; +import { sequentialize } from '../../../src/misc/sequentialize'; + +describe('throttle decorator', () => { + it('should call the function once at a time', async () => { + const foo = new Foo(); + const promises = []; + for (let i = 0; i < 10; i++) { + promises.push(foo.bar(1)); + } + await Promise.all(promises); + + expect(foo.calls).toBe(10); + }); + + it('should call the function once at a time for each object', async () => { + const foo = new Foo(); + const foo2 = new Foo(); + const promises = []; + for (let i = 0; i < 10; i++) { + promises.push(foo.bar(1)); + promises.push(foo2.bar(1)); + } + await Promise.all(promises); + + expect(foo.calls).toBe(10); + expect(foo2.calls).toBe(10); + }); + + it('should call the function limit at a time', async () => { + const foo = new Foo(); + const promises = []; + for (let i = 0; i < 10; i++) { + promises.push(foo.baz(1)); + } + await Promise.all(promises); + + expect(foo.calls).toBe(10); + }); + + it('should call the function limit at a time for each object', async () => { + const foo = new Foo(); + const foo2 = new Foo(); + const promises = []; + for (let i = 0; i < 10; i++) { + promises.push(foo.baz(1)); + promises.push(foo2.baz(1)); + } + await Promise.all(promises); + + expect(foo.calls).toBe(10); + expect(foo2.calls).toBe(10); + }); + + it('should work together with sequentialize', async () => { + const foo = new Foo(); + const promises = []; + for (let i = 0; i < 10; i++) { + promises.push(foo.qux(Math.floor(i / 2) * 2)); + } + await Promise.all(promises); + + expect(foo.calls).toBe(5); + }); +}); + +class Foo { + calls = 0; + inflight = 0; + + @throttle(1, () => 'bar') + bar(a: number) { + this.calls++; + this.inflight++; + return new Promise((res) => { + setTimeout(() => { + expect(this.inflight).toBe(1); + this.inflight--; + res(a * 2); + }, Math.random() * 10); + }); + } + + @throttle(5, () => 'baz') + baz(a: number) { + this.calls++; + this.inflight++; + return new Promise((res) => { + setTimeout(() => { + expect(this.inflight).toBeLessThanOrEqual(5); + this.inflight--; + res(a * 3); + }, Math.random() * 10); + }); + } + + @sequentialize((args) => 'qux' + args[0]) + @throttle(1, () => 'qux') + qux(a: number) { + this.calls++; + this.inflight++; + return new Promise((res) => { + setTimeout(() => { + expect(this.inflight).toBe(1); + this.inflight--; + res(a * 3); + }, Math.random() * 10); + }); + } +} diff --git a/src/misc/throttle.ts b/src/misc/throttle.ts new file mode 100644 index 0000000000..ace5cc0e7e --- /dev/null +++ b/src/misc/throttle.ts @@ -0,0 +1,57 @@ +/** + * Use as a Decorator on async functions, it will limit how many times the function can be + * in-flight at a time. + * + * Calls beyond the limit will be queued, and run when one of the active calls finishes + */ +export function throttle(limit: number, throttleKey: (args: any[]) => string) { + return (target: any, propertyKey: string | symbol, + descriptor: TypedPropertyDescriptor<(...args: any[]) => Promise>) => { + const originalMethod: () => Promise = descriptor.value; + const allThrottles = new Map void>>>(); + + const getThrottles = (obj: any) => { + let throttles = allThrottles.get(obj); + if (throttles != null) { + return throttles; + } + throttles = new Map void>>(); + allThrottles.set(obj, throttles); + return throttles; + }; + + return { + value: function(...args: any[]) { + const throttles = getThrottles(this); + const argsThrottleKey = throttleKey(args); + let queue = throttles.get(argsThrottleKey); + if (!queue) { + queue = []; + throttles.set(argsThrottleKey, queue); + } + + return new Promise((resolve, reject) => { + const exec = () => { + originalMethod.apply(this, args) + .finally(() => { + queue.splice(queue.indexOf(exec), 1); + if (queue.length >= limit) { + queue[limit - 1](); + } else if (queue.length === 0) { + throttles.delete(argsThrottleKey); + if (throttles.size === 0) { + allThrottles.delete(this); + } + } + }) + .then(resolve, reject); + }; + queue.push(exec); + if (queue.length <= limit) { + exec(); + } + }); + }, + }; + }; +} diff --git a/src/services/audit.service.ts b/src/services/audit.service.ts index 7939149299..7b821afff4 100644 --- a/src/services/audit.service.ts +++ b/src/services/audit.service.ts @@ -2,6 +2,7 @@ import { ApiService } from '../abstractions/api.service'; import { AuditService as AuditServiceAbstraction } from '../abstractions/audit.service'; import { CryptoFunctionService } from '../abstractions/cryptoFunction.service'; +import { throttle } from '../misc/throttle'; import { Utils } from '../misc/utils'; import { BreachAccountResponse } from '../models/response/breachAccountResponse'; @@ -12,6 +13,7 @@ const PwnedPasswordsApi = 'https://api.pwnedpasswords.com/range/'; export class AuditService implements AuditServiceAbstraction { constructor(private cryptoFunctionService: CryptoFunctionService, private apiService: ApiService) { } + @throttle(100, () => 'passwordLeaked') async passwordLeaked(password: string): Promise { const hashBytes = await this.cryptoFunctionService.hash(password, 'sha1'); const hash = Utils.fromBufferToHex(hashBytes).toUpperCase();