Pooled Buffer allocator for NodeJS

example.ts

import BufferPool from "./pool";

// Create pool of 2MB blocks
const Pool = new BufferPool(2 * 1024 * 1024, 1000);

// Allocate uninitialized buffer
const unsafe = Pool.allocUnsafe(1000);

// Allocate zero filled buffer
const safe = Pool.alloc(1000);

// Concatenate
const safe_unsafe = Pool.concat([safe, unsafe]);

// Create from array
const fromArr = Pool.fromArrayLike([0,1,2,3]);

// Create from typed array
const fromInt8 = Pool.fromInt8Array(new Uint8Array([1,2,3,4]));

// Create from hex string
const fromHex = Pool.fromHex('deadbeef');

pool.ts

/**
 * Pooled buffer allocator
 * (c) 2022 Alex D.
 * (c) 2019 puzpuzpuz/nbufpool
 * (c) Node.js contributors
 */

export default class BufferPool {
    #blockSize: number;
    #poolSize: number;
    #registry: FinalizationRegistry<Buffer>;
    #blockPoolRef = new WeakRef<Array<Buffer>>([]);
    #rootSlice: Buffer;
    #rootSliceSymbol = Symbol('rootSlice');
    #offset: number;

    constructor(blockSize: number, poolSize: number) {
        this.#blockSize = blockSize;
        this.#poolSize = poolSize;
        this.#registry = new FinalizationRegistry(this.#reclaim.bind(this));
        this.#createBlock();
    }

    get blockSize() : number {
        return this.#blockSize;
    }

    get poolSize() : number {
        return this.#poolSize;
    }

    /**
     * Clone of Buffer.allocUnsafe function: allocate new 
     *  uninitialized buffer or make a slice from preallocated internal buffer
     * @param size Buffer size in bytes
     * @returns Buffer or slice of internal buffer
     */
    allocUnsafe(size: number) : Buffer {
        if (size < (this.#blockSize >>> 1)) {
            if (size > (this.#blockSize - this.#offset)) {
                this.#createBlock();
            }
            const slice = this.#rootSlice.subarray(this.#offset, this.#offset + size);
            // Keep strong reference to root slice to prevent collection
            //  until the last of slices becomes unavailable
            (<any>slice)[this.#rootSliceSymbol] = this.#rootSlice; // HACK: cast here
            this.#offset += size;
            this.#alignBlock();
            return slice;
        }
        return Buffer.allocUnsafeSlow(size);
    }

    /**
     * Allocate buffer with allocUnsafe and initialize it with zeros
     * @param size Buffer size in bytes
     * @returns Buffer or slice of internal buffer
     */
    alloc(size: number) : Buffer {
        return this.allocUnsafe(size).fill(0);
    }

    /**
     * Construct a copy of given buffer
     * @param from Original buffer
     * @returns A copy
     */
    allocCopy(from: Buffer) : Buffer {
        const copy = this.allocUnsafe(from.length);
        from.copy(copy);
        return copy;
    }

    /**
     * Construction of buffer from array of bytes
     * @param from Array like object
     * @returns Buffer
     */
    fromArrayLike(from: Array<number>) : Buffer {
        const copy = this.allocUnsafe(from.length);
        copy.set(from);
        return copy;
    }

    /**
     * Construction of buffer from typed array of bytes
     * @param from Uint8Array or Int8Array
     * @returns Buffer
     */
    fromInt8Array(from: Uint8Array | Int8Array) : Buffer {
        const copy = this.allocUnsafe(from.length);
        new Uint8Array(copy.buffer).set(new Uint8Array(from.buffer))
        return copy;
    }

    /**
     * Construction of buffer from hex encoded string
     * @param from Hex string
     * @returns Buffer
     */
    fromHex(from: string) : Buffer {
        const buffer = this.allocUnsafe(from.length / 2 | 0);
        const bytesWritten = buffer.write(from, 'hex');
        return buffer.subarray(0, bytesWritten);
    }

    /**
     * Clone of Buffer.concat() function
     * See https://github.com/nodejs/node/blob/c5e74bb50f443e5c623c9736abcc512c8e56fb9f/lib/buffer.js#L545 for details
     * @param list Array of buffers to be concatenated
     * @param length (Optional) Number of bytes in resulting buffer
     */
    concat(list: Array<Buffer>, length?: number) : Buffer {
        let _length = length;
        if (undefined === _length) {
            _length = 0;
            for (let i = 0; i < list.length; i++) {
                if (list[i].length) {
                  _length += list[i].length;
                }
            }
        }

        const buffer = this.allocUnsafe(_length);

        let pos = 0;
        for (let i = 0; i < list.length; i++) {
          const buf = list[i];
          pos += buf.copy(buffer, pos, 0, buf.length);
        }

        // Note: `length` is always equal to `buffer.length` at this point
        if (pos < _length) {
            // Zero-fill the remaining bytes if the specified `length` was more than
            // the actual total length, i.e. if we have some remaining allocated bytes
            // there were not initialized.
            buffer.fill(0, pos, _length);
        }

        return buffer;
    }

    #createBlock() : void {
        let block = this.#takeFromPool() || Buffer.allocUnsafeSlow(this.#blockSize);
        this.#rootSlice = block.subarray();
        this.#registry.register(this.#rootSlice, block);
        this.#offset = 0;
    }

    #takeFromPool() : Buffer | undefined {
        const blockPool = this.#blockPoolRef.deref();
        if (undefined === blockPool) {
            return undefined;
        }
        return blockPool.pop();
    }


    #reclaim(buf: Buffer) : void {
        const blockPool = this.#blockPoolRef.deref() || [];
        this.#blockPoolRef = new WeakRef<Array<Buffer>>(blockPool);
        if (blockPool.length >= this.#poolSize) {
            return;
        }
        blockPool.push(buf);
    }


    #alignBlock() : void {
        // ensure aligned slices
        if (0 !== (this.#offset & 0x7)) {
            this.#offset |= 0x7;
            this.#offset++;
        }
    }
}