Skip to content
This repository has been archived by the owner on Apr 17, 2023. It is now read-only.

Commit

Permalink
feat(Rest): custom bucket support
Browse files Browse the repository at this point in the history
  • Loading branch information
didinele committed Jan 20, 2022
1 parent d53973f commit 7b0f400
Show file tree
Hide file tree
Showing 7 changed files with 177 additions and 67 deletions.
69 changes: 69 additions & 0 deletions libs/rest/src/fetcher/BaseBucket.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,69 @@
import { DiscordFetchOptions } from './Fetch';
import type { Rest } from '../struct';

export interface BucketConstructor {
makeRoute(method: string, url: string): string;
new (rest: Rest, route: string): BaseBucket;
}

/**
* Base bucket class - used to deal with all of the HTTP semantics
*/
export abstract class BaseBucket {
/**
* Time after a Bucket is destroyed if unused
*/
public static readonly BUCKET_TTL = 1e4;

/**
* Creates a simple API route representation (e.g. /users/:id), used as an identifier for each bucket.
*
* Credit to https://github.com/abalabahaha/eris
*/
public static makeRoute(method: string, url: string) {
let route = url
.replace(/\/([a-z-]+)\/(?:[0-9]{17,19})/g, (match, p) => (['channels', 'guilds', 'webhook'].includes(p) ? match : `/${p}/:id`))
.replace(/\/invites\/[\w\d-]{2,}/g, '/invites/:code')
.replace(/\/reactions\/[^/]+/g, '/reactions/:id')
.replace(/^\/webhooks\/(\d+)\/[A-Za-z0-9-_]{64,}/, '/webhooks/$1/:token')
.replace(/\?.*$/, '');

// Message deletes have their own rate limit
if (method === 'delete' && route.endsWith('/messages/:id')) {
route = method + route;
}

// In this case, /channels/[idHere]/messages is correct,
// however /channels/[idHere] is not. we need "/channels/:id"
if (/^\/channels\/[0-9]{17,19}$/.test(route)) {
route = route.replace(/[0-9]{17,19}/, ':id');
}

return route;
}

protected readonly _destroyTimeout: NodeJS.Timeout;

public ['constructor']!: typeof BaseBucket;

public constructor(
public readonly rest: Rest,
public readonly route: string
) {
// This is in the base constructor for backwards compatibility - in the future it'll be only in the Bucket class
this._destroyTimeout = setTimeout(() => this.rest.buckets.delete(this.route), this.constructor.BUCKET_TTL).unref();
}

/**
* Shortcut for the manager mutex
*/
public get mutex() {
return this.rest.mutex;
}

/**
* Makes a request to Discord
* @param req Request options
*/
public abstract make<T, D, Q>(req: DiscordFetchOptions<D, Q>): Promise<T>;
}
61 changes: 4 additions & 57 deletions libs/rest/src/struct/Bucket.ts → libs/rest/src/fetcher/Bucket.ts
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import { discordFetch, DiscordFetchOptions } from '../Fetch';
import { discordFetch, DiscordFetchOptions } from './Fetch';
import { CordisRestError, HTTPError } from '../Error';
import type { Rest } from './Rest';
import { BaseBucket } from './BaseBucket';

