Capture bandwidth switch (#14)

* Handle bitrate change by finding the stream with closest bitrate as peer

* Convert video id into bitrate when creating peer or changing bitrate

* Try to fix prometheus panic

* Revert metrics label name change

* minor fixes.

* bitrate selector.

* skip if moving to the same stream.

* no closure for getting target bitrate.

* fix: high res switch to lo video, stream bitrate out of range

* revert dev config change.

* white space.

Co-authored-by: Aleksandar Sukovic <aleksandar.sukovic@gmail.com>

internal/capture/buckets.go

@@ -2,6 +2,8 @@ package capture
 
 import (
 	"errors"
+	"fmt"
+	"math"
 
 	"github.com/rs/zerolog"
 	"github.com/rs/zerolog/log"
@@ -36,17 +38,17 @@ func (m *BucketsManagerCtx) shutdown() {
 }
 
 func (m *BucketsManagerCtx) destroyAll() {
-	for _, video := range m.streams {
-		if video.Started() {
-			video.destroyPipeline()
+	for _, stream := range m.streams {
+		if stream.Started() {
+			stream.destroyPipeline()
 		}
 	}
 }
 
 func (m *BucketsManagerCtx) recreateAll() error {
-	for _, video := range m.streams {
-		if video.Started() {
-			err := video.createPipeline()
+	for _, stream := range m.streams {
+		if stream.Started() {
+			err := stream.createPipeline()
 			if err != nil && !errors.Is(err, types.ErrCapturePipelineAlreadyExists) {
 				return err
 			}
@@ -65,22 +67,39 @@ func (m *BucketsManagerCtx) Codec() codec.RTPCodec {
 }
 
 func (m *BucketsManagerCtx) SetReceiver(receiver types.Receiver) error {
-	receiver.OnVideoIdChange(func(videoID string) error {
-		videoStream, ok := m.streams[videoID]
+	receiver.OnBitrateChange(func(bitrate int) error {
+		stream, ok := m.findNearestStream(bitrate)
 		if !ok {
-			return types.ErrWebRTCVideoNotFound
+			return fmt.Errorf("no stream found for bitrate %d", bitrate)
 		}
 
-		return receiver.SetStream(videoStream)
+		return receiver.SetStream(stream)
 	})
 
-	// TODO: Save receiver.
 	return nil
 }
 
+func (m *BucketsManagerCtx) findNearestStream(bitrate int) (ss *StreamSinkManagerCtx, ok bool) {
+	minDiff := math.MaxInt
+	for _, s := range m.streams {
+		streamBitrate, err := s.Bitrate()
+		if err != nil {
+			m.logger.Error().Err(err).Msgf("failed to get bitrate for stream %s", s.ID())
+			continue
+		}
+
+		diffAbs := int(math.Abs(float64(bitrate - streamBitrate)))
+
+		if diffAbs < minDiff {
+			minDiff, ss = diffAbs, s
+		}
+	}
+	ok = ss != nil
+	return
+}
+
 func (m *BucketsManagerCtx) RemoveReceiver(receiver types.Receiver) error {
-	// TODO: Unsubribe from OnVideoIdChange.
-	// TODO: Remove receiver.
+	receiver.OnBitrateChange(nil)
 	receiver.RemoveStream()
 	return nil
 }

internal/capture/manager.go

@@ -16,6 +16,7 @@ import (
 type CaptureManagerCtx struct {
 	logger  zerolog.Logger
 	desktop types.DesktopManager
+	config  *config.Capture
 
 	// sinks
 	broadcast  *BroacastManagerCtx
@@ -66,13 +67,18 @@ func New(desktop types.DesktopManager, config *config.Capture) *CaptureManagerCt
 			Str("pipeline", pipeline).
 			Msg("syntax check for video stream pipeline passed")
 
+		getVideoBitrate := pipelineConf.GetBitrateFn(desktop.GetScreenSize)
+		if err != nil {
+			logger.Panic().Err(err).Msg("unable to get video bitrate")
+		}
 		// append to videos
-		videos[video_id] = streamSinkNew(config.VideoCodec, createPipeline, video_id)
+		videos[video_id] = streamSinkNew(config.VideoCodec, createPipeline, video_id, getVideoBitrate)
 	}
 
 	return &CaptureManagerCtx{
 		logger:  logger,
 		desktop: desktop,
+		config:  config,
 
 		// sinks
 		broadcast: broadcastNew(func(url string) (string, error) {
@@ -132,7 +138,7 @@ func New(desktop types.DesktopManager, config *config.Capture) *CaptureManagerCt
 					"! %s "+
 					"! appsink name=appsink", config.AudioDevice, config.AudioCodec.Pipeline,
 			), nil
-		}, "audio"),
+		}, "audio", nil),
 		video: bucketsNew(config.VideoCodec, videos, config.VideoIDs),
 
 		// sources
@@ -242,6 +248,15 @@ func (manager *CaptureManagerCtx) Shutdown() error {
 	return nil
 }
 
+func (manager *CaptureManagerCtx) GetBitrateFromVideoID(videoID string) (int, error) {
+	cfg, ok := manager.config.VideoPipelines[videoID]
+	if !ok {
+		return 0, fmt.Errorf("video config not found for %s", videoID)
+	}
+
+	return cfg.GetBitrateFn(manager.desktop.GetScreenSize)()
+}
+
 func (manager *CaptureManagerCtx) Broadcast() types.BroadcastManager {
 	return manager.broadcast
 }

internal/capture/streamsink.go

@@ -19,6 +19,9 @@ import (
 var moveSinkListenerMu = sync.Mutex{}
 
 type StreamSinkManagerCtx struct {
+	id         string
+	getBitrate func() (int, error)
+
 	logger zerolog.Logger
 	mu     sync.Mutex
 	wg     sync.WaitGroup
@@ -37,13 +40,16 @@ type StreamSinkManagerCtx struct {
 	pipelinesActive  prometheus.Gauge
 }
 
-func streamSinkNew(codec codec.RTPCodec, pipelineFn func() (string, error), video_id string) *StreamSinkManagerCtx {
+func streamSinkNew(codec codec.RTPCodec, pipelineFn func() (string, error), id string, getBitrate func() (int, error)) *StreamSinkManagerCtx {
 	logger := log.With().
 		Str("module", "capture").
 		Str("submodule", "stream-sink").
-		Str("video_id", video_id).Logger()
+		Str("id", id).Logger()
 
 	manager := &StreamSinkManagerCtx{
+		id:         id,
+		getBitrate: getBitrate,
+
 		logger:     logger,
 		codec:      codec,
 		pipelineFn: pipelineFn,
@@ -56,7 +62,7 @@ func streamSinkNew(codec codec.RTPCodec, pipelineFn func() (string, error), vide
 			Subsystem: "capture",
 			Help:      "Current number of listeners for a pipeline.",
 			ConstLabels: map[string]string{
-				"video_id":   video_id,
+				"video_id":   id,
 				"codec_name": codec.Name,
 				"codec_type": codec.Type.String(),
 			},
@@ -68,7 +74,7 @@ func streamSinkNew(codec codec.RTPCodec, pipelineFn func() (string, error), vide
 			Help:      "Total number of created pipelines.",
 			ConstLabels: map[string]string{
 				"submodule":  "streamsink",
-				"video_id":   video_id,
+				"video_id":   id,
 				"codec_name": codec.Name,
 				"codec_type": codec.Type.String(),
 			},
@@ -80,7 +86,7 @@ func streamSinkNew(codec codec.RTPCodec, pipelineFn func() (string, error), vide
 			Help:      "Total number of active pipelines.",
 			ConstLabels: map[string]string{
 				"submodule":  "streamsink",
-				"video_id":   video_id,
+				"video_id":   id,
 				"codec_name": codec.Name,
 				"codec_type": codec.Type.String(),
 			},
@@ -103,6 +109,18 @@ func (manager *StreamSinkManagerCtx) shutdown() {
 	manager.wg.Wait()
 }
 
+func (manager *StreamSinkManagerCtx) ID() string {
+	return manager.id
+}
+
+func (manager *StreamSinkManagerCtx) Bitrate() (int, error) {
+	if manager.getBitrate == nil {
+		return 0, nil
+	}
+	// recalculate bitrate every time, take screen resolution (and fps) into account
+	return manager.getBitrate()
+}
+
 func (manager *StreamSinkManagerCtx) Codec() codec.RTPCodec {
 	return manager.codec
 }