Prevent a decoder stall from causing corruption of queued decode units

moonlight-common/src/com/limelight/nvstream/av/DecodeUnit.java

@@ -1,7 +1,10 @@
 package com.limelight.nvstream.av;
 
+import java.util.HashSet;
 import java.util.List;
 
+import com.limelight.nvstream.av.video.VideoPacket;
+
 public class DecodeUnit {
 	public static final int TYPE_UNKNOWN = 0;
 	public static final int TYPE_H264 = 1;
@@ -16,11 +19,13 @@ public class DecodeUnit {
 	private int frameNumber;
 	private long receiveTimestamp;
 	private int flags;
+	private HashSet<VideoPacket> backingPackets; 
 	
 	public DecodeUnit() {
 	}
 	
-	public void initialize(int type, List<ByteBufferDescriptor> bufferList, int dataLength, int frameNumber, long receiveTimestamp, int flags)
+	public void initialize(int type, List<ByteBufferDescriptor> bufferList, int dataLength,
+			int frameNumber, long receiveTimestamp, int flags, HashSet<VideoPacket> backingPackets)
 	{
 		this.type = type;
 		this.bufferList = bufferList;
@@ -28,6 +33,7 @@ public class DecodeUnit {
 		this.frameNumber = frameNumber;
 		this.receiveTimestamp = receiveTimestamp;
 		this.flags = flags;
+		this.backingPackets = backingPackets;
 	}
 	
 	public int getType()
@@ -59,4 +65,14 @@ public class DecodeUnit {
 	{
 		return flags;
 	}
+	
+	// Internal use only
+	public HashSet<VideoPacket> getBackingPackets() {
+		return backingPackets;
+	}
+	
+	// Internal use only
+	public void clearBackingPackets() {
+		backingPackets.clear();
+	}
 }

moonlight-common/src/com/limelight/nvstream/av/PopulatedBufferList.java

@@ -6,8 +6,12 @@ public class PopulatedBufferList<T> {
 	private ArrayBlockingQueue<T> populatedList;
 	private ArrayBlockingQueue<T> freeList;
 	
+	private BufferFactory factory;
+	
 	@SuppressWarnings("unchecked")
 	public PopulatedBufferList(int maxQueueSize, BufferFactory factory) {
+		this.factory = factory;
+		
 		this.populatedList = new ArrayBlockingQueue<T>(maxQueueSize, false);
 		this.freeList = new ArrayBlockingQueue<T>(maxQueueSize, false);
 		
@@ -25,13 +29,14 @@ public class PopulatedBufferList<T> {
 	}
 	
 	public void freePopulatedObject(T object) {
+		factory.cleanupObject(object);
 		freeList.add(object);
 	}
 	
 	public void clearPopulatedObjects() {
 		T object;
 		while ((object = populatedList.poll()) != null) {
-			freeList.add(object);
+			freePopulatedObject(object);
 		}
 	}
 	
@@ -49,5 +54,6 @@ public class PopulatedBufferList<T> {
 	
 	public static interface BufferFactory {
 		public Object createFreeBuffer();
+		public void cleanupObject(Object o);
 	}
 }

moonlight-common/src/com/limelight/nvstream/av/audio/AudioDepacketizer.java

@@ -32,6 +32,10 @@ public class AudioDepacketizer {
 				public Object createFreeBuffer() {
 					return new ByteBufferDescriptor(new byte[OpusDecoder.getMaxOutputShorts()*2], 0, OpusDecoder.getMaxOutputShorts()*2);
 				}
+
+				public void cleanupObject(Object o) {
+					// Nothing to do
+				}
 			});
 		}
 	}

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

@@ -1,5 +1,6 @@
 package com.limelight.nvstream.av.video;
 
+import java.util.HashSet;
 import java.util.LinkedList;
 
 import com.limelight.LimeLog;
@@ -14,6 +15,7 @@ public class VideoDepacketizer {
 	// Current frame state
 	private LinkedList<ByteBufferDescriptor> avcFrameDataChain = null;
 	private int avcFrameDataLength = 0;
+	private HashSet<VideoPacket> packetSet = null;
 	
 	// Sequencing state
 	private int lastPacketInStream = 0;
@@ -43,6 +45,16 @@ public class VideoDepacketizer {
 			public Object createFreeBuffer() {
 				return new DecodeUnit();
 			}
+
+			public void cleanupObject(Object o) {
+				DecodeUnit du = (DecodeUnit) o;
+				
+				// Disassociate video packets from this DU
+				for (VideoPacket pkt : du.getBackingPackets()) {
+					pkt.decodeUnitRefCount.decrementAndGet();
+				}
+				du.clearBackingPackets();
+			}
 		});
 	}
 	
@@ -54,6 +66,13 @@ public class VideoDepacketizer {
 	
 	private void cleanupAvcFrameState()
 	{
+		if (packetSet != null) {
+			for (VideoPacket pkt : packetSet) {
+				pkt.decodeUnitRefCount.decrementAndGet();
+			}
+			packetSet = null;
+		}
+		
 		avcFrameDataChain = null;
 		avcFrameDataLength = 0;
 	}
@@ -95,7 +114,10 @@ public class VideoDepacketizer {
 			
 			// Initialize the free DU
 			du.initialize(DecodeUnit.TYPE_H264, avcFrameDataChain,
-					avcFrameDataLength, frameNumber, frameStartTime, flags);
+					avcFrameDataLength, frameNumber, frameStartTime, flags, packetSet);
+			
+			// Packets now owned by the DU
+			packetSet = null;
 			
 			controlListener.connectionReceivedFrame(frameNumber);
 			
@@ -136,6 +158,7 @@ public class VideoDepacketizer {
 						// Setup state for the new NAL
 						avcFrameDataChain = new LinkedList<ByteBufferDescriptor>();
 						avcFrameDataLength = 0;
+						packetSet = new HashSet<VideoPacket>();
 						
 						if (cachedSpecialDesc.data[cachedSpecialDesc.offset+cachedSpecialDesc.length] == 0x65) {
 							// This is the NALU code for I-frame data
@@ -190,6 +213,10 @@ public class VideoDepacketizer {
 			if (isDecodingH264 && avcFrameDataChain != null)
 			{
 				ByteBufferDescriptor data = new ByteBufferDescriptor(location.data, start, location.offset-start);
+				
+				if (packetSet.add(packet)) {
+					packet.decodeUnitRefCount.incrementAndGet();
+				}
 
 				// Add a buffer descriptor describing the NAL data in this packet
 				avcFrameDataChain.add(data);
@@ -205,11 +232,17 @@ public class VideoDepacketizer {
 			frameStartTime = System.currentTimeMillis();
 			avcFrameDataChain = new LinkedList<ByteBufferDescriptor>();
 			avcFrameDataLength = 0;
+			packetSet = new HashSet<VideoPacket>();
 		}
 		
 		// Add the payload data to the chain
 		avcFrameDataChain.add(new ByteBufferDescriptor(location));
 		avcFrameDataLength += location.length;
+		
+		// The receive thread can't use this until we're done with it
+		if (packetSet.add(packet)) {
+			packet.decodeUnitRefCount.incrementAndGet();
+		}
 	}
 	
 	private static boolean isFirstPacket(int flags) {
@@ -331,7 +364,7 @@ public class VideoDepacketizer {
 			addInputDataSlow(packet, cachedReassemblyDesc);
 		}
 		else
-		{	
+		{
 			// Everything else can take the fast path
 			addInputDataFast(packet, cachedReassemblyDesc, firstPacket);
 		}
@@ -374,7 +407,7 @@ public class VideoDepacketizer {
 	}
 	
 	public void freeDecodeUnit(DecodeUnit du)
-	{
+	{	
 		decodedUnits.freePopulatedObject(du);
 	}
 }

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

@@ -2,6 +2,7 @@ package com.limelight.nvstream.av.video;
 
 import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
+import java.util.concurrent.atomic.AtomicInteger;
 
 import com.limelight.nvstream.av.ByteBufferDescriptor;
 import com.limelight.nvstream.av.RtpPacket;
@@ -19,6 +20,8 @@ public class VideoPacket implements RtpPacketFields {
 	
 	private short rtpSequenceNumber;
 	
+	AtomicInteger decodeUnitRefCount = new AtomicInteger();
+	
 	public static final int FLAG_CONTAINS_PIC_DATA = 0x1;
 	public static final int FLAG_EOF = 0x2;
 	public static final int FLAG_SOF = 0x4;

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

@@ -10,6 +10,7 @@ import java.net.Socket;
 import java.net.SocketException;
 import java.util.LinkedList;
 
+import com.limelight.LimeLog;
 import com.limelight.nvstream.NvConnectionListener;
 import com.limelight.nvstream.StreamConfiguration;
 import com.limelight.nvstream.av.ConnectionStatusListener;
@@ -203,6 +204,7 @@ public class VideoStream {
 
 				byte[] buffer;
 				DatagramPacket packet = new DatagramPacket(new byte[1], 1); // Placeholder array
+				int iterationStart;
 				while (!isInterrupted())
 				{
 					try {
@@ -227,9 +229,20 @@ public class VideoStream {
 								depacketizer.addInputData(queuedPacket);
 							}
 						}
-						
-						// The ring is large enough to account for the maximum queued packets
-						ringIndex = (ringIndex + 1) % VIDEO_RING_SIZE;
+
+						// Go to the next free element in the ring
+						iterationStart = ringIndex; 
+						do {
+							ringIndex = (ringIndex + 1) % VIDEO_RING_SIZE;
+							if (ringIndex == iterationStart) {								
+								// Reinitialize the video ring since they're all being used
+								LimeLog.warning("Packet ring wrapped around!");
+								for (int i = 0; i < VIDEO_RING_SIZE; i++) {
+									ring[i] = new VideoPacket(new byte[requiredBufferSize]);
+								}
+								break;
+							}
+						} while (ring[ringIndex].decodeUnitRefCount.get() != 0);
 					} catch (IOException e) {
 						listener.connectionTerminated(e);
 						return;