Improve memory performance of collector

Avoid unnecessary slice creation when adding a record to a set. We implement this change for both data record and template records, but in practice most of the gains will come from data records. Memory allocations (in bytes) are reduced by around 15% for BenchmarkCollector. Signed-off-by: Antonin Bas <[email protected]>
vmware · Oct 5, 2023 · b3d6e26 · b3d6e26
1 parent c5e0992
commit b3d6e26
Show file tree

Hide file tree

Showing 3 changed files with 66 additions and 23 deletions.
diff --git a/pkg/collector/process.go b/pkg/collector/process.go
@@ -280,7 +280,7 @@ func (cp *CollectingProcess) decodeTemplateSet(templateBuffer *bytes.Buffer, obs
 			return nil, err
 		}
 	}
-	err := templateSet.AddRecord(elementsWithValue, templateID)
+	err := templateSet.AddRecordV2(elementsWithValue, templateID)
 	if err != nil {
 		return nil, err
 	}
@@ -300,7 +300,7 @@ func (cp *CollectingProcess) decodeDataSet(dataBuffer *bytes.Buffer, obsDomainID
 	}
 
 	for dataBuffer.Len() > 0 {
-		elements := make([]entities.InfoElementWithValue, len(template))
+		elements := make([]entities.InfoElementWithValue, len(template), len(template)+cp.numExtraElements)
 		for i, element := range template {
 			var length int
 			if element.Len == entities.VariableLength { // string
@@ -312,7 +312,7 @@ func (cp *CollectingProcess) decodeDataSet(dataBuffer *bytes.Buffer, obsDomainID
 				return nil, err
 			}
 		}
-		err = dataSet.AddRecordWithExtraElements(elements, cp.numExtraElements, templateID)
+		err = dataSet.AddRecordV2(elements, templateID)
 		if err != nil {
 			return nil, err
 		}

diff --git a/pkg/entities/record.go b/pkg/entities/record.go
@@ -68,6 +68,24 @@ func NewDataRecord(id uint16, numElements, numExtraElements int, isDecoding bool
 	}
 }
 
+func NewDataRecordFromElements(id uint16, elements []InfoElementWithValue, isDecoding bool) *dataRecord {
+	length := 0
+	if !isDecoding {
+		for idx := range elements {
+			length += elements[idx].GetLength()
+		}
+	}
+	return &dataRecord{
+		baseRecord{
+			fieldCount:         uint16(len(elements)),
+			templateID:         id,
+			isDecoding:         isDecoding,
+			len:                length,
+			orderedElementList: elements,
+		},
+	}
+}
+
 type templateRecord struct {
 	baseRecord
 	// Minimum data record length required to be sent for this template.
@@ -91,6 +109,25 @@ func NewTemplateRecord(id uint16, numElements int, isDecoding bool) *templateRec
 	}
 }
 
+func NewTemplateRecordFromElements(id uint16, elements []InfoElementWithValue, isDecoding bool) *templateRecord {
+	r := &templateRecord{
+		baseRecord{
+			buffer:             make([]byte, 4),
+			fieldCount:         uint16(len(elements)),
+			templateID:         id,
+			isDecoding:         isDecoding,
+			orderedElementList: elements,
+		},
+		0,
+		len(elements),
+	}
+	for idx := range elements {
+		infoElement := elements[idx].GetInfoElement()
+		r.addInfoElement(infoElement)
+	}
+	return r
+}
+
 func (b *baseRecord) GetTemplateID() uint16 {
 	return b.templateID
 }
@@ -205,12 +242,7 @@ func (t *templateRecord) PrepareRecord() error {
 	return nil
 }
 
