From f4d632df2cb0a6109229c363c498ec304e0c912d Mon Sep 17 00:00:00 2001 From: a <57275181+tuxpainter@users.noreply.github.com> Date: Thu, 24 Sep 2026 03:58:39 -0500 Subject: [PATCH] feat: create capture subscription api to support more transports (#699) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: introduce encoded media subscriptions * chore: preserve sample field comments * perf: reduce encoded sample dispatch copies * docs: define encoded sample buffer ownership * fix: avoid pipeline recreation errors Reuse pipelines started by a new subscription during resize. * fix: synchronize stream bitrate state Protect bitrate updates, reads, and resets with one mutex. * simplify * group stream sink fields * fix: remove unnecessary blank line in removeSubscription method * remove defensive check * remove noop --------- Co-authored-by: Miroslav Šedivý --- server/internal/capture/manager.go | 6 +- server/internal/capture/streamselector.go | 29 ++- server/internal/capture/streamsink.go | 250 ++++++++++++++------- server/internal/capture/streamsink_test.go | 232 +++++++++++++++++++ server/internal/webrtc/peer.go | 4 +- server/internal/webrtc/track.go | 91 ++++---- server/pkg/gst/gst.c | 11 +- server/pkg/gst/gst.go | 6 +- server/pkg/gst/gst.h | 2 +- server/pkg/types/capture.go | 49 ++-- 10 files changed, 505 insertions(+), 175 deletions(-) create mode 100644 server/internal/capture/streamsink_test.go diff --git a/server/internal/capture/manager.go b/server/internal/capture/manager.go index 5872c90f..adeac6c6 100644 --- a/server/internal/capture/manager.go +++ b/server/internal/capture/manager.go @@ -33,7 +33,7 @@ type CaptureManagerCtx struct { func New(desktop types.DesktopManager, config *config.Capture) *CaptureManagerCtx { logger := log.With().Str("module", "capture").Logger() - videos := map[string]types.StreamSinkManager{} + videos := map[string]*StreamSinkManagerCtx{} for video_id, cnf := range config.VideoPipelines { pipelineConf := cnf @@ -257,11 +257,11 @@ func (manager *CaptureManagerCtx) Screencast() types.ScreencastManager { return manager.screencast } -func (manager *CaptureManagerCtx) Audio() types.StreamSinkManager { +func (manager *CaptureManagerCtx) Audio() types.EncodedStream { return manager.audio } -func (manager *CaptureManagerCtx) Video() types.StreamSelectorManager { +func (manager *CaptureManagerCtx) Video() types.EncodedStreamSelector { return manager.video } diff --git a/server/internal/capture/streamselector.go b/server/internal/capture/streamselector.go index d646e56c..e17fe869 100644 --- a/server/internal/capture/streamselector.go +++ b/server/internal/capture/streamselector.go @@ -1,7 +1,6 @@ package capture import ( - "errors" "sort" "github.com/rs/zerolog" @@ -14,11 +13,11 @@ import ( type StreamSelectorManagerCtx struct { logger zerolog.Logger codec codec.RTPCodec - streams map[string]types.StreamSinkManager + streams map[string]*StreamSinkManagerCtx streamIDs []string } -func streamSelectorNew(codec codec.RTPCodec, streams map[string]types.StreamSinkManager, streamIDs []string) *StreamSelectorManagerCtx { +func streamSelectorNew(codec codec.RTPCodec, streams map[string]*StreamSinkManagerCtx, streamIDs []string) *StreamSelectorManagerCtx { logger := log.With(). Str("module", "capture"). Str("submodule", "stream-selector"). @@ -40,17 +39,17 @@ func (manager *StreamSelectorManagerCtx) shutdown() { func (manager *StreamSelectorManagerCtx) destroyPipelines() { for _, stream := range manager.streams { - if stream.Started() { - stream.DestroyPipeline() + if stream.started() { + stream.destroyPipeline() } } } func (manager *StreamSelectorManagerCtx) recreatePipelines() error { for _, stream := range manager.streams { - if stream.Started() { - err := stream.CreatePipeline() - if err != nil && !errors.Is(err, types.ErrCapturePipelineAlreadyExists) { + if stream.started() { + err := stream.createPipeline() + if err != nil { return err } } @@ -66,12 +65,12 @@ func (manager *StreamSelectorManagerCtx) Codec() codec.RTPCodec { return manager.codec } -func (manager *StreamSelectorManagerCtx) GetStream(selector types.StreamSelector) (types.StreamSinkManager, bool) { +func (manager *StreamSelectorManagerCtx) GetStream(selector types.StreamSelector) (types.EncodedStream, bool) { // select stream by ID if selector.ID != "" { // select lower stream if selector.Type == types.StreamSelectorTypeLower { - var lastStream types.StreamSinkManager + var lastStream types.EncodedStream for i := len(manager.streamIDs) - 1; i >= 0; i-- { streamID := manager.streamIDs[i] if streamID == selector.ID { @@ -88,7 +87,7 @@ func (manager *StreamSelectorManagerCtx) GetStream(selector types.StreamSelector // select higher stream if selector.Type == types.StreamSelectorTypeHigher { - var lastStream types.StreamSinkManager + var lastStream types.EncodedStream for _, streamID := range manager.streamIDs { if streamID == selector.ID { return lastStream, lastStream != nil @@ -121,7 +120,7 @@ func (manager *StreamSelectorManagerCtx) GetStream(selector types.StreamSelector streamID := manager.streamIDs[i] stream := manager.streams[streamID] // if stream should be considered in calculation - considered := stream.Bitrate() != 0 && stream.Started() + considered := stream.Bitrate() != 0 && stream.started() if considered && stream.Bitrate() < selector.Bitrate { return stream, true } @@ -136,7 +135,7 @@ func (manager *StreamSelectorManagerCtx) GetStream(selector types.StreamSelector for _, streamID := range manager.streamIDs { stream := manager.streams[streamID] // if stream should be considered in calculation - considered := stream.Bitrate() != 0 && stream.Started() + considered := stream.Bitrate() != 0 && stream.started() if considered && stream.Bitrate() > selector.Bitrate { return stream, true } @@ -158,7 +157,7 @@ func (manager *StreamSelectorManagerCtx) GetStream(selector types.StreamSelector } // TODO: This is a very naive implementation, we should use a binary search instead. -func (manager *StreamSelectorManagerCtx) nearestBitrate(bitrate uint64) types.StreamSinkManager { +func (manager *StreamSelectorManagerCtx) nearestBitrate(bitrate uint64) types.EncodedStream { type streamDiff struct { id string bitrateDiff int @@ -181,7 +180,7 @@ func (manager *StreamSelectorManagerCtx) nearestBitrate(bitrate uint64) types.St for _, stream := range manager.streams { // if stream should be considered in calculation - considered := stream.Bitrate() != 0 && stream.Started() + considered := stream.Bitrate() != 0 && stream.started() if !considered { continue } diff --git a/server/internal/capture/streamsink.go b/server/internal/capture/streamsink.go index 79ba17ff..404128b4 100644 --- a/server/internal/capture/streamsink.go +++ b/server/internal/capture/streamsink.go @@ -2,8 +2,8 @@ package capture import ( "errors" - "reflect" "sync" + "sync/atomic" "time" "github.com/prometheus/client_golang/prometheus" @@ -17,30 +17,44 @@ import ( "github.com/m1k1o/neko/server/pkg/types/codec" ) -var moveSinkListenerMu = sync.Mutex{} +var moveSinkSubscriptionMu sync.Mutex + +type sinkPipeline interface { + Sample() chan types.Sample + AttachAppsink(string) + Play() + Destroy() + EmitVideoKeyframe() bool +} type StreamSinkManagerCtx struct { - id string + // Configuration set before use. + id string + logger zerolog.Logger + codec codec.RTPCodec + waitForKf bool // wait for a keyframe before sending samples + pipelineFn func() (string, error) + pipelineFactory func(string) (sinkPipeline, error) - // wait for a keyframe before sending samples - waitForKf bool + // Serializes subscription add, remove, and move operations. + mu sync.Mutex + // Serializes pipeline creation and destruction. + pipelineMu sync.Mutex + pipeline sinkPipeline + wg sync.WaitGroup // tracks pipeline sample readers + + // Serializes listener map updates and snapshot publication. + listenersMu sync.Mutex + listeners map[*streamSubscription]types.SampleConsumer + listenersKf map[*streamSubscription]types.SampleConsumer // keyframe lobby + listenersSnapshot atomic.Value // read without listenersMu + + // Protects bitrate accounting, reads, and resets. + bitrateMu sync.Mutex bitrate uint64 brBuckets map[int]float64 - logger zerolog.Logger - mu sync.Mutex - wg sync.WaitGroup - - codec codec.RTPCodec - pipeline gst.Pipeline - pipelineMu sync.Mutex - pipelineFn func() (string, error) - - listeners map[uintptr]types.SampleListener - listenersKf map[uintptr]types.SampleListener // keyframe lobby - listenersMu sync.Mutex - // metrics currentListeners prometheus.Gauge totalBytes prometheus.Counter @@ -66,9 +80,12 @@ func streamSinkNew(codec codec.RTPCodec, pipelineFn func() (string, error), id s logger: logger, codec: codec, pipelineFn: pipelineFn, + pipelineFactory: func(src string) (sinkPipeline, error) { + return gst.CreatePipeline(src) + }, - listeners: map[uintptr]types.SampleListener{}, - listenersKf: map[uintptr]types.SampleListener{}, + listeners: map[*streamSubscription]types.SampleConsumer{}, + listenersKf: map[*streamSubscription]types.SampleConsumer{}, // metrics currentListeners: promauto.NewGauge(prometheus.GaugeOpts{ @@ -118,6 +135,7 @@ func streamSinkNew(codec codec.RTPCodec, pipelineFn func() (string, error), id s }, }), } + manager.listenersSnapshot.Store([]types.SampleConsumer(nil)) return manager } @@ -132,9 +150,10 @@ func (manager *StreamSinkManagerCtx) shutdown() { for key := range manager.listenersKf { delete(manager.listenersKf, key) } + manager.updateListenersSnapshotLocked() manager.listenersMu.Unlock() - manager.DestroyPipeline() + manager.destroyPipeline() manager.wg.Wait() } @@ -143,6 +162,9 @@ func (manager *StreamSinkManagerCtx) ID() string { } func (manager *StreamSinkManagerCtx) Bitrate() uint64 { + manager.bitrateMu.Lock() + defer manager.bitrateMu.Unlock() + return manager.bitrate } @@ -150,10 +172,76 @@ func (manager *StreamSinkManagerCtx) Codec() codec.RTPCodec { return manager.codec } +type streamSubscription struct { + mu sync.Mutex + stream *StreamSinkManagerCtx + consumer types.SampleConsumer + closed bool +} + +func (subscription *streamSubscription) Stream() types.EncodedStream { + subscription.mu.Lock() + defer subscription.mu.Unlock() + + if subscription.closed { + return nil + } + return subscription.stream +} + +func (subscription *streamSubscription) Switch(stream types.EncodedStream) error { + subscription.mu.Lock() + defer subscription.mu.Unlock() + + if subscription.closed { + return errors.New("stream subscription is closed") + } + + target, ok := stream.(*StreamSinkManagerCtx) + if !ok || target == nil { + return errors.New("target stream is incompatible with this subscription") + } + if target == subscription.stream { + return nil + } + + if err := subscription.stream.moveSubscriptionTo(subscription, target); err != nil { + return err + } + subscription.stream = target + return nil +} + +func (subscription *streamSubscription) Close() error { + subscription.mu.Lock() + defer subscription.mu.Unlock() + + if subscription.closed { + return nil + } + + subscription.stream.removeSubscription(subscription) + subscription.closed = true + subscription.stream = nil + return nil +} + +func (manager *StreamSinkManagerCtx) Subscribe(consumer types.SampleConsumer) (types.StreamSubscription, error) { + if consumer == nil { + return nil, errors.New("sample consumer cannot be nil") + } + + subscription := &streamSubscription{stream: manager, consumer: consumer} + if err := manager.addSubscription(subscription); err != nil { + return nil, err + } + return subscription, nil +} + func (manager *StreamSinkManagerCtx) start() error { if len(manager.listeners)+len(manager.listenersKf) == 0 { - err := manager.CreatePipeline() - if err != nil && !errors.Is(err, types.ErrCapturePipelineAlreadyExists) { + err := manager.createPipeline() + if err != nil { return err } @@ -165,13 +253,12 @@ func (manager *StreamSinkManagerCtx) start() error { func (manager *StreamSinkManagerCtx) stop() { if len(manager.listeners)+len(manager.listenersKf) == 0 { - manager.DestroyPipeline() + manager.destroyPipeline() manager.logger.Info().Msgf("last listener, stopping") } } -func (manager *StreamSinkManagerCtx) addListener(listener types.SampleListener) { - ptr := reflect.ValueOf(listener).Pointer() +func (manager *StreamSinkManagerCtx) addSubscriptionToMaps(subscription *streamSubscription) { emitKeyframe := false manager.listenersMu.Lock() @@ -179,15 +266,16 @@ func (manager *StreamSinkManagerCtx) addListener(listener types.SampleListener) // if this is the first listener, we need to emit a keyframe emitKeyframe = len(manager.listenersKf) == 0 // if we're waiting for a keyframe, add it to the keyframe lobby - manager.listenersKf[ptr] = listener + manager.listenersKf[subscription] = subscription.consumer } else { // otherwise, add it as a regular listener - manager.listeners[ptr] = listener + manager.listeners[subscription] = subscription.consumer + manager.updateListenersSnapshotLocked() } manager.listenersMu.Unlock() - manager.logger.Debug().Interface("ptr", ptr).Msgf("adding listener") - manager.currentListeners.Set(float64(manager.ListenersCount())) + manager.logger.Debug().Interface("subscription", subscription).Msg("adding subscription") + manager.currentListeners.Set(float64(manager.subscriptionsCount())) // if we will be waiting for a keyframe, emit one now if manager.pipeline != nil && emitKeyframe { @@ -195,72 +283,50 @@ func (manager *StreamSinkManagerCtx) addListener(listener types.SampleListener) } } -func (manager *StreamSinkManagerCtx) removeListener(listener types.SampleListener) { - ptr := reflect.ValueOf(listener).Pointer() - +func (manager *StreamSinkManagerCtx) removeSubscriptionFromMaps(subscription *streamSubscription) { manager.listenersMu.Lock() - delete(manager.listeners, ptr) - delete(manager.listenersKf, ptr) // if it's a keyframe listener, remove it too + delete(manager.listeners, subscription) + delete(manager.listenersKf, subscription) + manager.updateListenersSnapshotLocked() manager.listenersMu.Unlock() - manager.logger.Debug().Interface("ptr", ptr).Msgf("removing listener") - manager.currentListeners.Set(float64(manager.ListenersCount())) + manager.logger.Debug().Interface("subscription", subscription).Msg("removing subscription") + manager.currentListeners.Set(float64(manager.subscriptionsCount())) } -func (manager *StreamSinkManagerCtx) AddListener(listener types.SampleListener) error { +func (manager *StreamSinkManagerCtx) addSubscription(subscription *streamSubscription) error { manager.mu.Lock() defer manager.mu.Unlock() - if listener == nil { - return errors.New("listener cannot be nil") - } - // start if stopped if err := manager.start(); err != nil { return err } // add listener - manager.addListener(listener) + manager.addSubscriptionToMaps(subscription) return nil } -func (manager *StreamSinkManagerCtx) RemoveListener(listener types.SampleListener) error { +func (manager *StreamSinkManagerCtx) removeSubscription(subscription *streamSubscription) { manager.mu.Lock() defer manager.mu.Unlock() - if listener == nil { - return errors.New("listener cannot be nil") - } - - // remove listener - manager.removeListener(listener) + manager.removeSubscriptionFromMaps(subscription) // stop if started manager.stop() - - return nil } -// moving listeners between streams ensures, that target pipeline is running -// before listener is added, and stops source pipeline if there are 0 listeners -func (manager *StreamSinkManagerCtx) MoveListenerTo(listener types.SampleListener, stream types.StreamSinkManager) error { - if listener == nil { - return errors.New("listener cannot be nil") - } - - targetStream, ok := stream.(*StreamSinkManagerCtx) - if !ok { - return errors.New("target stream manager does not support moving listeners") - } - +// moveSubscriptionTo starts the target before moving delivery and stopping the source. +func (manager *StreamSinkManagerCtx) moveSubscriptionTo(subscription *streamSubscription, targetStream *StreamSinkManagerCtx) error { // we need to acquire both mutextes, from source stream and from target stream // in order to do that safely (without possibility of deadlock) we need third // global mutex, that ensures atomic locking // lock global mutex - moveSinkListenerMu.Lock() + moveSinkSubscriptionMu.Lock() // lock source stream manager.mu.Lock() @@ -271,7 +337,7 @@ func (manager *StreamSinkManagerCtx) MoveListenerTo(listener types.SampleListene defer targetStream.mu.Unlock() // unlock global mutex - moveSinkListenerMu.Unlock() + moveSinkSubscriptionMu.Unlock() // start if stopped if err := targetStream.start(); err != nil { @@ -279,8 +345,8 @@ func (manager *StreamSinkManagerCtx) MoveListenerTo(listener types.SampleListene } // swap listeners - manager.removeListener(listener) - targetStream.addListener(listener) + manager.removeSubscriptionFromMaps(subscription) + targetStream.addSubscriptionToMaps(subscription) // stop if started manager.stop() @@ -288,23 +354,23 @@ func (manager *StreamSinkManagerCtx) MoveListenerTo(listener types.SampleListene return nil } -func (manager *StreamSinkManagerCtx) ListenersCount() int { +func (manager *StreamSinkManagerCtx) subscriptionsCount() int { manager.listenersMu.Lock() defer manager.listenersMu.Unlock() return len(manager.listeners) + len(manager.listenersKf) } -func (manager *StreamSinkManagerCtx) Started() bool { - return manager.ListenersCount() > 0 +func (manager *StreamSinkManagerCtx) started() bool { + return manager.subscriptionsCount() > 0 } -func (manager *StreamSinkManagerCtx) CreatePipeline() error { +func (manager *StreamSinkManagerCtx) createPipeline() error { manager.pipelineMu.Lock() defer manager.pipelineMu.Unlock() if manager.pipeline != nil { - return types.ErrCapturePipelineAlreadyExists + return nil } pipelineStr, err := manager.pipelineFn() @@ -317,7 +383,7 @@ func (manager *StreamSinkManagerCtx) CreatePipeline() error { Str("src", pipelineStr). Msgf("creating pipeline") - manager.pipeline, err = gst.CreatePipeline(pipelineStr) + manager.pipeline, err = manager.pipelineFactory(pipelineStr) if err != nil { return err } @@ -347,6 +413,9 @@ func (manager *StreamSinkManagerCtx) CreatePipeline() error { } func (manager *StreamSinkManagerCtx) saveSampleBitrate(timestamp time.Time, delta float64) { + manager.bitrateMu.Lock() + defer manager.bitrateMu.Unlock() + // get unix timestamp in seconds sec := timestamp.Unix() // last bucket is timestamp rounded to 3 seconds - 1 second @@ -357,7 +426,7 @@ func (manager *StreamSinkManagerCtx) saveSampleBitrate(timestamp time.Time, delt next := int((sec + 1) % 3) if manager.brBuckets[next] != 0 { - // update bitrate, TODO: atomic? + // update bitrate manager.bitrate = uint64(manager.brBuckets[last]) // empty next bucket manager.brBuckets[next] = 0 @@ -368,36 +437,41 @@ func (manager *StreamSinkManagerCtx) saveSampleBitrate(timestamp time.Time, delt } func (manager *StreamSinkManagerCtx) onSample(sample types.Sample) { - manager.listenersMu.Lock() - // save to metrics length := float64(sample.Length) manager.totalBytes.Add(length) manager.saveSampleBitrate(sample.Timestamp, length) // if is not delta unit -> it can be decoded independently -> it is a keyframe - if manager.waitForKf && !sample.DeltaUnit && len(manager.listenersKf) > 0 { + if manager.waitForKf && !sample.DeltaUnit { + manager.listenersMu.Lock() // if current sample is a keyframe, move listeners from // keyframe lobby to actual listeners map and clear lobby for k, v := range manager.listenersKf { manager.listeners[k] = v } - manager.listenersKf = make(map[uintptr]types.SampleListener) + if len(manager.listenersKf) > 0 { + manager.listenersKf = make(map[*streamSubscription]types.SampleConsumer) + manager.updateListenersSnapshotLocked() + } + manager.listenersMu.Unlock() } - // copy listeners before releasing lock to avoid holding it during dispatch - listeners := make([]types.SampleListener, 0, len(manager.listeners)) - for _, l := range manager.listeners { - listeners = append(listeners, l) - } - manager.listenersMu.Unlock() - + listeners := manager.listenersSnapshot.Load().([]types.SampleConsumer) for _, l := range listeners { l.WriteSample(sample) } } -func (manager *StreamSinkManagerCtx) DestroyPipeline() { +func (manager *StreamSinkManagerCtx) updateListenersSnapshotLocked() { + listeners := make([]types.SampleConsumer, 0, len(manager.listeners)) + for _, listener := range manager.listeners { + listeners = append(listeners, listener) + } + manager.listenersSnapshot.Store(listeners) +} + +func (manager *StreamSinkManagerCtx) destroyPipeline() { manager.pipelineMu.Lock() defer manager.pipelineMu.Unlock() @@ -411,6 +485,8 @@ func (manager *StreamSinkManagerCtx) DestroyPipeline() { manager.pipelinesActive.Set(0) + manager.bitrateMu.Lock() manager.brBuckets = make(map[int]float64) manager.bitrate = 0 + manager.bitrateMu.Unlock() } diff --git a/server/internal/capture/streamsink_test.go b/server/internal/capture/streamsink_test.go new file mode 100644 index 00000000..98f33e72 --- /dev/null +++ b/server/internal/capture/streamsink_test.go @@ -0,0 +1,232 @@ +package capture + +import ( + "errors" + "fmt" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/m1k1o/neko/server/pkg/types" + "github.com/m1k1o/neko/server/pkg/types/codec" +) + +var testStreamID atomic.Uint64 + +type fakeSinkPipeline struct { + samples chan types.Sample + onDestroy func() +} + +func (p *fakeSinkPipeline) Sample() chan types.Sample { return p.samples } +func (p *fakeSinkPipeline) AttachAppsink(string) {} +func (p *fakeSinkPipeline) Play() {} +func (p *fakeSinkPipeline) EmitVideoKeyframe() bool { return true } +func (p *fakeSinkPipeline) Destroy() { + p.onDestroy() + close(p.samples) +} + +type testSink struct { + *StreamSinkManagerCtx + created, destroyed int +} + +func newTestSink(t *testing.T, c codec.RTPCodec) *testSink { + t.Helper() + id := fmt.Sprintf("test-%d", testStreamID.Add(1)) + s := &testSink{StreamSinkManagerCtx: streamSinkNew(c, func() (string, error) { return id, nil }, id)} + s.pipelineFactory = func(string) (sinkPipeline, error) { + s.created++ + return &fakeSinkPipeline{make(chan types.Sample), func() { s.destroyed++ }}, nil + } + t.Cleanup(s.shutdown) + return s +} + +func (s *testSink) check(t *testing.T, created, destroyed int, started bool) { + t.Helper() + if s.created != created || s.destroyed != destroyed || s.started() != started { + t.Fatalf("lifecycle = (%d, %d, %v), want (%d, %d, %v)", s.created, s.destroyed, s.started(), created, destroyed, started) + } +} + +func checkError(t *testing.T, err error) { + t.Helper() + if err != nil { + t.Fatal(err) + } +} + +func subscribe(t *testing.T, s types.EncodedStream, consumer types.SampleConsumer) types.StreamSubscription { + t.Helper() + sub, err := s.Subscribe(consumer) + checkError(t, err) + t.Cleanup(func() { checkError(t, sub.Close()) }) + return sub +} + +type sampleCounter struct{ atomic.Int64 } + +func (c *sampleCounter) WriteSample(types.Sample) { c.Add(1) } + +func TestSubscriptionLifecycle(t *testing.T) { + s := newTestSink(t, codec.Opus()) + if _, err := s.Subscribe(nil); err == nil { + t.Fatal("Subscribe(nil) succeeded") + } + s.check(t, 0, 0, false) + + a := subscribe(t, s, &sampleCounter{}) + s.check(t, 1, 0, true) + b := subscribe(t, s, &sampleCounter{}) + s.check(t, 1, 0, true) + checkError(t, a.Close()) + s.check(t, 1, 0, true) + checkError(t, b.Close()) + s.check(t, 1, 1, false) + checkError(t, b.Close()) + s.check(t, 1, 1, false) +} + +func TestSampleDispatch(t *testing.T) { + for _, c := range []codec.RTPCodec{codec.Opus(), codec.VP8()} { + t.Run(c.Name, func(t *testing.T) { + s := newTestSink(t, c) + counter := &sampleCounter{} + subscribe(t, s, counter) + for i, delta := range []bool{true, false, true} { + s.onSample(types.Sample{DeltaUnit: delta}) + want := int64(i + 1) + if c.IsVideo() { + want-- + } + if got := counter.Load(); got != want { + t.Fatalf("sample %d: deliveries = %d, want %d", i, got, want) + } + } + sample := types.Sample{Timestamp: time.Unix(0, 0), DeltaUnit: true} + if allocs := testing.AllocsPerRun(100, func() { s.onSample(sample) }); allocs != 0 { + t.Fatalf("allocations per dispatch = %f, want 0", allocs) + } + }) + } +} + +func TestSubscriptionSwitch(t *testing.T) { + for _, fail := range []bool{false, true} { + t.Run(fmt.Sprintf("startup-fails=%v", fail), func(t *testing.T) { + source, target := newTestSink(t, codec.VP8()), newTestSink(t, codec.VP8()) + factory := target.pipelineFactory + target.pipelineFactory = func(src string) (sinkPipeline, error) { + if source.destroyed != 0 { + t.Error("source stopped before target startup") + } + if fail { + return nil, errors.New("target startup failed") + } + return factory(src) + } + counter := &sampleCounter{} + sub := subscribe(t, source, counter) + err := sub.Switch(target.StreamSinkManagerCtx) + if fail { + if err == nil || sub.Stream() != source.StreamSinkManagerCtx { + t.Fatal("failed switch did not preserve the source subscription") + } + source.check(t, 1, 0, true) + target.check(t, 0, 0, false) + source.onSample(types.Sample{}) + } else { + checkError(t, err) + if sub.Stream() != target.StreamSinkManagerCtx { + t.Fatal("successful switch did not select the target") + } + source.check(t, 1, 1, false) + target.check(t, 1, 0, true) + target.onSample(types.Sample{DeltaUnit: true}) + if counter.Load() != 0 { + t.Fatal("switched consumer received a delta before a keyframe") + } + target.onSample(types.Sample{}) + } + if counter.Load() != 1 { + t.Fatal("consumer did not receive the keyframe") + } + }) + } +} + +func TestPipelineRecreation(t *testing.T) { + for _, newViewer := range []bool{false, true} { + t.Run(fmt.Sprintf("new-viewer=%v", newViewer), func(t *testing.T) { + s := newTestSink(t, codec.VP8()) + selector := streamSelectorNew(codec.VP8(), map[string]*StreamSinkManagerCtx{s.ID(): s.StreamSinkManagerCtx}, []string{s.ID()}) + if newViewer { + selector.destroyPipelines() + } + counter := &sampleCounter{} + sub := subscribe(t, s, counter) + if !newViewer { + selector.destroyPipelines() + } + if sub.Stream() != s.StreamSinkManagerCtx || !s.started() { + t.Fatal("resize lost the logical subscription") + } + checkError(t, selector.recreatePipelines()) + if newViewer { + s.check(t, 1, 0, true) + } else { + s.check(t, 2, 1, true) + } + s.onSample(types.Sample{}) + if counter.Load() != 1 { + t.Fatal("subscription stopped delivering after recreation") + } + }) + } +} + +func TestConcurrentBitrateUpdates(t *testing.T) { + s := newTestSink(t, codec.Opus()) + s.saveSampleBitrate(time.Unix(2, 0), 1) + var wg sync.WaitGroup + for range 2 { + wg.Go(func() { + for range 1000 { + s.saveSampleBitrate(time.Unix(3, 0), 1) + } + }) + } + wg.Wait() + s.saveSampleBitrate(time.Unix(4, 0), 0) + if got := s.Bitrate(); got != 2000 { + t.Fatalf("bitrate = %d, want 2000", got) + } +} + +func TestBitrateDuringPipelineRestart(t *testing.T) { + s := newTestSink(t, codec.Opus()) + done := make(chan struct{}) + var wg sync.WaitGroup + wg.Go(func() { + for sec := int64(0); ; sec++ { + select { + case <-done: + return + default: + s.saveSampleBitrate(time.Unix(sec, 0), 1) + } + } + }) + defer func() { + close(done) + wg.Wait() + }() + for range 20 { + checkError(t, s.createPipeline()) + _ = s.Bitrate() + s.destroyPipeline() + } +} diff --git a/server/internal/webrtc/peer.go b/server/internal/webrtc/peer.go index 5bf00a10..1bc502ce 100644 --- a/server/internal/webrtc/peer.go +++ b/server/internal/webrtc/peer.go @@ -28,8 +28,8 @@ type WebRTCPeerCtx struct { estimator cc.BandwidthEstimator estimateTrend *utils.TrendDetector // stream selectors - video types.StreamSelectorManager - audio types.StreamSinkManager + video types.EncodedStreamSelector + audio types.EncodedStream // tracks & channels audioTrack *Track videoTrack *Track diff --git a/server/internal/webrtc/track.go b/server/internal/webrtc/track.go index d5a37d70..39b6bcab 100644 --- a/server/internal/webrtc/track.go +++ b/server/internal/webrtc/track.go @@ -20,10 +20,13 @@ type Track struct { rtcpCh chan []rtcp.Packet sample chan types.Sample + done chan struct{} - paused bool - stream types.StreamSinkManager - streamMu sync.Mutex + paused bool + stream types.EncodedStream + subscription types.StreamSubscription + streamMu sync.Mutex + shutdownOnce sync.Once } type trackOption func(*Track) @@ -46,6 +49,7 @@ func NewTrack(logger zerolog.Logger, codec codec.RTPCodec, connection *webrtc.Pe track: track, rtcpCh: nil, sample: make(chan types.Sample, 2), + done: make(chan struct{}), } for _, opt := range opts { @@ -64,8 +68,10 @@ func NewTrack(logger zerolog.Logger, codec codec.RTPCodec, connection *webrtc.Pe } func (t *Track) Shutdown() { - t.RemoveStream() - close(t.sample) + t.shutdownOnce.Do(func() { + t.RemoveStream() + close(t.done) + }) } func (t *Track) rtcpReader(sender *webrtc.RTPSender) { @@ -91,26 +97,27 @@ func (t *Track) rtcpReader(sender *webrtc.RTPSender) { func (t *Track) sampleReader() { for { - sample, ok := <-t.sample - if !ok { + select { + case <-t.done: t.logger.Debug().Msg("track sample reader closed") return - } + case sample := <-t.sample: + err := t.track.WriteSample(media.Sample{ + Data: sample.Data, + Duration: sample.Duration, + Timestamp: sample.Timestamp, + }) - err := t.track.WriteSample(media.Sample{ - Data: sample.Data, - Duration: sample.Duration, - Timestamp: sample.Timestamp, - }) - - if err != nil && !errors.Is(err, io.ErrClosedPipe) { - t.logger.Warn().Err(err).Msg("failed to write sample to track") + if err != nil && !errors.Is(err, io.ErrClosedPipe) { + t.logger.Warn().Err(err).Msg("failed to write sample to track") + } } } } func (t *Track) WriteSample(sample types.Sample) { select { + case <-t.done: case t.sample <- sample: default: t.logger.Trace().Msg("dropping sample: track channel full") @@ -119,7 +126,7 @@ func (t *Track) WriteSample(sample types.Sample) { // --- stream --- -func (t *Track) SetStream(stream types.StreamSinkManager) (bool, error) { +func (t *Track) SetStream(stream types.EncodedStream) (bool, error) { t.streamMu.Lock() defer t.streamMu.Unlock() @@ -134,14 +141,16 @@ func (t *Track) SetStream(stream types.StreamSinkManager) (bool, error) { return true, nil } - var err error - if t.stream != nil { - err = t.stream.MoveListenerTo(t, stream) + if t.subscription != nil { + if err := t.subscription.Switch(stream); err != nil { + return false, err + } } else { - err = stream.AddListener(t) - } - if err != nil { - return false, err + subscription, err := stream.Subscribe(t) + if err != nil { + return false, err + } + t.subscription = subscription } t.stream = stream @@ -152,21 +161,21 @@ func (t *Track) RemoveStream() { t.streamMu.Lock() defer t.streamMu.Unlock() - // if there is no stream, or paused we don't need to remove the listener - if t.stream == nil || t.paused { - t.stream = nil + if t.stream == nil { return } - err := t.stream.RemoveListener(t) - if err != nil { - t.logger.Warn().Err(err).Msg("failed to remove listener from stream") + if t.subscription != nil { + if err := t.subscription.Close(); err != nil { + t.logger.Warn().Err(err).Msg("failed to close stream subscription") + } + t.subscription = nil } t.stream = nil } -func (t *Track) Stream() (types.StreamSinkManager, bool) { +func (t *Track) Stream() (types.EncodedStream, bool) { t.streamMu.Lock() defer t.streamMu.Unlock() @@ -185,15 +194,21 @@ func (t *Track) SetPaused(paused bool) { return } - var err error if paused { - err = t.stream.RemoveListener(t) + if t.subscription != nil { + if err := t.subscription.Close(); err != nil { + t.logger.Warn().Err(err).Msg("failed to pause stream subscription") + return + } + t.subscription = nil + } } else { - err = t.stream.AddListener(t) - } - if err != nil { - t.logger.Warn().Err(err).Msg("failed to change listener state") - return + subscription, err := t.stream.Subscribe(t) + if err != nil { + t.logger.Warn().Err(err).Msg("failed to resume stream subscription") + return + } + t.subscription = subscription } t.paused = paused diff --git a/server/pkg/gst/gst.c b/server/pkg/gst/gst.c index f1aa006f..0d1166a4 100644 --- a/server/pkg/gst/gst.c +++ b/server/pkg/gst/gst.c @@ -87,18 +87,19 @@ static GstFlowReturn gstreamer_send_new_sample_handler(GstElement *object, gpoin GstPipelineCtx *ctx = (GstPipelineCtx *)user_data; GstSample *sample = NULL; GstBuffer *buffer = NULL; - gpointer copy = NULL; - gsize copy_size = 0; + GstMapInfo map = {0}; g_signal_emit_by_name(object, "pull-sample", &sample); if (sample) { buffer = gst_sample_get_buffer(sample); - if (buffer) { - gst_buffer_extract_dup(buffer, 0, gst_buffer_get_size(buffer), ©, ©_size); - goHandlePipelineBuffer(ctx->pipelineId, copy, copy_size, + if (buffer && gst_buffer_map(buffer, &map, GST_MAP_READ)) { + goHandlePipelineBuffer(ctx->pipelineId, map.data, (int)map.size, + GST_BUFFER_PTS(buffer), + GST_BUFFER_DTS(buffer), GST_BUFFER_DURATION(buffer), GST_BUFFER_FLAG_IS_SET(buffer, GST_BUFFER_FLAG_DELTA_UNIT) ); + gst_buffer_unmap(buffer, &map); } gst_sample_unref(sample); } diff --git a/server/pkg/gst/gst.go b/server/pkg/gst/gst.go index 1ce2016c..25125e5f 100644 --- a/server/pkg/gst/gst.go +++ b/server/pkg/gst/gst.go @@ -211,9 +211,7 @@ func CheckElement(element string) error { } //export goHandlePipelineBuffer -func goHandlePipelineBuffer(pipelineID C.int, buf C.gpointer, bufLen C.int, duration C.guint64, deltaUnit C.gboolean) { - defer C.g_free(buf) - +func goHandlePipelineBuffer(pipelineID C.int, buf C.gpointer, bufLen C.int, pts C.guint64, dts C.guint64, duration C.guint64, deltaUnit C.gboolean) { pipelinesLock.Lock() pipeline, ok := pipelines[int(pipelineID)] pipelinesLock.Unlock() @@ -223,6 +221,8 @@ func goHandlePipelineBuffer(pipelineID C.int, buf C.gpointer, bufLen C.int, dura Data: C.GoBytes(unsafe.Pointer(buf), bufLen), Length: int(bufLen), Timestamp: time.Now(), + PTS: time.Duration(pts), + DTS: time.Duration(dts), Duration: time.Duration(duration), DeltaUnit: deltaUnit == C.TRUE, } diff --git a/server/pkg/gst/gst.h b/server/pkg/gst/gst.h index bdbd0347..73349f0f 100644 --- a/server/pkg/gst/gst.h +++ b/server/pkg/gst/gst.h @@ -23,7 +23,7 @@ typedef struct GstPipelineCtx { GstElement *appsrc; } GstPipelineCtx; -extern void goHandlePipelineBuffer(int pipelineId, void *buffer, int bufferLen, guint64 duration, gboolean deltaUnit); +extern void goHandlePipelineBuffer(int pipelineId, void *buffer, int bufferLen, guint64 pts, guint64 dts, guint64 duration, gboolean deltaUnit); extern void goPipelineLog(int pipelineId, char *level, char *msg); GstPipelineCtx *gstreamer_pipeline_create(char *pipelineStr, int pipelineId, GError **error); diff --git a/server/pkg/types/capture.go b/server/pkg/types/capture.go index febae9db..49e615df 100644 --- a/server/pkg/types/capture.go +++ b/server/pkg/types/capture.go @@ -20,19 +20,37 @@ var ( type Sample struct { // timing information Timestamp time.Time - Duration time.Duration + // PTS, DTS, and Duration use -1 for GST_CLOCK_TIME_NONE. + PTS time.Duration + DTS time.Duration + Duration time.Duration // metadata DeltaUnit bool // this unit cannot be decoded independently. // buffer length Length int - // buffer with encoded media + // Data is immutable and remains valid while referenced. Consumers may retain + // it after WriteSample returns. Data []byte } -type SampleListener interface { +type SampleConsumer interface { WriteSample(Sample) } +type EncodedStream interface { + ID() string + Codec() codec.RTPCodec + Bitrate() uint64 + + Subscribe(SampleConsumer) (StreamSubscription, error) +} + +type StreamSubscription interface { + Stream() EncodedStream + Switch(EncodedStream) error + Close() error +} + type BroadcastManager interface { Start(url string) error Stop() @@ -103,27 +121,16 @@ type StreamSelector struct { Bitrate uint64 `json:"bitrate"` } -type StreamSelectorManager interface { +type EncodedStreamSelector interface { IDs() []string Codec() codec.RTPCodec - GetStream(selector StreamSelector) (StreamSinkManager, bool) + GetStream(selector StreamSelector) (EncodedStream, bool) } -type StreamSinkManager interface { - ID() string - Codec() codec.RTPCodec - Bitrate() uint64 - - AddListener(listener SampleListener) error - RemoveListener(listener SampleListener) error - MoveListenerTo(listener SampleListener, targetStream StreamSinkManager) error - - ListenersCount() int - Started() bool - - CreatePipeline() error - DestroyPipeline() +type EncodedMediaSource interface { + Audio() EncodedStream + Video() EncodedStreamSelector } type StreamSrcManager interface { @@ -137,13 +144,13 @@ type StreamSrcManager interface { } type CaptureManager interface { + EncodedMediaSource + Start() Shutdown() error Broadcast() BroadcastManager Screencast() ScreencastManager - Audio() StreamSinkManager - Video() StreamSelectorManager Webcam() StreamSrcManager Microphone() StreamSrcManager