Refactor EOF packet logic

Author: NicolasHug

src/torchcodec/_core/BetaCudaDeviceInterface.cpp

@@ -331,28 +331,33 @@ int BetaCudaDeviceInterface::streamPropertyChange(CUVIDEOFORMAT* videoFormat) {
 // Moral equivalent of avcodec_send_packet(). Here, we pass the AVPacket down to
 // the NVCUVID parser.
 int BetaCudaDeviceInterface::sendPacket(ReferenceAVPacket& packet) {
+  TORCH_CHECK(
+      packet.get() && packet->data && packet->size > 0,
+      "sendPacket received an empty packet, this is unexpected, please report.");
+
+  applyBSF(packet);
+
   CUVIDSOURCEDATAPACKET cuvidPacket = {};
+  cuvidPacket.payload = packet->data;
+  cuvidPacket.payload_size = packet->size;
+  cuvidPacket.flags = CUVID_PKT_TIMESTAMP;
+  cuvidPacket.timestamp = packet->pts;
 
-  if (packet.get() && packet->data && packet->size > 0) {
-    applyBSF(packet);
+  return sendCuvidPacket(cuvidPacket);
+}
 
-    // Regular packet with data
-    cuvidPacket.payload = packet->data;
-    cuvidPacket.payload_size = packet->size;
-    cuvidPacket.flags = CUVID_PKT_TIMESTAMP;
-    cuvidPacket.timestamp = packet->pts;
+int BetaCudaDeviceInterface::sendEOFPacket() {
+  CUVIDSOURCEDATAPACKET cuvidPacket = {};
+  cuvidPacket.flags = CUVID_PKT_ENDOFSTREAM;
+  eofSent_ = true;
 
-  } else {
-    // End of stream packet
-    cuvidPacket.flags = CUVID_PKT_ENDOFSTREAM;
-    eofSent_ = true;
-  }
+  return sendCuvidPacket(cuvidPacket);
+}
 
+int BetaCudaDeviceInterface::sendCuvidPacket(
+    CUVIDSOURCEDATAPACKET& cuvidPacket) {
   CUresult result = cuvidParseVideoData(videoParser_, &cuvidPacket);
-  if (result != CUDA_SUCCESS) {
-    return AVERROR_EXTERNAL;
-  }
-  return AVSUCCESS;
+  return result == CUDA_SUCCESS ? AVSUCCESS : AVERROR_EXTERNAL;
 }
 
 void BetaCudaDeviceInterface::applyBSF(ReferenceAVPacket& packet) {
@@ -541,17 +546,7 @@ UniqueAVFrame BetaCudaDeviceInterface::convertCudaFrameToAVFrame(
 void BetaCudaDeviceInterface::flush() {
   isFlushing_ = true;
 
-  // TODONVDEC P0: simplify flushing and "eofSent_" logic. We should just have a
-  // "sendEofPacket()" function that does the right thing, instead of setting
-  // CUVID_PKT_ENDOFSTREAM in different places.
-  if (!eofSent_) {
-    CUVIDSOURCEDATAPACKET cuvidPacket = {};
-    cuvidPacket.flags = CUVID_PKT_ENDOFSTREAM;
-    CUresult result = cuvidParseVideoData(videoParser_, &cuvidPacket);
-    if (result == CUDA_SUCCESS) {
-      eofSent_ = true;
-    }
-  }
+  sendEOFPacket();
 
   isFlushing_ = false;
 

src/torchcodec/_core/BetaCudaDeviceInterface.h

@@ -54,6 +54,7 @@ class BetaCudaDeviceInterface : public DeviceInterface {
   }
 
   int sendPacket(ReferenceAVPacket& packet) override;
+  int sendEOFPacket() override;
   int receiveFrame(UniqueAVFrame& avFrame) override;
   void flush() override;
 
@@ -63,6 +64,7 @@ class BetaCudaDeviceInterface : public DeviceInterface {
   int frameReadyInDisplayOrder(CUVIDPARSERDISPINFO* dispInfo);
 
  private:
+  int sendCuvidPacket(CUVIDSOURCEDATAPACKET& cuvidPacket);
   // Apply bitstream filter, modifies packet in-place
   void applyBSF(ReferenceAVPacket& packet);
   void initializeBSF(

src/torchcodec/_core/DeviceInterface.h

@@ -86,6 +86,14 @@ class DeviceInterface {
     return AVERROR(ENOSYS);
   }
 
+  // Send an EOF packet to flush the decoder
+  // Returns AVSUCCESS on success, or other AVERROR on failure
+  virtual int sendEOFPacket() {
+    TORCH_CHECK(
+        false, "Send EOF packet not implemented for this device interface");
+    return AVERROR(ENOSYS);
+  }
+
   // Moral equivalent of avcodec_receive_frame()
   // Returns AVSUCCESS on success, AVERROR(EAGAIN) if no frame ready,
   // AVERROR_EOF if end of stream, or other AVERROR on failure

src/torchcodec/_core/SingleStreamDecoder.cpp

@@ -1202,12 +1202,7 @@ UniqueAVFrame SingleStreamDecoder::decodeAVFrame(
       if (status == AVERROR_EOF) {
         // End of file reached. We must drain the decoder
         if (useCustomInterface) {
-          // TODONVDEC P0: Re-think this. This should be simpler.
-          AutoAVPacket eofAutoPacket;
-          ReferenceAVPacket eofPacket(eofAutoPacket);
-          eofPacket->data = nullptr;
-          eofPacket->size = 0;
-          status = deviceInterface_->sendPacket(eofPacket);
+          status = deviceInterface_->sendEOFPacket();
         } else {
           status = avcodec_send_packet(
               streamInfo.codecContext.get(),