quinn-proto: deduplicate before defragment if necessary

Author: geieredgar

quinn-proto/src/connection/assembler.rs

@@ -1,5 +1,5 @@
 use std::{
-    cmp::{max, Ordering},
+    cmp::{max, min, Ordering},
     collections::{binary_heap::PeekMut, BinaryHeap},
     mem,
 };
@@ -19,6 +19,7 @@ pub(crate) struct Assembler {
     /// length of the contiguous prefix of the stream which has been consumed by the application,
     /// aka the stream offset.
     bytes_read: u64,
+    end: u64,
 }
 
 impl Assembler {
@@ -175,6 +176,7 @@ impl Assembler {
             allocation_size,
             bytes.len()
         );
+        self.end = max(self.end, offset + bytes.len() as u64);
         if let State::Unordered { ref mut recvd } = self.state {
             // Discard duplicate data
             for duplicate in recvd.replace(offset..offset + bytes.len() as u64) {
@@ -216,9 +218,13 @@ impl Assembler {
         // of memory allocated. This limits over-allocation in proportion to the
         // buffered data. The constants are chosen somewhat arbitrarily and try to
         // balance between defragmentation overhead and over-allocation.
-        let over_allocation = self.allocated - self.buffered;
-        let threshold = max(self.buffered * 3 / 2, 32 * 1024);
+        let buffered = min(self.buffered, (self.end - self.bytes_read) as usize);
+        let over_allocation = self.allocated - buffered;
+        let threshold = max(buffered * 3 / 2, 32 * 1024);
         if over_allocation > threshold {
+            if self.state.is_ordered() && self.buffered > buffered {
+                self.deduplicate();
+            }
             self.defragment()
         }
     }