reset hw_frame to fix memory release

commaai · Nov 14, 2021 · cc74800 · cc74800
1 parent 834b90d
commit cc74800
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 14 deletions.
diff --git a/selfdrive/ui/replay/framereader.cc b/selfdrive/ui/replay/framereader.cc
@@ -15,10 +15,9 @@ int readFunction(void *opaque, uint8_t *buf, int buf_size) {
 
 FrameReader::FrameReader(bool local_cache, int chunk_size, int retries) : FileReader(local_cache, chunk_size, retries) {
   input_ctx = avformat_alloc_context();
-  av_frame_ = av_frame_alloc();
-  rgb_frame_ = av_frame_alloc();
-  yuv_frame_ = av_frame_alloc();
-  hw_frame = av_frame_alloc();
+  av_frame_.reset(av_frame_alloc());
+  rgb_frame_.reset(av_frame_alloc());
+  yuv_frame_.reset(av_frame_alloc());
 }
 
 FrameReader::~FrameReader() {
@@ -30,10 +29,6 @@ FrameReader::~FrameReader() {
   if (input_ctx) avformat_close_input(&input_ctx);
   if (hw_device_ctx) av_buffer_unref(&hw_device_ctx);
 
-  if (av_frame_) av_frame_free(&av_frame_);
-  if (rgb_frame_) av_frame_free(&rgb_frame_);
-  if (yuv_frame_) av_frame_free(&yuv_frame_);
-  if (hw_frame) av_frame_free(&hw_frame);
   if (rgb_sws_ctx_) sws_freeContext(rgb_sws_ctx_);
   if (yuv_sws_ctx_) sws_freeContext(yuv_sws_ctx_);
 
@@ -181,19 +176,21 @@ AVFrame *FrameReader::decodeFrame(AVPacket *pkt) {
     printf("Error sending a packet for decoding\n");
     return nullptr;
   }
-  ret = avcodec_receive_frame(decoder_ctx, av_frame_);
+
+  ret = avcodec_receive_frame(decoder_ctx, av_frame_.get());
   if (ret != 0) {
     return nullptr;
   }
 
   if (av_frame_->format == hw_pix_fmt) {
-    if ((ret = av_hwframe_transfer_data(hw_frame, av_frame_, 0)) < 0) {
+    hw_frame.reset(av_frame_alloc());
+    if ((ret = av_hwframe_transfer_data(hw_frame.get(), av_frame_.get(), 0)) < 0) {
       printf("error transferring the data from GPU to CPU\n");
       return nullptr;
     }
-    return hw_frame;
+    return hw_frame.get();
   } else {
-    return av_frame_;
+    return av_frame_.get();
   }
 }
 

diff --git a/selfdrive/ui/replay/framereader.h b/selfdrive/ui/replay/framereader.h
@@ -1,5 +1,6 @@
 #pragma once
 
+#include <memory>
 #include <string>
 #include <vector>
 
@@ -12,6 +13,10 @@ extern "C" {
 #include <libavutil/imgutils.h>
 }
 
+struct AVFrameDeleter {
+  void operator()(AVFrame* frame) const { av_frame_free(&frame); }
+};
+
 class FrameReader : protected FileReader {
 public:
   FrameReader(bool local_cache = false, int chunk_size = -1, int retries = 0);
@@ -38,7 +43,7 @@ class FrameReader : protected FileReader {
   };
   std::vector<Frame> frames_;
   SwsContext *rgb_sws_ctx_ = nullptr, *yuv_sws_ctx_ = nullptr;
-  AVFrame *av_frame_, *rgb_frame_, *yuv_frame_;
+  std::unique_ptr<AVFrame, AVFrameDeleter>av_frame_, rgb_frame_, yuv_frame_, hw_frame;
   AVFormatContext *input_ctx = nullptr;
   AVCodecContext *decoder_ctx = nullptr;
   int key_frames_count_ = 0;
@@ -48,5 +53,4 @@ class FrameReader : protected FileReader {
   AVHWDeviceType hw_device_type = AV_HWDEVICE_TYPE_NONE;
   AVPixelFormat hw_pix_fmt = AV_PIX_FMT_NONE;
   AVBufferRef *hw_device_ctx = nullptr;
-  AVFrame *hw_frame = nullptr;
 };