Support first packet reduction in H264 packetizer


Bug: webrtc:9680
Change-Id: I73c9a5acecdf8dd82347be602bbfd7412c9610c5
Reviewed-on: https://webrtc-review.googlesource.com/99804
Commit-Queue: Danil Chapovalov <danilchap@webrtc.org>
Reviewed-by: Ilya Nikolaevskiy <ilnik@webrtc.org>
Cr-Commit-Position: refs/heads/master@{#24823}
diff --git a/modules/rtp_rtcp/source/rtp_format_h264.cc b/modules/rtp_rtcp/source/rtp_format_h264.cc
index 0422595..9793eba 100644
--- a/modules/rtp_rtcp/source/rtp_format_h264.cc
+++ b/modules/rtp_rtcp/source/rtp_format_h264.cc
@@ -84,14 +84,10 @@
     PayloadSizeLimits limits,
     H264PacketizationMode packetization_mode,
     const RTPFragmentationHeader& fragmentation)
-    : max_payload_len_(limits.max_payload_len),
-      // TODO(bugs.webrtc.org/9680): Do not ignore first_packet_reduction_len.
-      last_packet_reduction_len_(limits.last_packet_reduction_len),
-      num_packets_left_(0) {
+    : limits_(limits), num_packets_left_(0) {
   // Guard against uninitialized memory in packetization_mode.
   RTC_CHECK(packetization_mode == H264PacketizationMode::NonInterleaved ||
             packetization_mode == H264PacketizationMode::SingleNalUnit);
-  RTC_CHECK_GT(limits.max_payload_len, limits.last_packet_reduction_len);
 
   for (int i = 0; i < fragmentation.fragmentationVectorSize; ++i) {
     const uint8_t* buffer =
@@ -188,14 +184,16 @@
         ++i;
         break;
       case H264PacketizationMode::NonInterleaved:
-        size_t fragment_len = input_fragments_[i].length;
-        if (i + 1 == input_fragments_.size()) {
-          // Pretend that last fragment is larger instead of making last packet
-          // smaller.
-          fragment_len += last_packet_reduction_len_;
-        }
-        if (fragment_len > max_payload_len_) {
-          PacketizeFuA(i);
+        int fragment_len = input_fragments_[i].length;
+        int single_packet_capacity = limits_.max_payload_len;
+        if (i == 0)
+          single_packet_capacity -= limits_.first_packet_reduction_len;
+        if (i + 1 == input_fragments_.size())
+          single_packet_capacity -= limits_.last_packet_reduction_len;
+
+        if (fragment_len > single_packet_capacity) {
+          if (!PacketizeFuA(i))
+            return false;
           ++i;
         } else {
           i = PacketizeStapA(i);
@@ -206,63 +204,47 @@
   return true;
 }
 
-void RtpPacketizerH264::PacketizeFuA(size_t fragment_index) {
+bool RtpPacketizerH264::PacketizeFuA(size_t fragment_index) {
   // Fragment payload into packets (FU-A).
-  // Strip out the original header and leave room for the FU-A header.
   const Fragment& fragment = input_fragments_[fragment_index];
-  bool is_last_fragment = fragment_index + 1 == input_fragments_.size();
+
+  PayloadSizeLimits limits = limits_;
+  // Leave room for the FU-A header.
+  limits.max_payload_len -= kFuAHeaderSize;
+  // Ignore first/last packet reductions unless it is first/last fragment.
+  if (fragment_index != 0)
+    limits.first_packet_reduction_len = 0;
+  if (fragment_index != input_fragments_.size() - 1)
+    limits.last_packet_reduction_len = 0;
+
+  // Strip out the original header.
   size_t payload_left = fragment.length - kNalHeaderSize;
-  size_t offset = kNalHeaderSize;
-  size_t per_packet_capacity = max_payload_len_ - kFuAHeaderSize;
+  int offset = kNalHeaderSize;
 
-  // Instead of making the last packet smaller we pretend that all packets are
-  // of the same size but we write additional virtual payload to the last
-  // packet.
-  size_t extra_len = is_last_fragment ? last_packet_reduction_len_ : 0;
+  std::vector<int> payload_sizes = SplitAboutEqually(payload_left, limits);
+  if (payload_sizes.empty())
+    return false;
 
-  // Integer divisions with rounding up. Minimal number of packets to fit all
-  // payload and virtual payload.
-  size_t num_packets = (payload_left + extra_len + (per_packet_capacity - 1)) /
-                       per_packet_capacity;
-  // Bytes per packet. Average rounded down.
-  size_t payload_per_packet = (payload_left + extra_len) / num_packets;
-  // We make several first packets to be 1 bytes smaller than the rest.
-  // i.e 14 bytes splitted in 4 packets would be 3+3+4+4.
-  size_t num_larger_packets = (payload_left + extra_len) % num_packets;
-
-  num_packets_left_ += num_packets;
-  while (payload_left > 0) {
-    // Increase payload per packet at the right time.
-    if (num_packets == num_larger_packets)
-      ++payload_per_packet;
-    size_t packet_length = payload_per_packet;
-    if (payload_left <= packet_length) {  // Last portion of the payload
-      packet_length = payload_left;
-      // One additional packet may be used for extensions in the last packet.
-      // Together with last payload packet there may be at most 2 of them.
-      RTC_DCHECK_LE(num_packets, 2);
-      if (num_packets == 2) {
-        // Whole payload fits in the first num_packets-1 packets but extra
-        // packet is used for virtual payload. Leave at least one byte of data
-        // for the last packet.
-        --packet_length;
-      }
-    }
+  for (size_t i = 0; i < payload_sizes.size(); ++i) {
+    int packet_length = payload_sizes[i];
     RTC_CHECK_GT(packet_length, 0);
     packets_.push(PacketUnit(Fragment(fragment.buffer + offset, packet_length),
-                             offset - kNalHeaderSize == 0,
-                             payload_left == packet_length, false,
-                             fragment.buffer[0]));
+                             /*first_fragment=*/i == 0,
+                             /*last_fragment=*/i == payload_sizes.size() - 1,
+                             false, fragment.buffer[0]));
     offset += packet_length;
     payload_left -= packet_length;
-    --num_packets;
   }
+  num_packets_left_ += payload_sizes.size();
   RTC_CHECK_EQ(0, payload_left);
+  return true;
 }
 
 size_t RtpPacketizerH264::PacketizeStapA(size_t fragment_index) {
   // Aggregate fragments into one packet (STAP-A).
-  size_t payload_size_left = max_payload_len_;
+  size_t payload_size_left = limits_.max_payload_len;
+  if (fragment_index == 0)
+    payload_size_left -= limits_.first_packet_reduction_len;
   int aggregated_fragments = 0;
   size_t fragment_headers_length = 0;
   const Fragment* fragment = &input_fragments_[fragment_index];
@@ -271,7 +253,7 @@
   while (payload_size_left >= fragment->length + fragment_headers_length &&
          (fragment_index + 1 < input_fragments_.size() ||
           payload_size_left >= fragment->length + fragment_headers_length +
-                                   last_packet_reduction_len_)) {
+                                   limits_.last_packet_reduction_len)) {
     RTC_CHECK_GT(fragment->length, 0);
     packets_.push(PacketUnit(*fragment, aggregated_fragments == 0, false, true,
                              fragment->buffer[0]));
@@ -299,16 +281,18 @@
 
 bool RtpPacketizerH264::PacketizeSingleNalu(size_t fragment_index) {
   // Add a single NALU to the queue, no aggregation.
-  size_t payload_size_left = max_payload_len_;
+  size_t payload_size_left = limits_.max_payload_len;
+  if (fragment_index == 0)
+    payload_size_left -= limits_.first_packet_reduction_len;
   if (fragment_index + 1 == input_fragments_.size())
-    payload_size_left -= last_packet_reduction_len_;
+    payload_size_left -= limits_.last_packet_reduction_len;
   const Fragment* fragment = &input_fragments_[fragment_index];
   if (payload_size_left < fragment->length) {
     RTC_LOG(LS_ERROR) << "Failed to fit a fragment to packet in SingleNalu "
                          "packetization mode. Payload size left "
                       << payload_size_left << ", fragment length "
                       << fragment->length << ", packet capacity "
-                      << max_payload_len_;
+                      << limits_.max_payload_len;
     return false;
   }
   RTC_CHECK_GT(fragment->length, 0u);
@@ -333,25 +317,20 @@
     packets_.pop();
     input_fragments_.pop_front();
   } else if (packet.aggregated) {
-    bool is_last_packet = num_packets_left_ == 1;
-    NextAggregatePacket(rtp_packet, is_last_packet);
+    NextAggregatePacket(rtp_packet);
   } else {
     NextFragmentPacket(rtp_packet);
   }
-  RTC_DCHECK_LE(rtp_packet->payload_size(), max_payload_len_);
-  if (packets_.empty()) {
-    RTC_DCHECK_LE(rtp_packet->payload_size(),
-                  max_payload_len_ - last_packet_reduction_len_);
-  }
   rtp_packet->SetMarker(packets_.empty());
   --num_packets_left_;
   return true;
 }
 
-void RtpPacketizerH264::NextAggregatePacket(RtpPacketToSend* rtp_packet,
-                                            bool last) {
-  uint8_t* buffer = rtp_packet->AllocatePayload(
-      last ? max_payload_len_ - last_packet_reduction_len_ : max_payload_len_);
+void RtpPacketizerH264::NextAggregatePacket(RtpPacketToSend* rtp_packet) {
+  // Reserve maximum available payload, set actual payload size later.
+  size_t payload_capacity = rtp_packet->FreeCapacity();
+  RTC_CHECK_GE(payload_capacity, kNalHeaderSize);
+  uint8_t* buffer = rtp_packet->AllocatePayload(payload_capacity);
   RTC_DCHECK(buffer);
   PacketUnit* packet = &packets_.front();
   RTC_CHECK(packet->first_fragment);
@@ -361,6 +340,7 @@
   bool is_last_fragment = packet->last_fragment;
   while (packet->aggregated) {
     const Fragment& fragment = packet->source_fragment;
+    RTC_CHECK_LE(index + kLengthFieldSize + fragment.length, payload_capacity);
     // Add NAL unit length field.
     ByteWriter<uint16_t>::WriteBigEndian(&buffer[index], fragment.length);
     index += kLengthFieldSize;
diff --git a/modules/rtp_rtcp/source/rtp_format_h264.h b/modules/rtp_rtcp/source/rtp_format_h264.h
index 57d9926..73e4087 100644
--- a/modules/rtp_rtcp/source/rtp_format_h264.h
+++ b/modules/rtp_rtcp/source/rtp_format_h264.h
@@ -78,14 +78,14 @@
   };
 
   bool GeneratePackets(H264PacketizationMode packetization_mode);
-  void PacketizeFuA(size_t fragment_index);
+  bool PacketizeFuA(size_t fragment_index);
   size_t PacketizeStapA(size_t fragment_index);
   bool PacketizeSingleNalu(size_t fragment_index);
-  void NextAggregatePacket(RtpPacketToSend* rtp_packet, bool last);
+
+  void NextAggregatePacket(RtpPacketToSend* rtp_packet);
   void NextFragmentPacket(RtpPacketToSend* rtp_packet);
 
-  const size_t max_payload_len_;
-  const size_t last_packet_reduction_len_;
+  const PayloadSizeLimits limits_;
   size_t num_packets_left_;
   std::deque<Fragment> input_fragments_;
   std::queue<PacketUnit> packets_;
diff --git a/modules/rtp_rtcp/source/rtp_format_h264_unittest.cc b/modules/rtp_rtcp/source/rtp_format_h264_unittest.cc
index fabaa93..edf907a 100644
--- a/modules/rtp_rtcp/source/rtp_format_h264_unittest.cc
+++ b/modules/rtp_rtcp/source/rtp_format_h264_unittest.cc
@@ -144,6 +144,28 @@
 }
 
 TEST_P(RtpPacketizerH264ModeTest,
+       SingleNaluFirstPacketReductionAppliesOnlyToFirstFragment) {
+  RtpPacketizer::PayloadSizeLimits limits;
+  limits.max_payload_len = 200;
+  limits.first_packet_reduction_len = 5;
+  const size_t fragments[] = {195, 200, 200};
+
+  RTPFragmentationHeader fragmentation = CreateFragmentation(fragments);
+  rtc::Buffer frame = CreateFrame(fragmentation);
+
+  RtpPacketizerH264 packetizer(frame, limits, GetParam(), fragmentation);
+  std::vector<RtpPacketToSend> packets = FetchAllPackets(&packetizer);
+
+  ASSERT_THAT(packets, SizeIs(3));
+  const uint8_t* next_fragment = frame.data();
+  EXPECT_THAT(packets[0].payload(), ElementsAreArray(next_fragment, 195));
+  next_fragment += 195;
+  EXPECT_THAT(packets[1].payload(), ElementsAreArray(next_fragment, 200));
+  next_fragment += 200;
+  EXPECT_THAT(packets[2].payload(), ElementsAreArray(next_fragment, 200));
+}
+
+TEST_P(RtpPacketizerH264ModeTest,
        SingleNaluLastPacketReductionAppliesOnlyToLastFragment) {
   RtpPacketizer::PayloadSizeLimits limits;
   limits.max_payload_len = 200;
@@ -165,6 +187,21 @@
   EXPECT_THAT(packets[2].payload(), ElementsAreArray(next_fragment, 195));
 }
 
+TEST_P(RtpPacketizerH264ModeTest,
+       SingleNaluFirstAndLastPacketReductionSumsForSinglePacket) {
+  RtpPacketizer::PayloadSizeLimits limits;
+  limits.max_payload_len = 200;
+  limits.first_packet_reduction_len = 20;
+  limits.last_packet_reduction_len = 30;
+  rtc::Buffer frame = CreateFrame(150);
+
+  RtpPacketizerH264 packetizer(frame, limits, GetParam(),
+                               NoFragmentation(frame));
+  std::vector<RtpPacketToSend> packets = FetchAllPackets(&packetizer);
+
+  EXPECT_THAT(packets, SizeIs(1));
+}
+
 INSTANTIATE_TEST_CASE_P(
     PacketMode,
     RtpPacketizerH264ModeTest,
@@ -225,6 +262,31 @@
   EXPECT_EQ(packets[2].payload_size(), 0x123u);
 }
 
+TEST(RtpPacketizerH264Test, StapARespectsFirstPacketReduction) {
+  RtpPacketizer::PayloadSizeLimits limits;
+  limits.max_payload_len = 1000;
+  limits.first_packet_reduction_len = 100;
+  const size_t kFirstFragmentSize =
+      limits.max_payload_len - limits.first_packet_reduction_len;
+  size_t fragments[] = {kFirstFragmentSize, 2, 2};
+  RTPFragmentationHeader fragmentation = CreateFragmentation(fragments);
+  rtc::Buffer frame = CreateFrame(fragmentation);
+
+  RtpPacketizerH264 packetizer(
+      frame, limits, H264PacketizationMode::NonInterleaved, fragmentation);
+  std::vector<RtpPacketToSend> packets = FetchAllPackets(&packetizer);
+
+  ASSERT_THAT(packets, SizeIs(2));
+  // Expect 1st packet is single nalu.
+  EXPECT_THAT(packets[0].payload(),
+              ElementsAreArray(frame.data(), kFirstFragmentSize));
+  // Expect 2nd packet is aggregate of last two fragments.
+  const uint8_t* tail = frame.data() + kFirstFragmentSize;
+  EXPECT_THAT(packets[1].payload(), ElementsAre(kStapA,                  //
+                                                0, 2, tail[0], tail[1],  //
+                                                0, 2, tail[2], tail[3]));
+}
+
 TEST(RtpPacketizerH264Test, StapARespectsLastPacketReduction) {
   RtpPacketizer::PayloadSizeLimits limits;
   limits.max_payload_len = 1000;
@@ -356,6 +418,13 @@
   EXPECT_THAT(TestFua(1200, limits), ElementsAre(600, 600));
 }
 
+TEST(RtpPacketizerH264Test, FUAWithFirstPacketReduction) {
+  RtpPacketizer::PayloadSizeLimits limits;
+  limits.max_payload_len = 1200;
+  limits.first_packet_reduction_len = 4;
+  EXPECT_THAT(TestFua(1198, limits), ElementsAre(597, 601));
+}
+
 TEST(RtpPacketizerH264Test, FUAWithLastPacketReduction) {
   RtpPacketizer::PayloadSizeLimits limits;
   limits.max_payload_len = 1200;
@@ -363,6 +432,14 @@
   EXPECT_THAT(TestFua(1198, limits), ElementsAre(601, 597));
 }
 
+TEST(RtpPacketizerH264Test, FUAWithFirstAndLastPacketReduction) {
+  RtpPacketizer::PayloadSizeLimits limits;
+  limits.max_payload_len = 1199;
+  limits.first_packet_reduction_len = 100;
+  limits.last_packet_reduction_len = 100;
+  EXPECT_THAT(TestFua(1000, limits), ElementsAre(500, 500));
+}
+
 TEST(RtpPacketizerH264Test, FUAEvenSize) {
   RtpPacketizer::PayloadSizeLimits limits;
   limits.max_payload_len = 1200;