/**
* Data held to represent ratelimit state for a Bucket
Expand All @@ -13,62 +13,9 @@ export interface RatelimitData {
}

/**
* Represents a rate limiting bucket for Discord's API
* Simple, default sequential bucket
*/
export class Bucket {
public static readonly BUCKET_TTL = 1e4;

/**
* Creates a simple API route representation (e.g. /users/:id), used as an identifier for each bucket.
*
* Credit to https://github.com/abalabahaha/eris
*/
public static makeRoute(method: string, url: string) {
let route = url
.replace(/\/([a-z-]+)\/(?:[0-9]{17,19})/g, (match, p) => (['channels', 'guilds', 'webhook'].includes(p) ? match : `/${p}/:id`))
.replace(/\/invites\/[\w\d-]{2,}/g, '/invites/:code')
.replace(/\/reactions\/[^/]+/g, '/reactions/:id')
.replace(/^\/webhooks\/(\d+)\/[A-Za-z0-9-_]{64,}/, '/webhooks/$1/:token')
.replace(/\?.*$/, '');

// Message deletes have their own rate limit
if (method === 'delete' && route.endsWith('/messages/:id')) {
route = method + route;
}

// In this case, /channels/[idHere]/messages is correct,
// however /channels/[idHere] is not. we need "/channels/:id"
if (/^\/channels\/[0-9]{17,19}$/.test(route)) {
route = route.replace(/[0-9]{17,19}/, ':id');
}

return route;
}

private readonly _destroyTimeout: NodeJS.Timeout;

/**
* @param rest The rest manager using this bucket instance
* @param route The identifier of this bucket
*/
public constructor(
public readonly rest: Rest,
public readonly route: string
) {
this._destroyTimeout = setTimeout(() => this.rest.buckets.delete(this.route), Bucket.BUCKET_TTL).unref();
}

/**
* Shortcut for the manager mutex
*/
public get mutex() {
return this.rest.mutex;
}

/**
* Makes a request to Discord
* @param req Request options
*/
export class Bucket extends BaseBucket {
public async make<T, D, Q>(req: DiscordFetchOptions<D, Q>): Promise<T> {
this._destroyTimeout.refresh();

Expand Down
8 changes: 5 additions & 3 deletions libs/rest/src/Fetch.ts → libs/rest/src/fetcher/Fetch.ts
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ import FormData from 'form-data';
import { URLSearchParams } from 'url';
import AbortController from 'abort-controller';
import { RouteBases } from 'discord-api-types/v9';
import type { Readable } from 'stream';

/**
* Represents a file that can be sent to Discord
Expand Down Expand Up @@ -37,15 +38,16 @@ export interface DiscordFetchOptions<D = RequestBodyData, Q = StringRecord> {
isRetryAfterRatelimit: boolean;
query?: Q | string;
files?: File[];
data?: D;
data?: D | Readable;
domain?: string;
}

/**
* Makes the actual HTTP request
* @param options Options for the request
*/
export const discordFetch = async <D, Q>(options: DiscordFetchOptions<D, Q>) => {
let { path, method, headers, controller, query, files, data } = options;
let { path, method, headers, controller, query, files, data, domain = RouteBases.api } = options;

let queryString: string | null = null;
if (query) {
Expand All @@ -59,7 +61,7 @@ export const discordFetch = async <D, Q>(options: DiscordFetchOptions<D, Q>) =>
).toString();
}

const url = `${RouteBases.api}${path}${queryString ? `?${queryString}` : ''}`;
const url = `${domain}${path}${queryString ? `?${queryString}` : ''}`;

let body: string | FormData;
if (files?.length) {
Expand Down
57 changes: 57 additions & 0 deletions libs/rest/src/fetcher/ProxyBucket.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,57 @@
import { discordFetch, DiscordFetchOptions } from './Fetch';
import { CordisRestError, HTTPError } from '../Error';
import { BaseBucket } from './BaseBucket';
import type { Rest } from '../struct';

/**
* Data held to represent ratelimit state for a Bucket
*/
export interface RatelimitData {
global: boolean;
limit: number;
timeout: number;
remaining: number;
}

/**
* Unconventional Bucket implementation that will hijack all requests (i.e. there is no seperate bucket depending on the route)
*
* This is meant for proxying requests, but will not handle any ratelimiting and will entirely ignore mutexes
*/
export class ProxyBucket extends BaseBucket {
public static override makeRoute() {
return 'proxy';
}

public constructor(
rest: Rest,
route: string
) {
super(rest, route);
// This shouldn't be needed - but for backwards compatibility BaseBucket sets this timeout still
clearTimeout(this._destroyTimeout);
}

public async make<T, D, Q>(req: DiscordFetchOptions<D, Q>): Promise<T> {
let timeout: NodeJS.Timeout;
if (req.implicitAbortBehavior) {
timeout = setTimeout(() => req.controller.abort(), this.rest.abortAfter);
}

const res = await discordFetch(req).finally(() => clearTimeout(timeout));

if (res.status === 429) {
return Promise.reject(new CordisRestError('rateLimited', `${req.method.toUpperCase()} ${req.path}`));
} else if (res.status >= 500 && res.status < 600) {
return Promise.reject(new CordisRestError('internal', `${req.method.toUpperCase()} ${req.path}`));
} else if (!res.ok) {
return Promise.reject(new HTTPError(res.clone(), await res.text()));
}

if (res.headers.get('content-type')?.startsWith('application/json')) {
return res.json() as Promise<T>;
}

return res.blob() as Promise<unknown> as Promise<T>;
}
}
4 changes: 4 additions & 0 deletions libs/rest/src/fetcher/index.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
export * from './BaseBucket';
export * from './Bucket';
export * from './Fetch';
export * from './ProxyBucket';
2 changes: 1 addition & 1 deletion libs/rest/src/index.ts
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
export * from './fetcher';
export * from './mutex';
export * from './struct';
export * from './Constants';
export * from './Error';
export * from './Fetch';
43 changes: 37 additions & 6 deletions libs/rest/src/struct/Rest.ts
Original file line number Diff line number Diff line change
@@ -1,19 +1,31 @@
import { Bucket, RatelimitData } from './Bucket';
import {
BaseBucket,
BucketConstructor,
Bucket,
RatelimitData,
DiscordFetchOptions,
File,
RequestBodyData,
StringRecord
} from '../fetcher';
import { USER_AGENT } from '../Constants';
import { EventEmitter } from 'events';
import { Headers, Response } from 'node-fetch';
import { Mutex, MemoryMutex } from '../mutex';
import AbortController from 'abort-controller';
import { CordisRestError, HTTPError } from '../Error';
import { halt } from '@cordis/common';
import type { DiscordFetchOptions, File, RequestBodyData, StringRecord } from '../Fetch';
import type { Readable } from 'stream';
import { RouteBases } from 'discord-api-types/v9';

/**
* Options for constructing a rest manager
*/
export interface RestOptions {
/**
* How many times to retry making a request before giving up
*
* Tip: If using ProxyBucket you should probably set this to 1 depending on your proxy server's implementation
*/
retries?: number;
/**
Expand All @@ -34,6 +46,14 @@ export interface RestOptions {
* Overwrites the default for `{@link RequestOptions.cacheTime}`
*/
cacheTime?: number;
/**
* Bucket constructor to use
*/
bucket?: BucketConstructor;
/**
* Overwrites the default domain used for every request
*/
domain?: string;
}

export interface Rest {
Expand Down Expand Up @@ -114,7 +134,7 @@ export interface RequestOptions<D, Q> {
/**
* Body to send, if any
*/
data?: D;
data?: D | Readable;
/**
* Wether or not this request should be re-attempted after a ratelimit is waited out
*/
Expand All @@ -132,6 +152,10 @@ export interface RequestOptions<D, Q> {
* @default 10000
*/
cacheTime?: number;
/**
* Overwrites the domain used for this request - not taking into account the option passed into {@link RestOptions}
*/
domain?: string;
}

/**
Expand All @@ -151,13 +175,15 @@ export class Rest extends EventEmitter {
/**
* Current active rate limiting Buckets
*/
public readonly buckets = new Map<string, Bucket>();
public readonly buckets = new Map<string, BaseBucket>();

public readonly retries: number;
public readonly abortAfter: number;
public readonly mutex: Mutex;
public readonly retryAfterRatelimit: boolean;
public readonly cacheTime: number;
public readonly bucket: BucketConstructor;
public readonly domain: string;

/**
* @param auth Your bot's Discord token
Expand All @@ -174,26 +200,30 @@ export class Rest extends EventEmitter {
mutex = new MemoryMutex(),
retryAfterRatelimit = true,
cacheTime = 10000,
bucket = Bucket,
domain = RouteBases.api
} = options;

this.retries = retries;
this.abortAfter = abortAfter;
this.mutex = mutex;
this.retryAfterRatelimit = retryAfterRatelimit;
this.cacheTime = cacheTime;
this.bucket = bucket;
this.domain = domain;
}

/**
* Prepares a request to Discord, associating it to the correct Bucket and attempting to prevent rate limits
* @param options Options needed for making a request; only the path is required
*/
public async make<T, D = RequestBodyData, Q = StringRecord>(options: RequestOptions<D, Q>): Promise<T> {
const route = Bucket.makeRoute(options.method, options.path);
const route = this.bucket.makeRoute(options.method, options.path);

let bucket = this.buckets.get(route);

if (!bucket) {
bucket = new Bucket(this, route);
bucket = new this.bucket(this, route);
this.buckets.set(route, bucket);
}

Expand All @@ -209,6 +239,7 @@ export class Rest extends EventEmitter {
}

options.cacheTime ??= this.cacheTime;
options.domain ??= this.domain;

let isRetryAfterRatelimit = false;

Expand Down

0 comments on commit 7b0f400

Please sign in to comment.