dxva2: use mp_image pool for d3d surfaces

This is required so that the individual surfaces can pass beyond the dxva2
decoder and be passed to the vo.

This also adds additional data to mp_image->planes[0] for IMGFMT_DXVA2, which is
required for maintaining and releasing the surface even if the decoder code is
uninited.

The IDirectXVideoDecoder itself is encapsulated together with its surface pool
and configuration in a dxva2_decoder structure whose creation and destruction is
managed by talloc.

4 files changed, 222 insertions, 131 deletions

diff --git a/video/decode/dxva2.c b/video/decode/dxva2.c
index 8cdda9f11b..8bfb75193d 100644
--- a/video/decode/dxva2.c
+++ b/video/decode/dxva2.c
@@ -25,8 +25,6 @@
 
 #include <stdint.h>
 
-#include <d3d9.h>
-#include <dxva2api.h>
 #include <ks.h>
 
 #include <libavcodec/dxva2.h>
@@ -36,6 +34,7 @@
 #include "common/av_common.h"
 #include "osdep/windows_utils.h"
 #include "video/fmt-conversion.h"
+#include "video/dxva2.h"
 #include "video/mp_image_pool.h"
 #include "video/hwdec.h"
 #include "video/d3d.h"
@@ -97,10 +96,11 @@ static const dxva2_mode dxva2_modes[] = {
 
 #undef MODE
 
-typedef struct surface_info {
-    int used;
-    uint64_t age;
-} surface_info;
+struct dxva2_decoder {
+    DXVA2_ConfigPictureDecode config;
+    IDirectXVideoDecoder      *decoder;
+    struct mp_image_pool      *pool;
+};
 
 typedef struct DXVA2Context {
     struct mp_log *log;
@@ -114,54 +114,18 @@ typedef struct DXVA2Context {
     IDirect3DDevice9            *d3d9device;
     IDirect3DDeviceManager9     *d3d9devmgr;
     IDirectXVideoDecoderService *decoder_service;
-    IDirectXVideoDecoder        *decoder;
-
-    DXVA2_ConfigPictureDecode   decoder_config;
+    struct dxva2_decoder        *decoder;
 
-    LPDIRECT3DSURFACE9          *surfaces;
-    surface_info                *surface_infos;
-    uint32_t                    num_surfaces;
-    uint64_t                    surface_age;
-
-    struct mp_image_pool       *sw_pool;
+    struct mp_image_pool        *sw_pool;
 } DXVA2Context;
 
-typedef struct DXVA2SurfaceWrapper {
-    DXVA2Context         *ctx;
-    LPDIRECT3DSURFACE9   surface;
-    IDirectXVideoDecoder *decoder;
-} DXVA2SurfaceWrapper;
-
-static void dxva2_destroy_decoder(struct lavc_ctx *s)
-{
-    DXVA2Context *ctx = s->hwdec_priv;
-    int i;
-
-    if (ctx->surfaces) {
-        for (i = 0; i < ctx->num_surfaces; i++) {
-            if (ctx->surfaces[i])
-                IDirect3DSurface9_Release(ctx->surfaces[i]);
-        }
-    }
-    av_freep(&ctx->surfaces);
-    av_freep(&ctx->surface_infos);
-    ctx->num_surfaces = 0;
-    ctx->surface_age  = 0;
-
-    if (ctx->decoder) {
-        IDirectXVideoDecoder_Release(ctx->decoder);
-        ctx->decoder = NULL;
-    }
-}
-
 static void dxva2_uninit(struct lavc_ctx *s)
 {
     DXVA2Context *ctx = s->hwdec_priv;
     if (!ctx)
         return;
 
-    if (ctx->decoder)
-        dxva2_destroy_decoder(s);
+    talloc_free(ctx->decoder);
 
     if (ctx->decoder_service)
         IDirectXVideoDecoderService_Release(ctx->decoder_service);
@@ -189,59 +153,14 @@ static void dxva2_uninit(struct lavc_ctx *s)
     s->hwdec_priv = NULL;
 }
 
-static void dxva2_release_img(void *ptr)
-{
-    DXVA2SurfaceWrapper *w   = ptr;
-    DXVA2Context        *ctx = w->ctx;
-    int i;
-
-    for (i = 0; i < ctx->num_surfaces; i++) {
-        if (ctx->surfaces[i] == w->surface) {
-            ctx->surface_infos[i].used = 0;
-            break;
-        }
-    }
-    IDirect3DSurface9_Release(w->surface);
-    IDirectXVideoDecoder_Release(w->decoder);
-    av_free(w);
-}
-
-static struct mp_image *dxva2_allocate_image(struct lavc_ctx *s,
-                                             int img_w, int img_h)
+static struct mp_image *dxva2_allocate_image(struct lavc_ctx *s, int w, int h)
 {
     DXVA2Context *ctx = s->hwdec_priv;
 
-    int i, old_unused = -1;
-    for (i = 0; i < ctx->num_surfaces; i++) {
-        surface_info *info = &ctx->surface_infos[i];
-        if (!info->used && (old_unused == -1 || info->age < ctx->surface_infos[old_unused].age))
-            old_unused = i;
-    }
-    if (old_unused == -1) {
-        MP_ERR(ctx, "No free DXVA2 surface!\n");
-        return NULL;
-    }
-    i = old_unused;
-
-    DXVA2SurfaceWrapper *w = av_mallocz(sizeof(*w));
-    if (!w)
-        return NULL;
-
-    w->ctx     = ctx;
-    w->surface = ctx->surfaces[i];;
-    IDirect3DSurface9_AddRef(w->surface);
-    w->decoder = ctx->decoder;
-    IDirectXVideoDecoder_AddRef(w->decoder);
-
-    ctx->surface_infos[i].used = 1;
-    ctx->surface_infos[i].age  = ctx->surface_age++;
-
-    struct mp_image mpi = {0};
-    mp_image_setfmt(&mpi, IMGFMT_DXVA2);
-    mp_image_set_size(&mpi, img_w, img_h);
-    mpi.planes[3] = (void *)w->surface;
-
-    return mp_image_new_custom_ref(&mpi, w, dxva2_release_img);
+    struct mp_image *img = mp_image_pool_get(ctx->decoder->pool, IMGFMT_DXVA2, w, h);
+    if (!img)
+        MP_ERR(ctx, "Failed to allocate additional DXVA2 surface.\n");
+    return img;
 }
 
 static void copy_nv12(struct mp_image *dest, uint8_t *src_bits,
@@ -262,7 +181,7 @@ static struct mp_image *dxva2_retrieve_image(struct lavc_ctx *s,
                                              struct mp_image *img)
 {
     DXVA2Context       *ctx = s->hwdec_priv;
-    LPDIRECT3DSURFACE9 surface = (LPDIRECT3DSURFACE9)img->planes[3];
+    LPDIRECT3DSURFACE9 surface = d3d9_surface_in_mp_image(img);
     D3DSURFACE_DESC    surfaceDesc;
     D3DLOCKED_RECT     LockedRect;
     HRESULT            hr;
@@ -469,20 +388,30 @@ static int dxva2_get_decoder_configuration(struct lavc_ctx *s,
     return 0;
 }
 
+static void dxva2_destroy_decoder(void *arg)
+{
+    struct dxva2_decoder *decoder = arg;
+    if (decoder->decoder)
+        IDirectXVideoDecoder_Release(decoder->decoder);
+}
+
 static int dxva2_create_decoder(struct lavc_ctx *s, int w, int h,
                                 enum AVCodecID codec_id, int profile)
 {
     DXVA2Context *ctx = s->hwdec_priv;
     struct dxva_context *dxva_ctx = s->avctx->hwaccel_context;
+    void *tmp = talloc_new(NULL);
     GUID *guid_list = NULL;
     unsigned guid_count = 0, i, j;
     GUID device_guid = GUID_NULL;
     D3DFORMAT target_format = 0;
     DXVA2_VideoDesc desc = { 0 };
-    DXVA2_ConfigPictureDecode config;
     HRESULT hr;
-    int surface_alignment;
-    int ret;
+    struct dxva2_decoder *decoder;
+    int surface_alignment, num_surfaces;
+    struct mp_image **imgs;
+    LPDIRECT3DSURFACE9 *surfaces;
+    int ret = -1;
 
     hr = IDirectXVideoDecoderService_GetDecoderDeviceGuids(ctx->decoder_service, &guid_count, &guid_list);
     if (FAILED(hr)) {
@@ -556,8 +485,10 @@ static int dxva2_create_decoder(struct lavc_ctx *s, int w, int h,
     desc.SampleHeight = h;
     desc.Format       = target_format;
 
-    ret = dxva2_get_decoder_configuration(s, codec_id, &device_guid, &desc, &config);
-    if (ret < 0) {
+    decoder = talloc_zero(tmp, struct dxva2_decoder);
+    talloc_set_destructor(decoder, dxva2_destroy_decoder);
+    if (dxva2_get_decoder_configuration(s, codec_id, &device_guid, &desc,
+                                        &decoder->config) < 0) {
         goto fail;
     }
 
@@ -572,50 +503,53 @@ static int dxva2_create_decoder(struct lavc_ctx *s, int w, int h,
     else
         surface_alignment = 16;
 
-    ctx->num_surfaces = hwdec_get_max_refs(s) + ADDITIONAL_SURFACES;
-
-    ctx->surfaces      = av_mallocz(ctx->num_surfaces * sizeof(*ctx->surfaces));
-    ctx->surface_infos = av_mallocz(ctx->num_surfaces * sizeof(*ctx->surface_infos));
+    num_surfaces = hwdec_get_max_refs(s) + ADDITIONAL_SURFACES;
 
-    if (!ctx->surfaces || !ctx->surface_infos) {
-        MP_ERR(ctx, "Unable to allocate surface arrays\n");
-        goto fail;
-    }
+    decoder->pool = talloc_steal(decoder, mp_image_pool_new(num_surfaces));
+    dxva2_pool_set_allocator(decoder->pool, ctx->decoder_service,
+                             target_format, surface_alignment);
 
-    hr = IDirectXVideoDecoderService_CreateSurface(ctx->decoder_service,
-                                                   FFALIGN(w, surface_alignment),
-                                                   FFALIGN(h, surface_alignment),
-                                                   ctx->num_surfaces - 1,
-                                                   target_format, D3DPOOL_DEFAULT, 0,
-                                                   DXVA2_VideoDecoderRenderTarget,
-                                                   ctx->surfaces, NULL);
-    if (FAILED(hr)) {
-        MP_ERR(ctx, "Failed to create %d video surfaces\n", ctx->num_surfaces);
-        goto fail;
+    // Preallocate images from the pool so the surfaces can be used to create
+    // the decoder and passed to ffmpeg in the dxva_ctx. The mp_images
+    // themselves will be freed (returned to the pool) along with the temporary
+    // talloc context on exit from this function.
+    imgs     = talloc_array(tmp,            struct mp_image *, num_surfaces);
+    surfaces = talloc_array(decoder->pool, LPDIRECT3DSURFACE9, num_surfaces);
+    for (i = 0; i < num_surfaces; i++) {
+        imgs[i] = talloc_steal(
+            imgs, mp_image_pool_get(decoder->pool, IMGFMT_DXVA2, w, h));
+        surfaces[i] = d3d9_surface_in_mp_image(imgs[i]);
     }
 
     hr = IDirectXVideoDecoderService_CreateVideoDecoder(ctx->decoder_service, &device_guid,
-                                                        &desc, &config, ctx->surfaces,
-                                                        ctx->num_surfaces, &ctx->decoder);
+                                                        &desc, &decoder->config, surfaces,
+                                                        num_surfaces, &decoder->decoder);
+
     if (FAILED(hr)) {
         MP_ERR(ctx, "Failed to create DXVA2 video decoder\n");
         goto fail;
     }
 
-    ctx->decoder_config = config;
+    // According to ffmpeg_dxva2.c, the surfaces must not outlive the
+    // IDirectXVideoDecoder they were used to create. This adds a reference for
+    // each one of them, which is released on final mp_image destruction.
+    for (i = 0; i < num_surfaces; i++)
+        dxva2_img_ref_decoder(imgs[i], decoder->decoder);
 
-    dxva_ctx->cfg           = &ctx->decoder_config;
-    dxva_ctx->decoder       = ctx->decoder;
-    dxva_ctx->surface       = ctx->surfaces;
-    dxva_ctx->surface_count = ctx->num_surfaces;
+    // Pass required information on to ffmpeg.
+    dxva_ctx->cfg           = &decoder->config;
+    dxva_ctx->decoder       = decoder->decoder;
+    dxva_ctx->surface       = surfaces;
+    dxva_ctx->surface_count = num_surfaces;
 
     if (IsEqualGUID(&device_guid, &DXVADDI_Intel_ModeH264_E))
         dxva_ctx->workaround |= FF_DXVA2_WORKAROUND_INTEL_CLEARVIDEO;
 
-    return 0;
+    ctx->decoder = talloc_steal(NULL, decoder);
+    ret = 0;
 fail:
-    dxva2_destroy_decoder(s);
-    return -1;
+    talloc_free(tmp);
+    return ret;
 }
 
 static int dxva2_init_decoder(struct lavc_ctx *s, int w, int h)
@@ -635,8 +569,8 @@ static int dxva2_init_decoder(struct lavc_ctx *s, int w, int h)
         return -1;
     }
 
-    if (ctx->decoder)
-        dxva2_destroy_decoder(s);
+    talloc_free(ctx->decoder);
+    ctx->decoder = NULL;
 
     if (dxva2_create_decoder(s, w, h, codec, profile) < 0) {
         MP_ERR(ctx, "Error creating the DXVA2 decoder\n");
diff --git a/video/dxva2.c b/video/dxva2.c
new file mode 100644
index 0000000000..cf0bf3d08c
--- /dev/null
+++ b/video/dxva2.c
@@ -0,0 +1,122 @@
+/*
+ * This file is part of mpv.
+ *
+ * mpv is free software; you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation; either version 2 of the License, or
+ * (at your option) any later version.
+ *
+ * mpv is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License along
+ * with mpv.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+#include <assert.h>
+
+#include "common/av_common.h"
+#include "dxva2.h"
+#include "mp_image.h"
+#include "img_format.h"
+#include "mp_image_pool.h"
+
+struct dxva2_surface {
+    HMODULE d3dlib;
+    HMODULE dxva2lib;
+
+    IDirectXVideoDecoder *decoder;
+    LPDIRECT3DSURFACE9   surface;
+};
+
+LPDIRECT3DSURFACE9 d3d9_surface_in_mp_image(struct mp_image *mpi)
+{
+    return mpi && mpi->imgfmt == IMGFMT_DXVA2 ?
+        (LPDIRECT3DSURFACE9)mpi->planes[3] : NULL;
+}
+
+void dxva2_img_ref_decoder(struct mp_image *mpi, IDirectXVideoDecoder *decoder)
+{
+    assert(mpi->imgfmt == IMGFMT_DXVA2);
+    struct dxva2_surface *surface = (struct dxva2_surface *)mpi->planes[0];
+    if (surface->decoder)
+        IDirectXVideoDecoder_Release(surface->decoder);
+    surface->decoder = decoder;
+    IDirectXVideoDecoder_AddRef(surface->decoder);
+}
+
+static void dxva2_pool_release_img(void *arg)
+{
+    struct dxva2_surface *surface = arg;
+    if (surface->surface)
+        IDirect3DSurface9_Release(surface->surface);
+
+    if (surface->decoder)
+        IDirectXVideoDecoder_Release(surface->decoder);
+
+    if (surface->dxva2lib)
+        FreeLibrary(surface->dxva2lib);
+
+    if (surface->d3dlib)
+        FreeLibrary(surface->d3dlib);
+
+    talloc_free(surface);
+}
+
+struct pool_alloc_ctx {
+    IDirectXVideoDecoderService *decoder_service;
+    D3DFORMAT target_format;
+    int surface_alignment;
+};
+
+static struct mp_image *dxva2_pool_alloc_img(void *arg, int fmt, int w, int h)
+{
+    if (fmt != IMGFMT_DXVA2)
+        return NULL;
+    struct dxva2_surface *surface = talloc_zero(NULL, struct dxva2_surface);
+
+    // Add additional references to the libraries which might otherwise be freed
+    // before the surface, which is observed to lead to bad behaviour
+    surface->d3dlib   = LoadLibrary(L"d3d9.dll");
+    surface->dxva2lib = LoadLibrary(L"dxva2.dll");
+    if (!surface->d3dlib || !surface->dxva2lib)
+        goto fail;
+
+    struct pool_alloc_ctx *alloc_ctx = arg;
+    HRESULT hr = IDirectXVideoDecoderService_CreateSurface(
+        alloc_ctx->decoder_service,
+        FFALIGN(w, alloc_ctx->surface_alignment),
+        FFALIGN(h, alloc_ctx->surface_alignment),
+        0, alloc_ctx->target_format, D3DPOOL_DEFAULT, 0,
+        DXVA2_VideoDecoderRenderTarget,
+        &surface->surface, NULL);
+    if (FAILED(hr))
+        goto fail;
+
+    struct mp_image mpi = {0};
+    mp_image_setfmt(&mpi, IMGFMT_DXVA2);
+    mp_image_set_size(&mpi, w, h);
+    mpi.planes[0] = (void *)surface;
+    mpi.planes[3] = (void *)surface->surface;
+
+    return mp_image_new_custom_ref(&mpi, surface, dxva2_pool_release_img);
+fail:
+    dxva2_pool_release_img(surface);
+    return NULL;
+}
+
+void dxva2_pool_set_allocator(struct mp_image_pool *pool,
+                              IDirectXVideoDecoderService *decoder_service,
+                              D3DFORMAT target_format, int surface_alignment)
+{
+    struct pool_alloc_ctx *alloc_ctx = talloc_ptrtype(pool, alloc_ctx);
+    *alloc_ctx =  (struct pool_alloc_ctx){
+        decoder_service   = decoder_service,
+        target_format     = target_format,
+        surface_alignment = surface_alignment
+    };
+    mp_image_pool_set_allocator(pool, dxva2_pool_alloc_img, alloc_ctx);
+    mp_image_pool_set_lru(pool);
+}
diff --git a/video/dxva2.h b/video/dxva2.h
new file mode 100644
index 0000000000..4acb9c99d4
--- /dev/null
+++ b/video/dxva2.h
@@ -0,0 +1,34 @@
+/*
+ * This file is part of mpv.
+ *
+ * mpv is free software; you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation; either version 2 of the License, or
+ * (at your option) any later version.
+ *
+ * mpv is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License along
+ * with mpv.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+#ifndef MPV_DXVA2_H
+#define MPV_DXVA2_H
+
+#include <d3d9.h>
+#include <dxva2api.h>
+
+struct mp_image;
+struct mp_image_pool;
+
+LPDIRECT3DSURFACE9 d3d9_surface_in_mp_image(struct mp_image *mpi);
+void dxva2_img_ref_decoder(struct mp_image *mpi, IDirectXVideoDecoder *decoder);
+
+void dxva2_pool_set_allocator(struct mp_image_pool *pool,
+                              IDirectXVideoDecoderService *decoder_service,
+                              D3DFORMAT target_format, int surface_alignment);
+
+#endif
diff --git a/wscript_build.py b/wscript_build.py
index 090ac96181..75dc763f35 100644
--- a/wscript_build.py
+++ b/wscript_build.py
@@ -280,6 +280,7 @@ def build(ctx):
         ( "video/mp_image.c" ),
         ( "video/mp_image_pool.c" ),
         ( "video/sws_utils.c" ),
+        ( "video/dxva2.c",                       "dxva2-hwaccel" ),
         ( "video/vaapi.c",                       "vaapi" ),
         ( "video/vdpau.c",                       "vdpau" ),
         ( "video/vdpau_mixer.c",                 "vdpau" ),