avcodec: better handling of packet loss

1 files changed, 178 insertions, 51 deletions

diff --git a/modules/avcodec/decode.c b/modules/avcodec/decode.c
index a88960d..77e6e77 100644
--- a/modules/avcodec/decode.c
+++ b/modules/avcodec/decode.c
@@ -23,19 +23,49 @@
 #endif
 
 
+enum {
+	DECODE_MAXSZ = 524288,
+};
+
+
 struct viddec_state {
 	AVCodec *codec;
 	AVCodecContext *ctx;
 	AVFrame *pict;
 	struct mbuf *mb;
 	bool got_keyframe;
+	size_t frag_start;
+	bool frag;
+	uint16_t frag_seq;
+
+	struct {
+		unsigned n_key;
+		unsigned n_lost;
+	} stats;
 };
 
 
+static inline int16_t seq_diff(uint16_t x, uint16_t y)
+{
+	return (int16_t)(y - x);
+}
+
+
+static inline void fragment_rewind(struct viddec_state *vds)
+{
+	vds->mb->pos = vds->frag_start;
+	vds->mb->end = vds->frag_start;
+}
+
+
 static void destructor(void *arg)
 {
 	struct viddec_state *st = arg;
 
+	debug("avcodec: decoder stats"
+	      " (keyframes:%u, lost_fragments:%u)\n",
+	      st->stats.n_key, st->stats.n_lost);
+
 	mem_deref(st->mb);
 
 	if (st->ctx) {
@@ -139,24 +169,16 @@ int decode_update(struct viddec_state **vdsp, const struct vidcodec *vc,
 }
 
 
-static int ffdecode(struct viddec_state *st, struct vidframe *frame,
-		    bool eof, struct mbuf *src)
+static int ffdecode(struct viddec_state *st, struct vidframe *frame)
 {
-	int i, got_picture, ret, err;
-
-	/* assemble packets in "mbuf" */
-	err = mbuf_write_mem(st->mb, mbuf_buf(src), mbuf_get_left(src));
-	if (err)
-		return err;
-
-	if (!eof)
-		return 0;
+	int i, got_picture, ret;
+	int err = 0;
 
 	st->mb->pos = 0;
 
 	if (!st->got_keyframe) {
-		err = EPROTO;
-		goto out;
+		debug("avcodec: waiting for key frame ..\n");
+		return 0;
 	}
 
 #if LIBAVCODEC_VERSION_INT >= ((57<<16)+(37<<8)+100)
@@ -166,10 +188,11 @@ static int ffdecode(struct viddec_state *st, struct vidframe *frame,
 
 		av_init_packet(&avpkt);
 		avpkt.data = st->mb->buf;
-		avpkt.size = (int)mbuf_get_left(st->mb);
+		avpkt.size = (int)st->mb->end;
 
 		ret = avcodec_send_packet(st->ctx, &avpkt);
 		if (ret < 0) {
+			warning("avcodec_send_packet error ret=%d\n", ret);
 			err = EBADMSG;
 			goto out;
 		}
@@ -179,6 +202,7 @@ static int ffdecode(struct viddec_state *st, struct vidframe *frame,
 			goto out;
 		}
 		else if (ret < 0) {
+			warning("avcodec_receive_frame error ret=%d\n", ret);
 			err = EBADMSG;
 			goto out;
 		}
@@ -190,14 +214,14 @@ static int ffdecode(struct viddec_state *st, struct vidframe *frame,
 #elif LIBAVCODEC_VERSION_INT <= ((52<<16)+(23<<8)+0)
 	ret = avcodec_decode_video(st->ctx, st->pict, &got_picture,
 				   st->mb->buf,
-				   (int)mbuf_get_left(st->mb));
+				   (int)st->mb->end);
 #else
 	do {
 		AVPacket avpkt;
 
 		av_init_packet(&avpkt);
 		avpkt.data = st->mb->buf;
-		avpkt.size = (int)mbuf_get_left(st->mb);
+		avpkt.size = (int)st->mb->end;
 
 		ret = avcodec_decode_video2(st->ctx, st->pict,
 					    &got_picture, &avpkt);
@@ -209,8 +233,6 @@ static int ffdecode(struct viddec_state *st, struct vidframe *frame,
 		goto out;
 	}
 
-	mbuf_skip_to_end(src);
-
 	if (got_picture) {
 
 #if LIBAVCODEC_VERSION_INT >= ((53<<16)+(5<<8)+0)
@@ -245,14 +267,12 @@ static int ffdecode(struct viddec_state *st, struct vidframe *frame,
 	}
 
  out:
-	if (eof)
-		mbuf_rewind(st->mb);
-
 	return err;
 }
 
 
-static int h264_decode(struct viddec_state *st, bool *intra, struct mbuf *src)
+int decode_h264(struct viddec_state *st, struct vidframe *frame,
+		bool *intra, bool marker, uint16_t seq, struct mbuf *src)
 {
 	struct h264_hdr h264_hdr;
 	const uint8_t nal_seq[3] = {0, 0, 1};
@@ -264,32 +284,40 @@ static int h264_decode(struct viddec_state *st, bool *intra, struct mbuf *src)
 	if (err)
 		return err;
 
+#if 0
+	re_printf("avcodec: decode: %s %s type=%2d  \n",
+		  marker ? "[M]" : "   ",
+		  h264_is_keyframe(h264_hdr.type) ? "<KEY>" : "     ",
+		  h264_hdr.type);
+#endif
+
 	if (h264_hdr.f) {
 		info("avcodec: H264 forbidden bit set!\n");
 		return EBADMSG;
 	}
 
+	if (st->frag && h264_hdr.type != H264_NAL_FU_A) {
+		debug("avcodec: lost fragments; discarding previous NAL\n");
+		fragment_rewind(st);
+		st->frag = false;
+		++st->stats.n_lost;
+	}
+
 	/* handle NAL types */
 	if (1 <= h264_hdr.type && h264_hdr.type <= 23) {
 
-		if (!st->got_keyframe) {
-			switch (h264_hdr.type) {
-
-			case H264_NAL_PPS:
-			case H264_NAL_SPS:
-				st->got_keyframe = true;
-				break;
-			}
-		}
-
 		if (h264_is_keyframe(h264_hdr.type))
 			*intra = true;
 
+		--src->pos;
+
 		/* prepend H.264 NAL start sequence */
-		mbuf_write_mem(st->mb, nal_seq, 3);
+		err  = mbuf_write_mem(st->mb, nal_seq, 3);
 
-		/* encode NAL header back to buffer */
-		err = h264_hdr_encode(&h264_hdr, st->mb);
+		err |= mbuf_write_mem(st->mb, mbuf_buf(src),
+				      mbuf_get_left(src));
+		if (err)
+			goto out;
 	}
 	else if (H264_NAL_FU_A == h264_hdr.type) {
 		struct h264_fu fu;
@@ -300,6 +328,16 @@ static int h264_decode(struct viddec_state *st, bool *intra, struct mbuf *src)
 		h264_hdr.type = fu.type;
 
 		if (fu.s) {
+			if (st->frag) {
+				debug("avcodec: lost fragments;"
+				      " ignoring NAL\n");
+				fragment_rewind(st);
+				++st->stats.n_lost;
+			}
+
+			st->frag_start = st->mb->pos;
+			st->frag = true;
+
 			if (h264_is_keyframe(fu.type))
 				*intra = true;
 
@@ -309,37 +347,75 @@ static int h264_decode(struct viddec_state *st, bool *intra, struct mbuf *src)
 			/* encode NAL header back to buffer */
 			err = h264_hdr_encode(&h264_hdr, st->mb);
 		}
+		else {
+			if (!st->frag) {
+				debug("avcodec: ignoring fragment\n");
+				++st->stats.n_lost;
+				return 0;
+			}
+
+			if (seq_diff(st->frag_seq, seq) != 1) {
+				debug("avcodec: lost fragments detected\n");
+				fragment_rewind(st);
+				st->frag = false;
+				++st->stats.n_lost;
+				return 0;
+			}
+		}
+
+		err = mbuf_write_mem(st->mb, mbuf_buf(src),
+				     mbuf_get_left(src));
+		if (err)
+			goto out;
+
+		if (fu.e)
+			st->frag = false;
+
+		st->frag_seq = seq;
 	}
 	else {
 		warning("avcodec: unknown NAL type %u\n", h264_hdr.type);
 		return EBADMSG;
 	}
 
-	return err;
-}
+	if (*intra) {
+		st->got_keyframe = true;
+		++st->stats.n_key;
+	}
 
+	if (!marker) {
 
-int decode_h264(struct viddec_state *st, struct vidframe *frame,
-		bool *intra, bool eof, uint16_t seq, struct mbuf *src)
-{
-	int err;
-
-	(void)seq;
+		if (st->mb->end > DECODE_MAXSZ) {
+			warning("avcodec: decode buffer size exceeded\n");
+			err = ENOMEM;
+			goto out;
+		}
 
-	if (!src)
 		return 0;
+	}
 
-	err = h264_decode(st, intra, src);
+	if (st->frag) {
+		err = EPROTO;
+		goto out;
+	}
+
+	err = ffdecode(st, frame);
 	if (err)
-		return err;
+		goto out;
 
-	return ffdecode(st, frame, eof, src);
+ out:
+	mbuf_rewind(st->mb);
+	st->frag = false;
+
+	return err;
 }
 
 
 int decode_mpeg4(struct viddec_state *st, struct vidframe *frame,
-		 bool *intra, bool eof, uint16_t seq, struct mbuf *src)
+		 bool *intra, bool marker, uint16_t seq, struct mbuf *src)
 {
+	int err;
+
 	if (!src)
 		return 0;
 
@@ -350,7 +426,30 @@ int decode_mpeg4(struct viddec_state *st, struct vidframe *frame,
 	/* let the decoder handle this */
 	st->got_keyframe = true;
 
-	return ffdecode(st, frame, eof, src);
+	err = mbuf_write_mem(st->mb, mbuf_buf(src),
+			     mbuf_get_left(src));
+	if (err)
+		goto out;
+
+	if (!marker) {
+
+		if (st->mb->end > DECODE_MAXSZ) {
+			warning("avcodec: decode buffer size exceeded\n");
+			err = ENOMEM;
+			goto out;
+		}
+
+		return 0;
+	}
+
+	err = ffdecode(st, frame);
+	if (err)
+		goto out;
+
+ out:
+	mbuf_rewind(st->mb);
+
+	return err;
 }
 
 
@@ -386,8 +485,9 @@ int decode_h263(struct viddec_state *st, struct vidframe *frame,
 
 	if (!hdr.i) {
 		st->got_keyframe = true;
-		if (st->mb->pos == 0)
+		if (st->mb->pos == 0) {
 			*intra = true;
+		}
 	}
 
 #if 0
@@ -425,5 +525,32 @@ int decode_h263(struct viddec_state *st, struct vidframe *frame,
 		st->mb->buf[st->mb->end - 1] |= sbyte;
 	}
 
-	return ffdecode(st, frame, marker, src);
+	err = mbuf_write_mem(st->mb, mbuf_buf(src),
+			     mbuf_get_left(src));
+	if (err)
+		goto out;
+
+	if (!marker) {
+
+		if (st->mb->end > DECODE_MAXSZ) {
+			warning("avcodec: decode buffer size exceeded\n");
+			err = ENOMEM;
+			goto out;
+		}
+
+		return 0;
+	}
+
+	if (!hdr.i) {
+		++st->stats.n_key;
+	}
+
+	err = ffdecode(st, frame);
+	if (err)
+		goto out;
+
+ out:
+	mbuf_rewind(st->mb);
+
+	return err;
 }