vaapi: replace image pool implementation with mp_image_pool

Although I at first thought it would be better to have a separate
implementation for hwaccels because the difference to software images
are too large, it turns out you can actually save some code with it.

Note that the old implementation had a small memory management bug. This
got painted over in commit 269c1e1, but is hereby solved properly.

Also note that I couldn't test vf_vavpp.c (due to lack of hardware), and
I hope I didn't accidentally break it.

5 files changed, 152 insertions, 241 deletions

diff --git a/video/decode/vaapi.c b/video/decode/vaapi.c
index 63c180c5a0..82855c04eb 100644
--- a/video/decode/vaapi.c
+++ b/video/decode/vaapi.c
@@ -64,9 +64,7 @@ struct priv {
     struct vaapi_context *va_context;
     struct vaapi_context va_context_storage;
 
-    VASurfaceID surfaces[MAX_SURFACES];
-
-    struct va_surface_pool *pool;
+    struct mp_image_pool *pool;
     int rt_format;
 
     struct mp_image_pool *sw_pool;
@@ -162,19 +160,26 @@ static int is_direct_mapping(VADisplay display)
 // We achieve this by reserving surfaces in the pool as needed.
 // Releasing surfaces is necessary after filling the surface id list so
 // that reserved surfaces can be reused for decoding.
-static bool preallocate_surfaces(struct lavc_ctx *ctx, int num, int w, int h)
+static bool preallocate_surfaces(struct lavc_ctx *ctx, int num, int w, int h,
+                                 VASurfaceID out_surfaces[MAX_SURFACES])
 {
     struct priv *p = ctx->hwdec_priv;
-    if (!va_surface_pool_reserve(p->pool, num, w, h)) {
-        MP_ERR(p, "Could not allocate surfaces.\n");
-        return false;
-    }
-    for (int i = 0; i < num; i++) {
-        struct va_surface *s = va_surface_pool_get(p->pool, w, h);
-        p->surfaces[i] = s->id;
-        va_surface_release(s);
+    assert(num <= MAX_SURFACES);
+    struct mp_image *reserve[MAX_SURFACES] = {0};
+    bool res = true;
+
+    for (int n = 0; n < num; n++) {
+        reserve[n] = mp_image_pool_get(p->pool, IMGFMT_VAAPI, w, h);
+        out_surfaces[n] = va_surface_id_in_mp_image(reserve[n]);
+        if (out_surfaces[n] == VA_INVALID_ID) {
+            MP_ERR(p, "Could not allocate surfaces.\n");
+            res = false;
+            break;
+        }
     }
-    return true;
+    for (int i = 0; i < num; i++)
+        talloc_free(reserve[i]);
+    return res;
 }
 
 static void destroy_decoder(struct lavc_ctx *ctx)
@@ -191,8 +196,7 @@ static void destroy_decoder(struct lavc_ctx *ctx)
         p->va_context->config_id = VA_INVALID_ID;
     }
 
-    for (int n = 0; n < MAX_SURFACES; n++)
-        p->surfaces[n] = VA_INVALID_ID;
+    mp_image_pool_clear(p->pool);
 }
 
 static bool has_profile(VAProfile *va_profiles, int num_profiles, VAProfile p)
@@ -251,7 +255,8 @@ static int init_decoder(struct lavc_ctx *ctx, int fmt, int w, int h)
         goto error;
     }
 
