Avoid stalling when packet sizes exceed burst size

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

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")
+				}
+			})
+		}
+	})
 }