diff --git a/frame_sorter.go b/frame_sorter.go index 417cdaffae6..125ddbf3699 100644 --- a/frame_sorter.go +++ b/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) + } } } diff --git a/frame_sorter_benchmark_test.go b/frame_sorter_benchmark_test.go new file mode 100644 index 00000000000..9c7ccb581e6 --- /dev/null +++ b/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) + } + } +} diff --git a/frame_sorter_gap_benchmark_test.go b/frame_sorter_gap_benchmark_test.go new file mode 100644 index 00000000000..131a8d6d8e9 --- /dev/null +++ b/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) + } + } +} diff --git a/frame_sorter_test.go b/frame_sorter_test.go index 8043b566796..0d5071da642 100644 --- a/frame_sorter_test.go +++ b/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) +} diff --git a/gap_set.go b/gap_set.go new file mode 100644 index 00000000000..b2b704d72d6 --- /dev/null +++ b/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 +} diff --git a/gap_set_benchmark_test.go b/gap_set_benchmark_test.go new file mode 100644 index 00000000000..bdb0ee38399 --- /dev/null +++ b/gap_set_benchmark_test.go @@ -0,0 +1,91 @@ +package quic + +import ( + "testing" + + "github.com/apernet/quic-go/internal/protocol" + "github.com/apernet/quic-go/internal/utils" +) + +func BenchmarkGapSetMatchMaxGaps(b *testing.B) { + set := newGapSetWithSparseGaps(b, protocol.MaxStreamFrameSorterGaps) + matches := make([]utils.ByteInterval, 0, protocol.MaxStreamFrameSorterGaps) + query := utils.ByteInterval{Start: 0, End: protocol.MaxByteCount} + b.ReportAllocs() + b.ResetTimer() + + for b.Loop() { + matches = set.MatchInto(query, matches[:0]) + if len(matches) != protocol.MaxStreamFrameSorterGaps { + b.Fatalf("matched %d gaps, want %d", len(matches), protocol.MaxStreamFrameSorterGaps) + } + } +} + +func BenchmarkGapSetUpdateMaxGaps(b *testing.B) { + set := newGapSetWithSparseGaps(b, protocol.MaxStreamFrameSorterGaps) + middle := protocol.MaxStreamFrameSorterGaps / 2 + start := protocol.ByteCount(middle * 10) + oldGap := utils.ByteInterval{Start: start, End: start + 5} + newGap := utils.ByteInterval{Start: start + 1, End: start + 5} + b.ReportAllocs() + b.ResetTimer() + + for b.Loop() { + if !set.UpdatePreservingOrder(oldGap, newGap) { + b.Fatal("failed to move gap boundary forward") + } + if !set.UpdatePreservingOrder(newGap, oldGap) { + b.Fatal("failed to restore gap boundary") + } + } +} + +func BenchmarkGapSetInsertMaxGapsShuffled(b *testing.B) { + b.ReportAllocs() + b.SetBytes(int64(protocol.MaxStreamFrameSorterGaps * 16)) + + for b.Loop() { + set := &gapSet{hint: -1} + for step := range protocol.MaxStreamFrameSorterGaps { + index := (step * 7919) % protocol.MaxStreamFrameSorterGaps + start := protocol.ByteCount(index * 10) + set.Insert(utils.ByteInterval{Start: start, End: start + 5}) + } + if set.Len() != protocol.MaxStreamFrameSorterGaps { + b.Fatalf("gap count = %d, want %d", set.Len(), protocol.MaxStreamFrameSorterGaps) + } + } +} + +func BenchmarkGapSetInsertDeleteMaxGapsShuffled(b *testing.B) { + b.ReportAllocs() + b.SetBytes(int64(protocol.MaxStreamFrameSorterGaps * 32)) + + for b.Loop() { + set := &gapSet{hint: -1} + for step := range protocol.MaxStreamFrameSorterGaps { + index := (step * 7919) % protocol.MaxStreamFrameSorterGaps + start := protocol.ByteCount(index * 10) + set.Insert(utils.ByteInterval{Start: start, End: start + 5}) + } + for step := range protocol.MaxStreamFrameSorterGaps { + index := (step * 6151) % protocol.MaxStreamFrameSorterGaps + start := protocol.ByteCount(index * 10) + set.Delete(utils.ByteInterval{Start: start, End: start + 5}) + } + if set.Len() != 0 { + b.Fatalf("gap count = %d, want 0", set.Len()) + } + } +} + +func newGapSetWithSparseGaps(tb testing.TB, count int) *gapSet { + tb.Helper() + set := &gapSet{hint: -1} + for index := range count { + start := protocol.ByteCount(index * 10) + set.Insert(utils.ByteInterval{Start: start, End: start + 5}) + } + return set +} diff --git a/gap_set_blocks.go b/gap_set_blocks.go new file mode 100644 index 00000000000..6829407a13b --- /dev/null +++ b/gap_set_blocks.go @@ -0,0 +1,233 @@ +package quic + +import ( + "github.com/apernet/quic-go/internal/protocol" + "github.com/apernet/quic-go/internal/utils" +) + +func (s *gapSet) matchBlocks(cond utils.ByteInterval, dst []utils.ByteInterval) []utils.ByteInterval { + blockIndex := s.firstBlockEndingAtOrAfter(cond.Start) + if blockIndex == len(s.blocks) { + return dst + } + s.hint = blockIndex + for ; blockIndex < len(s.blocks); blockIndex++ { + block := s.blocks[blockIndex] + index := 0 + if blockIndex == s.hint { + low, high := 0, block.count + for low < high { + middle := int(uint(low+high) >> 1) + if block.gaps[middle].End < cond.Start { + low = middle + 1 + } else { + high = middle + } + } + index = low + } + for ; index < block.count; index++ { + gap := block.gaps[index] + if gap.Start > cond.End { + return dst + } + dst = append(dst, gap) + } + } + return dst +} + +func (s *gapSet) firstBlockEndingAtOrAfter(start protocol.ByteCount) int { + if s.hint >= 0 && s.hint < len(s.blocks) { + block := s.blocks[s.hint] + previousEndsBefore := s.hint == 0 || s.blocks[s.hint-1].last().End < start + if previousEndsBefore && block.last().End >= start { + return s.hint + } + } + low, high := 0, len(s.blocks) + for low < high { + middle := int(uint(low+high) >> 1) + if s.blocks[middle].last().End < start { + low = middle + 1 + } else { + high = middle + } + } + return low +} + +func (s *gapSet) insertBlock(value utils.ByteInterval) { + blockIndex := s.blockForValue(value) + block := s.blocks[blockIndex] + index := gapLowerBound(block.values(), value) + if index < block.count && value.Comp(block.gaps[index]) == 0 { + block.gaps[index] = value + s.hint = blockIndex + return + } + if block.count == gapBlockCapacity { + blockIndex, block, index = s.splitBlock(blockIndex, index) + } + copy(block.gaps[index+1:block.count+1], block.gaps[index:block.count]) + block.gaps[index] = value + block.count++ + s.length++ + s.hint = blockIndex +} + +func (s *gapSet) splitBlock(blockIndex, insertionIndex int) (int, *gapBlock, int) { + block := s.blocks[blockIndex] + next := new(gapBlock) + next.count = gapBlockCapacity - gapBlockSplit + copy(next.gaps[:], block.gaps[gapBlockSplit:gapBlockCapacity]) + for index := gapBlockSplit; index < gapBlockCapacity; index++ { + block.gaps[index] = utils.ByteInterval{} + } + block.count = gapBlockSplit + + s.blocks = append(s.blocks, nil) + copy(s.blocks[blockIndex+2:], s.blocks[blockIndex+1:]) + s.blocks[blockIndex+1] = next + if insertionIndex >= gapBlockSplit { + return blockIndex + 1, next, insertionIndex - gapBlockSplit + } + return blockIndex, block, insertionIndex +} + +func (s *gapSet) deleteBlock(value utils.ByteInterval) { + blockIndex, index, found := s.locate(value) + if !found { + return + } + block := s.blocks[blockIndex] + copy(block.gaps[index:block.count-1], block.gaps[index+1:block.count]) + block.count-- + block.gaps[block.count] = utils.ByteInterval{} + s.length-- + + if s.length <= gapInlineCapacity { + s.demote() + return + } + if block.count == 0 { + s.removeBlock(blockIndex) + return + } + s.mergeSparseBlock(blockIndex) +} + +func (s *gapSet) updateBlock(oldValue, newValue utils.ByteInterval) bool { + blockIndex, index, found := s.locate(oldValue) + if !found { + return false + } + direction := newValue.Comp(oldValue) + if direction < 0 { + if predecessor, ok := s.predecessor(blockIndex, index); ok && newValue.Comp(predecessor) <= 0 { + return false + } + } else if direction > 0 { + if successor, ok := s.successor(blockIndex, index); ok && newValue.Comp(successor) >= 0 { + return false + } + } + s.blocks[blockIndex].gaps[index] = newValue + s.hint = blockIndex + return true +} + +func (s *gapSet) locate(value utils.ByteInterval) (int, int, bool) { + blockIndex := s.blockForValue(value) + block := s.blocks[blockIndex] + index := gapLowerBound(block.values(), value) + found := index < block.count && value.Comp(block.gaps[index]) == 0 + return blockIndex, index, found +} + +func (s *gapSet) blockForValue(value utils.ByteInterval) int { + if s.hint >= 0 && s.hint < len(s.blocks) { + previousBefore := s.hint == 0 || s.blocks[s.hint-1].last().Comp(value) < 0 + if previousBefore && s.blocks[s.hint].last().Comp(value) >= 0 { + return s.hint + } + } + low, high := 0, len(s.blocks) + for low < high { + middle := int(uint(low+high) >> 1) + if s.blocks[middle].last().Comp(value) < 0 { + low = middle + 1 + } else { + high = middle + } + } + if low == len(s.blocks) { + return low - 1 + } + return low +} + +func (s *gapSet) predecessor(blockIndex, index int) (utils.ByteInterval, bool) { + if index > 0 { + return s.blocks[blockIndex].gaps[index-1], true + } + if blockIndex > 0 { + return s.blocks[blockIndex-1].last(), true + } + return utils.ByteInterval{}, false +} + +func (s *gapSet) successor(blockIndex, index int) (utils.ByteInterval, bool) { + block := s.blocks[blockIndex] + if index+1 < block.count { + return block.gaps[index+1], true + } + if blockIndex+1 < len(s.blocks) { + return s.blocks[blockIndex+1].first(), true + } + return utils.ByteInterval{}, false +} + +func (s *gapSet) mergeSparseBlock(blockIndex int) { + block := s.blocks[blockIndex] + if block.count >= gapBlockSplit { + s.hint = blockIndex + return + } + if blockIndex+1 < len(s.blocks) { + next := s.blocks[blockIndex+1] + if block.count+next.count <= gapBlockCapacity { + copy(block.gaps[block.count:], next.values()) + block.count += next.count + s.removeBlock(blockIndex + 1) + s.hint = blockIndex + return + } + } + if blockIndex > 0 { + previous := s.blocks[blockIndex-1] + if previous.count+block.count <= gapBlockCapacity { + copy(previous.gaps[previous.count:], block.values()) + previous.count += block.count + s.removeBlock(blockIndex) + s.hint = blockIndex - 1 + return + } + } + s.hint = blockIndex +} + +func (s *gapSet) removeBlock(index int) { + copy(s.blocks[index:], s.blocks[index+1:]) + last := len(s.blocks) - 1 + s.blocks[last] = nil + s.blocks = s.blocks[:last] + if len(s.blocks) == 0 { + s.hint = -1 + return + } + if index == len(s.blocks) { + index-- + } + s.hint = index +} diff --git a/gap_set_test.go b/gap_set_test.go new file mode 100644 index 00000000000..e11b3f37c5a --- /dev/null +++ b/gap_set_test.go @@ -0,0 +1,159 @@ +package quic + +import ( + "math/rand/v2" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/apernet/quic-go/internal/protocol" + "github.com/apernet/quic-go/internal/utils" + "github.com/apernet/quic-go/internal/utils/tree" +) + +func newTestGapSet() *gapSet { + return &gapSet{hint: -1} +} + +func (s *gapSet) Head() *utils.ByteInterval { + if s.length == 0 { + return nil + } + if s.blocks == nil { + return &s.inline[0] + } + return &s.blocks[0].gaps[0] +} + +func (s *gapSet) Values() []utils.ByteInterval { + values := make([]utils.ByteInterval, 0, s.length) + if s.blocks == nil { + return append(values, s.inline[:s.length]...) + } + for _, block := range s.blocks { + values = append(values, block.values()...) + } + return values +} + +func TestGapSetMatchesTouchingIntervalsInOrder(t *testing.T) { + set := newTestGapSet() + set.Insert(utils.ByteInterval{Start: 40, End: 50}) + set.Insert(utils.ByteInterval{Start: 0, End: 10}) + set.Insert(utils.ByteInterval{Start: 20, End: 30}) + + prefix := []utils.ByteInterval{{Start: 100, End: 110}} + got := set.MatchInto(utils.ByteInterval{Start: 10, End: 40}, prefix) + require.Equal(t, []utils.ByteInterval{ + {Start: 100, End: 110}, + {Start: 0, End: 10}, + {Start: 20, End: 30}, + {Start: 40, End: 50}, + }, got) +} + +func TestGapSetUpdatePreservingOrderRejectsCrossingNeighbor(t *testing.T) { + set := newTestGapSet() + for _, interval := range []utils.ByteInterval{ + {Start: 0, End: 10}, + {Start: 20, End: 30}, + {Start: 40, End: 50}, + } { + set.Insert(interval) + } + + require.True(t, set.UpdatePreservingOrder( + utils.ByteInterval{Start: 20, End: 30}, + utils.ByteInterval{Start: 21, End: 31}, + )) + require.False(t, set.UpdatePreservingOrder( + utils.ByteInterval{Start: 21, End: 31}, + utils.ByteInterval{Start: 41, End: 45}, + )) + require.False(t, set.UpdatePreservingOrder( + utils.ByteInterval{Start: 60, End: 70}, + utils.ByteInterval{Start: 61, End: 71}, + )) + require.Equal(t, []utils.ByteInterval{ + {Start: 0, End: 10}, + {Start: 21, End: 31}, + {Start: 40, End: 50}, + }, set.Values()) +} + +func TestGapSetPromotesAndKeepsSortedValues(t *testing.T) { + set := newTestGapSet() + count := gapBlockCapacity*3 + 5 + for index := count - 1; index >= 0; index-- { + start := protocol.ByteCount(index * 4) + set.Insert(utils.ByteInterval{Start: start, End: start + 2}) + } + + require.Equal(t, count, set.Len()) + for index, interval := range set.Values() { + require.Equal(t, protocol.ByteCount(index*4), interval.Start) + } + + for index := 0; index < count; index += 3 { + start := protocol.ByteCount(index * 4) + set.Delete(utils.ByteInterval{Start: start, End: start + 2}) + } + require.Equal(t, count-(count+2)/3, set.Len()) + require.Equal(t, protocol.ByteCount(4), set.Head().Start) +} + +func TestGapSetUpdatesAcrossLeafBoundaryAndShrinksToInlineSize(t *testing.T) { + set := newTestGapSet() + for index := range gapBlockCapacity*2 + 3 { + start := protocol.ByteCount(index * 10) + set.Insert(utils.ByteInterval{Start: start, End: start + 5}) + } + + oldBoundary := set.blocks[0].last() + newBoundary := utils.ByteInterval{Start: oldBoundary.Start + 1, End: oldBoundary.End} + require.True(t, set.UpdatePreservingOrder(oldBoundary, newBoundary)) + require.Contains(t, set.Values(), newBoundary) + + for set.Len() > gapInlineCapacity { + set.Delete(set.Values()[0]) + } + values := set.Values() + require.Len(t, values, gapInlineCapacity) + require.Equal(t, values, set.MatchInto(utils.ByteInterval{Start: 0, End: protocol.MaxByteCount}, nil)) +} + +func TestGapSetMatchesAVLReferenceAfterMutations(t *testing.T) { + set := newTestGapSet() + reference := tree.New[utils.ByteInterval]() + const count = 257 + order := rand.New(rand.NewPCG(7, 11)).Perm(count) + for _, index := range order { + start := protocol.ByteCount(index * 10) + interval := utils.ByteInterval{Start: start, End: start + 5} + set.Insert(interval) + reference.Insert(interval) + } + + for index := 0; index < count; index += 5 { + start := protocol.ByteCount(index * 10) + oldInterval := utils.ByteInterval{Start: start, End: start + 5} + newInterval := utils.ByteInterval{Start: start + 1, End: start + 5} + require.True(t, set.UpdatePreservingOrder(oldInterval, newInterval)) + reference.Delete(oldInterval) + reference.Insert(newInterval) + } + for index := 2; index < count; index += 7 { + start := protocol.ByteCount(index * 10) + interval := utils.ByteInterval{Start: start, End: start + 5} + set.Delete(interval) + reference.Delete(interval) + } + + require.Equal(t, reference.Values(), set.Values()) + rng := rand.New(rand.NewPCG(13, 17)) + for range 1_000 { + start := protocol.ByteCount(rng.IntN(count * 10)) + query := utils.ByteInterval{Start: start, End: start + protocol.ByteCount(rng.IntN(40))} + require.Equal(t, reference.Match(query), set.MatchInto(query, nil), "query %v", query) + } +} diff --git a/receive_stream_benchmark_test.go b/receive_stream_benchmark_test.go new file mode 100644 index 00000000000..1147999bdb5 --- /dev/null +++ b/receive_stream_benchmark_test.go @@ -0,0 +1,235 @@ +package quic + +import ( + "errors" + "fmt" + "io" + "testing" + + "github.com/apernet/quic-go/internal/monotime" + "github.com/apernet/quic-go/internal/protocol" + "github.com/apernet/quic-go/internal/wire" +) + +const ( + receiveStreamBenchmarkFrameCount = 256 + receiveStreamBenchmarkFrameSize = 1200 + receiveStreamBenchmarkReadSize = 16 * 1024 + receiveStreamBenchmarkWindow = 32 +) + +type receiveStreamBenchmarkScenario struct { + frames []*wire.StreamFrame + order []int + readBuffer []byte +} + +type receiveStreamBenchmarkResult struct { + bytesRead protocol.ByteCount + flowBytesRead protocol.ByteCount + highestReceived protocol.ByteCount + finalOffset protocol.ByteCount + callbacks int + completions int + err error +} + +type receiveStreamBenchmarkFlowController struct { + bytesRead protocol.ByteCount + highestReceived protocol.ByteCount + finalOffset protocol.ByteCount +} + +func (f *receiveStreamBenchmarkFlowController) SendWindowSize() protocol.ByteCount { + return protocol.MaxByteCount +} + +func (*receiveStreamBenchmarkFlowController) UpdateSendWindow(protocol.ByteCount) bool { return false } +func (*receiveStreamBenchmarkFlowController) AddBytesSent(protocol.ByteCount) {} +func (*receiveStreamBenchmarkFlowController) GetWindowUpdate(monotime.Time) protocol.ByteCount { + return 0 +} + +func (f *receiveStreamBenchmarkFlowController) AddBytesRead(n protocol.ByteCount) (bool, bool) { + f.bytesRead += n + return false, false +} + +func (f *receiveStreamBenchmarkFlowController) UpdateHighestReceived( + offset protocol.ByteCount, + final bool, + _ monotime.Time, +) error { + if offset > f.highestReceived { + f.highestReceived = offset + } + if final { + f.finalOffset = offset + } + return nil +} + +func (*receiveStreamBenchmarkFlowController) Abandon() {} +func (*receiveStreamBenchmarkFlowController) IsNewlyBlocked() bool { return false } + +type receiveStreamBenchmarkSender struct { + completions int +} + +func (*receiveStreamBenchmarkSender) onHasConnectionData() {} +func (*receiveStreamBenchmarkSender) onHasStreamData(protocol.StreamID, *SendStream) { +} +func (*receiveStreamBenchmarkSender) onHasStreamControlFrame(protocol.StreamID, streamControlFrameGetter) { +} +func (s *receiveStreamBenchmarkSender) onStreamCompleted(protocol.StreamID) { + s.completions++ +} + +func BenchmarkReceiveStreamHandleAndReadSequential(b *testing.B) { + benchmarkReceiveStreamHandleAndRead(b, newReceiveStreamBenchmarkScenario(false)) +} + +func BenchmarkReceiveStreamHandleAndReadReorderedWindow32(b *testing.B) { + benchmarkReceiveStreamHandleAndRead(b, newReceiveStreamBenchmarkScenario(true)) +} + +func newReceiveStreamBenchmarkScenario(reordered bool) *receiveStreamBenchmarkScenario { + frames := make([]*wire.StreamFrame, receiveStreamBenchmarkFrameCount) + for frameIndex := range frames { + offset := frameIndex * receiveStreamBenchmarkFrameSize + data := make([]byte, receiveStreamBenchmarkFrameSize) + for dataIndex := range data { + data[dataIndex] = byte((offset + dataIndex) % 251) + } + frames[frameIndex] = &wire.StreamFrame{ + Offset: protocol.ByteCount(offset), + Data: data, + Fin: frameIndex == receiveStreamBenchmarkFrameCount-1, + } + } + + order := make([]int, 0, receiveStreamBenchmarkFrameCount) + for windowStart := 0; windowStart < receiveStreamBenchmarkFrameCount; windowStart += receiveStreamBenchmarkWindow { + if reordered { + for frameIndex := windowStart + receiveStreamBenchmarkWindow/2; frameIndex < windowStart+receiveStreamBenchmarkWindow; frameIndex++ { + order = append(order, frameIndex) + } + } + for frameIndex := windowStart; frameIndex < windowStart+receiveStreamBenchmarkWindow/2; frameIndex++ { + order = append(order, frameIndex) + } + if !reordered { + for frameIndex := windowStart + receiveStreamBenchmarkWindow/2; frameIndex < windowStart+receiveStreamBenchmarkWindow; frameIndex++ { + order = append(order, frameIndex) + } + } + } + + return &receiveStreamBenchmarkScenario{ + frames: frames, + order: order, + readBuffer: make([]byte, receiveStreamBenchmarkReadSize), + } +} + +func benchmarkReceiveStreamHandleAndRead(b *testing.B, scenario *receiveStreamBenchmarkScenario) { + b.Helper() + wantBytes := protocol.ByteCount(receiveStreamBenchmarkFrameCount * receiveStreamBenchmarkFrameSize) + + validation := scenario.run(true) + if validation.err != nil { + b.Fatalf("validating receive stream benchmark: %v", validation.err) + } + if validation.bytesRead != wantBytes || validation.flowBytesRead != wantBytes { + b.Fatalf("validation read %d bytes, flow controller recorded %d; want %d", validation.bytesRead, validation.flowBytesRead, wantBytes) + } + if validation.highestReceived != wantBytes || validation.finalOffset != wantBytes { + b.Fatalf("validation highest offset %d, final offset %d; want %d", validation.highestReceived, validation.finalOffset, wantBytes) + } + if validation.callbacks != receiveStreamBenchmarkFrameCount { + b.Fatalf("validation released %d frame callbacks; want %d", validation.callbacks, receiveStreamBenchmarkFrameCount) + } + if validation.completions != 1 { + b.Fatalf("validation completed stream %d times; want 1", validation.completions) + } + + b.ReportAllocs() + b.SetBytes(int64(wantBytes)) + for b.Loop() { + result := scenario.run(false) + if result.err != nil { + b.Fatalf("running receive stream benchmark: %v", result.err) + } + if result.bytesRead != wantBytes || result.flowBytesRead != wantBytes || result.completions != 1 { + b.Fatalf("benchmark lifecycle mismatch: read=%d flow=%d completions=%d", result.bytesRead, result.flowBytesRead, result.completions) + } + } +} + +func (s *receiveStreamBenchmarkScenario) run(validate bool) receiveStreamBenchmarkResult { + flowController := &receiveStreamBenchmarkFlowController{} + sender := &receiveStreamBenchmarkSender{} + stream := newReceiveStream(42, sender, flowController) + now := monotime.Now() + var bytesRead protocol.ByteCount + var callbacks int + + for windowStart := 0; windowStart < len(s.order); windowStart += receiveStreamBenchmarkWindow { + for _, frameIndex := range s.order[windowStart : windowStart+receiveStreamBenchmarkWindow] { + frame := s.frames[frameIndex] + if err := stream.handleStreamFrame(frame, now); err != nil { + return receiveStreamBenchmarkResult{err: fmt.Errorf("handling frame %d: %w", frameIndex, err)} + } + if validate { + entry, ok := stream.frameQueue.queue[frame.Offset] + if !ok { + return receiveStreamBenchmarkResult{err: fmt.Errorf("frame %d missing from sorter queue", frameIndex)} + } + originalCallback := entry.DoneCb + entry.DoneCb = func() { + callbacks++ + if originalCallback != nil { + originalCallback() + } + } + stream.frameQueue.queue[frame.Offset] = entry + } + } + + windowBytesRemaining := receiveStreamBenchmarkWindow * receiveStreamBenchmarkFrameSize + for windowBytesRemaining > 0 { + readSize := min(windowBytesRemaining, len(s.readBuffer)) + n, err := stream.Read(s.readBuffer[:readSize]) + isFinalRead := bytesRead+protocol.ByteCount(n) == protocol.ByteCount(receiveStreamBenchmarkFrameCount*receiveStreamBenchmarkFrameSize) + if n != readSize { + return receiveStreamBenchmarkResult{err: fmt.Errorf("read %d bytes at offset %d; want %d", n, bytesRead, readSize)} + } + if isFinalRead { + if !errors.Is(err, io.EOF) { + return receiveStreamBenchmarkResult{err: fmt.Errorf("final read error: %w", err)} + } + } else if err != nil { + return receiveStreamBenchmarkResult{err: fmt.Errorf("reading at offset %d: %w", bytesRead, err)} + } + if validate { + for dataIndex, value := range s.readBuffer[:n] { + want := byte((int(bytesRead) + dataIndex) % 251) + if value != want { + return receiveStreamBenchmarkResult{err: fmt.Errorf("byte at offset %d is %d; want %d", int(bytesRead)+dataIndex, value, want)} + } + } + } + bytesRead += protocol.ByteCount(n) + windowBytesRemaining -= n + } + } + + return receiveStreamBenchmarkResult{ + bytesRead: bytesRead, + flowBytesRead: flowController.bytesRead, + highestReceived: flowController.highestReceived, + finalOffset: flowController.finalOffset, + callbacks: callbacks, + completions: sender.completions, + } +}