-    if (!preallocate_surfaces(ctx, num_surfaces, w, h)) {
+    VASurfaceID surfaces[MAX_SURFACES];
+    if (!preallocate_surfaces(ctx, num_surfaces, w, h, surfaces)) {
         MP_ERR(p, "Could not allocate surfaces.\n");
         goto error;
     }
@@ -287,7 +292,7 @@ static int init_decoder(struct lavc_ctx *ctx, int fmt, int w, int h)
 
     status = vaCreateContext(p->display, p->va_context->config_id,
                              w, h, VA_PROGRESSIVE,
-                             p->surfaces, num_surfaces,
+                             surfaces, num_surfaces,
                              &p->va_context->context_id);
     if (!CHECK_VA_STATUS(p, "vaCreateContext()"))
         goto error;
@@ -303,19 +308,13 @@ static struct mp_image *allocate_image(struct lavc_ctx *ctx, int format,
 {
     struct priv *p = ctx->hwdec_priv;
 
-    struct va_surface *s = va_surface_pool_get(p->pool, w, h);
-    if (s) {
-        for (int n = 0; n < MAX_SURFACES; n++) {
-            if (p->surfaces[n] == s->id)
-                return va_surface_wrap(s);
-        }
-        va_surface_release(s);
-    }
-    MP_ERR(p, "Insufficient number of surfaces.\n");
-    return NULL;
+    struct mp_image *img =
+        mp_image_pool_get_no_alloc(p->pool, IMGFMT_VAAPI, w, h);
+    if (!img)
+        MP_ERR(p, "Insufficient number of surfaces.\n");
+    return img;
 }
 
-
 static void destroy_va_dummy_ctx(struct priv *p)
 {
     if (p->x11_display)
@@ -355,7 +354,6 @@ static void uninit(struct lavc_ctx *ctx)
         return;
 
     destroy_decoder(ctx);
-    va_surface_pool_release(p->pool);
 
     if (p->x11_display)
         destroy_va_dummy_ctx(p);
@@ -382,7 +380,8 @@ static int init_with_vactx(struct lavc_ctx *ctx, struct mp_vaapi_ctx *vactx)
     }
 
     p->display = p->ctx->display;
-    p->pool = va_surface_pool_alloc(p->ctx, p->rt_format);
+    p->pool = talloc_steal(p, mp_image_pool_new(MAX_SURFACES));
+    va_pool_set_allocator(p->pool, p->ctx, p->rt_format);
     p->sw_pool = talloc_steal(p, mp_image_pool_new(17));
 
     p->va_context->display = p->display;
diff --git a/video/filter/vf_vavpp.c b/video/filter/vf_vavpp.c
index 7ebb8f3eb0..a1a7d3adf0 100644
--- a/video/filter/vf_vavpp.c
+++ b/video/filter/vf_vavpp.c
@@ -23,6 +23,7 @@
 #include "vf.h"
 #include "video/vaapi.h"
 #include "video/hwdec.h"
+#include "video/mp_image_pool.h"
 
 static bool check_error(struct vf_instance *vf, VAStatus status, const char *msg)
 {
@@ -59,7 +60,8 @@ struct vf_priv_s {
     VADisplay display;
     struct mp_vaapi_ctx *va;
     struct pipeline pipe;
-    struct va_surface_pool *pool;
+    struct mp_image_pool *pool;
+    int current_rt_format;
 };
 
 static const struct vf_priv_s vf_priv_default = {
@@ -127,13 +129,16 @@ static struct mp_image *render(struct vf_instance *vf, struct va_surface *in,
     struct vf_priv_s *p = vf->priv;
     if (!p->pipe.filters || !in)
         return NULL;
-    struct va_surface *out = va_surface_pool_get(p->pool, in->w, in->h);
-    if (!out)
+    struct mp_image *img = mp_image_pool_get(p->pool, IMGFMT_VAAPI, in->w, in->h);
+    if (!img)
         return NULL;
     enum {Begun = 1, Rendered = 2};
     int state = 0;
     do { // not a loop, just for break
-        VAStatus status = vaBeginPicture(p->display, p->context, out->id);
+        VASurfaceID id = va_surface_id_in_mp_image(img);
+        if (id == VA_INVALID_ID)
+            break;
+        VAStatus status = vaBeginPicture(p->display, p->context, id);
         if (!check_error(vf, status, "vaBeginPicture()"))
             break;
         state |= Begun;
@@ -167,8 +172,8 @@ static struct mp_image *render(struct vf_instance *vf, struct va_surface *in,
     if (state & Begun)
         vaEndPicture(p->display, p->context);
     if (state & Rendered)
-        return va_surface_wrap(out);
-    va_surface_release(out);
+        return img;
+    talloc_free(img);
     return NULL;
 }
 
@@ -203,13 +208,13 @@ static int process(struct vf_instance *vf, struct mp_image *in,
 static struct mp_image *upload(struct vf_instance *vf, struct mp_image *in)
 {
     struct vf_priv_s *p = vf->priv;
-    struct va_surface *surface =
-        va_surface_pool_get_by_imgfmt(p->pool, in->imgfmt, in->w, in->h);
-    if (!surface)
-        surface = va_surface_pool_get(p->pool, in->w, in->h); // dummy
-    else
-        va_surface_upload(surface, in);
-    struct mp_image *out = va_surface_wrap(surface);
+    struct mp_image *out = mp_image_pool_get(p->pool, IMGFMT_VAAPI, in->w, in->h);
+    if (!out)
+        return NULL;
+    if (va_surface_upload_image(out, in) < 0) {
+        talloc_free(out);
+        return NULL;
+    }
     mp_image_copy_attributes(out, in);
     return out;
 }
@@ -219,14 +224,18 @@ static int filter_ext(struct vf_instance *vf, struct mp_image *in)
     struct vf_priv_s *p = vf->priv;
     struct va_surface *surface = va_surface_in_mp_image(in);
     const int rt_format = surface ? surface->rt_format : VA_RT_FORMAT_YUV420;
-    if (!p->pool || va_surface_pool_rt_format(p->pool) != rt_format) {
-        va_surface_pool_release(p->pool);
-        p->pool = va_surface_pool_alloc(p->va, rt_format);
+    if (!p->pool || p->current_rt_format != rt_format) {
+        talloc_free(p->pool);
+        p->pool = mp_image_pool_new(20);
+        va_pool_set_allocator(p->pool, p->va, rt_format);
+        p->current_rt_format = rt_format;
     }
     if (!surface) {
         struct mp_image *tmp = upload(vf, in);
         talloc_free(in);
         in = tmp;
+        if (!in)
+            return -1;
     }
 
     struct mp_image *out1, *out2;
@@ -267,7 +276,7 @@ static void uninit(struct vf_instance *vf)
         vaDestroyConfig(p->display, p->config);
     free(p->pipe.forward.surfaces);
     free(p->pipe.backward.surfaces);
-    va_surface_pool_release(p->pool);
+    talloc_free(p->pool);
 }
 
 static int query_format(struct vf_instance *vf, unsigned int imgfmt)
diff --git a/video/out/vo_vaapi.c b/video/out/vo_vaapi.c
index 8b40afe44c..fc9e93b2c4 100644
--- a/video/out/vo_vaapi.c
+++ b/video/out/vo_vaapi.c
@@ -32,6 +32,7 @@
 #include "common/msg.h"
 #include "video/out/vo.h"
 #include "video/memcpy_pic.h"
+#include "video/mp_image_pool.h"
 #include "sub/osd.h"
 #include "sub/img_convert.h"
 #include "x11_common.h"
@@ -90,10 +91,10 @@ struct priv {
     struct vaapi_osd_part    osd_parts[MAX_OSD_PARTS];
     bool                     osd_screen;
 
-    struct va_surface_pool  *pool;
+    struct mp_image_pool    *pool;
     struct va_image_formats *va_image_formats;
 
-    struct va_surface       *black_surface;
+    struct mp_image         *black_surface;
 
     VAImageFormat           *va_subpic_formats;
     unsigned int            *va_subpic_flags;
@@ -123,19 +124,20 @@ static void free_video_specific(struct priv *p)
 {
     flush_output_surfaces(p);
 
-    va_surface_releasep(&p->black_surface);
+    mp_image_unrefp(&p->black_surface);
 
     for (int n = 0; n < MAX_OUTPUT_SURFACES; n++)
         mp_image_unrefp(&p->swdec_surfaces[n]);
+
+    mp_image_pool_clear(p->pool);
 }
 
 static bool alloc_swdec_surfaces(struct priv *p, int w, int h, int imgfmt)
 {
     free_video_specific(p);
     for (int i = 0; i < MAX_OUTPUT_SURFACES; i++) {
-        p->swdec_surfaces[i] =
-            va_surface_pool_get_wrapped(p->pool, imgfmt, w, h);
-        if (!p->swdec_surfaces[i])
+        p->swdec_surfaces[i] = mp_image_pool_get(p->pool, IMGFMT_VAAPI, w, h);
+        if (va_surface_image_alloc_imgfmt(p->swdec_surfaces[i], imgfmt) < 0)
             return false;
     }
     return true;
@@ -188,17 +190,16 @@ static bool render_to_screen(struct priv *p, struct mp_image *mpi)
             int w = p->image_params.w, h = p->image_params.h;
             // 4:2:0 should work everywhere
             int fmt = IMGFMT_420P;
-            p->black_surface =
-                va_surface_pool_get_by_imgfmt(p->pool, fmt, w, h);
+            p->black_surface = mp_image_pool_get(p->pool, IMGFMT_VAAPI, w, h);
             if (p->black_surface) {
                 struct mp_image *img = mp_image_alloc(fmt, w, h);
                 mp_image_clear(img, 0, 0, w, h);
-                if (!va_surface_upload(p->black_surface, img))
-                    va_surface_releasep(&p->black_surface);
+                if (va_surface_upload_image(p->black_surface, img) < 0)
+                    mp_image_unrefp(&p->black_surface);
                 talloc_free(img);
             }
         }
-        surface = va_surface_id(p->black_surface);
+        surface = va_surface_id_in_mp_image(p->black_surface);
     }
 
     int fields = mpi ? mpi->fields : 0;
@@ -271,14 +272,13 @@ static void draw_image(struct vo *vo, struct mp_image *mpi)
     struct priv *p = vo->priv;
 
     if (mpi->imgfmt != IMGFMT_VAAPI) {
-        struct mp_image *wrapper = p->swdec_surfaces[p->output_surface];
-        struct va_surface *surface = va_surface_in_mp_image(wrapper);
-        if (!surface || !va_surface_upload(surface, mpi)) {
+        struct mp_image *dst = p->swdec_surfaces[p->output_surface];
+        if (!dst || va_surface_upload_image(dst, mpi) < 0) {
             MP_WARN(vo, "Could not upload surface.\n");
             return;
         }
-        mp_image_copy_attributes(wrapper, mpi);
-        mpi = wrapper;
+        mp_image_copy_attributes(dst, mpi);
+        mpi = dst;
     }
 
     mp_image_setrefp(&p->output_surfaces[p->output_surface], mpi);
@@ -551,7 +551,7 @@ static void uninit(struct vo *vo)
     struct priv *p = vo->priv;
 
     free_video_specific(p);
-    va_surface_pool_release(p->pool);
+    talloc_free(p->pool);
 
     for (int n = 0; n < MAX_OSD_PARTS; n++) {
         struct vaapi_osd_part *part = &p->osd_parts[n];
@@ -584,7 +584,8 @@ static int preinit(struct vo *vo)
         return -1;
     }
 
-    p->pool = va_surface_pool_alloc(p->mpvaapi, VA_RT_FORMAT_YUV420);
+    p->pool = mp_image_pool_new(MAX_OUTPUT_SURFACES + 3);
+    va_pool_set_allocator(p->pool, p->mpvaapi, VA_RT_FORMAT_YUV420);
     p->va_image_formats = p->mpvaapi->image_formats;
 
     int max_subpic_formats = vaMaxNumSubpictureFormats(p->display);
diff --git a/video/vaapi.c b/video/vaapi.c
index 0bdffe3c72..edb16ffd48 100644
--- a/video/vaapi.c
+++ b/video/vaapi.c
@@ -159,75 +159,19 @@ VAImageFormat *va_image_format_from_imgfmt(const struct va_image_formats *format
     return NULL;
 }
 
-static void va_surface_destroy(struct va_surface *surface);
-
-struct va_surface_pool {
-    struct mp_log *log;
-    struct mp_vaapi_ctx *ctx;
-    VADisplay display;
-    int rt_format;
-    int num_surfaces, lru_counter;
-    struct va_surface **surfaces;
-};
-
 typedef struct va_surface_priv {
     struct mp_vaapi_ctx *ctx;
     VADisplay display;
     VAImage image;       // used for software decoding case
     bool is_derived;     // is image derived by vaDeriveImage()?
-    bool is_used;        // referenced
-    bool is_dead;        // used, but deallocate VA objects as soon as possible
-    int  order;          // for LRU allocation
 } va_surface_priv_t;
 
-struct va_surface_pool * va_surface_pool_alloc(struct mp_vaapi_ctx *ctx, int rt_format)
-{
-    struct va_surface_pool *pool = talloc_ptrtype(NULL, pool);
-    *pool = (struct va_surface_pool) {
-        .ctx = ctx,
-        .log = ctx->log,
-        .display = ctx->display,
-        .rt_format = rt_format
-    };
-    return pool;
-}
-
-
-void va_surface_pool_release(struct va_surface_pool *pool)
-{
-    if (!pool)
-        return;
-    va_surface_pool_clear(pool);
-    talloc_free(pool);
-}
-
-void va_surface_pool_releasep(struct va_surface_pool **pool) {
-    if (!pool)
-        return;
-    va_surface_pool_release(*pool);
-    *pool = NULL;
-}
-
-void va_surface_pool_clear(struct va_surface_pool *pool)
-{
-    for (int i=0; i<pool->num_surfaces; ++i) {
-        struct va_surface *s = pool->surfaces[i];
-        if (s->p->is_used)
-            s->p->is_dead = true;
-        else
-            va_surface_destroy(s);
-    }
-    talloc_free(pool->surfaces);
-    pool->num_surfaces = 0;
-}
-
-void va_surface_destroy(struct va_surface *surface)
+static void va_surface_destroy(struct va_surface *surface)
 {
     if (!surface)
         return;
     if (surface->id != VA_INVALID_ID) {
         va_surface_priv_t *p = surface->p;
-        assert(!p->is_used);
         if (p->image.image_id != VA_INVALID_ID)
             vaDestroyImage(p->display, p->image.image_id);
         vaDestroySurfaces(p->display, &surface->id, 1);
@@ -235,93 +179,40 @@ void va_surface_destroy(struct va_surface *surface)
     talloc_free(surface);
 }
 
-void va_surface_release(struct va_surface *surface)
+static void release_va_surface(void *arg)
 {
-    if (!surface)
-        return;
-    surface->p->is_used = false;
-    if (surface->p->is_dead)
-        va_surface_destroy(surface);
+    struct va_surface *surface = arg;
+    va_surface_destroy(surface);
 }
 
-void va_surface_releasep(struct va_surface **surface)
-{
-    if (!surface)
-        return;
-    va_surface_release(*surface);
-    *surface = NULL;
-}
-
-static struct va_surface *va_surface_alloc(struct va_surface_pool *pool,
-                                           int w, int h)
+static struct mp_image *alloc_surface(struct mp_vaapi_ctx *ctx, int rt_format,
+                                      int w, int h)
 {
     VASurfaceID id = VA_INVALID_ID;
     VAStatus status;
-    status = vaCreateSurfaces(pool->display, w, h, pool->rt_format, 1, &id);
-    if (!CHECK_VA_STATUS(pool, "vaCreateSurfaces()"))
+    status = vaCreateSurfaces(ctx->display, w, h, rt_format, 1, &id);
+    if (!CHECK_VA_STATUS(ctx, "vaCreateSurfaces()"))
         return NULL;
 
     struct va_surface *surface = talloc_ptrtype(NULL, surface);
     if (!surface)
         return NULL;
 
-    MP_TARRAY_APPEND(NULL, pool->surfaces, pool->num_surfaces, surface);
     surface->id = id;
     surface->w = w;
     surface->h = h;
-    surface->rt_format = pool->rt_format;
+    surface->rt_format = rt_format;
     surface->p = talloc_zero(surface, va_surface_priv_t);
-    surface->p->ctx = pool->ctx;
-    surface->p->display = pool->display;
+    surface->p->ctx = ctx;
+    surface->p->display = ctx->display;
     surface->p->image.image_id = surface->p->image.buf = VA_INVALID_ID;
-    return surface;
-}
-
-struct mp_image *va_surface_pool_get_wrapped(struct va_surface_pool *pool,
-                                             int imgfmt, int w, int h)
-{
-    return va_surface_wrap(va_surface_pool_get_by_imgfmt(pool, imgfmt, w, h));
-}
-
-int va_surface_pool_rt_format(const struct va_surface_pool *pool)
-{
-    return pool->rt_format;
-}
 
-bool va_surface_pool_reserve(struct va_surface_pool *pool, int count,
-                             int w, int h)
-{
-    for (int i=0; i<pool->num_surfaces && count > 0; ++i) {
-        const struct va_surface *s = pool->surfaces[i];
-        if (s->w == w && s->h == h && !s->p->is_used)
-            --count;
-    }
-    while (count > 0) {
-        if (!va_surface_alloc(pool, w, h))
-            break;
-        --count;
-    }
-    return !count;
-}
-
-struct va_surface *va_surface_pool_get(struct va_surface_pool *pool,
-                                       int w, int h)
-{
-    struct va_surface *best = NULL;
-    for (int i=0; i<pool->num_surfaces; ++i) {
-        struct va_surface *s = pool->surfaces[i];
-        if (!s->p->is_used && s->w == w && s->h == h) {
-            if (!best || best->p->order > s->p->order)
-                best = s;
-        }
-    }
-    if (!best)
-        best = va_surface_alloc(pool, w, h);
-    if (best) {
-        best->p->is_used = true;
-        best->p->order = ++pool->lru_counter;
-    }
-    return best;
+    struct mp_image img = {0};
+    mp_image_setfmt(&img, IMGFMT_VAAPI);
+    mp_image_set_size(&img, surface->w, surface->h);
+    img.planes[0] = (uint8_t*)surface;
+    img.planes[3] = (uint8_t*)(uintptr_t)surface->id;
+    return mp_image_new_custom_ref(&img, surface, release_va_surface);
 }
 
 static void va_surface_image_destroy(struct va_surface *surface)
@@ -369,43 +260,20 @@ static VAImage *va_surface_image_alloc(struct va_surface *surface,
     return &surface->p->image;
 }
 
-
-
-struct va_surface *va_surface_pool_get_by_imgfmt(struct va_surface_pool *pool,
-                                                 int imgfmt, int w, int h)
+// img must be a VAAPI surface; make sure its internal VAImage is allocated
+// to a format corresponding to imgfmt (or return an error).
+int va_surface_image_alloc_imgfmt(struct mp_image *img, int imgfmt)
 {
-    if (imgfmt == IMGFMT_VAAPI)
-        return va_surface_pool_get(pool, w, h);
+    struct va_surface *surface = va_surface_in_mp_image(img);
+    if (!surface)
+        return -1;
     VAImageFormat *format =
-        va_image_format_from_imgfmt(pool->ctx->image_formats, imgfmt);
+        va_image_format_from_imgfmt(surface->p->ctx->image_formats, imgfmt);
     if (!format)
-        return NULL;
-    // WTF: no mapping from VAImageFormat -> VA_RT_FORMAT_
-    struct va_surface *surface = va_surface_pool_get(pool, w, h);
-    if (!surface)
-        return NULL;
-    if (va_surface_image_alloc(surface, format))
-        return surface;
-    va_surface_release(surface);
-    return NULL;
-}
-
-static void free_va_surface(void *arg)
-{
-    va_surface_release((struct va_surface*)arg);
-}
-
-struct mp_image *va_surface_wrap(struct va_surface *surface)
-{
-    if (!surface)
-        return NULL;
-
-    struct mp_image img = {0};
-    mp_image_setfmt(&img, IMGFMT_VAAPI);
-    mp_image_set_size(&img, surface->w, surface->h);
-    img.planes[0] = (uint8_t*)surface;
-    img.planes[3] = (uint8_t*)(uintptr_t)surface->id;
-    return mp_image_new_custom_ref(&img, surface, free_va_surface);
+        return -1;
+    if (!va_surface_image_alloc(surface, format))
+        return -1;
+    return 0;
 }
 
 VASurfaceID va_surface_id_in_mp_image(const struct mp_image *mpi)
@@ -461,10 +329,12 @@ bool va_image_unmap(struct mp_vaapi_ctx *ctx, VAImage *image)
 bool va_surface_upload(struct va_surface *surface, struct mp_image *mpi)
 {
     va_surface_priv_t *p = surface->p;
-    if (p->image.image_id == VA_INVALID_ID)
-        return false;
 
-    if (va_fourcc_to_imgfmt(p->image.format.fourcc) != mpi->imgfmt)
+    VAImageFormat *format =
+        va_image_format_from_imgfmt(p->ctx->image_formats, mpi->imgfmt);
+    if (!format)
+        return false;
+    if (!va_surface_image_alloc(surface, format))
         return false;
 
     struct mp_image img;
@@ -485,6 +355,18 @@ bool va_surface_upload(struct va_surface *surface, struct mp_image *mpi)
     return true;
 }
 
+// va_dst: copy destination, must be IMGFMT_VAAPI
+// sw_src: copy source, must be a software surface
+int va_surface_upload_image(struct mp_image *va_dst, struct mp_image *sw_src)
+{
+    struct va_surface *surface = va_surface_in_mp_image(va_dst);
+    if (!surface)
+        return -1;
+    if (!va_surface_upload(surface, sw_src))
+        return -1;
+    return 0;
+}
+
 static struct mp_image *try_download(struct va_surface *surface,
                                      VAImageFormat *format,
                                      struct mp_image_pool *pool)
@@ -549,3 +431,30 @@ struct mp_image *va_surface_download(struct va_surface *surface,
     return NULL;
 }
 
+struct pool_alloc_ctx {
+    struct mp_vaapi_ctx *vaapi;
+    int rt_format;
+};
+
+static struct mp_image *alloc_pool(void *pctx, int fmt, int w, int h)
+{
+    struct pool_alloc_ctx *alloc_ctx = pctx;
+    if (fmt != IMGFMT_VAAPI)
+        return NULL;
+
+    return alloc_surface(alloc_ctx->vaapi, alloc_ctx->rt_format, w, h);
+}
+
+// The allocator of the given image pool to allocate VAAPI surfaces, using
+// the given rt_format.
+void va_pool_set_allocator(struct mp_image_pool *pool, struct mp_vaapi_ctx *ctx,
+                           int rt_format)
+{
+    struct pool_alloc_ctx *alloc_ctx = talloc_ptrtype(pool, alloc_ctx);
+    *alloc_ctx = (struct pool_alloc_ctx){
+        .vaapi = ctx,
+        .rt_format = rt_format,
+    };
+    mp_image_pool_set_allocator(pool, alloc_pool, alloc_ctx);
+    mp_image_pool_set_lru(pool);
+}
diff --git a/video/vaapi.h b/video/vaapi.h
index 1136725af3..0396d01320 100644
--- a/video/vaapi.h
+++ b/video/vaapi.h
@@ -106,24 +106,17 @@ VAImageFormat *          va_image_format_from_imgfmt(const struct va_image_forma
 bool                     va_image_map(struct mp_vaapi_ctx *ctx, VAImage *image, struct mp_image *mpi);
 bool                     va_image_unmap(struct mp_vaapi_ctx *ctx, VAImage *image);
 
-struct va_surface_pool * va_surface_pool_alloc(struct mp_vaapi_ctx *ctx, int rt_format);
-void                     va_surface_pool_release(struct va_surface_pool *pool);
-void                     va_surface_pool_releasep(struct va_surface_pool **pool);
-void                     va_surface_pool_clear(struct va_surface_pool *pool);
-bool                     va_surface_pool_reserve(struct va_surface_pool *pool, int count, int w, int h);
-int                      va_surface_pool_rt_format(const struct va_surface_pool *pool);
-struct va_surface *      va_surface_pool_get(struct va_surface_pool *pool, int w, int h);
-struct va_surface *      va_surface_pool_get_by_imgfmt(struct va_surface_pool *pool, int imgfmt, int w, int h);
-struct mp_image *        va_surface_pool_get_wrapped(struct va_surface_pool *pool, int imgfmt, int w, int h);
-
-void                     va_surface_release(struct va_surface *surface);
-void                     va_surface_releasep(struct va_surface **surface);
+void va_pool_set_allocator(struct mp_image_pool *pool, struct mp_vaapi_ctx *ctx,
+                           int rt_format);
+
 struct va_surface *      va_surface_in_mp_image(struct mp_image *mpi);
-struct mp_image *        va_surface_wrap(struct va_surface *surface); // takes ownership
 VASurfaceID              va_surface_id(const struct va_surface *surface);
 VASurfaceID              va_surface_id_in_mp_image(const struct mp_image *mpi);
 bool                     va_surface_upload(struct va_surface *surface, struct mp_image *mpi);
 struct mp_image *        va_surface_download(struct va_surface *surface,
                                              struct mp_image_pool *pool);
 
+int va_surface_image_alloc_imgfmt(struct mp_image *img, int imgfmt);
+int va_surface_upload_image(struct mp_image *va_dst, struct mp_image *sw_src);
+
 #endif