From eaeba16a56e849ef892970facea935c25b44fc41 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tim-Philipp=20M=C3=BCller?= <tim.muller@collabora.co.uk>
Date: Fri, 16 Sep 2011 19:32:38 +0100
Subject: [PATCH] mpegpsmux: add "aggregate-gops" property to aggregate GOPs
 into buffer lists

Useful in connection with multifilesink, for example, to make sure
the new file always starts with a keyframe.
---
 gst/mpegpsmux/mpegpsmux.c | 77 +++++++++++++++++++++++++++++++++++++--
 gst/mpegpsmux/mpegpsmux.h |  3 ++
 2 files changed, 77 insertions(+), 3 deletions(-)

diff --git a/gst/mpegpsmux/mpegpsmux.c b/gst/mpegpsmux/mpegpsmux.c
index 19e23d93dd..6d07d023c0 100644
--- a/gst/mpegpsmux/mpegpsmux.c
+++ b/gst/mpegpsmux/mpegpsmux.c
@@ -54,9 +54,11 @@ GST_DEBUG_CATEGORY (mpegpsmux_debug);
 
 enum
 {
-  ARG_0
+  PROP_AGGREGATE_GOPS = 1
 };
 
+#define DEFAULT_AGGREGATE_GOPS FALSE
+
 static GstStaticPadTemplate mpegpsmux_sink_factory =
     GST_STATIC_PAD_TEMPLATE ("sink_%d",
     GST_PAD_SINK,
@@ -135,6 +137,10 @@ mpegpsmux_class_init (MpegPsMuxClass * klass)
   gstelement_class->release_pad = mpegpsmux_release_pad;
   gstelement_class->change_state = mpegpsmux_change_state;
 
+  g_object_class_install_property (gobject_class, PROP_AGGREGATE_GOPS,
+      g_param_spec_boolean ("aggregate-gops", "Aggregate GOPs",
+          "Whether to aggregate GOPs and push them out as buffer lists",
+          DEFAULT_AGGREGATE_GOPS, G_PARAM_READWRITE | G_PARAM_STATIC_STRINGS));
 }
 
 static void
@@ -171,6 +177,11 @@ mpegpsmux_dispose (GObject * object)
     mux->psmux = NULL;
   }
 
+  if (mux->gop_list != NULL) {
+    gst_buffer_list_unref (mux->gop_list);
+    mux->gop_list = NULL;
+  }
+
   GST_CALL_PARENT (G_OBJECT_CLASS, dispose, (object));
 }
 
