matroskaenc: Mux clusters better

Start them on keyframes when reasonable, and delay writing audio packets
to help ensure that there's audio samples available for the first frame in
clusters.

Patch by James Zern <jzern at google>

Originally committed as revision 23473 to svn://svn.ffmpeg.org/ffmpeg/trunk

4 files changed, 83 insertions, 22 deletions

diff --git a/libavformat/matroskaenc.c b/libavformat/matroskaenc.c
index 01e6e79f3f..2b2a4650da 100644
--- a/libavformat/matroskaenc.c
+++ b/libavformat/matroskaenc.c
@@ -81,6 +81,8 @@ typedef struct MatroskaMuxContext {
     mkv_track       *tracks;
 
     struct AVMD5    *md5_ctx;
+    unsigned int    audio_buffer_size;
+    AVPacket        cur_audio_pkt;
 } MatroskaMuxContext;
 
 
@@ -746,6 +748,10 @@ static int mkv_write_header(AVFormatContext *s)
     if (mkv->cues == NULL)
         return AVERROR(ENOMEM);
 
+    av_init_packet(&mkv->cur_audio_pkt);
+    mkv->cur_audio_pkt.size = 0;
+    mkv->audio_buffer_size  = 0;
+
     put_flush_packet(pb);
     return 0;
 }
@@ -861,7 +867,7 @@ static void mkv_flush_dynbuf(AVFormatContext *s)
     mkv->dyn_bc = NULL;
 }
 
-static int mkv_write_packet(AVFormatContext *s, AVPacket *pkt)
+static int mkv_write_packet_internal(AVFormatContext *s, AVPacket *pkt)
 {
     MatroskaMuxContext *mkv = s->priv_data;
     ByteIOContext *pb = s->pb;
@@ -910,9 +916,38 @@ static int mkv_write_packet(AVFormatContext *s, AVPacket *pkt)
         if (ret < 0) return ret;
     }
 
-    // start a new cluster every 5 MB or 5 sec, or 32k / 1 sec for streaming
-    if ((url_is_streamed(s->pb) && (url_ftell(pb) > 32*1024 || ts > mkv->cluster_pts + 1000))
-        ||  url_ftell(pb) > mkv->cluster_pos + 5*1024*1024 || ts > mkv->cluster_pts + 5000) {
+    mkv->duration = FFMAX(mkv->duration, ts + duration);
+    return 0;
+}
+
+static int mkv_copy_packet(MatroskaMuxContext *mkv, const AVPacket *pkt)
+{
+    uint8_t *data           = mkv->cur_audio_pkt.data;
+    mkv->cur_audio_pkt      = *pkt;
+    mkv->cur_audio_pkt.data = av_fast_realloc(data, &mkv->audio_buffer_size, pkt->size);
+    if (!mkv->cur_audio_pkt.data)
+        return AVERROR(ENOMEM);
+
+    memcpy(mkv->cur_audio_pkt.data, pkt->data, pkt->size);
+    mkv->cur_audio_pkt.size = pkt->size;
+    return 0;
+}
+
+static int mkv_write_packet(AVFormatContext *s, AVPacket *pkt)
+{
+    MatroskaMuxContext *mkv = s->priv_data;
+    ByteIOContext *pb = url_is_streamed(s->pb) ? mkv->dyn_bc : s->pb;
+    AVCodecContext *codec = s->streams[pkt->stream_index]->codec;
+    int ret, keyframe = !!(pkt->flags & AV_PKT_FLAG_KEY);
+    int64_t ts = mkv->tracks[pkt->stream_index].write_dts ? pkt->dts : pkt->pts;
+    int cluster_size = url_ftell(pb) - (url_is_streamed(s->pb) ? 0 : mkv->cluster_pos);
+
+    // start a new cluster every 5 MB or 5 sec, or 32k / 1 sec for streaming or
+    // after 4k and on a keyframe
+    if (mkv->cluster_pos &&
+        ((url_is_streamed(s->pb) && (cluster_size > 32*1024 || ts > mkv->cluster_pts + 1000))
+         ||                      cluster_size > 5*1024*1024 || ts > mkv->cluster_pts + 5000
+         || (codec->codec_type == AVMEDIA_TYPE_VIDEO && keyframe && cluster_size > 4*1024))) {
         av_log(s, AV_LOG_DEBUG, "Starting new cluster at offset %" PRIu64
                " bytes, pts %" PRIu64 "\n", url_ftell(pb), ts);
         end_ebml_master(pb, mkv->cluster);
@@ -921,8 +956,23 @@ static int mkv_write_packet(AVFormatContext *s, AVPacket *pkt)
             mkv_flush_dynbuf(s);
     }
 
-    mkv->duration = FFMAX(mkv->duration, ts + duration);
-    return 0;
+    // check if we have an audio packet cached
+    if (mkv->cur_audio_pkt.size > 0) {
+        ret = mkv_write_packet_internal(s, &mkv->cur_audio_pkt);
+        mkv->cur_audio_pkt.size = 0;
+        if (ret < 0) {
+            av_log(s, AV_LOG_ERROR, "Could not write cached audio packet ret:%d\n", ret);
+            return ret;
+        }
+    }
+
+    // buffer an audio packet to ensure the packet containing the video
+    // keyframe's timecode is contained in the same cluster for WebM
+    if (codec->codec_type == AVMEDIA_TYPE_AUDIO)
+        ret = mkv_copy_packet(mkv, pkt);
+    else
+        ret = mkv_write_packet_internal(s, pkt);
+    return ret;
 }
 
 static int mkv_write_trailer(AVFormatContext *s)
