Buffer pool

common/src/main/java/com/pedro/common/BufferPool.kt

@@ -0,0 +1,104 @@
+/*
+ * Copyright (C) 2024 pedroSG94.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.pedro.common
+
+/**
+ * Created by pedro on 29/7/26.
+ *
+ * Pool of reusable ByteArray to avoid allocating one per encoded frame.
+ *
+ * An encoded frame must be copied out of the MediaCodec output buffer before giving it back to
+ * the codec, so the copy is unavoidable, but the allocation is not. At 6Mbps a stream allocates
+ * around 750KB/s of short lived arrays, and all of it ends in the GC.
+ *
+ * Buffers are grouped in power of two size classes so a frame gets the smallest available array
+ * able to contain it. Arrays are handed out dirty (never zeroed) because the caller always
+ * writes the whole frame and reads it back limited to the frame size.
+ *
+ * Thread safe: acquired from the encoder thread and released from the sender thread.
+ */
+class BufferPool(
+  private val maxBuffersPerSizeClass: Int = 8,
+  private val maxRetainedBytes: Long = 16 * 1024 * 1024
+) {
+
+  companion object {
+    //1KB, audio frames are smaller but a smaller class is not worth the fragmentation
+    private const val MIN_SIZE_CLASS = 10
+    //4MB, bigger frames are allocated and discarded as before
+    private const val MAX_SIZE_CLASS = 22
+  }
+
+  private val buckets = arrayOfNulls<ArrayDeque<ByteArray>>(MAX_SIZE_CLASS - MIN_SIZE_CLASS + 1)
+  private val lock = Any()
+  private var retainedBytes = 0L
+
+  /**
+   * Return an array of at least [minSize] bytes. The content is undefined.
+   */
+  fun acquire(minSize: Int): ByteArray {
+    val sizeClass = sizeClassOf(minSize)
+    if (sizeClass > MAX_SIZE_CLASS) return ByteArray(minSize)
+    synchronized(lock) {
+      val buffer = buckets[sizeClass - MIN_SIZE_CLASS]?.removeLastOrNull()
+      if (buffer != null) {
+        retainedBytes -= buffer.size
+        return buffer
+      }
+    }
+    return ByteArray(1 shl sizeClass)
+  }
+
+  fun release(buffer: java.nio.ByteBuffer) {
+    if (buffer.hasArray()) release(buffer.array())
+  }
+
+  /**
+   * Give an array back to the pool. Arrays not created by [acquire], and arrays already released,
+   * are ignored, so the same array can never be handed out to two users at once.
+   */
+  fun release(buffer: ByteArray) {
+    val size = buffer.size
+    if (Integer.bitCount(size) != 1) return
+    val sizeClass = 31 - Integer.numberOfLeadingZeros(size)
+    if (sizeClass !in MIN_SIZE_CLASS..MAX_SIZE_CLASS) return
+    synchronized(lock) {
+      if (retainedBytes + size > maxRetainedBytes) return
+      val bucket = buckets[sizeClass - MIN_SIZE_CLASS]
+        ?: ArrayDeque<ByteArray>().also { buckets[sizeClass - MIN_SIZE_CLASS] = it }
+      if (bucket.size >= maxBuffersPerSizeClass) return
+      //a double release would put the same array twice in the bucket and corrupt one of its users
+      if (bucket.any { it === buffer }) return
+      bucket.addLast(buffer)
+      retainedBytes += size
+    }
+  }
+
+  fun clear() {
+    synchronized(lock) {
+      buckets.fill(null)
+      retainedBytes = 0
+    }
+  }
+
+  fun getRetainedBytes(): Long = synchronized(lock) { retainedBytes }
+
+  private fun sizeClassOf(size: Int): Int {
+    val value = maxOf(size, 1 shl MIN_SIZE_CLASS)
+    return 32 - Integer.numberOfLeadingZeros(value - 1)
+  }
+}

common/src/main/java/com/pedro/common/Extensions.kt

@@ -357,4 +357,12 @@ fun ByteBuffer.removeHeader(): ByteBuffer {
 
 fun ByteArray.writeUInt32(offset: Int, value: Int) {
   value.toUInt32().copyInto(this, offset)
-}
+}
+
+fun ByteBuffer.clone(data: ByteArray): ByteBuffer {
+  val length = limit()
+  val source = duplicate()
+  source.position(0)
+  source.get(data, 0, length)
+  return ByteBuffer.wrap(data, 0, length).slice()
+}

common/src/main/java/com/pedro/common/StreamBlockingQueue.kt

