Continue to use json-iterator to improve marshalling of SampleHistogramPair

Signed-off-by: Jeanette Tan <[email protected]>

Author: zenador

model/value_histogram.go

@@ -25,6 +25,7 @@ import (
 
 func init() {
 	jsoniter.RegisterTypeEncoderFunc("model.HistogramBucket", marshalHistogramBucketJSON, marshalHistogramBucketJSONIsEmpty)
+	jsoniter.RegisterTypeEncoderFunc("model.SampleHistogramPair", marshalSampleHistogramPairJSON, marshalSampleHistogramPairJSONIsEmpty)
 }
 
 type FloatString float64
@@ -136,20 +137,25 @@ type SampleHistogramPair struct {
 	Histogram *SampleHistogram
 }
 
+// marshalSampleHistogramPairJSON writes `[ts, "val"]`.
+func marshalSampleHistogramPairJSON(ptr unsafe.Pointer, stream *jsoniter.Stream) {
+	p := *((*SampleHistogramPair)(ptr))
+	stream.WriteArrayStart()
+	MarshalTimestamp(int64(p.Timestamp), stream)
+	stream.WriteMore()
+	MarshalHistogram(*p.Histogram, stream)
+	stream.WriteArrayEnd()
+}
+
+func marshalSampleHistogramPairJSONIsEmpty(ptr unsafe.Pointer) bool {
+	return false
+}
+
 func (s SampleHistogramPair) MarshalJSON() ([]byte, error) {
-	jsoni := jsoniter.ConfigCompatibleWithStandardLibrary
-	t, err := jsoni.Marshal(s.Timestamp)
-	if err != nil {
-		return nil, err
-	}
 	if s.Histogram == nil {
 		return nil, fmt.Errorf("histogram is nil")
 	}
-	v, err := jsoni.Marshal(s.Histogram)
-	if err != nil {
-		return nil, err
-	}
-	return []byte(fmt.Sprintf("[%s,%s]", t, v)), nil
+	return jsoniter.ConfigCompatibleWithStandardLibrary.Marshal(s)
 }
 
 func (s *SampleHistogramPair) UnmarshalJSON(buf []byte) error {

model/value_marshal.go

@@ -76,3 +76,31 @@ func MarshalHistogramBucket(b HistogramBucket, stream *jsoniter.Stream) {
 	MarshalValue(b.Count, stream)
 	stream.WriteArrayEnd()
 }
+
+// adapted from https://github.com/prometheus/prometheus/blob/main/web/api/v1/api.go
+func MarshalHistogram(h SampleHistogram, stream *jsoniter.Stream) {
+	stream.WriteObjectStart()
+	stream.WriteObjectField(`count`)
+	MarshalValue(h.Count, stream)
+	stream.WriteMore()
+	stream.WriteObjectField(`sum`)
+	MarshalValue(h.Sum, stream)
+
+	bucketFound := false
+	for _, bucket := range h.Buckets {
+		if bucket.Count == 0 {
+			continue // No need to expose empty buckets in JSON.
+		}
+		stream.WriteMore()
+		if !bucketFound {
+			stream.WriteObjectField(`buckets`)
+			stream.WriteArrayStart()
+		}
+		bucketFound = true
+		MarshalHistogramBucket(*bucket, stream)
+	}
+	if bucketFound {
+		stream.WriteArrayEnd()
+	}
+	stream.WriteObjectEnd()
+}