Add 100MB limit to parsed characters cache

52606616c4b7dac5f57d9e863543274c014bfdec

Cohee <18619528+Cohee1207@users.noreply.github.com>

Parents: afb4acc19b604fd736597d6ce3c090e4b29cd429

src/endpoints/characters.js+6 -2 Viewed

 import { AVATAR_WIDTH, AVATAR_HEIGHT } from '../constants.js';
 import { jsonParser, urlencodedParser } from '../express-common.js';
 import { deepMerge, humanizedISO8601DateTime, tryParse, extractFileFromZipBuffer, LimitedMap } from '../util.js';
 import { TavernCardValidator } from '../validator/TavernCardValidator.js';
 import { parse, write } from '../character-card-parser.js';
 import { readWorldInfoFile } from './worldinfo.js';

 const defaultAvatarPath = './public/img/ai4.png';
 // KV-store for parsed character data
 -const characterDataCache = new Map();
 +// 100 MB limit. Would take roughly 3000 characters to reach this limit
 +const characterDataCache = new LimitedMap(1024 * 1024 * 100);
 // Some Android devices require tighter memory management
 const isAndroid = process.platform === 'android';

     try {
         // Reset the cache
         for (const key of characterDataCache.keys()) {
 +            if (Buffer.isBuffer(inputFile)) {
 +                break;
 +            }
             if (key.startsWith(inputFile)) {
                 characterDataCache.delete(key);
                 break;

src/util.js+179 -0 Viewed

         return false;
+    }
+}
++
 +/**
 + * LimitedMap class that limits the memory usage of string values.
 + */
 +export class LimitedMap {
 +    /**
 +     * Creates an instance of LimitedMap.
 +     * @param {number} maxMemoryInBytes - The maximum allowed memory in bytes for string values.
 +     */
 +    constructor(maxMemoryInBytes) {
 +        if (typeof maxMemoryInBytes !== 'number' || maxMemoryInBytes <= 0) {
 +            throw new Error('maxMemoryInBytes must be a positive number');
 +        }
 +        this.maxMemory = maxMemoryInBytes;
 +        this.currentMemory = 0;
 +        this.map = new Map();
 +        this.queue = [];
 +    }
++
 +    /**
 +     * Estimates the memory usage of a string in bytes.
 +     * Assumes each character occupies 2 bytes (UTF-16).
 +     * @param {string} str
 +     * @returns {number}
 +     */
 +    static estimateStringSize(str) {
 +        return str.length * 2;
 +    }
++
 +    /**
 +     * Adds or updates a key-value pair in the map.
 +     * If adding the new value exceeds the memory limit, evicts oldest entries.
 +     * @param {string} key
 +     * @param {string} value
 +     */
 +    set(key, value) {
 +        if (typeof key !== 'string' || typeof value !== 'string') {
 +            return;
 +        }
++
 +        const newValueSize = LimitedMap.estimateStringSize(value);
++
 +        // If the new value itself exceeds the max memory, reject it
 +        if (newValueSize > this.maxMemory) {
 +            return;
 +        }
++
 +        // Check if the key already exists to adjust memory accordingly
 +        if (this.map.has(key)) {
 +            const oldValue = this.map.get(key);
 +            const oldValueSize = LimitedMap.estimateStringSize(oldValue);
 +            this.currentMemory -= oldValueSize;
 +            // Remove the key from its current position in the queue
 +            const index = this.queue.indexOf(key);
 +            if (index > -1) {
 +                this.queue.splice(index, 1);
 +            }
 +        }
++
 +        // Evict oldest entries until there's enough space
 +        while (this.currentMemory + newValueSize > this.maxMemory && this.queue.length > 0) {
 +            const oldestKey = this.queue.shift();
 +            const oldestValue = this.map.get(oldestKey);
 +            const oldestValueSize = LimitedMap.estimateStringSize(oldestValue);
 +            this.map.delete(oldestKey);
 +            this.currentMemory -= oldestValueSize;
 +        }
++
 +        // After eviction, check again if there's enough space
 +        if (this.currentMemory + newValueSize > this.maxMemory) {
 +            return;
 +        }
++
 +        // Add the new key-value pair
 +        this.map.set(key, value);
 +        this.queue.push(key);
 +        this.currentMemory += newValueSize;
 +    }
++
 +    /**
 +     * Retrieves the value associated with the given key.
 +     * @param {string} key
 +     * @returns {string | undefined}
 +     */
 +    get(key) {
 +        return this.map.get(key);
 +    }
++
 +    /**
 +     * Checks if the map contains the given key.
 +     * @param {string} key
 +     * @returns {boolean}
 +     */
 +    has(key) {
 +        return this.map.has(key);
 +    }
++
 +    /**
 +     * Deletes the key-value pair associated with the given key.
 +     * @param {string} key
 +     * @returns {boolean} - Returns true if the key was found and deleted, else false.
 +     */
 +    delete(key) {
 +        if (!this.map.has(key)) {
 +            return false;
 +        }
 +        const value = this.map.get(key);
 +        const valueSize = LimitedMap.estimateStringSize(value);
 +        this.map.delete(key);
 +        this.currentMemory -= valueSize;
++
 +        // Remove the key from the queue
 +        const index = this.queue.indexOf(key);
 +        if (index > -1) {
 +            this.queue.splice(index, 1);
 +        }
++
 +        return true;
 +    }
++
 +    /**
 +     * Clears all entries from the map.
 +     */
 +    clear() {
 +        this.map.clear();
 +        this.queue = [];
 +        this.currentMemory = 0;
 +    }
++
 +    /**
 +     * Returns the number of key-value pairs in the map.
 +     * @returns {number}
 +     */
 +    size() {
 +        return this.map.size;
 +    }
++
 +    /**
 +     * Returns the current memory usage in bytes.
 +     * @returns {number}
 +     */
 +    totalMemory() {
 +        return this.currentMemory;
 +    }
++
 +    /**
 +     * Returns an iterator over the keys in the map.
 +     * @returns {IterableIterator<string>}
 +     */
 +    keys() {
 +        return this.map.keys();
 +    }
++
 +    /**
 +     * Returns an iterator over the values in the map.
 +     * @returns {IterableIterator<string>}
 +     */
 +    values() {
 +        return this.map.values();
 +    }
++
 +    /**
 +     * Iterates over the map in insertion order.
 +     * @param {Function} callback - Function to execute for each element.
 +     */
 +    forEach(callback) {
 +        this.map.forEach((value, key) => {
 +            callback(value, key, this);
 +        });
 +    }
++
 +    /**
 +     * Makes the LimitedMap iterable.
 +     * @returns {Iterator} - Iterator over [key, value] pairs.
 +     */
 +    [Symbol.iterator]() {
 +        return this.map[Symbol.iterator]();
 +    }
 +}