p-ratelimit
Version:
Promise-based utility to make sure you don’t call rate-limited APIs too quickly.
125 lines • 5.39 kB
JavaScript
"use strict";
var __awaiter = (this && this.__awaiter) || function (thisArg, _arguments, P, generator) {
function adopt(value) { return value instanceof P ? value : new P(function (resolve) { resolve(value); }); }
return new (P || (P = Promise))(function (resolve, reject) {
function fulfilled(value) { try { step(generator.next(value)); } catch (e) { reject(e); } }
function rejected(value) { try { step(generator["throw"](value)); } catch (e) { reject(e); } }
function step(result) { result.done ? resolve(result.value) : adopt(result.value).then(fulfilled, rejected); }
step((generator = generator.apply(thisArg, _arguments || [])).next());
});
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.RedisQuotaManager = void 0;
const util_1 = require("../util");
const quotaManager_1 = require("./quotaManager");
/** QuotaManager that coordinates rate limits across servers. */
class RedisQuotaManager extends quotaManager_1.QuotaManager {
/**
* @param channelQuota the overall quota to be split among all clients
* @param channelName unique name for this quota - the Redis pub/sub channel name
* @param client a Redis client (or a pair of Redis clients if using a specialty Redis library)
* @param heartbeatInterval how often to ping the Redis channel (milliseconds)
*/
constructor(channelQuota, channelName, client, heartbeatInterval = 30000) {
// start with 0 concurrency so jobs don’t run until we’re ready
super(Object.assign({}, channelQuota, {
concurrency: channelQuota.fastStart ? channelQuota.concurrency : 0
}));
this.channelQuota = channelQuota;
this.heartbeatInterval = heartbeatInterval;
this.uniqueId = util_1.uniqueId();
this.pingsReceived = new Map();
this.heartbeatTimer = null;
this._ready = Boolean(channelQuota.fastStart);
this.channelName = `ratelimit-${channelName}`;
const clients = Array.isArray(client) ? client : [client];
if (clients.length === 1) {
this.client = clients[0];
if (typeof this.client['duplicate'] !== 'function') {
const msg = '[p-ratelimit RedisQuotaManager] Your Redis client does not ' +
'support the client.duplicate() function. Please provide an array of two ' +
'clients instead.';
throw new Error(msg);
}
this.pubSubClient = this.client['duplicate']();
}
else {
this.client = clients[0];
this.pubSubClient = clients[1];
}
this.register();
}
/** true once the Quota Manager has discovered its peers and calculated its quota */
get ready() {
return this._ready;
}
/** Join the client pool, coordinated by the shared channel on Redis */
register() {
return __awaiter(this, void 0, void 0, function* () {
this.pingsReceived.set(this.uniqueId, Date.now());
this.pubSubClient.on('message', (channel, message) => this.message(channel, message));
yield util_1.promisify(this.pubSubClient['subscribe'].bind(this.pubSubClient))(this.channelName);
this.ping();
if (!this.channelQuota.fastStart) {
yield util_1.sleep(3000);
}
yield this.updateQuota();
this._ready = true;
this.heartbeatTimer = setInterval(() => this.heartbeat(), this.heartbeatInterval);
});
}
/** Send a ping to the shared Redis channel */
ping() {
this.client['publish'](this.channelName, JSON.stringify(this.uniqueId));
}
/** Receive client pings */
message(channel, message) {
if (channel !== this.channelName) {
return;
}
let uniqueId;
try {
uniqueId = JSON.parse(message);
}
catch (_a) {
console.error(`invalid JSON on Redis pub/sub channel ${channel}: ${message}`);
return;
}
const newClient = !this.pingsReceived.has(uniqueId);
this.pingsReceived.set(uniqueId, Date.now());
if (newClient) {
this.ping();
if (this.ready) {
this.updateQuota();
}
}
}
/** Remove outdated clients */
removeOutdatedClients() {
const ancient = Date.now() - this.heartbeatInterval * 3;
const expired = [...this.pingsReceived].filter(([k, v]) => v <= ancient);
expired.forEach(([k, v]) => this.pingsReceived.delete(k));
}
/** Calculate our portion of the overall channel quota */
updateQuota() {
this.removeOutdatedClients();
if (!this.pingsReceived.size) {
return;
}
const newQuota = Object.assign({}, this.channelQuota);
newQuota.rate = Math.floor(newQuota.rate / this.pingsReceived.size);
if (newQuota.concurrency) {
newQuota.concurrency = Math.floor(newQuota.concurrency / this.pingsReceived.size);
}
this._quota = newQuota;
}
/** Let the others know we’re here */
heartbeat() {
this.ping();
if (this.ready) {
this.updateQuota();
}
}
}
exports.RedisQuotaManager = RedisQuotaManager;
//# sourceMappingURL=redisQuotaManager.js.map