feat: create capture subscription api to support more transports (#699)

* 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ý <sedivy.miro@gmail.com>
This commit is contained in:
a
2026-09-24 03:58:39 -05:00
committed by GitHub
parent 6bc8350213
commit f4d632df2c
10 changed files with 505 additions and 175 deletions
+3 -3
View File
@@ -33,7 +33,7 @@ type CaptureManagerCtx struct {
func New(desktop types.DesktopManager, config *config.Capture) *CaptureManagerCtx { func New(desktop types.DesktopManager, config *config.Capture) *CaptureManagerCtx {
logger := log.With().Str("module", "capture").Logger() logger := log.With().Str("module", "capture").Logger()
videos := map[string]types.StreamSinkManager{} videos := map[string]*StreamSinkManagerCtx{}
for video_id, cnf := range config.VideoPipelines { for video_id, cnf := range config.VideoPipelines {
pipelineConf := cnf pipelineConf := cnf
@@ -257,11 +257,11 @@ func (manager *CaptureManagerCtx) Screencast() types.ScreencastManager {
return manager.screencast return manager.screencast
} }
func (manager *CaptureManagerCtx) Audio() types.StreamSinkManager { func (manager *CaptureManagerCtx) Audio() types.EncodedStream {
return manager.audio return manager.audio
} }
func (manager *CaptureManagerCtx) Video() types.StreamSelectorManager { func (manager *CaptureManagerCtx) Video() types.EncodedStreamSelector {
return manager.video return manager.video
} }
+14 -15
View File
@@ -1,7 +1,6 @@
package capture package capture
import ( import (
"errors"
"sort" "sort"
"github.com/rs/zerolog" "github.com/rs/zerolog"
@@ -14,11 +13,11 @@ import (
type StreamSelectorManagerCtx struct { type StreamSelectorManagerCtx struct {
logger zerolog.Logger logger zerolog.Logger
codec codec.RTPCodec codec codec.RTPCodec
streams map[string]types.StreamSinkManager streams map[string]*StreamSinkManagerCtx
streamIDs []string 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(). logger := log.With().
Str("module", "capture"). Str("module", "capture").
Str("submodule", "stream-selector"). Str("submodule", "stream-selector").
@@ -40,17 +39,17 @@ func (manager *StreamSelectorManagerCtx) shutdown() {
func (manager *StreamSelectorManagerCtx) destroyPipelines() { func (manager *StreamSelectorManagerCtx) destroyPipelines() {
for _, stream := range manager.streams { for _, stream := range manager.streams {
if stream.Started() { if stream.started() {
stream.DestroyPipeline() stream.destroyPipeline()
} }
} }
} }
func (manager *StreamSelectorManagerCtx) recreatePipelines() error { func (manager *StreamSelectorManagerCtx) recreatePipelines() error {
for _, stream := range manager.streams { for _, stream := range manager.streams {
if stream.Started() { if stream.started() {
err := stream.CreatePipeline() err := stream.createPipeline()
if err != nil && !errors.Is(err, types.ErrCapturePipelineAlreadyExists) { if err != nil {
return err return err
} }
} }
@@ -66,12 +65,12 @@ func (manager *StreamSelectorManagerCtx) Codec() codec.RTPCodec {
return manager.codec 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 // select stream by ID
if selector.ID != "" { if selector.ID != "" {
// select lower stream // select lower stream
if selector.Type == types.StreamSelectorTypeLower { if selector.Type == types.StreamSelectorTypeLower {
var lastStream types.StreamSinkManager var lastStream types.EncodedStream
for i := len(manager.streamIDs) - 1; i >= 0; i-- { for i := len(manager.streamIDs) - 1; i >= 0; i-- {
streamID := manager.streamIDs[i] streamID := manager.streamIDs[i]
if streamID == selector.ID { if streamID == selector.ID {
@@ -88,7 +87,7 @@ func (manager *StreamSelectorManagerCtx) GetStream(selector types.StreamSelector
// select higher stream // select higher stream
if selector.Type == types.StreamSelectorTypeHigher { if selector.Type == types.StreamSelectorTypeHigher {
var lastStream types.StreamSinkManager var lastStream types.EncodedStream
for _, streamID := range manager.streamIDs { for _, streamID := range manager.streamIDs {
if streamID == selector.ID { if streamID == selector.ID {
return lastStream, lastStream != nil return lastStream, lastStream != nil
@@ -121,7 +120,7 @@ func (manager *StreamSelectorManagerCtx) GetStream(selector types.StreamSelector
streamID := manager.streamIDs[i] streamID := manager.streamIDs[i]
stream := manager.streams[streamID] stream := manager.streams[streamID]
// if stream should be considered in calculation // 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 { if considered && stream.Bitrate() < selector.Bitrate {
return stream, true return stream, true
} }
@@ -136,7 +135,7 @@ func (manager *StreamSelectorManagerCtx) GetStream(selector types.StreamSelector
for _, streamID := range manager.streamIDs { for _, streamID := range manager.streamIDs {
stream := manager.streams[streamID] stream := manager.streams[streamID]
// if stream should be considered in calculation // 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 { if considered && stream.Bitrate() > selector.Bitrate {
return stream, true 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. // 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 { type streamDiff struct {
id string id string
bitrateDiff int bitrateDiff int
@@ -181,7 +180,7 @@ func (manager *StreamSelectorManagerCtx) nearestBitrate(bitrate uint64) types.St
for _, stream := range manager.streams { for _, stream := range manager.streams {
// if stream should be considered in calculation // if stream should be considered in calculation
considered := stream.Bitrate() != 0 && stream.Started() considered := stream.Bitrate() != 0 && stream.started()
if !considered { if !considered {
continue continue
} }
+163 -87
View File
@@ -2,8 +2,8 @@ package capture
import ( import (
"errors" "errors"
"reflect"
"sync" "sync"
"sync/atomic"
"time" "time"
"github.com/prometheus/client_golang/prometheus" "github.com/prometheus/client_golang/prometheus"
@@ -17,30 +17,44 @@ import (
"github.com/m1k1o/neko/server/pkg/types/codec" "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 { 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 // Serializes subscription add, remove, and move operations.
waitForKf bool 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 bitrate uint64
brBuckets map[int]float64 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 // metrics
currentListeners prometheus.Gauge currentListeners prometheus.Gauge
totalBytes prometheus.Counter totalBytes prometheus.Counter
@@ -66,9 +80,12 @@ func streamSinkNew(codec codec.RTPCodec, pipelineFn func() (string, error), id s
logger: logger, logger: logger,
codec: codec, codec: codec,
pipelineFn: pipelineFn, pipelineFn: pipelineFn,
pipelineFactory: func(src string) (sinkPipeline, error) {
return gst.CreatePipeline(src)
},
listeners: map[uintptr]types.SampleListener{}, listeners: map[*streamSubscription]types.SampleConsumer{},
listenersKf: map[uintptr]types.SampleListener{}, listenersKf: map[*streamSubscription]types.SampleConsumer{},
// metrics // metrics
currentListeners: promauto.NewGauge(prometheus.GaugeOpts{ 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 return manager
} }
@@ -132,9 +150,10 @@ func (manager *StreamSinkManagerCtx) shutdown() {
for key := range manager.listenersKf { for key := range manager.listenersKf {
delete(manager.listenersKf, key) delete(manager.listenersKf, key)
} }
manager.updateListenersSnapshotLocked()
manager.listenersMu.Unlock() manager.listenersMu.Unlock()
manager.DestroyPipeline() manager.destroyPipeline()
manager.wg.Wait() manager.wg.Wait()
} }
@@ -143,6 +162,9 @@ func (manager *StreamSinkManagerCtx) ID() string {
} }
func (manager *StreamSinkManagerCtx) Bitrate() uint64 { func (manager *StreamSinkManagerCtx) Bitrate() uint64 {
manager.bitrateMu.Lock()
defer manager.bitrateMu.Unlock()
return manager.bitrate return manager.bitrate
} }
@@ -150,10 +172,76 @@ func (manager *StreamSinkManagerCtx) Codec() codec.RTPCodec {
return manager.codec 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 { func (manager *StreamSinkManagerCtx) start() error {
if len(manager.listeners)+len(manager.listenersKf) == 0 { if len(manager.listeners)+len(manager.listenersKf) == 0 {
err := manager.CreatePipeline() err := manager.createPipeline()
if err != nil && !errors.Is(err, types.ErrCapturePipelineAlreadyExists) { if err != nil {
return err return err
} }
@@ -165,13 +253,12 @@ func (manager *StreamSinkManagerCtx) start() error {
func (manager *StreamSinkManagerCtx) stop() { func (manager *StreamSinkManagerCtx) stop() {
if len(manager.listeners)+len(manager.listenersKf) == 0 { if len(manager.listeners)+len(manager.listenersKf) == 0 {
manager.DestroyPipeline() manager.destroyPipeline()
manager.logger.Info().Msgf("last listener, stopping") manager.logger.Info().Msgf("last listener, stopping")
} }
} }
func (manager *StreamSinkManagerCtx) addListener(listener types.SampleListener) { func (manager *StreamSinkManagerCtx) addSubscriptionToMaps(subscription *streamSubscription) {
ptr := reflect.ValueOf(listener).Pointer()
emitKeyframe := false emitKeyframe := false
manager.listenersMu.Lock() 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 // if this is the first listener, we need to emit a keyframe
emitKeyframe = len(manager.listenersKf) == 0 emitKeyframe = len(manager.listenersKf) == 0
// if we're waiting for a keyframe, add it to the keyframe lobby // if we're waiting for a keyframe, add it to the keyframe lobby
manager.listenersKf[ptr] = listener manager.listenersKf[subscription] = subscription.consumer
} else { } else {
// otherwise, add it as a regular listener // otherwise, add it as a regular listener
manager.listeners[ptr] = listener manager.listeners[subscription] = subscription.consumer
manager.updateListenersSnapshotLocked()
} }
manager.listenersMu.Unlock() manager.listenersMu.Unlock()
manager.logger.Debug().Interface("ptr", ptr).Msgf("adding listener") manager.logger.Debug().Interface("subscription", subscription).Msg("adding subscription")
manager.currentListeners.Set(float64(manager.ListenersCount())) manager.currentListeners.Set(float64(manager.subscriptionsCount()))
// if we will be waiting for a keyframe, emit one now // if we will be waiting for a keyframe, emit one now
if manager.pipeline != nil && emitKeyframe { if manager.pipeline != nil && emitKeyframe {
@@ -195,72 +283,50 @@ func (manager *StreamSinkManagerCtx) addListener(listener types.SampleListener)
} }
} }
func (manager *StreamSinkManagerCtx) removeListener(listener types.SampleListener) { func (manager *StreamSinkManagerCtx) removeSubscriptionFromMaps(subscription *streamSubscription) {
ptr := reflect.ValueOf(listener).Pointer()
manager.listenersMu.Lock() manager.listenersMu.Lock()
delete(manager.listeners, ptr) delete(manager.listeners, subscription)
delete(manager.listenersKf, ptr) // if it's a keyframe listener, remove it too delete(manager.listenersKf, subscription)
manager.updateListenersSnapshotLocked()
manager.listenersMu.Unlock() manager.listenersMu.Unlock()
manager.logger.Debug().Interface("ptr", ptr).Msgf("removing listener") manager.logger.Debug().Interface("subscription", subscription).Msg("removing subscription")
manager.currentListeners.Set(float64(manager.ListenersCount())) manager.currentListeners.Set(float64(manager.subscriptionsCount()))
} }
func (manager *StreamSinkManagerCtx) AddListener(listener types.SampleListener) error { func (manager *StreamSinkManagerCtx) addSubscription(subscription *streamSubscription) error {
manager.mu.Lock() manager.mu.Lock()
defer manager.mu.Unlock() defer manager.mu.Unlock()
if listener == nil {
return errors.New("listener cannot be nil")
}
// start if stopped // start if stopped
if err := manager.start(); err != nil { if err := manager.start(); err != nil {
return err return err
} }
// add listener // add listener
manager.addListener(listener) manager.addSubscriptionToMaps(subscription)
return nil return nil
} }
func (manager *StreamSinkManagerCtx) RemoveListener(listener types.SampleListener) error { func (manager *StreamSinkManagerCtx) removeSubscription(subscription *streamSubscription) {
manager.mu.Lock() manager.mu.Lock()
defer manager.mu.Unlock() defer manager.mu.Unlock()
if listener == nil { manager.removeSubscriptionFromMaps(subscription)
return errors.New("listener cannot be nil")
}
// remove listener
manager.removeListener(listener)
// stop if started // stop if started
manager.stop() manager.stop()
return nil
} }
// moving listeners between streams ensures, that target pipeline is running // moveSubscriptionTo starts the target before moving delivery and stopping the source.
// before listener is added, and stops source pipeline if there are 0 listeners func (manager *StreamSinkManagerCtx) moveSubscriptionTo(subscription *streamSubscription, targetStream *StreamSinkManagerCtx) error {
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")
}
// we need to acquire both mutextes, from source stream and from target stream // 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 // in order to do that safely (without possibility of deadlock) we need third
// global mutex, that ensures atomic locking // global mutex, that ensures atomic locking
// lock global mutex // lock global mutex
moveSinkListenerMu.Lock() moveSinkSubscriptionMu.Lock()
// lock source stream // lock source stream
manager.mu.Lock() manager.mu.Lock()
@@ -271,7 +337,7 @@ func (manager *StreamSinkManagerCtx) MoveListenerTo(listener types.SampleListene
defer targetStream.mu.Unlock() defer targetStream.mu.Unlock()
// unlock global mutex // unlock global mutex
moveSinkListenerMu.Unlock() moveSinkSubscriptionMu.Unlock()
// start if stopped // start if stopped
if err := targetStream.start(); err != nil { if err := targetStream.start(); err != nil {
@@ -279,8 +345,8 @@ func (manager *StreamSinkManagerCtx) MoveListenerTo(listener types.SampleListene
} }
// swap listeners // swap listeners
manager.removeListener(listener) manager.removeSubscriptionFromMaps(subscription)
targetStream.addListener(listener) targetStream.addSubscriptionToMaps(subscription)
// stop if started // stop if started
manager.stop() manager.stop()
@@ -288,23 +354,23 @@ func (manager *StreamSinkManagerCtx) MoveListenerTo(listener types.SampleListene
return nil return nil
} }
func (manager *StreamSinkManagerCtx) ListenersCount() int { func (manager *StreamSinkManagerCtx) subscriptionsCount() int {
manager.listenersMu.Lock() manager.listenersMu.Lock()
defer manager.listenersMu.Unlock() defer manager.listenersMu.Unlock()
return len(manager.listeners) + len(manager.listenersKf) return len(manager.listeners) + len(manager.listenersKf)
} }
func (manager *StreamSinkManagerCtx) Started() bool { func (manager *StreamSinkManagerCtx) started() bool {
return manager.ListenersCount() > 0 return manager.subscriptionsCount() > 0
} }
func (manager *StreamSinkManagerCtx) CreatePipeline() error { func (manager *StreamSinkManagerCtx) createPipeline() error {
manager.pipelineMu.Lock() manager.pipelineMu.Lock()
defer manager.pipelineMu.Unlock() defer manager.pipelineMu.Unlock()
if manager.pipeline != nil { if manager.pipeline != nil {
return types.ErrCapturePipelineAlreadyExists return nil
} }
pipelineStr, err := manager.pipelineFn() pipelineStr, err := manager.pipelineFn()
@@ -317,7 +383,7 @@ func (manager *StreamSinkManagerCtx) CreatePipeline() error {
Str("src", pipelineStr). Str("src", pipelineStr).
Msgf("creating pipeline") Msgf("creating pipeline")
manager.pipeline, err = gst.CreatePipeline(pipelineStr) manager.pipeline, err = manager.pipelineFactory(pipelineStr)
if err != nil { if err != nil {
return err return err
} }
@@ -347,6 +413,9 @@ func (manager *StreamSinkManagerCtx) CreatePipeline() error {
} }
func (manager *StreamSinkManagerCtx) saveSampleBitrate(timestamp time.Time, delta float64) { func (manager *StreamSinkManagerCtx) saveSampleBitrate(timestamp time.Time, delta float64) {
manager.bitrateMu.Lock()
defer manager.bitrateMu.Unlock()
// get unix timestamp in seconds // get unix timestamp in seconds
sec := timestamp.Unix() sec := timestamp.Unix()
// last bucket is timestamp rounded to 3 seconds - 1 second // 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) next := int((sec + 1) % 3)
if manager.brBuckets[next] != 0 { if manager.brBuckets[next] != 0 {
// update bitrate, TODO: atomic? // update bitrate
manager.bitrate = uint64(manager.brBuckets[last]) manager.bitrate = uint64(manager.brBuckets[last])
// empty next bucket // empty next bucket
manager.brBuckets[next] = 0 manager.brBuckets[next] = 0
@@ -368,36 +437,41 @@ func (manager *StreamSinkManagerCtx) saveSampleBitrate(timestamp time.Time, delt
} }
func (manager *StreamSinkManagerCtx) onSample(sample types.Sample) { func (manager *StreamSinkManagerCtx) onSample(sample types.Sample) {
manager.listenersMu.Lock()
// save to metrics // save to metrics
length := float64(sample.Length) length := float64(sample.Length)
manager.totalBytes.Add(length) manager.totalBytes.Add(length)
manager.saveSampleBitrate(sample.Timestamp, length) manager.saveSampleBitrate(sample.Timestamp, length)
// if is not delta unit -> it can be decoded independently -> it is a keyframe // 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 // if current sample is a keyframe, move listeners from
// keyframe lobby to actual listeners map and clear lobby // keyframe lobby to actual listeners map and clear lobby
for k, v := range manager.listenersKf { for k, v := range manager.listenersKf {
manager.listeners[k] = v 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 := manager.listenersSnapshot.Load().([]types.SampleConsumer)
listeners := make([]types.SampleListener, 0, len(manager.listeners))
for _, l := range manager.listeners {
listeners = append(listeners, l)
}
manager.listenersMu.Unlock()
for _, l := range listeners { for _, l := range listeners {
l.WriteSample(sample) 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() manager.pipelineMu.Lock()
defer manager.pipelineMu.Unlock() defer manager.pipelineMu.Unlock()
@@ -411,6 +485,8 @@ func (manager *StreamSinkManagerCtx) DestroyPipeline() {
manager.pipelinesActive.Set(0) manager.pipelinesActive.Set(0)
manager.bitrateMu.Lock()
manager.brBuckets = make(map[int]float64) manager.brBuckets = make(map[int]float64)
manager.bitrate = 0 manager.bitrate = 0
manager.bitrateMu.Unlock()
} }
+232
View File
@@ -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()
}
}
+2 -2
View File
@@ -28,8 +28,8 @@ type WebRTCPeerCtx struct {
estimator cc.BandwidthEstimator estimator cc.BandwidthEstimator
estimateTrend *utils.TrendDetector estimateTrend *utils.TrendDetector
// stream selectors // stream selectors
video types.StreamSelectorManager video types.EncodedStreamSelector
audio types.StreamSinkManager audio types.EncodedStream
// tracks & channels // tracks & channels
audioTrack *Track audioTrack *Track
videoTrack *Track videoTrack *Track
+53 -38
View File
@@ -20,10 +20,13 @@ type Track struct {
rtcpCh chan []rtcp.Packet rtcpCh chan []rtcp.Packet
sample chan types.Sample sample chan types.Sample
done chan struct{}
paused bool paused bool
stream types.StreamSinkManager stream types.EncodedStream
streamMu sync.Mutex subscription types.StreamSubscription
streamMu sync.Mutex
shutdownOnce sync.Once
} }
type trackOption func(*Track) type trackOption func(*Track)
@@ -46,6 +49,7 @@ func NewTrack(logger zerolog.Logger, codec codec.RTPCodec, connection *webrtc.Pe
track: track, track: track,
rtcpCh: nil, rtcpCh: nil,
sample: make(chan types.Sample, 2), sample: make(chan types.Sample, 2),
done: make(chan struct{}),
} }
for _, opt := range opts { for _, opt := range opts {
@@ -64,8 +68,10 @@ func NewTrack(logger zerolog.Logger, codec codec.RTPCodec, connection *webrtc.Pe
} }
func (t *Track) Shutdown() { func (t *Track) Shutdown() {
t.RemoveStream() t.shutdownOnce.Do(func() {
close(t.sample) t.RemoveStream()
close(t.done)
})
} }
func (t *Track) rtcpReader(sender *webrtc.RTPSender) { func (t *Track) rtcpReader(sender *webrtc.RTPSender) {
@@ -91,26 +97,27 @@ func (t *Track) rtcpReader(sender *webrtc.RTPSender) {
func (t *Track) sampleReader() { func (t *Track) sampleReader() {
for { for {
sample, ok := <-t.sample select {
if !ok { case <-t.done:
t.logger.Debug().Msg("track sample reader closed") t.logger.Debug().Msg("track sample reader closed")
return 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{ if err != nil && !errors.Is(err, io.ErrClosedPipe) {
Data: sample.Data, t.logger.Warn().Err(err).Msg("failed to write sample to track")
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")
} }
} }
} }
func (t *Track) WriteSample(sample types.Sample) { func (t *Track) WriteSample(sample types.Sample) {
select { select {
case <-t.done:
case t.sample <- sample: case t.sample <- sample:
default: default:
t.logger.Trace().Msg("dropping sample: track channel full") t.logger.Trace().Msg("dropping sample: track channel full")
@@ -119,7 +126,7 @@ func (t *Track) WriteSample(sample types.Sample) {
// --- stream --- // --- stream ---
func (t *Track) SetStream(stream types.StreamSinkManager) (bool, error) { func (t *Track) SetStream(stream types.EncodedStream) (bool, error) {
t.streamMu.Lock() t.streamMu.Lock()
defer t.streamMu.Unlock() defer t.streamMu.Unlock()
@@ -134,14 +141,16 @@ func (t *Track) SetStream(stream types.StreamSinkManager) (bool, error) {
return true, nil return true, nil
} }
var err error if t.subscription != nil {
if t.stream != nil { if err := t.subscription.Switch(stream); err != nil {
err = t.stream.MoveListenerTo(t, stream) return false, err
}
} else { } else {
err = stream.AddListener(t) subscription, err := stream.Subscribe(t)
} if err != nil {
if err != nil { return false, err
return false, err }
t.subscription = subscription
} }
t.stream = stream t.stream = stream
@@ -152,21 +161,21 @@ func (t *Track) RemoveStream() {
t.streamMu.Lock() t.streamMu.Lock()
defer t.streamMu.Unlock() defer t.streamMu.Unlock()
// if there is no stream, or paused we don't need to remove the listener if t.stream == nil {
if t.stream == nil || t.paused {
t.stream = nil
return return
} }
err := t.stream.RemoveListener(t) if t.subscription != nil {
if err != nil { if err := t.subscription.Close(); err != nil {
t.logger.Warn().Err(err).Msg("failed to remove listener from stream") t.logger.Warn().Err(err).Msg("failed to close stream subscription")
}
t.subscription = nil
} }
t.stream = nil t.stream = nil
} }
func (t *Track) Stream() (types.StreamSinkManager, bool) { func (t *Track) Stream() (types.EncodedStream, bool) {
t.streamMu.Lock() t.streamMu.Lock()
defer t.streamMu.Unlock() defer t.streamMu.Unlock()
@@ -185,15 +194,21 @@ func (t *Track) SetPaused(paused bool) {
return return
} }
var err error
if paused { 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 { } else {
err = t.stream.AddListener(t) subscription, err := t.stream.Subscribe(t)
} if err != nil {
if err != nil { t.logger.Warn().Err(err).Msg("failed to resume stream subscription")
t.logger.Warn().Err(err).Msg("failed to change listener state") return
return }
t.subscription = subscription
} }
t.paused = paused t.paused = paused
+6 -5
View File
@@ -87,18 +87,19 @@ static GstFlowReturn gstreamer_send_new_sample_handler(GstElement *object, gpoin
GstPipelineCtx *ctx = (GstPipelineCtx *)user_data; GstPipelineCtx *ctx = (GstPipelineCtx *)user_data;
GstSample *sample = NULL; GstSample *sample = NULL;
GstBuffer *buffer = NULL; GstBuffer *buffer = NULL;
gpointer copy = NULL; GstMapInfo map = {0};
gsize copy_size = 0;
g_signal_emit_by_name(object, "pull-sample", &sample); g_signal_emit_by_name(object, "pull-sample", &sample);
if (sample) { if (sample) {
buffer = gst_sample_get_buffer(sample); buffer = gst_sample_get_buffer(sample);
if (buffer) { if (buffer && gst_buffer_map(buffer, &map, GST_MAP_READ)) {
gst_buffer_extract_dup(buffer, 0, gst_buffer_get_size(buffer), &copy, &copy_size); goHandlePipelineBuffer(ctx->pipelineId, map.data, (int)map.size,
goHandlePipelineBuffer(ctx->pipelineId, copy, copy_size, GST_BUFFER_PTS(buffer),
GST_BUFFER_DTS(buffer),
GST_BUFFER_DURATION(buffer), GST_BUFFER_DURATION(buffer),
GST_BUFFER_FLAG_IS_SET(buffer, GST_BUFFER_FLAG_DELTA_UNIT) GST_BUFFER_FLAG_IS_SET(buffer, GST_BUFFER_FLAG_DELTA_UNIT)
); );
gst_buffer_unmap(buffer, &map);
} }
gst_sample_unref(sample); gst_sample_unref(sample);
} }
+3 -3
View File
@@ -211,9 +211,7 @@ func CheckElement(element string) error {
} }
//export goHandlePipelineBuffer //export goHandlePipelineBuffer
func goHandlePipelineBuffer(pipelineID C.int, buf C.gpointer, bufLen C.int, duration C.guint64, deltaUnit C.gboolean) { func goHandlePipelineBuffer(pipelineID C.int, buf C.gpointer, bufLen C.int, pts C.guint64, dts C.guint64, duration C.guint64, deltaUnit C.gboolean) {
defer C.g_free(buf)
pipelinesLock.Lock() pipelinesLock.Lock()
pipeline, ok := pipelines[int(pipelineID)] pipeline, ok := pipelines[int(pipelineID)]
pipelinesLock.Unlock() 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), Data: C.GoBytes(unsafe.Pointer(buf), bufLen),
Length: int(bufLen), Length: int(bufLen),
Timestamp: time.Now(), Timestamp: time.Now(),
PTS: time.Duration(pts),
DTS: time.Duration(dts),
Duration: time.Duration(duration), Duration: time.Duration(duration),
DeltaUnit: deltaUnit == C.TRUE, DeltaUnit: deltaUnit == C.TRUE,
} }
+1 -1
View File
@@ -23,7 +23,7 @@ typedef struct GstPipelineCtx {
GstElement *appsrc; GstElement *appsrc;
} GstPipelineCtx; } 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); extern void goPipelineLog(int pipelineId, char *level, char *msg);
GstPipelineCtx *gstreamer_pipeline_create(char *pipelineStr, int pipelineId, GError **error); GstPipelineCtx *gstreamer_pipeline_create(char *pipelineStr, int pipelineId, GError **error);
+28 -21
View File
@@ -20,19 +20,37 @@ var (
type Sample struct { type Sample struct {
// timing information // timing information
Timestamp time.Time 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 // metadata
DeltaUnit bool // this unit cannot be decoded independently. DeltaUnit bool // this unit cannot be decoded independently.
// buffer length // buffer length
Length int Length int
// buffer with encoded media // Data is immutable and remains valid while referenced. Consumers may retain
// it after WriteSample returns.
Data []byte Data []byte
} }
type SampleListener interface { type SampleConsumer interface {
WriteSample(Sample) 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 { type BroadcastManager interface {
Start(url string) error Start(url string) error
Stop() Stop()
@@ -103,27 +121,16 @@ type StreamSelector struct {
Bitrate uint64 `json:"bitrate"` Bitrate uint64 `json:"bitrate"`
} }
type StreamSelectorManager interface { type EncodedStreamSelector interface {
IDs() []string IDs() []string
Codec() codec.RTPCodec Codec() codec.RTPCodec
GetStream(selector StreamSelector) (StreamSinkManager, bool) GetStream(selector StreamSelector) (EncodedStream, bool)
} }
type StreamSinkManager interface { type EncodedMediaSource interface {
ID() string Audio() EncodedStream
Codec() codec.RTPCodec Video() EncodedStreamSelector
Bitrate() uint64
AddListener(listener SampleListener) error
RemoveListener(listener SampleListener) error
MoveListenerTo(listener SampleListener, targetStream StreamSinkManager) error
ListenersCount() int
Started() bool
CreatePipeline() error
DestroyPipeline()
} }
type StreamSrcManager interface { type StreamSrcManager interface {
@@ -137,13 +144,13 @@ type StreamSrcManager interface {
} }
type CaptureManager interface { type CaptureManager interface {
EncodedMediaSource
Start() Start()
Shutdown() error Shutdown() error
Broadcast() BroadcastManager Broadcast() BroadcastManager
Screencast() ScreencastManager Screencast() ScreencastManager
Audio() StreamSinkManager
Video() StreamSelectorManager
Webcam() StreamSrcManager Webcam() StreamSrcManager
Microphone() StreamSrcManager Microphone() StreamSrcManager