Properly free AVFormatContext and streams

Author: NicolasHug

src/torchcodec/decoders/_core/FFMPEGCommon.h

@@ -49,9 +49,12 @@ struct Deleter {
 };
 
 // Unique pointers for FFMPEG structures.
-using UniqueAVFormatContext = std::unique_ptr<
+using UniqueAVFormatContextForDecoding = std::unique_ptr<
     AVFormatContext,
     Deleterp<AVFormatContext, void, avformat_close_input>>;
+using UniqueAVFormatContextForEncoding = std::unique_ptr<
+    AVFormatContext,
+    Deleter<AVFormatContext, void, avformat_free_context>>;
 using UniqueAVCodecContext = std::unique_ptr<
     AVCodecContext,
     Deleterp<AVCodecContext, void, avcodec_free_context>>;

src/torchcodec/decoders/_core/VideoDecoder.cpp

@@ -41,7 +41,7 @@ int64_t secondsToClosestPts(double seconds, const AVRational& timeBase) {
 }
 
 struct AVInput {
-  UniqueAVFormatContext formatContext;
+  UniqueAVFormatContextForDecoding formatContext;
   std::unique_ptr<AVIOBytesContext> ioBytesContext;
 };
 
@@ -1695,10 +1695,7 @@ FrameDims getHeightAndWidthFromOptionsOrAVFrame(
       videoStreamOptions.width.value_or(avFrame.width));
 }
 
-Encoder::~Encoder() {
-  // TODO NEED TO CALL THIS
-  //   avformat_free_context(avFormatContext_.get());
-}
+Encoder::~Encoder() {}
 
 Encoder::Encoder(int sampleRate, std::string_view fileName)
     : sampleRate_(sampleRate) {
@@ -1756,6 +1753,9 @@ Encoder::Encoder(int sampleRate, std::string_view fileName)
       avCodecContext_->frame_size,
       ". Cannot encode. This should probably never happen?");
 
+  // We're allocating the stream here. Streams are meant to be freed by
+  // avformat_free_context(avFormatContext), which we call in the
+  // avFormatContext_'s destructor.
   avStream_ = avformat_new_stream(avFormatContext_.get(), NULL);
   TORCH_CHECK(avStream_ != nullptr, "Couldn't create new stream.");
   avcodec_parameters_from_context(avStream_->codecpar, avCodecContext_.get());
@@ -1829,10 +1829,10 @@ torch::Tensor Encoder::encode(const torch::Tensor& wf) {
     avFrame->pts += avFrame->nb_samples;
     numEncodedSamples += numSamplesToEncode;
   }
-  encode_inner_loop(autoAVPacket, nullptr); // flush
-
   TORCH_CHECK(numEncodedSamples == numSamples, "Hmmmmmm something went wrong.");
 
+  encode_inner_loop(autoAVPacket, nullptr); // flush
+
   ffmpegRet = av_write_trailer(avFormatContext_.get());
   TORCH_CHECK(
       ffmpegRet == AVSUCCESS,

src/torchcodec/decoders/_core/VideoDecoder.h

@@ -453,7 +453,7 @@ class VideoDecoder {
 
   SeekMode seekMode_;
   ContainerMetadata containerMetadata_;
-  UniqueAVFormatContext formatContext_;
+  UniqueAVFormatContextForDecoding formatContext_;
   std::map<int, StreamInfo> streamInfos_;
   const int NO_ACTIVE_STREAM = -2;
   int activeStreamIndex_ = NO_ACTIVE_STREAM;
@@ -559,7 +559,7 @@ class Encoder {
  private:
   void encode_inner_loop(AutoAVPacket& autoAVPacket, AVFrame* avFrame);
 
-  UniqueAVFormatContext avFormatContext_;
+  UniqueAVFormatContextForEncoding avFormatContext_;
   UniqueAVCodecContext avCodecContext_;
   AVStream* avStream_;