reduce allocations by cheating

Author: lavalamp

fieldpath/serialize-pe.go

@@ -111,7 +111,11 @@ var (
 // SerializePathElement serializes a path element
 func SerializePathElement(pe PathElement) (string, error) {
 	buf := strings.Builder{}
-	stream := writePool.BorrowStream(&buf)
+	return serializePathElementUsingSB(&buf, pe)
+}
+
+func serializePathElementUsingSB(sb *strings.Builder, pe PathElement) (string, error) {
+	stream := writePool.BorrowStream(sb)
 	defer writePool.ReturnStream(stream)
 	switch {
 	case pe.FieldName != nil:
@@ -138,6 +142,12 @@ func SerializePathElement(pe PathElement) (string, error) {
 	default:
 		return "", errors.New("invalid PathElement")
 	}
+	b := stream.Buffer()
 	err := stream.Flush()
-	return buf.String(), err
+	// Help jsoniter manage its buffers--without this, the next
+	// use of the stream is likely to require an allocation. Look
+	// at the jsoniter stream code to understand why. They were probably
+	// optimizing for folks using the buffer directly.
+	stream.SetBuffer(b[:0])
+	return sb.String(), err
 }

fieldpath/serialize.go

@@ -19,22 +19,72 @@ package fieldpath
 import (
 	"bytes"
 	"io"
+	"strings"
+	"unsafe"
 
 	jsoniter "github.com/json-iterator/go"
 )
 
 func (s *Set) ToJSON() ([]byte, error) {
 	buf := bytes.Buffer{}
-	stream := jsoniter.NewStream(jsoniter.ConfigCompatibleWithStandardLibrary, &buf, 4096)
+	err := s.ToJSONStream(&buf)
+	if err != nil {
+		return nil, err
+	}
+	return buf.Bytes(), nil
+}
+
+func (s *Set) ToJSONStream(w io.Writer) error {
+	stream := writePool.BorrowStream(w)
+	defer writePool.ReturnStream(stream)
+
+	var r reusableBuilder
 
 	stream.WriteObjectStart()
-	s.emitContents_v1(false, stream)
+	err := s.emitContents_v1(false, stream, &r)
+	if err != nil {
+		return err
+	}
 	stream.WriteObjectEnd()
-	err := stream.Flush()
-	return buf.Bytes(), err
+	return stream.Flush()
+}
+
+func manageMemory(stream *jsoniter.Stream) error {
+	// Help jsoniter manage its buffers--without this, it does a bunch of
+	// alloctaions that are not necessary. They were probably optimizing
+	// for folks using the buffer directly.
+	b := stream.Buffer()
+	if len(b) > 4096 || cap(b)-len(b) < 2048 {
+		if err := stream.Flush(); err != nil {
+			return err
+		}
+		stream.SetBuffer(b[:0])
+	}
+	return nil
+}
+
+// reusableBuilder is a copy of strings.StringBuilder, which supports copy-free
+// []byte -> string transformation. This offers a reset() method which mutates
+// internal data, permitting us to reuse the same memory over and over. This
+// only makes sense if the string is going to be copied first;, otherwise we'd
+// be mutating the string, which is very bad.
+//
+// We could remove this completely if we modified jsoniter to permit writing a
+// string from a []byte.
+//
+// This MUST match the strings.Builder definition EXACTLY, otherwise reset()
+// will do something very bad to the heap.
+type reusableBuilder struct {
+	addr *strings.Builder
+	buf  []byte
+}
+
+func (r *reusableBuilder) reset() *strings.Builder {
+	r.buf = r.buf[:0]
+	return (*strings.Builder)(unsafe.Pointer(r))
 }
 
-func (s *Set) emitContents_v1(includeSelf bool, stream *jsoniter.Stream) {
+func (s *Set) emitContents_v1(includeSelf bool, stream *jsoniter.Stream, r *reusableBuilder) error {
 	mi, ci := 0, 0
 	first := true
 	preWrite := func() {
@@ -51,24 +101,28 @@ func (s *Set) emitContents_v1(includeSelf bool, stream *jsoniter.Stream) {
 
 		if mpe.Less(cpe) {
 			preWrite()
-			str, _ := SerializePathElement(mpe)
+			str, _ := serializePathElementUsingSB(r.reset(), mpe)
 			stream.WriteObjectField(str)
 			stream.WriteEmptyObject()
 			mi++
 		} else if cpe.Less(mpe) {
 			preWrite()
-			str, _ := SerializePathElement(cpe)
+			str, _ := serializePathElementUsingSB(r.reset(), cpe)
 			stream.WriteObjectField(str)
 			stream.WriteObjectStart()
-			s.Children.members[ci].set.emitContents_v1(false, stream)
+			if err := s.Children.members[ci].set.emitContents_v1(false, stream, r); err != nil {
+				return err
+			}
 			stream.WriteObjectEnd()
 			ci++
 		} else {
 			preWrite()
-			str, _ := SerializePathElement(cpe)
+			str, _ := serializePathElementUsingSB(r.reset(), cpe)
 			stream.WriteObjectField(str)
 			stream.WriteObjectStart()
-			s.Children.members[ci].set.emitContents_v1(true, stream)
+			if err := s.Children.members[ci].set.emitContents_v1(true, stream, r); err != nil {
+				return err
+			}
 			stream.WriteObjectEnd()
 			mi++
 			ci++
@@ -79,7 +133,7 @@ func (s *Set) emitContents_v1(includeSelf bool, stream *jsoniter.Stream) {
 		mpe := s.Members.members[mi]
 
 		preWrite()
-		str, _ := SerializePathElement(mpe)
+		str, _ := serializePathElementUsingSB(r.reset(), mpe)
 		stream.WriteObjectField(str)
 		stream.WriteEmptyObject()
 		mi++
@@ -89,10 +143,12 @@ func (s *Set) emitContents_v1(includeSelf bool, stream *jsoniter.Stream) {
 		cpe := s.Children.members[ci].pathElement
 
 		preWrite()
-		str, _ := SerializePathElement(cpe)
+		str, _ := serializePathElementUsingSB(r.reset(), cpe)
 		stream.WriteObjectField(str)
 		stream.WriteObjectStart()
-		s.Children.members[ci].set.emitContents_v1(false, stream)
+		if err := s.Children.members[ci].set.emitContents_v1(false, stream, r); err != nil {
+			return err
+		}
 		stream.WriteObjectEnd()
 		ci++
 	}
@@ -102,10 +158,12 @@ func (s *Set) emitContents_v1(includeSelf bool, stream *jsoniter.Stream) {
 		stream.WriteObjectField(".")
 		stream.WriteEmptyObject()
 	}
+	return manageMemory(stream)
 }
 
 // FromJSON clears s and reads a JSON formatted set structure.
 func (s *Set) FromJSON(r io.Reader) error {
+	// The iterator pool is completely useless for memory management, grrr.
 	iter := jsoniter.Parse(jsoniter.ConfigCompatibleWithStandardLibrary, r, 4096)
 
 	found, _ := readIter_v1(iter)