Implement pacing interceptor

go.mod

@@ -1,6 +1,6 @@
 module github.com/pion/interceptor
 
-go 1.20
+go 1.21
 
 require (
 	github.com/pion/logging v0.2.3

pkg/pacing/interceptor.go

@@ -0,0 +1,144 @@
+package pacing
+
+import (
+	"context"
+	"errors"
+	"log"
+	"sync"
+	"time"
+
+	"github.com/pion/interceptor"
+	"github.com/pion/logging"
+	"github.com/pion/rtp"
+)
+
+type Pacer interface {
+	Budget(t time.Time) int
+	OnSent(t time.Time, size int)
+}
+
+type PacerFactory func() Pacer
+
+type InterceptorFactory struct {
+	pf   PacerFactory
+	opts []Option
+}
+
+type Option func(*Interceptor) error
+
+func NewInterceptor(pf PacerFactory) (*InterceptorFactory, error) {
+	return &InterceptorFactory{
+		pf: pf,
+	}, nil
+}
+
+func Interval(interval time.Duration) Option {
+	return func(i *Interceptor) error {
+		i.interval = interval
+		return nil
+	}
+}
+
+func QueueSize(size int) Option {
+	return func(i *Interceptor) error {
+		i.queueSize = size
+		return nil
+	}
+}
+
+func (f *InterceptorFactory) NewInterceptor(id string) (interceptor.Interceptor, error) {
+	ctx, cancel := context.WithCancel(context.Background())
+	i := &Interceptor{
+		NoOp:     interceptor.NoOp{},
+		lock:     sync.Mutex{},
+		log:      logging.NewDefaultLoggerFactory().NewLogger("pacer_interceptor"),
+		interval: 5 * time.Millisecond,
+		pacer:    f.pf(),
+		queue:    nil,
+		close:    ctx,
+		cancelFn: cancel,
+	}
+	for _, opt := range f.opts {
+		if err := opt(i); err != nil {
+			return nil, err
+		}
+	}
+	i.queue = make(chan packetToSend, i.queueSize)
+	go i.run()
+	return i, nil
+}
+
+type Interceptor struct {
+	interceptor.NoOp
+	lock      sync.Mutex
+	log       logging.LeveledLogger
+	interval  time.Duration
+	pacer     Pacer
+	queueSize int
+	queue     chan packetToSend
+	close     context.Context
+	cancelFn  context.CancelFunc
+}
+
+type packetToSend struct {
+	header     *rtp.Header
+	payload    []byte
+	attributes interceptor.Attributes
+	writer     interceptor.RTPWriter
+}
+
+// BindLocalStream implements interceptor.Interceptor.
+func (i *Interceptor) BindLocalStream(info *interceptor.StreamInfo, writer interceptor.RTPWriter) interceptor.RTPWriter {
+	return interceptor.RTPWriterFunc(func(header *rtp.Header, payload []byte, attributes interceptor.Attributes) (int, error) {
+		ch := header.Clone()
+		buf := make([]byte, len(payload))
+		if n := copy(buf, payload); n != len(payload) {
+			return n, errors.New("copied wrong payload length")
+		}
+		select {
+		case i.queue <- packetToSend{
+			header:     &ch,
+			payload:    buf,
+			attributes: attributes,
+			writer:     writer,
+		}:
+		default:
+			return header.MarshalSize() + len(payload), errors.New("pacer dropped packet due to queue overflow")
+		}
+		return header.MarshalSize() + len(payload), nil
+	})
+}
+
+// Close implements interceptor.Interceptor.
+func (i *Interceptor) Close() error {
+	i.cancelFn()
+	return nil
+}
+
+func (i *Interceptor) run() {
+	ticker := time.NewTicker(i.interval)
+	for {
+		select {
+		case <-i.close.Done():
+			return
+		case <-ticker.C:
+			for pkt := range i.queue {
+				size := pkt.header.MarshalSize() + len(pkt.payload)
+				now := time.Now()
+				budget := i.pacer.Budget(now)
+				if budget < size {
+					log.Printf("budget too small: %v, queuesize=%v", budget, len(i.queue))
+				}
+				n, err := pkt.writer.Write(pkt.header, pkt.payload, pkt.attributes)
+				if err != nil {
+					log.Printf("error while writing packet: %v", err)
+				}
+				if n != size {
+					log.Printf("copied wrong payload length")
+				}
+				i.pacer.OnSent(now, size)
+			}
+			ticker.Reset(i.interval)
+		}
+	}
+}

pkg/pacing/leaky_bucket_pacer.go

@@ -0,0 +1,62 @@
+package pacing
+
+import (
+	"math"
+	"sync"
+	"time"
+)
+
+type LeakyBucket struct {
+	lock       sync.Mutex
+	burstSize  int
+	rate       int
+	lastSent   time.Time
+	lastBudget int
+}
+
+func NewLeakyBucketPacer(burstSize int, rate int) *LeakyBucket {
+	return &LeakyBucket{
+		lock:       sync.Mutex{},
+		burstSize:  burstSize,
+		rate:       0,
+		lastSent:   time.Time{},
+		lastBudget: 0,
+	}
+}
+
+func (b *LeakyBucket) setRate(rate int) {
+	b.lock.Lock()
+	defer b.lock.Unlock()
+
+	b.rate = rate
+}
+
+func (b *LeakyBucket) onSent(t time.Time, size int) {
+	budget := b.budget(t)
+
+	b.lock.Lock()
+	defer b.lock.Unlock()
+
+	if size > budget {
+		b.lastBudget = 0
+	} else {
+		b.lastBudget = budget - size
+	}
+	b.lastSent = t
+}
+
+func (b *LeakyBucket) budget(t time.Time) int {
+	b.lock.Lock()
+	defer b.lock.Unlock()
+
+	if b.lastSent.IsZero() {
+		return b.burstSize
+	}
+	td := t.Sub(b.lastSent)
+	budget := b.lastBudget + 8*int(float64(b.rate)*td.Seconds())
+	if budget < 0 {
+		budget = math.MaxInt
+	}
+	budget = min(budget, b.burstSize)
+	return budget
+}