Performance optimization: using NextBytesNoCopy where possible (#17)

* Performance optimization: using NextBytesFast where possible * follow up * BenchmarkDemuxer_NextData added * NextBytesFast -> NextBytesNoCopy Co-authored-by: Ilya Barbashov <[email protected]>
asticode · Mar 27, 2021 · 077af8e · 077af8e
1 parent df3794c
commit 077af8e
Show file tree

Hide file tree

Showing 19 changed files with 191 additions and 42 deletions.
diff --git a/data_eit.go b/data_eit.go
@@ -36,7 +36,7 @@ func parseEITSection(i *astikit.BytesIterator, offsetSectionsEnd int, tableIDExt
 
 	// Get next 2 bytes
 	var bs []byte
-	if bs, err = i.NextBytes(2); err != nil {
+	if bs, err = i.NextBytesNoCopy(2); err != nil {
 		err = fmt.Errorf("astits: fetching next bytes failed: %w", err)
 		return
 	}
@@ -45,7 +45,7 @@ func parseEITSection(i *astikit.BytesIterator, offsetSectionsEnd int, tableIDExt
 	d.TransportStreamID = uint16(bs[0])<<8 | uint16(bs[1])
 
 	// Get next 2 bytes
-	if bs, err = i.NextBytes(2); err != nil {
+	if bs, err = i.NextBytesNoCopy(2); err != nil {
 		err = fmt.Errorf("astits: fetching next bytes failed: %w", err)
 		return
 	}
@@ -75,7 +75,7 @@ func parseEITSection(i *astikit.BytesIterator, offsetSectionsEnd int, tableIDExt
 	// Loop until end of section data is reached
 	for i.Offset() < offsetSectionsEnd {
 		// Get next 2 bytes
-		if bs, err = i.NextBytes(2); err != nil {
+		if bs, err = i.NextBytesNoCopy(2); err != nil {
 			err = fmt.Errorf("astits: fetching next bytes failed: %w", err)
 			return
 		}

diff --git a/data_nit.go b/data_nit.go
@@ -35,7 +35,7 @@ func parseNITSection(i *astikit.BytesIterator, tableIDExtension uint16) (d *NITD
 
 	// Get next bytes
 	var bs []byte
-	if bs, err = i.NextBytes(2); err != nil {
+	if bs, err = i.NextBytesNoCopy(2); err != nil {
 		err = fmt.Errorf("astits: fetching next bytes failed: %w", err)
 		return
 	}
@@ -50,7 +50,7 @@ func parseNITSection(i *astikit.BytesIterator, tableIDExtension uint16) (d *NITD
 		ts := &NITDataTransportStream{}
 
 		// Get next bytes
-		if bs, err = i.NextBytes(2); err != nil {
+		if bs, err = i.NextBytesNoCopy(2); err != nil {
 			err = fmt.Errorf("astits: fetching next bytes failed: %w", err)
 			return
 		}
@@ -59,7 +59,7 @@ func parseNITSection(i *astikit.BytesIterator, tableIDExtension uint16) (d *NITD
 		ts.TransportStreamID = uint16(bs[0])<<8 | uint16(bs[1])
 
 		// Get next bytes
-		if bs, err = i.NextBytes(2); err != nil {
+		if bs, err = i.NextBytesNoCopy(2); err != nil {
 			err = fmt.Errorf("astits: fetching next bytes failed: %w", err)
 			return
 		}

diff --git a/data_pat.go b/data_pat.go
@@ -32,7 +32,7 @@ func parsePATSection(i *astikit.BytesIterator, offsetSectionsEnd int, tableIDExt
 	for i.Offset() < offsetSectionsEnd {
 		// Get next bytes
 		var bs []byte
-		if bs, err = i.NextBytes(4); err != nil {
+		if bs, err = i.NextBytesNoCopy(4); err != nil {
 			err = fmt.Errorf("astits: fetching next bytes failed: %w", err)
 			return
 		}

diff --git a/data_pat_test.go b/data_pat_test.go
@@ -35,7 +35,7 @@ func TestParsePATSection(t *testing.T) {
 	assert.NoError(t, err)
 }
 
-func TestWritePatSection(t *testing.T) {
+func TestWritePATSection(t *testing.T) {
 	bw := &bytes.Buffer{}
 	w := astikit.NewBitsWriter(astikit.BitsWriterOptions{Writer: bw})
 	n, err := writePATSection(w, pat)
@@ -44,3 +44,25 @@ func TestWritePatSection(t *testing.T) {
 	assert.Equal(t, n, bw.Len())
 	assert.Equal(t, patBytes(), bw.Bytes())
 }
+
+func BenchmarkParsePATSection(b *testing.B) {
+	b.ReportAllocs()
+	bs := patBytes()
+
+	for i := 0; i < b.N; i++ {
+		parsePATSection(astikit.NewBytesIterator(bs), len(bs), uint16(1))
+	}
+}
+
+func BenchmarkWritePATSection(b *testing.B) {
+	b.ReportAllocs()
+
+	bw := &bytes.Buffer{}
+	bw.Grow(1024)
+	w := astikit.NewBitsWriter(astikit.BitsWriterOptions{Writer: bw})
+
+	for i := 0; i < b.N; i++ {
+		bw.Reset()
+		writePATSection(w, pat)
+	}
+}
diff --git a/data_pes.go b/data_pes.go
@@ -162,7 +162,7 @@ func parsePESHeader(i *astikit.BytesIterator) (h *PESHeader, dataStart, dataEnd
 
 	// Get next bytes
 	var bs []byte
-	if bs, err = i.NextBytes(2); err != nil {
+	if bs, err = i.NextBytesNoCopy(2); err != nil {
 		err = fmt.Errorf("astits: fetching next bytes failed: %w", err)
 		return
 	}
@@ -276,7 +276,7 @@ func parsePESOptionalHeader(i *astikit.BytesIterator) (h *PESOptionalHeader, dat
 	// ES rate
 	if h.HasESRate {
 		var bs []byte
-		if bs, err = i.NextBytes(3); err != nil {
+		if bs, err = i.NextBytesNoCopy(3); err != nil {
 			err = fmt.Errorf("astits: fetching next bytes failed: %w", err)
 			return
 		}
@@ -304,7 +304,7 @@ func parsePESOptionalHeader(i *astikit.BytesIterator) (h *PESOptionalHeader, dat
 	// CRC
 	if h.HasCRC {
 		var bs []byte
-		if bs, err = i.NextBytes(2); err != nil {
+		if bs, err = i.NextBytesNoCopy(2); err != nil {
 			err = fmt.Errorf("astits: fetching next bytes failed: %w", err)
 			return
 		}
@@ -347,7 +347,7 @@ func parsePESOptionalHeader(i *astikit.BytesIterator) (h *PESOptionalHeader, dat
 		// Program packet sequence counter
 		if h.HasProgramPacketSequenceCounter {
 			var bs []byte
-			if bs, err = i.NextBytes(2); err != nil {
+			if bs, err = i.NextBytesNoCopy(2); err != nil {
 				err = fmt.Errorf("astits: fetching next bytes failed: %w", err)
 				return
 			}
@@ -359,7 +359,7 @@ func parsePESOptionalHeader(i *astikit.BytesIterator) (h *PESOptionalHeader, dat
 		// P-STD buffer
 		if h.HasPSTDBuffer {
 			var bs []byte
-			if bs, err = i.NextBytes(2); err != nil {
+			if bs, err = i.NextBytesNoCopy(2); err != nil {
 				err = fmt.Errorf("astits: fetching next bytes failed: %w", err)
 				return
 			}
@@ -405,7 +405,7 @@ func parseDSMTrickMode(i byte) (m *DSMTrickMode) {
 // parsePTSOrDTS parses a PTS or a DTS
 func parsePTSOrDTS(i *astikit.BytesIterator) (cr *ClockReference, err error) {
 	var bs []byte
-	if bs, err = i.NextBytes(5); err != nil {
+	if bs, err = i.NextBytesNoCopy(5); err != nil {
 		err = fmt.Errorf("astits: fetching next bytes failed: %w", err)
 		return
 	}
@@ -416,7 +416,7 @@ func parsePTSOrDTS(i *astikit.BytesIterator) (cr *ClockReference, err error) {
 // parseESCR parses an ESCR
 func parseESCR(i *astikit.BytesIterator) (cr *ClockReference, err error) {
 	var bs []byte
-	if bs, err = i.NextBytes(6); err != nil {
+	if bs, err = i.NextBytesNoCopy(6); err != nil {
 		err = fmt.Errorf("astits: fetching next bytes failed: %w", err)
 		return
 	}

diff --git a/data_pes_test.go b/data_pes_test.go
@@ -485,3 +485,25 @@ func TestWritePESOptionalHeader(t *testing.T) {
 		})
 	}
 }
+
+func BenchmarkParsePESData(b *testing.B) {
+	bss := make([][]byte, len(pesTestCases))
+
+	for ti, tc := range pesTestCases {
+		buf := bytes.Buffer{}
+		w := astikit.NewBitsWriter(astikit.BitsWriterOptions{Writer: &buf})
+		tc.headerBytesFunc(w, true, true)
+		tc.optionalHeaderBytesFunc(w, true, true)
+		tc.bytesFunc(w, true, true)
+		bss[ti] = buf.Bytes()
+	}
+
+	for ti, tc := range pesTestCases {
+		b.Run(tc.name, func(b *testing.B) {
+			b.ReportAllocs()
+			for i := 0; i < b.N; i++ {
+				parsePESData(astikit.NewBytesIterator(bss[ti]))
+			}
+		})
+	}
+}
diff --git a/data_pmt.go b/data_pmt.go
@@ -58,7 +58,7 @@ func parsePMTSection(i *astikit.BytesIterator, offsetSectionsEnd int, tableIDExt
 
 	// Get next bytes
 	var bs []byte
-	if bs, err = i.NextBytes(2); err != nil {
+	if bs, err = i.NextBytesNoCopy(2); err != nil {
 		err = fmt.Errorf("astits: fetching next bytes failed: %w", err)
 		return
 	}
@@ -88,7 +88,7 @@ func parsePMTSection(i *astikit.BytesIterator, offsetSectionsEnd int, tableIDExt
 		e.StreamType = StreamType(b)
 
 		// Get next bytes
-		if bs, err = i.NextBytes(2); err != nil {
+		if bs, err = i.NextBytesNoCopy(2); err != nil {
 			err = fmt.Errorf("astits: fetching next bytes failed: %w", err)
 			return
 		}

diff --git a/data_pmt_test.go b/data_pmt_test.go
@@ -49,3 +49,25 @@ func TestWritePMTSection(t *testing.T) {
 	assert.Equal(t, n, buf.Len())
 	assert.Equal(t, pmtBytes(), buf.Bytes())
 }
+
+func BenchmarkParsePMTSection(b *testing.B) {
+	b.ReportAllocs()
+	bs := pmtBytes()
+
+	for i := 0; i < b.N; i++ {
+		parsePMTSection(astikit.NewBytesIterator(bs), len(bs), uint16(1))
+	}
+}
+
+func BenchmarkWritePMTSection(b *testing.B) {
+	b.ReportAllocs()
+
+	bw := &bytes.Buffer{}
+	bw.Grow(1024)
+	w := astikit.NewBitsWriter(astikit.BitsWriterOptions{Writer: bw})
+
+	for i := 0; i < b.N; i++ {
+		bw.Reset()
+		writePMTSection(w, pmt)
+	}
+}
diff --git a/data_psi.go b/data_psi.go
@@ -165,7 +165,7 @@ func parsePSISection(i *astikit.BytesIterator) (s *PSISection, stop bool, err er
 			// Get CRC32 data
 			i.Seek(offsetStart)
 			var crc32Data []byte
-			if crc32Data, err = i.NextBytes(offsetSectionsEnd - offsetStart); err != nil {
+			if crc32Data, err = i.NextBytesNoCopy(offsetSectionsEnd - offsetStart); err != nil {
 				err = fmt.Errorf("astits: fetching next bytes failed: %w", err)
 				return
 			}
@@ -189,7 +189,7 @@ func parsePSISection(i *astikit.BytesIterator) (s *PSISection, stop bool, err er
 // parseCRC32 parses a CRC32
 func parseCRC32(i *astikit.BytesIterator) (c uint32, err error) {
 	var bs []byte
-	if bs, err = i.NextBytes(4); err != nil {
+	if bs, err = i.NextBytesNoCopy(4); err != nil {
 		err = fmt.Errorf("astits: fetching next bytes failed: %w", err)
 		return
 	}
@@ -229,7 +229,7 @@ func parsePSISectionHeader(i *astikit.BytesIterator) (h *PSISectionHeader, offse
 
 	// Get next bytes
 	var bs []byte
-	if bs, err = i.NextBytes(2); err != nil {
+	if bs, err = i.NextBytesNoCopy(2); err != nil {
 		err = fmt.Errorf("astits: fetching next bytes failed: %w", err)
 		return
 	}
@@ -358,7 +358,7 @@ func parsePSISectionSyntaxHeader(i *astikit.BytesIterator) (h *PSISectionSyntaxH
 
 	// Get next 2 bytes
 	var bs []byte
-	if bs, err = i.NextBytes(2); err != nil {
+	if bs, err = i.NextBytesNoCopy(2); err != nil {
 		err = fmt.Errorf("astits: fetching next bytes failed: %w", err)
 		return
 	}

diff --git a/data_psi_test.go b/data_psi_test.go
@@ -374,3 +374,10 @@ func TestWritePSIData(t *testing.T) {
 		})
 	}
 }
+
+func BenchmarkParsePSIData(b *testing.B) {
+	b.ReportAllocs()
+	for i := 0; i < b.N; i++ {
+		parsePSIData(astikit.NewBytesIterator(psiBytes()))
+	}
+}
diff --git a/data_sdt.go b/data_sdt.go
@@ -42,7 +42,7 @@ func parseSDTSection(i *astikit.BytesIterator, offsetSectionsEnd int, tableIDExt
 
 	// Get next bytes
 	var bs []byte
-	if bs, err = i.NextBytes(2); err != nil {
+	if bs, err = i.NextBytesNoCopy(2); err != nil {
 		err = fmt.Errorf("astits: fetching next bytes failed: %w", err)
 		return
 	}
@@ -59,7 +59,7 @@ func parseSDTSection(i *astikit.BytesIterator, offsetSectionsEnd int, tableIDExt
 		s := &SDTDataService{}
 
 		// Get next bytes
-		if bs, err = i.NextBytes(2); err != nil {
+		if bs, err = i.NextBytesNoCopy(2); err != nil {
 			err = fmt.Errorf("astits: fetching next bytes failed: %w", err)
 			return
 		}

diff --git a/demuxer_test.go b/demuxer_test.go
@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"context"
 	"fmt"
+	"io"
 	"testing"
 
 	"github.com/asticode/go-astikit"
@@ -99,3 +100,28 @@ func TestDemuxerRewind(t *testing.T) {
 	assert.Equal(t, 0, len(dmx.packetPool.b))
 	assert.Nil(t, dmx.packetBuffer)
 }
+
+func BenchmarkDemuxer_NextData(b *testing.B) {
+	b.ReportAllocs()
+
+	buf := &bytes.Buffer{}
+	w := astikit.NewBitsWriter(astikit.BitsWriterOptions{Writer: buf})
+	bs := psiBytes()
+	b1, _ := packet(PacketHeader{ContinuityCounter: uint8(0), PayloadUnitStartIndicator: true, PID: PIDPAT}, PacketAdaptationField{}, bs[:147], true)
+	w.Write(b1)
+	b2, _ := packet(PacketHeader{ContinuityCounter: uint8(1), PID: PIDPAT}, PacketAdaptationField{}, bs[147:], true)
+	w.Write(b2)
+
+	r := bytes.NewReader(buf.Bytes())
+
+	for i := 0; i < b.N; i++ {
+		r.Seek(0, io.SeekStart)
+		dmx := NewDemuxer(context.Background(), r)
+
+		for _, s := range psi.Sections {
+			if !s.Header.TableID.isUnknown() {
+				dmx.NextData()
+			}
+		}
+	}
+}