Fix IDR frame NALU drop race condition

moonlight-common/src/main/java/com/limelight/nvstream/av/video/VideoDecoderRenderer.java

@@ -7,7 +7,10 @@ public abstract class VideoDecoderRenderer {
 
 	public abstract void stop();
 
-	public abstract int submitDecodeUnit(byte[] frameData, int frameLength, int frameNumber, long receiveTimeMs);
+	// This is called once for each frame-start NALU. This means it will be called several times
+	// for an IDR frame which contains several parameter sets and the I-frame data.
+	public abstract int submitDecodeUnit(byte[] decodeUnitData, int decodeUnitLength, int decodeUnitType,
+										 int frameNumber, long receiveTimeMs);
 	
 	public abstract void cleanup();
 

moonlight-common/src/main/java/com/limelight/nvstream/jni/MoonBridge.java

@@ -13,6 +13,11 @@ public class MoonBridge {
     public static final int VIDEO_FORMAT_H264 = 1;
     public static final int VIDEO_FORMAT_H265 = 2;
 
+    public static final int BUFFER_TYPE_PICDATA = 0;
+    public static final int BUFFER_TYPE_SPS = 1;
+    public static final int BUFFER_TYPE_PPS = 2;
+    public static final int BUFFER_TYPE_VPS = 3;
+
     public static final int CAPABILITY_DIRECT_SUBMIT = 1;
     public static final int CAPABILITY_REFERENCE_FRAME_INVALIDATION_AVC = 2;
     public static final int CAPABILITY_REFERENCE_FRAME_INVALIDATION_HEVC = 4;
@@ -60,9 +65,12 @@ public class MoonBridge {
         }
     }
 
-    public static int bridgeDrSubmitDecodeUnit(byte[] frameData, int frameLength, int frameNumber, long receiveTimeMs) {
+    public static int bridgeDrSubmitDecodeUnit(byte[] decodeUnitData, int decodeUnitLength,
+                                               int decodeUnitType,
+                                               int frameNumber, long receiveTimeMs) {
         if (videoRenderer != null) {
-            return videoRenderer.submitDecodeUnit(frameData, frameLength, frameNumber, receiveTimeMs);
+            return videoRenderer.submitDecodeUnit(decodeUnitData, decodeUnitLength,
+                    decodeUnitType, frameNumber, receiveTimeMs);
         }
         else {
             return DR_OK;

moonlight-common/src/main/jni/moonlight-core/callbacks.c

@@ -81,7 +81,7 @@ Java_com_limelight_nvstream_jni_MoonBridge_init(JNIEnv *env, jobject class) {
     BridgeDrStartMethod = (*env)->GetStaticMethodID(env, class, "bridgeDrStart", "()V");
     BridgeDrStopMethod = (*env)->GetStaticMethodID(env, class, "bridgeDrStop", "()V");
     BridgeDrCleanupMethod = (*env)->GetStaticMethodID(env, class, "bridgeDrCleanup", "()V");
-    BridgeDrSubmitDecodeUnitMethod = (*env)->GetStaticMethodID(env, class, "bridgeDrSubmitDecodeUnit", "([BIIJ)I");
+    BridgeDrSubmitDecodeUnitMethod = (*env)->GetStaticMethodID(env, class, "bridgeDrSubmitDecodeUnit", "([BIIIJ)I");
     BridgeArInitMethod = (*env)->GetStaticMethodID(env, class, "bridgeArInit", "(I)I");
     BridgeArStartMethod = (*env)->GetStaticMethodID(env, class, "bridgeArStart", "()V");
     BridgeArStopMethod = (*env)->GetStaticMethodID(env, class, "bridgeArStop", "()V");
@@ -152,6 +152,7 @@ void BridgeDrCleanup(void) {
 
 int BridgeDrSubmitDecodeUnit(PDECODE_UNIT decodeUnit) {
     JNIEnv* env = GetThreadEnv();
+    int ret;
 
     if ((*env)->ExceptionCheck(env)) {
         return DR_OK;
@@ -169,14 +170,33 @@ int BridgeDrSubmitDecodeUnit(PDECODE_UNIT decodeUnit) {
     currentEntry = decodeUnit->bufferList;
     offset = 0;
     while (currentEntry != NULL) {
+        // Submit parameter set NALUs separately from picture data
+        if (currentEntry->bufferType != BUFFER_TYPE_PICDATA) {
+            // Use the beginning of the buffer each time since this is a separate
+            // invocation of the decoder each time.
+            (*env)->SetByteArrayRegion(env, DecodedFrameBuffer, 0, currentEntry->length, (jbyte*)currentEntry->data);
+
+            ret = (*env)->CallStaticIntMethod(env, GlobalBridgeClass, BridgeDrSubmitDecodeUnitMethod,
+                                              DecodedFrameBuffer, currentEntry->length, currentEntry->bufferType,
+                                              decodeUnit->frameNumber, decodeUnit->receiveTimeMs);
+            if ((*env)->ExceptionCheck(env)) {
+                return DR_OK;
+            }
+            else if (ret != DR_OK) {
+                return ret;
+            }
+        }
+        else {
             (*env)->SetByteArrayRegion(env, DecodedFrameBuffer, offset, currentEntry->length, (jbyte*)currentEntry->data);
             offset += currentEntry->length;
+        }
 
         currentEntry = currentEntry->next;
     }
 
     return (*env)->CallStaticIntMethod(env, GlobalBridgeClass, BridgeDrSubmitDecodeUnitMethod,
-                                       DecodedFrameBuffer, decodeUnit->fullLength, decodeUnit->frameNumber,
+                                       DecodedFrameBuffer, offset, BUFFER_TYPE_PICDATA,
+                                       decodeUnit->frameNumber,
                                        decodeUnit->receiveTimeMs);
 }
 

moonlight-common/src/main/jni/moonlight-core/moonlight-common-c

@@ -1 +1 @@
-Subproject commit 314a5937f4c2f4a9299ab5a00dba3e73c7384161
+Subproject commit 1c386a898731ae07168d52c835657409b54f7790