wip: reduce nr allocations

kubernetes-sigs · Apr 11, 2024 · b7e17b3 · b7e17b3
1 parent d23b92a
commit b7e17b3
Show file tree

Hide file tree

Showing 6 changed files with 153 additions and 92 deletions.
diff --git a/fieldpath/path.go b/fieldpath/path.go
@@ -80,7 +80,7 @@ func (fp Path) Copy() Path {
 
 // MakePath constructs a Path. The parts may be PathElements, ints, strings.
 func MakePath(parts ...interface{}) (Path, error) {
-	var fp Path
+	fp := make(Path, 0, len(parts))
 	for _, p := range parts {
 		switch t := p.(type) {
 		case PathElement:

diff --git a/fieldpath/serialize-pe.go b/fieldpath/serialize-pe.go
@@ -17,11 +17,10 @@ limitations under the License.
 package fieldpath
 
 import (
+	"bytes"
 	"errors"
 	"fmt"
-	"io"
 	"strconv"
-	"strings"
 
 	"sigs.k8s.io/structured-merge-diff/v4/value"
 )
@@ -57,7 +56,7 @@ var (
 func DeserializePathElement(s string) (PathElement, error) {
 	b := []byte(s)
 	if len(b) < 2 {
-		return PathElement{}, errors.New("key must be 2 characters long:")
+		return PathElement{}, errors.New("key must be 2 characters long")
 	}
 	typeSep, b := b[:2], b[2:]
 	if typeSep[1] != peSepBytes[0] {
@@ -99,41 +98,48 @@ func DeserializePathElement(s string) (PathElement, error) {
 
 // SerializePathElement serializes a path element
 func SerializePathElement(pe PathElement) (string, error) {
-	buf := strings.Builder{}
+	buf := bytes.Buffer{}
 	err := serializePathElementToWriter(&buf, pe)
 	return buf.String(), err
 }
 
-func serializePathElementToWriter(w io.Writer, pe PathElement) error {
+type noNewlineWriter struct {
+	*bytes.Buffer
+}
+
+func (w noNewlineWriter) Write(p []byte) (n int, err error) {
+	if len(p) > 0 && p[len(p)-1] == '\n' {
+		p = p[:len(p)-1]
+	}
+	return w.Buffer.Write(p)
+}
+
+func serializePathElementToWriter(w *bytes.Buffer, pe PathElement) error {
 	switch {
 	case pe.FieldName != nil:
 		if _, err := w.Write(peFieldSepBytes); err != nil {
 			return err
 		}
-		fmt.Fprintf(w, "%s", *pe.FieldName)
+		w.WriteString(*pe.FieldName)
 	case pe.Key != nil:
 		if _, err := w.Write(peKeySepBytes); err != nil {
 			return err
 		}
-		jsonVal, err := value.FieldListToJSON(*pe.Key)
-		if err != nil {
+		if err := value.FieldListToJSON(*pe.Key, w); err != nil {
 			return err
 		}
-		fmt.Fprintf(w, "%s", jsonVal)
 	case pe.Value != nil:
 		if _, err := w.Write(peValueSepBytes); err != nil {
 			return err
 		}
-		jsonVal, err := value.ToJSON(*pe.Value)
-		if err != nil {
+		if err := value.ToJSON(*pe.Value, w); err != nil {
 			return err
 		}
-		fmt.Fprintf(w, "%s", jsonVal)
 	case pe.Index != nil:
 		if _, err := w.Write(peIndexSepBytes); err != nil {
 			return err
 		}
-		fmt.Fprintf(w, "%d", *pe.Index)
+		w.WriteString(strconv.Itoa(*pe.Index))
 	default:
 		return errors.New("invalid PathElement")
 	}

diff --git a/fieldpath/serialize.go b/fieldpath/serialize.go
@@ -21,30 +21,36 @@ import (
 	gojson "encoding/json"
 	"fmt"
 	"io"
+	"sort"
+	"unsafe"
 
 	json "sigs.k8s.io/json"
 )
 
 func (s *Set) ToJSON() ([]byte, error) {
 	buf := bytes.Buffer{}
-	err := s.ToJSONStream(&buf)
+	err := s.emitContentsV1(false, &buf, &reusableBuilder{})
 	if err != nil {
 		return nil, err
 	}
 	return buf.Bytes(), nil
 }
 
 func (s *Set) ToJSONStream(w io.Writer) error {
-	err := s.emitContentsV1(false, w)
+	jsonRaw, err := s.ToJSON()
 	if err != nil {
 		return err
 	}
-	return nil
+
+	_, err = w.Write(jsonRaw)
+	return err
 }
 
 type orderedMapItemWriter struct {
-	w        io.Writer
+	w        *bytes.Buffer
 	hasItems bool
+
+	builder *reusableBuilder
 }
 
 // writeKey writes a key to the writer, including a leading comma if necessary.
@@ -69,21 +75,43 @@ func (om *orderedMapItemWriter) writeKey(key []byte) error {
 	return nil
 }
 
+type reusableBuilder struct {
+	bytes.Buffer
+}
+
+func (r *reusableBuilder) reset() *bytes.Buffer {
+	r.Reset()
+	return &r.Buffer
+}
+
+func (r *reusableBuilder) unsafeString() string {
+	b := r.Bytes()
+	return *(*string)(unsafe.Pointer(&b))
+}
+
 // writePathKey writes a path element as a key to the writer, including a leading comma if necessary.
 // The path will be serialized as a JSON string (including quotes) and passed to writeKey.
 // After writing the key, the caller should write the encoded value, e.g. using
 // writeEmptyValue or by directly writing the value to the writer.
 func (om *orderedMapItemWriter) writePathKey(pe PathElement) error {
-	pev, err := SerializePathElement(pe)
-	if err != nil {
+	if om.hasItems {
+		if _, err := om.w.Write([]byte{','}); err != nil {
+			return err
+		}
+	}
+
+	if err := serializePathElementToWriter(om.builder.reset(), pe); err != nil {
 		return err
 	}
-	key, err := gojson.Marshal(pev)
-	if err != nil {
+	if err := gojson.NewEncoder(noNewlineWriter{om.w}).Encode(om.builder.unsafeString()); err != nil {
 		return err
 	}
 
-	return om.writeKey(key)
+	if _, err := om.w.Write([]byte{':'}); err != nil {
+		return err
+	}
+	om.hasItems = true
+	return nil
 }
 
 // writeEmptyValue writes an empty JSON object to the writer.
@@ -95,8 +123,8 @@ func (om orderedMapItemWriter) writeEmptyValue() error {
 	return nil
 }
 
-func (s *Set) emitContentsV1(includeSelf bool, w io.Writer) error {
-	om := orderedMapItemWriter{w: w}
+func (s *Set) emitContentsV1(includeSelf bool, w *bytes.Buffer, r *reusableBuilder) error {
+	om := orderedMapItemWriter{w: w, builder: r}
 	mi, ci := 0, 0
 
 	if _, err := om.w.Write([]byte{'{'}); err != nil {
@@ -129,7 +157,7 @@ func (s *Set) emitContentsV1(includeSelf bool, w io.Writer) error {
 			if err := om.writePathKey(cpe); err != nil {
 				return err
 			}
-			if err := s.Children.members[ci].set.emitContentsV1(c == 0, om.w); err != nil {
+			if err := s.Children.members[ci].set.emitContentsV1(c == 0, om.w, r); err != nil {
 				return err
 			}
 
@@ -160,7 +188,7 @@ func (s *Set) emitContentsV1(includeSelf bool, w io.Writer) error {
 		if err := om.writePathKey(cpe); err != nil {
 			return err
 		}
-		if err := s.Children.members[ci].set.emitContentsV1(false, om.w); err != nil {
+		if err := s.Children.members[ci].set.emitContentsV1(false, om.w, r); err != nil {
 			return err
 		}
 
@@ -237,34 +265,33 @@ func readIterV1(data []byte) (children *Set, isMember bool, err error) {
 				children = &Set{}
 			}
 
+			// Append the member to the members list, we will sort it later
 			m := &children.Members.members
-			// Since we expect that most of the time these will have been
-			// serialized in the right order, we just verify that and append.
-			appendOK := len(*m) == 0 || (*m)[len(*m)-1].Less(pe)
-			if appendOK {
-				*m = append(*m, pe)
-			} else {
-				children.Members.Insert(pe)
-			}
+			*m = append(*m, pe)
 		}
 
 		if v.target != nil {
 			if children == nil {
 				children = &Set{}
 			}
 
-			// Since we expect that most of the time these will have been
-			// serialized in the right order, we just verify that and append.
+			// Append the child to the children list, we will sort it later
 			m := &children.Children.members
-			appendOK := len(*m) == 0 || (*m)[len(*m)-1].pathElement.Less(pe)
-			if appendOK {
-				*m = append(*m, setNode{pe, v.target})
-			} else {
-				*children.Children.Descend(pe) = *v.target
-			}
+			*m = append(*m, setNode{pe, v.target})
 		}
 	}
 
+	// Sort the members and children
+	if children != nil {
+		sort.Slice(children.Members.members, func(i, j int) bool {
+			return children.Members.members[i].Less(children.Members.members[j])
+		})
+
+		sort.Slice(children.Children.members, func(i, j int) bool {
+			return children.Children.members[i].pathElement.Less(children.Children.members[j].pathElement)
+		})
+	}
+
 	if children == nil {
 		isMember = true
 	}

diff --git a/fieldpath/set_test.go b/fieldpath/set_test.go
@@ -90,18 +90,20 @@ func BenchmarkFieldSet(b *testing.B) {
 		}
 		randOperand := func() *Set { return operands[rand.Intn(len(operands))] }
 
-		b.Run(fmt.Sprintf("insert-%v", here.size), func(b *testing.B) {
-			b.ReportAllocs()
-			for i := 0; i < b.N; i++ {
-				makeSet()
-			}
-		})
-		b.Run(fmt.Sprintf("has-%v", here.size), func(b *testing.B) {
-			b.ReportAllocs()
-			for i := 0; i < b.N; i++ {
-				randOperand().Has(randomPathMaker.makePath(here.minPathLen, here.maxPathLen))
-			}
-		})
+		/*
+			b.Run(fmt.Sprintf("insert-%v", here.size), func(b *testing.B) {
+				b.ReportAllocs()
+				for i := 0; i < b.N; i++ {
+					makeSet()
+				}
+			})
+			b.Run(fmt.Sprintf("has-%v", here.size), func(b *testing.B) {
+				b.ReportAllocs()
+				for i := 0; i < b.N; i++ {
+					randOperand().Has(randomPathMaker.makePath(here.minPathLen, here.maxPathLen))
+				}
+			})
+		*/
 		b.Run(fmt.Sprintf("serialize-%v", here.size), func(b *testing.B) {
 			b.ReportAllocs()
 			for i := 0; i < b.N; i++ {
@@ -116,36 +118,38 @@ func BenchmarkFieldSet(b *testing.B) {
 			}
 		})
 
-		b.Run(fmt.Sprintf("union-%v", here.size), func(b *testing.B) {
-			b.ReportAllocs()
-			for i := 0; i < b.N; i++ {
-				randOperand().Union(randOperand())
-			}
-		})
-		b.Run(fmt.Sprintf("intersection-%v", here.size), func(b *testing.B) {
-			b.ReportAllocs()
-			for i := 0; i < b.N; i++ {
-				randOperand().Intersection(randOperand())
-			}
-		})
-		b.Run(fmt.Sprintf("difference-%v", here.size), func(b *testing.B) {
-			b.ReportAllocs()
-			for i := 0; i < b.N; i++ {
-				randOperand().Difference(randOperand())
-			}
-		})
-		b.Run(fmt.Sprintf("recursive-difference-%v", here.size), func(b *testing.B) {
-			b.ReportAllocs()
-			for i := 0; i < b.N; i++ {
-				randOperand().RecursiveDifference(randOperand())
-			}
-		})
-		b.Run(fmt.Sprintf("leaves-%v", here.size), func(b *testing.B) {
-			b.ReportAllocs()
-			for i := 0; i < b.N; i++ {
-				randOperand().Leaves()
-			}
-		})
+		/*
+			b.Run(fmt.Sprintf("union-%v", here.size), func(b *testing.B) {
+				b.ReportAllocs()
+				for i := 0; i < b.N; i++ {
+					randOperand().Union(randOperand())
+				}
+			})
+			b.Run(fmt.Sprintf("intersection-%v", here.size), func(b *testing.B) {
+				b.ReportAllocs()
+				for i := 0; i < b.N; i++ {
+					randOperand().Intersection(randOperand())
+				}
+			})
+			b.Run(fmt.Sprintf("difference-%v", here.size), func(b *testing.B) {
+				b.ReportAllocs()
+				for i := 0; i < b.N; i++ {
+					randOperand().Difference(randOperand())
+				}
+			})
+			b.Run(fmt.Sprintf("recursive-difference-%v", here.size), func(b *testing.B) {
+				b.ReportAllocs()
+				for i := 0; i < b.N; i++ {
+					randOperand().RecursiveDifference(randOperand())
+				}
+			})
+			b.Run(fmt.Sprintf("leaves-%v", here.size), func(b *testing.B) {
+				b.ReportAllocs()
+				for i := 0; i < b.N; i++ {
+					randOperand().Leaves()
+				}
+			})
+		*/
 	}
 }
 

diff --git a/value/fields.go b/value/fields.go
@@ -17,6 +17,7 @@ limitations under the License.
 package value
 
 import (
+	"bytes"
 	gojson "encoding/json"
 	"sort"
 	"strings"
@@ -50,12 +51,23 @@ func FieldListFromJSON(input []byte) (FieldList, error) {
 }
 
 // FieldListToJSON is a helper function for producing a JSON document.
-func FieldListToJSON(v FieldList) ([]byte, error) {
-	m := make(map[string]interface{}, len(v))
-	for _, f := range v {
-		m[f.Name] = f.Value.Unstructured()
+func FieldListToJSON(v FieldList, w *bytes.Buffer) error {
+	w.WriteByte('{')
+	encoder := gojson.NewEncoder(noNewlineWriter{w})
+	for i, f := range v {
+		if err := encoder.Encode(f.Name); err != nil {
+			return err
+		}
+		w.WriteByte(':')
+		if err := encoder.Encode(f.Value.Unstructured()); err != nil {
+			return err
+		}
+		if i < len(v)-1 {
+			w.WriteByte(',')
+		}
 	}
-	return gojson.Marshal(m)
+	w.WriteByte('}')
+	return nil
 }
 
 // Sort sorts the field list by Name.