Prevent concurrent access of stats worker map (#666)

pkg/telemetry/telemetryserviceinternal.go

@@ -2,6 +2,7 @@ package telemetry
 
 import (
 	"context"
+	"sync"
 
 	"github.com/gammazero/workerpool"
 
@@ -27,7 +28,9 @@ type telemetryServiceInternal struct {
 	webhookPool *workerpool.WorkerPool
 
 	// one worker per participant
-	workers map[livekit.ParticipantID]*StatsWorker
+	workersMu  sync.RWMutex
+	workers    []*StatsWorker
+	workersIdx map[livekit.ParticipantID]int
 
 	analytics AnalyticsService
 }
@@ -36,7 +39,7 @@ func NewTelemetryServiceInternal(notifier webhook.Notifier, analytics AnalyticsS
 	return &telemetryServiceInternal{
 		notifier:    notifier,
 		webhookPool: workerpool.New(maxWebhookWorkers),
-		workers:     make(map[livekit.ParticipantID]*StatsWorker),
+		workersIdx:  make(map[livekit.ParticipantID]int),
 		analytics:   analytics,
 	}
 }
@@ -57,8 +60,7 @@ func (t *telemetryServiceInternal) TrackStats(streamType livekit.StreamType, par
 	}
 	prometheus.IncrementRTCP(direction, nacks, plis, firs)
 
-	w := t.workers[participantID]
-	if w != nil {
+	if w := t.getStatsWorker(participantID); w != nil {
 		w.OnTrackStat(trackID, streamType, stat)
 	}
 }
@@ -88,7 +90,24 @@ func (t *telemetryServiceInternal) Report(ctx context.Context, stats []*livekit.
 }
 
 func (t *telemetryServiceInternal) SendAnalytics() {
-	for _, worker := range t.workers {
-		worker.Update()
+	t.workersMu.RLock()
+	workers := t.workers
+	t.workersMu.RUnlock()
+
+	for _, worker := range workers {
+		if worker != nil {
+			worker.Update()
+		}
 	}
 }
+
+func (t *telemetryServiceInternal) getStatsWorker(participantID livekit.ParticipantID) *StatsWorker {
+	t.workersMu.RLock()
+	defer t.workersMu.RUnlock()
+
+	if idx, ok := t.workersIdx[participantID]; ok {
+		return t.workers[idx]
+	}
+
+	return nil
+}

pkg/telemetry/telemetryserviceinternalevents.go

@@ -45,9 +45,14 @@ func (t *telemetryServiceInternal) RoomEnded(ctx context.Context, room *livekit.
 	})
 }
 
-func (t *telemetryServiceInternal) ParticipantJoined(ctx context.Context, room *livekit.Room,
-	participant *livekit.ParticipantInfo, clientInfo *livekit.ClientInfo, clientMeta *livekit.AnalyticsClientMeta) {
-	t.workers[livekit.ParticipantID(participant.Sid)] = newStatsWorker(
+func (t *telemetryServiceInternal) ParticipantJoined(
+	ctx context.Context,
+	room *livekit.Room,
+	participant *livekit.ParticipantInfo,
+	clientInfo *livekit.ClientInfo,
+	clientMeta *livekit.AnalyticsClientMeta,
+) {
+	w := newStatsWorker(
 		ctx,
 		t,
 		livekit.RoomID(room.Sid),
@@ -55,6 +60,24 @@ func (t *telemetryServiceInternal) ParticipantJoined(ctx context.Context, room *
 		livekit.ParticipantID(participant.Sid),
 		livekit.ParticipantIdentity(participant.Identity),
 	)
+	t.workersMu.Lock()
+	var free = false
+	for idx, w := range t.workers {
+		if w != nil {
+			continue
+		}
+
+		free = true
+		t.workersIdx[livekit.ParticipantID(participant.Sid)] = idx
+		t.workers[idx] = w
+		break
+	}
+
+	if !free {
+		t.workersIdx[livekit.ParticipantID(participant.Sid)] = len(t.workers)
+		t.workers = append(t.workers, w)
+	}
+	t.workersMu.Unlock()
 
 	prometheus.AddParticipant()
 
@@ -89,11 +112,18 @@ func (t *telemetryServiceInternal) ParticipantActive(ctx context.Context, room *
 }
 
 func (t *telemetryServiceInternal) ParticipantLeft(ctx context.Context, room *livekit.Room, participant *livekit.ParticipantInfo) {
-	if w := t.workers[livekit.ParticipantID(participant.Sid)]; w != nil {
+	w := t.getStatsWorker(livekit.ParticipantID(participant.Sid))
+	if w != nil {
 		w.Close()
-		delete(t.workers, livekit.ParticipantID(participant.Sid))
 	}
 
+	t.workersMu.Lock()
+	if idx, ok := t.workersIdx[livekit.ParticipantID(participant.Sid)]; ok {
+		delete(t.workersIdx, livekit.ParticipantID(participant.Sid))
+		t.workers[idx] = nil
+	}
+	t.workersMu.Unlock()
+
 	prometheus.SubParticipant()
 
 	t.notifyEvent(ctx, &livekit.WebhookEvent{
@@ -169,7 +199,7 @@ func (t *telemetryServiceInternal) TrackMaxSubscribedVideoQuality(ctx context.Co
 func (t *telemetryServiceInternal) TrackUnpublished(ctx context.Context, participantID livekit.ParticipantID, identity livekit.ParticipantIdentity, track *livekit.TrackInfo, ssrc uint32) {
 	roomID := livekit.RoomID("")
 	roomName := livekit.RoomName("")
-	w := t.workers[participantID]
+	w := t.getStatsWorker(participantID)
 	if w != nil {
 		roomID = w.roomID
 		roomName = w.roomName
@@ -259,8 +289,7 @@ func (t *telemetryServiceInternal) RecordingEnded(ctx context.Context, ri *livek
 }
 
 func (t *telemetryServiceInternal) getRoomDetails(participantID livekit.ParticipantID) (livekit.RoomID, livekit.RoomName) {
-	w := t.workers[participantID]
-	if w != nil {
+	if w := t.getStatsWorker(participantID); w != nil {
 		return w.roomID, w.roomName
 	}
 	return "", ""