express-rate-limit MemoryStore pool size test.js

// 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'))