/**
 * If consumer wants more memory than available, we grow the buffer by this ratio.
 */
const GROW_RATIO = 1.1;
export { AllocSource };
var AllocSource;
(function (AllocSource) {
    AllocSource["PERSISTENT_CHECKPOINTS_CACHE_VALIDATORS"] = "persistent_checkpoints_cache_validators";
    AllocSource["PERSISTENT_CHECKPOINTS_CACHE_STATE"] = "persistent_checkpoints_cache_state";
    AllocSource["ARCHIVE_STATE"] = "archive_state";
})(AllocSource || (AllocSource = {}));
/**
 * A simple implementation to manage a single buffer.
 * This is initially used for state serialization at every epoch and for state reload.
 * We can enhance and use this for other purposes in the future.
 */
export class BufferPool {
    buffer;
    inUse = false;
    currentKey;
    metrics = null;
    constructor(size, metrics = null) {
        this.buffer = new Uint8Array(Math.floor(size * GROW_RATIO));
        this.currentKey = 0;
        if (metrics) {
            this.metrics = metrics.bufferPool;
            metrics.bufferPool.length.addCollect(() => {
                metrics.bufferPool.length.set(this.buffer.length);
            });
        }
    }
    get length() {
        return this.buffer.length;
    }
    /**
     * Returns a buffer of the given size with all 0.
     * If the buffer is already in use, return null.
     * Grow the buffer if the requested size is larger than the current buffer.
     */
    alloc(size, source) {
        return this.doAlloc(size, source, false);
    }
    /**
     * Same to alloc() but the buffer is not zeroed.
     */
    allocUnsafe(size, source) {
        return this.doAlloc(size, source, true);
    }
    doAlloc(size, source, isUnsafe = false) {
        if (this.inUse) {
            this.metrics?.misses.inc({ source });
            return null;
        }
        this.inUse = true;
        this.metrics?.hits.inc({ source });
        this.currentKey += 1;
        if (size > this.buffer.length) {
            this.metrics?.grows.inc();
            this.buffer = new Uint8Array(Math.floor(size * GROW_RATIO));
        }
        const bytes = this.buffer.subarray(0, size);
        if (!isUnsafe) {
            bytes.fill(0);
        }
        return new BufferWithKey(bytes, this.currentKey, this);
    }
    /**
     * Marks the buffer as free.
     */
    free(key) {
        if (key === this.currentKey) {
            this.inUse = false;
        }
    }
}
export class BufferWithKey {
    buffer;
    key;
    pool;
    constructor(buffer, key, pool) {
        this.buffer = buffer;
        this.key = key;
        this.pool = pool;
    }
    [Symbol.dispose]() {
        this.pool.free(this.key);
    }
}
//# sourceMappingURL=bufferPool.js.map