From 28b5ce612f3a1f2fd4ed363fbf4be0ac06b1c16c Mon Sep 17 00:00:00 2001
From: Julian Rother <julianr@fsmpi.rwth-aachen.de>
Date: Sun, 26 Nov 2017 21:20:12 +0100
Subject: [PATCH] transcode: Restructured mainloop to fix flushing

---
 transcode.c | 96 +++++++++++++++++++++++++++++++++++------------------
 1 file changed, 64 insertions(+), 32 deletions(-)

diff --git a/transcode.c b/transcode.c
index cdeeac0..d985230 100644
--- a/transcode.c
+++ b/transcode.c
@@ -176,6 +176,54 @@ static void setup_output_stream(char *stream, AVStream *st, AVCodecContext *enc,
 	avcodec_parameters_from_context(st->codecpar, enc);
 }
 
+static void flush_encoder(AVFormatContext *mux, AVFilterContext *sink, AVCodecContext *enc, int idx)
+{
+	int err;
+	AVPacket pkt;
+	pkt.data = 0;
+	pkt.size = 0;
+	av_init_packet(&pkt);
+	while (!avcodec_receive_packet(enc, &pkt))
+	{
+		pkt.stream_index = idx;
+		av_packet_rescale_ts(&pkt, sink->inputs[0]->time_base,
+				mux->streams[idx]->time_base);
+		if (err = av_interleaved_write_frame(mux, &pkt))
+			job_failed("Could not write frame: %s", av_err2str(err));
+	}
+}
+
+static void flush_filtergraph(AVFormatContext *mux, AVFilterContext **sinks, AVCodecContext **encs)
+{
+	int i;
+	AVFrame *frame;
+	if (!(frame = av_frame_alloc()))
+		exit(99);
+	for (i = 0; i < mux->nb_streams; i ++)
+		while (av_buffersink_get_frame(sinks[i], frame) >= 0)
+		{
+			avcodec_send_frame(encs[i], frame);
+			av_frame_unref(frame);
+			flush_encoder(mux, sinks[i], encs[i], i);
+		}
+	av_frame_free(&frame);
+}
+
+static void filtergraph_send(AVFilterContext *src, AVCodecContext *dec, AVPacket *pkt)
+{
+	int err;
+	AVFrame *frame;
+	if (!dec || !src)
+		return;
+	if (!(frame = av_frame_alloc()))
+		exit(99);
+	avcodec_send_packet(dec, pkt);
+	while (!avcodec_receive_frame(dec, frame))
+		if ((err = av_buffersrc_add_frame(src, frame)) < 0)
+			job_failed("Could not insert frame into filter graph: %s", av_err2str(err));
+	av_frame_free(&frame);
+}
+
 int main(int argc, char *argv[])
 {
 	int err, i, progress, _progress;
@@ -186,7 +234,6 @@ int main(int argc, char *argv[])
 	AVFilterGraph *fg;
 	AVFilterInOut *inpads, *outpads;
 	AVPacket pkt;
-	AVFrame *frame;
 	AVStream *stream;
 	AVDictionary *opts;
 	if (argc != 5)
@@ -194,10 +241,9 @@ int main(int argc, char *argv[])
 	av_register_all();
 	avfilter_register_all();
 	init_avlogbuf();
-	memset(&pkt, 0, sizeof(pkt));
+	pkt.data = 0;
+	pkt.size = 0;
 	av_init_packet(&pkt);
-	if (!(frame = av_frame_alloc()))
-		return 99;
 	if (!(fg = avfilter_graph_alloc()))
 		return 99;
 
@@ -253,6 +299,7 @@ int main(int argc, char *argv[])
 
 	progress = 0;
 	ping_job(jobid, "running", 0);
+
 	while (!av_read_frame(demux, &pkt))
 	{
 		i = pkt.stream_index;
@@ -263,42 +310,27 @@ int main(int argc, char *argv[])
 					jescape(get_avlogbuf()));
 			progress = _progress;
 		}
-		if (!decs[i])
-			continue;
-		avcodec_send_packet(decs[i], &pkt);
-		while (!avcodec_receive_frame(decs[i], frame))
-			if ((err = av_buffersrc_add_frame(srcs[i], frame)) < 0)
-				job_failed("Could not insert frame into filter graph: %s", av_err2str(err));
-		for (i = 0; i < mux->nb_streams; i ++)
-			while (av_buffersink_get_frame(sinks[i], frame) >= 0)
-			{
-				avcodec_send_frame(encs[i], frame);
-				av_frame_unref(frame);
-				while (!avcodec_receive_packet(encs[i], &pkt))
-				{
-					pkt.stream_index = i;
-					av_packet_rescale_ts(&pkt, sinks[i]->inputs[0]->time_base,
-							mux->streams[i]->time_base);
-					if (err = av_interleaved_write_frame(mux, &pkt))
-						job_failed("Could not write frame: %s", av_err2str(err));
-				}
-			}
+		filtergraph_send(srcs[i], decs[i], &pkt);
+		flush_filtergraph(mux, sinks, encs);
 	}
 
-	avformat_close_input(&demux);
 	/* Flush */
+	for (i = 0; i < demux->nb_streams; i ++)
+	{
+		if (decs[i])
+			continue;
+		filtergraph_send(srcs[i], decs[i], 0);
+		!av_buffersrc_add_frame(srcs[i], 0);
+	}
+	flush_filtergraph(mux, sinks, encs);
 	for (i = 0; i < mux->nb_streams; i ++)
 	{
 		avcodec_send_frame(encs[i], 0);
-		while (!avcodec_receive_packet(encs[i], &pkt))
-		{
-			pkt.stream_index = i;
-			av_packet_rescale_ts(&pkt, sinks[i]->inputs[0]->time_base,
-					mux->streams[i]->time_base);
-			av_interleaved_write_frame(mux, &pkt);
-		}
+		flush_encoder(mux, sinks[i], encs[i], i);
 	}
 	av_interleaved_write_frame(mux, 0);
+
+	avformat_close_input(&demux);
 	if (err = av_write_trailer(mux))
 		job_failed("Error writing trailer to temporary file", av_err2str(err));
 	avio_closep(&mux->pb);
-- 
GitLab