Struct UnifiedMemoryPool 

pub struct UnifiedMemoryPool { /* private fields */ }

A bump-allocated pool carved from a single underlying UnifiedBuffer.

Allocations succeed until the slab is exhausted; the pool does not reclaim freed regions (callers maintain their own life cycle by dropping PoolAllocations, which mark the bump pointer as logically released only in the all-released-at-end discipline used by ephemeral Wasm instances).

Pool teardown contract

Slabs handed out as Wasm linear memory (via crate::wasm_memory::TensorWasmMemoryCreator) preserve the monotonic bump invariant for the lifetime of the pool: the PoolAllocation drop guard returned by Self::allocate is intentionally std::mem::forget-ed by the linear-memory creator, so the bump pointer never rewinds while pooled instances are alive. The bump pointer only rewinds via an explicit Self::reset call, which requires &mut self (every other Arc<UnifiedMemoryPool> clone dropped) AND live == 0 (every issued PoolAllocation AND every pool-backed Wasm linear memory dropped).

As of audit T6, crate::wasm_memory::PooledLinearMemory’s own Drop mirrors the mem::forget(alloc) effect by calling Self::release from its destructor: the bump pointer is not rewound (so any base_ptr held by other pool consumers stays valid), but the live counter is decremented so a future reset() can run once every issued memory has been dropped. Before T6 the counter stayed elevated forever once any pool-backed Wasm memory was issued, permanently blocking reset(); that left operators with no path to recycle a slab across tenants short of dropping the pool itself.

Operators that want per-tenant resets must still satisfy both gates (drop every Arc<UnifiedMemoryPool> clone except their own, and drop every issued memory) — alternately they can continue to use one pool per tenant and drop the pool wholesale. The pinned behaviour lives in crates/tensor-wasm-mem/tests/pool_teardown_contract.rs and the call-site rationale is recorded inline above the std::mem::forget(alloc) line in TensorWasmMemoryCreator::new_memory.

Implementations

impl UnifiedMemoryPool

pub fn new(capacity: usize) -> Result<Self, UnifiedError>

Create a pool that owns capacity bytes on the default device.

pub fn new_on( capacity: usize, device_id: DeviceId, ) -> Result<Self, UnifiedError>

Create a pool that owns capacity bytes on the named device.

pub fn capacity(&self) -> usize

Slab capacity in bytes.

pub fn device_id(&self) -> DeviceId

Device that the underlying slab is anchored to.

pub fn remaining(&self) -> usize

Bytes still available for new allocations.

pub fn live_allocations(&self) -> usize

Outstanding allocation count.

pub fn issued_total(&self) -> u64

Total bytes issued since the pool was created (or last reset).

pub fn allocate( &self, size: usize, align: usize, ) -> Result<PoolAllocation<'_>, UnifiedError>

Allocate size bytes aligned to align (must be a power of two).

Returns Err(UnifiedError::TooLarge { requested, limit }) when the slab is exhausted (carrying the requested size and the pool capacity so downstream layers can plumb the figures into TensorWasmError::MemoryExhausted without parsing strings), or Err(UnifiedError::Allocation(...)) when align is zero / not a power of two / exceeds the maximum alignment cap.

pub fn reset(&mut self) -> Result<(), UnifiedError>

Reset the bump pointer back to zero. Safe to call only when there are no outstanding PoolAllocations; returns an error otherwise.

Why &mut self (audit T4)

reset rewinds the bump pointer so the next allocate will hand out byte ranges that overlap with regions previously issued. If reset took &self, a caller could legitimately hold a &[u8] derived from an earlier allocation (or a raw pointer obtained via the slab base) at the moment of reset, and a subsequent allocate would alias that borrow with freshly-issued memory — a use-after-rewind UB hazard that is invisible to the borrow checker because the bump pointer move only requires a shared borrow of the interior Mutex. Requiring &mut self reflects the actual aliasing contract: no other borrow of self (and therefore no live PoolAllocation, no slab_ptr-derived raw pointer that has been turned into a slice, etc.) may coexist with a reset. Callers holding the pool through Arc<UnifiedMemoryPool> must either drop every clone but one (so Arc::get_mut succeeds) or switch to a per-tenant pool — see the Wasm-backing teardown contract below.

Teardown contract for Wasm-backing pools (audit T6)

Pool slabs that served a crate::wasm_memory::TensorWasmMemoryCreator-issued pool-backed Wasm linear memory (PooledLinearMemory) remain resettable: although PoolAllocation drop guards are intentionally leaked at carve time to keep the bump pointer monotonic (see the std::mem::forget(alloc) site in TensorWasmMemoryCreator::new_memory), PooledLinearMemory’s own Drop impl now mirrors the leak by calling Self::release when the linear memory itself is torn down. The live_allocations counter therefore returns to zero once every issued pool-backed Wasm memory has been dropped, at which point reset succeeds (the existing live > 0 guard below is unchanged; what changed is that live is now actually allowed to reach zero again). The bump pointer is NOT rewound by Drop — monotonic-bump semantics are preserved between explicit reset calls, so any in-flight reads through a base_ptr carved out earlier remain valid until reset is explicitly invoked (and reset is gated by &mut self, which the type system refuses while any other Arc<UnifiedMemoryPool> keepalive exists).

Operators wanting per-tenant resets can therefore either: (a) drop every issued memory then call reset on a uniquely-owned pool, or (b) continue to use one pool per tenant and drop the pool wholesale at tenant teardown.