@@ -46,9 +46,19 @@ class StreamBlockingQueue(var capacity: Int) {
         cacheQueue.drainTo(destiny.cacheQueue)
     }
 
-    fun clear() {
-        queue.clear()
-        cacheQueue.clear()
+    /**
+     * @param onRemove called for every discarded frame, used to recycle pooled buffers.
+     */
+    fun clear(onRemove: ((MediaFrame) -> Unit)? = null) {
+        if (onRemove != null) {
+            val removed = mutableListOf<MediaFrame>()
+            queue.drainTo(removed)
+            cacheQueue.drainTo(removed)
+            removed.forEach(onRemove)
+        } else {
+            queue.clear()
+            cacheQueue.clear()
+        }
         startTs = 0L
         cacheTimeFilled.set(false)
     }

common/src/main/java/com/pedro/common/base/BaseSender.kt

@@ -2,8 +2,10 @@ package com.pedro.common.base
 
 import android.util.Log
 import com.pedro.common.BitrateManager
+import com.pedro.common.BufferPool
 import com.pedro.common.ConnectChecker
 import com.pedro.common.StreamBlockingQueue
+import com.pedro.common.clone
 import com.pedro.common.frame.MediaFrame
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
@@ -13,6 +15,7 @@ import kotlinx.coroutines.cancelAndJoin
 import kotlinx.coroutines.delay
 import kotlinx.coroutines.isActive
 import kotlinx.coroutines.launch
+import kotlinx.coroutines.runInterruptible
 import java.nio.ByteBuffer
 import java.util.concurrent.atomic.AtomicLong
 
@@ -24,7 +27,8 @@ abstract class BaseSender(
     @Volatile
     protected var running = false
 
-    protected val queue = StreamBlockingQueue(400)
+    private val queue = StreamBlockingQueue(400)
+    private val bufferPool = BufferPool()
 
     protected val audioFramesSent = AtomicLong(0)
     protected val videoFramesSent = AtomicLong(0)
@@ -44,9 +48,13 @@ abstract class BaseSender(
     protected abstract suspend fun onRun()
     protected abstract suspend fun stopImp(clear: Boolean = true)
 
-    fun sendMediaFrame(mediaFrame: MediaFrame) {
-        if (running && !queue.trySend(mediaFrame)) {
-            when (mediaFrame.type) {
+    fun sendMediaFrame(buffer: ByteBuffer, info: MediaFrame.Info, type: MediaFrame.Type) {
+        if (!running) return
+        val data = bufferPool.acquire(buffer.limit())
+        val mediaFrame = MediaFrame(buffer.clone(data), info, type)
+        if (!queue.trySend(mediaFrame)) {
+            bufferPool.release(mediaFrame.data)
+            when (type) {
                 MediaFrame.Type.VIDEO -> {
                     Log.i(TAG, "Video frame discarded")
                     droppedVideoFrames.incrementAndGet()
@@ -59,9 +67,23 @@ abstract class BaseSender(
         }
     }
 
+    /**
+     * Take the next frame from the queue, hand it to [consume] and recycle its buffer once
+     * consumed. Senders must read frames only through this method, the buffer is reused right
+     * after [consume] returns.
+     */
+    protected suspend fun consumeFrame(consume: suspend (MediaFrame) -> Unit) {
+        val mediaFrame = runInterruptible { queue.take() }
+        try {
+            consume(mediaFrame)
+        } finally {
+            bufferPool.release(mediaFrame.data)
+        }
+    }
+
     fun start() {
         bitrateManager.reset()
-        queue.clear()
+        queue.clear { bufferPool.release(it.data) }
         running = true
         job = scope.launch {
             val bitrateTask = async {
@@ -86,7 +108,8 @@ abstract class BaseSender(
         resetBytesSend()
         job?.cancelAndJoin()
         job = null
-        queue.clear()
+        queue.clear { bufferPool.release(it.data) }
+        bufferPool.clear()
     }
 
     @Throws(IllegalArgumentException::class)
@@ -110,7 +133,7 @@ abstract class BaseSender(
     fun getItemsInCache(): Int = queue.getSize()
 
     fun clearCache() {
-        queue.clear()
+        queue.clear { bufferPool.release(it.data) }
     }
 
     fun getSentAudioFrames(): Long = audioFramesSent.get()

common/src/test/java/com/pedro/common/BufferPoolTest.kt

@@ -0,0 +1,116 @@
+package com.pedro.common
+
+import org.junit.Assert.assertArrayEquals
+import org.junit.Assert.assertEquals
+import org.junit.Assert.assertNotSame
+import org.junit.Assert.assertSame
+import org.junit.Assert.assertTrue
+import org.junit.Test
+
+class BufferPoolTest {
+
+  @Test
+  fun `GIVEN a size WHEN acquire THEN return a power of two array able to contain it`() {
+    val pool = BufferPool()
+    val sizes = listOf(0, 1, 500, 1024, 1025, 4096, 100000)
+    sizes.forEach { size ->
+      val buffer = pool.acquire(size)
+      assertTrue("$size does not fit in ${buffer.size}", buffer.size >= size)
+      assertEquals("${buffer.size} is not a power of two", 1, Integer.bitCount(buffer.size))
+      assertTrue("${buffer.size} is smaller than the min size class", buffer.size >= 1024)
+    }
+  }
+
+  @Test
+  fun `GIVEN a released buffer WHEN acquire the same size class THEN reuse the same instance`() {
+    val pool = BufferPool()
+    val buffer = pool.acquire(3000)
+    assertEquals(4096, buffer.size)
+    pool.release(buffer)
+    //any size inside the same class must get the very same array back
+    assertSame(buffer, pool.acquire(2049))
+  }
+
+  @Test
+  fun `GIVEN an empty pool WHEN acquire twice without releasing THEN return different instances`() {
+    val pool = BufferPool()
+    val first = pool.acquire(1024)
+    val second = pool.acquire(1024)
+    assertNotSame(first, second)
+  }
+
+  @Test
+  fun `GIVEN a foreign array WHEN release THEN ignore it`() {
+    val pool = BufferPool()
+    //not a power of two so it was never created by acquire
+    val foreign = ByteArray(3000)
+    pool.release(foreign)
+    assertEquals(0, pool.getRetainedBytes())
+    assertNotSame(foreign, pool.acquire(3000))
+  }
+
+  @Test
+  fun `GIVEN a size over the max class WHEN acquire and release THEN allocate exact and do not retain`() {
+    val pool = BufferPool()
+    val size = 8 * 1024 * 1024
+    val buffer = pool.acquire(size)
+    assertEquals(size, buffer.size)
+    pool.release(buffer)
+    assertEquals(0, pool.getRetainedBytes())
+  }
+
+  @Test
+  fun `GIVEN more buffers than the limit WHEN release THEN retain only the allowed ones`() {
+    val pool = BufferPool(maxBuffersPerSizeClass = 2)
+    val buffers = List(5) { pool.acquire(1024) }
+    buffers.forEach { pool.release(it) }
+    assertEquals(2048, pool.getRetainedBytes())
+  }
+
+  @Test
+  fun `GIVEN a retained bytes limit WHEN release THEN stop retaining once reached`() {
+    val pool = BufferPool(maxRetainedBytes = 4096)
+    val buffers = List(8) { pool.acquire(1024) }
+    buffers.forEach { pool.release(it) }
+    assertEquals(4096, pool.getRetainedBytes())
+  }
+
+  @Test
+  fun `GIVEN the same buffer WHEN release twice THEN retain it only once`() {
+    val pool = BufferPool()
+    val buffer = pool.acquire(1024)
+    pool.release(buffer)
+    pool.release(buffer)
+    assertEquals(1024, pool.getRetainedBytes())
+    //the second acquire must not get the very same array than the first one
+    assertSame(buffer, pool.acquire(1024))
+    assertNotSame(buffer, pool.acquire(1024))
+  }
+
+  @Test
+  fun `GIVEN retained buffers WHEN clear THEN drop all of them`() {
+    val pool = BufferPool()
+    val buffer = pool.acquire(1024)
+    pool.release(buffer)
+    assertEquals(1024, pool.getRetainedBytes())
+    pool.clear()
+    assertEquals(0, pool.getRetainedBytes())
+    assertNotSame(buffer, pool.acquire(1024))
+  }
+
+  @Test
+  fun `GIVEN a dirty reused buffer WHEN write a shorter frame THEN only the frame bytes are read`() {
+    val pool = BufferPool()
+    val first = pool.acquire(1024)
+    //fill it with garbage the next user must not see
+    first.fill(0xFF.toByte())
+    pool.release(first)
+
+    val frame = byteArrayOf(1, 2, 3, 4, 5)
+    val reused = pool.acquire(frame.size)
+    assertSame(first, reused)
+    frame.copyInto(reused)
+    //a consumer limited to the frame size never observes the stale bytes
+    assertArrayEquals(frame, reused.copyOfRange(0, frame.size))
+  }
+}