diff --git a/pkg/rtc/migrationdatacache.go b/pkg/rtc/migrationdatacache.go index 8c44eac61..45a689b7f 100644 --- a/pkg/rtc/migrationdatacache.go +++ b/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 +} diff --git a/pkg/rtc/migrationdatacache_test.go b/pkg/rtc/migrationdatacache_test.go index 046b2ddfc..a753c75f9 100644 --- a/pkg/rtc/migrationdatacache_test.go +++ b/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)) +} diff --git a/pkg/rtc/participant.go b/pkg/rtc/participant.go index daa68dde4..28ffccbc5 100644 --- a/pkg/rtc/participant.go +++ b/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, ) }