From bd25c2738e18a57a7dfcd0085d093f6213b31ea1 Mon Sep 17 00:00:00 2001
From: Seungha Yang <seungha@centricular.com>
Date: Sat, 16 Sep 2023 23:26:27 +0900
Subject: [PATCH] nvdecoder: Copy output frame if needed

Even if decoder is negotiated with CUDA memory feature, if downstream
proposed no buffer pool, assume that the pool size is unknown.
And disable zero-copy if there's no more free output surface.
Or, in case of reverse playback, always copy frames.

Part-of: <https://gitlab.freedesktop.org/gstreamer/gstreamer/-/merge_requests/5338>
---
 .../sys/nvcodec/gstnvdecobject.cpp            | 12 +++++++++
 .../sys/nvcodec/gstnvdecobject.h              |  2 ++
 .../sys/nvcodec/gstnvdecoder.cpp              | 25 ++++++++++++++++++-
 3 files changed, 38 insertions(+), 1 deletion(-)

diff --git a/subprojects/gst-plugins-bad/sys/nvcodec/gstnvdecobject.cpp b/subprojects/gst-plugins-bad/sys/nvcodec/gstnvdecobject.cpp
index 4feb67f3f8..0c588934fc 100644
--- a/subprojects/gst-plugins-bad/sys/nvcodec/gstnvdecobject.cpp
+++ b/subprojects/gst-plugins-bad/sys/nvcodec/gstnvdecobject.cpp
@@ -547,6 +547,18 @@ gst_nv_dec_object_export_surface (GstNvDecObject * object,
   return GST_FLOW_OK;
 }
 
+guint
+gst_nv_dec_object_get_num_free_surfaces (GstNvDecObject * object)
+{
+  GstNvDecObjectPrivate *priv = object->priv;
+  std::lock_guard < std::mutex > lk (priv->lock);
+
+  if (object->num_mapped >= object->create_info.ulNumOutputSurfaces)
+    return 0;
+
+  return object->create_info.ulNumOutputSurfaces - object->num_mapped;
+}
+
 static gboolean
 gst_nv_dec_surface_dispose (GstNvDecSurface * surf)
 {
diff --git a/subprojects/gst-plugins-bad/sys/nvcodec/gstnvdecobject.h b/subprojects/gst-plugins-bad/sys/nvcodec/gstnvdecobject.h
index 57b9c5c5d2..2a28582c9f 100644
--- a/subprojects/gst-plugins-bad/sys/nvcodec/gstnvdecobject.h
+++ b/subprojects/gst-plugins-bad/sys/nvcodec/gstnvdecobject.h
@@ -79,6 +79,8 @@ GstFlowReturn    gst_nv_dec_object_export_surface (GstNvDecObject * object,
                                                    GstCudaStream * stream,
                                                    GstMemory ** memory);
 
+guint            gst_nv_dec_object_get_num_free_surfaces (GstNvDecObject * object);
+
 GType gst_nv_dec_surface_get_type (void);
 
 static inline GstNvDecSurface *
diff --git a/subprojects/gst-plugins-bad/sys/nvcodec/gstnvdecoder.cpp b/subprojects/gst-plugins-bad/sys/nvcodec/gstnvdecoder.cpp
index d98511d4e7..750baf9c1f 100644
--- a/subprojects/gst-plugins-bad/sys/nvcodec/gstnvdecoder.cpp
+++ b/subprojects/gst-plugins-bad/sys/nvcodec/gstnvdecoder.cpp
@@ -91,6 +91,7 @@ struct _GstNvDecoder
   gboolean configured;
   guint downstream_min_buffers;
   guint num_output_surfaces;
+  gboolean wait_on_pool_full;
 
   GMutex lock;
 
@@ -726,6 +727,7 @@ gst_nv_decoder_output_picture (GstNvDecoder * decoder,
   GstFlowReturn ret = GST_FLOW_OK;
   GstNvDecSurface *surface;
   GstCudaStream *stream;
+  gboolean can_export = FALSE;
 
   if (picture->discont_state) {
     if (!gst_nv_decoder_negotiate (decoder, videodec, picture->discont_state)) {
@@ -754,14 +756,32 @@ gst_nv_decoder_output_picture (GstNvDecoder * decoder,
     goto error;
   }
 
-  if (decoder->output_type == GST_NV_DECODER_OUTPUT_TYPE_CUDA &&
+  if (videodec->input_segment.rate > 0 &&
+      decoder->output_type == GST_NV_DECODER_OUTPUT_TYPE_CUDA &&
       (guint) decoder->create_info.ulNumOutputSurfaces >=
       decoder->downstream_min_buffers) {
+    if (decoder->wait_on_pool_full) {
+      can_export = TRUE;
+    } else {
+      guint num_free_surfaces =
+          gst_nv_dec_object_get_num_free_surfaces (decoder->object);
+
+      /* If downstream didn't propose pool but we have free surfaces */
+      if (num_free_surfaces > 0)
+        can_export = TRUE;
+      else
+        GST_LOG_OBJECT (decoder, "No more free output surface, need copy");
+    }
+  }
+
+  if (can_export) {
     GstMemory *mem;
     GstCudaMemory *cmem;
     GstBuffer *buf;
     GstVideoInfo *info = &decoder->info;
 
+    GST_LOG_OBJECT (decoder, "Exporting output surface without copy");
+
     ret = gst_nv_dec_object_export_surface (decoder->object,
         surface, stream, &mem);
     if (ret != GST_FLOW_OK) {
@@ -1542,6 +1562,9 @@ gst_nv_decoder_ensure_cuda_pool (GstNvDecoder * decoder, GstQuery * query)
     if (outcaps)
       gst_video_info_from_caps (&vinfo, outcaps);
     size = (guint) vinfo.size;
+    decoder->wait_on_pool_full = FALSE;
+  } else {
+    decoder->wait_on_pool_full = TRUE;
   }
 
   config = gst_buffer_pool_get_config (pool);