iridium/neko
2
2
Fork 0
mirror of https://github.com/m1k1o/neko.git synced 2024-07-24 14:40:50 +12:00
Code Issues Packages Projects Releases Wiki Activity

webrtc refactor peer track.

Browse Source
This commit is contained in:
Miroslav Šedivý 2021-09-27 00:50:49 +02:00
parent beac1cb088
commit 9d4d5766ef
4 changed files with 219 additions and 209 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

136
internal/capture/stream.go
View File

@ -1,6 +1,7 @@
package capture
import (
"errors"
"reflect"
"sync"
@ -17,14 +18,18 @@ type StreamManagerCtx struct {
mu sync.Mutex
wg sync.WaitGroup
codec codec.RTPCodec
pipelineStr func() string
pipeline *gst.Pipeline
pipelineMu sync.Mutex
pipelineStr func() string
sample chan types.Sample
sampleStop chan interface{}
sampleUpdate chan interface{}
listeners map[uintptr]*func(sample types.Sample)
emitMu sync.Mutex
emitUpdate chan bool
emitStop chan bool
started bool
listenersMu sync.Mutex
listenersCount uint32
}
func streamNew(codec codec.RTPCodec, pipelineStr func() string, video_id string) *StreamManagerCtx {
@ -37,10 +42,9 @@ func streamNew(codec codec.RTPCodec, pipelineStr func() string, video_id string)
logger: logger,
codec: codec,
pipelineStr: pipelineStr,
sampleStop: make(chan interface{}),
sampleUpdate: make(chan interface{}),
listeners: map[uintptr]*func(sample types.Sample){},
emitUpdate: make(chan bool),
emitStop: make(chan bool),
started: false,
}
manager.wg.Add(1)
@ -51,17 +55,17 @@ func streamNew(codec codec.RTPCodec, pipelineStr func() string, video_id string)
for {
select {
case <-manager.emitStop:
case <-manager.sampleStop:
manager.logger.Debug().Msg("stopped emitting samples")
return
case <-manager.emitUpdate:
case <-manager.sampleUpdate:
manager.logger.Debug().Msg("update emitting samples")
case sample := <-manager.sample:
manager.emitMu.Lock()
manager.listenersMu.Lock()
for _, emit := range manager.listeners {
(*emit)(sample)
}
manager.emitMu.Unlock()
manager.listenersMu.Unlock()
}
}
}()
@ -72,15 +76,15 @@ func streamNew(codec codec.RTPCodec, pipelineStr func() string, video_id string)
func (manager *StreamManagerCtx) shutdown() {
manager.logger.Info().Msgf("shutdown")
manager.emitMu.Lock()
manager.listenersMu.Lock()
for key := range manager.listeners {
delete(manager.listeners, key)
}
manager.emitMu.Unlock()
manager.listenersMu.Unlock()
manager.destroyPipeline()
manager.emitStop <- true
close(manager.sampleStop)
manager.wg.Wait()
}
@ -88,63 +92,78 @@ func (manager *StreamManagerCtx) Codec() codec.RTPCodec {
return manager.codec
}
func (manager *StreamManagerCtx) AddListener(listener *func(sample types.Sample)) {
manager.emitMu.Lock()
defer manager.emitMu.Unlock()
if listener != nil {
ptr := reflect.ValueOf(listener).Pointer()
manager.listeners[ptr] = listener
manager.logger.Debug().Interface("ptr", ptr).Msgf("adding listener")
func (manager *StreamManagerCtx) NewListener(listener *func(sample types.Sample)) (addListener func(), err error) {
if listener == nil {
return addListener, errors.New("listener cannot be nil")
}
manager.mu.Lock()
defer manager.mu.Unlock()
if manager.listenersCount == 0 {
err := manager.createPipeline()
if err != nil && !errors.Is(err, types.ErrCapturePipelineAlreadyExists) {
return addListener, err
}
manager.listenersCount++
manager.logger.Info().Msgf("first listener, starting")
}
return func() {
ptr := reflect.ValueOf(listener).Pointer()
manager.listenersMu.Lock()
manager.listeners[ptr] = listener
manager.listenersMu.Unlock()
manager.logger.Debug().Interface("ptr", ptr).Msgf("adding listener")
}, nil
}
func (manager *StreamManagerCtx) RemoveListener(listener *func(sample types.Sample)) {
manager.emitMu.Lock()
defer manager.emitMu.Unlock()
if listener != nil {
ptr := reflect.ValueOf(listener).Pointer()
delete(manager.listeners, ptr)
manager.logger.Debug().Interface("ptr", ptr).Msgf("removing listener")
if listener == nil {
return
}
ptr := reflect.ValueOf(listener).Pointer()
manager.listenersMu.Lock()
delete(manager.listeners, ptr)
manager.listenersMu.Unlock()
manager.logger.Debug().Interface("ptr", ptr).Msgf("removing listener")
go func() {
manager.mu.Lock()
defer manager.mu.Unlock()
if manager.listenersCount == 1 {
manager.destroyPipeline()
manager.listenersCount = 0
manager.logger.Info().Msgf("last listener, stopping")
}
}()
}
func (manager *StreamManagerCtx) ListenersCount() int {
manager.emitMu.Lock()
defer manager.emitMu.Unlock()
manager.listenersMu.Lock()
defer manager.listenersMu.Unlock()
return len(manager.listeners)
}
func (manager *StreamManagerCtx) Start() error {
manager.mu.Lock()
defer manager.mu.Unlock()
err := manager.createPipeline()
if err != nil {
return err
}
manager.logger.Info().Msgf("start")
manager.started = true
return nil
}
func (manager *StreamManagerCtx) Stop() {
manager.mu.Lock()
defer manager.mu.Unlock()
manager.logger.Info().Msgf("stop")
manager.started = false
manager.destroyPipeline()
}
func (manager *StreamManagerCtx) Started() bool {
return manager.started
manager.mu.Lock()
defer manager.mu.Unlock()
return manager.listenersCount > 0
}
func (manager *StreamManagerCtx) createPipeline() error {
manager.pipelineMu.Lock()
defer manager.pipelineMu.Unlock()
if manager.pipeline != nil {
return types.ErrCapturePipelineAlreadyExists
}
@ -166,11 +185,14 @@ func (manager *StreamManagerCtx) createPipeline() error {
manager.pipeline.Start()
manager.sample = manager.pipeline.Sample
manager.emitUpdate <- true
manager.sampleUpdate <- struct{}{}
return nil
}
func (manager *StreamManagerCtx) destroyPipeline() {
manager.pipelineMu.Lock()
defer manager.pipelineMu.Unlock()
if manager.pipeline == nil {
return
}

8
internal/types/capture.go
View File

@ -35,12 +35,12 @@ type ScreencastManager interface {
type StreamManager interface {
Codec() codec.RTPCodec
AddListener(listener *func(sample Sample))
// starts pipeline if was not running before and returns register function
NewListener(listener *func(sample Sample)) (addListener func(), err error)
// stops pipeline if it was last listener
RemoveListener(listener *func(sample Sample))
ListenersCount() int
Start() error
Stop()
ListenersCount() int
Started() bool
}

165
internal/webrtc/manager.go
View File

@ -1,15 +1,11 @@
package webrtc
import (
"errors"
"fmt"
"io"
"strings"
"sync"
"time"
"github.com/pion/webrtc/v3"
"github.com/pion/webrtc/v3/pkg/media"
"github.com/rs/zerolog"
"github.com/rs/zerolog/log"
@ -39,13 +35,10 @@ func New(desktop types.DesktopManager, capture types.CaptureManager, config *con
capture: capture,
curImage: cursor.NewImage(desktop),
curPosition: cursor.NewPosition(desktop),
participants: 0,
}
}
type WebRTCManagerCtx struct {
mu sync.Mutex
logger zerolog.Logger
config *config.WebRTC
@ -53,33 +46,9 @@ type WebRTCManagerCtx struct {
capture types.CaptureManager
curImage *cursor.ImageCtx
curPosition *cursor.PositionCtx
audioTrack *webrtc.TrackLocalStaticSample
audioListener func(sample types.Sample)
participants uint32
}
func (manager *WebRTCManagerCtx) Start() {
var err error
// create audio track
audio := manager.capture.Audio()
manager.audioTrack, err = webrtc.NewTrackLocalStaticSample(audio.Codec().Capability, "audio", "stream")
if err != nil {
manager.logger.Panic().Err(err).Msg("unable to create audio track")
}
manager.audioListener = func(sample types.Sample) {
if err := manager.audioTrack.WriteSample(media.Sample(sample)); err != nil {
if errors.Is(err, io.ErrClosedPipe) {
// The peerConnection has been closed.
return
}
manager.logger.Warn().Err(err).Msg("audio pipeline failed to write")
}
}
audio.AddListener(&manager.audioListener)
manager.curImage.Start()
manager.logger.Info().
@ -97,9 +66,6 @@ func (manager *WebRTCManagerCtx) Shutdown() error {
manager.curImage.Shutdown()
manager.curPosition.Shutdown()
audio := manager.capture.Audio()
audio.RemoveListener(&manager.audioListener)
return nil
}
@ -112,6 +78,9 @@ func (manager *WebRTCManagerCtx) CreatePeer(session types.Session, videoID strin
logger := manager.logger.With().Str("session_id", session.ID()).Logger()
logger.Info().Msg("creating webrtc peer")
// all audios must have the same codec
audioStream := manager.capture.Audio()
// all videos must have the same codec
videoStream, ok := manager.capture.Video(videoID)
if !ok {
@ -119,8 +88,8 @@ func (manager *WebRTCManagerCtx) CreatePeer(session types.Session, videoID strin
}
connection, err := manager.newPeerConnection([]codec.RTPCodec{
audioStream.Codec(),
videoStream.Codec(),
manager.capture.Audio().Codec(),
}, logger)
if err != nil {
return nil, err
@ -142,79 +111,32 @@ func (manager *WebRTCManagerCtx) CreatePeer(session types.Session, videoID strin
})
}
// create video track
videoTrack, err := webrtc.NewTrackLocalStaticSample(videoStream.Codec().Capability, "video", "stream")
// audio track
audioTrack, err := manager.newPeerTrack(audioStream, logger)
if err != nil {
return nil, err
}
videoListener := func(sample types.Sample) {
if err := videoTrack.WriteSample(media.Sample(sample)); err != nil {
if errors.Is(err, io.ErrClosedPipe) {
// The peerConnection has been closed.
return
}
logger.Warn().Err(err).Msg("video pipeline failed to write")
}
}
manager.mu.Lock()
// should be stream started
if videoStream.ListenersCount() == 0 {
if err := videoStream.Start(); err != nil {
return nil, err
}
}
videoStream.AddListener(&videoListener)
// start audio, when first participant connects
if !manager.capture.Audio().Started() {
if err := manager.capture.Audio().Start(); err != nil {
manager.logger.Panic().Err(err).Msg("unable to start audio stream")
}
}
manager.participants = manager.participants + 1
manager.mu.Unlock()
changeVideo := func(videoID string) error {
newVideoStream, ok := manager.capture.Video(videoID)
if !ok {
return types.ErrWebRTCVideoNotFound
}
// should be new stream started
if newVideoStream.ListenersCount() == 0 {
if err := newVideoStream.Start(); err != nil {
return err
}
}
// switch videoListeners
videoStream.RemoveListener(&videoListener)
newVideoStream.AddListener(&videoListener)
// should be old stream stopped
if videoStream.ListenersCount() == 0 {
videoStream.Stop()
}
videoStream = newVideoStream
return nil
}
rtpAudio, err := connection.AddTrack(manager.audioTrack)
audioTrack.AddToConnection(connection)
if err != nil {
return nil, err
}
rtpVideo, err := connection.AddTrack(videoTrack)
// video track
videoTrack, err := manager.newPeerTrack(videoStream, logger)
if err != nil {
return nil, err
}
videoTrack.AddToConnection(connection)
if err != nil {
return nil, err
}
// data channel
dataChannel, err := connection.CreateDataChannel("data", nil)
if err != nil {
return nil, err
@ -224,7 +146,14 @@ func (manager *WebRTCManagerCtx) CreatePeer(session types.Session, videoID strin
logger: logger,
connection: connection,
dataChannel: dataChannel,
changeVideo: changeVideo,
changeVideo: func(videoID string) error {
videoStream, ok := manager.capture.Video(videoID)
if !ok {
return types.ErrWebRTCVideoNotFound
}
return videoTrack.SetStream(videoStream)
},
iceTrickle: manager.config.ICETrickle,
}
@ -252,29 +181,9 @@ func (manager *WebRTCManagerCtx) CreatePeer(session types.Session, videoID strin
webrtc.PeerConnectionStateFailed:
connection.Close()
case webrtc.PeerConnectionStateClosed:
manager.mu.Lock()
session.SetWebRTCConnected(peer, false)
videoStream.RemoveListener(&videoListener)
// should be stream stopped
if videoStream.ListenersCount() == 0 {
videoStream.Stop()
}
// decrease participants
manager.participants = manager.participants - 1
// stop audio, if last participant disonnects
if manager.participants <= 0 {
manager.participants = 0
if manager.capture.Audio().Started() {
manager.capture.Audio().Stop()
}
}
manager.mu.Unlock()
videoTrack.RemoveStream()
audioTrack.RemoveStream()
}
})
@ -310,24 +219,6 @@ func (manager *WebRTCManagerCtx) CreatePeer(session types.Session, videoID strin
}
})
go func() {
rtcpBuf := make([]byte, 1500)
for {
if _, _, err := rtpAudio.Read(rtcpBuf); err != nil {
return
}
}
}()
go func() {
rtcpBuf := make([]byte, 1500)
for {
if _, _, err := rtpVideo.Read(rtcpBuf); err != nil {
return
}
}
}()
session.SetWebRTCPeer(peer)
return peer.CreateOffer(false)
}

97
internal/webrtc/peertrack.go Normal file
View File

@ -0,0 +1,97 @@
package webrtc
import (
"demodesk/neko/internal/types"
"errors"
"io"
"sync"
"github.com/pion/webrtc/v3"
"github.com/pion/webrtc/v3/pkg/media"
"github.com/rs/zerolog"
)
func (manager *WebRTCManagerCtx) newPeerTrack(stream types.StreamManager, logger zerolog.Logger) (*PeerTrack, error) {
codec := stream.Codec()
id := codec.Type.String()
track, err := webrtc.NewTrackLocalStaticSample(codec.Capability, id, "stream")
if err != nil {
return nil, err
}
logger = logger.With().Str("id", id).Logger()
peer := &PeerTrack{
logger: logger,
track: track,
listener: func(sample types.Sample) {
err := track.WriteSample(media.Sample(sample))
if err != nil && errors.Is(err, io.ErrClosedPipe) {
logger.Warn().Err(err).Msg("pipeline failed to write")
}
},
}
peer.SetStream(stream)
return peer, nil
}
type PeerTrack struct {
logger zerolog.Logger
track *webrtc.TrackLocalStaticSample
listener func(sample types.Sample)
streamMu sync.Mutex
stream types.StreamManager
}
func (peer *PeerTrack) SetStream(stream types.StreamManager) error {
peer.streamMu.Lock()
defer peer.streamMu.Unlock()
// prepare new listener
addListener, err := stream.NewListener(&peer.listener)
if err != nil {
return err
}
// remove previous listener (in case it existed)
if peer.stream != nil {
peer.stream.RemoveListener(&peer.listener)
}
// add new listener
addListener()
peer.stream = stream
return nil
}
func (peer *PeerTrack) RemoveStream() {
peer.streamMu.Lock()
defer peer.streamMu.Unlock()
if peer.stream != nil {
peer.stream.RemoveListener(&peer.listener)
}
}
func (peer *PeerTrack) AddToConnection(connection *webrtc.PeerConnection) error {
sender, err := connection.AddTrack(peer.track)
if err != nil {
return err
}
go func() {
rtcpBuf := make([]byte, 1500)
for {
if _, _, err := sender.Read(rtcpBuf); err != nil {
return
}
}
}()
return nil
}