From 31a054b79ddbcfb0d9fdc5103e8ee37bebc39522 Mon Sep 17 00:00:00 2001 From: Alessandro Ros Date: Tue, 18 Aug 2026 12:49:38 +0200 Subject: [PATCH] moq: impose maximum size on pending reordered bytes (#6112) Decrease the maximum memory that clients can take by imposing a maximum size of 100MB on the pending reordered bytes. --- internal/protocols/moq/reorderer/reorderer.go | 38 ++++++++++++---- .../protocols/moq/reorderer/reorderer_test.go | 43 ++++++++++++++++--- internal/servers/moq/inbound_track.go | 5 ++- internal/servers/moq/session.go | 5 ++- 4 files changed, 72 insertions(+), 19 deletions(-) diff --git a/internal/protocols/moq/reorderer/reorderer.go b/internal/protocols/moq/reorderer/reorderer.go index d09607e2..8ccc9b77 100644 --- a/internal/protocols/moq/reorderer/reorderer.go +++ b/internal/protocols/moq/reorderer/reorderer.go @@ -11,13 +11,23 @@ import ( // Reorderer is a subgroup reorderer. type Reorderer struct { - MaxReordered int - Parent logger.Writer + MaxReordered int + MaxPendingBytes int + Parent logger.Writer - initialized bool - mu sync.Mutex - curGroupID uint64 - pending map[uint64]*subgroup.SubGroup + initialized bool + mu sync.Mutex + curGroupID uint64 + pending map[uint64]*subgroup.SubGroup + pendingBytes int +} + +func subGroupPayloadSize(sg *subgroup.SubGroup) int { + n := 0 + for _, obj := range sg.Objects { + n += len(obj.Payload) + } + return n } // Initialize initializes the reorderer. @@ -45,7 +55,11 @@ func (r *Reorderer) Push(sg *subgroup.SubGroup) ([]*subgroup.SubGroup, error) { return []*subgroup.SubGroup{sg}, nil default: + if prev, ok := r.pending[sg.Header.GroupID]; ok { + r.pendingBytes -= subGroupPayloadSize(prev) + } r.pending[sg.Header.GroupID] = sg + r.pendingBytes += subGroupPayloadSize(sg) diff := sg.Header.GroupID - r.curGroupID @@ -56,11 +70,17 @@ func (r *Reorderer) Push(sg *subgroup.SubGroup) ([]*subgroup.SubGroup, error) { } } - if countInRange == diff { + switch { + case countInRange == diff: return r.flushUpTo(sg.Header.GroupID), nil - } else if len(r.pending) > r.MaxReordered { + + case len(r.pending) > r.MaxReordered: r.Parent.Log(logger.Warn, "too many reordered subgroups, flushing") return r.flushUpTo(sg.Header.GroupID), nil + + case r.pendingBytes > r.MaxPendingBytes: + r.Parent.Log(logger.Warn, "too many reordered bytes, flushing") + return r.flushUpTo(sg.Header.GroupID), nil } } @@ -79,6 +99,7 @@ func (r *Reorderer) flushUpTo(maxGroupID uint64) []*subgroup.SubGroup { out := make([]*subgroup.SubGroup, 0, len(ids)) for _, id := range ids { out = append(out, r.pending[id]) + r.pendingBytes -= subGroupPayloadSize(r.pending[id]) delete(r.pending, id) } @@ -90,6 +111,7 @@ func (r *Reorderer) flushUpTo(maxGroupID uint64) []*subgroup.SubGroup { break } out = append(out, next) + r.pendingBytes -= subGroupPayloadSize(next) delete(r.pending, r.curGroupID+1) r.curGroupID++ } diff --git a/internal/protocols/moq/reorderer/reorderer_test.go b/internal/protocols/moq/reorderer/reorderer_test.go index aab7693e..0936966a 100644 --- a/internal/protocols/moq/reorderer/reorderer_test.go +++ b/internal/protocols/moq/reorderer/reorderer_test.go @@ -20,8 +20,17 @@ func makeSG(groupID uint64) *subgroup.SubGroup { } } +func makeSGWithPayload(groupID uint64, size int) *subgroup.SubGroup { + return &subgroup.SubGroup{ + Header: subgroup.Header{GroupID: groupID}, + Objects: []subgroup.Object{{ + Payload: make([]byte, size), + }}, + } +} + func TestReordererFirstPush(t *testing.T) { - r := &reorderer.Reorderer{MaxReordered: 5, Parent: nopLogger{}} + r := &reorderer.Reorderer{MaxReordered: 5, MaxPendingBytes: 100000, Parent: nopLogger{}} r.Initialize() sg0 := makeSG(0) @@ -31,7 +40,7 @@ func TestReordererFirstPush(t *testing.T) { } func TestReordererInOrder(t *testing.T) { - r := &reorderer.Reorderer{MaxReordered: 5, Parent: nopLogger{}} + r := &reorderer.Reorderer{MaxReordered: 5, MaxPendingBytes: 100000, Parent: nopLogger{}} r.Initialize() sg0 := makeSG(0) @@ -51,7 +60,7 @@ func TestReordererInOrder(t *testing.T) { } func TestReordererStale(t *testing.T) { - r := &reorderer.Reorderer{MaxReordered: 5, Parent: nopLogger{}} + r := &reorderer.Reorderer{MaxReordered: 5, MaxPendingBytes: 100000, Parent: nopLogger{}} r.Initialize() sg5 := makeSG(5) @@ -71,7 +80,7 @@ func TestReordererStale(t *testing.T) { } func TestReordererOutOfOrderPending(t *testing.T) { - r := &reorderer.Reorderer{MaxReordered: 5, Parent: nopLogger{}} + r := &reorderer.Reorderer{MaxReordered: 5, MaxPendingBytes: 100000, Parent: nopLogger{}} r.Initialize() sg0 := makeSG(0) @@ -86,7 +95,7 @@ func TestReordererOutOfOrderPending(t *testing.T) { } func TestReordererOutOfOrderFilled(t *testing.T) { - r := &reorderer.Reorderer{MaxReordered: 5, Parent: nopLogger{}} + r := &reorderer.Reorderer{MaxReordered: 5, MaxPendingBytes: 100000, Parent: nopLogger{}} r.Initialize() sg0 := makeSG(0) @@ -109,7 +118,7 @@ func TestReordererOutOfOrderFilled(t *testing.T) { } func TestReordererDrainAfterFill(t *testing.T) { //nolint:dupl - r := &reorderer.Reorderer{MaxReordered: 5, Parent: nopLogger{}} + r := &reorderer.Reorderer{MaxReordered: 5, MaxPendingBytes: 100000, Parent: nopLogger{}} r.Initialize() sg0 := makeSG(0) @@ -136,7 +145,7 @@ func TestReordererDrainAfterFill(t *testing.T) { //nolint:dupl } func TestReordererMaxReordered(t *testing.T) { //nolint:dupl - r := &reorderer.Reorderer{MaxReordered: 2, Parent: nopLogger{}} + r := &reorderer.Reorderer{MaxReordered: 2, MaxPendingBytes: 100000, Parent: nopLogger{}} r.Initialize() sg0 := makeSG(0) @@ -163,3 +172,23 @@ func TestReordererMaxReordered(t *testing.T) { //nolint:dupl require.NoError(t, err) require.Equal(t, []*subgroup.SubGroup{sg2, sg4, sg5}, out) } + +func TestReordererMaxPendingBytes(t *testing.T) { + r := &reorderer.Reorderer{MaxReordered: 10, MaxPendingBytes: 100, Parent: nopLogger{}} + r.Initialize() + + sg0 := makeSG(0) + out, err := r.Push(sg0) + require.NoError(t, err) + require.Equal(t, []*subgroup.SubGroup{sg0}, out) + + sg2 := makeSGWithPayload(2, 60) + out, err = r.Push(sg2) + require.NoError(t, err) + require.Nil(t, out) + + sg4 := makeSGWithPayload(4, 60) + out, err = r.Push(sg4) + require.NoError(t, err) + require.Equal(t, []*subgroup.SubGroup{sg2, sg4}, out) +} diff --git a/internal/servers/moq/inbound_track.go b/internal/servers/moq/inbound_track.go index f80bfda3..ad03a7de 100644 --- a/internal/servers/moq/inbound_track.go +++ b/internal/servers/moq/inbound_track.go @@ -15,8 +15,9 @@ type inboundTrack struct { func (t *inboundTrack) initialize() { t.reorderer = &reorderer.Reorderer{ - MaxReordered: maxReorderedSubGroups, - Parent: t.parent, + MaxReordered: maxReorderedSubGroups, + MaxPendingBytes: maxReorderedPendingBytes, + Parent: t.parent, } t.reorderer.Initialize() } diff --git a/internal/servers/moq/session.go b/internal/servers/moq/session.go index 2114a798..1c02fcf2 100644 --- a/internal/servers/moq/session.go +++ b/internal/servers/moq/session.go @@ -34,8 +34,9 @@ import ( ) const ( - maxReorderedSubGroups = 50 - maxCatalogTracks = 50 + maxReorderedSubGroups = 50 + maxReorderedPendingBytes = 100 * 1024 * 1024 + maxCatalogTracks = 50 ) func findAuthorizationToken(parameters []parameter.Parameter) *parameter.AuthorizationToken {