From b1c14b035742711a60a82c8d56c7bc0751b2aeb8 Mon Sep 17 00:00:00 2001
From: Seungha Yang <seungha@centricular.com>
Date: Mon, 20 Feb 2023 02:10:07 +0900
Subject: [PATCH] nvencoder: Fix b-frame encoding on Linux

On Windows, Win32 event handle is used to wait for encoded output,
but it's not available on Linux. We should delay bitstream locking
if encoder returns "need-more-input"

Part-of: <https://gitlab.freedesktop.org/gstreamer/gstreamer/-/merge_requests/4004>
---
 .../sys/nvcodec/gstnvencobject.cpp            | 36 +++++++++++++++++--
 .../sys/nvcodec/gstnvencobject.h              |  2 ++
 2 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/subprojects/gst-plugins-bad/sys/nvcodec/gstnvencobject.cpp b/subprojects/gst-plugins-bad/sys/nvcodec/gstnvencobject.cpp
index 51187b8e13..6463e4afba 100644
--- a/subprojects/gst-plugins-bad/sys/nvcodec/gstnvencobject.cpp
+++ b/subprojects/gst-plugins-bad/sys/nvcodec/gstnvencobject.cpp
@@ -367,6 +367,7 @@ GstNvEncObject::InitSession (NV_ENC_INITIALIZE_PARAMS * params,
   }
 
   task_size_ = task_size;
+  lookahead_ = params->encodeConfig->rcParams.lookaheadDepth;
   initialized_ = true;
 
 out:
@@ -490,8 +491,34 @@ GstNvEncObject::Encode (GstVideoCodecFrame * codec_frame,
       active_resource_queue_.insert (task->resource);
   }
 
-  task_queue_.push (task);
-  cond_.notify_all ();
+  /* On Windows and if async encoding is enabled, output thread will wait
+   * for completion event. But on Linux, async encoding is not supported.
+   * So, we should wait for NV_ENC_SUCCESS in case of sync mode
+   * (it would introduce latency though).
+   * Otherwise nvEncLockBitstream() will return error */
+  if (params.completionEvent) {
+    /* Windows only path */
+    task_queue_.push (task);
+    cond_.notify_all ();
+  } else {
+    pending_task_queue_.push (task);
+    if (status == NV_ENC_SUCCESS) {
+      bool notify = false;
+
+      /* XXX: nvEncLockBitstream() will return NV_ENC_ERR_INVALID_PARAM
+       * if lookahead is enabled. See also
+       * https://gitlab.freedesktop.org/gstreamer/gst-plugins-bad/-/merge_requests/494
+       */
+      while (pending_task_queue_.size() > lookahead_) {
+        notify = true;
+        task_queue_.push (pending_task_queue_.front ());
+        pending_task_queue_.pop ();
+      }
+
+      if (notify)
+        cond_.notify_all ();
+    }
+  }
 
   return NV_ENC_SUCCESS;
 }
@@ -531,6 +558,11 @@ GstNvEncObject::Drain (GstNvEncTask * task)
     break;
   } while (true);
 
+  while (!pending_task_queue_.empty ()) {
+    task_queue_.push (pending_task_queue_.front ());
+    pending_task_queue_.pop ();
+  }
+
   task_queue_.push (task);
   cond_.notify_all ();
 
diff --git a/subprojects/gst-plugins-bad/sys/nvcodec/gstnvencobject.h b/subprojects/gst-plugins-bad/sys/nvcodec/gstnvencobject.h
index 866a7ced80..0902450950 100644
--- a/subprojects/gst-plugins-bad/sys/nvcodec/gstnvencobject.h
+++ b/subprojects/gst-plugins-bad/sys/nvcodec/gstnvencobject.h
@@ -233,6 +233,7 @@ private:
   /* list of GstNvEncResource in task_queue */
   std::set <GstNvEncResource *> active_resource_queue_;
   std::queue <GstNvEncTask *> task_queue_;
+  std::queue <GstNvEncTask *> pending_task_queue_;
   std::queue <GstNvEncTask *> empty_task_queue_;
   gint64 user_token_;
   GstCudaContext *context_ = nullptr;
@@ -245,6 +246,7 @@ private:
   bool initialized_ = false;
   bool flushing_ = false;
   guint task_size_ = 0;
+  guint lookahead_ = 0;
 
   NV_ENC_DEVICE_TYPE device_type_ = NV_ENC_DEVICE_TYPE_CUDA;
   NV_ENC_BUFFER_FORMAT buffer_format_ = NV_ENC_BUFFER_FORMAT_UNDEFINED;