Close avio object *before* encode destruction. And flush, too. (#773)

Author: NicolasHug

src/torchcodec/_core/Encoder.cpp

@@ -99,8 +99,18 @@ AVSampleFormat findBestOutputSampleFormat(const AVCodec& avCodec) {
 } // namespace
 
 AudioEncoder::~AudioEncoder() {
-  if (avFormatContext_ && avFormatContext_->pb && !avioContextHolder_) {
-    avio_close(avFormatContext_->pb);
+  close_avio();
+}
+
+void AudioEncoder::close_avio() {
+  if (avFormatContext_ && avFormatContext_->pb) {
+    avio_flush(avFormatContext_->pb);
+
+    if (!avioContextHolder_) {
+      avio_close(avFormatContext_->pb);
+      // avoids closing again in destructor, which would segfault.
+      avFormatContext_->pb = nullptr;
+    }
   }
 }
 
@@ -308,6 +318,8 @@ void AudioEncoder::encode() {
       status == AVSUCCESS,
       "Error in: av_write_trailer",
       getFFMPEGErrorStringFromErrorCode(status));
+
+  close_avio();
 }
 
 UniqueAVFrame AudioEncoder::maybeConvertAVFrame(const UniqueAVFrame& avFrame) {

src/torchcodec/_core/Encoder.h

@@ -33,6 +33,7 @@ class AudioEncoder {
   void encodeFrame(AutoAVPacket& autoAVPacket, const UniqueAVFrame& avFrame);
   void maybeFlushSwrBuffers(AutoAVPacket& autoAVPacket);
   void flushBuffers();
+  void close_avio();
 
   UniqueEncodingAVFormatContext avFormatContext_;
   UniqueAVCodecContext avCodecContext_;