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

webrtc/modules/rtp_rtcp/source/rtp_format_vp9.cc

Index: webrtc/modules/rtp_rtcp/source/rtp_format_vp9.cc
diff --git a/webrtc/modules/rtp_rtcp/source/rtp_format_vp9.cc b/webrtc/modules/rtp_rtcp/source/rtp_format_vp9.cc
index 4c8f7954f4ccb97358b176233044e1955c18937e..bd4cf9043111344d5d56769b8125dea69a6a11ac 100644
--- a/webrtc/modules/rtp_rtcp/source/rtp_format_vp9.cc
+++ b/webrtc/modules/rtp_rtcp/source/rtp_format_vp9.cc
@@ -30,10 +30,6 @@ namespace {
 // Length of VP9 payload descriptors' fixed part.
 const size_t kFixedPayloadDescriptorBytes = 1;
 
-// Packet fragmentation mode. If true, packets are split into (almost) equal
-// sizes. Otherwise, as many bytes as possible are fit into one packet.
-const bool kBalancedMode = true;
-
 const uint32_t kReservedBitValue0 = 0;
 
 uint8_t TemporalIdxField(const RTPVideoHeaderVP9& hdr, uint8_t def) {
@@ -462,29 +458,16 @@ bool ParseSsData(rtc::BitBuffer* parser, RTPVideoHeaderVP9* vp9) {
   }
   return true;
 }
-
-// Gets the size of next payload chunk to send. Returns 0 on error.
-size_t CalcNextSize(size_t max_length, size_t rem_bytes) {
-  if (max_length == 0 || rem_bytes == 0) {
-    return 0;
-  }
-  if (kBalancedMode) {
-    size_t num_frags = std::ceil(static_cast<double>(rem_bytes) / max_length);
-    return static_cast<size_t>(
-        static_cast<double>(rem_bytes) / num_frags + 0.5);
-  }
-  return max_length >= rem_bytes ? rem_bytes : max_length;
-}
 }  // namespace
 
-
 RtpPacketizerVp9::RtpPacketizerVp9(const RTPVideoHeaderVP9& hdr,
-                                   size_t max_payload_length)
+                                   size_t max_payload_length,
+                                   size_t last_packet_extension_len)
     : hdr_(hdr),
       max_payload_length_(max_payload_length),
       payload_(nullptr),
-      payload_size_(0) {
-}
+      payload_size_(0),
+      last_packet_extension_len_(last_packet_extension_len) {}
 
 RtpPacketizerVp9::~RtpPacketizerVp9() {
 }
@@ -520,45 +503,78 @@ void RtpPacketizerVp9::SetPayloadData(
   GeneratePackets();
 }
 
+// Splits payload in minimal number of roughly equal in size packets.
 void RtpPacketizerVp9::GeneratePackets() {
   if (max_payload_length_ < PayloadDescriptorLength(hdr_) + 1) {
-    LOG(LS_ERROR) << "Payload header and one payload byte won't fit.";
+    LOG(LS_ERROR) << "Payload header and one payload byte won't fit in packet.";
+    return;
+  }
+  if (max_payload_length_ < PayloadDescriptorLengthMinusSsData(hdr_) + 1 +
+                                last_packet_extension_len_) {
+    LOG(LS_ERROR) << "Payload header, extensions, and one payload byte won't "
+                     "fit in the last packet.";
     return;
   }
+  // Payload, extensions and SS hdr data in the first packet together.
+  size_t total_bytes =
+      payload_size_ + last_packet_extension_len_ + SsDataLength(hdr_);
+  // Accounting for VP9 headers of the same length in each packet.
+  size_t capacity =
+      max_payload_length_ - PayloadDescriptorLengthMinusSsData(hdr_);
+  // Integer divisions rounding up.
+  size_t num_packets = (total_bytes + capacity - 1) / capacity;
+  size_t per_packet_data = (total_bytes + num_packets - 1) / num_packets;
+  // Last smaller_packets packets are 1 byte smaller than the rest.
+  size_t smaller_packets = num_packets - total_bytes % num_packets;
+  // If all packets are the same size, we assume 0 packets are smaller.
+  if (smaller_packets == num_packets)
+    smaller_packets = 0;
   size_t bytes_processed = 0;
+  size_t packets_left = num_packets;
   while (bytes_processed < payload_size_) {
+    if (packets_left == smaller_packets)
+      per_packet_data--;
+    size_t packet_bytes = per_packet_data;
+    // First packet also has SS hdr data.
+    if (bytes_processed == 0)
+      packet_bytes -= SsDataLength(hdr_);
     size_t rem_bytes = payload_size_ - bytes_processed;
-    size_t rem_payload_len = max_payload_length_ -
-         (bytes_processed ? PayloadDescriptorLengthMinusSsData(hdr_)
-                          : PayloadDescriptorLength(hdr_));
-
-    size_t packet_bytes = CalcNextSize(rem_payload_len, rem_bytes);
-    if (packet_bytes == 0) {
-      LOG(LS_ERROR) << "Failed to generate VP9 packets.";
-      while (!packets_.empty())
-        packets_.pop();
-      return;
+    if (packet_bytes >= rem_bytes) {
+      packet_bytes = rem_bytes;
+      // All payload will fit into pre-last packet. Leave at least one byte for
+      // the last packet. Should only happen if extensions take at least half
+      // of available capacity.
+      if (packets_left == 2) {
+        RTC_DCHECK(last_packet_extension_len_ >= capacity / 2);
+        packet_bytes--;
+      }
     }
     QueuePacket(bytes_processed, packet_bytes, bytes_processed == 0,
                 rem_bytes == packet_bytes, &packets_);
+    packets_left--;
     bytes_processed += packet_bytes;
+    // Last packet should have enough space for extensions.
+    assert(packets_left > 0 ||
+           capacity - packet_bytes >= last_packet_extension_len_);
   }
   assert(bytes_processed == payload_size_);
 }
 
-bool RtpPacketizerVp9::NextPacket(RtpPacketToSend* packet, bool* last_packet) {
+size_t RtpPacketizerVp9::TotalPackets() {
+  return packets_.size();
+}
+
+bool RtpPacketizerVp9::NextPacket(RtpPacketToSend* packet) {
   RTC_DCHECK(packet);
-  RTC_DCHECK(last_packet);
   if (packets_.empty()) {
     return false;
   }
   PacketInfo packet_info = packets_.front();
   packets_.pop();
 
-  if (!WriteHeaderAndPayload(packet_info, packet)) {
+  if (!WriteHeaderAndPayload(packet_info, packet, packets_.empty())) {
     return false;
   }
-  *last_packet = packets_.empty();
   packet->SetMarker(packets_.empty() &&
                     (hdr_.spatial_idx == kNoSpatialIdx ||
                      hdr_.spatial_idx == hdr_.num_spatial_layers - 1));
@@ -602,8 +618,11 @@ bool RtpPacketizerVp9::NextPacket(RtpPacketToSend* packet, bool* last_packet) {
 //      +-+-+-+-+-+-+-+-+
 
 bool RtpPacketizerVp9::WriteHeaderAndPayload(const PacketInfo& packet_info,
-                                             RtpPacketToSend* packet) const {
-  uint8_t* buffer = packet->AllocatePayload(max_payload_length_);
+                                             RtpPacketToSend* packet,
+                                             bool last) const {
+  uint8_t* buffer = packet->AllocatePayload(
+      last ? max_payload_length_ - last_packet_extension_len_
+           : max_payload_length_);
   RTC_DCHECK(buffer);
   size_t header_length;
   if (!WriteHeader(packet_info, buffer, &header_length))