Squashed 'livekit-server/' content from commit 154b4d26

git-subtree-dir: livekit-server git-subtree-split: 154b4d26b769c68a03c096124094b97bf61a996f
2026-06-25 14:35:28 +09:00
commit 0da97ebd21
339 changed files with 114111 additions and 0 deletions
@@ -0,0 +1,109 @@
+package datachannel
+
+import (
+	"sync"
+	"time"
+
+	"github.com/gammazero/deque"
+
+	"github.com/livekit/protocol/utils/mono"
+)
+
+const (
+	BitrateDuration = 2 * time.Second
+	BitrateWindow   = 100 * time.Millisecond
+)
+
+// BitrateCalculator calculates bitrate over sliding window
+type BitrateCalculator struct {
+	lock           sync.Mutex
+	windowDuration time.Duration
+	duration       time.Duration
+
+	windows deque.Deque[bitrateWindow]
+	active  bitrateWindow
+
+	bytes              int
+	lastBufferedAmount int
+	start              time.Time
+}
+
+func NewBitrateCalculator(duration time.Duration, window time.Duration) *BitrateCalculator {
+	windowCnt := int((duration + (window - 1)) / window)
+	if windowCnt == 0 {
+		windowCnt = 1
+	}
+	now := mono.Now()
+	c := &BitrateCalculator{
+		duration:       duration,
+		windowDuration: window,
+		start:          now,
+		active:         bitrateWindow{start: now},
+	}
+	c.windows.SetBaseCap(windowCnt + 1)
+
+	return c
+}
+
+func (c *BitrateCalculator) AddBytes(bytes int, bufferedAmout int, ts time.Time) {
+	c.lock.Lock()
+	defer c.lock.Unlock()
+
+	bytes -= bufferedAmout - c.lastBufferedAmount
+	if bytes < 0 {
+		// it is possible that internal buffering (non-data like DCEP packet from webrtc) caused bytes to be negative
+		bytes = 0
+	}
+	c.lastBufferedAmount = bufferedAmout
+	if ts.Sub(c.active.start) >= c.windowDuration {
+		c.windows.PushBack(c.active)
+		c.active.start = ts
+		c.active.bytes = 0
+
+		for c.windows.Len() > 0 {
+			// pop expired windows
+			if w := c.windows.Front(); ts.Sub(w.start) > (c.duration + c.windowDuration) {
+				c.bytes -= w.bytes
+				c.windows.PopFront()
+			} else {
+				c.start = w.start
+				break
+			}
+		}
+		if c.windows.Len() == 0 {
+			c.start = ts
+			c.bytes = 0
+		}
+	}
+	c.bytes += bytes
+	c.active.bytes += bytes
+
+}
+
+func (c *BitrateCalculator) Bitrate(ts time.Time) (int, bool) {
+	return c.bitrate(ts, false)
+}
+
+func (c *BitrateCalculator) ForceBitrate(ts time.Time) (int, bool) {
+	return c.bitrate(ts, true)
+}
+
+func (c *BitrateCalculator) bitrate(ts time.Time, force bool) (int, bool) {
+	c.lock.Lock()
+	defer c.lock.Unlock()
+	duration := ts.Sub(c.start)
+	if duration < c.windowDuration {
+		if force {
+			duration = c.windowDuration
+		} else {
+			return 0, false
+		}
+	}
+
+	return c.bytes * 8 * 1000 / int(duration.Milliseconds()), true
+}
+
+type bitrateWindow struct {
+	start time.Time
+	bytes int
+}
@@ -0,0 +1,36 @@
+package datachannel
+
+import (
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestBitrateCalculator(t *testing.T) {
+	c := NewBitrateCalculator(BitrateDuration, BitrateWindow)
+	require.NotNil(t, c)
+
+	t0 := time.Now()
+	c.AddBytes(100, 0, t0)
+	// bytes buffered
+	c.AddBytes(100, 100, t0.Add(50*time.Millisecond))
+	bitrate, ok := c.Bitrate(t0.Add(50 * time.Millisecond))
+	require.Equal(t, 0, bitrate)
+	require.False(t, ok)
+	// 50 bytes sent (50 bytes buffer flushed)
+	c.AddBytes(100, 50, t0.Add(time.Second))
+
+	// 250 bytes sent in 1 second
+	bitrate, ok = c.Bitrate(t0.Add(time.Second))
+	require.Equal(t, 2000, bitrate)
+	require.True(t, ok)
+
+	// silence for long time
+	t1 := t0.Add(2 * BitrateDuration)
+	// 150 bytes sent (50 bytes buffer flushed)
+	c.AddBytes(100, 0, t1)
+	bitrate, ok = c.Bitrate(t1.Add(time.Second))
+	require.Equal(t, 1200, bitrate)
+	require.True(t, ok)
+}
@@ -0,0 +1,80 @@
+package datachannel
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"time"
+
+	"github.com/pion/datachannel"
+
+	"github.com/livekit/protocol/utils/mono"
+)
+
+const (
+	singleWriteTimeout = 50 * time.Millisecond
+)
+
+var ErrDataDroppedBySlowReader = errors.New("data dropped by slow reader")
+
+type BufferedAmountGetter interface {
+	BufferedAmount() uint64
+}
+
+type DataChannelWriter[T BufferedAmountGetter] struct {
+	bufferGetter  T
+	rawDC         datachannel.ReadWriteCloserDeadliner
+	slowThreshold int
+	rate          *BitrateCalculator
+}
+
+// NewDataChannelWriter creates a new DataChannelWriter by detaching the data channel, when
+// writing to the datachanel times out, it will block and retry if the receiver's bitrate is
+// above the slowThreshold or drop the data if it's below the threshold. If the slowThreshold
+// is 0, it will never retry on write timeout.
+func NewDataChannelWriter[T BufferedAmountGetter](bufferGetter T, rawDC datachannel.ReadWriteCloserDeadliner, slowThreshold int) *DataChannelWriter[T] {
+	var rate *BitrateCalculator
+	if slowThreshold > 0 {
+		rate = NewBitrateCalculator(BitrateDuration, BitrateWindow)
+	}
+	return &DataChannelWriter[T]{
+		bufferGetter:  bufferGetter,
+		rawDC:         rawDC,
+		slowThreshold: slowThreshold,
+		rate:          rate,
+	}
+}
+
+func (w *DataChannelWriter[T]) BufferedAmountGetter() T {
+	return w.bufferGetter
+}
+
+func (w *DataChannelWriter[T]) Write(p []byte) (n int, err error) {
+	for {
+		err = w.rawDC.SetWriteDeadline(time.Now().Add(singleWriteTimeout))
+		if err != nil {
+			return 0, err
+		}
+		n, err = w.rawDC.Write(p)
+		if w.slowThreshold == 0 {
+			return
+		}
+
+		now := mono.Now()
+		w.rate.AddBytes(n, int(w.bufferGetter.BufferedAmount()), now)
+		// retry if the write timed out on a non-slow receiver
+		if errors.Is(err, context.DeadlineExceeded) {
+			if bitrate, ok := w.rate.Bitrate(now); !ok || bitrate >= w.slowThreshold {
+				continue
+			} else {
+				err = fmt.Errorf("%w: bitrate %d, threshold %d", ErrDataDroppedBySlowReader, bitrate, w.slowThreshold)
+			}
+		}
+
+		return
+	}
+}
+
+func (w *DataChannelWriter[T]) Close() error {
+	return w.rawDC.Close()
+}
@@ -0,0 +1,94 @@
+package datachannel
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/pion/datachannel"
+	"github.com/pion/transport/v3/deadline"
+	"github.com/stretchr/testify/require"
+)
+
+func TestDataChannelWriter(t *testing.T) {
+	mockDC := newMockDataChannelWriter()
+	// slow threshold is 1000B/s
+	w := NewDataChannelWriter(mockDC, mockDC, 8000)
+	require.Equal(t, mockDC, w.BufferedAmountGetter())
+	buf := make([]byte, 2000)
+	// write 2000 bytes so it should not drop in 2 seconds
+	t0 := time.Now()
+	n, err := w.Write(buf)
+	require.NoError(t, err)
+	require.Equal(t, 2000, n)
+
+	t1 := time.Now()
+	mockDC.SetNextWriteCompleteAt(t0.Add(time.Second))
+	n, err = w.Write(buf[:10])
+	require.NoError(t, err)
+	require.Equal(t, 10, n)
+	require.GreaterOrEqual(t, time.Since(t1), time.Second)
+
+	// bitrate below slow threshold(2000bytes/3sec), should drop by timeout
+	mockDC.SetNextWriteCompleteAt(t0.Add(3 * time.Second))
+	n, err = w.Write(buf[:1000])
+	require.ErrorIs(t, err, ErrDataDroppedBySlowReader, err)
+	require.Equal(t, 0, n)
+}
+
+func TestDataChannelWriter_NoSlowThreshold(t *testing.T) {
+	mockDC := newMockDataChannelWriter()
+	w := NewDataChannelWriter(mockDC, mockDC, 0)
+	buf := make([]byte, 2000)
+	n, err := w.Write(buf)
+	require.NoError(t, err)
+	require.Equal(t, 2000, n)
+	mockDC.SetNextWriteCompleteAt(time.Now().Add(singleWriteTimeout / 2))
+	n, err = w.Write(buf[:10])
+	require.NoError(t, err)
+	require.Equal(t, 10, n)
+
+	// slow threshold is 0, should not block & retry
+	mockDC.SetNextWriteCompleteAt(time.Now().Add(singleWriteTimeout * 2))
+	n, err = w.Write(buf[:1000])
+	require.ErrorIs(t, err, context.DeadlineExceeded, err)
+	require.Equal(t, 0, n)
+}
+
+type mockDataChannelWriter struct {
+	datachannel.ReadWriteCloserDeadliner
+	nextWriteCompleteAt time.Time
+	deadline            *deadline.Deadline
+}
+
+func newMockDataChannelWriter() *mockDataChannelWriter {
+	return &mockDataChannelWriter{
+		deadline: deadline.New(),
+	}
+}
+
+func (m *mockDataChannelWriter) BufferedAmount() uint64 {
+	return 0
+}
+
+func (m *mockDataChannelWriter) Write(b []byte) (int, error) {
+	wait := time.Until(m.nextWriteCompleteAt)
+	if wait <= 0 {
+		return len(b), nil
+	}
+	select {
+	case <-m.deadline.Done():
+		return 0, m.deadline.Err()
+	case <-time.After(wait):
+		return len(b), nil
+	}
+}
+
+func (m *mockDataChannelWriter) SetWriteDeadline(t time.Time) error {
+	m.deadline.Set(t)
+	return nil
+}
+
+func (m *mockDataChannelWriter) SetNextWriteCompleteAt(t time.Time) {
+	m.nextWriteCompleteAt = t
+}