qsv: fix --start-at frames issue and dx surfaces leaks (#2533)

5 files changed, 110 insertions, 94 deletions

diff --git a/libhb/enc_qsv.c b/libhb/enc_qsv.c
index 06f3ee27e..63474ab66 100644
--- a/libhb/enc_qsv.c
+++ b/libhb/enc_qsv.c
@@ -117,23 +117,6 @@ struct hb_work_private_s
     hb_list_t          * loaded_plugins;
 };
 
-int hb_qsv_find_surface_idx(QSVMid *mids, int nb_mids, QSVMid *mid)
-{
-    if(mids)
-    {
-        QSVMid *m = &mids[0];
-        if (m->texture != mid->texture)
-            return -1;
-        int i;
-        for (i = 0; i < nb_mids; i++) {
-            m = &mids[i];
-            if ( m->handle == mid->handle )
-                return i;
-        }
-    }
-    return -1;
-}
-
 static void hb_qsv_add_new_dts(hb_list_t *list, int64_t new_dts)
 {
     if (list != NULL)
@@ -2098,31 +2081,9 @@ static int qsv_enc_work(hb_work_private_t *pv,
 
             if (sts == MFX_ERR_MORE_DATA)
             {
-                if(surface)
+                if(!pv->is_sys_mem && surface)
                 {
-                    // release surface
-                    if(!pv->is_sys_mem)
-                    {
-                        QSVMid *mid = surface->Data.MemId;
-                        int ret = hb_qsv_find_surface_idx(hb_enc_qsv_frames_ctx.mids, hb_enc_qsv_frames_ctx.nb_mids, mid);
-                        if (ret < 0)
-                        {
-                            ret = hb_qsv_find_surface_idx(hb_enc_qsv_frames_ctx.mids2, hb_enc_qsv_frames_ctx.nb_mids, mid);
-                            if (ret < 0)
-                            {
-                                hb_error("encqsv: Surface with MemId=%p has not been found in the pool\n", mid);
-                                return -1;
-                            }
-                            else
-                            {
-                                ff_qsv_atomic_dec(&hb_enc_qsv_frames_ctx.pool2[ret]);
-                            }
-                        }
-                        else
-                        {
-                            ff_qsv_atomic_dec(&hb_enc_qsv_frames_ctx.pool[ret]);
-                        }
-                    }
+                    hb_qsv_release_surface_from_pool(surface->Data.MemId);
                 }
 
                 if (qsv_atom != NULL)
@@ -2196,32 +2157,10 @@ static int qsv_enc_work(hb_work_private_t *pv,
                 /* perform a sync operation to get the output bitstream */
                 hb_qsv_wait_on_sync(qsv_ctx, task->stage);
 
-                // release surface
-                if(!pv->is_sys_mem)
+                mfxFrameSurface1 *surface = task->stage->in.p_surface;
+                if(!pv->is_sys_mem && surface)
                 {
-                    mfxFrameSurface1 *surface_to_release = task->stage->in.p_surface;
-                    if(surface_to_release)
-                    {
-                        QSVMid *mid = surface_to_release->Data.MemId;
-                        int ret = hb_qsv_find_surface_idx(hb_enc_qsv_frames_ctx.mids, hb_enc_qsv_frames_ctx.nb_mids, mid);
-                        if (ret < 0)
-                        {
-                            ret = hb_qsv_find_surface_idx(hb_enc_qsv_frames_ctx.mids2, hb_enc_qsv_frames_ctx.nb_mids, mid);
-                            if (ret < 0)
-                            {
-                                hb_error("encqsv: Surface with MemId=%p has not been found in the pool\n", mid);
-                                return -1;
-                            }
-                            else
-                            {
-                                ff_qsv_atomic_dec(&hb_enc_qsv_frames_ctx.pool2[ret]);
-                            }
-                        }
-                        else
-                        {
-                            ff_qsv_atomic_dec(&hb_enc_qsv_frames_ctx.pool[ret]);
-                        }
-                    }
+                    hb_qsv_release_surface_from_pool(surface->Data.MemId);
                 }
 
                 if (task->bs->DataLength > 0)
@@ -2255,7 +2194,6 @@ int encqsvWork(hb_work_object_t *w, hb_buffer_t **buf_in, hb_buffer_t **buf_out)
     hb_work_private_t *pv = w->private_data;
     hb_buffer_t *in       = *buf_in;
     hb_job_t *job         = pv->job;
-
     while (qsv_enc_init(pv) >= 2)
     {
         hb_qsv_sleep(1); // encoding not initialized, wait and repeat the call
@@ -2310,29 +2248,12 @@ int encqsvWork(hb_work_object_t *w, hb_buffer_t **buf_in, hb_buffer_t **buf_out)
         else
         {
             // Create black buffer in the begining of the encoding, usually first 2 frames
-            hb_qsv_get_free_surface_from_pool(&mid, &surface, HB_POOL_SURFACE_SIZE);
+            hb_qsv_get_free_surface_from_pool(HB_POOL_SURFACE_SIZE - HB_POOL_ENCODER_SIZE, HB_POOL_SURFACE_SIZE, &mid, &surface);
         }
 
         if(surface)
         {
-            int ret = hb_qsv_find_surface_idx(hb_enc_qsv_frames_ctx.mids, hb_enc_qsv_frames_ctx.nb_mids, mid);
-            if (ret < 0)
-            {
-                ret = hb_qsv_find_surface_idx(hb_enc_qsv_frames_ctx.mids2, hb_enc_qsv_frames_ctx.nb_mids, mid);
-                if (ret < 0)
-                {
-                    hb_error("encqsv: Surface with MemId=%p has not been found in the pool\n", mid);
-                    return -1;
-                }
-                else
-                {
-                    surface->Data.MemId = &hb_enc_qsv_frames_ctx.mids2[ret];
-                }
-            }
-            else
-            {
-                surface->Data.MemId = &hb_enc_qsv_frames_ctx.mids[ret];
-            }
+            hb_qsv_replace_surface_mid(mid, surface);
         }
         else
         {
@@ -2435,7 +2356,12 @@ int encqsvWork(hb_work_object_t *w, hb_buffer_t **buf_in, hb_buffer_t **buf_out)
     {
         goto fail;
     }
-
+#if HB_PROJECT_FEATURE_QSV
+    if (in->qsv_details.frame)
+    {
+        in->qsv_details.frame->data[3] = 0;
+    }
+#endif
     *buf_out = hb_buffer_list_clear(&pv->encoded_frames);
     return HB_WORK_OK;
 
diff --git a/libhb/fifo.c b/libhb/fifo.c
index 68585ee31..6a279fbb3 100644
--- a/libhb/fifo.c
+++ b/libhb/fifo.c
@@ -12,6 +12,7 @@
 #include "handbrake/handbrake.h"
 #if HB_PROJECT_FEATURE_QSV
 #include "handbrake/qsv_libav.h"
+#include "handbrake/qsv_common.h"
 #endif
 
 #ifndef SYS_DARWIN
@@ -724,6 +725,15 @@ void hb_buffer_close( hb_buffer_t ** _b )
 #if HB_PROJECT_FEATURE_QSV
         // Reclaim QSV resources before dropping the buffer.
         // when decoding without QSV, the QSV atom will be NULL.
+        if(b->qsv_details.frame)
+        {
+            mfxFrameSurface1 *surface = (mfxFrameSurface1*)b->qsv_details.frame->data[3];
+            if(surface)
+            {
+                hb_qsv_release_surface_from_pool(surface->Data.MemId);
+                b->qsv_details.frame->data[3] = 0;
+            }
+        }
         if (b->qsv_details.qsv_atom != NULL && b->qsv_details.ctx != NULL)
         {
             hb_qsv_stage *stage = hb_qsv_get_last_stage(b->qsv_details.qsv_atom);
diff --git a/libhb/handbrake/qsv_common.h b/libhb/handbrake/qsv_common.h
index 93d1f9609..99685ae34 100644
--- a/libhb/handbrake/qsv_common.h
+++ b/libhb/handbrake/qsv_common.h
@@ -227,6 +227,7 @@ typedef struct QSVFrame {
 } QSVFrame;
 
 #define HB_POOL_SURFACE_SIZE (64)
+#define HB_POOL_ENCODER_SIZE (8)
 
 typedef struct EncQSVFramesContext {
     AVBufferRef *hw_frames_ctx;
@@ -251,7 +252,9 @@ typedef struct EncQSVFramesContext {
 int hb_qsv_full_path_is_enabled(hb_job_t *job);
 AVBufferRef *hb_qsv_create_mids(AVBufferRef *hw_frames_ref);
 hb_buffer_t* hb_qsv_copy_frame(AVFrame *frame, hb_qsv_context *qsv_ctx);
-void hb_qsv_get_free_surface_from_pool(QSVMid **out_mid, mfxFrameSurface1 **out_surface, int pool_size);
+void hb_qsv_get_free_surface_from_pool(const int start_index, const int end_index, QSVMid **out_mid, mfxFrameSurface1 **out_surface);
+int hb_qsv_replace_surface_mid(const QSVMid *mid, mfxFrameSurface1 *surface);
+int hb_qsv_release_surface_from_pool(const QSVMid *mid);
 int hb_qsv_get_buffer(AVCodecContext *s, AVFrame *frame, int flags);
 enum AVPixelFormat hb_qsv_get_format(AVCodecContext *s, const enum AVPixelFormat *pix_fmts);
 int hb_qsv_preset_is_zero_copy_enabled(const hb_dict_t *job_dict);
diff --git a/libhb/qsv_common.c b/libhb/qsv_common.c
index 2c96886ed..0af10b9fb 100644
--- a/libhb/qsv_common.c
+++ b/libhb/qsv_common.c
@@ -2523,7 +2523,70 @@ static HRESULT unlock_device(
     return hr;
 }
 
-void hb_qsv_get_free_surface_from_pool(QSVMid **out_mid, mfxFrameSurface1 **out_surface, int pool_size)
+static int hb_qsv_find_surface_idx(const QSVMid *mids, const int nb_mids, const QSVMid *mid)
+{
+    if(mids)
+    {
+        const QSVMid *m = &mids[0];
+        if (m->texture != mid->texture)
+            return -1;
+        int i;
+        for (i = 0; i < nb_mids; i++) {
+            m = &mids[i];
+            if ( m->handle == mid->handle )
+                return i;
+        }
+    }
+    return -1;
+}
+
+int hb_qsv_replace_surface_mid(const QSVMid *mid, mfxFrameSurface1 *surface)
+{
+    int ret = hb_qsv_find_surface_idx(hb_enc_qsv_frames_ctx.mids, hb_enc_qsv_frames_ctx.nb_mids, mid);
+    if (ret < 0)
+    {
+        ret = hb_qsv_find_surface_idx(hb_enc_qsv_frames_ctx.mids2, hb_enc_qsv_frames_ctx.nb_mids, mid);
+        if (ret < 0)
+        {
+            hb_error("encqsv: Surface with MemId=%p has not been found in the pool\n", mid);
+            return -1;
+        }
+        else
+        {
+            surface->Data.MemId = &hb_enc_qsv_frames_ctx.mids2[ret];
+        }
+    }
+    else
+    {
+        surface->Data.MemId = &hb_enc_qsv_frames_ctx.mids[ret];
+    }
+    return 0;
+}
+
+int hb_qsv_release_surface_from_pool(const QSVMid *mid)
+{
+    int ret = hb_qsv_find_surface_idx(hb_enc_qsv_frames_ctx.mids, hb_enc_qsv_frames_ctx.nb_mids, mid);
+    if (ret < 0)
+    {
+        ret = hb_qsv_find_surface_idx(hb_enc_qsv_frames_ctx.mids2, hb_enc_qsv_frames_ctx.nb_mids, mid);
+        if (ret < 0)
+        {
+            hb_error("encqsv: Surface with MemId=%p has not been found in the pool\n", mid);
+            return -1;
+        }
+        else
+        {
+            ff_qsv_atomic_dec(&hb_enc_qsv_frames_ctx.pool2[ret]);
+        }
+    }
+    else
+    {
+        ff_qsv_atomic_dec(&hb_enc_qsv_frames_ctx.pool[ret]);
+    }
+    return 0;
+}
+
+void hb_qsv_get_free_surface_from_pool(const int start_index, const int end_index, QSVMid **out_mid, mfxFrameSurface1 **out_surface)
 {
     QSVMid *mid = NULL;
     mfxFrameSurface1 *output_surface = NULL;
@@ -2540,11 +2603,11 @@ void hb_qsv_get_free_surface_from_pool(QSVMid **out_mid, mfxFrameSurface1 **out_
     {
         if(count > 30)
         {
-            hb_qsv_sleep(100); // prevent hang when all surfaces all used
+            hb_qsv_sleep(10); // prevent hang when all surfaces all used
             count = 0;
         }
 
-        for(int i = 0; i < pool_size; i++)
+        for(int i = start_index; i < end_index; i++)
         {
             if(hb_enc_qsv_frames_ctx.pool[i] == 0)
             {
@@ -2560,7 +2623,7 @@ void hb_qsv_get_free_surface_from_pool(QSVMid **out_mid, mfxFrameSurface1 **out_
             }
         }
 
-        for(int i = 0; i < pool_size; i++)
+        for(int i = start_index; i < end_index; i++)
         {
             if(hb_enc_qsv_frames_ctx.pool2[i] == 0)
             {
@@ -2600,7 +2663,7 @@ hb_buffer_t* hb_qsv_copy_frame(AVFrame *frame, hb_qsv_context *qsv_ctx)
     QSVMid *mid = NULL;
     mfxFrameSurface1* output_surface = NULL;
 
-    hb_qsv_get_free_surface_from_pool(&mid, &output_surface, HB_POOL_SURFACE_SIZE - 2); // leave 2 empty surfaces in the pool for black buffers
+    hb_qsv_get_free_surface_from_pool(0, HB_POOL_SURFACE_SIZE - HB_POOL_ENCODER_SIZE, &mid, &output_surface);
 
     static const mfxHandleType handle_types[] = {
         MFX_HANDLE_VA_DISPLAY,
@@ -2917,11 +2980,21 @@ hb_buffer_t* hb_qsv_copy_frame(AVFrame *frame, hb_qsv_context *qsv_ctx)
     return NULL;
 }
 
-void hb_qsv_get_free_surface_from_pool(QSVMid **out_mid, mfxFrameSurface1 **out_surface, int pool_size)
+void hb_qsv_get_free_surface_from_pool(const int start_index, const int end_index, QSVMid **out_mid, mfxFrameSurface1 **out_surface)
 {
     return;
 }
 
+int hb_qsv_replace_surface_mid(const QSVMid *mid, mfxFrameSurface1 *surface)
+{
+    return -1;
+}
+
+int hb_qsv_release_surface_from_pool(const QSVMid *mid)
+{
+    return -1;
+}
+
 enum AVPixelFormat hb_qsv_get_format(AVCodecContext *s, const enum AVPixelFormat *pix_fmts)
 {
     return AV_PIX_FMT_NONE;
diff --git a/libhb/sync.c b/libhb/sync.c
index f01cbe32a..4e14732e4 100644
--- a/libhb/sync.c
+++ b/libhb/sync.c
@@ -12,6 +12,10 @@
 #include <stdio.h>
 #include "handbrake/audio_resample.h"
 
+#if HB_PROJECT_FEATURE_QSV
+#include "handbrake/qsv_common.h"
+#endif
+
 #define SYNC_MAX_VIDEO_QUEUE_LEN    40
 #define SYNC_MIN_VIDEO_QUEUE_LEN    20