Add bytes cap on migration data cache size (#4782)

pkg/rtc/migrationdatacache.go

@@ -14,9 +14,14 @@ const (
 	MigrationDataCacheStateDone
 )
 
+const (
+	migrationDataCacheMaxSize = 4 << 20 // 4 MiB
+)
+
 type MigrationDataCache struct {
 	lastSeq   uint32
 	pkts      []*livekit.DataPacket
+	size      int
 	state     MigrationDataCacheState
 	expiredAt time.Time
 }
@@ -30,10 +35,10 @@ func NewMigrationDataCache(lastSeq uint32, expiredAt time.Time) *MigrationDataCa
 
 // Add adds a message to the cache if there is a gap between the last sequence number and cached messages then return the cache State:
 //   - MigrationDataCacheStateWaiting: waiting for the next packet (lastSeq + 1) of last sequence from old node
-//   - MigrationDataCacheStateTimeout: the next packet is not received before the expiredAt, participant will
-//     continue to process the reliable messages, subscribers will see the gap after the publisher migration
+//   - MigrationDataCacheStateTimeout: the next packet is not received before the expiredAt or the cache is full, participant
+//     will continue to process the reliable messages, subscribers will see the gap after the publisher migration
 //   - MigrationDataCacheStateDone: the next packet is received, participant can continue to process the reliable messages
-func (c *MigrationDataCache) Add(pkt *livekit.DataPacket) MigrationDataCacheState {
+func (c *MigrationDataCache) Add(pkt *livekit.DataPacket, size int) MigrationDataCacheState {
 	if c.state == MigrationDataCacheStateDone || c.state == MigrationDataCacheStateTimeout {
 		return c.state
 	}
@@ -48,7 +53,8 @@ func (c *MigrationDataCache) Add(pkt *livekit.DataPacket) MigrationDataCacheStat
 	}
 
 	c.pkts = append(c.pkts, pkt)
-	if time.Now().After(c.expiredAt) {
+	c.size += size
+	if c.size >= migrationDataCacheMaxSize || time.Now().After(c.expiredAt) {
 		c.state = MigrationDataCacheStateTimeout
 	}
 	return c.state
@@ -57,3 +63,7 @@ func (c *MigrationDataCache) Add(pkt *livekit.DataPacket) MigrationDataCacheStat
 func (c *MigrationDataCache) Get() []*livekit.DataPacket {
 	return c.pkts
 }
+
+func (c *MigrationDataCache) Size() int {
+	return c.size
+}

pkg/rtc/migrationdatacache_test.go

@@ -14,25 +14,44 @@ func TestMigrationDataCache_Add(t *testing.T) {
 	cache := NewMigrationDataCache(10, expiredAt)
 
 	pkt1 := &livekit.DataPacket{Sequence: 9}
-	state := cache.Add(pkt1)
+	state := cache.Add(pkt1, 0)
 	require.Equal(t, MigrationDataCacheStateWaiting, state)
 	require.Empty(t, cache.Get())
 
 	pkt2 := &livekit.DataPacket{Sequence: 11}
-	state = cache.Add(pkt2)
+	state = cache.Add(pkt2, 0)
 	require.Equal(t, MigrationDataCacheStateDone, state)
 	require.Empty(t, cache.Get())
 
 	pkt3 := &livekit.DataPacket{Sequence: 12}
-	state = cache.Add(pkt3)
+	state = cache.Add(pkt3, 0)
 	require.Equal(t, MigrationDataCacheStateDone, state)
 	require.Empty(t, cache.Get())
 
 	cache2 := NewMigrationDataCache(20, time.Now().Add(10*time.Millisecond))
 	pkt4 := &livekit.DataPacket{Sequence: 22}
 	time.Sleep(20 * time.Millisecond)
-	state = cache2.Add(pkt4)
+	state = cache2.Add(pkt4, 0)
 	require.Equal(t, MigrationDataCacheStateTimeout, state)
 	require.Len(t, cache2.Get(), 1)
 	require.Equal(t, uint32(22), cache2.Get()[0].Sequence)
 }
+
+func TestMigrationDataCache_MaxSize(t *testing.T) {
+	// the cache should not grow past the size budget even if the expiry is far in the future
+	cache := NewMigrationDataCache(10, time.Now().Add(time.Minute))
+
+	pktSize := 1000
+	seq := uint32(12)
+	state := MigrationDataCacheStateWaiting
+	for ; state == MigrationDataCacheStateWaiting; seq++ {
+		state = cache.Add(&livekit.DataPacket{Sequence: seq}, pktSize)
+	}
+
+	require.Equal(t, MigrationDataCacheStateTimeout, state)
+	require.LessOrEqual(t, cache.Size(), migrationDataCacheMaxSize+pktSize)
+	require.Len(t, cache.Get(), migrationDataCacheMaxSize/pktSize+1)
+
+	// once full, further packets are dropped, including the continuous one
+	require.Equal(t, MigrationDataCacheStateTimeout, cache.Add(&livekit.DataPacket{Sequence: 11}, pktSize))
+}

pkg/rtc/participant.go

@@ -2447,7 +2447,7 @@ func (p *ParticipantImpl) onReceivedDataMessage(kind livekit.DataPacket_Kind, da
 		}
 
 		if migrationCache := p.reliableDataInfo.migrateInPubDataCache.Load(); migrationCache != nil {
-			switch migrationCache.Add(dp) {
+			switch migrationCache.Add(dp, len(data)) {
 			case MigrationDataCacheStateWaiting:
 				// waiting for the reliable sequence to continue from last node
 				return
@@ -2466,6 +2466,8 @@ func (p *ParticipantImpl) onReceivedDataMessage(kind livekit.DataPacket_Kind, da
 						"migration data cache timed out, handling cached messages", nil,
 						"cachedFirstSeq", cachedMsgs[0].Sequence,
 						"cachedLastSeq", cachedMsgs[len(cachedMsgs)-1].Sequence,
+						"cachedNum", len(cachedMsgs),
+						"cachedSize", migrationCache.Size(),
 						"lastPubReliableSeq", p.params.LastPubReliableSeq,
 					)
 				}