Pooled ByteBuffers (Fast TreeSet+Double-linked-list version)

ByteBufferSlicePool.kt

package adb.io

import adb.WeakReferenceMap
import java.lang.ref.Cleaner
import java.nio.ByteBuffer
import java.util.*

class ByteBufferSlicePool2(
    capacity: Int,
) : ByteBufferAllocator {
    private val primaryBuffer: ByteBuffer = ByteBuffer.allocateDirect(capacity)
    private val free = TreeSet<Segment>(compareBy { it.size }).apply { add(Segment(0, capacity)) }
    private val used = WeakReferenceMap<ByteBuffer, Segment>()
    private val cleaner = Cleaner.create()
    
    fun available() = synchronized(primaryBuffer) { free.maxBy { it.size }.size }
    fun remaining() = synchronized(primaryBuffer) { free.sumOf { it.size } }

    private fun Segment.splitFromStart(length: Int): Segment {
        val prefix = copy(size = length, free = false)
        start += length
        size -= length

        prev?.next = prefix
        prev = prefix
        prefix.next = if (size > 0) this else next

        return prefix
    }

    override fun alloc(size: Int): ByteBuffer = synchronized(primaryBuffer) {
        val source: Segment = free.tailSet(Segment(0, size)).firstOrNull()
            ?: error("No free regions with at least $size bytes available! $free")

        val slice = primaryBuffer.slice(source.start, size)

        val alloc = source.splitFromStart(size)

        free.remove(source)
        if (source.size > 0) {
            free.add(source)
        }

        // For normal free() cleanup
        used[slice] = alloc

        // For gc() cleanup
        cleaner.register(slice) {
            printLeakWarning(alloc)
            free(alloc)
        }

        return slice
    }

    override fun free(buffer: ByteBuffer): Unit = synchronized(primaryBuffer) {
        free(used.remove(buffer) ?: error("Buffer to free was not found in an alloc region!"))
    }

    private fun free(alloc: Segment): Unit = synchronized(primaryBuffer) {
        require(!alloc.free) { "Cannot free an unallocated segment!" }

        // Check for and merge with adjacent Free regions
        alloc.prev?.takeIf { it.free }?.also {
            alloc.prev = it.prev
            alloc.start = it.start
            alloc.size += it.size
            free.remove(it)
        }

        alloc.next?.takeIf { it.free }?.also {
            alloc.next = it.next
            alloc.size += it.size
            free.remove(it)
        }

        alloc.free = true
        free.add(alloc)
    }

    private fun printLeakWarning(alloc: Segment) {
        println("A byte buffer was not freed! ($alloc)")
    }

    private data class Segment(
        var start: Int,
        var size: Int,
        var prev: Segment? = null,
        var next: Segment? = null,
        var free: Boolean = true,
    ) {
        override fun hashCode() = start * 31
        override fun equals(other: Any?) = (other as? Segment)?.start == start
        override fun toString() = "Segment(start=$start, size=$size)"
    }
}