From 9a10a57b30cea6d16f38fb82baea3309af0dbffc Mon Sep 17 00:00:00 2001 From: David Zhao Date: Wed, 13 Jan 2021 17:38:40 -0800 Subject: [PATCH] single publisher integration test --- cmd/cli/client/client.go | 108 ++++++++++++++++++++++------------ cmd/cli/client/trackwriter.go | 27 +++++++++ cmd/cli/commands/rtc.go | 8 +-- test/integration_helpers.go | 18 +++++- test/integration_test.go | 53 ++++++++++++++++- 5 files changed, 166 insertions(+), 48 deletions(-) diff --git a/cmd/cli/client/client.go b/cmd/cli/client/client.go index d0e1e8c36..c6635b566 100644 --- a/cmd/cli/client/client.go +++ b/cmd/cli/client/client.go @@ -9,7 +9,6 @@ import ( "net/http" "net/url" "path/filepath" - "strings" "sync" "time" @@ -25,6 +24,7 @@ import ( ) type RTCClient struct { + id string conn *websocket.Conn PeerConn *webrtc.PeerConnection localTracks []webrtc.TrackLocal @@ -35,7 +35,7 @@ type RTCClient struct { iceConnected bool paused bool me *webrtc.MediaEngine // optional, populated only when receiving tracks - subscribedTracks map[string]*webrtc.TrackRemote + subscribedTracks map[string][]*webrtc.TrackRemote localParticipant *livekit.ParticipantInfo remoteParticipants map[string]*livekit.ParticipantInfo @@ -90,7 +90,7 @@ func NewRTCClient(conn *websocket.Conn) (*RTCClient, error) { lock: sync.Mutex{}, pendingCandidates: make([]*webrtc.ICECandidate, 0), localTracks: make([]webrtc.TrackLocal, 0), - subscribedTracks: make(map[string]*webrtc.TrackRemote), + subscribedTracks: make(map[string][]*webrtc.TrackRemote), remoteParticipants: make(map[string]*livekit.ParticipantInfo), reader: logRing, writer: logRing, @@ -158,6 +158,10 @@ func NewRTCClient(conn *websocket.Conn) (*RTCClient, error) { return c, nil } +func (c *RTCClient) ID() string { + return c.id +} + // create an offer for the server func (c *RTCClient) Run() error { go c.logLoop() @@ -186,6 +190,8 @@ func (c *RTCClient) Run() error { } switch msg := res.Message.(type) { case *livekit.SignalResponse_Join: + c.id = msg.Join.Participant.Sid + c.lock.Lock() for _, p := range msg.Join.OtherParticipants { c.remoteParticipants[p.Sid] = p @@ -256,7 +262,7 @@ func (c *RTCClient) Run() error { } func (c *RTCClient) WaitUntilConnected() error { - ctx, _ := context.WithTimeout(context.Background(), 2*time.Second) + ctx, _ := context.WithTimeout(context.Background(), 5*time.Second) for { select { case <-ctx.Done(): @@ -300,7 +306,7 @@ func (c *RTCClient) ReadResponse() (*livekit.SignalResponse, error) { } } -func (c *RTCClient) SubscribedTracks() map[string]*webrtc.TrackRemote { +func (c *RTCClient) SubscribedTracks() map[string][]*webrtc.TrackRemote { return c.subscribedTracks } @@ -309,6 +315,8 @@ func (c *RTCClient) RemoteParticipants() []*livekit.ParticipantInfo { } func (c *RTCClient) Stop() { + c.connected = false + c.iceConnected = false c.conn.Close() c.cancel() } @@ -340,11 +348,50 @@ func (c *RTCClient) SendIceCandidate(ic *webrtc.ICECandidate) error { }) } -func (c *RTCClient) AddTrack(path string, id string, label string) error { +func (c *RTCClient) AddTrack(track *webrtc.TrackLocalStaticSample, path string) (writer *TrackWriter, err error) { + trackType := livekit.TrackType_AUDIO + if track.Kind() == webrtc.RTPCodecTypeVideo { + trackType = livekit.TrackType_VIDEO + } + + if err = c.SendAddTrack(track.ID(), track.StreamID(), trackType); err != nil { + return + } + + c.lock.Lock() + defer c.lock.Unlock() + c.localTracks = append(c.localTracks, track) + + if _, err = c.PeerConn.AddTrack(track); err != nil { + return + } + + writer = NewTrackWriter(c.ctx, track, path) + + // write tracks only after ICE connectivity + if c.iceConnected { + err = writer.Start() + } else { + c.pendingTrackWriters = append(c.pendingTrackWriters, writer) + } + + return +} + +func (c *RTCClient) AddStaticTrack(mime string, id string, label string) (writer *TrackWriter, err error) { + track, err := webrtc.NewTrackLocalStaticSample(webrtc.RTPCodecCapability{MimeType: mime}, id, label) + if err != nil { + return + } + + return c.AddTrack(track, "") +} + +func (c *RTCClient) AddFileTrack(path string, id string, label string) (writer *TrackWriter, err error) { // determine file mime mime, ok := extMimeMapping[filepath.Ext(path)] if !ok { - return fmt.Errorf("%s has an unsupported extension", filepath.Base(path)) + return nil, fmt.Errorf("%s has an unsupported extension", filepath.Base(path)) } c.AppendLog("adding track", @@ -357,37 +404,10 @@ func (c *RTCClient) AddTrack(path string, id string, label string) error { label, ) if err != nil { - return err + return } - trackType := livekit.TrackType_AUDIO - if strings.HasPrefix(mime, "video") { - trackType = livekit.TrackType_VIDEO - } - - if err := c.SendAddTrack(id, label, trackType); err != nil { - return err - } - - c.lock.Lock() - defer c.lock.Unlock() - c.localTracks = append(c.localTracks, track) - - if _, err := c.PeerConn.AddTrack(track); err != nil { - return err - } - - tw := NewTrackWriter(c.ctx, track, path) - - // write tracks only after ICE connectivity - if c.iceConnected { - return tw.Start() - } else { - c.pendingTrackWriters = append(c.pendingTrackWriters, tw) - return nil - } - - return nil + return c.AddTrack(track, path) } // send AddTrack command to server to initiate server-side negotiation @@ -516,7 +536,17 @@ func (c *RTCClient) logLoop() { func (c *RTCClient) processTrack(track *webrtc.TrackRemote) { lastUpdate := time.Time{} - peerId, trackId := rtc.UnpackTrackId(track.ID()) + pId, trackId := rtc.UnpackTrackId(track.ID()) + c.lock.Lock() + c.subscribedTracks[pId] = append(c.subscribedTracks[pId], track) + c.lock.Unlock() + + defer func() { + c.lock.Lock() + c.subscribedTracks[pId] = funk.Without(c.subscribedTracks[pId], track).([]*webrtc.TrackRemote) + c.lock.Unlock() + }() + numBytes := 0 for { pkt, _, err := track.ReadRTP() @@ -532,8 +562,8 @@ func (c *RTCClient) processTrack(track *webrtc.TrackRemote) { } numBytes += pkt.MarshalSize() if time.Now().Sub(lastUpdate) > 30*time.Second { - c.AppendLog("consumed from peer", - "track", trackId, "peer", peerId, + c.AppendLog("consumed from participant", + "track", trackId, "participant", pId, "size", numBytes) lastUpdate = time.Now() } diff --git a/cmd/cli/client/trackwriter.go b/cmd/cli/client/trackwriter.go index e568d628e..27114f9bf 100644 --- a/cmd/cli/client/trackwriter.go +++ b/cmd/cli/client/trackwriter.go @@ -19,9 +19,11 @@ import ( // makes it easier to debug and create RTP streams type TrackWriter struct { ctx context.Context + cancel context.CancelFunc track *webrtc.TrackLocalStaticSample filePath string mime string + done chan bool ogg *oggreader.OggReader ivfheader *ivfreader.IVFFileHeader @@ -30,15 +32,23 @@ type TrackWriter struct { } func NewTrackWriter(ctx context.Context, track *webrtc.TrackLocalStaticSample, filePath string) *TrackWriter { + ctx, cancel := context.WithCancel(ctx) return &TrackWriter{ ctx: ctx, + cancel: cancel, track: track, filePath: filePath, mime: track.Codec().MimeType, + done: make(chan bool), } } func (w *TrackWriter) Start() error { + if w.filePath == "" { + go w.writeNull() + return nil + } + file, err := os.Open(w.filePath) if err != nil { return err @@ -70,6 +80,23 @@ func (w *TrackWriter) Start() error { return nil } +func (w *TrackWriter) Stop() { + w.cancel() +} + +func (w *TrackWriter) writeNull() { + defer w.onWriteComplete() + sample := media.Sample{Data: []byte{0x0, 0xff, 0xff, 0xff, 0xff}, Duration: time.Second} + for { + select { + case <-time.After(20 * time.Millisecond): + w.track.WriteSample(sample) + case <-w.ctx.Done(): + break + } + } +} + func (w *TrackWriter) writeOgg() { // Keep track of last granule, the difference is the amount of samples in the buffer var lastGranule uint64 diff --git a/cmd/cli/commands/rtc.go b/cmd/cli/commands/rtc.go index 41ba793d8..61b636515 100644 --- a/cmd/cli/commands/rtc.go +++ b/cmd/cli/commands/rtc.go @@ -97,10 +97,10 @@ func joinRoom(c *cli.Context) error { rc.OnConnected = func() { // add after connection, since we need proper publish track APIs if audioFile != "" { - rc.AddTrack(audioFile, "audio", filepath.Base(audioFile)) + rc.AddFileTrack(audioFile, "audio", filepath.Base(audioFile)) } if videoFile != "" { - rc.AddTrack(videoFile, "video", filepath.Base(videoFile)) + rc.AddFileTrack(videoFile, "video", filepath.Base(videoFile)) } } @@ -188,7 +188,7 @@ func handleAddMedia(rc *client.RTCClient, isAudio bool) error { mediaPath = ExpandUser(mediaPath) // TODO: see what the ID should be - err = rc.AddTrack(mediaPath, codecType.String(), filepath.Base(mediaPath)) + _, err = rc.AddFileTrack(mediaPath, codecType.String(), filepath.Base(mediaPath)) if err != nil { return err } @@ -206,7 +206,7 @@ func handleAddMedia(rc *client.RTCClient, isAudio bool) error { audioPath := mediaPath[0:len(mediaPath)-len(videoExt)] + ".ogg" if _, err = os.Stat(audioPath); err == nil { - err = rc.AddTrack(audioPath, codecType.String(), filepath.Base(audioPath)) + _, err = rc.AddFileTrack(audioPath, codecType.String(), filepath.Base(audioPath)) if err != nil { fmt.Printf("added audio track: %s\n", audioPath) } diff --git a/test/integration_helpers.go b/test/integration_helpers.go index 86242e8a4..c84a76e5f 100644 --- a/test/integration_helpers.go +++ b/test/integration_helpers.go @@ -4,13 +4,15 @@ import ( "context" "fmt" "net/http" + "sync" "testing" "time" + "github.com/stretchr/testify/assert" + "github.com/livekit/livekit-server/cmd/cli/client" "github.com/livekit/livekit-server/pkg/auth" "github.com/livekit/livekit-server/pkg/config" - "github.com/livekit/livekit-server/pkg/logger" "github.com/livekit/livekit-server/pkg/service" "github.com/livekit/livekit-server/proto/livekit" ) @@ -55,6 +57,19 @@ func withTimeout(t *testing.T, f func() bool) { } } +func waitUntilConnected(t *testing.T, clients ...*client.RTCClient) { + wg := sync.WaitGroup{} + for i := range clients { + c := clients[i] + wg.Add(1) + go func() { + assert.NoError(t, c.WaitUntilConnected()) + wg.Done() + }() + } + wg.Wait() +} + func createServer() *service.LivekitServer { var err error serverConfig, err = config.NewConfig("") @@ -119,7 +134,6 @@ func (p *StaticKeyProvider) NumKeys() int { func (p *StaticKeyProvider) GetSecret(key string) string { if key == testApiKey { - logger.GetLogger().Debugf("returning secret: %s", testApiSecret) return testApiSecret } return "" diff --git a/test/integration_test.go b/test/integration_test.go index af3ec66fe..e100b33eb 100644 --- a/test/integration_test.go +++ b/test/integration_test.go @@ -16,13 +16,60 @@ import ( func TestClientCouldConnect(t *testing.T) { c1 := createClient("c1") - assert.NoError(t, c1.WaitUntilConnected()) c2 := createClient("c2") - assert.NoError(t, c2.WaitUntilConnected()) + waitUntilConnected(t, c1, c2) // ensure they both see each other withTimeout(t, func() bool { - return len(c1.RemoteParticipants()) == 1 && len(c2.RemoteParticipants()) == 1 + if len(c1.RemoteParticipants()) == 0 || len(c2.RemoteParticipants()) == 0 { + return false + } + //assert.Equal() + return true + }) +} + +func TestSinglePublisher(t *testing.T) { + c1 := createClient("c1") + c2 := createClient("c2") + waitUntilConnected(t, c1, c2) + + // publish a track and ensure clients receive it ok + t1, err := c1.AddStaticTrack("audio/opus", "audio", "webcam") + assert.NoError(t, err) + defer t1.Stop() + t2, err := c1.AddStaticTrack("video/vp8", "video", "webcam") + assert.NoError(t, err) + defer t2.Stop() + + // a new client joins and should get the initial stream + c3 := createClient("c3") + + withTimeout(t, func() bool { + if len(c2.SubscribedTracks()) == 0 { + return false + } + // should have received two tracks + if len(c2.SubscribedTracks()[c1.ID()]) != 2 { + return false + } + + tr1 := c2.SubscribedTracks()[c1.ID()][0] + assert.Equal(t, "webcam", tr1.StreamID()) + return true + }) + + // ensure that new client that has joined also received tracks + waitUntilConnected(t, c3) + withTimeout(t, func() bool { + if len(c3.SubscribedTracks()) == 0 { + return false + } + // should have received two tracks + if len(c3.SubscribedTracks()[c1.ID()]) != 2 { + return false + } + return true }) }