nvenc: Ensure drain all frames on finish

To drain all queued encoding items, encoder should gracefully
wait the encoding thread without stealing queued items.
Otherwise, some input frames can be dropped.

1 files changed, 20 insertions, 18 deletions

diff --git a/sys/nvenc/gstnvbaseenc.c b/sys/nvenc/gstnvbaseenc.c
index 51e75eba8..3d991bafb 100644
--- a/sys/nvenc/gstnvbaseenc.c
+++ b/sys/nvenc/gstnvbaseenc.c
@@ -210,7 +210,8 @@ static void gst_nv_base_enc_get_property (GObject * object, guint prop_id,
 static void gst_nv_base_enc_finalize (GObject * obj);
 static GstCaps *gst_nv_base_enc_getcaps (GstVideoEncoder * enc,
     GstCaps * filter);
-static gboolean gst_nv_base_enc_stop_bitstream_thread (GstNvBaseEnc * nvenc);
+static gboolean gst_nv_base_enc_stop_bitstream_thread (GstNvBaseEnc * nvenc,
+    gboolean force);
 
 static void
 gst_nv_base_enc_class_init (GstNvBaseEncClass * klass)
@@ -508,7 +509,7 @@ gst_nv_base_enc_stop (GstVideoEncoder * enc)
 {
   GstNvBaseEnc *nvenc = GST_NV_BASE_ENC (enc);
 
-  gst_nv_base_enc_stop_bitstream_thread (nvenc);
+  gst_nv_base_enc_stop_bitstream_thread (nvenc, TRUE);
 
   gst_nv_base_enc_free_buffers (nvenc);
 
@@ -883,24 +884,27 @@ gst_nv_base_enc_start_bitstream_thread (GstNvBaseEnc * nvenc)
 }
 
 static gboolean
-gst_nv_base_enc_stop_bitstream_thread (GstNvBaseEnc * nvenc)
+gst_nv_base_enc_stop_bitstream_thread (GstNvBaseEnc * nvenc, gboolean force)
 {
   gpointer out_buf;
 
   if (nvenc->bitstream_thread == NULL)
     return TRUE;
 
-  /* FIXME */
-  GST_FIXME_OBJECT (nvenc, "stop bitstream reading thread properly");
-  g_async_queue_lock (nvenc->bitstream_queue);
-  g_async_queue_lock (nvenc->bitstream_pool);
-  while ((out_buf = g_async_queue_try_pop_unlocked (nvenc->bitstream_queue))) {
-    GST_INFO_OBJECT (nvenc, "stole bitstream buffer %p from queue", out_buf);
-    g_async_queue_push_unlocked (nvenc->bitstream_pool, out_buf);
+  if (force) {
+    g_async_queue_lock (nvenc->bitstream_queue);
+    g_async_queue_lock (nvenc->bitstream_pool);
+    while ((out_buf = g_async_queue_try_pop_unlocked (nvenc->bitstream_queue))) {
+      GST_INFO_OBJECT (nvenc, "stole bitstream buffer %p from queue", out_buf);
+      g_async_queue_push_unlocked (nvenc->bitstream_pool, out_buf);
+    }
+    g_async_queue_push_unlocked (nvenc->bitstream_queue, SHUTDOWN_COOKIE);
+    g_async_queue_unlock (nvenc->bitstream_pool);
+    g_async_queue_unlock (nvenc->bitstream_queue);
+  } else {
+    /* wait for encoder to drain the remaining buffers */
+    g_async_queue_push (nvenc->bitstream_queue, SHUTDOWN_COOKIE);
   }
-  g_async_queue_push_unlocked (nvenc->bitstream_queue, SHUTDOWN_COOKIE);
-  g_async_queue_unlock (nvenc->bitstream_pool);
-  g_async_queue_unlock (nvenc->bitstream_queue);
 
   /* temporary unlock, so other thread can find and push frame */
   GST_VIDEO_ENCODER_STREAM_UNLOCK (nvenc);
@@ -1902,12 +1906,10 @@ gst_nv_base_enc_finish (GstVideoEncoder * enc)
 {
   GstNvBaseEnc *nvenc = GST_NV_BASE_ENC (enc);
 
-  GST_FIXME_OBJECT (enc, "implement finish");
-
-  gst_nv_base_enc_drain_encoder (nvenc);
+  if (!gst_nv_base_enc_drain_encoder (nvenc))
+    return GST_FLOW_ERROR;
 
-  /* wait for encoder to output the remaining buffers */
-  gst_nv_base_enc_stop_bitstream_thread (nvenc);
+  gst_nv_base_enc_stop_bitstream_thread (nvenc, FALSE);
 
   return GST_FLOW_OK;
 }