Created
March 2, 2024 22:38
-
-
Save daneren2005/e035b940429b649f500ca37cecdccc7b to your computer and use it in GitHub Desktop.
Thread-safe MemPool
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import type { Pow2 } from './interfaces/pow2'; | |
import type { TypedArray } from './interfaces/typed-array'; | |
import { lock, unlock } from './lock/simple-lock'; | |
import { typedArray } from './utils/typedarray'; | |
const STATE_FREE = 0; | |
const STATE_USED = 1; | |
const STATE_TOP = 2; | |
const STATE_END = 3; | |
const STATE_ALIGN = 4; | |
const STATE_FLAGS = 5; | |
const STATE_MIN_SPLIT = 6; | |
const MASK_COMPACT = 1; | |
const MASK_SPLIT = 2; | |
const SIZEOF_STATE = 8 * 4; | |
const MEM_BLOCK_SIZE = 0; | |
const MEM_BLOCK_NEXT = 1; | |
const SIZEOF_MEM_BLOCK = 2 * 4; | |
// Copied from https://github.com/thi-ng/umbrella/blob/develop/packages/malloc/src/pool.ts | |
// Changes: Atomic.load/store on state, simple lock on malloc/free | |
export default class MemoryBuffer { | |
buf: ArrayBufferLike; | |
protected readonly start: number; | |
protected u8: Uint8Array; | |
protected u32: Uint32Array; | |
protected state: Uint32Array; | |
protected lock: Int32Array; | |
constructor(opts: Partial<MemoryBufferConfig> = {}) { | |
this.buf = opts.buf ? opts.buf : new ArrayBuffer(opts.size || 0x1000); | |
this.start = opts.start != null ? align(Math.max(opts.start, 0), 4) : 0; | |
this.u8 = new Uint8Array(this.buf); | |
this.u32 = new Uint32Array(this.buf); | |
this.state = new Uint32Array(this.buf, this.start, SIZEOF_STATE / 4); | |
this.lock = new Int32Array(this.buf, this.start + this.state.byteLength - 4, 1); | |
if(!opts.skipInitialization) { | |
const _align = opts.align || 8; | |
if(_align < 8) { | |
throw new Error(`invalid alignment: ${_align}, must be a pow2 and >= 8`); | |
} | |
const top = this.initialTop(_align); | |
const resolvedEnd = | |
opts.end != null | |
? Math.min(opts.end, this.buf.byteLength) | |
: this.buf.byteLength; | |
if(top >= resolvedEnd) { | |
throw new Error( | |
`insufficient address range (0x${this.start.toString( | |
16 | |
)} - 0x${resolvedEnd.toString(16)})` | |
); | |
} | |
this.align = _align; | |
this.doCompact = opts.compact !== false; | |
this.doSplit = opts.split !== false; | |
this.minSplit = opts.minSplit || 16; | |
this.end = resolvedEnd; | |
this.top = top; | |
this._free = 0; | |
this._used = 0; | |
} | |
} | |
stats(): Readonly<MemoryBufferStats> { | |
const listStats = (block: number) => { | |
let count = 0; | |
let size = 0; | |
while(block) { | |
count++; | |
size += this.blockSize(block); | |
block = this.blockNext(block); | |
if(block > this.end) { | |
console.error(`Trying to get stats for block past end of buffer: ${block} > ${this.end}`); | |
break; | |
} | |
} | |
return { count, size }; | |
}; | |
const free = listStats(this._free); | |
return { | |
free, | |
used: listStats(this._used), | |
top: this.top, | |
available: this.end - this.top + free.size, | |
total: this.buf.byteLength | |
}; | |
} | |
callocAs<T extends Type>(type: T, num: number, fill = 0) { | |
const block = this.mallocAs(type, num); | |
block && block.fill(fill); | |
return block; | |
} | |
mallocAs<T extends Type>(type: T, num: number) { | |
const addr = this.malloc(num * SIZEOF[type]); | |
return addr ? typedArray(type, this.buf, addr, num) : undefined; | |
} | |
calloc(bytes: number, fill = 0) { | |
const addr = this.malloc(bytes); | |
addr && this.u8.fill(fill, addr, addr + bytes); | |
return addr; | |
} | |
malloc(bytes: number) { | |
if(bytes <= 0) { | |
return 0; | |
} | |
lock(this.lock); | |
const paddedSize = align(bytes + SIZEOF_MEM_BLOCK, this.align); | |
const end = this.end; | |
let top = this.top; | |
let block = this._free; | |
let prev = 0; | |
while(block) { | |
const blockSize = this.blockSize(block); | |
const isTop = block + blockSize >= top; | |
if(isTop || blockSize >= paddedSize) { | |
let result = this.mallocTop( | |
block, | |
prev, | |
blockSize, | |
paddedSize, | |
isTop | |
); | |
unlock(this.lock); | |
return result; | |
} | |
prev = block; | |
block = this.blockNext(block); | |
} | |
block = top; | |
top = block + paddedSize; | |
if(top <= end) { | |
this.initBlock(block, paddedSize, this._used); | |
this._used = block; | |
this.top = top; | |
let result = blockDataAddress(block); | |
unlock(this.lock); | |
return result; | |
} | |
unlock(this.lock); | |
return 0; | |
} | |
private mallocTop( | |
block: number, | |
prev: number, | |
blockSize: number, | |
paddedSize: number, | |
isTop: boolean | |
) { | |
if(isTop && block + paddedSize > this.end) return 0; | |
if(prev) { | |
this.unlinkBlock(prev, block); | |
} else { | |
this._free = this.blockNext(block); | |
} | |
this.setBlockNext(block, this._used); | |
this._used = block; | |
if(isTop) { | |
this.top = block + this.setBlockSize(block, paddedSize); | |
} else if(this.doSplit) { | |
const excess = blockSize - paddedSize; | |
excess >= this.minSplit && | |
this.splitBlock(block, paddedSize, excess); | |
} | |
return blockDataAddress(block); | |
} | |
realloc(ptr: number, bytes: number) { | |
if(bytes <= 0) { | |
return 0; | |
} | |
const oldAddr = blockSelfAddress(ptr); | |
let newAddr = 0; | |
let block = this._used; | |
let blockEnd = 0; | |
while(block) { | |
if(block === oldAddr) { | |
[newAddr, blockEnd] = this.reallocBlock(block, bytes); | |
break; | |
} | |
block = this.blockNext(block); | |
} | |
// copy old block contents to new addr | |
if(newAddr && newAddr !== oldAddr) { | |
this.u8.copyWithin( | |
blockDataAddress(newAddr), | |
blockDataAddress(oldAddr), | |
blockEnd | |
); | |
} | |
return blockDataAddress(newAddr); | |
} | |
private reallocBlock(block: number, bytes: number) { | |
const blockSize = this.blockSize(block); | |
const blockEnd = block + blockSize; | |
const isTop = blockEnd >= this.top; | |
const paddedSize = align(bytes + SIZEOF_MEM_BLOCK, this.align); | |
// shrink & possibly split existing block | |
if(paddedSize <= blockSize) { | |
if(this.doSplit) { | |
const excess = blockSize - paddedSize; | |
if(excess >= this.minSplit) { | |
this.splitBlock(block, paddedSize, excess); | |
} else if(isTop) { | |
this.top = block + paddedSize; | |
} | |
} else if(isTop) { | |
this.top = block + paddedSize; | |
} | |
return [block, blockEnd]; | |
} | |
// try to enlarge block if current top | |
if(isTop && block + paddedSize < this.end) { | |
this.top = block + this.setBlockSize(block, paddedSize); | |
return [block, blockEnd]; | |
} | |
// fallback to free & malloc | |
this.free(block); | |
return [blockSelfAddress(this.malloc(bytes)), blockEnd]; | |
} | |
reallocArray<T extends TypedArray>(array: T, num: number): T | undefined { | |
if(array.buffer !== this.buf) { | |
return; | |
} | |
const addr = this.realloc( | |
array.byteOffset, | |
num * array.BYTES_PER_ELEMENT | |
); | |
return addr | |
? new (<any>array.constructor)(this.buf, addr, num) | |
: undefined; | |
} | |
free(ptrOrArray: number | TypedArray) { | |
let addr: number; | |
if(typeof ptrOrArray !== 'number') { | |
if(ptrOrArray.buffer !== this.buf) { | |
return false; | |
} | |
addr = ptrOrArray.byteOffset; | |
} else { | |
addr = ptrOrArray; | |
} | |
lock(this.lock); | |
addr = blockSelfAddress(addr); | |
let block = this._used; | |
let prev = 0; | |
while(block) { | |
if(block === addr) { | |
if(prev) { | |
this.unlinkBlock(prev, block); | |
} else { | |
this._used = this.blockNext(block); | |
} | |
this.insert(block); | |
this.doCompact && this.compact(); | |
unlock(this.lock); | |
return true; | |
} | |
prev = block; | |
block = this.blockNext(block); | |
} | |
unlock(this.lock); | |
return false; | |
} | |
freeAll() { | |
this._free = 0; | |
this._used = 0; | |
this.top = this.initialTop(); | |
} | |
release() { | |
delete (<any> this).u8; | |
delete (<any> this).u32; | |
delete (<any> this).state; | |
delete (<any> this).buf; | |
return true; | |
} | |
protected get align() { | |
return <Pow2> this.state[STATE_ALIGN]; | |
} | |
protected set align(x: Pow2) { | |
this.state[STATE_ALIGN] = x; | |
} | |
protected get end() { | |
return this.state[STATE_END]; | |
} | |
protected set end(x: number) { | |
this.state[STATE_END] = x; | |
} | |
protected get top() { | |
return Atomics.load(this.state, STATE_TOP); | |
} | |
protected set top(x: number) { | |
Atomics.store(this.state, STATE_TOP, x); | |
} | |
protected get _free() { | |
return Atomics.load(this.state, STATE_FREE); | |
} | |
protected set _free(block: number) { | |
Atomics.store(this.state, STATE_FREE, block); | |
} | |
protected get _used() { | |
return Atomics.load(this.state, STATE_USED); | |
} | |
protected set _used(block: number) { | |
Atomics.store(this.state, STATE_USED, block); | |
} | |
protected get doCompact() { | |
return !!(this.state[STATE_FLAGS] & MASK_COMPACT); | |
} | |
protected set doCompact(flag: boolean) { | |
flag | |
? (this.state[STATE_FLAGS] |= 1 << (MASK_COMPACT - 1)) | |
: (this.state[STATE_FLAGS] &= ~MASK_COMPACT); | |
} | |
protected get doSplit() { | |
return !!(this.state[STATE_FLAGS] & MASK_SPLIT); | |
} | |
protected set doSplit(flag: boolean) { | |
flag | |
? (this.state[STATE_FLAGS] |= 1 << (MASK_SPLIT - 1)) | |
: (this.state[STATE_FLAGS] &= ~MASK_SPLIT); | |
} | |
protected get minSplit() { | |
return this.state[STATE_MIN_SPLIT]; | |
} | |
protected set minSplit(x: number) { | |
if(x <= SIZEOF_MEM_BLOCK) { | |
throw new Error(`illegal min split threshold: ${x}, require at least ${ | |
SIZEOF_MEM_BLOCK + 1 | |
}`); | |
} | |
this.state[STATE_MIN_SPLIT] = x; | |
} | |
protected blockSize(block: number) { | |
return Atomics.load(this.u32, (block >> 2) + MEM_BLOCK_SIZE); | |
} | |
/** | |
* Sets & returns given block size. | |
* | |
* @param block - | |
* @param size - | |
*/ | |
protected setBlockSize(block: number, size: number) { | |
Atomics.store(this.u32, (block >> 2) + MEM_BLOCK_SIZE, size); | |
return size; | |
} | |
protected blockNext(block: number) { | |
return Atomics.load(this.u32, (block >> 2) + MEM_BLOCK_NEXT); | |
} | |
/** | |
* Sets block next pointer to `next`. Use zero to indicate list end. | |
* | |
* @param block - | |
*/ | |
protected setBlockNext(block: number, next: number) { | |
Atomics.store(this.u32, (block >> 2) + MEM_BLOCK_NEXT, next); | |
} | |
/** | |
* Initializes block header with given `size` and `next` pointer. Returns `block`. | |
* | |
* @param block - | |
* @param size - | |
* @param next - | |
*/ | |
protected initBlock(block: number, size: number, next: number) { | |
const idx = block >>> 2; | |
Atomics.store(this.u32, idx + MEM_BLOCK_SIZE, size); | |
Atomics.store(this.u32, idx + MEM_BLOCK_NEXT, next); | |
return block; | |
} | |
protected unlinkBlock(prev: number, block: number) { | |
this.setBlockNext(prev, this.blockNext(block)); | |
} | |
protected splitBlock(block: number, blockSize: number, excess: number) { | |
this.insert( | |
this.initBlock( | |
block + this.setBlockSize(block, blockSize), | |
excess, | |
0 | |
) | |
); | |
this.doCompact && this.compact(); | |
} | |
protected initialTop(_align = this.align) { | |
return ( | |
align(this.start + SIZEOF_STATE + SIZEOF_MEM_BLOCK, _align) - | |
SIZEOF_MEM_BLOCK | |
); | |
} | |
/** | |
* Traverses free list and attempts to recursively merge blocks | |
* occupying consecutive memory regions. Returns true if any blocks | |
* have been merged. Only called if `compact` option is enabled. | |
*/ | |
protected compact() { | |
let block = this._free; | |
let prev = 0; | |
let scan = 0; | |
let scanPrev: number; | |
let res = false; | |
while(block) { | |
scanPrev = block; | |
scan = this.blockNext(block); | |
while(scan && scanPrev + this.blockSize(scanPrev) === scan) { | |
// console.log("merge:", scan.addr, scan.size); | |
scanPrev = scan; | |
scan = this.blockNext(scan); | |
} | |
if(scanPrev !== block) { | |
const newSize = scanPrev - block + this.blockSize(scanPrev); | |
// console.log("merged size:", newSize); | |
this.setBlockSize(block, newSize); | |
const next = this.blockNext(scanPrev); | |
let tmp = this.blockNext(block); | |
while(tmp && tmp !== next) { | |
// console.log("release:", tmp.addr); | |
const tn = this.blockNext(tmp); | |
this.setBlockNext(tmp, 0); | |
tmp = tn; | |
} | |
this.setBlockNext(block, next); | |
res = true; | |
} | |
// re-adjust top if poss | |
if(block + this.blockSize(block) >= this.top) { | |
this.top = block; | |
prev | |
? this.unlinkBlock(prev, block) | |
: (this._free = this.blockNext(block)); | |
} | |
prev = block; | |
block = this.blockNext(block); | |
} | |
return res; | |
} | |
/** | |
* Inserts given block into list of free blocks, sorted by address. | |
* | |
* @param block - | |
*/ | |
protected insert(block: number) { | |
let ptr = this._free; | |
let prev = 0; | |
while(ptr) { | |
if(block <= ptr) break; | |
prev = ptr; | |
ptr = this.blockNext(ptr); | |
} | |
if(prev) { | |
this.setBlockNext(prev, block); | |
} else { | |
this._free = block; | |
} | |
this.setBlockNext(block, ptr); | |
} | |
} | |
/** | |
* Returns a block's data address, based on given alignment. | |
* | |
* @param blockAddress - | |
*/ | |
const blockDataAddress = (blockAddress: number) => blockAddress > 0 ? blockAddress + SIZEOF_MEM_BLOCK : 0; | |
/** | |
* Returns block start address for given data address and alignment. | |
* | |
* @param dataAddress - | |
*/ | |
const blockSelfAddress = (dataAddress: number) => dataAddress > 0 ? dataAddress - SIZEOF_MEM_BLOCK : 0; | |
const align = (addr: number, size: number) => (size--, addr + size & ~size); | |
const SIZEOF = { | |
u8: 1, | |
u8c: 1, | |
i8: 1, | |
u16: 2, | |
i16: 2, | |
u32: 4, | |
i32: 4, | |
i64: 8, | |
u64: 8, | |
f32: 4, | |
f64: 8 | |
}; | |
type Type = 'u8' | 'u8c' | 'i8' | 'u16' | 'i16' | 'u32' | 'i32' | 'f32' | 'f64'; | |
interface MemoryBufferConfig { | |
/** | |
* Backing ArrayBuffer (or SharedArrayBuffer). If not given, a new | |
* one will be created with given `size`. | |
*/ | |
buf: ArrayBufferLike; | |
/** | |
* Byte size for newly created ArrayBuffers (if `buf` is not given). | |
* | |
* @defaultValue 0x1000 (4KB) | |
*/ | |
size: number; | |
/** | |
* Anchor index (byte address) inside the array buffer. The MemPool | |
* stores its internal state from the given address and heap space | |
* starts at least 32 bytes later (depending on chosen `align` | |
* value). Unlike allocator state variables, `start`` cannot be | |
* saved inside the array buffer itself. If the ArrayBuffer is | |
* passed to other consumers they must use the same start value. | |
* MUST be multiple of 4. | |
* | |
* @defaultValue 0 | |
*/ | |
start: number; | |
/** | |
* Byte address (+1) of the end of the memory region managed by the | |
* {@link MemPool}. | |
* | |
* @defaultValue end of the backing ArrayBuffer | |
*/ | |
end: number; | |
/** | |
* Number of bytes to align memory blocks to. MUST be a power of 2 | |
* and >= 8. Use 16 if the pool is being used for allocating memory | |
* used in SIMD operations. | |
* | |
* @defaultValue 8 | |
*/ | |
align: Pow2; | |
/** | |
* Flag to configure memory block compaction. If true, | |
* adjoining free blocks (in terms of address space) will be merged | |
* to minimize fragementation. | |
* | |
* @defaultValue true | |
*/ | |
compact: boolean; | |
/** | |
* Flag to configure memory block splitting. If true, and when the | |
* allocator is re-using a previously freed block larger than the | |
* requested size, the block will be split to minimize wasted/unused | |
* memory. The splitting behavior can further customized via the | |
* `minSplit` option. | |
* | |
* @defaultValue true | |
*/ | |
split: boolean; | |
/** | |
* Only used if `split` behavior is enabled. Defines min number of | |
* excess bytes available in a block for memory block splitting to | |
* occur. | |
* | |
* @defaultValue 16, MUST be > 8 | |
*/ | |
minSplit: number; | |
/** | |
* Only needed when sharing the underlying ArrayBuffer. If true, the | |
* {@link MemPool} constructor will NOT initialize its internal state and | |
* assume the underlying ArrayBuffer has already been initialized by | |
* another {@link MemPool} instance. If this option is used, `buf` MUST be | |
* given. | |
* | |
* @defaultValue false | |
*/ | |
skipInitialization: boolean; | |
} | |
interface MemoryBufferStats { | |
/** | |
* Free block stats. | |
*/ | |
free: { count: number; size: number }; | |
/** | |
* Used block stats. | |
*/ | |
used: { count: number; size: number }; | |
/** | |
* Current top address. | |
*/ | |
top: number; | |
/** | |
* Bytes available | |
*/ | |
available: number; | |
/** | |
* Total pool size. | |
*/ | |
total: number; | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment