[K/N][interop] Made NativeMemoryAllocator lock-free

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, raw chunk ref,...padding...,diff to start,aligned data start,.....data.....]
-
+    fun alloc(size: Long, align: Int): Long {
-    // Chunk layout: [chunk size,...padding...,diff to start,aligned data start,.....data.....]
-    private fun allocNoLock(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)
+        val chunkWithSkippedHeader = chunkStart + ChunkHeaderSize
-        unsafe.putInt(alignedPtr - Int.SIZE_BYTES, (alignedPtr - ptr).toInt())
+        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 {
+    fun free(mem: Long) {
-        val paddedSize = alignUp(size, align)
+        val chunkStart = mem - unsafe.getInt(mem - Int.SIZE_BYTES)
-        val index = paddedSize / align - 1
+        val chunkSize = getChunkSize(chunkStart)
-        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<Long>()
-    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)
         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)
+        for (index in rawChunks.indices) {
-            unsafe.freeMemory(chunk)
+            val rawChunk = rawChunks[index]
-        rawChunks.clear()
+            if (rawChunk != 0L)
-    }
+                unsafe.freeMemory(rawChunk)
-
+            rawChunks[index] = 0L
-    private val unsafe = with(Unsafe::class.java.getDeclaredField("theUnsafe")) {
+        }
-        isAccessible = true
+        rawOffset = 0L
-        return@with this.get(null) as Unsafe
     }
 }