@@ -932,6 +982,16 @@ static int mkv_write_trailer(AVFormatContext *s)
     int64_t currentpos, second_seekhead, cuespos;
     int ret;
 
+    // check if we have an audio packet cached
+    if (mkv->cur_audio_pkt.size > 0) {
+        ret = mkv_write_packet_internal(s, &mkv->cur_audio_pkt);
+        mkv->cur_audio_pkt.size = 0;
+        if (ret < 0) {
+            av_log(s, AV_LOG_ERROR, "Could not write cached audio packet ret:%d\n", ret);
+            return ret;
+        }
+    }
+
     if (mkv->dyn_bc) {
         end_ebml_master(mkv->dyn_bc, mkv->cluster);
         mkv_flush_dynbuf(s);
@@ -970,6 +1030,7 @@ static int mkv_write_trailer(AVFormatContext *s)
     end_ebml_master(pb, mkv->segment);
     av_free(mkv->md5_ctx);
     av_free(mkv->tracks);
+    av_destruct_packet(&mkv->cur_audio_pkt);
     put_flush_packet(pb);
     return 0;
 }
diff --git a/tests/ref/acodec/pcm b/tests/ref/acodec/pcm
index 8e344087eb..7f11b2913a 100644
--- a/tests/ref/acodec/pcm
+++ b/tests/ref/acodec/pcm
@@ -22,11 +22,11 @@ stddev:    0.00 PSNR:999.99 bytes:  1058444/  1058444
 1058444 ./tests/data/acodec/pcm_s16le.wav
 95e54b261530a1bcf6de6fe3b21dc5f6 *./tests/data/pcm.acodec.out.wav
 stddev:    0.00 PSNR:999.99 bytes:  1058444/  1058444
-43440f8078f2383c0dd618ad606f6830 *./tests/data/acodec/pcm_s16be.mkv
+fca887459876a7fdc4b9580691acd835 *./tests/data/acodec/pcm_s16be.mkv
  1060693 ./tests/data/acodec/pcm_s16be.mkv
 95e54b261530a1bcf6de6fe3b21dc5f6 *./tests/data/pcm.acodec.out.wav
 stddev:    0.00 PSNR:999.99 bytes:  1058444/  1058444
-1e63166f1672b7eb00877c697c06743c *./tests/data/acodec/pcm_s16le.mkv
+2817e2a028a9117d4809604fa0f81a80 *./tests/data/acodec/pcm_s16le.mkv
  1060693 ./tests/data/acodec/pcm_s16le.mkv
 95e54b261530a1bcf6de6fe3b21dc5f6 *./tests/data/pcm.acodec.out.wav
 stddev:    0.00 PSNR:999.99 bytes:  1058444/  1058444
diff --git a/tests/ref/lavf/mkv b/tests/ref/lavf/mkv
index 213d38b758..23e1cbe827 100644
--- a/tests/ref/lavf/mkv
+++ b/tests/ref/lavf/mkv
@@ -1,3 +1,3 @@
-cf230e089e737ed9449fbfb3cb87da92 *./tests/data/lavf/lavf.mkv
-  320471 ./tests/data/lavf/lavf.mkv
+5e906a7003f893cf6481f186c92dc16c *./tests/data/lavf/lavf.mkv
+  320537 ./tests/data/lavf/lavf.mkv
 ./tests/data/lavf/lavf.mkv CRC=0x2a83e6b0
diff --git a/tests/ref/seek/lavf.mkv.ref b/tests/ref/seek/lavf.mkv.ref
index 5af4488d2c..9f447b3257 100644
--- a/tests/ref/seek/lavf.mkv.ref
+++ b/tests/ref/seek/lavf.mkv.ref
@@ -2,35 +2,35 @@ ret: 0         st: 0 flags:1 dts: 0.000000 pts: 0.000000 pos:    513 size: 27837
 ret: 0         st:-1 flags:0  ts:-1.000000
 ret: 0         st: 0 flags:1 dts: 0.000000 pts: 0.000000 pos:    513 size: 27837
 ret: 0         st:-1 flags:1  ts: 1.894167
-ret: 0         st: 0 flags:1 dts: 0.960000 pts: 0.960000 pos: 291902 size: 27834
+ret: 0         st: 0 flags:1 dts: 0.960000 pts: 0.960000 pos: 291934 size: 27834
 ret: 0         st: 0 flags:0  ts: 0.788000
