replay/framereader: no longer cache all AVPacket instances in memory (#32236)

tools/replay/framereader.cc

@@ -1,9 +1,8 @@
 #include "tools/replay/framereader.h"
-#include "tools/replay/util.h"
 
-#include <cassert>
+#include "common/util.h"
-#include <algorithm>
 #include "third_party/libyuv/include/libyuv.h"
+#include "tools/replay/util.h"
 
 #ifdef __APPLE__
 #define HW_DEVICE_TYPE AV_HWDEVICE_TYPE_VIDEOTOOLBOX
@@ -13,24 +12,8 @@
 #define HW_PIX_FMT AV_PIX_FMT_CUDA
 #endif
 
-namespace {
 
-struct buffer_data {
+namespace {
-  const uint8_t *data;
-  int64_t offset;
-  size_t size;
-};
-
-int readPacket(void *opaque, uint8_t *buf, int buf_size) {
-  struct buffer_data *bd = (struct buffer_data *)opaque;
-  assert(bd->offset <= bd->size);
-  buf_size = std::min((size_t)buf_size, (size_t)(bd->size - bd->offset));
-  if (!buf_size) return AVERROR_EOF;
-
-  memcpy(buf, bd->data + bd->offset, buf_size);
-  bd->offset += buf_size;
-  return buf_size;
-}
 
 enum AVPixelFormat get_hw_format(AVCodecContext *ctx, const enum AVPixelFormat *pix_fmts) {
   enum AVPixelFormat *hw_pix_fmt = reinterpret_cast<enum AVPixelFormat *>(ctx->opaque);
@@ -50,101 +33,61 @@ FrameReader::FrameReader() {
 }
 
 FrameReader::~FrameReader() {
-  for (AVPacket *pkt : packets) {
-    av_packet_free(&pkt);
-  }
-
   if (decoder_ctx) avcodec_free_context(&decoder_ctx);
   if (input_ctx) avformat_close_input(&input_ctx);
   if (hw_device_ctx) av_buffer_unref(&hw_device_ctx);
-
-  if (avio_ctx_) {
-    av_freep(&avio_ctx_->buffer);
-    avio_context_free(&avio_ctx_);
-  }
 }
 
 bool FrameReader::load(const std::string &url, bool no_hw_decoder, std::atomic<bool> *abort, bool local_cache, int chunk_size, int retries) {
-  FileReader f(local_cache, chunk_size, retries);
+  auto local_file_path = url.find("https://") == 0 ? cacheFilePath(url) : url;
-  std::string data = f.read(url, abort);
+  if (!util::file_exists(local_file_path)) {
-  if (data.empty()) {
+    FileReader f(local_cache, chunk_size, retries);
-    rWarning("URL %s returned no data", url.c_str());
+    if (f.read(url, abort).empty()) {
-    return false;
+      return false;
+    }
   }
-
+  return loadFromFile(local_file_path, no_hw_decoder, abort);
-  return load((std::byte *)data.data(), data.size(), no_hw_decoder, abort);
 }
 
-bool FrameReader::load(const std::byte *data, size_t size, bool no_hw_decoder, std::atomic<bool> *abort) {
+bool FrameReader::loadFromFile(const std::string &file, bool no_hw_decoder, std::atomic<bool> *abort) {
-  input_ctx = avformat_alloc_context();
+  if (avformat_open_input(&input_ctx, file.c_str(), nullptr, nullptr) != 0 ||
-  if (!input_ctx) {
+      avformat_find_stream_info(input_ctx, nullptr) < 0) {
-    rError("Error calling avformat_alloc_context");
+    rError("Failed to open input file or find video stream");
     return false;
   }
-
-  struct buffer_data bd = {
-    .data = (const uint8_t*)data,
-    .offset = 0,
-    .size = size,
-  };
-  const int avio_ctx_buffer_size = 64 * 1024;
-  unsigned char *avio_ctx_buffer = (unsigned char *)av_malloc(avio_ctx_buffer_size);
-  avio_ctx_ = avio_alloc_context(avio_ctx_buffer, avio_ctx_buffer_size, 0, &bd, readPacket, nullptr, nullptr);
-  input_ctx->pb = avio_ctx_;
-
   input_ctx->probesize = 10 * 1024 * 1024;  // 10MB
-  int ret = avformat_open_input(&input_ctx, nullptr, nullptr, nullptr);
-  if (ret != 0) {
-    char err_str[1024] = {0};
-    av_strerror(ret, err_str, std::size(err_str));
-    rError("Error loading video - %s", err_str);
-    return false;
-  }
-
-  ret = avformat_find_stream_info(input_ctx, nullptr);
-  if (ret < 0) {
-    rError("cannot find a video stream in the input file");
-    return false;
-  }
 
   AVStream *video = input_ctx->streams[0];
   const AVCodec *decoder = avcodec_find_decoder(video->codecpar->codec_id);
   if (!decoder) return false;
 
   decoder_ctx = avcodec_alloc_context3(decoder);
-  ret = avcodec_parameters_to_context(decoder_ctx, video->codecpar);
+  if (!decoder_ctx || avcodec_parameters_to_context(decoder_ctx, video->codecpar) != 0) {
-  if (ret != 0) return false;
+    rError("Failed to allocate or initialize codec context");
+    return false;
+  }
 
   width = (decoder_ctx->width + 3) & ~3;
   height = decoder_ctx->height;
 
-  if (has_hw_decoder && !no_hw_decoder) {
+  if (has_hw_decoder && !no_hw_decoder && !initHardwareDecoder(HW_DEVICE_TYPE)) {
-    if (!initHardwareDecoder(HW_DEVICE_TYPE)) {
+    rWarning("No device with hardware decoder found. fallback to CPU decoding.");
-      rWarning("No device with hardware decoder found. fallback to CPU decoding.");
-    }
   }
 
-  ret = avcodec_open2(decoder_ctx, decoder, nullptr);
+  if (avcodec_open2(decoder_ctx, decoder, nullptr) < 0) {
-  if (ret < 0) {
+    rError("Failed to open codec");
-    rError("avcodec_open2 failed %d", ret);
     return false;
   }
 
-  packets.reserve(60 * 20);  // 20fps, one minute
+  AVPacket pkt;
-  while (!(abort && *abort)) {
+  packets_info.reserve(60 * 20);  // 20fps, one minute
-    AVPacket *pkt = av_packet_alloc();
+  while (!(abort && *abort) && av_read_frame(input_ctx, &pkt) == 0) {
-    ret = av_read_frame(input_ctx, pkt);
+    packets_info.emplace_back(PacketInfo{.flags = pkt.flags, .pos = pkt.pos});
-    if (ret < 0) {
+    av_packet_unref(&pkt);
-      av_packet_free(&pkt);
-      valid_ = (ret == AVERROR_EOF);
-      break;
-    }
-    packets.push_back(pkt);
-    // some stream seems to contain no keyframes
-    key_frames_count_ += pkt->flags & AV_PKT_FLAG_KEY;
   }
-  valid_ = valid_ && !packets.empty();
+
-  return valid_;
+  avio_seek(input_ctx->pb, 0, SEEK_SET);
+  return !packets_info.empty();
 }
 
 bool FrameReader::initHardwareDecoder(AVHWDeviceType hw_device_type) {
@@ -176,8 +119,7 @@ bool FrameReader::initHardwareDecoder(AVHWDeviceType hw_device_type) {
 }
 
 bool FrameReader::get(int idx, VisionBuf *buf) {
-  assert(buf != nullptr);
+  if (!buf || idx < 0 || idx >= packets_info.size()) {
-  if (!valid_ || idx < 0 || idx >= packets.size()) {
     return false;
   }
   return decode(idx, buf);
@@ -185,24 +127,30 @@ bool FrameReader::get(int idx, VisionBuf *buf) {
 
 bool FrameReader::decode(int idx, VisionBuf *buf) {
   int from_idx = idx;
-  if (idx != prev_idx + 1 && key_frames_count_ > 1) {
+  if (idx != prev_idx + 1) {
     // seeking to the nearest key frame
     for (int i = idx; i >= 0; --i) {
-      if (packets[i]->flags & AV_PKT_FLAG_KEY) {
+      if (packets_info[i].flags & AV_PKT_FLAG_KEY) {
         from_idx = i;
         break;
       }
     }
+    avio_seek(input_ctx->pb, packets_info[from_idx].pos, SEEK_SET);
   }
   prev_idx = idx;
 
+  bool result = false;
+  AVPacket pkt;
   for (int i = from_idx; i <= idx; ++i) {
-    AVFrame *f = decodeFrame(packets[i]);
+    if (av_read_frame(input_ctx, &pkt) == 0) {
-    if (f && i == idx) {
+      AVFrame *f = decodeFrame(&pkt);
-      return copyBuffers(f, buf);
+      if (f && i == idx) {
+        result = copyBuffers(f, buf);
+      }
+      av_packet_unref(&pkt);
     }
   }
-  return false;
+  return result;
 }
 
 AVFrame *FrameReader::decodeFrame(AVPacket *pkt) {
@@ -232,7 +180,6 @@ AVFrame *FrameReader::decodeFrame(AVPacket *pkt) {
 }
 
 bool FrameReader::copyBuffers(AVFrame *f, VisionBuf *buf) {
-  assert(f != nullptr && buf != nullptr);
   if (hw_pix_fmt == HW_PIX_FMT) {
     for (int i = 0; i < height/2; i++) {
       memcpy(buf->y + (i*2 + 0)*buf->stride, f->data[0] + (i*2 + 0)*f->linesize[0], width);

tools/replay/framereader.h

@@ -22,11 +22,9 @@ public:
   ~FrameReader();
   bool load(const std::string &url, bool no_hw_decoder = false, std::atomic<bool> *abort = nullptr, bool local_cache = false,
             int chunk_size = -1, int retries = 0);
-  bool load(const std::byte *data, size_t size, bool no_hw_decoder = false, std::atomic<bool> *abort = nullptr);
+  bool loadFromFile(const std::string &file, bool no_hw_decoder = false, std::atomic<bool> *abort = nullptr);
   bool get(int idx, VisionBuf *buf);
-  int getYUVSize() const { return width * height * 3 / 2; }
+  size_t getFrameCount() const { return packets_info.size(); }
-  size_t getFrameCount() const { return packets.size(); }
-  bool valid() const { return valid_; }
 
   int width = 0, height = 0;
 
@@ -36,16 +34,17 @@ private:
   AVFrame * decodeFrame(AVPacket *pkt);
   bool copyBuffers(AVFrame *f, VisionBuf *buf);
 
-  std::vector<AVPacket*> packets;
   std::unique_ptr<AVFrame, AVFrameDeleter>av_frame_, hw_frame;
   AVFormatContext *input_ctx = nullptr;
   AVCodecContext *decoder_ctx = nullptr;
-  int key_frames_count_ = 0;
-  bool valid_ = false;
-  AVIOContext *avio_ctx_ = nullptr;
 
   AVPixelFormat hw_pix_fmt = AV_PIX_FMT_NONE;
   AVBufferRef *hw_device_ctx = nullptr;
   int prev_idx = -1;
+  struct PacketInfo {
+    int flags;
+    int64_t pos;
+  };
+  std::vector<PacketInfo> packets_info;
   inline static std::atomic<bool> has_hw_decoder = true;
 };