From 4e9e3f5f46ca158dfd54797c673e30249405aaea Mon Sep 17 00:00:00 2001
From: Jeanette Tan <jeanette.tan@grafana.com>
Date: Fri, 27 Jan 2023 18:11:46 +0800
Subject: [PATCH] Continue to use json-iterator to improve marshalling of
 SamplePair

Signed-off-by: Jeanette Tan <jeanette.tan@grafana.com>
---
 model/value_float.go     | 28 ++++++++++++++++++----------
 model/value_histogram.go | 12 ++----------
 model/value_marshal.go   | 20 ++++++++++++--------
 3 files changed, 32 insertions(+), 28 deletions(-)

diff --git a/model/value_float.go b/model/value_float.go
index 9266fb98..8b59571a 100644
--- a/model/value_float.go
+++ b/model/value_float.go
@@ -18,8 +18,15 @@ import (
 	"fmt"
 	"math"
 	"strconv"
+	"unsafe"
+
+	jsoniter "github.com/json-iterator/go"
 )
 
+func init() {
+	jsoniter.RegisterTypeEncoderFunc("model.SamplePair", marshalSamplePairJSON, marshalJSONIsEmpty)
+}
+
 var (
 	// ZeroSamplePair is the pseudo zero-value of SamplePair used to signal a
 	// non-existing sample pair. It is a SamplePair with timestamp Earliest and
@@ -71,17 +78,18 @@ type SamplePair struct {
 	Value     SampleValue
 }
 
-// MarshalJSON implements json.Marshaler.
+// marshalSamplePairJSON writes `[ts, "val"]`.
+func marshalSamplePairJSON(ptr unsafe.Pointer, stream *jsoniter.Stream) {
+	p := *((*SamplePair)(ptr))
+	stream.WriteArrayStart()
+	MarshalTimestamp(int64(p.Timestamp), stream)
+	stream.WriteMore()
+	MarshalValue(float64(p.Value), stream)
+	stream.WriteArrayEnd()
+}
+
 func (s SamplePair) MarshalJSON() ([]byte, error) {
-	t, err := json.Marshal(s.Timestamp)
-	if err != nil {
-		return nil, err
-	}
-	v, err := json.Marshal(s.Value)
-	if err != nil {
-		return nil, err
-	}
-	return []byte(fmt.Sprintf("[%s,%s]", t, v)), nil
+	return jsoniter.ConfigCompatibleWithStandardLibrary.Marshal(s)
 }
 
 // UnmarshalJSON implements json.Unmarshaler.
diff --git a/model/value_histogram.go b/model/value_histogram.go
index 0d2eaf25..cc221a88 100644
--- a/model/value_histogram.go
+++ b/model/value_histogram.go
@@ -24,8 +24,8 @@ import (
 )
 
 func init() {
-	jsoniter.RegisterTypeEncoderFunc("model.HistogramBucket", marshalHistogramBucketJSON, marshalHistogramBucketJSONIsEmpty)
-	jsoniter.RegisterTypeEncoderFunc("model.SampleHistogramPair", marshalSampleHistogramPairJSON, marshalSampleHistogramPairJSONIsEmpty)
+	jsoniter.RegisterTypeEncoderFunc("model.HistogramBucket", marshalHistogramBucketJSON, marshalJSONIsEmpty)
+	jsoniter.RegisterTypeEncoderFunc("model.SampleHistogramPair", marshalSampleHistogramPairJSON, marshalJSONIsEmpty)
 }
 
 type FloatString float64
@@ -63,10 +63,6 @@ func marshalHistogramBucketJSON(ptr unsafe.Pointer, stream *jsoniter.Stream) {
 	MarshalHistogramBucket(b, stream)
 }
 
-func marshalHistogramBucketJSONIsEmpty(ptr unsafe.Pointer) bool {
-	return false
-}
-
 func (s *HistogramBucket) UnmarshalJSON(buf []byte) error {
 	tmp := []interface{}{&s.Boundaries, &s.Lower, &s.Upper, &s.Count}
 	wantLen := len(tmp)
@@ -147,10 +143,6 @@ func marshalSampleHistogramPairJSON(ptr unsafe.Pointer, stream *jsoniter.Stream)
 	stream.WriteArrayEnd()
 }
 
-func marshalSampleHistogramPairJSONIsEmpty(ptr unsafe.Pointer) bool {
-	return false
-}
-
 func (s SampleHistogramPair) MarshalJSON() ([]byte, error) {
 	if s.Histogram == nil {
 		return nil, fmt.Errorf("histogram is nil")
diff --git a/model/value_marshal.go b/model/value_marshal.go
index c11d9a0f..0c7669e3 100644
--- a/model/value_marshal.go
+++ b/model/value_marshal.go
@@ -16,10 +16,15 @@ package model
 import (
 	"math"
 	"strconv"
+	"unsafe"
 
 	jsoniter "github.com/json-iterator/go"
 )
 
+func marshalJSONIsEmpty(ptr unsafe.Pointer) bool {
+	return false
+}
+
 // from https://github.com/prometheus/prometheus/blob/main/util/jsonutil/marshal.go
 // MarshalTimestamp marshals a point timestamp using the passed jsoniter stream.
 func MarshalTimestamp(t int64, stream *jsoniter.Stream) {
@@ -43,10 +48,9 @@ func MarshalTimestamp(t int64, stream *jsoniter.Stream) {
 	}
 }
 
-// adapted from https://github.com/prometheus/prometheus/blob/main/util/jsonutil/marshal.go
+// from https://github.com/prometheus/prometheus/blob/main/util/jsonutil/marshal.go
 // MarshalValue marshals a point value using the passed jsoniter stream.
-func MarshalValue(f FloatString, stream *jsoniter.Stream) {
-	v := float64(f)
+func MarshalValue(v float64, stream *jsoniter.Stream) {
 	stream.WriteRaw(`"`)
 	// Taken from https://github.com/json-iterator/go/blob/master/stream_float.go#L71 as a workaround
 	// to https://github.com/json-iterator/go/issues/365 (jsoniter, to follow json standard, doesn't allow inf/nan).
@@ -69,11 +73,11 @@ func MarshalHistogramBucket(b HistogramBucket, stream *jsoniter.Stream) {
 	stream.WriteArrayStart()
 	stream.WriteInt32(b.Boundaries)
 	stream.WriteMore()
-	MarshalValue(b.Lower, stream)
+	MarshalValue(float64(b.Lower), stream)
 	stream.WriteMore()
-	MarshalValue(b.Upper, stream)
+	MarshalValue(float64(b.Upper), stream)
 	stream.WriteMore()
-	MarshalValue(b.Count, stream)
+	MarshalValue(float64(b.Count), stream)
 	stream.WriteArrayEnd()
 }
 
@@ -81,10 +85,10 @@ func MarshalHistogramBucket(b HistogramBucket, stream *jsoniter.Stream) {
 func MarshalHistogram(h SampleHistogram, stream *jsoniter.Stream) {
 	stream.WriteObjectStart()
 	stream.WriteObjectField(`count`)
-	MarshalValue(h.Count, stream)
+	MarshalValue(float64(h.Count), stream)
 	stream.WriteMore()
 	stream.WriteObjectField(`sum`)
-	MarshalValue(h.Sum, stream)
+	MarshalValue(float64(h.Sum), stream)
 
 	bucketFound := false
 	for _, bucket := range h.Buckets {