Use shared pointer start

src/torchcodec/decoders/_core/FFMPEGCommon.cpp

@@ -10,6 +10,25 @@
 
 namespace facebook::torchcodec {
 
+AVPacketHandler::AVPacketHandler() : avPacket_(av_packet_alloc()) {
+  TORCH_CHECK(avPacket_ != nullptr, "Couldn't allocate avPacket.");
+}
+AVPacketHandler::~AVPacketHandler() {
+  av_packet_free(&avPacket_);
+}
+
+ReferencedAVPacket::ReferencedAVPacket(AVPacketHandler& shared)
+    : avPacket_(shared.avPacket_) {}
+ReferencedAVPacket::~ReferencedAVPacket() {
+  av_packet_unref(avPacket_);
+}
+AVPacket* ReferencedAVPacket::get() {
+  return avPacket_;
+}
+AVPacket* ReferencedAVPacket::operator->() {
+  return avPacket_;
+}
+
 AVCodecOnlyUseForCallingAVFindBestStream
 makeAVCodecOnlyUseForCallingAVFindBestStream(const AVCodec* codec) {
 #if LIBAVCODEC_VERSION_INT < AV_VERSION_INT(59, 18, 100)

src/torchcodec/decoders/_core/FFMPEGCommon.h

@@ -57,8 +57,6 @@ using UniqueAVCodecContext = std::unique_ptr<
     Deleterp<AVCodecContext, void, avcodec_free_context>>;
 using UniqueAVFrame =
     std::unique_ptr<AVFrame, Deleterp<AVFrame, void, av_frame_free>>;
-using UniqueAVPacket =
-    std::unique_ptr<AVPacket, Deleterp<AVPacket, void, av_packet_free>>;
 using UniqueAVFilterGraph = std::unique_ptr<
     AVFilterGraph,
     Deleterp<AVFilterGraph, void, avfilter_graph_free>>;
@@ -70,6 +68,43 @@ using UniqueAVIOContext = std::
 using UniqueSwsContext =
     std::unique_ptr<SwsContext, Deleter<SwsContext, void, sws_freeContext>>;
 
+// These 2 classes are meant to be used in tandem, like so:
+// ```
+// AVPacketHandler avPacketHandler;
+// while(...){
+//   ReferencedAVPacket packet(avPacketHandler);
+//   av_read_frame(..., packet.get());
+// }
+// ```
+// This achieves a few desirable things:
+// - Memory allocation of the underlying AVPacket happens only once, when the
+//   avPacketHandler created.
+// - av_packet_free() is called when avPacketHandler gets out of scope
+// - av_packet_unref() is automatically called when needed, i.e. at the end of
+//   each loop iteration (or when hitting break / continue). This prevents the
+//   risk of us forgetting to call it.
+class AVPacketHandler {
+  friend class ReferencedAVPacket;
+
+ private:
+  AVPacket* avPacket_;
+
+ public:
+  AVPacketHandler();
+  ~AVPacketHandler();
+};
+
+class ReferencedAVPacket {
+ private:
+  AVPacket* avPacket_;
+
+ public:
+  ReferencedAVPacket(AVPacketHandler& shared);
+  ~ReferencedAVPacket();
+  AVPacket* get();
+  AVPacket* operator->();
+};
+
 // av_find_best_stream is not const-correct before commit:
 // https://github.com/FFmpeg/FFmpeg/commit/46dac8cf3d250184ab4247809bc03f60e14f4c0c
 // which was released in FFMPEG version=5.0.3

src/torchcodec/decoders/_core/VideoDecoder.cpp

@@ -573,8 +573,10 @@ void VideoDecoder::scanFileAndUpdateMetadataAndIndex() {
     return;
   }
 
-  UniqueAVPacket packet(av_packet_alloc());
+  AVPacketHandler avPacketHandler;
   while (true) {
+    ReferencedAVPacket packet(avPacketHandler);
+
     // av_read_frame is a misleading name: it gets the next **packet**.
     int ffmpegStatus = av_read_frame(formatContext_.get(), packet.get());
 
@@ -589,7 +591,6 @@ void VideoDecoder::scanFileAndUpdateMetadataAndIndex() {
     }
 
     if (packet->flags & AV_PKT_FLAG_DISCARD) {
-      av_packet_unref(packet.get());
       continue;
     }
 
@@ -612,7 +613,6 @@ void VideoDecoder::scanFileAndUpdateMetadataAndIndex() {
       streams_[streamIndex].keyFrames.push_back(frameInfo);
     }
     streams_[streamIndex].allFrames.push_back(frameInfo);
-    av_packet_unref(packet.get());
   }
 
   // Set all per-stream metadata that requires knowing the content of all
@@ -806,7 +806,7 @@ VideoDecoder::RawDecodedOutput VideoDecoder::getDecodedOutputWithFilter(
   }
   // Need to get the next frame or error from PopFrame.
   UniqueAVFrame frame(av_frame_alloc());
-  UniqueAVPacket packet(av_packet_alloc());
+  AVPacketHandler avPacketHandler;
   int ffmpegStatus = AVSUCCESS;
   bool reachedEOF = false;
   int frameStreamIndex = -1;
@@ -846,6 +846,7 @@ VideoDecoder::RawDecodedOutput VideoDecoder::getDecodedOutputWithFilter(
       // pulling frames from its internal buffers.
       continue;
     }
+    ReferencedAVPacket packet(avPacketHandler);
     ffmpegStatus = av_read_frame(formatContext_.get(), packet.get());
     decodeStats_.numPacketsRead++;
     if (ffmpegStatus == AVERROR_EOF) {
@@ -863,7 +864,6 @@ VideoDecoder::RawDecodedOutput VideoDecoder::getDecodedOutputWithFilter(
         }
       }
       reachedEOF = true;
-      av_packet_unref(packet.get());
       continue;
     }
     if (ffmpegStatus < AVSUCCESS) {
@@ -873,7 +873,6 @@ VideoDecoder::RawDecodedOutput VideoDecoder::getDecodedOutputWithFilter(
     }
     if (activeStreamIndices_.count(packet->stream_index) == 0) {
       // This packet is not for any of the active streams.
-      av_packet_unref(packet.get());
       continue;
     }
     ffmpegStatus = avcodec_send_packet(
@@ -884,7 +883,6 @@ VideoDecoder::RawDecodedOutput VideoDecoder::getDecodedOutputWithFilter(
           getFFMPEGErrorStringFromErrorCode(ffmpegStatus));
     }
     decodeStats_.numPacketsSentToDecoder++;
-    av_packet_unref(packet.get());
   }
   if (ffmpegStatus < AVSUCCESS) {
     if (reachedEOF || ffmpegStatus == AVERROR_EOF) {