sfu/buffer: stop allocating per ingested packet (#4951)

* sfu/buffer: stop allocating per ingested packet

Buffer.Write allocated six objects for every RTP packet. They were
the parsed rtp.Packet, a copy of it in patchExtPacket, and the flow
state, whose address two rare log calls took. The rest were an
AbsCaptureTime, present on about one packet per second, and the
extension slice that pion grows with append.

Unmarshal into a pooled scratch packet. Give ExtPacket its own backing
rtp.Packet and an [8]rtp.Extension array, and copy the header into
them. Patch the payload in place. If the abs-capture-time extension is
present, parse it. Log a copy of the flow state.

Per packet, Write plus ReadExtended. Opus: 6 allocs, 448 B -> 0.
VP8: 6 allocs, 416 B -> 1 alloc, 32 B. Opus 240 -> 142 ns, VP8 269 -> 204 ns.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* sfu/buffer: point extension payloads at the bucket copy on read

Queued ExtPackets kept extension payload slices into the write input,
which the transport reuses for the next packet. ReadExtended now parses
the header again from the bucket copy, so extensions and CSRCs point
into RawPacket like the payload does.

The test writes every packet through one input buffer and checks the
audio level of each packet read back.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* sfu/buffer: test that a dropped padding packet keeps the shifted sequence number on read

The bucket stores the shifted number, so re-parsing the header on read
returns it. This test guards that.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

---------

Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>

pkg/sfu/buffer/buffer.go

@@ -157,7 +157,8 @@ func (b *Buffer) Bind(params webrtc.RTPParameters, codec webrtc.RTPCodecCapabili
 
 // Write adds an RTP Packet, ordering is not guaranteed, newer packets may arrive later
 func (b *Buffer) Write(pkt []byte) (n int, err error) {
-	var rtpPacket rtp.Packet
+	rtpPacket := rtpPacketFactory.Get().(*rtp.Packet)
+	defer rtpPacketFactory.Put(rtpPacket)
 	err = rtpPacket.Unmarshal(pkt)
 	if err != nil {
 		return
@@ -185,7 +186,7 @@ func (b *Buffer) Write(pkt []byte) (n int, err error) {
 	}
 
 	if b.streamInfoProbe != nil {
-		b.probeStreamInfoLocked(&rtpPacket)
+		b.probeStreamInfoLocked(rtpPacket)
 	}
 
 	// handle RTX packet
@@ -197,7 +198,7 @@ func (b *Buffer) Write(pkt []byte) (n int, err error) {
 			return
 		}
 
-		pb.writeRTX(&rtpPacket, now)
+		pb.writeRTX(rtpPacket, now)
 		return
 	}
 
@@ -235,7 +236,7 @@ func (b *Buffer) Write(pkt []byte) (n int, err error) {
 		return
 	}
 
-	rtcpPackets := b.calc(pkt, &rtpPacket, now, false, false)
+	rtcpPackets := b.calc(pkt, rtpPacket, now, false, false)
 	b.Unlock()
 
 	if len(rtcpPackets) != 0 {

pkg/sfu/buffer/buffer_base.go

@@ -18,6 +18,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"slices"
 	"strings"
 	"sync"
 	"time"
@@ -50,6 +51,13 @@ var (
 			return &ExtPacket{}
 		},
 	}
+
+	// unmarshal scratch, a reused packet keeps its Extensions capacity
+	rtpPacketFactory = &sync.Pool{
+		New: func() any {
+			return &rtp.Packet{}
+		},
+	}
 )
 
 func ReleaseExtPacket(extPkt *ExtPacket) {
@@ -78,6 +86,11 @@ type ExtPacket struct {
 	AbsCaptureTimeExt    *act.AbsCaptureTime
 	IsOutOfOrder         bool
 	IsBuffered           bool
+
+	// backing storage for Packet when the buffer creates the ext packet,
+	// extension and CSRC slices point into the write input until ReadExtended re-points them into RawPacket
+	pkt        rtp.Packet
+	extensions [8]rtp.Extension
 }
 
 type BufferProvider interface {
@@ -733,7 +746,8 @@ func (b *BufferBase) HandleIncomingPacketLocked(
 	oobSequenceNumber uint16,
 ) (uint64, error) {
 	if rtpPacket == nil {
-		rtpPacket = &rtp.Packet{}
+		rtpPacket = rtpPacketFactory.Get().(*rtp.Packet)
+		defer rtpPacketFactory.Put(rtpPacket)
 		if err := rtpPacket.Unmarshal(rawPkt); err != nil {
 			b.logger.Errorw("could not unmarshal RTP packet", err)
 			return 0, err
@@ -877,10 +891,12 @@ func (b *BufferBase) HandleIncomingPacketLocked(
 			if errors.Is(err, bucket.ErrPacketTooOld) {
 				packetTooOldCount := b.packetTooOldCount.Inc()
 				if (packetTooOldCount-1)%100 == 0 {
+					// log a copy so flowState stays on the stack
+					fs := flowState
 					b.logger.Warnw(
 						"could not add packet to bucket", err,
 						"count", packetTooOldCount,
-						"flowState", &flowState,
+						"flowState", &fs,
 						"snAdjustment", snAdjustment,
 						"incomingSequenceNumber", flowState.ExtSequenceNumber+snAdjustment,
 						"rtpStats", b.rtpStats,
@@ -890,9 +906,10 @@ func (b *BufferBase) HandleIncomingPacketLocked(
 					)
 				}
 			} else if err != bucket.ErrRTXPacket {
+				fs := flowState
 				b.logger.Warnw(
 					"could not add packet to bucket", err,
-					"flowState", &flowState,
+					"flowState", &fs,
 					"snAdjustment", snAdjustment,
 					"incomingSequenceNumber", flowState.ExtSequenceNumber+snAdjustment,
 					"rtpStats", b.rtpStats,
@@ -1025,7 +1042,7 @@ func (b *BufferBase) getExtPacket(
 		Arrival:           arrivalTime,
 		ExtSequenceNumber: flowState.ExtSequenceNumber,
 		ExtTimestamp:      flowState.ExtTimestamp,
-		Packet:            rtpPacket,
+		pkt:               *rtpPacket,
 		VideoLayer: VideoLayer{
 			Spatial:  InvalidLayerSpatial,
 			Temporal: InvalidLayerTemporal,
@@ -1033,6 +1050,10 @@ func (b *BufferBase) getExtPacket(
 		IsOutOfOrder: flowState.IsOutOfOrder,
 		IsBuffered:   isBuffered,
 	}
+	// rtpPacket may be reused scratch, copy what it backs
+	ep.pkt.Extensions = append(ep.extensions[:0], rtpPacket.Extensions...)
+	ep.pkt.CSRC = slices.Clone(rtpPacket.CSRC)
+	ep.Packet = &ep.pkt
 
 	if len(ep.Packet.Payload) == 0 {
 		// padding only packet, nothing else to do
@@ -1045,11 +1066,11 @@ func (b *BufferBase) getExtPacket(
 	}
 
 	if b.absCaptureTimeExtID != 0 {
-		extData := rtpPacket.GetExtension(b.absCaptureTimeExtID)
-
-		var actExt act.AbsCaptureTime
-		if err := actExt.Unmarshal(extData); err == nil {
-			ep.AbsCaptureTimeExt = &actExt
+		if extData := rtpPacket.GetExtension(b.absCaptureTimeExtID); len(extData) != 0 {
+			actExt := &act.AbsCaptureTime{}
+			if err := actExt.Unmarshal(extData); err == nil {
+				ep.AbsCaptureTimeExt = actExt
+			}
 		}
 	}
 
@@ -1201,16 +1222,19 @@ func (b *BufferBase) patchExtPacket(ep *ExtPacket, buf []byte) *ExtPacket {
 	}
 	ep.RawPacket = buf[:n]
 
-	// patch RTP packet to point payload to new buffer
-	pkt := *ep.Packet
-	payloadStart := ep.Packet.Header.MarshalSize()
+	// parse the header again from the new buffer so extension and CSRC slices point into it,
+	// then point the payload there too, ep owns its rtp.Packet
+	payloadStart, err := ep.Packet.Header.Unmarshal(buf[:n])
+	if err != nil {
+		b.logger.Warnw("could not unmarshal header", err, "sn", ep.Packet.SequenceNumber)
+		return nil
+	}
 	payloadEnd := payloadStart + len(ep.Packet.Payload)
 	if payloadEnd > n {
 		b.logger.Warnw("unexpected marshal size", nil, "max", n, "need", payloadEnd)
 		return nil
 	}
-	pkt.Payload = buf[payloadStart:payloadEnd]
-	ep.Packet = &pkt
+	ep.Packet.Payload = buf[payloadStart:payloadEnd]
 
 	return ep
 }

pkg/sfu/buffer/buffer_test.go

@@ -15,6 +15,7 @@
 package buffer
 
 import (
+	"encoding/binary"
 	"fmt"
 	"math"
 	"sync"
@@ -23,10 +24,14 @@ import (
 
 	"github.com/pion/rtcp"
 	"github.com/pion/rtp"
+	"github.com/pion/sdp/v3"
 	"github.com/pion/webrtc/v4"
 	"github.com/stretchr/testify/require"
 
 	"github.com/livekit/mediatransportutil/pkg/nack"
+
+	act "github.com/livekit/livekit-server/pkg/sfu/rtpextension/abscapturetime"
+	"github.com/livekit/livekit-server/pkg/sfu/utils"
 )
 
 var h265Codec = webrtc.RTPCodecParameters{
@@ -460,6 +465,138 @@ func TestCodecChange(t *testing.T) {
 	}
 }
 
+// Write unmarshals into reused scratch and callers reuse the input buffer, so each ExtPacket
+// must keep its own header and extensions, and Write plus ReadExtended must not allocate per packet
+func TestIngestNoAlloc(t *testing.T) {
+	testCases := []struct {
+		name      string
+		codec     webrtc.RTPCodecParameters
+		payload   []byte
+		maxAllocs float64
+	}{
+		{"opus", opusCodec, make([]byte, 100), 0},
+		// VP8 boxes the parsed descriptor into ExtPacket.Payload
+		{"vp8", vp8Codec, []byte{0x90, 0xe0, 0x80, 0x01, 0x05, 0x00, 0x01}, 1},
+	}
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			buff := NewBuffer(123, 100, 100)
+			require.NoError(t, buff.Bind(webrtc.RTPParameters{
+				HeaderExtensions: []webrtc.RTPHeaderExtensionParameter{
+					{URI: sdp.AudioLevelURI, ID: 1},
+					{URI: act.AbsCaptureTimeURI, ID: 2},
+				},
+				Codecs: []webrtc.RTPCodecParameters{tc.codec},
+			}, tc.codec.RTPCodecCapability, 0))
+			defer buff.Close()
+
+			marshal := func(sn uint16, level byte) []byte {
+				p := rtp.Packet{
+					Header: rtp.Header{
+						Version:        2,
+						PayloadType:    uint8(tc.codec.PayloadType),
+						SequenceNumber: sn,
+						Timestamp:      uint32(sn) * 960,
+						SSRC:           123,
+					},
+					Payload: tc.payload,
+				}
+				require.NoError(t, p.SetExtension(1, []byte{level}))
+				require.NoError(t, p.SetExtension(3, []byte{0, 0}))
+				raw, err := p.Marshal()
+				require.NoError(t, err)
+				return raw
+			}
+
+			// one input buffer for all writes, like the transport read loop
+			input := make([]byte, 1500)
+			write := func(sn uint16, level byte) {
+				n := copy(input, marshal(sn, level))
+				_, err := buff.Write(input[:n])
+				require.NoError(t, err)
+			}
+
+			readBuf := make([]byte, 1500)
+			write(1, 10)
+			write(2, 20)
+			for _, want := range []struct {
+				sn    uint16
+				level byte
+			}{{1, 10}, {2, 20}} {
+				ep, err := buff.ReadExtended(readBuf)
+				require.NoError(t, err)
+				require.Equal(t, want.sn, ep.Packet.SequenceNumber)
+				require.Equal(t, []byte{want.level}, ep.Packet.GetExtension(1))
+				require.Equal(t, tc.payload, ep.Packet.Payload)
+				require.Nil(t, ep.AbsCaptureTimeExt)
+				ReleaseExtPacket(ep)
+			}
+
+			raw := marshal(3, 30)
+			sn := uint16(3)
+			ingest := func() {
+				binary.BigEndian.PutUint16(raw[2:], sn)
+				binary.BigEndian.PutUint32(raw[4:], uint32(sn)*960)
+				sn++
+				_, err := buff.Write(raw)
+				require.NoError(t, err)
+				ep, err := buff.ReadExtended(readBuf)
+				require.NoError(t, err)
+				ReleaseExtPacket(ep)
+			}
+			allocs := testing.AllocsPerRun(1000, ingest)
+			if !utils.RaceEnabled {
+				require.LessOrEqual(t, allocs, tc.maxAllocs, "allocations per Write plus ReadExtended")
+			}
+		})
+	}
+}
+
+// a dropped padding only packet shifts later sequence numbers, ReadExtended must keep the shifted number
+// in the header and in RawPacket
+func TestIngestDroppedPaddingKeepsAdjustedSequenceNumber(t *testing.T) {
+	buff := NewBuffer(123, 100, 100)
+	require.NoError(t, buff.Bind(webrtc.RTPParameters{
+		Codecs: []webrtc.RTPCodecParameters{opusCodec},
+	}, opusCodec.RTPCodecCapability, 0))
+	defer buff.Close()
+
+	write := func(sn uint16, paddingOnly bool) {
+		p := rtp.Packet{
+			Header: rtp.Header{
+				Version:        2,
+				PayloadType:    uint8(opusCodec.PayloadType),
+				SequenceNumber: sn,
+				Timestamp:      uint32(sn) * 960,
+				SSRC:           123,
+			},
+		}
+		if paddingOnly {
+			p.Header.Padding = true
+			p.Header.PaddingSize = 20
+		} else {
+			p.Payload = make([]byte, 100)
+		}
+		raw, err := p.Marshal()
+		require.NoError(t, err)
+		_, err = buff.Write(raw)
+		require.NoError(t, err)
+	}
+
+	write(40, false)
+	write(41, true)
+	write(42, false)
+
+	readBuf := make([]byte, 1500)
+	for _, wantSN := range []uint16{40, 41} {
+		ep, err := buff.ReadExtended(readBuf)
+		require.NoError(t, err)
+		require.Equal(t, wantSN, ep.Packet.SequenceNumber)
+		require.Equal(t, wantSN, binary.BigEndian.Uint16(ep.RawPacket[2:]))
+		ReleaseExtPacket(ep)
+	}
+}
+
 func BenchmarkMemcpu(b *testing.B) {
 	buf := make([]byte, 1500*1500*10)
 	buf2 := make([]byte, 1500*1500*20)