Repository: pion/interceptor (go)

Issue:
NACK responder always responding with most recent packet
In the NACK responder, it tries to cache sent packets in order to respond to NACKs. However, it's storing a pointer to the rtp.Packet struct, which may be re-used by the sender process without the knowledge of NackResponder.

Adam from the community has observed symptoms where retransmitted packets contained payload that did not match the sequence number that was requested. (but a later sequence number that overrode the earlier one).

We need to be storing a copy of the packet to ensure that nack responder has full ownership over its data.

Pull request description:
Fix NACK interceptor sendBuffer overridden
Keeps a copy of packet in responder_interceptor, fixes #84

While trying to fix this, I was also thinking about how we might reduce copy/allocations throughout pion. Have implemented a version of this that works within intercepters (unexported for now). any feedback is welcome!

Patch:
diff --git a/AUTHORS.txt b/AUTHORS.txt
index ea21052..8f85a31 100644
--- a/AUTHORS.txt
+++ b/AUTHORS.txt
@@ -9,6 +9,7 @@ aler9 <46489434+aler9@users.noreply.github.com>
 Antoine Baché <antoine@tenten.app>
 Atsushi Watanabe <atsushi.w@ieee.org>
 boks1971 <raja.gobi@tutanota.com>
+David Zhao <david@davidzhao.com>
 Jonathan Müller <jonathan@fotokite.com>
 Mathis Engelbart <mathis.engelbart@gmail.com>
 Sean DuBois <sean@siobud.com>
diff --git a/pkg/nack/errors.go b/pkg/nack/errors.go
index bbfc773..66ed2b3 100644
--- a/pkg/nack/errors.go
+++ b/pkg/nack/errors.go
@@ -4,3 +4,5 @@ import "errors"
 
 // ErrInvalidSize is returned by newReceiveLog/newSendBuffer, when an incorrect buffer size is supplied.
 var ErrInvalidSize = errors.New("invalid buffer size")
