avoid sending to closed channels

This commit is contained in:
David Zhao
2021-01-26 21:45:00 -08:00
parent a065a01592
commit ed0b9db655
12 changed files with 195 additions and 46 deletions
+7 -3
View File
@@ -2,11 +2,13 @@ package routing
import (
"google.golang.org/protobuf/proto"
"github.com/livekit/livekit-server/pkg/utils"
)
type MessageChannel struct {
msgChan chan proto.Message
isClosed bool
isClosed utils.AtomicFlag
onClose func()
}
@@ -22,7 +24,7 @@ func (m *MessageChannel) OnClose(f func()) {
}
func (m *MessageChannel) WriteMessage(msg proto.Message) error {
if m.isClosed {
if m.isClosed.Get() {
return ErrChannelClosed
}
m.msgChan <- msg
@@ -34,7 +36,9 @@ func (m *MessageChannel) ReadChan() <-chan proto.Message {
}
func (m *MessageChannel) Close() {
m.isClosed = true
if !m.isClosed.TrySet(true) {
return
}
close(m.msgChan)
if m.onClose != nil {
m.onClose()
+10 -12
View File
@@ -2,11 +2,11 @@ package routing
import (
"context"
"sync"
"github.com/go-redis/redis/v8"
"google.golang.org/protobuf/proto"
"github.com/livekit/livekit-server/pkg/utils"
"github.com/livekit/livekit-server/proto/livekit"
)
@@ -69,8 +69,7 @@ type RedisSink struct {
rc *redis.Client
nodeId string
participantId string
isClosed bool
once sync.Once
isClosed utils.AtomicFlag
onClose func()
}
@@ -78,26 +77,25 @@ func NewRedisSink(rc *redis.Client, nodeId, participantId string) *RedisSink {
return &RedisSink{
rc: rc,
nodeId: nodeId,
once: sync.Once{},
participantId: participantId,
}
}
func (s *RedisSink) WriteMessage(msg proto.Message) error {
if s.isClosed {
if s.isClosed.Get() {
return ErrChannelClosed
}
return publishRouterMessage(s.rc, s.nodeId, s.participantId, msg)
}
func (s *RedisSink) Close() {
s.once.Do(func() {
publishRouterMessage(s.rc, s.nodeId, s.participantId, &livekit.EndSession{})
s.isClosed = true
if s.onClose != nil {
s.onClose()
}
})
if !s.isClosed.TrySet(true) {
return
}
publishRouterMessage(s.rc, s.nodeId, s.participantId, &livekit.EndSession{})
if s.onClose != nil {
s.onClose()
}
}
func (s *RedisSink) OnClose(f func()) {
+11 -1
View File
@@ -261,18 +261,28 @@ func (r *RedisRouter) statsWorker() {
func (r *RedisRouter) subscribeWorker() {
sub := r.rc.Subscribe(redisCtx, nodeChannel(r.currentNode.Id))
defer func() {
logger.Debugw("finishing redis subscribeWorker", "node", r.currentNode.Id)
}()
logger.Debugw("starting redis subscribeWorker", "node", r.currentNode.Id)
for r.ctx.Err() == nil {
obj, err := sub.Receive(r.ctx)
if err != nil {
logger.Warnw("error receiving redis message", "error", err)
// TODO: retry? ignore? at a minimum need to sleep here to retry
time.Sleep(100 * time.Millisecond)
time.Sleep(time.Second)
continue
}
if obj == nil {
return
}
msg, ok := obj.(*redis.Message)
if !ok {
continue
}
logger.Debugw("received redis message", "message", msg, "node", r.currentNode.Id)
rm := livekit.RouterMessage{}
err = proto.Unmarshal([]byte(msg.Payload), &rm)
pId := rm.ParticipantId
+7 -6
View File
@@ -15,6 +15,7 @@ import (
"github.com/livekit/livekit-server/pkg/logger"
"github.com/livekit/livekit-server/pkg/rtc/types"
"github.com/livekit/livekit-server/pkg/sfu"
"github.com/livekit/livekit-server/pkg/utils"
"github.com/livekit/livekit-server/proto/livekit"
)
@@ -42,7 +43,7 @@ type MediaTrack struct {
onClose func()
// channel to send RTCP packets to the source
rtcpCh chan []rtcp.Packet
rtcpCh *utils.CalmChannel
lock sync.RWMutex
once sync.Once
// map of target participantId -> DownTrack
@@ -53,7 +54,7 @@ type MediaTrack struct {
lastPLI time.Time
}
func NewMediaTrack(trackId string, pId string, rtcpCh chan []rtcp.Packet, track *webrtc.TrackRemote, receiver types.Receiver) *MediaTrack {
func NewMediaTrack(trackId string, pId string, rtcpCh *utils.CalmChannel, track *webrtc.TrackRemote, receiver types.Receiver) *MediaTrack {
t := &MediaTrack{
id: trackId,
participantId: pId,
@@ -198,7 +199,7 @@ func (t *MediaTrack) RemoveAllSubscribers() {
t.downtracks = make(map[string]types.DownTrack)
}
func (t *MediaTrack) sendDownTrackBindingReports(participantId string, rtcpCh chan<- []rtcp.Packet) {
func (t *MediaTrack) sendDownTrackBindingReports(participantId string, rtcpCh *utils.CalmChannel) {
var sd []rtcp.SourceDescriptionChunk
t.lock.RLock()
@@ -221,7 +222,7 @@ func (t *MediaTrack) sendDownTrackBindingReports(participantId string, rtcpCh ch
batch := pkts
i := 0
for {
rtcpCh <- batch
rtcpCh.Write(batch)
if i > 5 {
return
}
@@ -277,7 +278,7 @@ func (t *MediaTrack) forwardRTPWorker() {
rtcpPkts := []rtcp.Packet{
&rtcp.PictureLossIndication{SenderSSRC: uint32(t.ssrc), MediaSSRC: pkt.SSRC},
}
t.rtcpCh <- rtcpPkts
t.rtcpCh.Write(rtcpPkts)
t.lastPLI = time.Now()
} else if err != nil {
logger.Warnw("could not forward packet to participant",
@@ -351,6 +352,6 @@ func (t *MediaTrack) handleRTCP(dt *sfu.DownTrack, rtcpBuf []byte) {
}
if len(fwdPkts) > 0 {
t.rtcpCh <- fwdPkts
t.rtcpCh.Write(fwdPkts)
}
}
+7 -2
View File
@@ -64,7 +64,12 @@ func TestMissingKeyFrames(t *testing.T) {
time.Sleep(testWaitDuration)
select {
case pkts := <-mt.rtcpCh:
case o := <-mt.rtcpCh.ReadChan():
if o == nil {
return
}
pkts, ok := o.([]rtcp.Packet)
assert.True(t, ok)
assert.Len(t, pkts, 1, "a single RTCP packet should be returned")
assert.IsType(t, &rtcp.PictureLossIndication{}, pkts[0])
default:
@@ -87,7 +92,7 @@ func newMediaTrackWithReceiver() *MediaTrack {
muted: false,
kind: livekit.TrackType_VIDEO,
codec: webrtc.RTPCodecParameters{},
rtcpCh: make(chan []rtcp.Packet, 5),
rtcpCh: utils.NewCalmChannel(5),
lock: sync.RWMutex{},
once: sync.Once{},
downtracks: map[string]types.DownTrack{},
+15 -7
View File
@@ -41,7 +41,7 @@ type ParticipantImpl struct {
mediaEngine *webrtc.MediaEngine
name string
state livekit.ParticipantInfo_State
rtcpCh chan []rtcp.Packet
rtcpCh *utils.CalmChannel
subscribedTracks map[string][]*sfu.DownTrack
// publishedTracks that participant is publishing
publishedTracks map[string]types.PublishedTrack
@@ -88,7 +88,7 @@ func NewParticipant(participantId, name string, pc types.PeerConnection, rs rout
receiverConfig: receiverConfig,
ctx: ctx,
cancel: cancel,
rtcpCh: make(chan []rtcp.Packet, 50),
rtcpCh: utils.NewCalmChannel(50),
subscribedTracks: make(map[string][]*sfu.DownTrack),
state: livekit.ParticipantInfo_JOINING,
lock: sync.RWMutex{},
@@ -163,7 +163,7 @@ func (p *ParticipantImpl) IsReady() bool {
return p.state == livekit.ParticipantInfo_JOINED || p.state == livekit.ParticipantInfo_ACTIVE
}
func (p *ParticipantImpl) RTCPChan() chan<- []rtcp.Packet {
func (p *ParticipantImpl) RTCPChan() *utils.CalmChannel {
return p.rtcpCh
}
@@ -331,8 +331,9 @@ func (p *ParticipantImpl) Close() error {
p.onClose(p)
}
p.cancel()
close(p.rtcpCh)
return p.peerConn.Close()
p.peerConn.Close()
p.rtcpCh.Close()
return nil
}
// Subscribes otherPeer to all of the publishedTracks
@@ -482,7 +483,7 @@ func (p *ParticipantImpl) updateState(state livekit.ParticipantInfo_State) {
}
oldState := p.state
p.state = state
logger.Debugw("updating participant state", "state", state.String())
logger.Debugw("updating participant state", "state", state.String(), "participant", p.ID())
if p.onStateChange != nil {
go func() {
defer Recover()
@@ -612,7 +613,14 @@ func (p *ParticipantImpl) downTracksRTCPWorker() {
func (p *ParticipantImpl) rtcpSendWorker() {
defer Recover()
// read from rtcpChan
for pkts := range p.rtcpCh {
for o := range p.rtcpCh.ReadChan() {
if o == nil {
return
}
pkts, ok := o.([]rtcp.Packet)
if !ok {
logger.Errorw("unexpected obj from rtcpCh", "object", o)
}
//for _, pkt := range pkts {
// logger.Debugw("writing RTCP", "packet", pkt)
//}
+14 -3
View File
@@ -7,6 +7,7 @@ import (
"github.com/pion/webrtc/v3"
"github.com/livekit/livekit-server/pkg/logger"
"github.com/livekit/livekit-server/pkg/utils"
)
const (
@@ -20,10 +21,10 @@ type ReceiverImpl struct {
track *webrtc.TrackRemote
buffer *buffer.Buffer
rtcpReader *buffer.RTCPReader
rtcpChan chan []rtcp.Packet
rtcpChan *utils.CalmChannel
}
func NewReceiver(rtcpCh chan []rtcp.Packet, rtpReceiver *webrtc.RTPReceiver, track *webrtc.TrackRemote, config ReceiverConfig) *ReceiverImpl {
func NewReceiver(rtcpCh *utils.CalmChannel, rtpReceiver *webrtc.RTPReceiver, track *webrtc.TrackRemote, config ReceiverConfig) *ReceiverImpl {
r := &ReceiverImpl{
rtpReceiver: rtpReceiver,
rtcpChan: rtcpCh,
@@ -34,8 +35,10 @@ func NewReceiver(rtcpCh chan []rtcp.Packet, rtpReceiver *webrtc.RTPReceiver, tra
// when we have feedback for the sender, send through the rtcp channel
r.buffer.OnFeedback(func(fb []rtcp.Packet) {
RecoverSilent()
// rtcpChan could be closed
if r.rtcpChan != nil {
r.rtcpChan <- fb
r.rtcpChan.Write(fb)
}
})
@@ -68,6 +71,14 @@ func NewReceiver(rtcpCh chan []rtcp.Packet, rtpReceiver *webrtc.RTPReceiver, tra
return r
}
func (r *ReceiverImpl) Close() {
r.rtcpChan = nil
r.buffer.OnFeedback(nil)
r.buffer.Close()
r.rtpReceiver.Stop()
r.rtcpReader.Close()
}
// PacketBuffer interface, retrieves a packet from buffer and deserializes
// it's possible that the packet can't be found, or the connection has been closed (io.EOF)
func (r *ReceiverImpl) GetBufferedPacket(pktBuf []byte, sn uint16, snOffset uint16) (p rtp.Packet, err error) {
+2 -1
View File
@@ -8,6 +8,7 @@ import (
"github.com/pion/webrtc/v3"
"github.com/livekit/livekit-server/pkg/sfu"
"github.com/livekit/livekit-server/pkg/utils"
"github.com/livekit/livekit-server/proto/livekit"
)
@@ -51,7 +52,7 @@ type Participant interface {
State() livekit.ParticipantInfo_State
IsReady() bool
ToProto() *livekit.ParticipantInfo
RTCPChan() chan<- []rtcp.Packet
RTCPChan() *utils.CalmChannel
AddTrack(clientId, name string, trackType livekit.TrackType)
Answer(sdp webrtc.SessionDescription) (answer webrtc.SessionDescription, err error)
+11 -11
View File
@@ -6,8 +6,8 @@ import (
"github.com/livekit/livekit-server/pkg/rtc/types"
"github.com/livekit/livekit-server/pkg/sfu"
"github.com/livekit/livekit-server/pkg/utils"
"github.com/livekit/livekit-server/proto/livekit"
"github.com/pion/rtcp"
webrtc "github.com/pion/webrtc/v3"
)
@@ -150,15 +150,15 @@ type FakeParticipant struct {
peerConnectionReturnsOnCall map[int]struct {
result1 types.PeerConnection
}
RTCPChanStub func() chan<- []rtcp.Packet
RTCPChanStub func() *utils.CalmChannel
rTCPChanMutex sync.RWMutex
rTCPChanArgsForCall []struct {
}
rTCPChanReturns struct {
result1 chan<- []rtcp.Packet
result1 *utils.CalmChannel
}
rTCPChanReturnsOnCall map[int]struct {
result1 chan<- []rtcp.Packet
result1 *utils.CalmChannel
}
RemoveDownTrackStub func(string, *sfu.DownTrack)
removeDownTrackMutex sync.RWMutex
@@ -991,7 +991,7 @@ func (fake *FakeParticipant) PeerConnectionReturnsOnCall(i int, result1 types.Pe
}{result1}
}
func (fake *FakeParticipant) RTCPChan() chan<- []rtcp.Packet {
func (fake *FakeParticipant) RTCPChan() *utils.CalmChannel {
fake.rTCPChanMutex.Lock()
ret, specificReturn := fake.rTCPChanReturnsOnCall[len(fake.rTCPChanArgsForCall)]
fake.rTCPChanArgsForCall = append(fake.rTCPChanArgsForCall, struct {
@@ -1015,32 +1015,32 @@ func (fake *FakeParticipant) RTCPChanCallCount() int {
return len(fake.rTCPChanArgsForCall)
}
func (fake *FakeParticipant) RTCPChanCalls(stub func() chan<- []rtcp.Packet) {
func (fake *FakeParticipant) RTCPChanCalls(stub func() *utils.CalmChannel) {
fake.rTCPChanMutex.Lock()
defer fake.rTCPChanMutex.Unlock()
fake.RTCPChanStub = stub
}
func (fake *FakeParticipant) RTCPChanReturns(result1 chan<- []rtcp.Packet) {
func (fake *FakeParticipant) RTCPChanReturns(result1 *utils.CalmChannel) {
fake.rTCPChanMutex.Lock()
defer fake.rTCPChanMutex.Unlock()
fake.RTCPChanStub = nil
fake.rTCPChanReturns = struct {
result1 chan<- []rtcp.Packet
result1 *utils.CalmChannel
}{result1}
}
func (fake *FakeParticipant) RTCPChanReturnsOnCall(i int, result1 chan<- []rtcp.Packet) {
func (fake *FakeParticipant) RTCPChanReturnsOnCall(i int, result1 *utils.CalmChannel) {
fake.rTCPChanMutex.Lock()
defer fake.rTCPChanMutex.Unlock()
fake.RTCPChanStub = nil
if fake.rTCPChanReturnsOnCall == nil {
fake.rTCPChanReturnsOnCall = make(map[int]struct {
result1 chan<- []rtcp.Packet
result1 *utils.CalmChannel
})
}
fake.rTCPChanReturnsOnCall[i] = struct {
result1 chan<- []rtcp.Packet
result1 *utils.CalmChannel
}{result1}
}
+2
View File
@@ -200,6 +200,8 @@ func (r *RoomManager) rtcSessionWorker(room *rtc.Room, participant types.Partici
)
}()
defer rtc.Recover()
logger.Debugw("starting RTC session", "node", r.currentNode.Id)
for {
select {
case <-time.After(time.Millisecond * 100):
+63
View File
@@ -0,0 +1,63 @@
package utils
import (
"errors"
"sync/atomic"
)
var (
ErrChannelClosed = errors.New("cannot write to closed channel")
)
type AtomicFlag struct {
val uint32
}
// set flag to value if existing flag is different, otherwise return
func (b *AtomicFlag) TrySet(bVal bool) bool {
var v uint32
if bVal {
v = 1
}
old := b.val
// value is the same, nochanges
if old == v {
return false
}
return atomic.CompareAndSwapUint32(&b.val, old, v)
}
func (b *AtomicFlag) Get() bool {
return b.val == 1
}
// a channel that ignores writes when it closes instead of panic
type CalmChannel struct {
channel chan interface{}
isClosed AtomicFlag
}
func NewCalmChannel(size int) *CalmChannel {
return &CalmChannel{
channel: make(chan interface{}, size),
}
}
func (c *CalmChannel) Close() {
if !c.isClosed.TrySet(true) {
return
}
close(c.channel)
}
func (c *CalmChannel) ReadChan() <-chan interface{} {
return c.channel
}
func (c *CalmChannel) Write(o interface{}) error {
if c.isClosed.Get() {
return ErrChannelClosed
}
c.channel <- o
return nil
}
+46
View File
@@ -0,0 +1,46 @@
package utils
import (
"sync"
"sync/atomic"
"testing"
"github.com/stretchr/testify/assert"
)
func TestAtomicFlag(t *testing.T) {
t.Run("single thread, basic functionality", func(t *testing.T) {
f := AtomicFlag{}
assert.False(t, f.TrySet(false))
assert.False(t, f.Get())
assert.True(t, f.TrySet(true))
assert.True(t, f.Get())
})
t.Run("only one thread should succeed", func(t *testing.T) {
var numSets uint32
f := AtomicFlag{}
wg := sync.WaitGroup{}
for i := 0; i < 10; i++ {
wg.Add(1)
go func() {
if f.TrySet(true) {
atomic.AddUint32(&numSets, 1)
}
wg.Done()
}()
}
wg.Wait()
assert.Equal(t, 1, int(numSets))
})
}
func TestCalmChannel(t *testing.T) {
t.Run("should receive nil after close", func(t *testing.T) {
c := NewCalmChannel(1)
c.Close()
assert.Equal(t, nil, <-c.ReadChan())
})
}