refactor(core): make the endpoint responsible for closing the frame

Author: ThibaultBee

core/src/androidTest/java/io/github/thibaultbee/streampack/core/elements/endpoints/DummyEndpoint.kt

@@ -63,13 +63,14 @@ class DummyEndpoint : IEndpointInternal {
         _isOpenFlow.emit(false)
     }
 
-    override suspend fun write(frame: Frame, streamPid: Int) {
+    override suspend fun write(frame: Frame, streamPid: Int, onFrameProcessed: (() -> Unit)) {
         Log.i(TAG, "write: $frame")
         _frameFlow.emit(frame)
         when {
             frame.isAudio -> numOfAudioFramesWritten++
             frame.isVideo -> numOfVideoFramesWritten++
         }
+        onFrameProcessed()
     }
 
     override fun addStreams(streamConfigs: List<CodecConfig>): Map<CodecConfig, Int> {

core/src/main/java/io/github/thibaultbee/streampack/core/elements/endpoints/CombineEndpoint.kt

@@ -23,15 +23,24 @@ import io.github.thibaultbee.streampack.core.elements.utils.combineStates
 import io.github.thibaultbee.streampack.core.elements.utils.extensions.intersect
 import io.github.thibaultbee.streampack.core.logger.Logger
 import io.github.thibaultbee.streampack.core.pipelines.outputs.encoding.EncodingPipelineOutputDispatcherProvider
+import io.github.thibaultbee.streampack.core.pipelines.utils.MultiThrowable
+import kotlinx.coroutines.CompletableDeferred
+import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Deferred
+import kotlinx.coroutines.SupervisorJob
 import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.launch
 
 
 /**
  * Combines multiple endpoints into one.
  *
  * @param endpoints Endpoints to combine
+ * @param coroutineDispatcher Coroutine dispatcher to use for frame writing
  */
-fun CombineEndpoint(vararg endpoints: IEndpointInternal) = CombineEndpoint(endpoints.toList())
+fun CombineEndpoint(vararg endpoints: IEndpointInternal, coroutineDispatcher: CoroutineDispatcher) =
+    CombineEndpoint(endpoints.toList(), coroutineDispatcher)
 
 /**
  * Combines multiple endpoints into one.
@@ -42,9 +51,17 @@ fun CombineEndpoint(vararg endpoints: IEndpointInternal) = CombineEndpoint(endpo
  *
  * For specific behavior like reconnecting your remote endpoint, you can create a custom endpoint that
  * inherits from [CombineEndpoint] and override [open], [close], [startStream], [stopStream].
+ *
+ * @param endpointInternals List of endpoints to combine
+ * @param coroutineDispatcher Coroutine dispatcher to use for frame writing
  */
-open class CombineEndpoint(protected val endpointInternals: List<IEndpointInternal>) :
+open class CombineEndpoint(
+    protected val endpointInternals: List<IEndpointInternal>,
+    coroutineDispatcher: CoroutineDispatcher
+) :
     IEndpointInternal {
+    private val coroutineScope = CoroutineScope(SupervisorJob() + coroutineDispatcher)
+
     /**
      * Internal map of endpoint streamId to real streamIds
      */
@@ -185,31 +202,36 @@ open class CombineEndpoint(protected val endpointInternals: List<IEndpointIntern
      *
      * If all endpoints write fails, it throws the exception of the first endpoint that failed.
      */
-    override suspend fun write(frame: Frame, streamPid: Int) {
-        val currentBufferPos = frame.buffer.position()
-        var numOfThrowable = 0
-        var throwable: Throwable? = null
+    override suspend fun write(frame: Frame, streamPid: Int, onFrameProcessed: (() -> Unit)) {
+        val throwables = mutableListOf<Throwable>()
 
-        endpointInternals.forEach { endpoint ->
+        /**
+         * Track the number of frames written and processed to call onFrameProcessed only once
+         * when all endpoints have processed the frame.
+         */
+        val deferreds = mutableListOf<Deferred<*>>()
+
+        endpointInternals.filter { it.isOpenFlow.value }.forEach { endpoint ->
             try {
-                if (endpoint.isOpenFlow.value) {
-                    val endpointStreamId = endpointsToStreamIdsMap[Pair(endpoint, streamPid)]!!
-                    endpoint.write(frame, endpointStreamId)
+                val deferred = CompletableDeferred<Unit>()
+                val duplicatedFrame = frame.copy(rawBuffer = frame.rawBuffer.duplicate())
+                val endpointStreamId = endpointsToStreamIdsMap[Pair(endpoint, streamPid)]!!
 
-                    // Reset buffer position to write frame to next endpoint
-                    frame.buffer.position(currentBufferPos)
-                }
+                deferreds += deferred
+                endpoint.write(duplicatedFrame, endpointStreamId, { deferred.complete(Unit) })
             } catch (t: Throwable) {
-                Logger.e(TAG, "Failed to write frame to endpoint $endpoint", t)
-                if (throwable == null) {
-                    throwable = t
-                }
-                numOfThrowable++
+                Logger.e(TAG, "Failed to get stream id for endpoint $endpoint", t)
+                throwables += t
             }
         }
 
-        if (numOfThrowable == endpointInternals.size) {
-            throw throwable!!
+        coroutineScope.launch {
+            deferreds.forEach { it.await() }
+            onFrameProcessed()
+        }
+
+        if (throwables.isNotEmpty()) {
+            throw MultiThrowable(throwables)
         }
     }
 
@@ -245,7 +267,8 @@ class CombineEndpointFactory(private val endpointFactory: List<IEndpointInternal
         dispatcherProvider: EncodingPipelineOutputDispatcherProvider
     ): IEndpointInternal {
         return CombineEndpoint(
-            endpointFactory.map { it.create(context, dispatcherProvider) }
+            endpointFactory.map { it.create(context, dispatcherProvider) },
+            dispatcherProvider.defaultDispatcher
         )
     }
 }

core/src/main/java/io/github/thibaultbee/streampack/core/elements/endpoints/DualEndpoint.kt

@@ -18,6 +18,7 @@ package io.github.thibaultbee.streampack.core.elements.endpoints
 import android.content.Context
 import io.github.thibaultbee.streampack.core.configuration.mediadescriptor.MediaDescriptor
 import io.github.thibaultbee.streampack.core.pipelines.outputs.encoding.EncodingPipelineOutputDispatcherProvider
+import kotlinx.coroutines.CoroutineDispatcher
 
 /**
  * An implementation of [CombineEndpoint] that combines two endpoints.
@@ -29,11 +30,15 @@ import io.github.thibaultbee.streampack.core.pipelines.outputs.encoding.Encoding
  *
  * @param mainEndpoint the main endpoint
  * @param secondEndpoint the second endpoint
+ * @param coroutineDispatcher the coroutine dispatcher
  */
 open class DualEndpoint(
-    private val mainEndpoint: IEndpointInternal, private val secondEndpoint: IEndpointInternal
+    private val mainEndpoint: IEndpointInternal,
+    private val secondEndpoint: IEndpointInternal,
+    coroutineDispatcher: CoroutineDispatcher
 ) : CombineEndpoint(
-    listOf(secondEndpoint, mainEndpoint)
+    listOf(secondEndpoint, mainEndpoint),
+    coroutineDispatcher
 ) {
     /**
      * Opens the [mainEndpoint].
@@ -91,7 +96,8 @@ class DualEndpointFactory(
     ): IEndpointInternal {
         return DualEndpoint(
             mainEndpointFactory.create(context, dispatcherProvider),
-            secondEndpointFactory.create(context, dispatcherProvider)
+            secondEndpointFactory.create(context, dispatcherProvider),
+            dispatcherProvider.defaultDispatcher
         )
     }
 }

core/src/main/java/io/github/thibaultbee/streampack/core/elements/endpoints/DynamicEndpoint.kt

@@ -98,7 +98,8 @@ open class DynamicEndpoint(
 
     override fun addStream(streamConfig: CodecConfig) = endpoint.addStream(streamConfig)
 
-    override suspend fun write(frame: Frame, streamPid: Int) = endpoint.write(frame, streamPid)
+    override suspend fun write(frame: Frame, streamPid: Int, onFrameProcessed: () -> Unit) =
+        endpoint.write(frame, streamPid, onFrameProcessed)
 
     override suspend fun startStream() = endpoint.startStream()
 

core/src/main/java/io/github/thibaultbee/streampack/core/elements/endpoints/IEndpoint.kt

@@ -39,12 +39,18 @@ interface IEndpointInternal : IEndpoint, SuspendStreamable,
     /**
      * Writes a [Frame] to the [IEndpointInternal].
      *
+     * The [onFrameProcessed] callback must be called when the frame has been processed and the [Frame.rawBuffer] is not used anymore.
+     * The [IEndpointInternal] must called [onFrameProcessed] even if the frame is dropped or it somehow crashes.
+     * Also, once [onFrameProcessed] is called, the [Frame.rawBuffer] must not be used anymore by the [IEndpointInternal].
+     *
      * @param frame the [Frame] to write
      * @param streamPid the stream id the [Frame] belongs to
+     * @param onFrameProcessed a callback called when the [Frame.rawBuffer] is not used anymore
      */
     suspend fun write(
         frame: Frame,
-        streamPid: Int
+        streamPid: Int,
+        onFrameProcessed: (() -> Unit)
     )
 
     /**

core/src/main/java/io/github/thibaultbee/streampack/core/elements/endpoints/MediaMuxerEndpoint.kt

@@ -140,40 +140,47 @@ class MediaMuxerEndpoint(
     }
 
     override suspend fun write(
-        frame: Frame, streamPid: Int
+        frame: Frame, streamPid: Int, onFrameProcessed: () -> Unit
     ) = withContext(ioDispatcher) {
         mutex.withLock {
-            if (state != State.STARTED && state != State.PENDING_START) {
-                Logger.w(TAG, "Trying to write while not started. Current state: $state")
-                return@withContext
-            }
+            try {
+                if (state != State.STARTED && state != State.PENDING_START) {
+                    Logger.w(TAG, "Trying to write while not started. Current state: $state")
+                    return@withContext
+                }
 
-            val mediaMuxer = requireNotNull(mediaMuxer) { "MediaMuxer is not initialized" }
+                val mediaMuxer = requireNotNull(mediaMuxer) { "MediaMuxer is not initialized" }
 
-            if ((state == State.PENDING_START) && (streamIdToTrackId.size < numOfStreams)) {
-                addTrack(mediaMuxer, streamPid, frame.format)
-                if (streamIdToTrackId.size == numOfStreams) {
-                    mediaMuxer.start()
-                    setState(State.STARTED)
+                if ((state == State.PENDING_START) && (streamIdToTrackId.size < numOfStreams)) {
+                    addTrack(mediaMuxer, streamPid, frame.format)
+                    if (streamIdToTrackId.size == numOfStreams) {
+                        mediaMuxer.start()
+                        setState(State.STARTED)
+                    }
                 }
-            }
 
-            if (state == State.STARTED) {
-                val trackId = streamIdToTrackId[streamPid]
-                    ?: throw IllegalStateException("Could not find trackId for streamPid $streamPid: ${frame.format}")
-                val info = BufferInfo().apply {
-                    set(
-                        0,
-                        frame.buffer.remaining(),
-                        frame.ptsInUs,
-                        if (frame.isKeyFrame) BUFFER_FLAG_KEY_FRAME else 0
-                    )
-                }
-                try {
-                    mediaMuxer.writeSampleData(trackId, frame.buffer, info)
-                } catch (e: IllegalStateException) {
-                    Logger.w(TAG, "MediaMuxer is in an illegal state. ${e.message}")
+                if (state == State.STARTED) {
+                    val trackId = streamIdToTrackId[streamPid]
+                        ?: throw IllegalStateException("Could not find trackId for streamPid $streamPid: ${frame.format}")
+                    val info = BufferInfo().apply {
+                        set(
+                            0,
+                            frame.buffer.remaining(),
+                            frame.ptsInUs,
+                            if (frame.isKeyFrame) BUFFER_FLAG_KEY_FRAME else 0
+                        )
+                    }
+                    try {
+                        mediaMuxer.writeSampleData(trackId, frame.buffer, info)
+                    } catch (e: IllegalStateException) {
+                        Logger.w(TAG, "MediaMuxer is in an illegal state. ${e.message}")
+                    }
                 }
+            } catch (t: Throwable) {
+                Logger.e(TAG, "Error while writing frame: ${t.message}")
+                throw t
+            } finally {
+                onFrameProcessed()
             }
         }
     }

core/src/main/java/io/github/thibaultbee/streampack/core/elements/endpoints/composites/CompositeEndpoint.kt

@@ -76,8 +76,9 @@ class CompositeEndpoint(
 
     override suspend fun write(
         frame: Frame,
-        streamPid: Int
-    ) = muxer.write(frame, streamPid)
+        streamPid: Int,
+        onFrameProcessed: (() -> Unit)
+    ) = muxer.write(frame, streamPid, onFrameProcessed)
 
     override fun addStreams(streamConfigs: List<CodecConfig>): Map<CodecConfig, Int> {
         mutex.tryLock()

core/src/main/java/io/github/thibaultbee/streampack/core/elements/endpoints/composites/muxers/IMuxerInternal.kt

@@ -32,7 +32,7 @@ interface IMuxerInternal :
         fun onOutputFrame(packet: Packet)
     }
 
-    fun write(frame: Frame, streamPid: Int)
+    fun write(frame: Frame, streamPid: Int, onFrameProcessed: () -> Unit)
 
     fun addStreams(streamsConfig: List<CodecConfig>): Map<CodecConfig, Int>
 

core/src/main/java/io/github/thibaultbee/streampack/core/elements/endpoints/composites/muxers/flv/FlvMuxer.kt

@@ -41,42 +41,46 @@ class FlvMuxer(
     override val streamConfigs: List<CodecConfig>
         get() = streams.map { it.config }
 
-    override fun write(frame: Frame, streamPid: Int) {
+    override fun write(frame: Frame, streamPid: Int, onFrameProcessed: () -> Unit) {
         synchronized(this) {
-            if (!hasFirstFrame) {
-                /**
-                 * Wait for first video frame to start (only if video is present)
-                 */
-                if (hasVideo) {
-                    // Expected first video key frame
-                    if (frame.isVideo && frame.isKeyFrame) {
+            try {
+                if (!hasFirstFrame) {
+                    /**
+                     * Wait for first video frame to start (only if video is present)
+                     */
+                    if (hasVideo) {
+                        // Expected first video key frame
+                        if (frame.isVideo && frame.isKeyFrame) {
+                            startUpTime = frame.ptsInUs
+                            hasFirstFrame = true
+                        } else {
+                            // Drop
+                            return
+                        }
+                    } else {
+                        // Audio only
                         startUpTime = frame.ptsInUs
                         hasFirstFrame = true
-                    } else {
-                        // Drop
-                        return
                     }
-                } else {
-                    // Audio only
-                    startUpTime = frame.ptsInUs
-                    hasFirstFrame = true
                 }
-            }
-        }
 
-        if (frame.ptsInUs < startUpTime!!) {
-            return
-        }
+                if (frame.ptsInUs < startUpTime!!) {
+                    return
+                }
 
-        frame.ptsInUs -= startUpTime!!
-        val stream = streams[streamPid]
-        val sendHeader = stream.sendHeader
-        stream.sendHeader = false
-        val flvTags = AVTagsFactory(frame, stream.config, sendHeader).build()
-        flvTags.forEach {
-            listener?.onOutputFrame(
-                Packet(it.write(), frame.ptsInUs)
-            )
+                frame.ptsInUs -= startUpTime!!
+                val stream = streams[streamPid]
+                val sendHeader = stream.sendHeader
+                stream.sendHeader = false
+                val flvTags = AVTagsFactory(frame, stream.config, sendHeader).build()
+                flvTags.forEach {
+                    listener?.onOutputFrame(
+                        Packet(it.write(), frame.ptsInUs)
+                    )
+                }
+            } finally {
+                onFrameProcessed()
+            }
         }
     }
 

core/src/main/java/io/github/thibaultbee/streampack/core/elements/endpoints/composites/muxers/mp4/Mp4Muxer.kt

@@ -56,13 +56,17 @@ class Mp4Muxer(
 
     override val streamConfigs: List<CodecConfig>
         get() = tracks.map { it.config }
-    
-    override fun write(frame: Frame, streamPid: Int) {
+
+    override fun write(frame: Frame, streamPid: Int, onFrameProcessed: () -> Unit) {
         synchronized(this) {
-            if (segmenter!!.mustWriteSegment(frame)) {
-                writeSegment()
+            try {
+                if (segmenter!!.mustWriteSegment(frame)) {
+                    writeSegment()
+                }
+                currentSegment!!.add(frame, streamPid)
+            } finally {
+                onFrameProcessed()
             }
-            currentSegment!!.add(frame, streamPid)
         }
     }