+
+var errPacketReleased = errors.New("could not retain packet, already released")
diff --git a/pkg/nack/responder_interceptor.go b/pkg/nack/responder_interceptor.go
index 396efba..9b49f11 100644
--- a/pkg/nack/responder_interceptor.go
+++ b/pkg/nack/responder_interceptor.go
@@ -17,9 +17,10 @@ type ResponderInterceptorFactory struct {
 // NewInterceptor constructs a new ResponderInterceptor
 func (r *ResponderInterceptorFactory) NewInterceptor(id string) (interceptor.Interceptor, error) {
 	i := &ResponderInterceptor{
-		size:    8192,
-		log:     logging.NewDefaultLoggerFactory().NewLogger("nack_responder"),
-		streams: map[uint32]*localStream{},
+		size:      8192,
+		log:       logging.NewDefaultLoggerFactory().NewLogger("nack_responder"),
+		streams:   map[uint32]*localStream{},
+		packetMan: newPacketManager(),
 	}
 
 	for _, opt := range r.opts {
@@ -38,8 +39,9 @@ func (r *ResponderInterceptorFactory) NewInterceptor(id string) (interceptor.Int
 // ResponderInterceptor responds to nack feedback messages
 type ResponderInterceptor struct {
 	interceptor.NoOp
-	size uint16
-	log  logging.LeveledLogger
+	size      uint16
+	log       logging.LeveledLogger
+	packetMan *packetManager
 
 	streams   map[uint32]*localStream
 	streamsMu sync.Mutex
@@ -98,7 +100,11 @@ func (n *ResponderInterceptor) BindLocalStream(info *interceptor.StreamInfo, wri
 	n.streamsMu.Unlock()
 
 	return interceptor.RTPWriterFunc(func(header *rtp.Header, payload []byte, attributes interceptor.Attributes) (int, error) {
-		sendBuffer.add(&rtp.Packet{Header: *header, Payload: payload})
+		pkt, err := n.packetMan.NewPacket(header, payload)
+		if err != nil {
+			return 0, err
+		}
+		sendBuffer.add(pkt)
 		return writer.Write(header, payload, attributes)
 	})
 }
@@ -121,9 +127,10 @@ func (n *ResponderInterceptor) resendPackets(nack *rtcp.TransportLayerNack) {
 	for i := range nack.Nacks {
 		nack.Nacks[i].Range(func(seq uint16) bool {
 			if p := stream.sendBuffer.get(seq); p != nil {
-				if _, err := stream.rtpWriter.Write(&p.Header, p.Payload, interceptor.Attributes{}); err != nil {
+				if _, err := stream.rtpWriter.Write(p.Header(), p.Payload(), interceptor.Attributes{}); err != nil {
 					n.log.Warnf("failed resending nacked packet: %+v", err)
 				}
+				p.Release()
 			}
 
 			return true
diff --git a/pkg/nack/retainable_packet.go b/pkg/nack/retainable_packet.go
new file mode 100644
index 0000000..89e9cb6
--- /dev/null
+++ b/pkg/nack/retainable_packet.go
@@ -0,0 +1,105 @@
+package nack
+
+import (
+	"io"
+	"sync"
+
+	"github.com/pion/rtp"
+)
+
+const maxPayloadLen = 1460
+
+type packetManager struct {
+	headerPool  *sync.Pool
+	payloadPool *sync.Pool
+}
+
+func newPacketManager() *packetManager {
+	return &packetManager{
+		headerPool: &sync.Pool{
+			New: func() interface{} {
+				return &rtp.Header{}
+			},
+		},
+		payloadPool: &sync.Pool{
+			New: func() interface{} {
+				buf := make([]byte, maxPayloadLen)
+				return &buf
+			},
+		},
+	}
+}
+
+func (m *packetManager) NewPacket(header *rtp.Header, payload []byte) (*retainablePacket, error) {
+	if len(payload) > maxPayloadLen {
+		return nil, io.ErrShortBuffer
+	}
+
+	p := &retainablePacket{
+		onRelease: m.releasePacket,
+		// new packets have retain count of 1
+		count: 1,
+	}
+
+	p.header = m.headerPool.Get().(*rtp.Header)
+	*p.header = header.Clone()
+
+	if payload != nil {
+		p.buffer = m.payloadPool.Get().(*[]byte)
+		size := copy(*p.buffer, payload)
+		p.payload = (*p.buffer)[:size]
+	}
+
+	return p, nil
+}
+
+func (m *packetManager) releasePacket(header *rtp.Header, payload *[]byte) {
+	m.headerPool.Put(header)
+	if payload != nil {
+		m.payloadPool.Put(payload)
+	}
+}
+
+type retainablePacket struct {
+	onRelease func(*rtp.Header, *[]byte)
+
+	countMu sync.Mutex
+	count   int
+
+	header  *rtp.Header
+	buffer  *[]byte
+	payload []byte
+}
+
+func (p *retainablePacket) Header() *rtp.Header {
+	return p.header
+}
+
+func (p *retainablePacket) Payload() []byte {
+	return p.payload
+}
+
+func (p *retainablePacket) Retain() error {
+	p.countMu.Lock()
+	defer p.countMu.Unlock()
+	if p.count == 0 {
+		// already released
+		return errPacketReleased
+	}
+	p.count++
+	return nil
+}
+
+func (p *retainablePacket) Release() {
+	p.countMu.Lock()
+	defer p.countMu.Unlock()
+	p.count--
+
+	if p.count == 0 {
+		// release back to pool
+		p.onRelease(p.header, p.buffer)
+		p.header = nil
+		p.buffer = nil
+		p.payload = nil
+	}
+}
diff --git a/pkg/nack/send_buffer.go b/pkg/nack/send_buffer.go
index 95dbf65..8e36075 100644
--- a/pkg/nack/send_buffer.go
+++ b/pkg/nack/send_buffer.go
@@ -3,8 +3,6 @@ package nack
 import (
 	"fmt"
 	"sync"
-
-	"github.com/pion/rtp"
 )
 
 const (
@@ -12,7 +10,7 @@ const (
 )
 
 type sendBuffer struct {
-	packets   []*rtp.Packet
+	packets   []*retainablePacket
 	size      uint16
 	lastAdded uint16
 	started   bool
@@ -36,16 +34,16 @@ func newSendBuffer(size uint16) (*sendBuffer, error) {
 	}
 
 	return &sendBuffer{
-		packets: make([]*rtp.Packet, size),
+		packets: make([]*retainablePacket, size),
 		size:    size,
 	}, nil
 }
 
-func (s *sendBuffer) add(packet *rtp.Packet) {
+func (s *sendBuffer) add(packet *retainablePacket) {
 	s.m.Lock()
 	defer s.m.Unlock()
 
-	seq := packet.SequenceNumber
+	seq := packet.Header().SequenceNumber
 	if !s.started {
 		s.packets[seq%s.size] = packet
 		s.lastAdded = seq
@@ -58,15 +56,25 @@ func (s *sendBuffer) add(packet *rtp.Packet) {
 		return
 	} else if diff < uint16SizeHalf {
 		for i := s.lastAdded + 1; i != seq; i++ {
-			s.packets[i%s.size] = nil
+			idx := i % s.size
+			prevPacket := s.packets[idx]
+			if prevPacket != nil {
+				prevPacket.Release()
+			}
+			s.packets[idx] = nil
 		}
 	}
 
-	s.packets[seq%s.size] = packet
+	idx := seq % s.size
+	prevPacket := s.packets[idx]
+	if prevPacket != nil {
+		prevPacket.Release()
+	}
+	s.packets[idx] = packet
 	s.lastAdded = seq
 }
 
-func (s *sendBuffer) get(seq uint16) *rtp.Packet {
+func (s *sendBuffer) get(seq uint16) *retainablePacket {
 	s.m.RLock()
 	defer s.m.RUnlock()
 
@@ -79,5 +87,15 @@ func (s *sendBuffer) get(seq uint16) *rtp.Packet {
 		return nil
 	}
 
-	return s.packets[seq%s.size]
+	pkt := s.packets[seq%s.size]
+	if pkt != nil {
+		if pkt.Header().SequenceNumber != seq {
+			return nil
+		}
+		// already released
+		if err := pkt.Retain(); err != nil {
+			return nil
+		}
+	}
+	return pkt
 }


Test patch:
diff --git a/pkg/nack/send_buffer_test.go b/pkg/nack/send_buffer_test.go
index d107195..751f6c0 100644
--- a/pkg/nack/send_buffer_test.go
+++ b/pkg/nack/send_buffer_test.go
@@ -4,20 +4,23 @@ import (
 	"testing"
 
 	"github.com/pion/rtp"
-	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func TestSendBuffer(t *testing.T) {
+	pm := newPacketManager()
 	for _, start := range []uint16{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 511, 512, 513, 32767, 32768, 32769, 65527, 65528, 65529, 65530, 65531, 65532, 65533, 65534, 65535} {
 		start := start
 
 		sb, err := newSendBuffer(8)
-		assert.NoError(t, err)
+		require.NoError(t, err)
 
 		add := func(nums ...uint16) {
 			for _, n := range nums {
 				seq := start + n
-				sb.add(&rtp.Packet{Header: rtp.Header{SequenceNumber: seq}})
+				pkt, err := pm.NewPacket(&rtp.Header{SequenceNumber: seq}, nil)
+				require.NoError(t, err)
+				sb.add(pkt)
 			}
 		}
 
@@ -30,9 +33,10 @@ func TestSendBuffer(t *testing.T) {
 					t.Errorf("packet not found: %d", seq)
 					continue
 				}
-				if packet.SequenceNumber != seq {
-					t.Errorf("packet for %d returned with incorrect SequenceNumber: %d", seq, packet.SequenceNumber)
+				if packet.Header().SequenceNumber != seq {
+					t.Errorf("packet for %d returned with incorrect SequenceNumber: %d", seq, packet.Header().SequenceNumber)
 				}
+				packet.Release()
 			}
 		}
 		assertNOTGet := func(nums ...uint16) {
@@ -41,7 +45,7 @@ func TestSendBuffer(t *testing.T) {
 				seq := start + n
 				packet := sb.get(seq)
 				if packet != nil {
-					t.Errorf("packet found for %d: %d", seq, packet.SequenceNumber)
+					t.Errorf("packet found for %d: %d", seq, packet.Header().SequenceNumber)
 				}
 			}
 		}
@@ -63,20 +67,52 @@ func TestSendBuffer(t *testing.T) {
 	}
 }
 
+func TestSendBuffer_Overridden(t *testing.T) {
+	// override original packet content and get
+	pm := newPacketManager()
+	sb, err := newSendBuffer(1)
+	require.NoError(t, err)
+	require.Equal(t, uint16(1), sb.size)
+
+	originalBytes := []byte("originalContent")
+	pkt, err := pm.NewPacket(&rtp.Header{SequenceNumber: 1}, originalBytes)
+	require.NoError(t, err)
+	sb.add(pkt)
+
+	// change payload
+	copy(originalBytes, "altered")
+	retrieved := sb.get(1)
+	require.NotNil(t, retrieved)
+	require.Equal(t, "originalContent", string(retrieved.Payload()))
+	retrieved.Release()
+	require.Equal(t, 1, retrieved.count)
+
+	// ensure original packet is released
+	pkt, err = pm.NewPacket(&rtp.Header{SequenceNumber: 2}, originalBytes)
+	require.NoError(t, err)
+	sb.add(pkt)
+	require.Equal(t, 0, retrieved.count)
+
+	require.Nil(t, sb.get(1))
+}
+
 // this test is only useful when being run with the race detector, it won't fail otherwise:
 //
 //     go test -race ./pkg/nack/
 func TestSendBuffer_Race(t *testing.T) {
+	pm := newPacketManager()
 	for _, start := range []uint16{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 511, 512, 513, 32767, 32768, 32769, 65527, 65528, 65529, 65530, 65531, 65532, 65533, 65534, 65535} {
 		start := start
 
 		sb, err := newSendBuffer(8)
-		assert.NoError(t, err)
+		require.NoError(t, err)
 
 		add := func(nums ...uint16) {
 			for _, n := range nums {
 				seq := start + n
-				sb.add(&rtp.Packet{Header: rtp.Header{SequenceNumber: seq}})
+				pkt, err := pm.NewPacket(&rtp.Header{SequenceNumber: seq}, nil)
+				require.NoError(t, err)
+				sb.add(pkt)
 			}
 		}