single publisher integration test

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()
 		}

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

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)
 		}

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 ""

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
 	})
 }