diff --git a/kotlin-native/utilities/basic-utils/src/main/kotlin/NativeMemoryAllocator.kt b/kotlin-native/utilities/basic-utils/src/main/kotlin/NativeMemoryAllocator.kt index 29cd94b0552..0d18894fd58 100644 --- a/kotlin-native/utilities/basic-utils/src/main/kotlin/NativeMemoryAllocator.kt +++ b/kotlin-native/utilities/basic-utils/src/main/kotlin/NativeMemoryAllocator.kt @@ -63,9 +63,14 @@ private const val BigChunksSizeAlignment = 4096 private const val MaxSmallSize = ChunkBucketSize * SmallChunksSizeAlignment private const val MaxMediumSize = ChunkBucketSize * MediumChunksSizeAlignment private const val MaxBigSize = ChunkBucketSize * BigChunksSizeAlignment -private const val ChunkHeaderSize = 2 * Int.SIZE_BYTES // chunk size + alignment hop size. +private const val ChunkHeaderSize = 3 * Int.SIZE_BYTES // chunk size + raw chunk ref + alignment hop size. -private const val RawChunkSize: Long = 4L * 1024 * 1024 +private const val RawChunkSizeBits = 22 // 4MB +private const val RawChunkSize = 1L shl RawChunkSizeBits +private const val ChunkSizeAlignmentBits = 3 // All chunk sizes are aligned to at least 8. +private const val RawChunkOffsetBits = RawChunkSizeBits - ChunkSizeAlignmentBits +private const val MinChunkSize = 8 +private const val MaxRawChunksCount = 1024 // 4GB in total. class NativeMemoryAllocator { companion object { @@ -76,70 +81,76 @@ class NativeMemoryAllocator { private fun alignUp(x: Long, align: Int) = (x + align - 1) and (align - 1).toLong().inv() private fun alignUp(x: Int, align: Int) = (x + align - 1) and (align - 1).inv() + private val unsafe = with(Unsafe::class.java.getDeclaredField("theUnsafe")) { + isAccessible = true + return@with this.get(null) as Unsafe + } + + private val longArrayBaseOffset = unsafe.arrayBaseOffset(LongArray::class.java).toLong() + + private val rawOffsetFieldOffset = unsafe.objectFieldOffset(this::class.java.getDeclaredField("rawOffset")) + + @JvmInline + private value class ChunkRef(val value: Int) { + val offset get() = (value and ((1 shl RawChunkOffsetBits) - 1)) shl ChunkSizeAlignmentBits + + val index get() = (value ushr RawChunkOffsetBits) - 1 + + companion object { + init { + // Ensure that pair (index, offset) fits in 32-bit integer. + check(MaxRawChunksCount < 1L shl (32 - RawChunkOffsetBits)) + } + + fun encode(index: Int, offset: Int) = ChunkRef(((index + 1) shl RawChunkOffsetBits) or (offset ushr ChunkSizeAlignmentBits)) + + val Invalid = ChunkRef(0) + } + } + + // Timestamps here solve the ABA problem. + @JvmInline + private value class ChunkRefWithTimestamp(val value: Long) { + val chunkRef get() = ChunkRef(value.toInt()) + + val timestamp get() = (value ushr 32).toInt() + + companion object { + fun encode(chunkRef: Int, timestamp: Int) = ChunkRefWithTimestamp(chunkRef.toLong() or (timestamp.toLong() shl 32)) + } + } + + private fun getChunkSize(chunk: Long) = unsafe.getInt(chunk) + private fun setChunkSize(chunk: Long, size: Int) = unsafe.putInt(chunk, size) + + private fun getChunkRef(chunk: Long): ChunkRef = ChunkRef(unsafe.getInt(chunk + Int.SIZE_BYTES /* skip chunk size */)) + private fun setChunkRef(chunk: Long, chunkRef: ChunkRef) = unsafe.putInt(chunk + Int.SIZE_BYTES /* skip chunk size */, chunkRef.value) + private val smallChunks = LongArray(ChunkBucketSize) private val mediumChunks = LongArray(ChunkBucketSize) private val bigChunks = LongArray(ChunkBucketSize) - private val lock = Any() - fun alloc(size: Long, align: Int) = synchronized(lock) { allocNoLock(size, align) } - - // Chunk layout: [chunk size,...padding...,diff to start,aligned data start,.....data.....] - private fun allocNoLock(size: Long, align: Int): Long { + // Chunk layout: [chunk size, raw chunk ref,...padding...,diff to start,aligned data start,.....data.....] + fun alloc(size: Long, align: Int): Long { val totalChunkSize = ChunkHeaderSize + size + align - val ptr = ChunkHeaderSize + when { + val chunkStart = when { totalChunkSize <= MaxSmallSize -> allocFromFreeList(totalChunkSize.toInt(), SmallChunksSizeAlignment, smallChunks) totalChunkSize <= MaxMediumSize -> allocFromFreeList(totalChunkSize.toInt(), MediumChunksSizeAlignment, mediumChunks) totalChunkSize <= MaxBigSize -> allocFromFreeList(totalChunkSize.toInt(), BigChunksSizeAlignment, bigChunks) else -> unsafe.allocateMemory(totalChunkSize).also { // The actual size is not used. Just put value bigger than the biggest threshold. - unsafe.putInt(it, Int.MAX_VALUE) + setChunkSize(it, Int.MAX_VALUE) } } - val alignedPtr = alignUp(ptr, align) - unsafe.putInt(alignedPtr - Int.SIZE_BYTES, (alignedPtr - ptr).toInt()) + val chunkWithSkippedHeader = chunkStart + ChunkHeaderSize + val alignedPtr = alignUp(chunkWithSkippedHeader, align) + unsafe.putInt(alignedPtr - Int.SIZE_BYTES, (alignedPtr - chunkStart).toInt()) return alignedPtr } - private fun allocFromFreeList(size: Int, align: Int, freeList: LongArray): Long { - val paddedSize = alignUp(size, align) - val index = paddedSize / align - 1 - val chunk = freeList[index] - val ptr = if (chunk == 0L) - allocRaw(paddedSize) - else { - val nextChunk = unsafe.getLong(chunk) - freeList[index] = nextChunk - chunk - } - unsafe.putInt(ptr, paddedSize) - return ptr - } - - private fun freeToFreeList(paddedSize: Int, align: Int, freeList: LongArray, chunk: Long) { - require(paddedSize > 0 && paddedSize % align == 0) - val index = paddedSize / align - 1 - unsafe.putLong(chunk, freeList[index]) - freeList[index] = chunk - } - - private val rawChunks = mutableListOf() - private var rawOffset = 0 - - private fun allocRaw(size: Int): Long { - if (rawChunks.isEmpty() || rawOffset + size > RawChunkSize) { - val newRawChunk = unsafe.allocateMemory(RawChunkSize) - rawChunks.add(newRawChunk) - rawOffset = size - return newRawChunk - } - return (rawChunks.last() + rawOffset).also { rawOffset += size } - } - - fun free(mem: Long) = synchronized(lock) { freeNoLock(mem) } - - private fun freeNoLock(mem: Long) { - val chunkStart = mem - ChunkHeaderSize - unsafe.getInt(mem - Int.SIZE_BYTES) - val chunkSize = unsafe.getInt(chunkStart) + fun free(mem: Long) { + val chunkStart = mem - unsafe.getInt(mem - Int.SIZE_BYTES) + val chunkSize = getChunkSize(chunkStart) when { chunkSize <= MaxSmallSize -> freeToFreeList(chunkSize, SmallChunksSizeAlignment, smallChunks, chunkStart) chunkSize <= MaxMediumSize -> freeToFreeList(chunkSize, MediumChunksSizeAlignment, mediumChunks, chunkStart) @@ -148,19 +159,90 @@ class NativeMemoryAllocator { } } + private val rawChunks = LongArray(MaxRawChunksCount) + + private fun ChunkRef.dereference() = rawChunks[index] + offset + + private val rawChunksLock = Any() + + @Volatile + private var rawOffset = 0L + + private fun allocRaw(size: Int): Long { + require(size % MinChunkSize == 0) { "Sizes should be multiples of $MinChunkSize" } + while (true) { + val offset = rawOffset + val remainedInCurChunk = (alignUp(offset + 1, RawChunkSize.toInt()) - offset).toInt() + val newOffset = offset + if (remainedInCurChunk >= size) size else remainedInCurChunk + size + if (!unsafe.compareAndSwapLong(this, rawOffsetFieldOffset, offset, newOffset)) continue + val dataStartOffset = newOffset - size + val rawChunkIndex = (dataStartOffset ushr RawChunkSizeBits).toInt() + val rawChunkOffset = (dataStartOffset and (RawChunkSize - 1)).toInt() + var rawChunk = rawChunks[rawChunkIndex] + if (rawChunk == 0L) { + synchronized(rawChunksLock) { + rawChunk = rawChunks[rawChunkIndex] + if (rawChunk == 0L) { + rawChunk = unsafe.allocateMemory(RawChunkSize) + rawChunks[rawChunkIndex] = rawChunk + } + } + } + val ptr = rawChunk + rawChunkOffset + setChunkRef(ptr, ChunkRef.encode(rawChunkIndex, rawChunkOffset)) + return ptr + } + } + + private fun allocFromFreeList(size: Int, align: Int, freeList: LongArray): Long { + val paddedSize = alignUp(size, align) + val index = paddedSize / align - 1 + val ptr: Long + while (true) { + val chunkRefWithTimestamp = ChunkRefWithTimestamp(freeList[index]) + val chunkRef = chunkRefWithTimestamp.chunkRef + if (chunkRef == ChunkRef.Invalid) { + ptr = allocRaw(paddedSize) + break + } else { + val chunk = chunkRef.dereference() + val nextChunkRef = unsafe.getInt(chunk) + val nextChunkRefWithTimestamp = ChunkRefWithTimestamp.encode(nextChunkRef, chunkRefWithTimestamp.timestamp + 1) + if (unsafe.compareAndSwapLong(freeList, longArrayBaseOffset + index * Long.SIZE_BYTES, + chunkRefWithTimestamp.value, nextChunkRefWithTimestamp.value)) { + ptr = chunk + break + } + } + } + setChunkSize(ptr, paddedSize) + return ptr + } + + private fun freeToFreeList(paddedSize: Int, align: Int, freeList: LongArray, chunk: Long) { + require(paddedSize > 0 && paddedSize % align == 0) + val index = paddedSize / align - 1 + do { + val nextChunkRefWithTimestamp = ChunkRefWithTimestamp(freeList[index]) + unsafe.putInt(chunk, nextChunkRefWithTimestamp.chunkRef.value) + val chunkRef = getChunkRef(chunk) + val chunkRefWithTimestamp = ChunkRefWithTimestamp.encode(chunkRef.value, nextChunkRefWithTimestamp.timestamp + 1) + } while (!unsafe.compareAndSwapLong(freeList, longArrayBaseOffset + index * Long.SIZE_BYTES, + nextChunkRefWithTimestamp.value, chunkRefWithTimestamp.value)) + } + internal fun freeAll() { for (i in 0 until ChunkBucketSize) { smallChunks[i] = 0L mediumChunks[i] = 0L bigChunks[i] = 0L } - for (chunk in rawChunks) - unsafe.freeMemory(chunk) - rawChunks.clear() - } - - private val unsafe = with(Unsafe::class.java.getDeclaredField("theUnsafe")) { - isAccessible = true - return@with this.get(null) as Unsafe + for (index in rawChunks.indices) { + val rawChunk = rawChunks[index] + if (rawChunk != 0L) + unsafe.freeMemory(rawChunk) + rawChunks[index] = 0L + } + rawOffset = 0L } } \ No newline at end of file