diff --git a/pkg/pacing/interceptor.go b/pkg/pacing/interceptor.go index bd19971e..e894a80f 100644 --- a/pkg/pacing/interceptor.go +++ b/pkg/pacing/interceptor.go @@ -35,7 +35,7 @@ type Option func(*Interceptor) error // the interceptor factory. func InitialRate(rate int) Option { return func(i *Interceptor) error { - i.initialRate = rate + i.rate = rate return nil } @@ -109,19 +109,21 @@ func (f *InterceptorFactory) NewInterceptor(id string) (interceptor.Interceptor, defer f.lock.Unlock() interceptor := &Interceptor{ - NoOp: interceptor.NoOp{}, - initialRate: 1_000_000, - interval: 5 * time.Millisecond, - queueSize: 1_000_000, + NoOp: interceptor.NoOp{}, + interval: 5 * time.Millisecond, + queueSize: 1_000_000, pacerFactory: func(initialRate, burst int) pacer { return newRateLimitPacer(initialRate, burst) }, - limit: nil, - queue: nil, - closed: make(chan struct{}), - wg: sync.WaitGroup{}, - id: id, - onClose: f.remove, + limit: nil, + queue: nil, + rateLock: sync.Mutex{}, + rate: 1_000_000, + mtu: 1500, + closed: make(chan struct{}), + wg: sync.WaitGroup{}, + id: id, + onClose: f.remove, } for _, opt := range f.opts { if err := opt(interceptor); err != nil { @@ -133,8 +135,8 @@ func (f *InterceptorFactory) NewInterceptor(id string) (interceptor.Interceptor, } interceptor.log = interceptor.loggerFactory.NewLogger("pacer_interceptor") interceptor.limit = interceptor.pacerFactory( - interceptor.initialRate, - burst(interceptor.initialRate, interceptor.interval), + interceptor.rate, + burst(interceptor.rate, interceptor.interval, 8*interceptor.mtu), ) interceptor.queue = make(chan packet, interceptor.queueSize) @@ -157,7 +159,6 @@ type Interceptor struct { loggerFactory logging.LoggerFactory // config - initialRate int interval time.Duration queueSize int pacerFactory pacerFactory @@ -166,6 +167,10 @@ type Interceptor struct { limit pacer queue chan packet + rateLock sync.Mutex + rate int + mtu int + // shutdown closed chan struct{} wg sync.WaitGroup @@ -174,18 +179,34 @@ type Interceptor struct { } // burst calculates the minimal burst size required to reach the given rate and -// pacing interval. -func burst(rate int, interval time.Duration) int { +// pacing interval. The burst is never smaller than minBurst, so that a packet +// of minBurst bits can always be sent. +func burst(rate int, interval time.Duration, minBurst int) int { if interval <= 0 { interval = time.Millisecond } - return max(8*1500, int(float64(rate)*interval.Seconds())) + return max(minBurst, int(float64(rate)*interval.Seconds())) } // setRate updates the pacing rate and burst of the rate limiter. func (i *Interceptor) setRate(r int) { - i.limit.SetRate(r, burst(r, i.interval)) + i.rateLock.Lock() + defer i.rateLock.Unlock() + + i.rate = r + i.limit.SetRate(r, burst(r, i.interval, 8*i.mtu)) +} + +func (i *Interceptor) growMTU(bytes int) { + i.rateLock.Lock() + defer i.rateLock.Unlock() + + if bytes <= i.mtu { + return + } + i.mtu = bytes + i.limit.SetRate(i.rate, burst(i.rate, i.interval, 8*bytes)) } // BindLocalStream implements interceptor.Interceptor. @@ -235,7 +256,7 @@ func (i *Interceptor) loop() { for { select { case now := <-ticker.C: - for len(queue) > 0 && i.limit.Budget(now) > 8*float64(queue[0].len()) { + for len(queue) > 0 && i.limit.Budget(now) >= 8*float64(queue[0].len()) { i.limit.AllowN(now, 8*queue[0].len()) var next packet next, queue = queue[0], queue[1:] @@ -244,6 +265,7 @@ func (i *Interceptor) loop() { } } case pkt := <-i.queue: + i.growMTU(pkt.len()) queue = append(queue, pkt) case <-i.closed: return diff --git a/pkg/pacing/interceptor_test.go b/pkg/pacing/interceptor_test.go index d31eb26c..5a6f2c8c 100644 --- a/pkg/pacing/interceptor_test.go +++ b/pkg/pacing/interceptor_test.go @@ -4,6 +4,7 @@ package pacing import ( + "fmt" "sync" "testing" "time" @@ -62,20 +63,23 @@ func TestBurst(t *testing.T) { name string rate int interval time.Duration + minBurst int expected int }{ - {"sub_millisecond_interval", 1_000_000, 500 * time.Microsecond, minBurst}, - {"sub_millisecond_interval_high_rate", 100_000_000, 500 * time.Microsecond, 50_000}, - {"zero_interval_defaults_to_1ms", 100_000_000, 0, 100_000}, - {"negative_interval_defaults_to_1ms", 100_000_000, -time.Second, 100_000}, - {"rate_below_min_burst", 300_000, 5 * time.Millisecond, minBurst}, - {"divides_evenly", 3_000_000, 5 * time.Millisecond, 15_000}, - {"does_not_divide_evenly", 3_000_000, 7 * time.Millisecond, 21_000}, - {"long_interval", 3_000_000, 33 * time.Millisecond, 99_000}, - {"zero_rate", 0, 5 * time.Millisecond, minBurst}, + {"sub_millisecond_interval", 1_000_000, 500 * time.Microsecond, minBurst, minBurst}, + {"sub_millisecond_interval_high_rate", 100_000_000, 500 * time.Microsecond, minBurst, 50_000}, + {"zero_interval_defaults_to_1ms", 100_000_000, 0, minBurst, 100_000}, + {"negative_interval_defaults_to_1ms", 100_000_000, -time.Second, minBurst, 100_000}, + {"rate_below_min_burst", 300_000, 5 * time.Millisecond, minBurst, minBurst}, + {"divides_evenly", 3_000_000, 5 * time.Millisecond, minBurst, 15_000}, + {"does_not_divide_evenly", 3_000_000, 7 * time.Millisecond, minBurst, 21_000}, + {"long_interval", 3_000_000, 33 * time.Millisecond, minBurst, 99_000}, + {"zero_rate", 0, 5 * time.Millisecond, minBurst, minBurst}, + {"grown_min_burst_wins", 300_000, 5 * time.Millisecond, 8 * 4000, 8 * 4000}, + {"rate_above_grown_min_burst", 100_000_000, 5 * time.Millisecond, 8 * 4000, 500_000}, } { t.Run(cc.name, func(t *testing.T) { - assert.Equal(t, cc.expected, burst(cc.rate, cc.interval)) + assert.Equal(t, cc.expected, burst(cc.rate, cc.interval, cc.minBurst)) }) } } @@ -98,6 +102,43 @@ func TestInterceptor(t *testing.T) { assert.Equal(t, 12000, mp.burst) }) + t.Run("grows_burst", func(t *testing.T) { + mp := &mockPacer{} + factory := NewInterceptor( + setPacerFactory(func(int, int) pacer { + return mp + }), + InitialRate(300_000), + Interval(5*time.Millisecond), + ) + + created, err := factory.NewInterceptor("id") + assert.NoError(t, err) + pacer, ok := created.(*Interceptor) + assert.True(t, ok) + defer func() { + assert.NoError(t, pacer.Close()) + }() + + pacer.growMTU(1500) + mp.lock.Lock() + assert.Equal(t, 0, mp.burst) + mp.lock.Unlock() + + pacer.growMTU(4000) + mp.lock.Lock() + assert.Equal(t, 300_000, mp.rate) + assert.Equal(t, 8*4000, mp.burst) + mp.lock.Unlock() + + pacer.growMTU(1500) + factory.SetRate("id", 600_000) + mp.lock.Lock() + assert.Equal(t, 600_000, mp.rate) + assert.Equal(t, 8*4000, mp.burst) + mp.lock.Unlock() + }) + t.Run("paces_packets", func(t *testing.T) { mp := &mockPacer{ rate: 0, @@ -167,4 +208,40 @@ func TestInterceptor(t *testing.T) { case <-time.After(10 * time.Millisecond): } }) + + // A packet is only sent once the budget covers its full size, but the + // budget never exceeds the burst size. Packets at or above the burst size + // must still be sent, instead of blocking the queue forever. + t.Run("sends_packets_at_and_above_burst_size", func(t *testing.T) { + // At 300 kbps and a 5ms pacing interval the rate alone allows a burst + // of only 1500 bits. + for _, size := range []int{20, 1499, 1500, 1501, 4000} { + t.Run(fmt.Sprintf("%d_bytes", size), func(t *testing.T) { + i := NewInterceptor( + InitialRate(300_000), + Interval(5*time.Millisecond), + ) + + pacer, err := i.NewInterceptor("") + assert.NoError(t, err) + + stream := test.NewMockStream(&interceptor.StreamInfo{}, pacer) + defer func() { + assert.NoError(t, stream.Close()) + }() + + hdr := rtp.Header{} + assert.NoError(t, stream.WriteRTP(&rtp.Packet{ + Header: hdr, + Payload: make([]byte, size-hdr.MarshalSize()), + })) + + select { + case <-stream.WrittenRTP(): + case <-time.After(2 * time.Second): + assert.Fail(t, "no RTP packet written") + } + }) + } + }) }