Repository: segmentio/encoding
Pull request #36 (2020-05-11 20:59:57)

Add SetTrustRawMessage Encoder method
## Description
This change adds a new `SetTrustRawMessage` method to the `Encoder` struct. If enabled, this causes the resulting the encoder to not check that `json.RawMessage` values are valid JSON. These checks are not needed in cases where the values are known to be valid JSON, e.g. because they were created from a `json.Marshal` into a `map[string]json.RawMessage`. 

Based on some profiling in the Segment replays system, these checks can get pretty expensive when marshalling out large, heavily nested objects.


Patch:
diff --git a/json/encode.go b/json/encode.go
index 75641bd..d955c7f 100644
--- a/json/encode.go
+++ b/json/encode.go
@@ -607,9 +607,16 @@ func (e encoder) encodeRawMessage(b []byte, p unsafe.Pointer) ([]byte, error) {
 		return append(b, "null"...), nil
 	}
 
-	s, _, err := parseValue(v)
-	if err != nil {
-		return b, &UnsupportedValueError{Value: reflect.ValueOf(v), Str: err.Error()}
+	var s []byte
+
+	if (e.flags & TrustRawMessage) != 0 {
+		s = v
+	} else {
+		var err error
+		s, _, err = parseValue(v)
+		if err != nil {
+			return b, &UnsupportedValueError{Value: reflect.ValueOf(v), Str: err.Error()}
+		}
 	}
 
 	if (e.flags & EscapeHTML) != 0 {
diff --git a/json/json.go b/json/json.go
index c9f8967..8d775e2 100644
--- a/json/json.go
+++ b/json/json.go
@@ -64,6 +64,11 @@ const (
 	// encoding JSON (this matches the behavior of the standard encoding/json
 	// package).
 	SortMapKeys
+
+	// TrustRawMessage is a performance optimization flag to skip value
+	// checking of raw messages. It should only be used if the values are
+	// known to be valid json (e.g., they were created by json.Unmarshal).
+	TrustRawMessage
 )
 
 // ParseFlags is a type used to represent configuration options that can be
@@ -437,6 +442,17 @@ func (enc *Encoder) SetSortMapKeys(on bool) {
 	}
 }
 
+// SetTrustRawMessage skips value checking when encoding a raw json message. It should only
+// be used if the values are known to be valid json, e.g. because they were originally created
+// by json.Unmarshal.
+func (enc *Encoder) SetTrustRawMessage(on bool) {
+	if on {
+		enc.flags |= TrustRawMessage
+	} else {
+		enc.flags &= ^TrustRawMessage
+	}
+}
+
 var encoderBufferPool = sync.Pool{
 	New: func() interface{} { return &encoderBuffer{data: make([]byte, 0, 4096)} },
 }


Test patch:
diff --git a/json/json_test.go b/json/json_test.go
index 8b07a7d..08a91a6 100644
--- a/json/json_test.go
+++ b/json/json_test.go
@@ -1535,3 +1535,48 @@ func TestGithubIssue28(t *testing.T) {
 	}
 
 }
+
+func TestSetTrustRawMessage(t *testing.T) {
+	buf := &bytes.Buffer{}
+	enc := NewEncoder(buf)
+	enc.SetTrustRawMessage(true)
+
+	// "Good" values are encoded in the regular way
+	m := map[string]json.RawMessage{
+		"k": json.RawMessage(`"value"`),
+	}
+	if err := enc.Encode(m); err != nil {
+		t.Error(err)
+	}
+
+	b := buf.Bytes()
+	exp := []byte(`{"k":"value"}`)
+	exp = append(exp, '\n')
+	if bytes.Compare(exp, b) != 0 {
+		t.Error(
+			"unexpected encoding:",
+			"expected", exp,
+			"got", b,
+		)
+	}
+
+	// "Bad" values are encoded without checking and throwing an error
+	buf.Reset()
+	m = map[string]json.RawMessage{
+		"k": json.RawMessage(`bad"value`),
+	}
+	if err := enc.Encode(m); err != nil {
+		t.Error(err)
+	}
+
+	b = buf.Bytes()
+	exp = []byte(`{"k":bad"value}`)
+	exp = append(exp, '\n')
+	if bytes.Compare(exp, b) != 0 {
+		t.Error(
+			"unexpected encoding:",
+			"expected", exp,
+			"got", b,
+		)
+	}
+}