Repository: pion/mediadevices
Pull request #123 (2020-03-26 03:20:31)

Add bitrate measurement
This helper is going to be useful for writing unit tests for codec package, especially with rate limiting feature.


Patch:
diff --git a/internal/codec/measurement.go b/internal/codec/measurement.go
new file mode 100644
index 0000000..c14e45b
--- /dev/null
+++ b/internal/codec/measurement.go
@@ -0,0 +1,41 @@
+package codec
+
+import (
+	"io"
+	"time"
+
+	mio "github.com/pion/mediadevices/pkg/io"
+)
+
+// MeasureBitRate measures average bitrate after dur by reading r as fast as possible
+func MeasureBitRate(r io.Reader, dur time.Duration) (float64, error) {
+	var n, totalBytes int
+	var err error
+	buf := make([]byte, 1024)
+	start := time.Now()
+	now := start
+	end := now.Add(dur)
+	for now.Before(end) {
+		n, err = r.Read(buf)
+		if err != nil {
+			if e, ok := err.(*mio.InsufficientBufferError); ok {
+				buf = make([]byte, 2*e.RequiredSize)
+				continue
+			}
+
+			if err == io.EOF {
+				totalBytes += n
+				break
+			}
+
+			return 0, err
+		}
+
+		totalBytes += n
+		now = time.Now()
+	}
+
+	elapsed := time.Now().Sub(start).Seconds()
+	avg := float64(totalBytes*8) / elapsed
+	return avg, nil
+}


Test patch:
diff --git a/internal/codec/measurement_test.go b/internal/codec/measurement_test.go
new file mode 100644
index 0000000..0b2d711
--- /dev/null
+++ b/internal/codec/measurement_test.go
@@ -0,0 +1,97 @@
+package codec
+
+import (
+	"io"
+	"sync"
+	"testing"
+	"time"
+)
+
+func TestMeasureBitRateStatic(t *testing.T) {
+	r, w := io.Pipe()
+	dur := time.Second * 5
+	dataSize := 1000
+	var precision float64 = 8 // 1 byte
+	var wg sync.WaitGroup
+	wg.Add(1)
+	done := make(chan struct{})
+
+	go func() {
+		data := make([]byte, dataSize)
+		// Make sure that this goroutine is synchronized with main goroutine
+		wg.Done()
+		ticker := time.NewTicker(time.Second)
+
+		for {
+			select {
+			case <-ticker.C:
+				w.Write(data)
+			case <-done:
+				w.Close()
+				return
+			}
+		}
+	}()
+
+	wg.Wait()
+	bitrate, err := MeasureBitRate(r, dur)
+	if err != nil {
+		t.Error(err)
+	}
+	done <- struct{}{}
+
+	expected := float64(dataSize * 8)
+	if bitrate < expected-precision || bitrate > expected+precision {
+		t.Fatalf("expected: %f (with %f precision), but got %f", expected, precision, bitrate)
+	}
+}
+
+func TestMeasureBitRateDynamic(t *testing.T) {
+	r, w := io.Pipe()
+	dur := time.Second * 5
+	dataSize := 1000
+	var precision float64 = 8 // 1 byte
+	var wg sync.WaitGroup
+	wg.Add(1)
+	done := make(chan struct{})
+
+	go func() {
+		data := make([]byte, dataSize)
+		wg.Done()
+		ticker := time.NewTicker(time.Millisecond * 500)
+		var count int
+
+		for {
+			select {
+			case <-ticker.C:
+				w.Write(data)
+				count++
+				// Wait until 4 slow ticks, which is also equal to 2 seconds
+				if count == 4 {
+					ticker.Stop()
+					// Speed up the tick by 2 times for the rest
+					ticker = time.NewTicker(time.Millisecond * 250)
+				}
+			case <-done:
+				w.Close()
+				return
+			}
+		}
+	}()
+
+	wg.Wait()
+	bitrate, err := MeasureBitRate(r, dur)
+	if err != nil {
+		t.Error(err)
+	}
+	done <- struct{}{}
+
+	// Measure the expected bitrate using the number of ticks or write * data size / duration in seconds
+	// For the first 2 seconds: #ticks = 2000ms / 500ms = 4
+	// For the last 3 seconds: #ticks = 3000ms / 250ms = 12
+	// So, in 5 seconds, there will be 16 ticks
+	expected := float64(16*(dataSize*8)) / 5
+	if bitrate < expected-precision || bitrate > expected+precision {
+		t.Fatalf("expected: %f (with %f precision), but got %f", expected, precision, bitrate)
+	}
+}