From a9850b03ca5485192c925c27893b6a519d72e66d Mon Sep 17 00:00:00 2001
From: Atsushi Watanabe <atsushi.w@ieee.org>
Date: Fri, 31 May 2024 15:59:30 +0900
Subject: [PATCH] SampleBuilder: add memory leak test

Test that the input RTP packets are unreferenced
after all samples are popped.
---
 pkg/media/samplebuilder/samplebuilder_test.go | 41 +++++++++++++++++++
 1 file changed, 41 insertions(+)

diff --git a/pkg/media/samplebuilder/samplebuilder_test.go b/pkg/media/samplebuilder/samplebuilder_test.go
index 3f888c79ec..c20a6797ee 100644
--- a/pkg/media/samplebuilder/samplebuilder_test.go
+++ b/pkg/media/samplebuilder/samplebuilder_test.go
@@ -5,6 +5,8 @@ package samplebuilder
 
 import (
 	"fmt"
+	"runtime"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -529,6 +531,45 @@ func TestSampleBuilderData(t *testing.T) {
 	assert.Equal(t, j, 0x1FFFF)
 }
 
+func TestSampleBuilderPacketUnreference(t *testing.T) {
+	s := New(10, &fakeDepacketizer{
+		headChecker: true,
+	}, 1)
+
+	var refs int64
+	finalizer := func(*rtp.Packet) {
+		atomic.AddInt64(&refs, -1)
+	}
+
+	for i := 0; i < 0x20000; i++ {
+		atomic.AddInt64(&refs, 1)
+		p := rtp.Packet{
+			Header: rtp.Header{
+				SequenceNumber: uint16(i),
+				Timestamp:      uint32(i + 42),
+			},
+			Payload: []byte{byte(i)},
+		}
+		runtime.SetFinalizer(&p, finalizer)
+		s.Push(&p)
+		for {
+			sample := s.Pop()
+			if sample == nil {
+				break
+			}
+		}
+	}
+
+	runtime.GC()
+	time.Sleep(10 * time.Millisecond)
+
+	remainedRefs := atomic.LoadInt64(&refs)
+	runtime.KeepAlive(s)
+
+	// only the last packet should be still referenced
+	assert.Equal(t, int64(1), remainedRefs)
+}
+
 func TestSampleBuilder_Flush(t *testing.T) {
 	s := New(50, &fakeDepacketizer{
 		headChecker: true,