TODO Allocator: realloc fixes and tests
diff --git a/libraries/stdlib/wasm/src/kotlin/wasm/unsafe/MemoryAllocation.kt b/libraries/stdlib/wasm/src/kotlin/wasm/unsafe/MemoryAllocation.kt index df0643f..9052c7b 100644 --- a/libraries/stdlib/wasm/src/kotlin/wasm/unsafe/MemoryAllocation.kt +++ b/libraries/stdlib/wasm/src/kotlin/wasm/unsafe/MemoryAllocation.kt
@@ -76,7 +76,7 @@ } @UnsafeWasmMemoryApi -private class MemorySlot(val ptr: Pointer, val size: UInt) { +private data class MemorySlot(val ptr: Pointer, val size: UInt) { // TODO in the usages, this is only ever used where we know the direction of the only possible successful one-way merge. So could optimize it based on that fun tryMerge(b: MemorySlot): MemorySlot? { val a = this @@ -112,103 +112,147 @@ } } -// NOTE: design choice for now: store all the info here (i.e., in WasmGC structs), instead of trying to be clever and use headers in linear memory or similar. +/** + * Returns size but possibly lengthened to align to an implementation-defined alignment. + * + * This means it represents the actual size of any allocation made with the size parameter. + * + * The alignment is currently 8, as it's currently the maximum needed for the Wasm component model canonical ABI. + */ +private fun realAllocationSize(size: UInt): UInt { + val alignment = 8u + + // round up the size to a multiple of 8, so that all addresses are always guaranteed to be aligned to at least 8 + // by adding 7, we're guaranteed to: + // - if size mod 8 == 0: NOT cross the divisible-by-8 boundary + // - if size mod 8 != 0: cross the divisible-by-8 boundary exactly once. + // so after adding 7, just shave off the lower bits under 7 + return (size + (alignment - 1u)) and (alignment - 1u).inv() + // (size + 7u ) & 0xFFFFFFF8u +} + + +// NOTES: +// - design choice for now: store all the info here (i.e., in WasmGC structs), instead of trying to be clever and use headers in linear memory or similar. +// - NOT thread-safe, would need synchronization if not used in a single-threaded environment +// - NOT reentrant, i.e., cannot call any member functions of this, from within any member functions of this // TODO think about this again, but I don't see a reason to use headers if we have other memory available anyway @UnsafeWasmMemoryApi private object FreeList { - val alignment = 8u + // TODO remove + fun debugDump(): String = buildString { + appendLine("FreeList:") + for (it in list) { + appendLine(" ${it.ptr.address} - ${it.ptr.address + it.size}") + } + } // TODO right now, this is the default list implementation, which I assume is an array list. That's not really optimal, because it requires copying around stuff when the number of free slots change, and we can't make use of O(1) element access val list = mutableListOf<MemorySlot>( MemorySlot(Pointer(0u), ((1u shl 31) - 1u)) ) + /** + * TODO(REVIEW) can we get rid of this / make it test only? Should we keep it even in production? + * + * Make sure to not call call free / alloc from anywhere inside free / alloc + */ + private var isAlreadyOperating = false + // NOTE: freeing is when things are merged back together @PublishedApi internal fun free(allocatedSlot: MemorySlot): Unit { - // need to find the slots that this lies in between, in terms of start address - // NOTE: we assume the allocatedSlot does not overlap with anything in the free list, that wouldn't make sense, by definition, allocatedSlot is not free - val minusInsertionPointMinusOne = list.binarySearch { - // because we assume it can't overlap, we know that size is irrelevant here: we'll get the index of the insertion point from this function, and inserting there will not lead to overlap - (it.ptr.address.toLong() - allocatedSlot.ptr.address.toLong()).toInt() - } + check(!isAlreadyOperating) { "Cannot call free from within the allocator" } + isAlreadyOperating = true + try { + check(allocatedSlot.size == realAllocationSize(allocatedSlot.size)) { "Slot to free clearly does not originate from allocated slot: alignment is wrong" } - require(minusInsertionPointMinusOne != 0) { "Slot to free can't already be in the free list" } + // need to find the slots that this lies in between, in terms of start address + // NOTE: we assume the allocatedSlot does not overlap with anything in the free list, that wouldn't make sense, by definition, allocatedSlot is not free + val minusInsertionPointMinusOne = list.binarySearch { + // because we assume it can't overlap, we know that size is irrelevant here: we'll get the index of the insertion point from this function, and inserting there will not lead to overlap + (it.ptr.address.toLong() - allocatedSlot.ptr.address.toLong()).toInt() + } - // convert back to actual insertion point - val insertionPointIndex = -(minusInsertionPointMinusOne + 1) + require(minusInsertionPointMinusOne != 0) { "Double-free: slot to free can't already be in the free list; slot to free: $allocatedSlot; " + FreeList.debugDump() } - // before we insert, try to merge - val leftElement = list.getOrNull(insertionPointIndex - 1) - val rightElement = list.getOrNull(insertionPointIndex) - check( - (leftElement?.ptr?.address ?: UInt.MIN_VALUE) < allocatedSlot.ptr.address && - allocatedSlot.ptr.address < (rightElement?.ptr?.address ?: UInt.MAX_VALUE) - ) { "Binary search has gone wrong" } + // convert back to actual insertion point + val insertionPointIndex = -(minusInsertionPointMinusOne + 1) - // once we start merging anything, the left and right slots might become adjacent, and need to be merged themselves - fun tryMergeLeftAndRight() { - // need to access left and right again here, as they can change during the function + // before we insert, try to merge val leftElement = list.getOrNull(insertionPointIndex - 1) val rightElement = list.getOrNull(insertionPointIndex) + check( + leftElement?.ptr?.address?.let { it < allocatedSlot.ptr.address } ?: true && + rightElement?.ptr?.address?.let { allocatedSlot.ptr.address < it } ?: true + ) { "Binary search has gone wrong" } - if (leftElement == null || rightElement == null) - return + // once we start merging anything, the left and right slots might become adjacent, and need to be merged themselves + fun tryMergeLeftAndRight() { + // need to access left and right again here, as they can change during the function + val leftElement = list.getOrNull(insertionPointIndex - 1) + val rightElement = list.getOrNull(insertionPointIndex) - val successfulMerge = leftElement.tryMerge(rightElement) - if (successfulMerge != null) { - list[insertionPointIndex - 1] = successfulMerge - list.removeAt(insertionPointIndex) + if (leftElement == null || rightElement == null) + return + + val successfulMerge = leftElement.tryMerge(rightElement) + if (successfulMerge != null) { + list[insertionPointIndex - 1] = successfulMerge + list.removeAt(insertionPointIndex) + return + } + } + + val successfulMergeLeft = leftElement?.tryMerge(allocatedSlot) + if (successfulMergeLeft != null) { + list[insertionPointIndex - 1] = successfulMergeLeft + // now that we merged, left and right might be adjacent + tryMergeLeftAndRight() return } - } - val successfulMergeLeft = leftElement?.tryMerge(allocatedSlot) - if (successfulMergeLeft != null) { - list[insertionPointIndex - 1] = successfulMergeLeft - // now that we merged, left and right might be adjacent - tryMergeLeftAndRight() - return - } + val successfulMergeRight = rightElement?.tryMerge(allocatedSlot) + if (successfulMergeRight != null) { + list[insertionPointIndex] = successfulMergeRight + tryMergeLeftAndRight() + return + } - val successfulMergeRight = rightElement?.tryMerge(allocatedSlot) - if (successfulMergeRight != null) { - list[insertionPointIndex] = successfulMergeRight - tryMergeLeftAndRight() - return + // otherwise, we couldn't merge with either side, so by definition there's no overlap, and we just need to insert a new slot + list.add(insertionPointIndex, allocatedSlot) + } finally { + isAlreadyOperating = false } - - // otherwise, we couldn't merge with either side, so by definition there's no overlap, and we just need to insert a new slot - list.add(insertionPointIndex, allocatedSlot) } @PublishedApi internal fun allocate(size: UInt): MemorySlot { - // round up the size to a multiple of 8, so that all addresses are always guaranteed to be aligned to at least 8 - // by adding 7, we're guaranteed to: - // - if size mod 8 == 0: NOT cross the divisible-by-8 boundary - // - if size mod 8 != 0: cross the divisible-by-8 boundary exactly once. - // so after adding 7, just shave off the lower bits under 7 - val roundedSize = (size + (alignment - 1u)) and (alignment - 1u).inv() - // equation for now: (size + 7u ) & 0xFFFFFFF8u + check(!isAlreadyOperating) { "Cannot call free from within the allocator" } + isAlreadyOperating = true + try { + val alignedSize = realAllocationSize(size) - val slotIndex = list.indexOfFirst { it.size >= roundedSize } - if (slotIndex == -1) - throw OutOfMemoryError("Out of linear memory. All available address space (2gb) is used.") + val slotIndex = list.indexOfFirst { it.size >= alignedSize } + if (slotIndex == -1) + throw OutOfMemoryError("Out of linear memory. All available address space (2gb) is used.") - val slot = list[slotIndex] + val slot = list[slotIndex] - if (slot.size == roundedSize) { - list.removeAt(slotIndex) - return slot - } else { - // in this case, split the slot into 2, return the left part, and reinsert the right part - val allocatedSlot = MemorySlot(slot.ptr, roundedSize) - val reinsertedSlot = MemorySlot(slot.ptr + roundedSize, slot.size - roundedSize) + if (slot.size == alignedSize) { + list.removeAt(slotIndex) + return slot + } else { + // in this case, split the slot into 2, return the left part, and reinsert the right part + val allocatedSlot = MemorySlot(slot.ptr, alignedSize) + val reinsertedSlot = MemorySlot(slot.ptr + alignedSize, slot.size - alignedSize) - list[slotIndex] = reinsertedSlot - return allocatedSlot + list[slotIndex] = reinsertedSlot + return allocatedSlot + } + } finally { + isAlreadyOperating = false } } } @@ -238,8 +282,6 @@ // TODO go back to UInt for this too check(size >= 0) { "size must be >= 0" } - // Pad available address to align it to 8 - // 8 is a max alignment number currently needed for Wasm component model canonical ABI val result = FreeList.allocate(size.toUInt()) check(result.ptr.address % 8u == 0u) { "result must be 8-byte aligned" } @@ -269,6 +311,8 @@ return createAllocatorInTheNewScope() } + // NOTE: we don't expose a free() function directly, to a) make it harder to write use-after-free's, and b) not expose MemorySlot / FreeList beyond this file + @PublishedApi internal fun destroy() { // TODO once we figure out the cabi realloc frees, also actually free this @@ -287,9 +331,10 @@ /** * WebAssembly Component Model Canonical ABI realloc implementation. - * This function is intended to be exported to a Component Model and must not be called directly. - * Memory allocated by this function must be freed - * by calling [freeAllComponentModelReallocAllocatedMemory] before calling any [withScopedMemoryAllocator]. + * This function is intended to be exported for Component Model support and must not be called directly. + * + * Memory allocated by this function must be freed by calling this function again, + * with the original pointer, the original size, and new size 0. */ @OptIn(UnsafeWasmMemoryApi::class) @ComponentModelInternalApi @@ -305,24 +350,92 @@ } val allocator = reallocAllocator!! - val newAllocation = allocator.allocate(newSize) - if (originalSize < newSize && originalPtr + originalSize == newAllocation.address.toInt()) { - // in that case, don't need to copy because we just grew at the same point + // to address the correct slot, we must extend the original size to be aligned, as that will be the internal size of the slot + val originalAllocationSize = realAllocationSize(originalSize.toUInt()) + + if (newSize == 0) { + + // TODO this is an easy way to get the program to trap, if it's misused. Any possible guardrails against this? + FreeList.free(MemorySlot(Pointer(originalPtr.toUInt()), originalAllocationSize)) + // TODO figure out a fitting return value here. + return 0 + } + + val newAllocationSize = realAllocationSize(newSize.toUInt()) + + // cases: + // 1. size doesn't change + // 2. allocation shrinks + // 3. allocation grows + // 3a. fresh allocation (original size was 0) + // NOTE: this would technically be handled by case 3b, but it's simpler to handle it separately + // 3b. allocation grows in place + // 3c. allocation grows elsewhere, needs copy + + if (newAllocationSize == originalAllocationSize) // case 1 + return originalPtr + + // case 2: shrinking, i.e., the new size is smaller than the old size: nothing to do except free a portion + if (newAllocationSize < originalAllocationSize) { + // NOTE: because we're only subtracting aligned sizes, the result will still be aligned + FreeList.free(MemorySlot(Pointer(originalPtr.toUInt() + newAllocationSize), originalAllocationSize - newAllocationSize)) return originalPtr } - // otherwise, we'll have to copy the old data, if the old size wasn't zero - if (originalSize > 0) - wasm_memory_copy(newAllocation.address.toInt(), originalPtr, minOf(originalSize, newSize)) + // case 3: growing, i.e., we need to do some actual allocation + val newAllocation = allocator.allocate(newSize) + + // case 3a: the original size was 0, we're done + if (originalSize == 0) + return newAllocation.address.toInt() + + // case 3b: we can grow the allocation in place + if (originalAllocationSize <= newAllocationSize && originalPtr.toUInt() + originalAllocationSize == newAllocation.address) { + // in that case, don't need to copy data from the old allocation because we just grew at the same point + // BUT: Because we grew, we're actually reusing the original allocation with its original aligned size. + // But at this moment, we just have one big allocation with size originalSizeAligned + newSizeAligned. + // So free the difference. + val startOfOverallocatedMemory = originalPtr.toUInt() + newAllocationSize + val overallocatedSize = originalAllocationSize // we allocated as if we didn't have the original allocation + FreeList.free(MemorySlot(Pointer(startOfOverallocatedMemory), overallocatedSize)) + // NOTE: allocating and then freeing again might seem overcomplicated; the obvious alternative would be to allocate twice in a row instead. The reason not to allocate twice is as follows: + // if the allocator ever changes, and stops giving out contiguous memory, this code path, as it stands, will simply stop being used, and nothing will break. + // While the allocation does occur contiguously, the free is also contiguous and doesn't perform any complex logic, because all that changes is the start address of the free list block that we're allocating from, the list itself is not modified. + // + // If we instead allocated twice, in case we can't grow the original allocation in place, we're implicitly relying on being able to perform the second allocation in place. This isn't always true, and would create further complications in these cases, by having to free the "failed" allocation first. Thus, overallocating plus freeing is safer than allocating incrementally. + // Conversely, this implementation suffers from sometimes not being able to grow an allocation in place, when the initial overallocatedSize is too large. However, this only results in an additional copy, instead of a semantics change. + + return originalPtr + } + + // case 3c: we now know the allocation grew (and the original allocation size was non-zero), and couldn't grow in place, so we have to copy the old data + // as this is only for useful bytes, we use the sizes that are given out to the application here, not the allocation sizes + wasm_memory_copy(newAllocation.address.toInt(), originalPtr, minOf(originalSize, newSize)) + // also free the old allocation (from which we copied), which is now useless + FreeList.free(MemorySlot(Pointer(originalPtr.toUInt()), originalAllocationSize)) return newAllocation.address.toInt() } /** - * Frees memory allocated by all previous calls of [componentModelRealloc]. + * Frees memory allocated by all previous calls of [componentModelRealloc]. + * + * NOTE: This function is incompatible with freeing memory manually through `componentModelRealloc(ptr, size, 0)` calls, as this will result in a double-free. + * TODO(REVIEW): Try to automatically handle these cases? Would make everything a bit uglier, but also reduce the chances of people running into double-frees. */ @OptIn(UnsafeWasmMemoryApi::class) +@Deprecated("Freeing all cabi_realloc-allocated memory is incompatible with the WASI preview 1 to preview 2 adapter which uses cabi_realloc allocated memory in a persistent way (never frees it, expects it to essentially have static storage duration). This means that the use of this function will always cause use-after-free UB in the adapter.") @ComponentModelInternalApi public fun freeAllComponentModelReallocAllocatedMemory() { - // empty for now, will be unused soon + // TODO need to reenable this before merging +// if (reallocAllocator != null) { +// reallocAllocator!!.destroy() +// reallocAllocator = null +// } +} + +// TODO(REVIEW) is this okay? is only for testing purposes +@OptIn(UnsafeWasmMemoryApi::class) +internal fun dumpFreeList(): String { + return FreeList.debugDump() }
diff --git a/libraries/stdlib/wasm/test/unsafe/ComponentModelReallocTest.kt b/libraries/stdlib/wasm/test/unsafe/ComponentModelReallocTest.kt index daa10cf..1bffad1 100644 --- a/libraries/stdlib/wasm/test/unsafe/ComponentModelReallocTest.kt +++ b/libraries/stdlib/wasm/test/unsafe/ComponentModelReallocTest.kt
@@ -5,20 +5,55 @@ @OptIn(UnsafeWasmMemoryApi::class, ComponentModelInternalApi::class) class ReallocTest { + var freelist: String = "" + + @BeforeTest + fun saveFreelist() { + freelist = dumpFreeList() + } + + @AfterTest + fun compareFreelist() { + assertEquals(freelist, dumpFreeList()) + } + + @Test + fun reallocFreeAllTest(){ + componentModelRealloc(0, 0, 100) + @Suppress("DEPRECATION") + freeAllComponentModelReallocAllocatedMemory() + } + @Test fun freshReallocTest() { val address1 = componentModelRealloc(0, 0, 10) val address2 = componentModelRealloc(0, 0, 10) - freeAllComponentModelReallocAllocatedMemory() assertNotEquals(address1, address2) + // free it, FIFO order + componentModelRealloc(address1, 10, 0) + componentModelRealloc(address2, 10, 0) + val address3 = componentModelRealloc(0, 0, 10) val address4 = componentModelRealloc(0, 0, 10) - freeAllComponentModelReallocAllocatedMemory() // After freeing memory, new reallocs should reuse the old memory assertEquals(address1, address3) - assertEquals(address4, address4) + assertEquals(address2, address4) + + // now free it in LIFO order + componentModelRealloc(address4, 10, 0) + componentModelRealloc(address3, 10, 0) + + // And again check that new allocations reuse the freed memory + val address5 = componentModelRealloc(0, 0, 10) + val address6 = componentModelRealloc(0, 0, 10) + + assertEquals(address1, address5) + assertEquals(address2, address6) + + componentModelRealloc(address5, 10, 0) + componentModelRealloc(address6, 10, 0) } @Test @@ -32,13 +67,14 @@ originalSize = (i + 1) * allocationStepSize, newSize = (i + 2) * allocationStepSize ) - assertEquals(newAddress1, address1) + assertEquals(address1, newAddress1) } val address2 = componentModelRealloc(0, 0, 10) assertTrue(address2 - address1 >= allocationStepSize * numReallocs) - freeAllComponentModelReallocAllocatedMemory() + componentModelRealloc(address1, numReallocs * allocationStepSize, 0) + componentModelRealloc(address2, 10, 0) } private fun writeNBytes(address: Int, n: Int, value: Byte) { @@ -62,7 +98,8 @@ val addrToClean = componentModelRealloc(0, 0, sizeToClean) writeNBytes(addrToClean, sizeToClean, 0.toByte()) assertBytesEquals(addrToClean, sizeToClean, 0.toByte()) - freeAllComponentModelReallocAllocatedMemory() + // free + componentModelRealloc(addrToClean, sizeToClean, 0) val address1 = componentModelRealloc(0, 0, bufferSize) writeNBytes(address1, bufferSize, 1.toByte()) @@ -82,7 +119,9 @@ assertTrue(address2new > address1new) assertBytesEquals(address2new, bufferSize, 2.toByte()) - freeAllComponentModelReallocAllocatedMemory() + // free + componentModelRealloc(address1new, bufferSize * 2, 0) + componentModelRealloc(address2new, bufferSize * 2, 0) } @Test @@ -93,22 +132,83 @@ assertTrue(reallocAddr.toUInt() > scopedAddr.address) assertTrue((reallocAddr.toUInt() - scopedAddr.address) >= 10u) - freeAllComponentModelReallocAllocatedMemory() + componentModelRealloc(reallocAddr, 10, 0) val scopedAddr2 = allocator.allocate(10) assertEquals(scopedAddr2.address.toInt(), reallocAddr) } } + @Test - fun creatingAllocatorsBeforeReallocIsFreedTest() { - componentModelRealloc(0, 0, 10) - assertFailsWith<IllegalStateException> { - withScopedMemoryAllocator { allocator -> - allocator.allocate(10) - } - } - freeAllComponentModelReallocAllocatedMemory() + fun reallocShrinkingDoesPartialFreeCorrectlyTest() { + // allocated a big chunk, then shrunk the allocation: should free the excess memory, and thus place the next allocation in there + + // so to get a benchmark, first allocate some memory normally, and do another allocation, this one will have to be reproduced exactly by the realloc shrink after + val v = componentModelRealloc(0, 0, 50) + val correct = componentModelRealloc(0, 0, 20) + + // now free these again, so we're back to the start state + componentModelRealloc(correct, 20, 0) + componentModelRealloc(v, 50, 0) + + // allocate and shrink + val overallocatedAddress = componentModelRealloc(0, 0, 100) + val shrunkenAddress = componentModelRealloc(overallocatedAddress, 100, 50) + assertEquals(overallocatedAddress, shrunkenAddress) + + // allocate again: when we allocate now, we should get the same address as in "correct" before + val actual = componentModelRealloc(0, 0, 20) + assertEquals(correct, actual) + + // free stuff + componentModelRealloc(actual, 20, 0) + componentModelRealloc(shrunkenAddress, 50, 0) } -} \ No newline at end of file + @Test + fun reallocGrowingInPlaceFreesRestTest() { + // get a benchmark: allocate 1000 bytes normally, then allocate 20. Those 20 should start at the same address as if we allocate 50, then grow in place by 50, then allocate 20 again + // NOTE: the allocation is this huge, to attempt to guarantee we get a consecutive one (probably at the end of everything) + val v = componentModelRealloc(0, 0, 1000) + val correct = componentModelRealloc(0, 0, 20) + // free these again to get back to the starting state + componentModelRealloc(correct, 20, 0) + componentModelRealloc(v, 1000, 0) + + val underallocatedAddress = componentModelRealloc(0, 0, 500) + val extendedAllocationAddress = componentModelRealloc(underallocatedAddress, 500, 1000) + assertEquals(underallocatedAddress, extendedAllocationAddress) + + // allocate again, we should be at the same state as if we had allocated 1000 immediately + val actual = componentModelRealloc(0, 0, 20) + assertEquals(correct, actual) + + // free + componentModelRealloc(extendedAllocationAddress, 1000, 0) + componentModelRealloc(actual, 20, 0) + } + + @Test + fun reallocGrowingWithCopyFreesOldTest() { + val ungrowableMemoryAddress = componentModelRealloc(0, 0, 10) + // now we'll do a "filler" allocation in the middle, just to make it impossible to grow the first one + val filler = componentModelRealloc(0, 0, 10) + + writeNBytes(ungrowableMemoryAddress, 10, 42.toByte()) + // now try to grow the ungrowable memory (which must result in a copy) + val newAddress = componentModelRealloc(ungrowableMemoryAddress, 10, 20) + // check the copy itself + assertNotEquals(ungrowableMemoryAddress, newAddress) + assertBytesEquals(newAddress, 10, 42.toByte()) + + // NOTE: implementation detail: for a fresh allocation we should go through the free list from the start, and thus get the same address again as the original ungrowableMemoryAddress, as that must be freed by the growing allocation above + val actual = componentModelRealloc(0, 0, 10) + assertEquals(ungrowableMemoryAddress, actual) + + // free + componentModelRealloc(newAddress, 20, 0) + componentModelRealloc(actual, 10, 0) + componentModelRealloc(filler, 10, 0) + } +}