diff --git a/pkg/codec/bitrate_tracker.go b/pkg/codec/bitrate_tracker.go new file mode 100644 index 0000000..dbcbcaf --- /dev/null +++ b/pkg/codec/bitrate_tracker.go @@ -0,0 +1,48 @@ +package codec + +import ( + "time" +) + +type BitrateTracker struct { + windowSize time.Duration + buffer []int + times []time.Time +} + +func NewBitrateTracker(windowSize time.Duration) *BitrateTracker { + return &BitrateTracker{ + windowSize: windowSize, + } +} + +func (bt *BitrateTracker) AddFrame(sizeBytes int, timestamp time.Time) { + bt.buffer = append(bt.buffer, sizeBytes) + bt.times = append(bt.times, timestamp) + + // Remove old entries outside the window + cutoff := timestamp.Add(-bt.windowSize) + i := 0 + for ; i < len(bt.times); i++ { + if bt.times[i].After(cutoff) { + break + } + } + bt.buffer = bt.buffer[i:] + bt.times = bt.times[i:] +} + +func (bt *BitrateTracker) GetBitrate() float64 { + if len(bt.times) < 2 { + return 0 + } + totalBytes := 0 + for _, b := range bt.buffer { + totalBytes += b + } + duration := bt.times[len(bt.times)-1].Sub(bt.times[0]).Seconds() + if duration <= 0 { + return 0 + } + return float64(totalBytes*8) / duration // bits per second +} diff --git a/pkg/codec/bitrate_tracker_test.go b/pkg/codec/bitrate_tracker_test.go new file mode 100644 index 0000000..ba3d2db --- /dev/null +++ b/pkg/codec/bitrate_tracker_test.go @@ -0,0 +1,19 @@ +package codec + +import ( + "math" + "testing" + "time" +) + +func TestBitrateTracker(t *testing.T) { + packetSize := 1000 + bt := NewBitrateTracker(time.Second) + bt.AddFrame(packetSize, time.Now()) + bt.AddFrame(packetSize, time.Now().Add(time.Millisecond*100)) + bt.AddFrame(packetSize, time.Now().Add(time.Millisecond*999)) + eps := float64(packetSize*8) / 10 + if got, want := bt.GetBitrate(), float64(packetSize*8)*3; math.Abs(got-want) > eps { + t.Fatalf("GetBitrate() = %v, want %v (|diff| <= %v)", got, want, eps) + } +}