Improve IPFIX Set implementation (#390)

* Optimize ResetSet to avoid memory allocations
* Add new interface method AddRecordV3 which takes a Record as a
  parameter (instead of a list of IEs). AddRecordV3 can avoid all memory
  allocations.

We also add a new GetRecordType method to the Record interface.

Signed-off-by: Antonin Bas <antonin.bas@broadcom.com>

pkg/entities/record.go

@@ -41,6 +41,7 @@ type Record interface {
 	GetRecordLength() int
 	GetMinDataRecordLen() uint16
 	GetElementMap() map[string]interface{}
+	GetRecordType() ContentType
 }
 
 type baseRecord struct {
@@ -235,6 +236,10 @@ func (d *dataRecord) AddInfoElement(element InfoElementWithValue) error {
 	return nil
 }
 
+func (d *dataRecord) GetRecordType() ContentType {
+	return Data
+}
+
 func (t *templateRecord) PrepareRecord() error {
 	// Add Template Record Header
 	binary.BigEndian.PutUint16(t.buffer[0:2], t.templateID)
@@ -287,3 +292,7 @@ func (t *templateRecord) GetRecordLength() int {
 func (t *templateRecord) GetMinDataRecordLen() uint16 {
 	return t.minDataRecLength
 }
+
+func (d *templateRecord) GetRecordType() ContentType {
+	return Template
+}

pkg/entities/set.go

@@ -34,6 +34,7 @@ const (
 	SetHeaderLen  int    = 4
 )
 
+// ContenType is used for both sets and records.
 type ContentType uint8
 
 const (
@@ -45,6 +46,8 @@ const (
 
 type Set interface {
 	PrepareSet(setType ContentType, templateID uint16) error
+	// Call ResetSet followed by a new call to PrepareSet in order to reuse an existing Set,
+	// instead of instantiating a new one.
 	ResetSet()
 	GetHeaderBuffer() []byte
 	GetSetLength() int
@@ -56,13 +59,19 @@ type Set interface {
 	// one. This can result in fewer memory allocations. The caller should not modify the
 	// contents of the slice after calling AddRecordV2.
 	AddRecordV2(elements []InfoElementWithValue, templateID uint16) error
+	// Unlike other AddRecord* variants, AddRecordV3 takes an actual existing Record as an input
+	// parameter, instead of a list of Information Elements with values. When calling
+	// AddRecordV3, the Set effectively takes ownership of the Record, and the Record should no
+	// longer be modified by the caller.
+	AddRecordV3(record Record) error
 	GetRecords() []Record
 	GetNumberOfRecords() uint32
 }
 
 type set struct {
 	headerBuffer []byte
 	setType      ContentType
+	templateID   uint16
 	records      []Record
 	isDecoding   bool
 	length       int
@@ -90,6 +99,7 @@ func (s *set) PrepareSet(setType ContentType, templateID uint16) error {
 	} else {
 		s.setType = setType
 	}
+	s.templateID = templateID
 	if !s.isDecoding {
 		// Create the set header and append it when encoding
 		s.createHeader(s.setType, templateID)
@@ -98,14 +108,17 @@ func (s *set) PrepareSet(setType ContentType, templateID uint16) error {
 }
 
 func (s *set) ResetSet() {
-	if !s.isDecoding {
-		s.headerBuffer = nil
-		s.headerBuffer = make([]byte, SetHeaderLen)
+	if s.isDecoding {
+		s.length = 0
+	} else {
+		clear(s.headerBuffer)
 		s.length = SetHeaderLen
 	}
 	s.setType = Undefined
-	s.records = nil
-	s.records = make([]Record, 0)
+	// Clear before shrinking the slice so that existing elements are eligible for garbage collection.
+	clear(s.records)
+	// Shrink the slice: the slice capacity is preserved.
+	s.records = s.records[:0]
 }
 
 func (s *set) GetHeaderBuffer() []byte {
@@ -174,6 +187,20 @@ func (s *set) AddRecordV2(elements []InfoElementWithValue, templateID uint16) er
 	return nil
 }
 
+func (s *set) AddRecordV3(record Record) error {
+	// Sanity check: we need to make sure that the record is allowed to be added.
+	recordType := record.GetRecordType()
+	if recordType != s.setType {
+		return fmt.Errorf("record and set types don't match")
+	}
+	if recordType == Data && record.GetTemplateID() != s.templateID {
+		return fmt.Errorf("all data records in the same data set must have the same template ID")
+	}
+	s.records = append(s.records, record)
+	s.length += record.GetRecordLength()
+	return nil
+}
+
 func (s *set) GetRecords() []Record {
 	return s.records
 }
@@ -193,7 +220,7 @@ func (s *set) createHeader(setType ContentType, templateID uint16) {
 // MakeTemplateSet is a convenience function which creates a template Set with a single Record.
 func MakeTemplateSet(templateID uint16, ies []*InfoElement) (*set, error) {
 	tempSet := NewSet(false)
-	if err := tempSet.PrepareSet(Template, templateID); err != nil {
+	if err := tempSet.PrepareSet(Template, TemplateSetID); err != nil {
 		return nil, err
 	}
 	elements := make([]InfoElementWithValue, len(ies))

pkg/entities/set_test.go

@@ -167,12 +167,49 @@ func TestMakeTemplateSet(t *testing.T) {
 }
 
 func TestMakeDataSet(t *testing.T) {
-	ie1 := NewIPAddressInfoElement(NewInfoElement("sourceIPv4Address", 8, 18, 0, 4), net.ParseIP("1.1.1.1"))
-	ie2 := NewIPAddressInfoElement(NewInfoElement("destinationIPv4Address", 12, 18, 0, 4), net.ParseIP("1.1.2.1"))
+	ie1 := NewIPAddressInfoElement(NewInfoElement("sourceIPv4Address", 8, 18, 0, 4), net.ParseIP(testIPv4Addr1))
+	ie2 := NewIPAddressInfoElement(NewInfoElement("destinationIPv4Address", 12, 18, 0, 4), net.ParseIP(testIPv4Addr2))
 	elements := []InfoElementWithValue{ie1, ie2}
 	s, err := MakeDataSet(testTemplateID, elements)
 	require.NoError(t, err)
 	assert.Equal(t, Data, s.setType)
 	require.Len(t, s.records, 1)
 	assert.Equal(t, elements, s.records[0].GetOrderedElementList())
 }
+
+func BenchmarkSet(b *testing.B) {
+	const (
+		numRecords = 10
+		templateID = 256
+	)
+
+	sourceIE := NewInfoElement("sourceIPv4Address", 8, 18, 0, 4)
+	destinationIE := NewInfoElement("destinationIPv4Address", 12, 18, 0, 4)
+
+	records := make([]Record, numRecords)
+	for idx := range records {
+		records[idx] = NewDataRecordFromElements(
+			templateID,
+			[]InfoElementWithValue{
+				NewIPAddressInfoElement(sourceIE, net.ParseIP(testIPv4Addr1)),
+				NewIPAddressInfoElement(destinationIE, net.ParseIP(testIPv4Addr2)),
+			},
+			false,
+		)
+	}
+
+	set := NewSet(false)
+
+	addRecords := func() {
+		for _, record := range records {
+			set.AddRecordV3(record)
+		}
+	}
+
+	b.ResetTimer()
+	for range b.N {
+		set.ResetSet()
+		set.PrepareSet(Data, templateID)
+		addRecords()
+	}
+}