Fix packetization logic to leave space for extensions in the last packet

webrtc/modules/rtp_rtcp/source/rtp_format_h264.cc

Index: webrtc/modules/rtp_rtcp/source/rtp_format_h264.cc
diff --git a/webrtc/modules/rtp_rtcp/source/rtp_format_h264.cc b/webrtc/modules/rtp_rtcp/source/rtp_format_h264.cc
index e190ea2f6aeea5f8e261cec86c6782477e867d32..11b13a67f93cdefc1fdfbae6bcf68411073a079c 100644
--- a/webrtc/modules/rtp_rtcp/source/rtp_format_h264.cc
+++ b/webrtc/modules/rtp_rtcp/source/rtp_format_h264.cc
@@ -79,8 +79,11 @@ bool ParseStapAStartOffsets(const uint8_t* nalu_ptr,
 }  // namespace
 
 RtpPacketizerH264::RtpPacketizerH264(size_t max_payload_len,
+                                     size_t last_packet_reduction_len,
                                      H264PacketizationMode packetization_mode)
     : max_payload_len_(max_payload_len),
+      last_packet_reduction_len_(last_packet_reduction_len),
+      total_packets_(0),
       packetization_mode_(packetization_mode) {
   // Guard against uninitialized memory in packetization_mode.
   RTC_CHECK(packetization_mode == H264PacketizationMode::NonInterleaved ||
@@ -95,7 +98,7 @@ RtpPacketizerH264::Fragment::Fragment(const uint8_t* buffer, size_t length)
 RtpPacketizerH264::Fragment::Fragment(const Fragment& fragment)
     : buffer(fragment.buffer), length(fragment.length) {}
 
-void RtpPacketizerH264::SetPayloadData(
+size_t RtpPacketizerH264::SetPayloadData(
     const uint8_t* payload_data,
     size_t payload_size,
     const RTPFragmentationHeader* fragmentation) {
@@ -164,6 +167,7 @@ void RtpPacketizerH264::SetPayloadData(
       input_fragments_.push_back(Fragment(buffer, length));
   }
   GeneratePackets();
+  return total_packets_;
 }
 
 void RtpPacketizerH264::GeneratePackets() {
@@ -174,7 +178,10 @@ void RtpPacketizerH264::GeneratePackets() {
         ++i;
         break;
       case H264PacketizationMode::NonInterleaved:
-        if (input_fragments_[i].length > max_payload_len_) {
+        if (input_fragments_[i].length > max_payload_len_ ||
+            (i + 1 == input_fragments_.size() &&
+             (input_fragments_[i].length + last_packet_reduction_len_ >
+              max_payload_len_))) {

[sprang_webrtc, 2017/05/17 13:10:07]

nit: Maybe you can clarify this a bit by having a temporary length variable, I'm
thinking something like

size_t fragment_length = input_fragments_[i].length;
if (i >= input_fragments_.size() - 1) {
  // Optionally comment this.
  fragment_length += last_packet_reduction_len_;
}
if (fragment_length > max_payload_len_) {
 fua
} else {
 stap
}

wdyt?

[ilnik, 2017/05/17 15:06:33]

Done.

           PacketizeFuA(i);
           ++i;
         } else {
@@ -189,26 +196,58 @@ void 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];
-
-  size_t fragment_length = fragment.length - kNalHeaderSize;
+  bool is_last_fragment = fragment_index + 1 == input_fragments_.size();
+  size_t payload_left = fragment.length - kNalHeaderSize;
   size_t offset = kNalHeaderSize;
   size_t bytes_available = max_payload_len_ - kFuAHeaderSize;
-  const size_t num_fragments =
-      (fragment_length + (bytes_available - 1)) / bytes_available;
-
-  const size_t avg_size = (fragment_length + num_fragments - 1) / num_fragments;
-  while (fragment_length > 0) {
-    size_t packet_length = avg_size;
-    if (fragment_length < avg_size)
-      packet_length = fragment_length;
+
+  // 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;
+
+  // Integer divisions with rounding up. Minimal number of packets to fit all
+  // payload and virtual payload.
+  size_t num_packets =
+      (payload_left + extra_len + (bytes_available - 1)) / bytes_available;
+  // Bytes per packet.
+  size_t payload_per_packet =
+      (payload_left + extra_len + num_packets - 1) / num_packets;

[sprang_webrtc, 2017/05/17 13:10:07]

nit: maybe (num_packets - 1) for consistency with bytes_available above

[ilnik, 2017/05/17 15:06:33]

Done.

+  // We make several first packets to be 1 bytes larger than the rest.
+  // i.e 14 bytes splitted in 4 packets would be 4+4+3+3.
+  size_t num_smaller_packets =
+      num_packets - (payload_left + extra_len) % num_packets;

[sprang_webrtc, 2017/05/17 13:10:06]

If you use num_larger_packets instead you can get rid of the if statement below,
and if you decrement num_larger_packets in the loop you can use that to
determine if packet_length should be payload_per_packet + 1.

[ilnik, 2017/05/17 15:06:33]

Done.

+  // If all the packets are of equal size, we assume they all are 1 byte larger.
+  if (num_smaller_packets == num_packets)
+    num_smaller_packets = 0;
+  while (payload_left > 0) {
+    // Reduce payload per packet at the right time.
+    if (num_packets == num_smaller_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 = packet_length - 1;
+      }
+    }
+    RTC_CHECK_GT(packet_length, 0);
     packets_.push(PacketUnit(Fragment(fragment.buffer + offset, packet_length),
                              offset - kNalHeaderSize == 0,
-                             fragment_length == packet_length, false,
+                             payload_left == packet_length, false,
                              fragment.buffer[0]));
     offset += packet_length;
-    fragment_length -= packet_length;
+    payload_left -= packet_length;
+    total_packets_++;

[sprang_webrtc, 2017/05/17 13:10:07]

nit: just do total_packet_ += num_packets before the loop

[ilnik, 2017/05/17 15:06:33]

Done.

+    num_packets--;

[sprang_webrtc, 2017/05/17 13:10:07]

nit: prefer pre-increments (here and otherwise)

[ilnik, 2017/05/17 15:06:33]

Done.

   }
-  RTC_CHECK_EQ(0, fragment_length);
+  RTC_CHECK_EQ(0, payload_left);
 }
 
 size_t RtpPacketizerH264::PacketizeStapA(size_t fragment_index) {
@@ -218,19 +257,17 @@ size_t RtpPacketizerH264::PacketizeStapA(size_t fragment_index) {
   size_t fragment_headers_length = 0;
   const Fragment* fragment = &input_fragments_[fragment_index];
   RTC_CHECK_GE(payload_size_left, fragment->length);
-  while (payload_size_left >= fragment->length + fragment_headers_length) {
+  total_packets_++;
+  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_)) {
     RTC_CHECK_GT(fragment->length, 0);
     packets_.push(PacketUnit(*fragment, aggregated_fragments == 0, false, true,
                              fragment->buffer[0]));
     payload_size_left -= fragment->length;
     payload_size_left -= fragment_headers_length;
 
-    // Next fragment.
-    ++fragment_index;
-    if (fragment_index == input_fragments_.size())
-      break;
-    fragment = &input_fragments_[fragment_index];
-
     fragment_headers_length = kLengthFieldSize;
     // If we are going to try to aggregate more fragments into this packet
     // we need to add the STAP-A NALU header and a length field for the first
@@ -238,7 +275,14 @@ size_t RtpPacketizerH264::PacketizeStapA(size_t fragment_index) {
     if (aggregated_fragments == 0)
       fragment_headers_length += kNalHeaderSize + kLengthFieldSize;
     ++aggregated_fragments;
+
+    // Next fragment.
+    ++fragment_index;
+    if (fragment_index == input_fragments_.size())
+      break;
+    fragment = &input_fragments_[fragment_index];
   }
+  RTC_CHECK_GT(aggregated_fragments, 0);
   packets_.back().last_fragment = true;
   return fragment_index;
 }
@@ -246,6 +290,8 @@ size_t RtpPacketizerH264::PacketizeStapA(size_t fragment_index) {
 void RtpPacketizerH264::PacketizeSingleNalu(size_t fragment_index) {
   // Add a single NALU to the queue, no aggregation.
   size_t payload_size_left = max_payload_len_;
+  if (fragment_index + 1 == input_fragments_.size())
+    payload_size_left -= last_packet_reduction_len_;
   const Fragment* fragment = &input_fragments_[fragment_index];
   RTC_CHECK_GE(payload_size_left, fragment->length)
       << "Payload size left " << payload_size_left << ", fragment length "
@@ -253,14 +299,12 @@ void RtpPacketizerH264::PacketizeSingleNalu(size_t fragment_index) {
   RTC_CHECK_GT(fragment->length, 0u);
   packets_.push(PacketUnit(*fragment, true /* first */, true /* last */,
                            false /* aggregated */, fragment->buffer[0]));
+  total_packets_++;
 }
 
-bool RtpPacketizerH264::NextPacket(RtpPacketToSend* rtp_packet,
-                                   bool* last_packet) {
+bool RtpPacketizerH264::NextPacket(RtpPacketToSend* rtp_packet) {
   RTC_DCHECK(rtp_packet);
-  RTC_DCHECK(last_packet);
   if (packets_.empty()) {
-    *last_packet = true;
     return false;
   }
 
@@ -274,19 +318,25 @@ bool RtpPacketizerH264::NextPacket(RtpPacketToSend* rtp_packet,
     input_fragments_.pop_front();
   } else if (packet.aggregated) {
     RTC_CHECK_EQ(H264PacketizationMode::NonInterleaved, packetization_mode_);
-    NextAggregatePacket(rtp_packet);
+    NextAggregatePacket(rtp_packet, total_packets_ == 1);
   } else {
     RTC_CHECK_EQ(H264PacketizationMode::NonInterleaved, packetization_mode_);
     NextFragmentPacket(rtp_packet);
   }
   RTC_DCHECK_LE(rtp_packet->payload_size(), max_payload_len_);
-  *last_packet = packets_.empty();
-  rtp_packet->SetMarker(*last_packet);
+  if (packets_.empty()) {
+    RTC_DCHECK_LE(rtp_packet->payload_size(),
+                  max_payload_len_ - last_packet_reduction_len_);
+  }
+  rtp_packet->SetMarker(packets_.empty());
+  total_packets_--;
   return true;
 }
 
-void RtpPacketizerH264::NextAggregatePacket(RtpPacketToSend* rtp_packet) {
-  uint8_t* buffer = rtp_packet->AllocatePayload(max_payload_len_);
+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_);
   RTC_DCHECK(buffer);
   PacketUnit* packet = &packets_.front();
   RTC_CHECK(packet->first_fragment);