gristlabs_grist-core/app/server/lib/MemoryPool.ts

115 lines
3.8 KiB
TypeScript
Raw Normal View History

import Deque from 'double-ended-queue';
/**
* Usage:
*
* OPTION 1, using a callback, which may be async (but doesn't have to be).
*
* await mpool.withReserved(initialSize, async (updateReservation) => {
* ...
* updateReservation(newSize); // if needed
* ...
* });
*
* OPTION 2, lower-level.
*
* Note: dispose() MUST be called (e.g. using try/finally). If not called, other work will
* eventually deadlock waiting for it.
*
* const memoryReservation = await mpool.waitAndReserve(initialSize);
* try {
* ...
* memoryReservation.updateReservation(newSize1); // if needed
* memoryReservation.updateReservation(newSize2); // if needed
* ...
* } finally {
* memoryReservation.dispose();
* }
*
* With both options, it's common for the initialSize to be a pool estimate. You may call
* updateReservation() to update it. If it lowers the estimate, other work may unblock. If it
* raises it, it may delay future work, but will have no impact on work that's already unblocked.
* So it's always safer for initialSize to be an overestimate.
*
* When it's hard to estimate initialSize in bytes, you may specify it as e.g.
* memPool.getTotalSize() / 20. This way at most 20 such parallel tasks may be unblocked at a
* time, and further ones will wait until some release their memory or revise down their estimate.
*/
export class MemoryPool {
private _reservedSize: number = 0;
private _queue = new Deque<MemoryAwaiter>();
constructor(private _totalSize: number) {}
public getTotalSize(): number { return this._totalSize; }
public getReservedSize(): number { return this._reservedSize; }
public getAvailableSize(): number { return this._totalSize - this._reservedSize; }
public isEmpty(): boolean { return this._reservedSize === 0; }
public hasSpace(size: number): boolean { return this._reservedSize + size <= this._totalSize; }
// To avoid failures, allow reserving more than totalSize when memory pool is empty.
public hasSpaceOrIsEmpty(size: number): boolean { return this.hasSpace(size) || this.isEmpty(); }
public numWaiting(): number { return this._queue.length; }
public async waitAndReserve(size: number): Promise<MemoryReservation> {
if (this.hasSpaceOrIsEmpty(size)) {
this._updateReserved(size);
} else {
await new Promise<void>(resolve => this._queue.push({size, resolve}));
}
return new MemoryReservation(size, this._updateReserved.bind(this));
}
public async withReserved(size: number, callback: (updateRes: UpdateReservation) => void|Promise<void>) {
const memRes = await this.waitAndReserve(size);
try {
return await callback(memRes.updateReservation.bind(memRes));
} finally {
memRes.dispose();
}
}
// Update the total size. Returns the old size. This is intended for testing.
public setTotalSize(newTotalSize: number): number {
const oldTotalSize = this._totalSize;
this._totalSize = newTotalSize;
this._checkWaiting();
return oldTotalSize;
}
private _checkWaiting() {
while (!this._queue.isEmpty() && this.hasSpaceOrIsEmpty(this._queue.peekFront()!.size)) {
const item = this._queue.shift()!;
this._updateReserved(item.size);
item.resolve();
}
}
private _updateReserved(sizeDelta: number): void {
this._reservedSize += sizeDelta;
this._checkWaiting();
}
}
type UpdateReservation = (sizeDelta: number) => void;
export class MemoryReservation {
constructor(private _size: number, private _updateReserved: UpdateReservation) {}
public updateReservation(newSize: number) {
this._updateReserved(newSize - this._size);
this._size = newSize;
}
public dispose() {
this.updateReservation(0);
this._updateReserved = undefined as any; // Make sure we don't keep using it after dispose
}
}
interface MemoryAwaiter {
size: number;
resolve: () => void;
}