element-web/src/utils/Singleflight.ts

/*
Copyright 2021 The Matrix.org Foundation C.I.C.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

import {EnhancedMap} from "./maps";

// Inspired by https://pkg.go.dev/golang.org/x/sync/singleflight

const keyMap = new EnhancedMap<Object, EnhancedMap<string, unknown>>();

/**
 * Access class to get a singleflight context. Singleflights execute a
 * function exactly once, unless instructed to forget about a result.
 */
export class Singleflight {
    private constructor() {
    }

    /**
     * A void marker to help with returning a value in a singleflight context.
     * If your code doesn't return anything, return this instead.
     */
    public static Void = Symbol("void");

    /**
     * Acquire a singleflight context.
     * @param {Object} instance An instance to associate the context with. Can be any object.
     * @param {string} key A string key relevant to that instance to namespace under.
     * @returns {SingleflightContext} Returns the context to execute the function.
     */
    public static for(instance: Object, key: string): SingleflightContext {
        if (!instance || !key) throw new Error("An instance and key must be supplied");
        return new SingleflightContext(instance, key);
    }

    /**
     * Forgets all results for a given instance.
     * @param {Object} instance The instance to forget about.
     */
    public static forgetAllFor(instance: Object) {
        keyMap.delete(instance);
    }

    /**
     * Forgets all cached results for all instances. Intended for use by tests.
     */
    public static forgetAll() {
        for (const k of keyMap.keys()) {
            keyMap.remove(k);
        }
    }
}

class SingleflightContext {
    public constructor(private instance: Object, private key: string) {
    }

    /**
     * Forget this particular instance and key combination, discarding the result.
     */
    public forget() {
        const map = keyMap.get(this.instance);
        if (!map) return;
        map.remove(this.key);
        if (!map.size) keyMap.remove(this.instance);
    }

    /**
     * Execute a function. If a result is already known, that will be returned instead
     * of executing the provided function. However, if no result is known then the function
     * will be called, with its return value cached. The function must return a value
     * other than `undefined` - take a look at Singleflight.Void if you don't have a return
     * to make.
     * @param {Function} fn The function to execute.
     * @returns The recorded value.
     */
    public do<T>(fn: () => T): T {
        const map = keyMap.getOrCreate(this.instance, new EnhancedMap<string, unknown>());

        // We have to manually getOrCreate() because we need to execute the fn
        let val = <T>map.get(this.key);
        if (val === undefined) {
            val = fn();
            map.set(this.key, val);
        }

        return val;
    }
}