Adds WithPacketHeadHandler to SampleBuilder

This option allows inspecting head packet for
each media.Sample and then lets the user return
their custom metadata. This might be useful
in case when you need to check whether the
given sample is a keyframe.

Author: digitalix

pkg/media/media.go

@@ -17,6 +17,7 @@ type Sample struct {
 	Duration           time.Duration
 	PacketTimestamp    uint32
 	PrevDroppedPackets uint16
+	Metadata           interface{}
 }
 
 // Writer defines an interface to handle

pkg/media/samplebuilder/samplebuilder.go

@@ -40,6 +40,9 @@ type SampleBuilder struct {
 
 	// number of packets forced to be dropped
 	droppedPackets uint16
+
+	// allows inspecting head packets of each sample and then returns a custom metadata
+	packetHeadHandler func(headPacket interface{}) interface{}
 }
 
 // New constructs a new SampleBuilder.
@@ -253,11 +256,16 @@ func (s *SampleBuilder) buildSample(purgingBuffers bool) *media.Sample {
 
 	// merge all the buffers into a sample
 	data := []byte{}
+	var metadata interface{}
 	for i := consume.head; i != consume.tail; i++ {
 		p, err := s.depacketizer.Unmarshal(s.buffer[i].Payload)
 		if err != nil {
 			return nil
 		}
+		if i == consume.head && s.packetHeadHandler != nil {
+			metadata = s.packetHeadHandler(s.depacketizer)
+		}
+
 		data = append(data, p...)
 	}
 	samples := afterTimestamp - sampleTimestamp
@@ -267,6 +275,7 @@ func (s *SampleBuilder) buildSample(purgingBuffers bool) *media.Sample {
 		Duration:           time.Duration((float64(samples)/float64(s.sampleRate))*secondToNanoseconds) * time.Nanosecond,
 		PacketTimestamp:    sampleTimestamp,
 		PrevDroppedPackets: s.droppedPackets,
+		Metadata:           metadata,
 	}
 
 	s.droppedPackets = 0
@@ -341,6 +350,14 @@ func WithPacketReleaseHandler(h func(*rtp.Packet)) Option {
 	}
 }
 
+// WithPacketHeadHandler set a head packet handler to allow inspecting
+// the packet to extract certain information and return as custom metadata
+func WithPacketHeadHandler(h func(headPacket interface{}) interface{}) Option {
+	return func(o *SampleBuilder) {
+		o.packetHeadHandler = h
+	}
+}
+
 // WithMaxTimeDelay ensures that packets that are too old in the buffer get
 // purged based on time rather than building up an extraordinarily long delay.
 func WithMaxTimeDelay(maxLateDuration time.Duration) Option {

pkg/media/samplebuilder/samplebuilder_test.go

@@ -389,6 +389,38 @@ func TestSampleBuilderWithPacketReleaseHandler(t *testing.T) {
 	}
 }
 
+func TestSampleBuilderWithPacketHeadHandler(t *testing.T) {
+	packets := []*rtp.Packet{
+		{Header: rtp.Header{SequenceNumber: 5000, Timestamp: 5}, Payload: []byte{0x01}},
+		{Header: rtp.Header{SequenceNumber: 5001, Timestamp: 5}, Payload: []byte{0x02}},
+		{Header: rtp.Header{SequenceNumber: 5002, Timestamp: 6}, Payload: []byte{0x01}},
+		{Header: rtp.Header{SequenceNumber: 5003, Timestamp: 6}, Payload: []byte{0x02}},
+		{Header: rtp.Header{SequenceNumber: 5004, Timestamp: 7}, Payload: []byte{0x01}},
+	}
+
+	headCount := 0
+	s := New(10, &fakeDepacketizer{}, 1, WithPacketHeadHandler(func(headPacket interface{}) interface{} {
+		headCount++
+		return true
+	}))
+
+	for _, pkt := range packets {
+		s.Push(pkt)
+	}
+
+	for {
+		sample := s.Pop()
+		if sample == nil {
+			break
+		}
+
+		assert.NotNil(t, sample.Metadata, "sample metadata shouldn't be nil")
+		assert.Equal(t, true, sample.Metadata, "sample metadata should've been set to true")
+	}
+
+	assert.Equal(t, 2, headCount, "two sample heads should have been inspected")
+}
+
 func TestPopWithTimestamp(t *testing.T) {
 	t.Run("Crash on nil", func(t *testing.T) {
 		s := New(0, &fakeDepacketizer{}, 1)