-func (t *templateRecord) AddInfoElement(element InfoElementWithValue) error {
-	infoElement := element.GetInfoElement()
-	// val could be used to specify smaller length than default? For now assert it to be nil
-	if !element.IsValueEmpty() {
-		return fmt.Errorf("template record cannot take element %v with non-empty value", infoElement.Name)
-	}
+func (t *templateRecord) addInfoElement(infoElement *InfoElement) {
 	initialLength := len(t.buffer)
 	// Add field specifier {elementID: uint16, elementLen: uint16}
 	addBytes := make([]byte, 4)
@@ -224,14 +256,23 @@ func (t *templateRecord) AddInfoElement(element InfoElementWithValue) error {
 		binary.BigEndian.PutUint32(addBytes, infoElement.EnterpriseId)
 		t.buffer = append(t.buffer, addBytes...)
 	}
-	t.orderedElementList[t.index] = element
-	t.index++
 	// Keep track of minimum data record length required for sanity check
 	if infoElement.Len == VariableLength {
 		t.minDataRecLength = t.minDataRecLength + 1
 	} else {
 		t.minDataRecLength = t.minDataRecLength + infoElement.Len
 	}
+}
+
+func (t *templateRecord) AddInfoElement(element InfoElementWithValue) error {
+	infoElement := element.GetInfoElement()
+	// val could be used to specify smaller length than default? For now assert it to be nil
+	if !element.IsValueEmpty() {
+		return fmt.Errorf("template record cannot take element %v with non-empty value", infoElement.Name)
+	}
+	t.addInfoElement(infoElement)
+	t.orderedElementList[t.index] = element
+	t.index++
 	return nil
 }
 

diff --git a/pkg/entities/set.go b/pkg/entities/set.go
@@ -49,6 +49,10 @@ type Set interface {
 	UpdateLenInHeader()
 	AddRecord(elements []InfoElementWithValue, templateID uint16) error
 	AddRecordWithExtraElements(elements []InfoElementWithValue, numExtraElements int, templateID uint16) error
+	// Unlike AddRecord, AddRecordV2 uses the elements slice directly, instead of creating a new
+	// 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
 	GetRecords() []Record
 	GetNumberOfRecords() uint32
 }
@@ -122,9 +126,13 @@ func (s *set) UpdateLenInHeader() {
 }
 
 func (s *set) AddRecord(elements []InfoElementWithValue, templateID uint16) error {
+	return s.AddRecordWithExtraElements(elements, 0, templateID)
+}
+
+func (s *set) AddRecordWithExtraElements(elements []InfoElementWithValue, numExtraElements int, templateID uint16) error {
 	var record Record
 	if s.setType == Data {
-		record = NewDataRecord(templateID, len(elements), 0, s.isDecoding)
+		record = NewDataRecord(templateID, len(elements), numExtraElements, s.isDecoding)
 	} else if s.setType == Template {
 		record = NewTemplateRecord(templateID, len(elements), s.isDecoding)
 		err := record.PrepareRecord()
@@ -134,8 +142,8 @@ func (s *set) AddRecord(elements []InfoElementWithValue, templateID uint16) erro
 	} else {
 		return fmt.Errorf("set type is not supported")
 	}
-	for _, element := range elements {
-		err := record.AddInfoElement(element)
+	for i := range elements {
+		err := record.AddInfoElement(elements[i])
 		if err != nil {
 			return err
 		}
@@ -145,25 +153,19 @@ func (s *set) AddRecord(elements []InfoElementWithValue, templateID uint16) erro
 	return nil
 }
 
-func (s *set) AddRecordWithExtraElements(elements []InfoElementWithValue, numExtraElements int, templateID uint16) error {
+func (s *set) AddRecordV2(elements []InfoElementWithValue, templateID uint16) error {
 	var record Record
 	if s.setType == Data {
-		record = NewDataRecord(templateID, len(elements), numExtraElements, s.isDecoding)
+		record = NewDataRecordFromElements(templateID, elements, s.isDecoding)
 	} else if s.setType == Template {
-		record = NewTemplateRecord(templateID, len(elements), s.isDecoding)
+		record = NewTemplateRecordFromElements(templateID, elements, s.isDecoding)
 		err := record.PrepareRecord()
 		if err != nil {
 			return err
 		}
 	} else {
 		return fmt.Errorf("set type is not supported")
 	}
-	for i := range elements {
-		err := record.AddInfoElement(elements[i])
-		if err != nil {
-			return err
-		}
-	}
 	s.records = append(s.records, record)
 	s.length += record.GetRecordLength()
 	return nil