-ret: 0         st: 0 flags:1 dts: 0.960000 pts: 0.960000 pos: 291902 size: 27834
+ret: 0         st: 0 flags:1 dts: 0.960000 pts: 0.960000 pos: 291934 size: 27834
 ret: 0         st: 0 flags:1  ts:-0.317000
 ret: 0         st: 0 flags:1 dts: 0.000000 pts: 0.000000 pos:    513 size: 27837
 ret: 0         st: 1 flags:0  ts: 2.577000
 ret:-EOF
 ret: 0         st: 1 flags:1  ts: 1.471000
-ret: 0         st: 1 flags:1 dts: 1.019000 pts: 1.019000 pos: 320175 size:   209
+ret: 0         st: 1 flags:1 dts: 1.019000 pts: 1.019000 pos: 320207 size:   209
 ret: 0         st:-1 flags:0  ts: 0.365002
-ret: 0         st: 0 flags:1 dts: 0.480000 pts: 0.480000 pos: 146687 size: 27925
+ret: 0         st: 0 flags:1 dts: 0.480000 pts: 0.480000 pos: 146703 size: 27925
 ret: 0         st:-1 flags:1  ts:-0.740831
 ret: 0         st: 0 flags:1 dts: 0.000000 pts: 0.000000 pos:    513 size: 27837
 ret: 0         st: 0 flags:0  ts: 2.153000
 ret:-EOF
 ret: 0         st: 0 flags:1  ts: 1.048000
-ret: 0         st: 0 flags:1 dts: 0.960000 pts: 0.960000 pos: 291902 size: 27834
+ret: 0         st: 0 flags:1 dts: 0.960000 pts: 0.960000 pos: 291934 size: 27834
 ret: 0         st: 1 flags:0  ts:-0.058000
 ret: 0         st: 0 flags:1 dts: 0.000000 pts: 0.000000 pos:    513 size: 27837
 ret: 0         st: 1 flags:1  ts: 2.836000
-ret: 0         st: 1 flags:1 dts: 1.019000 pts: 1.019000 pos: 320175 size:   209
+ret: 0         st: 1 flags:1 dts: 1.019000 pts: 1.019000 pos: 320207 size:   209
 ret: 0         st:-1 flags:0  ts: 1.730004
 ret:-EOF
 ret: 0         st:-1 flags:1  ts: 0.624171
-ret: 0         st: 0 flags:1 dts: 0.480000 pts: 0.480000 pos: 146687 size: 27925
+ret: 0         st: 0 flags:1 dts: 0.480000 pts: 0.480000 pos: 146703 size: 27925
 ret: 0         st: 0 flags:0  ts:-0.482000
 ret: 0         st: 0 flags:1 dts: 0.000000 pts: 0.000000 pos:    513 size: 27837
 ret: 0         st: 0 flags:1  ts: 2.413000
-ret: 0         st: 0 flags:1 dts: 0.960000 pts: 0.960000 pos: 291902 size: 27834
+ret: 0         st: 0 flags:1 dts: 0.960000 pts: 0.960000 pos: 291934 size: 27834
 ret: 0         st: 1 flags:0  ts: 1.307000
 ret:-EOF
 ret: 0         st: 1 flags:1  ts: 0.201000
@@ -38,16 +38,16 @@ ret: 0         st: 1 flags:1 dts: 0.183000 pts: 0.183000 pos:  72083 size:   209
 ret: 0         st:-1 flags:0  ts:-0.904994
 ret: 0         st: 0 flags:1 dts: 0.000000 pts: 0.000000 pos:    513 size: 27837
 ret: 0         st:-1 flags:1  ts: 1.989173
-ret: 0         st: 0 flags:1 dts: 0.960000 pts: 0.960000 pos: 291902 size: 27834
+ret: 0         st: 0 flags:1 dts: 0.960000 pts: 0.960000 pos: 291934 size: 27834
 ret: 0         st: 0 flags:0  ts: 0.883000
-ret: 0         st: 0 flags:1 dts: 0.960000 pts: 0.960000 pos: 291902 size: 27834
+ret: 0         st: 0 flags:1 dts: 0.960000 pts: 0.960000 pos: 291934 size: 27834
 ret: 0         st: 0 flags:1  ts:-0.222000
 ret: 0         st: 0 flags:1 dts: 0.000000 pts: 0.000000 pos:    513 size: 27837
 ret: 0         st: 1 flags:0  ts: 2.672000
 ret:-EOF
 ret: 0         st: 1 flags:1  ts: 1.566000
-ret: 0         st: 1 flags:1 dts: 1.019000 pts: 1.019000 pos: 320175 size:   209
+ret: 0         st: 1 flags:1 dts: 1.019000 pts: 1.019000 pos: 320207 size:   209
 ret: 0         st:-1 flags:0  ts: 0.460008
-ret: 0         st: 0 flags:1 dts: 0.480000 pts: 0.480000 pos: 146687 size: 27925
+ret: 0         st: 0 flags:1 dts: 0.480000 pts: 0.480000 pos: 146703 size: 27925
 ret: 0         st:-1 flags:1  ts:-0.645825
 ret: 0         st: 0 flags:1 dts: 0.000000 pts: 0.000000 pos:    513 size: 27837