Reduce frame sorter gap-tracking overhead

frame_sorter.go

@@ -5,7 +5,6 @@ import (
 
 	"github.com/apernet/quic-go/internal/protocol"
 	"github.com/apernet/quic-go/internal/utils"
-	"github.com/apernet/quic-go/internal/utils/tree"
 )
 
 // byteInterval is an interval from one ByteCount to the other
@@ -20,16 +19,20 @@ type frameSorterEntry struct {
 }
 
 type frameSorter struct {
-	queue   map[protocol.ByteCount]frameSorterEntry
-	readPos protocol.ByteCount
-	gapTree *tree.Btree[utils.ByteInterval]
+	queue       map[protocol.ByteCount]frameSorterEntry
+	readPos     protocol.ByteCount
+	gapTree     gapSet
+	matchedGaps []utils.ByteInterval
 }
 
+// Keep the common reorder window reusable without pinning pathological high-water allocations.
+const maxRetainedMatchedGaps = 64
+
 var errDuplicateStreamData = errors.New("duplicate stream data")
 
 func newFrameSorter() *frameSorter {
 	s := frameSorter{
-		gapTree: tree.New[utils.ByteInterval](),
+		gapTree: gapSet{hint: -1},
 		queue:   make(map[protocol.ByteCount]frameSorterEntry),
 	}
 	s.gapTree.Insert(utils.ByteInterval{Start: 0, End: protocol.MaxByteCount})
@@ -56,7 +59,12 @@ func (s *frameSorter) push(data []byte, offset protocol.ByteCount, doneCb func()
 	end := offset + protocol.ByteCount(len(data))
 	covInterval := utils.ByteInterval{Start: start, End: end}
 
-	gaps := s.gapTree.Match(covInterval)
+	gaps := s.gapTree.MatchInto(covInterval, s.matchedGaps[:0])
+	if cap(gaps) <= maxRetainedMatchedGaps {
+		s.matchedGaps = gaps
+	} else {
+		s.matchedGaps = nil
+	}
 
 	if len(gaps) == 0 {
 		// No overlap with any existing gap
@@ -121,16 +129,20 @@ func (s *frameSorter) push(data []byte, offset protocol.ByteCount, doneCb func()
 			// The frame covers the whole startGap. Delete the gap.
 			s.gapTree.Delete(startGap)
 		} else {
-			s.gapTree.Delete(startGap)
+			oldStartGap := startGap
 			startGap.Start = end
-			// Re-insert the gap, but with the new start.
-			s.gapTree.Insert(startGap)
+			if !s.gapTree.UpdatePreservingOrder(oldStartGap, startGap) {
+				s.gapTree.Delete(oldStartGap)
+				s.gapTree.Insert(startGap)
+			}
 		}
 	} else if !hasReplacedAtLeastOne {
-		s.gapTree.Delete(startGap)
+		oldStartGap := startGap
 		startGap.End = start
-		// Re-insert the gap, but with the new end.
-		s.gapTree.Insert(startGap)
+		if !s.gapTree.UpdatePreservingOrder(oldStartGap, startGap) {
+			s.gapTree.Delete(oldStartGap)
+			s.gapTree.Insert(startGap)
+		}
 		adjustedStartGapEnd = true
 	}
 
@@ -161,10 +173,12 @@ func (s *frameSorter) push(data []byte, offset protocol.ByteCount, doneCb func()
 			// The frame split the existing gap into two.
 			s.gapTree.Insert(utils.ByteInterval{Start: end, End: startGapEnd})
 		} else if !startGapEqualsEndGap {
-			s.gapTree.Delete(endGap)
+			oldEndGap := endGap
 			endGap.Start = end
-			// Re-insert the gap, but with the new start.
-			s.gapTree.Insert(endGap)
+			if !s.gapTree.UpdatePreservingOrder(oldEndGap, endGap) {
+				s.gapTree.Delete(oldEndGap)
+				s.gapTree.Insert(endGap)
+			}
 		}
 	}
 

frame_sorter_benchmark_test.go

@@ -0,0 +1,69 @@
+package quic
+
+import (
+	"testing"
+
+	"github.com/apernet/quic-go/internal/protocol"
+)
+
+const (
+	benchmarkFrameCount = 256
+	benchmarkFrameSize  = 1200
+)
+
+func BenchmarkFrameSorterPushOrdered(b *testing.B) {
+	payload := make([]byte, benchmarkFrameSize)
+	b.ReportAllocs()
+	b.SetBytes(int64(benchmarkFrameCount * benchmarkFrameSize))
+
+	for b.Loop() {
+		sorter := newFrameSorter()
+		for index := range benchmarkFrameCount {
+			offset := protocol.ByteCount(index * benchmarkFrameSize)
+			if err := sorter.Push(payload, offset, nil); err != nil {
+				b.Fatal(err)
+			}
+			poppedOffset, data, _ := sorter.Pop()
+			if poppedOffset != offset || len(data) != benchmarkFrameSize {
+				b.Fatalf("popped (%d, %d bytes), expected (%d, %d bytes)", poppedOffset, len(data), offset, benchmarkFrameSize)
+			}
+		}
+	}
+}
+
+func BenchmarkFrameSorterPushReordered(b *testing.B) {
+	payload := make([]byte, benchmarkFrameSize)
+	b.ReportAllocs()
+	b.SetBytes(int64(benchmarkFrameCount * benchmarkFrameSize))
+
+	for b.Loop() {
+		sorter := newFrameSorter()
+		var callbacks int
+		callback := func() { callbacks++ }
+		for windowStart := 0; windowStart < benchmarkFrameCount; windowStart += 32 {
+			for index := windowStart + 16; index < windowStart+32; index++ {
+				offset := protocol.ByteCount(index * benchmarkFrameSize)
+				if err := sorter.Push(payload, offset, callback); err != nil {
+					b.Fatal(err)
+				}
+			}
+			for index := windowStart; index < windowStart+16; index++ {
+				offset := protocol.ByteCount(index * benchmarkFrameSize)
+				if err := sorter.Push(payload, offset, callback); err != nil {
+					b.Fatal(err)
+				}
+			}
+			for index := windowStart; index < windowStart+32; index++ {
+				expectedOffset := protocol.ByteCount(index * benchmarkFrameSize)
+				poppedOffset, data, done := sorter.Pop()
+				if poppedOffset != expectedOffset || len(data) != benchmarkFrameSize || done == nil {
+					b.Fatalf("popped (%d, %d bytes), expected (%d, %d bytes)", poppedOffset, len(data), expectedOffset, benchmarkFrameSize)
+				}
+				done()
+			}
+		}
+		if callbacks != benchmarkFrameCount {
+			b.Fatalf("callbacks = %d, expected %d", callbacks, benchmarkFrameCount)
+		}
+	}
+}

frame_sorter_gap_benchmark_test.go

@@ -0,0 +1,52 @@
+package quic
+
+import (
+	"testing"
+
+	"github.com/apernet/quic-go/internal/protocol"
+)
+
+func BenchmarkFrameSorterPushManyGaps(b *testing.B) {
+	separator := make([]byte, benchmarkFrameSize)
+	partial := make([]byte, benchmarkFrameSize/2)
+	b.ReportAllocs()
+	b.SetBytes(int64(benchmarkFrameCount * (len(separator) + len(partial))))
+
+	for b.Loop() {
+		sorter := newFrameSorter()
+		for index := range benchmarkFrameCount {
+			offset := protocol.ByteCount((2*index + 1) * benchmarkFrameSize)
+			if err := sorter.Push(separator, offset, nil); err != nil {
+				b.Fatal(err)
+			}
+		}
+		for index := range benchmarkFrameCount {
+			offset := protocol.ByteCount(2 * index * benchmarkFrameSize)
+			if err := sorter.Push(partial, offset, nil); err != nil {
+				b.Fatal(err)
+			}
+		}
+		if sorter.gapTree.Len() != benchmarkFrameCount+1 {
+			b.Fatalf("gap count = %d, want %d", sorter.gapTree.Len(), benchmarkFrameCount+1)
+		}
+	}
+}
+
+func BenchmarkFrameSorterPushMaxGaps(b *testing.B) {
+	payload := make([]byte, 6)
+	b.ReportAllocs()
+	b.SetBytes(int64(protocol.MaxStreamFrameSorterGaps * len(payload)))
+
+	for b.Loop() {
+		sorter := newFrameSorter()
+		for index := range protocol.MaxStreamFrameSorterGaps {
+			offset := protocol.ByteCount(index * 7)
+			if err := sorter.Push(payload, offset, nil); err != nil {
+				b.Fatal(err)
+			}
+		}
+		if sorter.gapTree.Len() != protocol.MaxStreamFrameSorterGaps {
+			b.Fatalf("gap count = %d, want %d", sorter.gapTree.Len(), protocol.MaxStreamFrameSorterGaps)
+		}
+	}
+}

frame_sorter_test.go

@@ -1506,3 +1506,13 @@ func TestFrameSorterPeek(t *testing.T) {
 	p = make([]byte, 10)
 	require.ErrorIs(t, s.Peek(0, p), errTooLittleData)
 }
+
+func TestFrameSorterDoesNotRetainLargeGapMatchBuffer(t *testing.T) {
+	s := newFrameSorter()
+	for offset := 0; offset < 256; offset += 2 {
+		require.NoError(t, s.Push([]byte{0}, protocol.ByteCount(offset), nil))
+	}
+
+	require.NoError(t, s.Push(make([]byte, 256), 0, nil))
+	require.LessOrEqual(t, cap(s.matchedGaps), maxRetainedMatchedGaps)
+}

gap_set.go

@@ -0,0 +1,142 @@
+package quic
+
+import "github.com/apernet/quic-go/internal/utils"
+
+const (
+	gapInlineCapacity = 8
+	gapBlockCapacity  = 128
+	gapBlockSplit     = gapBlockCapacity / 2
+)
+
+type gapBlock struct {
+	count int
+	gaps  [gapBlockCapacity]utils.ByteInterval
+}
+
+func (b *gapBlock) values() []utils.ByteInterval {
+	return b.gaps[:b.count]
+}
+
+func (b *gapBlock) first() utils.ByteInterval {
+	return b.gaps[0]
+}
+
+func (b *gapBlock) last() utils.ByteInterval {
+	return b.gaps[b.count-1]
+}
+
+// gapSet stores disjoint intervals ordered by Start, which also makes End monotonic.
+// Small sets stay inline. Larger sets use dense blocks to reduce allocation and lookup costs.
+type gapSet struct {
+	length int
+	hint   int
+	inline [gapInlineCapacity]utils.ByteInterval
+	blocks []*gapBlock
+}
+
+func (s *gapSet) Len() int {
+	return s.length
+}
+
+func (s *gapSet) MatchInto(cond utils.ByteInterval, dst []utils.ByteInterval) []utils.ByteInterval {
+	if s.blocks != nil {
+		return s.matchBlocks(cond, dst)
+	}
+	for index := 0; index < s.length; index++ {
+		gap := s.inline[index]
+		if gap.End < cond.Start {
+			continue
+		}
+		if gap.Start > cond.End {
+			break
+		}
+		dst = append(dst, gap)
+	}
+	return dst
+}
+
+func (s *gapSet) Insert(value utils.ByteInterval) {
+	if s.blocks != nil {
+		s.insertBlock(value)
+		return
+	}
+	index := gapLowerBound(s.inline[:s.length], value)
+	if index < s.length && value.Comp(s.inline[index]) == 0 {
+		s.inline[index] = value
+		return
+	}
+	if s.length == gapInlineCapacity {
+		s.promote()
+		s.insertBlock(value)
+		return
+	}
+	copy(s.inline[index+1:s.length+1], s.inline[index:s.length])
+	s.inline[index] = value
+	s.length++
+}
+
+func (s *gapSet) Delete(value utils.ByteInterval) {
+	if s.blocks != nil {
+		s.deleteBlock(value)
+		return
+	}
+	index := gapLowerBound(s.inline[:s.length], value)
+	if index == s.length || value.Comp(s.inline[index]) != 0 {
+		return
+	}
+	copy(s.inline[index:s.length-1], s.inline[index+1:s.length])
+	s.length--
+	s.inline[s.length] = utils.ByteInterval{}
+}
+
+func (s *gapSet) UpdatePreservingOrder(oldValue, newValue utils.ByteInterval) bool {
+	if s.blocks != nil {
+		return s.updateBlock(oldValue, newValue)
+	}
+	index := gapLowerBound(s.inline[:s.length], oldValue)
+	if index == s.length || oldValue.Comp(s.inline[index]) != 0 {
+		return false
+	}
+	direction := newValue.Comp(oldValue)
+	if direction < 0 && index > 0 && newValue.Comp(s.inline[index-1]) <= 0 {
+		return false
+	}
+	if direction > 0 && index+1 < s.length && newValue.Comp(s.inline[index+1]) >= 0 {
+		return false
+	}
+	s.inline[index] = newValue
+	return true
+}
+
+func (s *gapSet) promote() {
+	block := new(gapBlock)
+	block.count = s.length
+	copy(block.gaps[:], s.inline[:s.length])
+	s.inline = [gapInlineCapacity]utils.ByteInterval{}
+	s.blocks = []*gapBlock{block}
+	s.hint = 0
+}
+
+func (s *gapSet) demote() {
+	var inline [gapInlineCapacity]utils.ByteInterval
+	position := 0
+	for _, block := range s.blocks {
+		position += copy(inline[position:], block.values())
+	}
+	s.inline = inline
+	s.blocks = nil
+	s.hint = -1
+}
+
+func gapLowerBound(values []utils.ByteInterval, target utils.ByteInterval) int {
+	low, high := 0, len(values)
+	for low < high {
+		middle := int(uint(low+high) >> 1)
+		if values[middle].Comp(target) < 0 {
+			low = middle + 1
+		} else {
+			high = middle
+		}
+	}
+	return low
+}