@@ -178,9 +189,12 @@ static void
 gst_mpegpsmux_set_property (GObject * object, guint prop_id,
     const GValue * value, GParamSpec * pspec)
 {
-/*  MpegPsMux *mux = GST_MPEG_PSMUX (object); */
+  MpegPsMux *mux = GST_MPEG_PSMUX (object);
 
   switch (prop_id) {
+    case PROP_AGGREGATE_GOPS:
+      mux->aggregate_gops = g_value_get_boolean (value);
+      break;
     default:
       G_OBJECT_WARN_INVALID_PROPERTY_ID (object, prop_id, pspec);
       break;
@@ -191,9 +205,12 @@ static void
 gst_mpegpsmux_get_property (GObject * object, guint prop_id,
     GValue * value, GParamSpec * pspec)
 {
-/*  MpegPsMux *mux = GST_MPEG_PSMUX (object); */
+  MpegPsMux *mux = GST_MPEG_PSMUX (object);
 
   switch (prop_id) {
+    case PROP_AGGREGATE_GOPS:
+      g_value_set_boolean (value, mux->aggregate_gops);
+      break;
     default:
       G_OBJECT_WARN_INVALID_PROPERTY_ID (object, prop_id, pspec);
       break;
@@ -456,6 +473,20 @@ mpegpsmux_choose_best_stream (MpegPsMux * mux)
   return best;
 }
 
+static GstFlowReturn
+mpegpsmux_push_gop_list (MpegPsMux * mux)
+{
+  GstFlowReturn flow;
+
+  g_assert (mux->gop_list != NULL);
+
+  GST_DEBUG_OBJECT (mux, "Sending pending GOP of %u buffers",
+      gst_buffer_list_n_groups (mux->gop_list));
+  flow = gst_pad_push_list (mux->srcpad, mux->gop_list);
+  mux->gop_list = NULL;
+  return flow;
+}
+
 static GstFlowReturn
 mpegpsmux_collected (GstCollectPads * pads, MpegPsMux * mux)
 {
@@ -506,8 +537,16 @@ mpegpsmux_collected (GstCollectPads * pads, MpegPsMux * mux)
           G_GINT64_FORMAT, GST_TIME_ARGS (best->cur_ts), pts);
     }
 
+    /* start of new GOP? */
     keyunit = !GST_BUFFER_FLAG_IS_SET (buf, GST_BUFFER_FLAG_DELTA_UNIT);
 
+    if (keyunit && best->stream_id == mux->video_stream_id
+        && mux->gop_list != NULL) {
+      ret = mpegpsmux_push_gop_list (mux);
+      if (ret != GST_FLOW_OK)
+        goto done;
+    }
+
     /* give the buffer to libpsmux for processing */
     psmux_stream_add_data (best->stream, GST_BUFFER_DATA (buf),
         GST_BUFFER_SIZE (buf), buf, pts, -1, keyunit);
@@ -526,12 +565,17 @@ mpegpsmux_collected (GstCollectPads * pads, MpegPsMux * mux)
   } else {
     /* FIXME: Drain all remaining streams */
     /* At EOS */
+    if (mux->gop_list != NULL)
+      mpegpsmux_push_gop_list (mux);
+
     if (psmux_write_end_code (mux->psmux)) {
       GST_WARNING_OBJECT (mux, "Writing MPEG PS Program end code failed.");
     }
     gst_pad_push_event (mux->srcpad, gst_event_new_eos ());
   }
 
+done:
+
   return ret;
 new_seg_fail:
   return GST_FLOW_ERROR;
@@ -606,6 +650,26 @@ mpegpsmux_release_pad (GstElement * element, GstPad * pad)
   gst_collect_pads_remove_pad (mux->collect, pad);
 }
 
+static void
+add_buffer_to_goplist (MpegPsMux * mux, GstBuffer * buf)
+{
+  GstBufferListIterator *it;
+
+  if (mux->gop_list == NULL)
+    mux->gop_list = gst_buffer_list_new ();
+
+  it = gst_buffer_list_iterate (mux->gop_list);
+
+  /* move iterator to end */
+  while (gst_buffer_list_iterator_next_group (it)) {
+    /* .. */
+  }
+
+  gst_buffer_list_iterator_add_group (it);
+  gst_buffer_list_iterator_add (it, buf);
+  gst_buffer_list_iterator_free (it);
+}
+
 static gboolean
 new_packet_cb (guint8 * data, guint len, void *user_data)
 {
@@ -626,7 +690,14 @@ new_packet_cb (guint8 * data, guint len, void *user_data)
 
   memcpy (GST_BUFFER_DATA (buf), data, len);
   GST_BUFFER_TIMESTAMP (buf) = mux->last_ts;
+
+  if (mux->aggregate_gops) {
+    add_buffer_to_goplist (mux, buf);
+    return TRUE;
+  }
+
   ret = gst_pad_push (mux->srcpad, buf);
+
   if (G_UNLIKELY (ret != GST_FLOW_OK)) {
     mux->last_flow_ret = ret;
     return FALSE;
diff --git a/gst/mpegpsmux/mpegpsmux.h b/gst/mpegpsmux/mpegpsmux.h
index 125f8f6742..6d994f456f 100644
--- a/gst/mpegpsmux/mpegpsmux.h
+++ b/gst/mpegpsmux/mpegpsmux.h
@@ -78,6 +78,9 @@ struct MpegPsMux {
   GstFlowReturn last_flow_ret;
   
   GstClockTime last_ts;
+
+  GstBufferList *gop_list;
+  gboolean       aggregate_gops;
 };
 
 struct MpegPsMuxClass  {