memo/src/async.ts

import type { AsyncFunc, AsyncResult, Params } from '@js/types';

export interface MemoAsyncOptions<T extends AsyncFunc> {
	/**
	 * The function to call and cache results from
	 */
	func: T;

	/**
	 * Time-to-live for records in the cache. After this amount of time,
	 * a value is considered "stale", and by default, will be evicted
	 * from the cache.
	 * 
	 * This value can be a number (ms), or a function which receives the
	 * params and result of the call and returns a number for each record.
	 */
	ttl: number | MemoAsyncTTLFunc<T>;

	/**
	 * Custom function for generating a cache key. If not provided, a
	 * built-in default will be used instead. (At the moment, that
	 * default is `JSON.stringify()`, which is not guranteed to be
	 * stable or particularly fast; Providing a better implementation
	 * for most functions is recommended)
	 */
	key?: MemoAsyncKeyFunc<T>;

	/**
	 * Callback to run when evicting a record from the cache (usually
	 * because it expired)
	 */
	on_evict?: MemoAsyncOnEvict<T>;

	/**
	 * If provided and non-zero, the cache will keep values in storage
	 * for this long (ms) past their expiration times. During this window,
	 * a subsequent request for the value will result in a new call to the
	 * source to refresh the cache, but the stale value in cache will still
	 * be served immediately to reduce cache misses
	 */
	stale_ttl?: number;

	/**
	 * If provided and non-zero, outstanding (not yet fulfilled) promises
	 * will be cached and returned to subsequent callers for the given
	 * amount of time (ms). If the promise does not resolve before the TTL
	 * passes, it will be evicted from the cache to allow for another attempt.
	 * If the promise rejects, it will be dropped from the cache. If the
	 * promise resolves, the actual result will then be stored
	 */
	promise_ttl?: number;
}

export interface MemoAsyncKeyFunc<T extends AsyncFunc> {
	(params: Params<T>) : string;
}

export interface MemoAsyncTTLFunc<T extends AsyncFunc> {
	(params: Params<T>, result: ReturnType<T>) : number;
}

export interface MemoAsyncOnEvict<T extends AsyncFunc> {
	(evicted: MemoAsyncStoredResult<T>) : any;
}

export interface MemoAsyncStoredResult<T extends AsyncFunc> {
	key: string;
	created: number;
	resolved: number;
	expires: number;
	promise: ReturnType<T>;
	result: AsyncResult<T>;
	hits: number;
	timer: ReturnType<typeof setTimeout>;
}

export function memo_async<T extends AsyncFunc>(opts: MemoAsyncOptions<T>) : T {
	const get_key = opts.key ?? default_key;
	const storage: Record<string, MemoAsyncStoredResult<T>> = Object.create(null);
	const refreshing: Record<string, MemoAsyncStoredResult<T>> = Object.create(null);

	function memoized(...params: Params<T>) : ReturnType<T> {
		const key = get_key(params);

		if (storage[key]) {
			storage[key].hits++;

			// If the record is stale (but still stored because of `stale_ttl`), start
			// a new refresh call in the background
			if (Date.now() >= storage[key].expires) {
				refresh_stale(storage[key], params);
			}

			return storage[key].result;
		}

		const promise = opts.func(...params) as ReturnType<T>;
		store_record(key, params, promise);
		return promise;
	}

	function evict(record: MemoAsyncStoredResult<T>) {
		if (storage[record.key] === record) {
			delete storage[record.key];

			if (refreshing[record.key]?.timer) {
				// If another call is already working on refreshing this record and
				// it has an active `promise_ttl` timer, promote it into the main
				// storage now
				storage[record.key] = refreshing[record.key];
				delete refreshing[record.key];
			}

			if (opts.on_evict) {
				opts.on_evict(record);
			}
		}
	}

	function refresh_stale(old_record: MemoAsyncStoredResult<T>, params: Params<T>) {
		const record: MemoAsyncStoredResult<T> = {
			key: old_record.key,
			created: Date.now(),
			resolved: null,
			expires: null,
			promise: opts.func(...params) as ReturnType<T>,
			result: null,
			hits: 0,
			timer: null,
		};

		refreshing[record.key] = record;

		if (opts.promise_ttl) {
			record.timer = setTimeout(() => on_timeout(record), opts.promise_ttl);
		}

		record.promise.then(
			(result) => on_refresh_resolve(record, params, result),
			(reason) => on_reject(record, reason)
		);
	}

	function store_record(key: string, params: Params<T>, promise: ReturnType<T>) {
		const record: MemoAsyncStoredResult<T> = {
			key,
			created: Date.now(),
			resolved: null,
			expires: null,
			promise,
			result: null,
			hits: 0,
			timer: null,
		};

		if (opts.promise_ttl) {
			storage[key] = record;
			record.timer = setTimeout(() => on_timeout(record), opts.promise_ttl);
		}

		promise.then(
			(result) => on_resolve(record, params, result),
			(reason) => on_reject(record, reason)
		);
	}

	function on_timeout(record: MemoAsyncStoredResult<T>) {
		record.timer = null;
		
		if (storage[record.key] === record) {
			delete storage[record.key];
		}
	}

	function on_resolve(record: MemoAsyncStoredResult<T>, params: Params<T>, result: AsyncResult<T>) {
		if (record.timer) {
			clearTimeout(record.timer);
		}

		const ttl = typeof opts.ttl === 'function'
			? opts.ttl(params, result)
			: opts.ttl;

		const evict_after = opts.stale_ttl
			? ttl + opts.stale_ttl
			: ttl;

		record.result = result;
		record.resolved = Date.now();
		record.expires = record.resolved + ttl;
		record.timer = setTimeout(() => evict(record), evict_after);
		storage[record.key] = record;
	}

	function on_refresh_resolve(record: MemoAsyncStoredResult<T>, params: Params<T>, result: AsyncResult<T>) {
		if (record.timer) {
			clearTimeout(record.timer);
		}

		const ttl = typeof opts.ttl === 'function'
			? opts.ttl(params, result)
			: opts.ttl;

		const evict_after = opts.stale_ttl
			? ttl + opts.stale_ttl
			: ttl;

		record.result = result;
		record.resolved = Date.now();
		record.expires = record.resolved + ttl;
		record.timer = setTimeout(() => evict(record), evict_after);

		if (storage[record.key]?.timer) {
			clearTimeout(storage[record.key].timer);
			evict(storage[record.key]);
		}

		storage[record.key] = record;
	}

	function on_reject(record: MemoAsyncStoredResult<T>, reason: any) {
		if (record.timer) {
			// If a timeout timer was set due to `promise_ttl`, clear it now
			clearTimeout(record.timer);
		}

		if (storage[record.key] === record) {
			delete storage[record.key];
		}
	}
	
	return memoized as unknown as T;
}

function default_key(...params: any[]) : string {
	return JSON.stringify(params);
}