1 files changed, 91 insertions, 0 deletions

diff --git a/video/out/vo.c b/video/out/vo.c
index 79fc4f3bb4..e52495e195 100644
--- a/video/out/vo.c
+++ b/video/out/vo.c
@@ -23,9 +23,12 @@
 #include <pthread.h>
 #include <math.h>
 
+#include <libavutil/buffer.h>
+
 #include "mpv_talloc.h"
 
 #include "config.h"
+#include "osdep/atomic.h"
 #include "osdep/timer.h"
 #include "osdep/threads.h"
 #include "misc/dispatch.h"
@@ -113,6 +116,8 @@ struct vo_internal {
     pthread_t thread;
     struct mp_dispatch_queue *dispatch;
 
+    atomic_ullong dr_in_flight;
+
     // --- The following fields are protected by lock
     pthread_mutex_t lock;
     pthread_cond_t wakeup;
@@ -955,6 +960,7 @@ static void *vo_thread(void *ptr)
     talloc_free(in->current_frame);
     in->current_frame = NULL;
     vo->driver->uninit(vo);
+    assert(atomic_load(&vo->in->dr_in_flight) == 0);
     return NULL;
 }
 
@@ -1258,3 +1264,88 @@ int lookup_keymap_table(const struct mp_keymap *map, int key)
         map++;
     return map->to;
 }
+
+struct free_dr_context {
+    struct vo *vo;
+    AVBufferRef *ref;
+};
+
+static void vo_thread_free(void *ptr)
+{
+    struct free_dr_context *ctx = ptr;
+
+    unsigned long long v = atomic_fetch_add(&ctx->vo->in->dr_in_flight, -1);
+    assert(v); // value before sub is 0 - unexpected underflow.
+
+    av_buffer_unref(&ctx->ref);
+    talloc_free(ctx);
+}
+
+static void free_dr_buffer_on_vo_thread(void *opaque, uint8_t *data)
+{
+    struct free_dr_context *ctx = opaque;
+
+    // The image could be unreffed even on the VO thread. In practice, this
+    // matters most on VO destruction.
+    if (pthread_equal(ctx->vo->in->thread, pthread_self())) {
+        vo_thread_free(ctx);
+    } else {
+        mp_dispatch_run(ctx->vo->in->dispatch, vo_thread_free, ctx);
+    }
+}
+
+struct get_image_cmd {
+    struct vo *vo;
+    int imgfmt, w, h, stride_align;
+    struct mp_image *res;
+};
+
+static void sync_get_image(void *ptr)
+{
+    struct get_image_cmd *cmd = ptr;
+    struct vo *vo = cmd->vo;
+
+    cmd->res = vo->driver->get_image(vo, cmd->imgfmt, cmd->w, cmd->h,
+                                     cmd->stride_align);
+    if (!cmd->res)
+        return;
+
+    // We require exactly 1 AVBufferRef.
+    assert(cmd->res->bufs[0]);
+    assert(!cmd->res->bufs[1]);
+
+    // Apply some magic to get it free'd on the VO thread as well. For this to
+    // work, we create a dummy-ref that aliases the original ref, which is why
+    // the original ref must be writable in the first place. (A newly allocated
+    // image should be always writable of course.)
+    assert(mp_image_is_writeable(cmd->res));
+
+    struct free_dr_context *ctx = talloc_zero(NULL, struct free_dr_context);
+    *ctx = (struct free_dr_context){
+        .vo = vo,
+        .ref = cmd->res->bufs[0],
+    };
+
+    AVBufferRef *new_ref = av_buffer_create(ctx->ref->data, ctx->ref->size,
+                                            free_dr_buffer_on_vo_thread, ctx, 0);
+    if (!new_ref)
+        abort(); // tiny malloc OOM
+
+    cmd->res->bufs[0] = new_ref;
+
+    atomic_fetch_add(&vo->in->dr_in_flight, 1);
+}
+
+struct mp_image *vo_get_image(struct vo *vo, int imgfmt, int w, int h,
+                              int stride_align)
+{
+    if (!vo->driver->get_image)
+        return NULL;
+
+    struct get_image_cmd cmd = {
+        .vo = vo,
+        .imgfmt = imgfmt, .w = w, .h = h, .stride_align = stride_align,
+    };
+    mp_dispatch_run(vo->in->dispatch, sync_get_image, &cmd);
+    return cmd.res;
+}