Skip to content

Instantly share code, notes, and snippets.

Show Gist options
  • Save nfriedly/fdba25bcf9f23d18fe8ac5bd1bad5dac to your computer and use it in GitHub Desktop.
Save nfriedly/fdba25bcf9f23d18fe8ac5bd1bad5dac to your computer and use it in GitHub Desktop.
// source/memory-store.ts
var MemoryStore = class {
/**
* Create a new MemoryStore with an optional custom poolSize
*
* Note that the windowMS option is passed to init() by express-rate-limit
*
* @param [options]
* @param [options.poolSize] - Maximum number of unused objects to keep around. Increase to reduce garbage collection.
*/
constructor({ poolSize } = {}) {
/**
* Maximum number of unused clients to keep in the pool
*/
this.poolSize = 100;
/**
* Confirmation that the keys incremented in once instance of MemoryStore
* cannot affect other instances.
*/
this.localKeys = true;
if (typeof poolSize === "number") {
this.poolSize = poolSize;
}
}
/**
* Method that initializes the store.
*
* @param options {Options} - The options used to setup the middleware.
*/
init(options) {
this.windowMs = options.windowMs;
this.previous = /* @__PURE__ */ new Map();
this.current = /* @__PURE__ */ new Map();
this.pool = [];
if (this.interval) {
clearTimeout(this.interval);
}
this.interval = setInterval(() => {
this.resetPrevious();
}, this.windowMs);
if (this.interval.unref)
this.interval.unref();
}
/**
* Method to increment a client's hit counter.
*
* @param key {string} - The identifier for a client.
*
* @returns {IncrementResponse} - The number of hits and reset time for that client.
*
* @public
*/
async increment(key) {
const client = this.getClient(key);
const now = Date.now();
if (client.resetTime.getTime() <= now) {
this.resetClient(client, now);
}
client.totalHits++;
return client;
}
/**
* Method to decrement a client's hit counter.
*
* @param key {string} - The identifier for a client.
*
* @public
*/
async decrement(key) {
const client = this.getClient(key);
if (client.totalHits > 1)
client.totalHits--;
}
/**
* Method to reset a client's hit counter.
*
* @param key {string} - The identifier for a client.
*
* @public
*/
async resetKey(key) {
this.current.delete(key);
this.previous.delete(key);
}
/**
* Method to reset everyone's hit counter.
*
* @public
*/
async resetAll() {
this.current.clear();
this.previous.clear();
}
/**
* Method to stop the timer (if currently running) and prevent any memory
* leaks.
*
* @public
*/
shutdown() {
clearInterval(this.interval);
void this.resetAll();
}
resetClient(client, now = Date.now()) {
client.totalHits = 0;
client.resetTime.setTime(now + this.windowMs);
}
/**
* Refill the pool, set previous to current, reset current
*/
resetPrevious() {
const temporary = this.previous;
this.previous = this.current;
let poolSpace = this.poolSize - this.pool.length;
for (const client of temporary.values()) {
if (poolSpace > 0) {
this.pool.push(client);
poolSpace--;
} else {
break;
}
}
temporary.clear();
this.current = temporary;
}
/**
* Retrieves or creates a client. Ensures it is in this.current
* @param key IP or other key
* @returns Client
*/
getClient(key) {
if (this.current.has(key)) {
return this.current.get(key);
}
let client;
if (this.previous.has(key)) {
client = this.previous.get(key);
} else if (this.pool.length > 0) {
client = this.pool.pop();
this.resetClient(client);
} else {
client = { totalHits: 0, resetTime: /* @__PURE__ */ new Date() };
this.resetClient(client);
}
this.current.set(key, client);
return client;
}
};
if (!global.gc) throw new Error('execute with --expose-gc')
// heavily weighted towards lower numbers
function weightedRandom(min, max) {
return min + Math.round(max / (Math.random() * max));
}
const keys = [];
for(let i=0;i<100000000;i++) {
// generate a bunch of IPv4-looking strings, but weighted so that low numbers occur more often than higher ones
keys.push(`${weightedRandom(1,255)}.${weightedRandom(1,255)}.${weightedRandom(1,255)}.${weightedRandom(1,255)}`)
}
console.log('created list of', keys.length, 'hits containing', (new Set(keys)).size, 'unique keys');
// runs a million hits, reports time taken in ms
async function runTest(poolSize) {
//console.log('running test with pool size', poolSize)
// try to start with a clean slate
if (global.gc) global.gc();
const store = new MemoryStore({poolSize})
store.init({windowMs: 60*60*1000})
const promises = new Array(1000000);
// warm-up
for(let i=0; i<10000; i++) {
const promise = store.increment(keys[i]);
promises[i] = promise;
if (i % 1000 == 0) {
store.resetPrevious()
}
}
await Promise.all(promises);
//console.log('warmup complete')
let peakPool = -Infinity;
let minPool = Infinity;
const len = keys.length;
const breakpoint = Math.floor(len/1000)
const start = Date.now()
for(let i=0; i<len; i++) {
const promise = store.increment(keys[i]);
promises[i] = promise;
if (i % breakpoint == 0 && i > 0) {
if (i > 10000 * 2) {
// pool doesn't get filed until the second call
minPool = Math.min(minPool, store.pool.length)
}
store.resetPrevious()
peakPool = Math.max(peakPool, store.pool.length)
}
}
await Promise.all(promises);
const time = Date.now() - start
console.log('full round complete, took', time, 'ms for pool size', poolSize, 'with pool size ranging from ', minPool, 'to', peakPool)
return time
}
async function runAll() {
await runTest(0)
await runTest(100)
await runTest(500)
await runTest(1000)
await runTest(5000)
await runTest(10000)
await runTest(30000)
// each Client object is about 152 bytes (in node 16.17.0), so 60k takes a little under 1MB of RAM
await runTest(60000)
await runTest(Infinity)
}
runAll().then(() => console.log('all tests done'))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment