neko/internal/capture/streamsink.go

415 lines
10 KiB
Go
Raw Normal View History

package capture
import (
2021-09-27 11:50:49 +13:00
"errors"
"reflect"
2021-02-15 02:40:17 +13:00
"sync"
"sync/atomic"
"time"
2022-06-15 10:23:16 +12:00
"github.com/prometheus/client_golang/prometheus"
"github.com/prometheus/client_golang/prometheus/promauto"
"github.com/rs/zerolog"
"github.com/rs/zerolog/log"
"github.com/demodesk/neko/pkg/gst"
"github.com/demodesk/neko/pkg/types"
"github.com/demodesk/neko/pkg/types/codec"
)
2021-12-02 08:30:18 +13:00
var moveSinkListenerMu = sync.Mutex{}
2021-12-02 08:30:18 +13:00
type StreamSinkManagerCtx struct {
id string
// wait for a keyframe before sending samples
waitForKf bool
bitrate uint64 // atomic
brBuckets map[int]float64
2021-09-27 11:50:49 +13:00
logger zerolog.Logger
mu sync.Mutex
wg sync.WaitGroup
2022-09-18 03:52:55 +12:00
codec codec.RTPCodec
2022-10-14 06:57:31 +13:00
pipeline gst.Pipeline
2022-09-18 03:52:55 +12:00
pipelineMu sync.Mutex
pipelineFn func() (string, error)
2021-09-27 11:50:49 +13:00
listeners map[uintptr]types.SampleListener
listenersKf map[uintptr]types.SampleListener // keyframe lobby
2021-10-02 00:46:10 +13:00
listenersMu sync.Mutex
2022-06-15 10:23:16 +12:00
// metrics
currentListeners prometheus.Gauge
totalBytes prometheus.Counter
2022-06-15 10:23:16 +12:00
pipelinesCounter prometheus.Counter
2022-06-19 11:03:16 +12:00
pipelinesActive prometheus.Gauge
}
func streamSinkNew(codec codec.RTPCodec, pipelineFn func() (string, error), id string) *StreamSinkManagerCtx {
2021-08-29 06:15:54 +12:00
logger := log.With().
Str("module", "capture").
2021-12-02 08:30:18 +13:00
Str("submodule", "stream-sink").
Str("id", id).Logger()
2021-08-29 06:15:54 +12:00
2021-12-02 08:30:18 +13:00
manager := &StreamSinkManagerCtx{
id: id,
// only wait for keyframes if the codec is video
waitForKf: codec.IsVideo(),
bitrate: 0,
brBuckets: map[int]float64{},
2022-09-18 03:52:55 +12:00
logger: logger,
codec: codec,
2022-09-18 03:52:55 +12:00
pipelineFn: pipelineFn,
listeners: map[uintptr]types.SampleListener{},
listenersKf: map[uintptr]types.SampleListener{},
2022-06-15 10:23:16 +12:00
// metrics
currentListeners: promauto.NewGauge(prometheus.GaugeOpts{
2022-06-19 10:51:17 +12:00
Name: "streamsink_listeners",
2022-06-15 10:23:16 +12:00
Namespace: "neko",
2022-06-19 10:51:17 +12:00
Subsystem: "capture",
2022-06-15 10:23:16 +12:00
Help: "Current number of listeners for a pipeline.",
ConstLabels: map[string]string{
"video_id": id,
"codec_name": codec.Name,
"codec_type": codec.Type.String(),
2022-06-15 10:23:16 +12:00
},
}),
2023-03-18 08:20:23 +13:00
totalBytes: promauto.NewCounter(prometheus.CounterOpts{
Name: "streamsink_bytes",
Namespace: "neko",
Subsystem: "capture",
Help: "Total number of bytes created by the pipeline.",
ConstLabels: map[string]string{
"video_id": id,
"codec_name": codec.Name,
"codec_type": codec.Type.String(),
},
}),
2022-06-15 10:23:16 +12:00
pipelinesCounter: promauto.NewCounter(prometheus.CounterOpts{
Name: "pipelines_total",
Namespace: "neko",
2022-06-19 10:51:17 +12:00
Subsystem: "capture",
2022-06-15 10:23:16 +12:00
Help: "Total number of created pipelines.",
ConstLabels: map[string]string{
2022-06-19 10:51:17 +12:00
"submodule": "streamsink",
"video_id": id,
"codec_name": codec.Name,
"codec_type": codec.Type.String(),
2022-06-15 10:23:16 +12:00
},
}),
2022-06-19 11:03:16 +12:00
pipelinesActive: promauto.NewGauge(prometheus.GaugeOpts{
Name: "pipelines_active",
Namespace: "neko",
Subsystem: "capture",
Help: "Total number of active pipelines.",
ConstLabels: map[string]string{
"submodule": "streamsink",
"video_id": id,
"codec_name": codec.Name,
"codec_type": codec.Type.String(),
2022-06-19 11:03:16 +12:00
},
}),
}
return manager
}
2021-12-02 08:30:18 +13:00
func (manager *StreamSinkManagerCtx) shutdown() {
2021-09-02 10:00:56 +12:00
manager.logger.Info().Msgf("shutdown")
2021-09-27 11:50:49 +13:00
manager.listenersMu.Lock()
for key := range manager.listeners {
delete(manager.listeners, key)
}
for key := range manager.listenersKf {
delete(manager.listenersKf, key)
}
2021-09-27 11:50:49 +13:00
manager.listenersMu.Unlock()
manager.DestroyPipeline()
2021-09-02 10:00:56 +12:00
manager.wg.Wait()
}
func (manager *StreamSinkManagerCtx) ID() string {
return manager.id
}
func (manager *StreamSinkManagerCtx) Bitrate() uint64 {
return atomic.LoadUint64(&manager.bitrate)
}
2021-12-02 08:30:18 +13:00
func (manager *StreamSinkManagerCtx) Codec() codec.RTPCodec {
return manager.codec
}
2021-12-02 08:30:18 +13:00
func (manager *StreamSinkManagerCtx) start() error {
if len(manager.listeners)+len(manager.listenersKf) == 0 {
err := manager.CreatePipeline()
2021-09-27 11:50:49 +13:00
if err != nil && !errors.Is(err, types.ErrCapturePipelineAlreadyExists) {
2021-10-02 00:46:10 +13:00
return err
2021-09-27 11:50:49 +13:00
}
2021-09-27 11:50:49 +13:00
manager.logger.Info().Msgf("first listener, starting")
}
2021-10-02 00:46:10 +13:00
return nil
}
2021-12-02 08:30:18 +13:00
func (manager *StreamSinkManagerCtx) stop() {
if len(manager.listeners)+len(manager.listenersKf) == 0 {
manager.DestroyPipeline()
2021-10-02 00:46:10 +13:00
manager.logger.Info().Msgf("last listener, stopping")
}
}
2021-09-27 11:50:49 +13:00
func (manager *StreamSinkManagerCtx) addListener(listener types.SampleListener) {
2021-10-02 00:46:10 +13:00
ptr := reflect.ValueOf(listener).Pointer()
emitKeyframe := false
2021-09-27 11:50:49 +13:00
2021-10-02 00:46:10 +13:00
manager.listenersMu.Lock()
if manager.waitForKf {
// 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
} else {
// otherwise, add it as a regular listener
manager.listeners[ptr] = listener
}
2021-10-02 00:46:10 +13:00
manager.listenersMu.Unlock()
2021-10-02 00:46:10 +13:00
manager.logger.Debug().Interface("ptr", ptr).Msgf("adding listener")
2022-06-15 10:23:16 +12:00
manager.currentListeners.Set(float64(manager.ListenersCount()))
// if we will be waiting for a keyframe, emit one now
if manager.pipeline != nil && emitKeyframe {
manager.pipeline.EmitVideoKeyframe()
}
}
func (manager *StreamSinkManagerCtx) removeListener(listener types.SampleListener) {
2021-09-27 11:50:49 +13:00
ptr := reflect.ValueOf(listener).Pointer()
2021-09-27 11:50:49 +13:00
manager.listenersMu.Lock()
delete(manager.listeners, ptr)
delete(manager.listenersKf, ptr) // if it's a keyframe listener, remove it too
2021-09-27 11:50:49 +13:00
manager.listenersMu.Unlock()
2021-09-27 11:50:49 +13:00
manager.logger.Debug().Interface("ptr", ptr).Msgf("removing listener")
2022-06-15 10:23:16 +12:00
manager.currentListeners.Set(float64(manager.ListenersCount()))
2021-10-02 00:46:10 +13:00
}
2021-02-06 00:49:02 +13:00
func (manager *StreamSinkManagerCtx) AddListener(listener types.SampleListener) error {
2021-09-27 12:17:25 +13:00
manager.mu.Lock()
2021-10-02 00:46:10 +13:00
defer manager.mu.Unlock()
2021-09-27 12:17:25 +13:00
2021-10-02 00:46:10 +13:00
if listener == nil {
return errors.New("listener cannot be nil")
}
2021-09-29 12:03:39 +13:00
2021-10-02 00:46:10 +13:00
// start if stopped
if err := manager.start(); err != nil {
return err
}
2021-02-06 00:41:02 +13:00
2021-10-02 00:46:10 +13:00
// add listener
manager.addListener(listener)
2021-09-29 11:21:53 +13:00
2021-10-02 00:46:10 +13:00
return nil
}
func (manager *StreamSinkManagerCtx) RemoveListener(listener types.SampleListener) error {
2021-10-02 00:46:10 +13:00
manager.mu.Lock()
defer manager.mu.Unlock()
if listener == nil {
return errors.New("listener cannot be nil")
}
// remove listener
manager.removeListener(listener)
// 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")
}
2021-10-02 00:46:10 +13:00
2021-12-02 08:30:18 +13:00
targetStream, ok := stream.(*StreamSinkManagerCtx)
2021-10-02 00:46:10 +13:00
if !ok {
return errors.New("target stream manager does not support moving listeners")
2021-10-02 00:46:10 +13:00
}
// 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
2021-12-02 08:30:18 +13:00
moveSinkListenerMu.Lock()
// lock source stream
manager.mu.Lock()
defer manager.mu.Unlock()
// lock target stream
targetStream.mu.Lock()
defer targetStream.mu.Unlock()
// unlock global mutex
2021-12-02 08:30:18 +13:00
moveSinkListenerMu.Unlock()
2021-09-29 12:03:39 +13:00
2021-10-02 00:46:10 +13:00
// start if stopped
if err := targetStream.start(); err != nil {
return err
}
// swap listeners
manager.removeListener(listener)
targetStream.addListener(listener)
// stop if started
manager.stop()
return nil
}
2021-12-02 08:30:18 +13:00
func (manager *StreamSinkManagerCtx) ListenersCount() int {
2021-09-27 11:50:49 +13:00
manager.listenersMu.Lock()
defer manager.listenersMu.Unlock()
2021-02-06 00:49:02 +13:00
return len(manager.listeners) + len(manager.listenersKf)
}
2021-12-02 08:30:18 +13:00
func (manager *StreamSinkManagerCtx) Started() bool {
2021-10-02 00:46:10 +13:00
return manager.ListenersCount() > 0
}
func (manager *StreamSinkManagerCtx) CreatePipeline() error {
2021-09-27 11:50:49 +13:00
manager.pipelineMu.Lock()
defer manager.pipelineMu.Unlock()
2021-02-06 00:49:02 +13:00
if manager.pipeline != nil {
2021-08-30 03:09:13 +12:00
return types.ErrCapturePipelineAlreadyExists
2021-02-06 00:49:02 +13:00
}
2022-09-18 03:52:55 +12:00
pipelineStr, err := manager.pipelineFn()
if err != nil {
return err
}
manager.logger.Info().
2021-12-02 08:30:18 +13:00
Str("codec", manager.codec.Name).
2021-02-10 08:36:22 +13:00
Str("src", pipelineStr).
Msgf("creating pipeline")
2021-02-10 08:36:22 +13:00
manager.pipeline, err = gst.CreatePipeline(pipelineStr)
if err != nil {
return err
}
2021-11-29 10:19:06 +13:00
manager.pipeline.AttachAppsink("appsink")
manager.pipeline.Play()
2022-01-23 05:58:05 +13:00
manager.wg.Add(1)
2022-02-15 06:41:47 +13:00
pipeline := manager.pipeline
2022-01-23 05:58:05 +13:00
go func() {
manager.logger.Debug().Msg("started emitting samples")
defer manager.wg.Done()
for {
2022-10-14 06:57:31 +13:00
sample, ok := <-pipeline.Sample()
2022-01-23 05:58:05 +13:00
if !ok {
manager.logger.Debug().Msg("stopped emitting samples")
return
}
manager.onSample(sample)
2022-01-23 05:58:05 +13:00
}
}()
2022-06-15 10:23:16 +12:00
manager.pipelinesCounter.Inc()
2022-06-19 11:03:16 +12:00
manager.pipelinesActive.Set(1)
return nil
}
func (manager *StreamSinkManagerCtx) saveSampleBitrate(timestamp time.Time, delta float64) {
// get unix timestamp in seconds
sec := timestamp.Unix()
// last bucket is timestamp rounded to 3 seconds - 1 second
last := int((sec - 1) % 3)
// current bucket is timestamp rounded to 3 seconds
curr := int(sec % 3)
// next bucket is timestamp rounded to 3 seconds + 1 second
next := int((sec + 1) % 3)
if manager.brBuckets[next] != 0 {
// atomic update bitrate
atomic.StoreUint64(&manager.bitrate, uint64(manager.brBuckets[last]))
// empty next bucket
manager.brBuckets[next] = 0
}
// add rate to current bucket
manager.brBuckets[curr] += delta
}
func (manager *StreamSinkManagerCtx) onSample(sample types.Sample) {
manager.listenersMu.Lock()
defer manager.listenersMu.Unlock()
// 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 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)
}
for _, l := range manager.listeners {
l.WriteSample(sample)
}
}
func (manager *StreamSinkManagerCtx) DestroyPipeline() {
2021-09-27 11:50:49 +13:00
manager.pipelineMu.Lock()
defer manager.pipelineMu.Unlock()
if manager.pipeline == nil {
return
}
2021-12-06 06:16:26 +13:00
manager.pipeline.Destroy()
2021-02-06 01:58:02 +13:00
manager.logger.Info().Msgf("destroying pipeline")
manager.pipeline = nil
2022-06-19 11:03:16 +12:00
manager.pipelinesActive.Set(0)
manager.brBuckets = make(map[int]float64)
atomic.StoreUint64(&manager.bitrate, 0)
}