Fix CVPixelBuffer/CMSampleBuffer ownership issues

src/platform/macos/av_img_t.h

@@ -10,10 +10,43 @@
 #include <CoreVideo/CoreVideo.h>
 
 namespace platf {
-  struct av_img_t: public img_t {
-    CVPixelBufferRef pixel_buffer = nullptr;
-    CMSampleBufferRef sample_buffer = nullptr;
+  struct av_sample_buf_t {
+    av_sample_buf_t(CMSampleBufferRef buf):
+        buf((CMSampleBufferRef) CFRetain(buf)) {}
 
-    ~av_img_t();
+    ~av_sample_buf_t() {
+      CFRelease(buf);
+    }
+
+    CMSampleBufferRef buf;
+  };
+
+  struct av_pixel_buf_t {
+    av_pixel_buf_t(CVPixelBufferRef buf):
+        buf((CVPixelBufferRef) CFRetain(buf)),
+        locked(false) {}
+
+    uint8_t *
+    lock() {
+      if (!locked) {
+        CVPixelBufferLockBaseAddress(buf, kCVPixelBufferLock_ReadOnly);
+      }
+      return (uint8_t *) CVPixelBufferGetBaseAddress(buf);
+    }
+
+    ~av_pixel_buf_t() {
+      if (locked) {
+        CVPixelBufferUnlockBaseAddress(buf, kCVPixelBufferLock_ReadOnly);
+      }
+      CFRelease(buf);
+    }
+
+    CVPixelBufferRef buf;
+    bool locked;
+  };
+
+  struct av_img_t: public img_t {
+    std::shared_ptr<av_sample_buf_t> sample_buffer;
+    std::shared_ptr<av_pixel_buf_t> pixel_buffer;
   };
 }  // namespace platf

src/platform/macos/display.mm

@@ -20,18 +20,6 @@ namespace fs = std::filesystem;
 namespace platf {
   using namespace std::literals;
 
-  av_img_t::~av_img_t() {
-    if (pixel_buffer != NULL) {
-      CVPixelBufferUnlockBaseAddress(pixel_buffer, 0);
-    }
-
-    if (sample_buffer != nullptr) {
-      CFRelease(sample_buffer);
-    }
-
-    data = nullptr;
-  }
-
   struct av_display_t: public display_t {
     AVVideo *av_capture;
     CGDirectDisplayID display_id;
@@ -43,11 +31,6 @@ namespace platf {
     capture_e
     capture(const push_captured_image_cb_t &push_captured_image_cb, const pull_free_image_cb_t &pull_free_image_cb, bool *cursor) override {
       auto signal = [av_capture capture:^(CMSampleBufferRef sampleBuffer) {
-        CFRetain(sampleBuffer);
-
-        CVPixelBufferRef pixelBuffer = CMSampleBufferGetImageBuffer(sampleBuffer);
-        CVPixelBufferLockBaseAddress(pixelBuffer, kCVPixelBufferLock_ReadOnly);
-
         std::shared_ptr<img_t> img_out;
         if (!pull_free_image_cb(img_out)) {
           // got interrupt signal
@@ -56,16 +39,11 @@ namespace platf {
         }
         auto av_img = std::static_pointer_cast<av_img_t>(img_out);
 
-        if (av_img->pixel_buffer != nullptr)
-          CVPixelBufferUnlockBaseAddress(av_img->pixel_buffer, 0);
-
-        if (av_img->sample_buffer != nullptr)
-          CFRelease(av_img->sample_buffer);
-
-        av_img->sample_buffer = sampleBuffer;
-        av_img->pixel_buffer = pixelBuffer;
-        img_out->data = (uint8_t *) CVPixelBufferGetBaseAddress(pixelBuffer);
+        CVPixelBufferRef pixelBuffer = CMSampleBufferGetImageBuffer(sampleBuffer);
 
+        av_img->sample_buffer = std::make_shared<av_sample_buf_t>(sampleBuffer);
+        av_img->pixel_buffer = std::make_shared<av_pixel_buf_t>(pixelBuffer);
+        img_out->data = av_img->pixel_buffer->lock();
 
         img_out->width = CVPixelBufferGetWidth(pixelBuffer);
         img_out->height = CVPixelBufferGetHeight(pixelBuffer);
@@ -117,23 +95,11 @@ namespace platf {
       auto signal = [av_capture capture:^(CMSampleBufferRef sampleBuffer) {
         auto av_img = (av_img_t *) img;
 
-        CFRetain(sampleBuffer);
-
         CVPixelBufferRef pixelBuffer = CMSampleBufferGetImageBuffer(sampleBuffer);
-        CVPixelBufferLockBaseAddress(pixelBuffer, kCVPixelBufferLock_ReadOnly);
-
-        // XXX: next_img->img should be moved to a smart pointer with
-        // the CFRelease as custom deallocator
-        if (av_img->pixel_buffer != nullptr)
-          CVPixelBufferUnlockBaseAddress(((av_img_t *) img)->pixel_buffer, 0);
-
-        if (av_img->sample_buffer != nullptr)
-          CFRelease(av_img->sample_buffer);
-
-        av_img->sample_buffer = sampleBuffer;
-        av_img->pixel_buffer = pixelBuffer;
-        img->data = (uint8_t *) CVPixelBufferGetBaseAddress(pixelBuffer);
 
+        av_img->sample_buffer = std::make_shared<av_sample_buf_t>(sampleBuffer);
+        av_img->pixel_buffer = std::make_shared<av_pixel_buf_t>(pixelBuffer);
+        img->data = av_img->pixel_buffer->lock();
 
         img->width = CVPixelBufferGetWidth(pixelBuffer);
         img->height = CVPixelBufferGetHeight(pixelBuffer);

src/platform/macos/nv12_zero_device.cpp

@@ -3,7 +3,6 @@
  * @brief todo
  */
 #include "src/platform/macos/nv12_zero_device.h"
-#include "src/platform/macos/av_img_t.h"
 
 #include "src/video.h"
 
@@ -18,8 +17,6 @@ namespace platf {
     av_frame_free(&frame);
   }
 
-  util::safe_ptr<AVFrame, free_frame> av_frame;
-
   int
   nv12_zero_device::convert(platf::img_t &img) {
     av_frame_make_writable(av_frame.get());
@@ -33,6 +30,10 @@ namespace platf {
       av_frame->data[i] = (uint8_t *) CVPixelBufferGetBaseAddressOfPlane(av_img->pixel_buffer->buf, i);
     }
 
+    // We just set data pointers to point into our CVPixelBuffer above, so we have to hold
+    // a reference to these buffers to keep them around until the AVFrame is done using them.
+    backing_img.sample_buffer = av_img->sample_buffer;
+    backing_img.pixel_buffer = av_img->pixel_buffer;
 
     return 0;
   }

src/platform/macos/nv12_zero_device.h

@@ -5,8 +5,13 @@
 #pragma once
 
 #include "src/platform/common.h"
+#include "src/platform/macos/av_img_t.h"
+
+struct AVFrame;
 
 namespace platf {
+  void
+  free_frame(AVFrame *frame);
 
   class nv12_zero_device: public avcodec_encode_device_t {
     // display holds a pointer to an av_video object. Since the namespaces of AVFoundation
@@ -27,6 +32,10 @@ namespace platf {
     convert(img_t &img);
     int
     set_frame(AVFrame *frame, AVBufferRef *hw_frames_ctx);
+
+  private:
+    util::safe_ptr<AVFrame, free_frame> av_frame;
+    av_img_t backing_img;
   };
 
 }  // namespace platf