Add support for H.264/HEVC bitstreams with multiple sets of parameter set NALUs

app/src/main/java/com/limelight/binding/video/MediaCodecDecoderRenderer.java

@@ -3,6 +3,7 @@ package com.limelight.binding.video;
 import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 import java.util.concurrent.LinkedBlockingQueue;
@@ -48,11 +49,10 @@ public class MediaCodecDecoderRenderer extends VideoDecoderRenderer implements C
     private MediaCodecInfo hevcDecoder;
     private MediaCodecInfo av1Decoder;
 
-    private byte[] vpsBuffer;
+    private final ArrayList<byte[]> vpsBuffers = new ArrayList<>();
-    private byte[] spsBuffer;
+    private final ArrayList<byte[]> spsBuffers = new ArrayList<>();
-    private byte[] ppsBuffer;
+    private final ArrayList<byte[]> ppsBuffers = new ArrayList<>();
     private boolean submittedCsd;
-    private boolean submitCsdNextCall;
     private byte[] currentHdrMetadata;
 
     private int nextInputBufferIndex = -1;
@@ -544,10 +544,9 @@ public class MediaCodecDecoderRenderer extends VideoDecoderRenderer implements C
 
         // After reconfiguration, we must resubmit CSD buffers
         submittedCsd = false;
-        submitCsdNextCall = false;
+        vpsBuffers.clear();
-        vpsBuffer = null;
+        spsBuffers.clear();
-        spsBuffer = null;
+        ppsBuffers.clear();
-        ppsBuffer = null;
 
         if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.LOLLIPOP) {
             // This will contain the actual accepted input format attributes
@@ -1414,6 +1413,13 @@ public class MediaCodecDecoderRenderer extends VideoDecoderRenderer implements C
             activeWindowVideoStats.frameLossEvents++;
         }
 
+        // Reset CSD data for each IDR frame
+        if (lastFrameNumber != frameNumber && frameType == MoonBridge.FRAME_TYPE_IDR) {
+            vpsBuffers.clear();
+            spsBuffers.clear();
+            ppsBuffers.clear();
+        }
+
         lastFrameNumber = frameNumber;
 
         // Flip stats windows roughly every second
@@ -1462,9 +1468,10 @@ public class MediaCodecDecoderRenderer extends VideoDecoderRenderer implements C
             activeWindowVideoStats.measurementStartTimestamp = SystemClock.uptimeMillis();
         }
 
-        long timestampUs;
+        boolean csdSubmittedForThisFrame = false;
-        int codecFlags = 0;
 
