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

Signed-off-by: Jeanette Tan <jeanette.tan@grafana.com>
---
 model/value_histogram.go | 26 ++++++++++++++++----------
 model/value_marshal.go   | 28 ++++++++++++++++++++++++++++
 2 files changed, 44 insertions(+), 10 deletions(-)

diff --git a/model/value_histogram.go b/model/value_histogram.go
index 5b9724f4..0d2eaf25 100644
--- a/model/value_histogram.go
+++ b/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 {
diff --git a/model/value_marshal.go b/model/value_marshal.go
index febc1399..c11d9a0f 100644
--- a/model/value_marshal.go
+++ b/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()
+}