Skip to content

Commit

Permalink
ref!(command): refactor throttling internally
Browse files Browse the repository at this point in the history
  • Loading branch information
Snazzah committed Oct 27, 2023
1 parent 56c27a7 commit 425c108
Show file tree
Hide file tree
Showing 3 changed files with 37 additions and 24 deletions.
34 changes: 24 additions & 10 deletions src/command.ts
Original file line number Diff line number Diff line change
Expand Up @@ -55,7 +55,7 @@ export class SlashCommand<T = any> {
/** The creator responsible for this command. */
readonly creator: BaseSlashCreator;

/** Current throttle objects for the command, mapped by user ID. */
/** @private */
private _throttles = new Map<string, ThrottleObject>();

/**
Expand Down Expand Up @@ -206,26 +206,32 @@ export class SlashCommand<T = any> {
onUnload(): any {}

/**
* Creates/obtains the throttle object for a user, if necessary.
* @param userID ID of the user to throttle for
* @private
* Called in order to throttle command usages before running.
* @param ctx The context being throttled
*/
throttle(userID: string): ThrottleObject | null {
async throttle(ctx: CommandContext): Promise<ThrottleResult | null> {
if (!this.throttling) return null;
const userID = ctx.user.id;

let throttle = this._throttles.get(userID);
if (!throttle) {
if (!throttle || throttle.start + this.throttling.duration * 1000 - Date.now() < 0) {
if (throttle) clearTimeout(throttle.timeout);
throttle = {
start: Date.now(),
usages: 0,
timeout: setTimeout(() => {
this._throttles.delete(userID);
}, this.throttling.duration * 1000)
timeout: setTimeout(() => this._throttles.delete(userID), this.throttling.duration * 1000)
};
this._throttles.set(userID, throttle);
}

return throttle;
// Return throttle result if the user has been throttled
if (throttle.usages + 1 > this.throttling.usages) {
const retryAfter = (throttle.start + this.throttling.duration * 1000 - Date.now()) / 1000;
return { retryAfter };
}
throttle.usages++;

return null;
}

/** Unloads the command. */
Expand Down Expand Up @@ -358,3 +364,11 @@ export interface ThrottleObject {
usages: number;
timeout: any;
}

/** @private */
export interface ThrottleResult {
start?: number;
limit?: number;
remaining?: number;
retryAfter: number;
}
8 changes: 3 additions & 5 deletions src/creator.ts
Original file line number Diff line number Diff line change
Expand Up @@ -632,16 +632,14 @@ export class BaseSlashCreator extends (EventEmitter as any as new () => TypedEve
}

// Throttle the command
const throttle = command.throttle(ctx.user.id);
if (throttle && command.throttling && throttle.usages + 1 > command.throttling.usages) {
const remaining = (throttle.start + command.throttling.duration * 1000 - Date.now()) / 1000;
const data = { throttle, remaining };
const throttleResult = await command.throttle(ctx);
if (throttleResult) {
const data = { throttle: throttleResult, remaining: throttleResult.retryAfter };
this.emit('commandBlock', command, ctx, 'throttling', data);
return command.onBlock(ctx, 'throttling', data);
}

// Run the command
if (throttle) throttle.usages++;
return this._runCommand(command, ctx);
}
}
Expand Down
19 changes: 10 additions & 9 deletions src/rest/sequentialBucket.ts
Original file line number Diff line number Diff line change
Expand Up @@ -154,15 +154,16 @@ export class SequentialBucket {
return this.#execute(request, ++attempts);
}

this.#handler.creator?.emit('rawREST', {
auth: request.options.auth ?? false,
body: request.options.body,
files: request.options.files,
latency: latency,
method: request.method,
response: res.clone() as any,
url: request.url
});
if (this.#handler.creator?.listenerCount('rawREST'))
this.#handler.creator.emit('rawREST', {
auth: request.options.auth ?? false,
body: request.options.body,
files: request.options.files,
latency: latency,
method: request.method,
response: res as any,
url: request.url
});

const retryAfter = this.#handle(request, res, latency);
this.#handler.creator?.emit(
Expand Down

0 comments on commit 425c108

Please sign in to comment.