+        // IDR frames require special handling for CSD buffer submission
+        if (frameType == MoonBridge.FRAME_TYPE_IDR) {
             // H264 SPS
             if (decodeUnitType == MoonBridge.BUFFER_TYPE_SPS && (videoFormat & MoonBridge.VIDEO_FORMAT_MASK_H264) != 0) {
                 numSpsIn++;
@@ -1591,65 +1598,86 @@ public class MediaCodecDecoderRenderer extends VideoDecoderRenderer implements C
                 // Annex B NALUs (including NALUs with escape sequences)
                 ByteBuffer escapedNalu = H264Utils.writeSPS(sps, decodeUnitLength);
 
-            // Batch this to submit together with PPS
+                // Construct the patched SPS
-            spsBuffer = new byte[startSeqLen + 1 + escapedNalu.limit()];
+                byte[] naluBuffer = new byte[startSeqLen + 1 + escapedNalu.limit()];
-            System.arraycopy(decodeUnitData, 0, spsBuffer, 0, startSeqLen + 1);
+                System.arraycopy(decodeUnitData, 0, naluBuffer, 0, startSeqLen + 1);
-            escapedNalu.get(spsBuffer, startSeqLen + 1, escapedNalu.limit());
+                escapedNalu.get(naluBuffer, startSeqLen + 1, escapedNalu.limit());
+
+                // Batch this to submit together with other CSD per AOSP docs
+                spsBuffers.add(naluBuffer);
                 return MoonBridge.DR_OK;
             }
             else if (decodeUnitType == MoonBridge.BUFFER_TYPE_VPS) {
                 numVpsIn++;
 
-            // Batch this to submit together with SPS and PPS per AOSP docs
+                // Batch this to submit together with other CSD per AOSP docs
-            vpsBuffer = new byte[decodeUnitLength];
+                byte[] naluBuffer = new byte[decodeUnitLength];
-            System.arraycopy(decodeUnitData, 0, vpsBuffer, 0, decodeUnitLength);
+                System.arraycopy(decodeUnitData, 0, naluBuffer, 0, decodeUnitLength);
+                vpsBuffers.add(naluBuffer);
                 return MoonBridge.DR_OK;
             }
             // Only the HEVC SPS hits this path (H.264 is handled above)
             else if (decodeUnitType == MoonBridge.BUFFER_TYPE_SPS) {
                 numSpsIn++;
 
-            // Batch this to submit together with VPS and PPS per AOSP docs
+                // Batch this to submit together with other CSD per AOSP docs
-            spsBuffer = new byte[decodeUnitLength];
+                byte[] naluBuffer = new byte[decodeUnitLength];
-            System.arraycopy(decodeUnitData, 0, spsBuffer, 0, decodeUnitLength);
+                System.arraycopy(decodeUnitData, 0, naluBuffer, 0, decodeUnitLength);
+                spsBuffers.add(naluBuffer);
                 return MoonBridge.DR_OK;
             }
             else if (decodeUnitType == MoonBridge.BUFFER_TYPE_PPS) {
                 numPpsIn++;
 
-            // If this is the first CSD blob or we aren't supporting
+                // Batch this to submit together with other CSD per AOSP docs
-            // fused IDR frames, we will submit the CSD blob in a
+                byte[] naluBuffer = new byte[decodeUnitLength];
-            // separate input buffer.
+                System.arraycopy(decodeUnitData, 0, naluBuffer, 0, decodeUnitLength);
+                ppsBuffers.add(naluBuffer);
+                return MoonBridge.DR_OK;
+            }
+            else if ((videoFormat & (MoonBridge.VIDEO_FORMAT_MASK_H264 | MoonBridge.VIDEO_FORMAT_MASK_H265)) != 0) {
+                // If this is the first CSD blob or we aren't supporting fused IDR frames, we will
+                // submit the CSD blob in a separate input buffer for each IDR frame.
                 if (!submittedCsd || !fusedIdrFrame) {
                     if (!fetchNextInputBuffer()) {
                         return MoonBridge.DR_NEED_IDR;
                     }
 
-                // When we get the PPS, submit the VPS and SPS together with
+                    // Submit all CSD when we receive the first non-CSD blob in an IDR frame
-                // the PPS, as required by AOSP docs on use of MediaCodec.
+                    for (byte[] vpsBuffer : vpsBuffers) {
-                if (vpsBuffer != null) {
                         nextInputBuffer.put(vpsBuffer);
                     }
-                if (spsBuffer != null) {
+                    for (byte[] spsBuffer : spsBuffers) {
                         nextInputBuffer.put(spsBuffer);
                     }
-
+                    for (byte[] ppsBuffer : ppsBuffers) {
-                // This is the CSD blob
+                        nextInputBuffer.put(ppsBuffer);
-                codecFlags |= MediaCodec.BUFFER_FLAG_CODEC_CONFIG;
-                timestampUs = 0;
                     }
-            else {
-                // Batch this to submit together with the next I-frame
-                ppsBuffer = new byte[decodeUnitLength];
-                System.arraycopy(decodeUnitData, 0, ppsBuffer, 0, decodeUnitLength);
 
-                // Next call will be I-frame data
+                    if (!queueNextInputBuffer(0, MediaCodec.BUFFER_FLAG_CODEC_CONFIG)) {
-                submitCsdNextCall = true;
+                        return MoonBridge.DR_NEED_IDR;
+                    }
 
-                return MoonBridge.DR_OK;
+                    // Remember that we already submitted CSD for this frame, so we don't do it
+                    // again in the fused IDR case below.
+                    csdSubmittedForThisFrame = true;
+
+                    // Remember that we submitted CSD globally for this MediaCodec instance
+                    submittedCsd = true;
+
+                    if (needsBaselineSpsHack) {
+                        needsBaselineSpsHack = false;
+
+                        if (!replaySps()) {
+                            return MoonBridge.DR_NEED_IDR;
+                        }
+
+                        LimeLog.info("SPS replay complete");
                     }
                 }
-        else {
+            }
+        }
+
         if (frameHostProcessingLatency != 0) {
             if (activeWindowVideoStats.minHostProcessingLatency != 0) {
                 activeWindowVideoStats.minHostProcessingLatency = (char) Math.min(activeWindowVideoStats.minHostProcessingLatency, frameHostProcessingLatency);
@@ -1675,36 +1703,34 @@ public class MediaCodecDecoderRenderer extends VideoDecoderRenderer implements C
             return MoonBridge.DR_NEED_IDR;
         }
 
-            if (submitCsdNextCall) {
+        int codecFlags = 0;
-                if (vpsBuffer != null) {
-                    nextInputBuffer.put(vpsBuffer);
-                }
-                if (spsBuffer != null) {
-                    nextInputBuffer.put(spsBuffer);
-                }
-                if (ppsBuffer != null) {
-                    nextInputBuffer.put(ppsBuffer);
-                }
-
-                submitCsdNextCall = false;
-            }
 
         if (frameType == MoonBridge.FRAME_TYPE_IDR) {
             codecFlags |= MediaCodec.BUFFER_FLAG_SYNC_FRAME;
+
+            // If we are using fused IDR frames, submit the CSD with each IDR frame
+            if (fusedIdrFrame && !csdSubmittedForThisFrame) {
+                for (byte[] vpsBuffer : vpsBuffers) {
+                    nextInputBuffer.put(vpsBuffer);
+                }
+                for (byte[] spsBuffer : spsBuffers) {
+                    nextInputBuffer.put(spsBuffer);
+                }
+                for (byte[] ppsBuffer : ppsBuffers) {
+                    nextInputBuffer.put(ppsBuffer);
+                }
+            }
         }
 
-            timestampUs = enqueueTimeMs * 1000;
+        long timestampUs = enqueueTimeMs * 1000;
-
         if (timestampUs <= lastTimestampUs) {
             // We can't submit multiple buffers with the same timestamp
             // so bump it up by one before queuing
             timestampUs = lastTimestampUs + 1;
         }
-
         lastTimestampUs = timestampUs;
 
         numFramesIn++;
-        }
 
         if (decodeUnitLength > nextInputBuffer.limit() - nextInputBuffer.position()) {
             IllegalArgumentException exception = new IllegalArgumentException(
@@ -1723,20 +1749,6 @@ public class MediaCodecDecoderRenderer extends VideoDecoderRenderer implements C
             return MoonBridge.DR_NEED_IDR;
         }
 
-        if ((codecFlags & MediaCodec.BUFFER_FLAG_CODEC_CONFIG) != 0) {
-            submittedCsd = true;
-
-            if (needsBaselineSpsHack) {
-                needsBaselineSpsHack = false;
-
-                if (!replaySps()) {
-                    return MoonBridge.DR_NEED_IDR;
-                }
-
-                LimeLog.info("SPS replay complete");
-            }
-        }
-
         return MoonBridge.DR_OK;
     }