From 00310b6a2ce5a95aac03a8db3de161595800eb00 Mon Sep 17 00:00:00 2001
From: Kamil Trzcinski <ayufan@ayufan.eu>
Date: Thu, 14 Apr 2022 09:54:20 +0200
Subject: [PATCH] Improve remuxer

---
 ffmpeg/remuxer.c   | 46 +++++++++++++++++++++++++++++++++++++++-------
 ffmpeg/remuxer.h   | 11 ++++++++---
 http/http_ffmpeg.c |  2 +-
 3 files changed, 48 insertions(+), 11 deletions(-)

diff --git a/ffmpeg/remuxer.c b/ffmpeg/remuxer.c
index 178ab4c..cf6320d 100644
--- a/ffmpeg/remuxer.c
+++ b/ffmpeg/remuxer.c
@@ -3,18 +3,22 @@
 
 #ifdef USE_FFMPEG
 static AVRational time_base = {1, 1000LL * 1000LL};
-static int avio_ctx_buffer_size = 4096;
+static unsigned avio_ctx_buffer_size = 4096;
 
 static int ffmpeg_remuxer_init_avcontext(AVFormatContext **context, ffmpeg_remuxer_t *remuxer, int output, int (*packet)(void *opaque, uint8_t *buf, int buf_size))
 {
   uint8_t *buffer = NULL;
   AVIOContext *avio = NULL;
   int ret = -1;
+
+  unsigned buffer_size = MAX(
+    output ? remuxer->write_buffer_size : remuxer->read_buffer_size,
+    avio_ctx_buffer_size);
   
-  buffer = av_malloc(avio_ctx_buffer_size);
+  buffer = av_malloc(buffer_size);
   if (!buffer)
     return AVERROR(ENOMEM);
-  avio = avio_alloc_context(buffer, avio_ctx_buffer_size, output, remuxer->opaque, output ? NULL : packet, output ? packet : NULL, NULL);
+  avio = avio_alloc_context(buffer, buffer_size, output, remuxer->opaque, output ? NULL : packet, output ? packet : NULL, NULL);
   if (!avio)
     goto error;
   if (output && (ret = avformat_alloc_output_context2(context, NULL, remuxer->video_format, NULL)) < 0)
@@ -112,7 +116,7 @@ int ffmpeg_remuxer_open(ffmpeg_remuxer_t *remuxer)
   return 0;
 }
 
-int ffmpeg_remuxer_close(ffmpeg_remuxer_t *remuxer)
+void ffmpeg_remuxer_close(ffmpeg_remuxer_t *remuxer)
 {
   if (remuxer->output_context)
     av_write_trailer(remuxer->output_context);
@@ -123,11 +127,16 @@ int ffmpeg_remuxer_close(ffmpeg_remuxer_t *remuxer)
   av_dict_free(&remuxer->output_opts);
 }
 
-int ffmpeg_remuxer_feed(ffmpeg_remuxer_t *remuxer)
+int ffmpeg_remuxer_feed(ffmpeg_remuxer_t *remuxer, int nframes)
 {
   int ret = 0;
+  int frames = 0;
 
   while (ret >= 0) {
+    if (nframes > 0 && frames >= nframes) {
+      break;
+    }
+
     ret = av_read_frame(remuxer->input_context, remuxer->packet);
     if (ret == AVERROR_EOF) {
       ret = 0;
@@ -143,6 +152,9 @@ int ffmpeg_remuxer_feed(ffmpeg_remuxer_t *remuxer)
       continue;
     }
 
+    remuxer->frames++;
+    frames++;
+
     AVStream *in_stream = remuxer->input_context->streams[remuxer->packet->stream_index];
     remuxer->packet->stream_index = 0;
     AVStream *out_stream = remuxer->output_context->streams[remuxer->packet->stream_index];
@@ -167,6 +179,22 @@ int ffmpeg_remuxer_feed(ffmpeg_remuxer_t *remuxer)
     }
   }
 
+  if (ret >= 0) {
+    return frames;
+  }
+
+  return ret;
+}
+
+int ffmpeg_remuxer_flush(ffmpeg_remuxer_t *remuxer)
+{
+  int ret = av_write_frame(remuxer->output_context, NULL);
+  if (ret == AVERROR_EOF) {
+    LOG_DEBUG(remuxer, "av_write_frame (flush): EOF", ret);
+  } else {
+    LOG_DEBUG(remuxer, "av_write_frame (flush): %08x", ret);
+  }
+
   return ret;
 }
 #else
@@ -175,13 +203,17 @@ int ffmpeg_remuxer_open(ffmpeg_remuxer_t *remuxer)
   return -1;
 }
 
-int ffmpeg_remuxer_feed(ffmpeg_remuxer_t *remuxer)
+int ffmpeg_remuxer_feed(ffmpeg_remuxer_t *remuxer, int nframes)
 {
   return -1;
 }
 
-int ffmpeg_remuxer_close(ffmpeg_remuxer_t *remuxer)
+int ffmpeg_remuxer_flush(ffmpeg_remuxer_t *remuxer)
 {
   return -1;
 }
+
+void ffmpeg_remuxer_close(ffmpeg_remuxer_t *remuxer)
+{
+}
 #endif
diff --git a/ffmpeg/remuxer.h b/ffmpeg/remuxer.h
index e462c39..4a5c5a9 100644
--- a/ffmpeg/remuxer.h
+++ b/ffmpeg/remuxer.h
@@ -19,6 +19,11 @@ typedef struct ffmpeg_remuxer_s {
   void *opaque;
   ffmpeg_data_packet read_packet;
   ffmpeg_data_packet write_packet;
+  unsigned read_buffer_size;
+  unsigned write_buffer_size;
+
+  uint64_t start_time;
+  unsigned frames;
 
 #ifdef USE_FFMPEG
   AVIOContext *input_avio;
@@ -29,10 +34,10 @@ typedef struct ffmpeg_remuxer_s {
   AVPacket *packet;
   AVDictionary *output_opts;
   int video_stream;
-  uint64_t start_time;
 #endif
 } ffmpeg_remuxer_t;
 
 int ffmpeg_remuxer_open(ffmpeg_remuxer_t *remuxer);
-int ffmpeg_remuxer_feed(ffmpeg_remuxer_t *remuxer);
-int ffmpeg_remuxer_close(ffmpeg_remuxer_t *remuxer);
+int ffmpeg_remuxer_feed(ffmpeg_remuxer_t *remuxer, int nframes);
+int ffmpeg_remuxer_flush(ffmpeg_remuxer_t *remuxer);
+void ffmpeg_remuxer_close(ffmpeg_remuxer_t *remuxer);
diff --git a/http/http_ffmpeg.c b/http/http_ffmpeg.c
index c9facd0..42770d2 100644
--- a/http/http_ffmpeg.c
+++ b/http/http_ffmpeg.c
@@ -94,7 +94,7 @@ static int http_ffmpeg_video_buf_part(buffer_lock_t *buf_lock, buffer_t *buf, in
 
   if ((ret = ffmpeg_remuxer_open(status->remuxer)) < 0)
     goto error;
-  if ((ret = ffmpeg_remuxer_feed(status->remuxer)) < 0)
+  if ((ret = ffmpeg_remuxer_feed(status->remuxer, 0)) < 0)
     goto error;
 
   ret = 1;