chore(core): ts: use byte buffer pool when allocating TS buffer

ThibaultBee · ThibaultBee · commit 3c4e83656770 · 2025-10-09T14:44:50.000+02:00
diff --git a/core/src/main/java/io/github/thibaultbee/streampack/core/elements/endpoints/composites/muxers/ts/TsMuxer.kt b/core/src/main/java/io/github/thibaultbee/streampack/core/elements/endpoints/composites/muxers/ts/TsMuxer.kt
@@ -33,11 +33,14 @@ import io.github.thibaultbee.streampack.core.elements.endpoints.composites.muxer
 import io.github.thibaultbee.streampack.core.elements.endpoints.composites.muxers.ts.utils.av.OpusControlHeader
 import io.github.thibaultbee.streampack.core.elements.utils.av.audio.aac.ADTSFrameWriter
 import io.github.thibaultbee.streampack.core.elements.utils.av.audio.aac.LATMFrameWriter
+import io.github.thibaultbee.streampack.core.elements.utils.pool.ByteBufferPool
 import java.nio.ByteBuffer
 import java.util.MissingFormatArgumentException
 import kotlin.random.Random
 
 class TsMuxer : IMuxerInternal {
+    private val byteBufferPool = ByteBufferPool(true)
+
     override val info by lazy { TSMuxerInfo }
     private val tsServices = mutableListOf<Service>()
     private val tsPes = mutableListOf<Pes>()
@@ -56,10 +59,10 @@ class TsMuxer : IMuxerInternal {
 
     private val tsId = Random.nextInt(Byte.MIN_VALUE.toInt(), Byte.MAX_VALUE.toInt()).toShort()
     private var pat = Pat(
-        listener, tsServices, tsId, packetCount = 0
+        byteBufferPool, listener, tsServices, tsId, packetCount = 0
     )
     private var sdt = Sdt(
-        listener, tsServices, tsId, packetCount = 0
+        byteBufferPool, listener, tsServices, tsId, packetCount = 0
     )
 
     override val streamConfigs: List<CodecConfig>
@@ -339,11 +342,12 @@ class TsMuxer : IMuxerInternal {
         service.pmt = service.pmt?.apply {
             versionNumber = (versionNumber + 1).toByte()
             streams = service.streams
-        } ?: Pmt(listener, service, service.streams, getNewPid())
+        } ?: Pmt(byteBufferPool, listener, service, service.streams, getNewPid())
 
         // Init PES
         newStreams.forEach {
             Pes(
+                byteBufferPool,
                 listener,
                 it,
                 service.pcrPid == it.pid,
@@ -419,10 +423,12 @@ class TsMuxer : IMuxerInternal {
         tsServices.forEach {
             removeStreams(it)
         }
+        byteBufferPool.clear()
     }
 
     override fun release() {
         tsServices.clear()
+        byteBufferPool.close()
     }
 
     /**
diff --git a/core/src/main/java/io/github/thibaultbee/streampack/core/elements/endpoints/composites/muxers/ts/packets/Pat.kt b/core/src/main/java/io/github/thibaultbee/streampack/core/elements/endpoints/composites/muxers/ts/packets/Pat.kt
@@ -18,16 +18,19 @@ package io.github.thibaultbee.streampack.core.elements.endpoints.composites.muxe
 import io.github.thibaultbee.streampack.core.elements.endpoints.composites.muxers.IMuxerInternal
 import io.github.thibaultbee.streampack.core.elements.endpoints.composites.muxers.ts.data.ITSElement
 import io.github.thibaultbee.streampack.core.elements.endpoints.composites.muxers.ts.data.Service
+import io.github.thibaultbee.streampack.core.elements.utils.pool.ByteBufferPool
 import java.nio.ByteBuffer
 import kotlin.experimental.or
 
 class Pat(
+    byteBufferPool: ByteBufferPool,
     listener: IMuxerInternal.IMuxerListener? = null,
     private val services: List<Service>,
     tsId: Short,
     versionNumber: Byte = 0,
     var packetCount: Int = 0,
 ) : Psi(
+    byteBufferPool,
     listener,
     PID,
     TID,
diff --git a/core/src/main/java/io/github/thibaultbee/streampack/core/elements/endpoints/composites/muxers/ts/packets/Pes.kt b/core/src/main/java/io/github/thibaultbee/streampack/core/elements/endpoints/composites/muxers/ts/packets/Pes.kt
@@ -23,12 +23,14 @@ import io.github.thibaultbee.streampack.core.elements.endpoints.composites.muxer
 import io.github.thibaultbee.streampack.core.elements.utils.TimeUtils
 import io.github.thibaultbee.streampack.core.elements.utils.extensions.isAudio
 import io.github.thibaultbee.streampack.core.elements.utils.extensions.isVideo
+import io.github.thibaultbee.streampack.core.elements.utils.pool.ByteBufferPool
 
 class Pes(
+    byteBufferPool: ByteBufferPool,
     muxerListener: IMuxerInternal.IMuxerListener? = null,
     val stream: Stream,
     private val hasPcr: Boolean,
-) : TS(muxerListener, stream.pid) {
+) : TS(byteBufferPool, muxerListener, stream.pid) {
     fun write(frame: Frame) {
         val programClockReference = if (hasPcr) {
             TimeUtils.currentTime()
@@ -48,7 +50,13 @@ class Pes(
             dts = frame.dtsInUs
         )
 
-        write(frame.buffer, adaptationField.toByteBuffer(), header.toByteBuffer(), true, frame.ptsInUs)
+        write(
+            frame.buffer,
+            adaptationField.toByteBuffer(),
+            header.toByteBuffer(),
+            true,
+            frame.ptsInUs
+        )
     }
 
     enum class StreamId(val value: Short) {
diff --git a/core/src/main/java/io/github/thibaultbee/streampack/core/elements/endpoints/composites/muxers/ts/packets/Pmt.kt b/core/src/main/java/io/github/thibaultbee/streampack/core/elements/endpoints/composites/muxers/ts/packets/Pmt.kt
@@ -24,15 +24,18 @@ import io.github.thibaultbee.streampack.core.elements.endpoints.composites.muxer
 import io.github.thibaultbee.streampack.core.elements.endpoints.composites.muxers.ts.data.Stream
 import io.github.thibaultbee.streampack.core.elements.utils.extensions.put
 import io.github.thibaultbee.streampack.core.elements.utils.extensions.putShort
+import io.github.thibaultbee.streampack.core.elements.utils.pool.ByteBufferPool
 import java.nio.ByteBuffer
 
 class Pmt(
+    byteBufferPool: ByteBufferPool,
     listener: IMuxerInternal.IMuxerListener? = null,
     private val service: Service,
     var streams: List<Stream>,
     pid: Short,
     versionNumber: Byte = 0,
 ) : Psi(
+    byteBufferPool,
     listener,
     pid,
     TID,
diff --git a/core/src/main/java/io/github/thibaultbee/streampack/core/elements/endpoints/composites/muxers/ts/packets/Psi.kt b/core/src/main/java/io/github/thibaultbee/streampack/core/elements/endpoints/composites/muxers/ts/packets/Psi.kt
@@ -15,12 +15,15 @@
  */
 package io.github.thibaultbee.streampack.core.elements.endpoints.composites.muxers.ts.packets
 
+import io.github.thibaultbee.streampack.core.elements.endpoints.composites.muxers.IMuxerInternal
 import io.github.thibaultbee.streampack.core.elements.endpoints.composites.muxers.ts.utils.CRC32
 import io.github.thibaultbee.streampack.core.elements.utils.extensions.put
+import io.github.thibaultbee.streampack.core.elements.utils.pool.ByteBufferPool
 import java.nio.ByteBuffer
 
 open class Psi(
-    listener: io.github.thibaultbee.streampack.core.elements.endpoints.composites.muxers.IMuxerInternal.IMuxerListener? = null,
+    byteBufferPool: ByteBufferPool,
+    listener: IMuxerInternal.IMuxerListener? = null,
     pid: Short,
     private val tableId: Byte,
     private val sectionSyntaxIndicator: Boolean = false,
@@ -29,7 +32,7 @@ open class Psi(
     var versionNumber: Byte = 0,
     private val sectionNumber: Byte = 0,
     private val lastSectionNumber: Byte = 0,
-) : TS(listener, pid) {
+) : TS(byteBufferPool, listener, pid) {
     companion object {
         const val CRC_SIZE = 4
         const val PSI_HEADER_SIZE = 9 // contains pointer_field
diff --git a/core/src/main/java/io/github/thibaultbee/streampack/core/elements/endpoints/composites/muxers/ts/packets/Sdt.kt b/core/src/main/java/io/github/thibaultbee/streampack/core/elements/endpoints/composites/muxers/ts/packets/Sdt.kt
@@ -15,22 +15,26 @@
  */
 package io.github.thibaultbee.streampack.core.elements.endpoints.composites.muxers.ts.packets
 
+import io.github.thibaultbee.streampack.core.elements.endpoints.composites.muxers.IMuxerInternal
 import io.github.thibaultbee.streampack.core.elements.endpoints.composites.muxers.ts.data.ITSElement
 import io.github.thibaultbee.streampack.core.elements.endpoints.composites.muxers.ts.data.Service
 import io.github.thibaultbee.streampack.core.elements.utils.extensions.put
 import io.github.thibaultbee.streampack.core.elements.utils.extensions.putShort
 import io.github.thibaultbee.streampack.core.elements.utils.extensions.putString
+import io.github.thibaultbee.streampack.core.elements.utils.pool.ByteBufferPool
 import java.nio.ByteBuffer
 
 
 class Sdt(
-    listener: io.github.thibaultbee.streampack.core.elements.endpoints.composites.muxers.IMuxerInternal.IMuxerListener? = null,
+    byteBufferPool: ByteBufferPool,
+    listener: IMuxerInternal.IMuxerListener? = null,
     private val services: List<Service>,
     tsId: Short,
     private val originalNetworkId: Short = 0xff01.toShort(),
     versionNumber: Byte = 0,
     var packetCount: Int = 0,
 ) : Psi(
+    byteBufferPool = byteBufferPool,
     listener,
     PID,
     TID,
diff --git a/core/src/main/java/io/github/thibaultbee/streampack/core/elements/endpoints/composites/muxers/ts/packets/TS.kt b/core/src/main/java/io/github/thibaultbee/streampack/core/elements/endpoints/composites/muxers/ts/packets/TS.kt
@@ -20,17 +20,19 @@ import io.github.thibaultbee.streampack.core.elements.endpoints.composites.muxer
 import io.github.thibaultbee.streampack.core.elements.endpoints.composites.muxers.ts.utils.MuxerConst
 import io.github.thibaultbee.streampack.core.elements.endpoints.composites.muxers.ts.utils.TSOutputCallback
 import io.github.thibaultbee.streampack.core.elements.utils.extensions.toInt
+import io.github.thibaultbee.streampack.core.elements.utils.pool.ByteBufferPool
 import java.nio.ByteBuffer
 import java.security.InvalidParameterException
 
 open class TS(
+     byteBufferPool: ByteBufferPool,
     listener: IMuxerInternal.IMuxerListener? = null,
     val pid: Short,
     private val transportErrorIndicator: Boolean = false,
     private val transportPriority: Boolean = false,
     private val transportScramblingControl: Byte = 0, // Not scrambled
     private var continuityCounter: Byte = 0,
-) : TSOutputCallback(listener) {
+) : TSOutputCallback(byteBufferPool, listener) {
 
     companion object {
         const val SYNC_BYTE: Byte = 0x47
@@ -52,7 +54,7 @@ open class TS(
 
         var packetIndicator = 0
 
-        val buffer = ByteBuffer.allocateDirect(PACKET_SIZE * MuxerConst.MAX_OUTPUT_PACKET_NUMBER)
+        val buffer = byteBufferPool.get(PACKET_SIZE * MuxerConst.MAX_OUTPUT_PACKET_NUMBER)
 
         while (payload?.hasRemaining() == true || adaptationFieldIndicator) {
             buffer.limit(buffer.position() + PACKET_SIZE)
diff --git a/core/src/main/java/io/github/thibaultbee/streampack/core/elements/endpoints/composites/muxers/ts/utils/TSOutputCallback.kt b/core/src/main/java/io/github/thibaultbee/streampack/core/elements/endpoints/composites/muxers/ts/utils/TSOutputCallback.kt
@@ -17,10 +17,15 @@ package io.github.thibaultbee.streampack.core.elements.endpoints.composites.muxe
 
 import io.github.thibaultbee.streampack.core.elements.data.SrtPacket
 import io.github.thibaultbee.streampack.core.elements.endpoints.composites.muxers.IMuxerInternal
+import io.github.thibaultbee.streampack.core.elements.utils.pool.ByteBufferPool
 
-open class TSOutputCallback(var listener: IMuxerInternal.IMuxerListener? = null) {
+open class TSOutputCallback(
+    protected val byteBufferPool: ByteBufferPool,
+    var listener: IMuxerInternal.IMuxerListener? = null
+) {
     protected fun writePacket(packet: SrtPacket) {
         packet.buffer.rewind()
         listener?.onOutputFrame(packet)
+        byteBufferPool.put(packet.buffer)
     }
 }
diff --git a/core/src/main/java/io/github/thibaultbee/streampack/core/elements/utils/pool/ByteBufferPool.kt b/core/src/main/java/io/github/thibaultbee/streampack/core/elements/utils/pool/ByteBufferPool.kt
@@ -46,9 +46,7 @@ class ByteBufferPool(private val isDirect: Boolean) : IBufferPool<ByteBuffer>, C
             buffers.tailMap(
                 capacity,
                 true
-            ).values
-                .mapNotNull(ArrayDeque<ByteBuffer>::removeFirstOrNull)
-                .firstOrNull()
+            ).values.firstNotNullOfOrNull(ArrayDeque<ByteBuffer>::removeFirstOrNull)
         }
         return if (buffer != null) {
             buffer.clear().limit(capacity)
diff --git a/core/src/test/java/io/github/thibaultbee/streampack/core/elements/endpoints/composites/muxers/ts/packets/PesTest.kt b/core/src/test/java/io/github/thibaultbee/streampack/core/elements/endpoints/composites/muxers/ts/packets/PesTest.kt
@@ -24,6 +24,7 @@ import io.github.thibaultbee.streampack.core.elements.endpoints.composites.muxer
 import io.github.thibaultbee.streampack.core.elements.utils.FakeFrames
 import io.github.thibaultbee.streampack.core.elements.utils.MockUtils
 import io.github.thibaultbee.streampack.core.elements.utils.ResourcesUtils
+import io.github.thibaultbee.streampack.core.elements.utils.pool.ByteBufferPool
 import org.junit.Test
 import java.io.File
 import java.nio.ByteBuffer
@@ -63,6 +64,7 @@ class PesTest {
 
         val expectedBuffers = readFrames(TEST_SAMPLES_DIR + "pes-video1")
         Pes(
+            ByteBufferPool(true),
             AssertEqualsBuffersMockMuxerListener(expectedBuffers),
             Stream(
                 VideoCodecConfig(
@@ -91,6 +93,7 @@ class PesTest {
 
         val expectedBuffers = readFrames(TEST_SAMPLES_DIR + "pes-audio1")
         Pes(
+            ByteBufferPool(true),
             AssertEqualsBuffersMockMuxerListener(expectedBuffers),
             Stream(AudioCodecConfig(), 256),
             true
@@ -114,6 +117,7 @@ class PesTest {
 
         val expectedBuffers = readFrames(TEST_SAMPLES_DIR + "pes-audio2")
         Pes(
+            ByteBufferPool(true),
             AssertEqualsBuffersMockMuxerListener(expectedBuffers),
             Stream(AudioCodecConfig(), 256),
             true
diff --git a/core/src/test/java/io/github/thibaultbee/streampack/core/elements/endpoints/composites/muxers/ts/packets/TransportStreamTest.kt b/core/src/test/java/io/github/thibaultbee/streampack/core/elements/endpoints/composites/muxers/ts/packets/TransportStreamTest.kt
@@ -15,7 +15,10 @@
  */
 package io.github.thibaultbee.streampack.core.elements.endpoints.composites.muxers.ts.packets
 
+import io.github.thibaultbee.streampack.core.elements.data.Packet
+import io.github.thibaultbee.streampack.core.elements.endpoints.composites.muxers.IMuxerInternal
 import io.github.thibaultbee.streampack.core.elements.endpoints.composites.muxers.ts.descriptors.AdaptationField
+import io.github.thibaultbee.streampack.core.elements.utils.pool.ByteBufferPool
 import org.junit.Assert.assertEquals
 import org.junit.Test
 import java.nio.ByteBuffer
@@ -26,9 +29,9 @@ class TransportStreamTest {
     class MockMuxerListener(
         private val expectedBuffer: ByteBuffer,
         private val numExpectedBuffer: Int = 1
-    ) : io.github.thibaultbee.streampack.core.elements.endpoints.composites.muxers.IMuxerInternal.IMuxerListener {
+    ) : IMuxerInternal.IMuxerListener {
         private var nBuffer = 0
-        override fun onOutputFrame(packet: io.github.thibaultbee.streampack.core.elements.data.Packet) {
+        override fun onOutputFrame(packet: Packet) {
             assertEquals(TS.PACKET_SIZE * numExpectedBuffer, packet.buffer.limit())
             assertEquals(TS.SYNC_BYTE, packet.buffer[0])
             for (i in 0 until min(packet.buffer.limit() - 4, expectedBuffer.remaining())) {
@@ -43,10 +46,10 @@ class TransportStreamTest {
     }
 
     class MockTransportStream(
-        muxerListener: io.github.thibaultbee.streampack.core.elements.endpoints.composites.muxers.IMuxerInternal.IMuxerListener,
+        muxerListener: IMuxerInternal.IMuxerListener,
         pid: Short = Random.nextInt().toShort()
     ) :
-        TS(muxerListener, pid) {
+        TS(ByteBufferPool(true), muxerListener, pid) {
         fun mockWrite(
             payload: ByteBuffer? = null,
             adaptationField: ByteBuffer? = null,
diff --git a/core/src/test/java/io/github/thibaultbee/streampack/core/elements/endpoints/composites/muxers/ts/tables/PatTest.kt b/core/src/test/java/io/github/thibaultbee/streampack/core/elements/endpoints/composites/muxers/ts/tables/PatTest.kt
@@ -21,6 +21,7 @@ import io.github.thibaultbee.streampack.core.elements.endpoints.composites.muxer
 import io.github.thibaultbee.streampack.core.elements.endpoints.composites.muxers.ts.packets.Pat
 import io.github.thibaultbee.streampack.core.elements.endpoints.composites.muxers.ts.packets.Pmt
 import io.github.thibaultbee.streampack.core.elements.endpoints.composites.muxers.ts.utils.AssertEqualsSingleBufferMockMuxerListener
+import io.github.thibaultbee.streampack.core.elements.utils.pool.ByteBufferPool
 import org.junit.Test
 
 class PatTest {
@@ -29,6 +30,7 @@ class PatTest {
      */
     @Test
     fun `simple pat test`() {
+        val byteBufferPool = ByteBufferPool(true)
         val expectedBuffer = TSResourcesUtils.readByteBuffer("pat.ts")
         val listener = AssertEqualsSingleBufferMockMuxerListener(expectedBuffer)
         val service = Service(
@@ -40,9 +42,9 @@ class PatTest {
             ),
             pcrPid = 4567,  // not used in this test
         )
-        Pmt(listener, service, emptyList(), 0x64).run { service.pmt = this }
+        Pmt(byteBufferPool, listener, service, emptyList(), 0x64).run { service.pmt = this }
 
-        Pat(listener, listOf(service), tsId = 0x437, versionNumber = 3).run {
+        Pat(byteBufferPool, listener, listOf(service), tsId = 0x437, versionNumber = 3).run {
             write()
         }
     }
diff --git a/core/src/test/java/io/github/thibaultbee/streampack/core/elements/endpoints/composites/muxers/ts/tables/SdtTest.kt b/core/src/test/java/io/github/thibaultbee/streampack/core/elements/endpoints/composites/muxers/ts/tables/SdtTest.kt
@@ -20,6 +20,7 @@ import io.github.thibaultbee.streampack.core.elements.endpoints.composites.muxer
 import io.github.thibaultbee.streampack.core.elements.endpoints.composites.muxers.ts.data.TSServiceInfo
 import io.github.thibaultbee.streampack.core.elements.endpoints.composites.muxers.ts.packets.Sdt
 import io.github.thibaultbee.streampack.core.elements.endpoints.composites.muxers.ts.utils.AssertEqualsSingleBufferMockMuxerListener
+import io.github.thibaultbee.streampack.core.elements.utils.pool.ByteBufferPool
 import org.junit.Test
 
 class SdtTest {
@@ -40,7 +41,7 @@ class SdtTest {
                 )
             )
         )
-        Sdt(listener, services, 0x0001, 0xff01.toShort()).run {
+        Sdt(ByteBufferPool(true), listener, services, 0x0001, 0xff01.toShort()).run {
             write()
         }
     }
