From 871fa29639220a98952751e545d7555215874d37 Mon Sep 17 00:00:00 2001
From: Thibault Saunier <tsaunier@igalia.com>
Date: Sat, 4 Jul 2020 12:29:06 -0400
Subject: [PATCH] streamsplitter/combiner: Drain encoder before switching
 branch

Otherwise we miht have frames queued in the encoder from the old branch
that do not get encoded/muxed when they should.

The implementation is a bit 'weird' but the rational and solution
is documented in the code.

Part-of: <https://gitlab.freedesktop.org/gstreamer/gst-plugins-base/-/merge_requests/751>
---
 gst/encoding/gststreamcombiner.c | 32 ++++++++++++++++++++++++++++++++
 gst/encoding/gststreamcombiner.h |  5 +++++
 gst/encoding/gststreamsplitter.c | 23 ++++++++++++++++++++++-
 3 files changed, 59 insertions(+), 1 deletion(-)

diff --git a/gst/encoding/gststreamcombiner.c b/gst/encoding/gststreamcombiner.c
index 22fff918f6..349df19e5e 100644
--- a/gst/encoding/gststreamcombiner.c
+++ b/gst/encoding/gststreamcombiner.c
@@ -136,8 +136,40 @@ gst_stream_combiner_sink_event (GstPad * pad, GstObject * parent,
   GST_DEBUG_OBJECT (pad, "Got event %s", GST_EVENT_TYPE_NAME (event));
 
   switch (GST_EVENT_TYPE (event)) {
+    case GST_EVENT_CUSTOM_DOWNSTREAM:
+      STREAMS_LOCK (stream_combiner);
+      if (gst_structure_has_name (gst_event_get_structure (event),
+              "start-draining-encoder")) {
+        GST_INFO_OBJECT (pad, "Starting to drain the encoder");
+        stream_combiner->draining_encoder = TRUE;
+      }
+      STREAMS_UNLOCK (stream_combiner);
+      break;
+    case GST_EVENT_FLUSH_START:
+      STREAMS_LOCK (stream_combiner);
+      if (stream_combiner->draining_encoder) {
+        GST_INFO_OBJECT (pad, "Discarding FLUSH_START as draining encoder");
+        gst_clear_event (&event);
+      }
+      STREAMS_UNLOCK (stream_combiner);
+      break;
+    case GST_EVENT_FLUSH_STOP:
+      STREAMS_LOCK (stream_combiner);
+      if (stream_combiner->draining_encoder) {
+        gst_clear_event (&event);
+        GST_INFO_OBJECT (stream_combiner, "Done draining the encoder.");
+      }
+      stream_combiner->draining_encoder = FALSE;
+      STREAMS_UNLOCK (stream_combiner);
+      break;
     case GST_EVENT_EOS:
       STREAMS_LOCK (stream_combiner);
+      if (stream_combiner->draining_encoder) {
+        STREAMS_UNLOCK (stream_combiner);
+        GST_INFO_OBJECT (stream_combiner, "Discarding EOS as draining encoder");
+        gst_clear_event (&event);
+        break;
+      }
       combiner_pad->is_eos = TRUE;
       if (!_all_sink_pads_eos (stream_combiner)) {
         gst_event_unref (event);
diff --git a/gst/encoding/gststreamcombiner.h b/gst/encoding/gststreamcombiner.h
index ba22a13ee1..f466731e4e 100644
--- a/gst/encoding/gststreamcombiner.h
+++ b/gst/encoding/gststreamcombiner.h
@@ -41,6 +41,11 @@ struct _GstStreamCombiner {
   GstPad *current;
   GList *sinkpads;
   guint32 cookie;
+  gboolean draining_encoder; /* TRUE when streamspliter informed us that it is
+                              * draining the encoder, meaning that we are
+                              * expecting to receive and discard an EOS, a flush
+                              * start, and then a flush_stop which implies the
+                              * draining is done. */
 
 };
 
diff --git a/gst/encoding/gststreamsplitter.c b/gst/encoding/gststreamsplitter.c
index f17904b3e4..4d22698074 100644
--- a/gst/encoding/gststreamsplitter.c
+++ b/gst/encoding/gststreamsplitter.c
@@ -415,6 +415,7 @@ gst_stream_splitter_sink_query (GstPad * pad, GstObject * parent,
 static gboolean
 gst_stream_splitter_sink_setcaps (GstPad * pad, GstCaps * caps)
 {
+  GstPad *prev = NULL;
   GstStreamSplitter *stream_splitter =
       (GstStreamSplitter *) GST_PAD_PARENT (pad);
   guint32 cookie;
@@ -433,6 +434,9 @@ resync:
   }
 
   res = FALSE;
+  prev =
+      stream_splitter->
+      current ? gst_object_ref (stream_splitter->current) : NULL;
   tmp = stream_splitter->srcpads;
   cookie = stream_splitter->cookie;
 
@@ -440,6 +444,7 @@ resync:
     GstPad *srcpad = (GstPad *) tmp->data;
     GstCaps *peercaps;
 
+    // GST_ERROR_OBJECT (srcpad, "Checking");
     STREAMS_UNLOCK (stream_splitter);
     peercaps = gst_pad_peer_query_caps (srcpad, NULL);
     if (peercaps) {
@@ -452,15 +457,31 @@ resync:
       goto resync;
 
     if (res) {
-      /* FIXME : we need to switch properly */
       GST_DEBUG_OBJECT (srcpad, "Setting caps on this pad was successful");
       stream_splitter->current = srcpad;
+      if (prev && prev != srcpad) {
+        /* When switching branches, we need to ensure that the branch it was
+         * pushing to before is drainned and the encoder pushes all its internal
+         * data, we do it by first letting downstream streamcombiner that we are
+         * working on it, and then send an EOS to ensure all the data is
+         * processed and cleanup the encoder with a flush start/flush stop
+         * dance. Those events are discarded in the stream combiner */
+        /* FIXME We should theoretically be able to use a DRAIN query here
+         * instead but the encoders do not react properly to them */
+        gst_pad_push_event (prev,
+            gst_event_new_custom (GST_EVENT_CUSTOM_DOWNSTREAM,
+                gst_structure_new_empty ("start-draining-encoder")));
+        gst_pad_push_event (prev, gst_event_new_eos ());
+        gst_pad_push_event (prev, gst_event_new_flush_start ());
+        gst_pad_push_event (prev, gst_event_new_flush_stop (FALSE));
+      }
       goto beach;
     }
     tmp = tmp->next;
   }
 
 beach:
+  gst_clear_object (&prev);
   STREAMS_UNLOCK (stream_splitter);
   return res;
 }