parallel writing for data packet broadcast (#1425)

pkg/rtc/room.go

@@ -5,6 +5,7 @@ import (
 	"errors"
 	"io"
 	"math"
+	"runtime"
 	"sort"
 	"sync"
 	"time"
@@ -30,6 +31,8 @@ const (
 	AudioLevelQuantization    = 8 // ideally power of 2 to minimize float decimal
 	invAudioLevelQuantization = 1.0 / AudioLevelQuantization
 	subscriberUpdateInterval  = 3 * time.Second
+
+	dataForwardLoadBalanceThreshold = 20
 )
 
 type broadcastOptions struct {
@@ -159,6 +162,10 @@ func (r *Room) GetParticipants() []types.LocalParticipant {
 	return participants
 }
 
+func (r *Room) GetLocalParticipants() []types.LocalParticipant {
+	return r.GetParticipants()
+}
+
 func (r *Room) GetActiveSpeakers() []*livekit.SpeakerInfo {
 	participants := r.GetParticipants()
 	speakers := make([]*livekit.SpeakerInfo, 0, len(participants))
@@ -810,42 +817,7 @@ func (r *Room) onParticipantUpdate(p types.LocalParticipant) {
 }
 
 func (r *Room) onDataPacket(source types.LocalParticipant, dp *livekit.DataPacket) {
-	dest := dp.GetUser().GetDestinationSids()
-	var dpData []byte
-
-	for _, op := range r.GetParticipants() {
-		if op.State() != livekit.ParticipantInfo_ACTIVE {
-			continue
-		}
-		if source != nil && op.ID() == source.ID() {
-			continue
-		}
-		if len(dest) > 0 {
-			found := false
-			for _, dID := range dest {
-				if op.ID() == livekit.ParticipantID(dID) {
-					found = true
-					break
-				}
-			}
-			if !found {
-				continue
-			}
-		}
-		if dpData == nil {
-			var err error
-			dpData, err = proto.Marshal(dp)
-			if err != nil {
-				r.Logger.Errorw("failed to marshal data packet", err)
-				return
-			}
-		}
-
-		err := op.SendDataPacket(dp, dpData)
-		if err != nil && !errors.Is(err, io.ErrClosedPipe) {
-			r.Logger.Infow("send data packet error", "error", err, "participant", op.Identity())
-		}
-	}
+	BroadcastDataPacketForRoom(r, source, dp, r.Logger)
 }
 
 func (r *Room) subscribeToExistingTracks(p types.LocalParticipant) {
@@ -1155,3 +1127,86 @@ func (r *Room) DebugInfo() map[string]interface{} {
 
 	return info
 }
+
+func BroadcastDataPacketForRoom(r types.Room, source types.LocalParticipant, dp *livekit.DataPacket, logger logger.Logger) {
+	dest := dp.GetUser().GetDestinationSids()
+	var dpData []byte
+
+	participants := r.GetLocalParticipants()
+	cap := len(dest)
+	if cap == 0 {
+		cap = len(participants)
+	}
+	destParticpants := make([]types.LocalParticipant, 0, cap)
+
+	for _, op := range participants {
+		if op.State() != livekit.ParticipantInfo_ACTIVE {
+			continue
+		}
+		if source != nil && op.ID() == source.ID() {
+			continue
+		}
+		if len(dest) > 0 {
+			found := false
+			for _, dID := range dest {
+				if op.ID() == livekit.ParticipantID(dID) {
+					found = true
+					break
+				}
+			}
+			if !found {
+				continue
+			}
+		}
+		if dpData == nil {
+			var err error
+			dpData, err = proto.Marshal(dp)
+			if err != nil {
+				logger.Errorw("failed to marshal data packet", err)
+				return
+			}
+		}
+		destParticpants = append(destParticpants, op)
+	}
+
+	if len(destParticpants) < dataForwardLoadBalanceThreshold {
+		for _, op := range destParticpants {
+			err := op.SendDataPacket(dp, dpData)
+			if err != nil && !errors.Is(err, io.ErrClosedPipe) {
+				logger.Infow("send data packet error", "error", err, "participant", op.Identity())
+			}
+		}
+		return
+	}
+
+	// parallel - enables much more efficient multi-core utilization
+	start := atomic.NewUint64(0)
+	end := uint64(len(destParticpants))
+
+	step := uint64(1)
+
+	var wg sync.WaitGroup
+	numCPU := runtime.NumCPU()
+	wg.Add(numCPU)
+	for p := 0; p < numCPU; p++ {
+		go func() {
+			defer wg.Done()
+			for {
+				n := start.Add(step)
+				if n >= end+step {
+					return
+				}
+
+				for i := n - step; i < n && i < end; i++ {
+					op := destParticpants[i]
+					err := op.SendDataPacket(dp, dpData)
+					if err != nil && !errors.Is(err, io.ErrClosedPipe) {
+						logger.Infow("send data packet error", "error", err, "participant", op.Identity())
+					}
+				}
+			}
+		}()
+	}
+	wg.Wait()
+
+}

pkg/rtc/types/interfaces.go

@@ -346,6 +346,7 @@ type Room interface {
 	SimulateScenario(participant LocalParticipant, scenario *livekit.SimulateScenario) error
 	UpdateVideoLayers(participant Participant, updateVideoLayers *livekit.UpdateVideoLayers) error
 	ResolveMediaTrackForSubscriber(subIdentity livekit.ParticipantIdentity, trackID livekit.TrackID) MediaResolverResult
+	GetLocalParticipants() []LocalParticipant
 }
 
 // MediaTrack represents a media track

pkg/rtc/types/typesfakes/fake_room.go

@@ -9,6 +9,16 @@ import (
 )
 
 type FakeRoom struct {
+	GetLocalParticipantsStub        func() []types.LocalParticipant
+	getLocalParticipantsMutex       sync.RWMutex
+	getLocalParticipantsArgsForCall []struct {
+	}
+	getLocalParticipantsReturns struct {
+		result1 []types.LocalParticipant
+	}
+	getLocalParticipantsReturnsOnCall map[int]struct {
+		result1 []types.LocalParticipant
+	}
 	IDStub        func() livekit.RoomID
 	iDMutex       sync.RWMutex
 	iDArgsForCall []struct {
@@ -108,6 +118,59 @@ type FakeRoom struct {
 	invocationsMutex sync.RWMutex
 }
 
+func (fake *FakeRoom) GetLocalParticipants() []types.LocalParticipant {
+	fake.getLocalParticipantsMutex.Lock()
+	ret, specificReturn := fake.getLocalParticipantsReturnsOnCall[len(fake.getLocalParticipantsArgsForCall)]
+	fake.getLocalParticipantsArgsForCall = append(fake.getLocalParticipantsArgsForCall, struct {
+	}{})
+	stub := fake.GetLocalParticipantsStub
+	fakeReturns := fake.getLocalParticipantsReturns
+	fake.recordInvocation("GetLocalParticipants", []interface{}{})
+	fake.getLocalParticipantsMutex.Unlock()
+	if stub != nil {
+		return stub()
+	}
+	if specificReturn {
+		return ret.result1
+	}
+	return fakeReturns.result1
+}
+
+func (fake *FakeRoom) GetLocalParticipantsCallCount() int {
+	fake.getLocalParticipantsMutex.RLock()
+	defer fake.getLocalParticipantsMutex.RUnlock()
+	return len(fake.getLocalParticipantsArgsForCall)
+}
+
+func (fake *FakeRoom) GetLocalParticipantsCalls(stub func() []types.LocalParticipant) {
+	fake.getLocalParticipantsMutex.Lock()
+	defer fake.getLocalParticipantsMutex.Unlock()
+	fake.GetLocalParticipantsStub = stub
+}
+
+func (fake *FakeRoom) GetLocalParticipantsReturns(result1 []types.LocalParticipant) {
+	fake.getLocalParticipantsMutex.Lock()
+	defer fake.getLocalParticipantsMutex.Unlock()
+	fake.GetLocalParticipantsStub = nil
+	fake.getLocalParticipantsReturns = struct {
+		result1 []types.LocalParticipant
+	}{result1}
+}
+
+func (fake *FakeRoom) GetLocalParticipantsReturnsOnCall(i int, result1 []types.LocalParticipant) {
+	fake.getLocalParticipantsMutex.Lock()
+	defer fake.getLocalParticipantsMutex.Unlock()
+	fake.GetLocalParticipantsStub = nil
+	if fake.getLocalParticipantsReturnsOnCall == nil {
+		fake.getLocalParticipantsReturnsOnCall = make(map[int]struct {
+			result1 []types.LocalParticipant
+		})
+	}
+	fake.getLocalParticipantsReturnsOnCall[i] = struct {
+		result1 []types.LocalParticipant
+	}{result1}
+}
+
 func (fake *FakeRoom) ID() livekit.RoomID {
 	fake.iDMutex.Lock()
 	ret, specificReturn := fake.iDReturnsOnCall[len(fake.iDArgsForCall)]
@@ -606,6 +669,8 @@ func (fake *FakeRoom) UpdateVideoLayersReturnsOnCall(i int, result1 error) {
 func (fake *FakeRoom) Invocations() map[string][][]interface{} {
 	fake.invocationsMutex.RLock()
 	defer fake.invocationsMutex.RUnlock()
+	fake.getLocalParticipantsMutex.RLock()
+	defer fake.getLocalParticipantsMutex.RUnlock()
 	fake.iDMutex.RLock()
 	defer fake.iDMutex.RUnlock()
 	fake.nameMutex.RLock()